text:
Base.summary(df::AbstractDataFrame) = # -> String
    @sprintf("%d×%d %s", size(df)..., typeof(df).name)

let
    local buffer = IOBuffer(Vector{UInt8}(undef, 80), read=true, write=true)
    global ourstrwidth

    """
        DataFrames.ourstrwidth(io::IO, x::Any)

    Determine the number of characters that would be used to print a value.
    """
    function ourstrwidth(io::IO, x::Any) # -> Int
        truncate(buffer, 0)
        ourshow(IOContext(buffer, :compact=>get(io, :compact, true)), x)
        textwidth(String(take!(buffer)))
    end
end

"""
    DataFrames.ourshow(io::IO, x::Any)

Render a value to an IO object. Unlike
`show`, render strings without surrounding quote marks.
"""
ourshow(io::IO, x::Any) =
    show(IOContext(io, :compact=>get(io, :compact, true), :typeinfo=>typeof(x)), x)
ourshow(io::IO, x::AbstractString) = escape_string(io, x, "")
ourshow(io::IO, x::Symbol) = ourshow(io, string(x))
ourshow(io::IO, x::Nothing) = nothing

"""Return compact string representation of type T"""
function compacttype(T::Type, maxwidth::Int=8)
    T === Any && return "Any"
    T === Missing && return "Missing"
    sT = string(T)
    length(sT) ≤ maxwidth && return sT
    if T >: Missing
        T = nonmissingtype(T)
        sT = string(T)
        suffix = "⍰"
        # handle the case when after removing Missing union type name is short
        length(sT) ≤ 8 && return sT * suffix
    else
        suffix = ""
    end
    T <: Union{CategoricalString, CategoricalValue} && return "Categorical…"*suffix
    # we abbreviate consistently to total of 8 characters
    match(Regex("^.\\w{0,$(7-length(suffix))}"), sT).match * "…"*suffix
end

"""
    DataFrames.getmaxwidths(df::AbstractDataFrame,
                            io::IO,
                            rowindices1::AbstractVector{Int},
                            rowindices2::AbstractVector{Int},
                            rowlabel::Symbol)

Calculate, for each column of an AbstractDataFrame, the maximum
string width used to render the name of that column, its type, and the
longest entry in that column -- among the rows of the data frame
will be rendered to IO. The widths for all columns are returned as a
vector.

Return a `Vector{Int}` giving the maximum string widths required to render
each column, including that column's name and type.

NOTE: The last entry of the result vector is the string width of the
implicit row ID column contained in every `AbstractDataFrame`.

# Arguments
- `df::AbstractDataFrame`: The data frame whose columns will be printed.
- `io::IO`: The `IO` to which `df` is to be printed
- `rowindices1::AbstractVector{Int}: A set of indices of the first
  chunk of the AbstractDataFrame that would be rendered to IO.
- `rowindices2::AbstractVector{Int}: A set of indices of the second
  chunk of the AbstractDataFrame that would be rendered to IO. Can
  be empty if the AbstractDataFrame would be printed without any
  ellipses.
- `rowlabel::AbstractString`: The label that will be used when rendered the
  numeric ID's of each row. Typically, this will be set to "Row".

# Examples
```jldoctest
julia> using DataFrames

julia> df = DataFrame(A = 1:3, B = ["x", "yy", "z"]);

julia> DataFrames.getmaxwidths(df, stdout, 1:1, 3:3, :Row)
3-element Array{Int64,1}:
 1
 1
 3
```
"""
function getmaxwidths(df::AbstractDataFrame,
                      io::IO,
                      rowindices1::AbstractVector{Int},
                      rowindices2::AbstractVector{Int},
                      rowlabel::Symbol,
                      rowid=nothing) # -> Vector{Int}
    maxwidths = Vector{Int}(undef, size(df, 2) + 1)

    undefstrwidth = ourstrwidth(io, Base.undef_ref_str)

    j = 1
    for (name, col) in eachcol(df, true)
        # (1) Consider length of column name
        maxwidth = ourstrwidth(io, name)

        # (2) Consider length of longest entry in that column
        for indices in (rowindices1, rowindices2), i in indices
            if isassigned(col, i)
                maxwidth = max(maxwidth, ourstrwidth(io, col[i]))
            else
                maxwidth = max(maxwidth, undefstrwidth)
            end
        end
        maxwidths[j] = max(maxwidth, ourstrwidth(io, compacttype(eltype(col))))
        j += 1
    end

    if rowid isa Nothing
        rowmaxwidth1 = isempty(rowindices1) ? 0 : ndigits(maximum(rowindices1))
        rowmaxwidth2 = isempty(rowindices2) ? 0 : ndigits(maximum(rowindices2))
        maxwidths[j] = max(max(rowmaxwidth1, rowmaxwidth2), ourstrwidth(io, rowlabel))
    else
        maxwidths[j] = max(ndigits(rowid), ourstrwidth(io, rowlabel))
    end

    return maxwidths
end

"""
    DataFrames.getprintedwidth(maxwidths::Vector{Int})

Given the maximum widths required to render each column of an
`AbstractDataFrame`, return the total number of characters
that would be required to render an entire row to an I/O stream.

NOTE: This width includes the whitespace and special characters used to
pretty print the `AbstractDataFrame`.

# Arguments
- `maxwidths::Vector{Int}`: The maximum width needed to render each
  column of an `AbstractDataFrame`.

# Examples
```jldoctest
julia> using DataFrames

julia> df = DataFrame(A = 1:3, B = ["x", "yy", "z"]);

julia> maxwidths = DataFrames.getmaxwidths(df, stdout, 1:1, 3:3, :Row)
3-element Array{Int64,1}:
 1
 1
 3

julia> DataFrames.getprintedwidth(maxwidths)
15
```
"""
function getprintedwidth(maxwidths::Vector{Int}) # -> Int
    # Include length of line-initial |
    totalwidth = 1
    for i in 1:length(maxwidths)
        # Include length of field + 2 spaces + trailing |
        totalwidth += maxwidths[i] + 3
    end
    return totalwidth
end

"""
    getchunkbounds(maxwidths::Vector{Int},
                   splitcols::Bool,
                   availablewidth::Int)

When rendering an `AbstractDataFrame` to a REPL window in chunks, each of
which will fit within the width of the REPL window, this function will
return the indices of the columns that should be included in each chunk.

NOTE: The resulting bounds should be interpreted as follows: the
i-th chunk bound is the index MINUS 1 of the first column in the
i-th chunk. The (i + 1)-th chunk bound is the EXACT index of the
last column in the i-th chunk. For example, the bounds [0, 3, 5]
imply that the first chunk contains columns 1-3 and the second chunk
contains columns 4-5.

# Arguments
- `maxwidths::Vector{Int}`: The maximum width needed to render each
  column of an AbstractDataFrame.
- `splitcols::Bool`: Whether to split printing in chunks of columns
  fitting the screen width rather than printing all columns in the same block.
- `availablewidth::Int`: The available width in the REPL.

# Examples
```jldoctest
julia> using DataFrames

julia> df = DataFrame(A = 1:3, B = ["x", "yy", "z"]);

julia> maxwidths = DataFrames.getmaxwidths(df, stdout, 1:1, 3:3, :Row)
3-element Array{Int64,1}:
 1
 1
 3

julia> DataFrames.getchunkbounds(maxwidths, true, displaysize()[2])
2-element Array{Int64,1}:
 0
 2
```
"""
function getchunkbounds(maxwidths::Vector{Int},
                        splitcols::Bool,
                        availablewidth::Int) # -> Vector{Int}
    ncols = length(maxwidths) - 1
    rowmaxwidth = maxwidths[ncols + 1]
    if splitcols
        chunkbounds = [0]
        # Include 2 spaces + 2 | characters for row/col label
        totalwidth = rowmaxwidth + 4
        for j in 1:ncols
            # Include 2 spaces + | character in per-column character count
            totalwidth += maxwidths[j] + 3
            if totalwidth > availablewidth
                push!(chunkbounds, j - 1)
                totalwidth = rowmaxwidth + 4 + maxwidths[j] + 3
            end
        end
        push!(chunkbounds, ncols)
    else
        chunkbounds = [0, ncols]
    end
    return chunkbounds
end

"""
    showrowindices(io::IO,
                   df::AbstractDataFrame,
                   rowindices::AbstractVector{Int},
                   maxwidths::Vector{Int},
                   leftcol::Int,
                   rightcol::Int,
                   rowid::Union{Int,Nothing})

Render a subset of rows and columns of an `AbstractDataFrame` to an
I/O stream. For chunked printing, this function is used to print a
single chunk, starting from the first indicated column and ending with
the last indicated column. Assumes that the maximum string widths
required for printing have been precomputed.

# Arguments
- `io::IO`: The I/O stream to which `df` will be printed.
- `df::AbstractDataFrame`: An AbstractDataFrame.
- `rowindices::AbstractVector{Int}`: The indices of the subset of rows
  that will be rendered to `io`.
- `maxwidths::Vector{Int}`: The pre-computed maximum string width
  required to render each column.
- `leftcol::Int`: The index of the first column in a chunk to be rendered.
- `rightcol::Int`: The index of the last column in a chunk to be rendered.
- `rowid`: Used to handle showing `DataFrameRow`.

# Examples
```jldoctest
julia> using DataFrames

julia> df = DataFrame(A = 1:3, B = ["x", "y", "z"]);

julia> DataFrames.showrowindices(stdout, df, 1:2, [1, 1, 5], 1, 2)
│ 1     │ 1 │ x │
│ 2     │ 2 │ y │
```
"""
function showrowindices(io::IO,
                        df::AbstractDataFrame,
                        rowindices::AbstractVector{Int},
                        maxwidths::Vector{Int},
                        leftcol::Int,
                        rightcol::Int,
                        rowid) # -> Void
    rowmaxwidth = maxwidths[end]

    for i in rowindices
        # Print row ID
        if rowid isa Nothing
            @printf io "│ %d" i
        else
            @printf io "│ %d" rowid
        end
        padding = rowmaxwidth - ndigits(rowid isa Nothing ? i : rowid)
        for _ in 1:padding
            write(io, ' ')
        end
        print(io, " │ ")
        # Print DataFrame entry
        for j in leftcol:rightcol
            strlen = 0
            if isassigned(df[!, j], i)
                s = df[i, j]
                strlen = ourstrwidth(io, s)
                if ismissing(s)
                    printstyled(io, s, color=:light_black)
                elseif s === nothing
                    strlen = 0
                else
                    ourshow(io, s)
                end
            else
                strlen = ourstrwidth(io, Base.undef_ref_str)
                ourshow(io, Base.undef_ref_str)
            end
            padding = maxwidths[j] - strlen
            for _ in 1:padding
                write(io, ' ')
            end
            if j == rightcol
                if i == rowindices[end]
                    print(io, " │")
                else
                    print(io, " │\n")
                end
            else
                print(io, " │ ")
            end
        end
    end
    return
end

"""
    showrows(io::IO,
             df::AbstractDataFrame,
             rowindices1::AbstractVector{Int},
             rowindices2::AbstractVector{Int},
             maxwidths::Vector{Int},
             splitcols::Bool = false,
             allcols::Bool = false,
             rowlabel::Symbol = :Row,
             displaysummary::Bool = true,
             rowid=nothing)

Render a subset of rows (possibly in chunks) of an `AbstractDataFrame` to an
I/O stream.

NOTE: The value of `maxwidths[end]` must be the string width of
`rowlabel`.

# Arguments
- `io::IO`: The I/O stream to which `df` will be printed.
- `df::AbstractDataFrame`: An AbstractDataFrame.
- `rowindices1::AbstractVector{Int}`: The indices of the first subset
  of rows to be rendered.
- `rowindices2::AbstractVector{Int}`: The indices of the second subset
  of rows to be rendered. An ellipsis will be printed before
  rendering this second subset of rows.
- `maxwidths::Vector{Int}`: The pre-computed maximum string width
  required to render each column.
- `allcols::Bool = false`: Whether to print all columns, rather than
  a subset that fits the device width.
- `splitcols::Bool`: Whether to split printing in chunks of columns fitting the screen width
  rather than printing all columns in the same block.
- `rowlabel::Symbol`: What label should be printed when rendering the
  numeric ID's of each row? Defaults to `:Row`.
- `displaysummary::Bool`: Should a brief string summary of the
  AbstractDataFrame be rendered to the I/O stream before printing the
  contents of the renderable rows? Defaults to `true`.
- `rowid = nothing`: Used to handle showing `DataFrameRow`

# Examples
julia> using DataFrames

julia> df = DataFrame(A = 1:3, B = ["x", "y", "z"]);

julia> DataFrames.showrows(stdout, df, 1:2, 3:3, [5, 6, 3], false, true, :Row, true)
3×2 DataFrame
│ Row │ A     │ B      │
│     │ Int64 │ String │
├─────┼───────┼────────┤
│ 1   │ 1     │ x      │
│ 2   │ 2     │ y      │
⋮
│ 3   │ 3     │ z      │
```
"""
function showrows(io::IO,
                  df::AbstractDataFrame,
                  rowindices1::AbstractVector{Int},
                  rowindices2::AbstractVector{Int},
                  maxwidths::Vector{Int},
                  splitcols::Bool = false,
                  allcols::Bool = false,
                  rowlabel::Symbol = :Row,
                  displaysummary::Bool = true,
                  rowid=nothing) # -> Void
    ncols = size(df, 2)

    if isempty(rowindices1)
        if displaysummary
            println(io, summary(df))
        end
        return
    end

    rowmaxwidth = maxwidths[ncols + 1]
    chunkbounds = getchunkbounds(maxwidths, splitcols, displaysize(io)[2])
    nchunks = allcols ? length(chunkbounds) - 1 : min(length(chunkbounds) - 1, 1)

    header = displaysummary ? summary(df) : ""
    if !allcols && length(chunkbounds) > 2
        header *= ". Omitted printing of $(chunkbounds[end] - chunkbounds[2]) columns"
    end
    println(io, header)

    for chunkindex in 1:nchunks
        leftcol = chunkbounds[chunkindex] + 1
        rightcol = chunkbounds[chunkindex + 1]

        # Print column names
        @printf io "│ %s" rowlabel
        padding = rowmaxwidth - ourstrwidth(io, rowlabel)
        for itr in 1:padding
            write(io, ' ')
        end
        print(io, " │ ")
        for j in leftcol:rightcol
            s = _names(df)[j]
            ourshow(io, s)
            padding = maxwidths[j] - ourstrwidth(io, s)
            for itr in 1:padding
                write(io, ' ')
            end
            if j == rightcol
                print(io, " │\n")
            else
                print(io, " │ ")
            end
        end

        # Print column types
        print(io, "│ ")
        padding = rowmaxwidth
        for itr in 1:padding
            write(io, ' ')
        end
        print(io, " │ ")
        for j in leftcol:rightcol
            s = compacttype(eltype(df[!, j]), maxwidths[j])
            printstyled(io, s, color=:light_black)
            padding = maxwidths[j] - ourstrwidth(io, s)
            for itr in 1:padding
                write(io, ' ')
            end
            if j == rightcol
                print(io, " │\n")
            else
                print(io, " │ ")
            end
        end

        # Print table bounding line
        write(io, '├')
        for itr in 1:(rowmaxwidth + 2)
            write(io, '─')
        end
        write(io, '┼')
        for j in leftcol:rightcol
            for itr in 1:(maxwidths[j] + 2)
                write(io, '─')
            end
            if j < rightcol
                write(io, '┼')
            else
                write(io, '┤')
            end
        end
        write(io, '\n')

        # Print main table body, potentially in two abbreviated sections
        showrowindices(io,
                       df,
                       rowindices1,
                       maxwidths,
                       leftcol,
                       rightcol,
                       rowid)

        if !isempty(rowindices2)
            print(io, "\n⋮\n")
            showrowindices(io,
                           df,
                           rowindices2,
                           maxwidths,
                           leftcol,
                           rightcol,
                           rowid)
        end

        # Print newlines to separate chunks
        if chunkindex < nchunks
            print(io, "\n\n")
        end
    end

    return
end

function _show(io::IO,
               df::AbstractDataFrame;
               allrows::Bool = !get(io, :limit, false),
               allcols::Bool = !get(io, :limit, false),
               splitcols = get(io, :limit, false),
               rowlabel::Symbol = :Row,
               summary::Bool = true,
               rowid=nothing)
    _check_consistency(df)
    nrows = size(df, 1)
    if rowid !== nothing
        if size(df, 2) == 0
            rowid = nothing
        elseif nrows != 1
            throw(ArgumentError("rowid may be passed only with a single row data frame"))
        end
    end
    dsize = displaysize(io)
    availableheight = dsize[1] - 7
    nrowssubset = fld(availableheight, 2)
    bound = min(nrowssubset - 1, nrows)
    if allrows || nrows <= availableheight
        rowindices1 = 1:nrows
        rowindices2 = 1:0
    else
        rowindices1 = 1:bound
        rowindices2 = max(bound + 1, nrows - nrowssubset + 1):nrows
    end
    maxwidths = getmaxwidths(df, io, rowindices1, rowindices2, rowlabel, rowid)
    width = getprintedwidth(maxwidths)
    showrows(io,
             df,
             rowindices1,
             rowindices2,
             maxwidths,
             splitcols,
             allcols,
             rowlabel,
             summary,
             rowid)
    return
end

"""
    show([io::IO,] df::AbstractDataFrame;
         allrows::Bool = !get(io, :limit, false),
         allcols::Bool = !get(io, :limit, false),
         allgroups::Bool = !get(io, :limit, false),
         splitcols::Bool = get(io, :limit, false),
         rowlabel::Symbol = :Row,
         summary::Bool = true)

Render a data frame to an I/O stream. The specific visual
representation chosen depends on the width of the display.

If `io` is omitted, the result is printed to `stdout`,
and `allrows`, `allcols` and `allgroups` default to `false`
while `splitcols` defaults to `true`.

# Arguments
- `io::IO`: The I/O stream to which `df` will be printed.
- `df::AbstractDataFrame`: The data frame to print.
- `allrows::Bool `: Whether to print all rows, rather than
  a subset that fits the device height. By default this is the case only if
  `io` does not have the `IOContext` property `limit` set.
- `allcols::Bool`: Whether to print all columns, rather than
  a subset that fits the device width. By default this is the case only if
  `io` does not have the `IOContext` property `limit` set.
- `allgroups::Bool`: Whether to print all groups rather than
  the first and last, when `df` is a `GroupedDataFrame`.
  By default this is the case only if `io` does not have the `IOContext` property
  `limit` set.
- `splitcols::Bool`: Whether to split printing in chunks of columns fitting the screen width
  rather than printing all columns in the same block. Only applies if `allcols` is `true`.
  By default this is the case only if `io` has the `IOContext` property `limit` set.
- `rowlabel::Symbol = :Row`: The label to use for the column containing row numbers.
- `summary::Bool = true`: Whether to print a brief string summary of the data frame.

# Examples
```jldoctest
julia> using DataFrames

julia> df = DataFrame(A = 1:3, B = ["x", "y", "z"]);

julia> show(df, allcols=true)
3×2 DataFrame
│ Row │ A     │ B      │
│     │ Int64 │ String │
├─────┼───────┼────────┤
│ 1   │ 1     │ x      │
│ 2   │ 2     │ y      │
│ 3   │ 3     │ z      │
```
"""
Base.show(io::IO,
          df::AbstractDataFrame;
          allrows::Bool = !get(io, :limit, false),
          allcols::Bool = !get(io, :limit, false),
          splitcols = get(io, :limit, false),
          rowlabel::Symbol = :Row,
          summary::Bool = true) =
    _show(io, df, allrows=allrows, allcols=allcols, splitcols=splitcols,
          rowlabel=rowlabel, summary=summary)

Base.show(df::AbstractDataFrame;
          allrows::Bool = !get(stdout, :limit, true),
          allcols::Bool = !get(stdout, :limit, true),
          splitcols = get(stdout, :limit, true),
          rowlabel::Symbol = :Row,
          summary::Bool = true) =
    show(stdout, df,
         allrows=allrows, allcols=allcols, splitcols=splitcols,
         rowlabel=rowlabel, summary=summary)


text:
using MuZero

include("../diagnostics/tac/game.jl")

l1 = Learner(Tac)
l1.config.save_interval = 1
l1.config.checkpoint_interval = 1
selfplay!(l1; limit=10)
fit!(l1, steps=1)

batch_tensor = randtrainingbatch(l1) |> tensor

train_net_initial_out1a = initial_inference(l1.test, batch_tensor.images) |> NetworkOutput
train_net_initial_out1b = initial_inference(l1.test, batch_tensor.images) |> NetworkOutput
train_net_initial_out1a == train_net_initial_out1b

save!(l1)

l2 = Learner(Tac; uuid=l1.uuid)
train_net_initial_out2 = initial_inference(l2.test, batch_tensor.images) |> NetworkOutput
train_net_initial_out1a == train_net_initial_out2

test_net = l1.train
import BSON
using Flux
BSON.@save "test_net.bson" test_net

d = BSON.load("test_net.bson")
test_net_same = d[:test_net]

test_out1 = initial_inference(test_net, batch_tensor.images) |> NetworkOutput
test_out2 = initial_inference(test_net_same, batch_tensor.images) |> NetworkOutput

using Luxor
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
learner = Learner(TicTacToe)
selfplay!(learner; limit=10)
game = learner |> randgame
game |> show
game |> svg
batch = learner |> randtrainingbatch
batch |> show

function cell_fn(x,y)
    bg_color = (2,2) == (x,y) ? "#CFE2FC" : nothing
    opacity = (1,1) == (x,y) ? 0.5 : nothing
    tr = (3,3) == (x,y) ? (value="●", color="green", size=18) : "tr"
    (value="O", color="red", bl="bl", tr=tr, bg_color=bg_color, opacity=opacity)
end
svg1 = GridSVG(Grid(3, 3, cell_fn), bg_color="#fff", title="State 0", header="to play: 1", canvas_height=310, canvas_width=240)
#cell_fn_2(x,y) = (value="X", color="#333", bl="bl", tr="tr")
svg2 = GridSVG(Grid(3, 3, cell_fn_2), bg_color="#fff", title="UCB Policy", header="r: 0.12", br_header="v: 0.89", canvas_height=310, canvas_width=240)
svg3 = GridSVG(Grid(3, 3, cell_fn_2), bg_color="#fff", title="target", header="r: 0.12", br_header="v: 0.89", canvas_height=310, canvas_width=240)
svg4 = GridSVG(Grid(3, 3, cell_fn_2), bg_color="#fff", title="target", header="r: 0.12", br_header="v: 0.89", canvas_height=310, canvas_width=240)
mat = [[svg1,svg2,svg3]]
g = GridSVG(Grid(3, 1), cellsize=200, show_border=false, canvas_height=300, canvas_width=600)
#g((x,y)->"X")
#function test_fn(x,y)
#    GridSVG(Grid(2, 2), bg_color="#ccc", canvas_height=20, canvas_width=20)
#end
#g(test_fn)
results_fn(x,y) = mat[y][x]
g(results_fn)

using Base64
using Serialization

base64encode(serialize, rand(UInt128))

# Dataframes Test
using MuZero
import Zygote
using DataFrames
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
config = make_tictactoe_config()
config.batch_size = 10
config.num_simulations = 10
config.num_unroll_steps = 1

replay_buffer = ReplayBuffer(config)
selfplay!(config, storage, replay_buffer; limit=10)
b = batch(replay_buffer, config)
state(b, 1)

using Serialization, Base64
rand(1:2, 3, 3) |> (x)->base64encode(serialize, x)
deserialize(IOBuffer(base64decode(test)))

df = DataFrame(
    "Inf Type"=>["Initial","Recurrent"],
    "State"=>[
        base64encode(serialize, rand(UInt32)),
        base64encode(serialize, rand(UInt32))],
    "Action"=>[nothing, 3],
    ""=>["=>","=>"],
    "Reward"=>Float32.([0,0]),
    "Value"=>Float32.([-1,1]),
    "π[1]"=>Float32.([0.08, 0.12]),
    "π[2]"=>Float32.([0.12, 0.12]),
    "π[3]"=>Float32.([0.20, 0.00]),
    "π[4]"=>Float32.([0.08, 0.20]),
    "π[5]"=>Float32.([0.00, 0.00]),
    "π[6]"=>Float32.([0.08, 0.12]),
    "π[7]"=>Float32.([0.12, 0.20]),
    "π[8]"=>Float32.([0.12, 0.12]),
    "π[9]"=>Float32.([0.20, 0.12]),
)
show(df)

# Full Training Test (needs GPU for reasonable training time)using MuZero
using MuZero
import Zygote
using Grids
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
using Base.Threads: @spawn, @threads
config = make_tictactoe_config()
config.batch_size = 128
config.num_actors = 3
config.weight_decay = 1e-4
config.lr_init=0.002f0
config.save_interval = 10000
config.num_simulations = 100
config.num_unroll_steps = 1

replay_buffer = ReplayBuffer(config)
losses, dynamics_stats = [], []
for i = 1:(config.num_actors)
    @spawn selfplay!(config, storage, replay_buffer)
end
policy_perfs = []; calc_policy_perf() = push!(policy_perfs, policy_perf())
value_perfs = [] ; calc_value_perf() = push!(value_perfs, value_perf())
reward_perfs = [] ; calc_reward_perf() = push!(reward_perfs, reward_perf())
cb = [(500, 1, [calc_reward_perf,calc_value_perf,calc_policy_perf]),]
@spawn fit!(config, storage, replay_buffer, losses=losses, callbacks=cb)

# Half Training Test (needs GPU for reasonable training time)
using MuZero
import Zygote
using Grids
using Plots
using Statistics
using Base.Threads: @spawn, @threads
include("../examples/tictactoe/state.jl")
include("../examples/tictactoe/game.jl")
config = make_tictactoe_config()
config.batch_size = 10
config.num_actors = 3
config.training_steps = 20
config.weight_decay = 1e-3
config.lr_init=0.002f0
config.save_interval = 10000
config.num_simulations = 10
config.num_unroll_steps = 1
config.training_steps = 12000

replay_buffer = ReplayBuffer(config)
losses, dynamics_stats = [], []
selfplay!(config, storage, replay_buffer; limit=32)
policy_perfs = []; calc_policy_perf() = push!(policy_perfs, policy_perf())
value_perfs = [] ; calc_value_perf() = push!(value_perfs, value_perf())
reward_perfs = [] ; calc_reward_perf() = push!(reward_perfs, reward_perf())

avgx(x,size) = mean.(collect(Iterators.partition(x, size)))

function plot_losses()
    plot(avgx(losses, 100), fmt = :png)
    savefig("losses-copy.png")
    cp("losses-copy.png", "losses.png", force=true)
    rm("losses-copy.png")
end

cb = [(100, 1, [plot_losses])]
res = @spawn fit!(config, storage, replay_buffer, losses=losses, callbacks=cb)
fetch(res)

using Grids

function cell_fn(x,y)
    "O"
end

s = Grid(3,3) |> Grids.svg
s(GridCellSVG(3, 3, "X"; val_bl="1"))
s(cell_fn)

g |> get_state |> svg

Learner(config)
Learner(config; uuid="20B4993D-D1B8-41B7-ABB2-9BC8C70EA5F9")
Learner(config; dir="~/models/MuZeroModels.jl/tictactoe/")
Learner(
    config;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="~/models/MuZeroModels.jl/tictactoe/"
)
Learner(;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de"
)
Learner(;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="~/models/MuZeroModels.jl/tictactoe/"
)

#
# Pre-existing path: No
# Pre-existing config: No
#
Learner(
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    dir="~/models/MuZeroModels.jl/tictactoe/"
)

learner = Learner(
    Tac;
    uuid="8bcecce1-c870-4e34-8ae9-df6433d033de",
    config=config,
    schedule=tacschedule!
)

ext: jl


text:
@testset "BLAS" begin

import Base.BLAS: dot, nrm2, asum, gemm, gemv, symm, symv, trmm, trmv, trsm, trsv

################################## Level 1 ##################################
let P = 10, Q = 6, rng = MersenneTwister(123456), N = 10

    # Utility random generators.
    sc, vP, vQ = ()->randn(rng), ()->randn(rng, P), ()->randn(rng, Q)

    # Unit-stride dot.
    @test check_errs(N, binary_ȲD(dot, 1, vP)..., sc, vP, vP)
    @test check_errs(N, binary_ȲD(dot, 2, vP)..., sc, vP, vP)
    @test check_errs(N, binary_ȲD_inplace(dot, 1, vP, vP())..., sc, vP, vP)
    @test check_errs(N, binary_ȲD_inplace(dot, 2, vP, vP())..., sc, vP, vP)

    # Strided dot.
    _x, _y = vP(), vQ()
    _dot2, _dot4 = x->dot(5, x, 2, _y, 1), y->dot(5, _x, 2, y, 1)
    _∇dot2 = (z, z̄, x)->∇(dot, Val{2}, (), z, z̄, 5, x, 2, _y, 1)
    _∇dot4 = (z, z̄, y)->∇(dot, Val{4}, (), z, z̄, 5, _x, 2, y, 1)
    @test check_errs(N, _dot2, _∇dot2, sc, vP, vP)
    @test check_errs(N, _dot4, _∇dot4, sc, vQ, vQ)

    # In-place strided dot.
    _δx, _δy = vP(), vQ()
    _∇dot2 = (z, z̄, x)->∇(copy(_δx), dot, Val{2}, (), z, z̄, 5, x, 2, _y, 1) - _δx
    _∇dot4 = (z, z̄, y)->∇(copy(_δy), dot, Val{4}, (), z, z̄, 5, _x, 2, y, 1) - _δy
    @test check_errs(N, _dot2, _∇dot2, sc, vP, vP)
    @test check_errs(N, _dot4, _∇dot4, sc, vQ, vQ)

    # Unit-stride nrm2.
    @test check_errs(N, unary_ȲD(nrm2)..., sc, vP, vP)
    @test check_errs(N, unary_ȲD_inplace(nrm2, vP())..., sc, vP, vP)

    # Arbitrary-stride nrm2.
    _nrm2 = x->nrm2(5, x, 2)
    _∇nrm2 = (y, ȳ, x)->∇(nrm2, Val{2}, (), y, ȳ, 5, x, 2)
    _∇nrm2_in_place = (y, ȳ, x)->∇(copy(_δx), nrm2, Val{2}, (), y, ȳ, 5, x, 2) - _δx
    @test check_errs(N, _nrm2, _∇nrm2, sc, vP, vP)
    @test check_errs(N, _nrm2, _∇nrm2_in_place, sc, vP, vP)

    # Unit-stride `asum`.
    @test check_errs(N, unary_ȲD(asum)..., sc, vP, vP)
    @test check_errs(N, unary_ȲD_inplace(asum, vP())..., sc, vP, vP)

    # Arbitrary-stride `asum`.
    _asum = x->asum(5, x, 2)
    _∇asum = (y, ȳ, x)->∇(asum, Val{2}, (), y, ȳ, 5, x, 2)
    _∇asum_in_place = (y, ȳ, x)->∇(copy(_δx), asum, Val{2}, (), y, ȳ, 5, x, 2) - _δx
    @test check_errs(N, _asum, _∇asum, sc, vP, vP)
    @test check_errs(N, _asum, _∇asum_in_place, sc, vP, vP)
end

################################## Level 2 ##################################
let P = 10, Q = 6, rng = MersenneTwister(123456), N = 10

    # Utility random generators.
    sc, vP, vQ = ()->randn(rng), ()->randn(rng, P), ()->randn(rng, Q)
    mPQ, mQP = ()->randn(rng, P, Q), ()->randn(rng, Q, P)
    mPP = ()->randn(rng, P, P)

    # gemv:
    for tA in ['T', 'N']
        _α, _x = sc(), vQ()
        _A = tA == 'N' ? mPQ() : mQP()
        A = tA == 'N' ? mPQ : mQP
        _δA, _δx = randn!(rng, similar(_A)), randn!(rng, similar(_x))

        # α != 1 tests.
        _gemv1 = α->gemv(tA, α, _A, _x)
        _gemv2 = A->gemv(tA, _α, A, _x)
        _gemv3 = x->gemv(tA, _α, _A, x)
        _∇gemv1 = (y, ȳ, α)->∇(gemv, Val{2}, (), y, ȳ, tA, α, _A, _x)
        _∇gemv2 = (y, ȳ, A)->∇(gemv, Val{3}, (), y, ȳ, tA, _α, A, _x)
        _∇gemv2_inp = (y, ȳ, A)->∇(copy(_δA), gemv, Val{3}, (), y, ȳ, tA, _α, A, _x) - _δA
        _∇gemv3 = (y, ȳ, x)->∇(gemv, Val{4}, (), y, ȳ, tA, _α, _A, x)
        _∇gemv3_inp = (y, ȳ, x)->∇(copy(_δx), gemv, Val{4}, (), y, ȳ, tA, _α, _A, x) - _δx

        @test check_errs(N, _gemv1, _∇gemv1, vP, sc, sc)
        @test check_errs(N, _gemv2, _∇gemv2, vP, A, A)
        @test check_errs(N, _gemv2, _∇gemv2_inp, vP, A, A)
        @test check_errs(N, _gemv3, _∇gemv3, vP, vQ, vQ)
        @test check_errs(N, _gemv3, _∇gemv3_inp, vP, vQ, vQ)

        # α = 1 tests.
        _gemv1, _gemv2 = A->gemv(tA, A, _x), x->gemv(tA, _A, x)
        _∇gemv1 = (y, ȳ, A)->∇(gemv, Val{2}, (), y, ȳ, tA, A, _x)
        _∇gemv1_inp = (y, ȳ, A)->∇(copy(_δA), gemv, Val{2}, (), y, ȳ, tA, A, _x) - _δA
        _∇gemv2 = (y, ȳ, x)->∇(gemv, Val{3}, (), y, ȳ, tA, _A, x)
        _∇gemv2_inp = (y, ȳ, x)->∇(copy(_δx), gemv, Val{3}, (), y, ȳ, tA, _A, x) - _δx

        @test check_errs(N, _gemv1, _∇gemv1, vP, A, A)
        @test check_errs(N, _gemv1, _∇gemv1_inp, vP, A, A)
        @test check_errs(N, _gemv2, _∇gemv2, vP, vQ, vQ)
        @test check_errs(N, _gemv2, _∇gemv2_inp, vP, vQ, vQ)
    end

    # symv:
    for ul in ['L', 'U']
        α_gen, A_gen, x_gen = sc, mPP, vP
        _α, _A, _x = α_gen(), A_gen(), x_gen()
        _δA, _δx = randn!(rng, similar(_A)), randn!(rng, similar(_x))

        # α != 1 tests.
        _symv_α = α->symv(ul, α, _A, _x)
        _symv_A = A->symv(ul, _α, A, _x)
        _symv_x = x->symv(ul, _α, _A, x)
        _∇symv_α = (y, ȳ, α)->∇(symv, Val{2}, (), y, ȳ, ul, α, _A, _x)
        _∇symv_A = (y, ȳ, A)->∇(symv, Val{3}, (), y, ȳ, ul, _α, A, _x)
        _∇symv_A_inp = (y, ȳ, A)->∇(copy(_δA), symv, Val{3}, (), y, ȳ, ul, _α, A, _x) - _δA
        _∇symv_x = (y, ȳ, x)->∇(symv, Val{4}, (), y, ȳ, ul, _α, _A, x)
        _∇symv_x_inp = (y, ȳ, x)->∇(copy(_δx), symv, Val{4}, (), y, ȳ, ul, _α, _A, x) - _δx

        @test check_errs(N, _symv_α, _∇symv_α, vP, sc, sc)
        @test check_errs(N, _symv_A, _∇symv_A, vP, A_gen, A_gen)
        @test check_errs(N, _symv_A, _∇symv_A_inp, vP, A_gen, A_gen)
        @test check_errs(N, _symv_x, _∇symv_x, vP, x_gen, x_gen)
        @test check_errs(N, _symv_x, _∇symv_x_inp, vP, x_gen, x_gen)

        # α = 1 tests.
        _symv_A = A->symv(ul, A, _x)
        _symv_x = x->symv(ul, _A, x)
        _∇symv_A = (y, ȳ, A)->∇(symv, Val{2}, (), y, ȳ, ul, A, _x)
        _∇symv_A_inp = (y, ȳ, A)->∇(copy(_δA), symv, Val{2}, (), y, ȳ, ul, A, _x) - _δA
        _∇symv_x = (y, ȳ, x)->∇(symv, Val{3}, (), y, ȳ, ul, _A, x)
        _∇symv_x_inp = (y, ȳ, x)->∇(copy(_δx), symv, Val{3}, (), y, ȳ, ul, _A, x) - _δx

        @test check_errs(N, _symv_A, _∇symv_A, vP, A_gen, A_gen)
        @test check_errs(N, _symv_A, _∇symv_A_inp, vP, A_gen, A_gen)
        @test check_errs(N, _symv_x, _∇symv_x, vP, x_gen, x_gen)
        @test check_errs(N, _symv_x, _∇symv_x_inp, vP, x_gen, x_gen)
    end

    for f in [trmv, trsv], ul in ['L', 'U'], tA in ['N', 'T'], dA in ['U', 'N']
        A_gen, x_gen = mPP, vP
        _A, _x = A_gen(), x_gen()

        _f_A = A->f(ul, tA, dA, A, _x)
        _f_x = x->f(ul, tA, dA, _A, x)
        _∇f_A = (y, ȳ, A)->∇(f, Val{4}, (), y, ȳ, ul, tA, dA, A, _x)
        _∇f_x = (y, ȳ, x)->∇(f, Val{5}, (), y, ȳ, ul, tA, dA, _A, x)

        @test check_errs(N, _f_A, _∇f_A, vP, A_gen, A_gen)
        @test check_errs(N, _f_x, _∇f_x, vP, x_gen, x_gen)
    end
end

################################## Level 3 ##################################
let P = 5, Q = 3, rng = MersenneTwister(123456), N = 10

    # Utility random generators.
    sc, mPP, mQQ = ()->randn(rng), ()->randn(rng, P, P), ()->randn(rng, Q, Q)
    mPQ, mQP = ()->randn(rng, P, Q), ()->randn(rng, Q, P)

    # gemm:
    for tA in ['T', 'N'], tB in ['T', 'N']

        # Generate conformal test matrices.
        A_gen = mPQ
        B_gen = tA == 'N' ?
                    (tB == 'N' ? mQP : mPQ) :
                    (tB == 'T' ? mQP : mPQ)
        C_gen = tA == 'N' ? mPP : mQQ

        # α != 1 tests.
        _α, _A, _B = sc(), A_gen(), B_gen()
        _δA, _δB = randn!(rng, similar(_A)), randn!(rng, similar(_B))
        _gemm_α = α->gemm(tA, tB, α, _A, _B)
        _gemm_A = A->gemm(tA, tB, _α, A, _B)
        _gemm_B = B->gemm(tA, tB, _α, _A, B)
        _∇gemm_α = (Y, Ȳ, α)->∇(gemm, Val{3}, (), Y, Ȳ, tA, tB, α, _A, _B)
        _∇gemm_A = (Y, Ȳ, A)->∇(gemm, Val{4}, (), Y, Ȳ, tA, tB, _α, A, _B)
        _∇gemm_B = (Y, Ȳ, B)->∇(gemm, Val{5}, (), Y, Ȳ, tA, tB, _α, _A, B)
        _∇gemm_A_inp = (Y, Ȳ, A)->∇(copy(_δA), gemm, Val{4}, (), Y, Ȳ, tA, tB, _α, A, _B) - _δA
        _∇gemm_B_inp = (Y, Ȳ, B)->∇(copy(_δB), gemm, Val{5}, (), Y, Ȳ, tA, tB, _α, _A, B) - _δB

        @test check_errs(N, _gemm_α, _∇gemm_α, C_gen, sc, sc)
        @test check_errs(N, _gemm_A, _∇gemm_A, C_gen, A_gen, A_gen)
        @test check_errs(N, _gemm_A, _∇gemm_A_inp, C_gen, A_gen, A_gen)
        @test check_errs(N, _gemm_B, _∇gemm_B, C_gen, B_gen, B_gen)
        @test check_errs(N, _gemm_B, _∇gemm_B_inp, C_gen, B_gen, B_gen)

        # α = 1 tests.
        _A, _B = A_gen(), B_gen()
        _gemm_A = A->gemm(tA, tB, A, _B)
        _gemm_B = B->gemm(tA, tB, _A, B)
        _∇gemm_A = (Y, Ȳ, A)->∇(gemm, Val{3}, (), Y, Ȳ, tA, tB, A, _B)
        _∇gemm_B = (Y, Ȳ, B)->∇(gemm, Val{4}, (), Y, Ȳ, tA, tB, _A, B)
        _∇gemm_A_inp = (Y, Ȳ, A)->∇(copy(_δA), gemm, Val{3}, (), Y, Ȳ, tA, tB, A, _B) - _δA
        _∇gemm_B_inp = (Y, Ȳ, B)->∇(copy(_δB), gemm, Val{4}, (), Y, Ȳ, tA, tB, _A, B) - _δB

        @test check_errs(N, _gemm_A, _∇gemm_A, C_gen, A_gen, A_gen)
        @test check_errs(N, _gemm_A, _∇gemm_A_inp, C_gen, A_gen, A_gen)
        @test check_errs(N, _gemm_B, _∇gemm_B, C_gen, B_gen, B_gen)
        @test check_errs(N, _gemm_B, _∇gemm_B_inp, C_gen, B_gen, B_gen)
    end

    # symm:
    for side in ['L', 'R'], ul in ['L', 'U']

        # Fixed qtts.
        _α, _A, _B = sc(), mPP(), side == 'L' ? mPQ() : mQP()
        _δA, _δB = randn!(rng, similar(_A)), randn!(rng, similar(_B))

        # α != 1 tests.
        _symm_α = α->symm(side, ul, α, _A, _B)
        _symm_A = A->symm(side, ul, _α, A, _B)
        _symm_B = B->symm(side, ul, _α, _A, B)
        _∇symm_α = (Y, Ȳ, α)->∇(symm, Val{3}, (), Y, Ȳ, side, ul, α, _A, _B)
        _∇symm_A = (Y, Ȳ, A)->∇(symm, Val{4}, (), Y, Ȳ, side, ul, _α, A, _B)
        _∇symm_B = (Y, Ȳ, B)->∇(symm, Val{5}, (), Y, Ȳ, side, ul, _α, _A, B)
        _∇symm_A_inp = (Y, Ȳ, A)->∇(copy(_δA), symm, Val{4}, (), Y, Ȳ, side, ul, _α, A, _B) - _δA
        _∇symm_B_inp = (Y, Ȳ, B)->∇(copy(_δB), symm, Val{5}, (), Y, Ȳ, side, ul, _α, _A, B) - _δB

        B_gen = side == 'L' ? mPQ : mQP
        @test check_errs(N, _symm_α, _∇symm_α, B_gen, sc, sc)
        @test check_errs(N, _symm_A, _∇symm_A, B_gen, mPP, mPP)
        @test check_errs(N, _symm_A, _∇symm_A_inp, B_gen, mPP, mPP)
        @test check_errs(N, _symm_B, _∇symm_B, B_gen, B_gen, B_gen)
        @test check_errs(N, _symm_B, _∇symm_B_inp, B_gen, B_gen, B_gen)

        # α = 1 tests.
        _symm_A = A->symm(side, ul, A, _B)
        _symm_B = B->symm(side, ul, _A, B)
        _∇symm_A = (Y, Ȳ, A)->∇(symm, Val{3}, (), Y, Ȳ, side, ul, A, _B)
        _∇symm_B = (Y, Ȳ, B)->∇(symm, Val{4}, (), Y, Ȳ, side, ul, _A, B)
        _∇symm_A_inp = (Y, Ȳ, A)->∇(copy(_δA), symm, Val{3}, (), Y, Ȳ, side, ul, A, _B) - _δA
        _∇symm_B_inp = (Y, Ȳ, B)->∇(copy(_δB), symm, Val{4}, (), Y, Ȳ, side, ul, _A, B) - _δB

        B_gen = side == 'L' ? mPQ : mQP
        @test check_errs(N, _symm_A, _∇symm_A, B_gen, mPP, mPP)
        @test check_errs(N, _symm_A, _∇symm_A_inp, B_gen, mPP, mPP)
        @test check_errs(N, _symm_B, _∇symm_B, B_gen, B_gen, B_gen)
        @test check_errs(N, _symm_B, _∇symm_B_inp, B_gen, B_gen, B_gen)
    end

    # trmm / trsm:
    for f in [trmm, trsm], side in ['L', 'R'], ul in ['L', 'U'], tA in ['N', 'T'], dA in ['U', 'N']
        A_gen = ()->mPP() + 1e-3I
        B_gen = side == 'L' ? mPQ : mQP
        C_gen = B_gen

        _α, _A, _B = sc(), A_gen(), B_gen()
        _trmm_α = α->f(side, ul, tA, dA, α, _A, _B)
        _trmm_A = A->f(side, ul, tA, dA, _α, A, _B)
        _trmm_B = B->f(side, ul, tA, dA, _α, _A, B)
        _∇trmm_α = (Y, Ȳ, α)->∇(f, Val{5}, (), Y, Ȳ, side, ul, tA, dA, α, _A, _B)
        _∇trmm_A = (Y, Ȳ, A)->∇(f, Val{6}, (), Y, Ȳ, side, ul, tA, dA, _α, A, _B)
        _∇trmm_B = (Y, Ȳ, B)->∇(f, Val{7}, (), Y, Ȳ, side, ul, tA, dA, _α, _A, B)

        @test check_errs(N, _trmm_α, _∇trmm_α, C_gen, sc, sc)
        @test check_errs(N, _trmm_A, _∇trmm_A, C_gen, A_gen, A_gen)
        @test check_errs(N, _trmm_B, _∇trmm_B, C_gen, B_gen, B_gen)
    end
end
end


text:
using DirectionFinding
using Test

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


text:
using AbInitioTopo

nstart = 7
nend = 50
myFile = "irvsp_sr3pbo.out"
compute_Z4(myFile, nstart, nend)
compute_Z2(myFile, nstart, nend)


text:
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %%
module Foo

const a = Float64[]

function prepare_workspace(m::Integer...)
    resize!(a, prod(m))
    reshape(a, m)
end

prepare_workspace(x::AbstractArray) = prepare_workspace(size(x)...)

function f(x)
    ws = prepare_workspace(x)
    ws .= 2x
end

end

@show Foo.f([1 2; 3 4])
@show Foo.a;

# %%

max_stars_repo_path: 0015/reshape example.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_path: 0015/reshape example.jl

max_issues_repo_name: genkuroki/public

max_forks_repo_path: 0015/reshape example.jl

max_forks_repo_name: genkuroki/public


text:
change_vector_eltype(S0::Type{<:Vector}, T) = S0.name.wrapper{T, 1}

convert_mat(M::Union{SparseMatrixCOO, SparseMatrixCSC}, T) =
  convert(typeof(M).name.wrapper{T, Int}, M)
convert_mat(M::Matrix, T) = convert(Matrix{T}, M)

function push_history_residuals!(
  res::ResidualsHistory{T},
  itd::IterData{T},
  pad::PreallocatedData{T},
  id::QM_IntData,
) where {T <: Real}
  push!(res.rbNormH, res.rbNorm)
  push!(res.rcNormH, res.rcNorm)
  push!(res.pddH, itd.pdd)
  push!(res.μH, itd.μ)

  bound_dist = zero(T)
  if id.nlow > 0 && id.nupp > 0
    bound_dist = min(minimum(itd.x_m_lvar), minimum(itd.uvar_m_x))
  elseif id.nlow > 0 && id.nupp == 0
    bound_dist = min(minimum(itd.x_m_lvar))
  elseif id.nlow == 0 && id.nupp > 0
    bound_dist = min(minimum(itd.uvar_m_x))
  end
  (id.nlow > 0 || id.nupp > 0) && push!(res.min_bound_distH, bound_dist)

  pad_type = typeof(pad)
  if pad_type <: PreallocatedDataAugmentedKrylov ||
     pad_type <: PreallocatedDataNewtonKrylov ||
     pad_type <: PreallocatedDataAugmentedStructured ||
     pad_type <: PreallocatedDataNewtonStructured
    push!(res.kiterH, niterations(pad.KS))
    push!(res.KresNormH, norm(res.Kres))
    push!(res.KresPNormH, @views norm(res.Kres[(id.nvar + 1):(id.nvar + id.ncon)]))
    push!(res.KresDNormH, @views norm(res.Kres[1:(id.nvar)]))
  elseif pad_type <: PreallocatedDataNormalKrylov || pad_type <: PreallocatedDataNormalStructured
    push!(res.kiterH, niterations(pad.KS))
    push!(res.KresNormH, norm(res.Kres))
  end
end

function get_kiter(pad::PreallocatedData)
  padT = typeof(pad)
  nprod =
    (
      padT <: PreallocatedDataNewtonKrylov ||
      padT <: PreallocatedDataNewtonStructured ||
      padT <: PreallocatedDataAugmentedKrylov ||
      padT <: PreallocatedDataAugmentedStructured ||
      padT <: PreallocatedDataNormalKrylov ||
      padT <: PreallocatedDataNormalStructured
    ) ? niterations(pad.KS) : zero(Int)
  return nprod
end

struct IntDataInit{I <: Integer}
  nvar::I
  ncon::I
  ilow::Vector{I}
  iupp::Vector{I}
  irng::Vector{I}
  ifix::Vector{I}
  jlow::Vector{I}
  jupp::Vector{I}
  jrng::Vector{I}
  jfix::Vector{I}
end

function get_multipliers(
  s_l::AbstractVector{T},
  s_u::AbstractVector{T},
  ilow::Vector{Int},
  iupp::Vector{Int},
  nvar::Int,
  y::AbstractVector{T},
  idi::IntDataInit{Int},
  nvarrm::Int,
) where {T <: Real}
  nlow, nupp, nrng = length(idi.ilow), length(idi.iupp), length(idi.irng)
  njlow, njupp, njrng = length(idi.jlow), length(idi.jupp), length(idi.jrng)

  S = typeof(y)
  if S <: Vector
    s_l_sp = SparseVector(nvar + nvarrm, ilow, s_l)
    s_u_sp = SparseVector(nvar + nvarrm, iupp, s_u)
  else
    s_l_sp, s_u_sp = fill!(S(undef, nvar), zero(T)), fill!(S(undef, nvar), zero(T))
    s_l_sp[ilow] .= s_l
    s_u_sp[iupp] .= s_u
  end

  if idi.nvar != nvar
    multipliers_L = s_l_sp[1:(idi.nvar)]
    multipliers_U = s_u_sp[1:(idi.nvar)]
  else
    multipliers_L = s_l_sp
    multipliers_U = s_u_sp
  end

  multipliers = fill!(S(undef, idi.ncon), zero(T))
  multipliers[idi.jfix] .= @views y[idi.jfix]
  multipliers[idi.jlow] .+= @views s_l[(nlow + nrng + 1):(nlow + nrng + njlow)]
  multipliers[idi.jupp] .-= @views s_u[(nupp + nrng + 1):(nupp + nrng + njupp)]
  multipliers[idi.jrng] .+=
    @views s_l[(nlow + nrng + njlow + 1):end] .- s_u[(nupp + nrng + njupp + 1):end]

  return multipliers, multipliers_L, multipliers_U
end


text:
using Documenter, UnROOT

makedocs(;
    modules=[UnROOT],
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true",
        assets=String[],
    ),
    pages=[
        "Introduction" => "index.md",
        "Quick Start" => "quickstart.md",
        "Internals" => "internals.md",
    ],
    repo="https://github.com/tamasgal/UnROOT.jl/blob/{commit}{path}#L{line}",
    sitename="UnROOT.jl",
    authors="Tamas Gal and contributors",
)

deploydocs(;
    repo="github.com/tamasgal/UnROOT.jl",
)


text:
## ----------------------------------------------------------------------------
# This script take the data from data/raw, fetched from
# Rath 2017 (https://pure.mpg.de/pubman/item/item_2508673_4)
# and convert it for compatibility with further modeling

using ProjAssistant
@quickactivate

@time begin
    import DataFrames: DataFrame
    import CSV

    import Chemostat_Rath2017
    const ChR = Chemostat_Rath2017
    const Rd = ChR.RathData
end

## ----------------------------------------------------------------------------
# Standard Medium 42 MAX-UB
# Standard medium original file
# This is the base of the feed medium
stdm_orig = CSV.read(Rd._stand_medium_raw_file(), DataFrame; delim = "\t");

## ----------------------------------------------------------------------------
# Converting all conc to mM
stdm_conv = DataFrame(stdm_orig);

# Converting all needed to mM
# C(g/L) / MM(g/mol) * 1e3 = mM
idx = findfirst(stdm_conv.id .== "HCO3")
stdm_conv[idx, [:conc, :unit]] .= [(stdm_orig[idx, :conc] / 84) * 1e3, "mM"]
idx = findfirst(stdm_conv.id .== "GAL")
stdm_conv[idx, [:conc, :unit]] .= [(stdm_orig[idx, :conc] / 180) * 1e3, "mM"]
idx = findfirst(stdm_conv.id .== "GLC")
stdm_conv[idx, [:conc, :unit]] .= [(stdm_orig[idx, :conc] / 180) * 1e3, "mM"]

# pH = -log(C(H+)), C(H+) = 10^(-pH) * 1e3 = mM
idx = findfirst(stdm_conv.id .== "H")
stdm_conv[idx, [:conc, :unit]] .= [10.0^(-stdm_orig[idx, :conc]) * 1e3, "mM"]

# C(μM) * 1e-3 = mM
for met in ["GLU", "ALA", "ARG", "ASN", "ASP", "CYS", "GLY",
            "HIS", "ILE", "LEU", "LAC", "LYS", "MET", "PHE",
            "PRO", "SER", "THR", "TRP", "TYR", "VAL"]
    local idx = findfirst(stdm_conv.id .== met)
    stdm_conv[idx, [:conc, :unit]] .= [stdm_orig[idx, :conc] * 1e-3, "mM"]
end
# -

# Saving
CSV.write(Rd._stand_medium_proc_file(), stdm_conv, delim = "\t")
println("created $(relpath(Rd._stand_medium_proc_file()))!!!")

## ----------------------------------------------------------------------------
# Cont cult data

# Cell density ρ = 0.25 pgDW/μm³ from Niklas(2011) https://doi.org/10.1007/s00449-010-0502-y.
ρ = 0.25

## ----------------------------------------------------------------------------
# Data Files
# Original files
for exp in Rd.EXPS

    # file names
    # filename = "$(Rd.RATH_CONT_CUL_DATA_FILE_SUFFIX)_$(exp).tsv"

    # orig_file_path = Rd.RATH_CONT_CUL_DATA_ORIG_FILES[exp]
    orig_file_path = Rd._cont_culture_raw_file(exp)

    if !isfile(orig_file_path)
        error("$(orig_file_path) not found!!!")
    end

    # Original data
    orig_data = CSV.read(orig_file_path, DataFrame; delim = "\t")

    # Converted Data
    # conv_data = DataFrame([String, Float64, Float64, String], [:id, :val, :err, :unit])
    conv_data = DataFrame(id = String[], val = Float64[], err = Float64[], unit = String[])

    # Converting cell volume per media volume (CVv) to Cell mass density (gDW/ L)
    # μL/ mL * 1e9     = μm^3/ mL                    (1 μL = 1e9 μm^3)
    # μm^3/ mL * ρ     = pgDW/ mL                    (ρ = 0.25 pgDW/ μm^3) Niklas(2011)
    # pgDW/ mL * 1e-12 = gDW/ mL                     (1 pg = 1 g * 1e-12)
    # gDW/ mL * 1e3    = gDW/ L                      (1 mL = 1e-3 L)
    # Xv = cvrath * 1e9 * ρ * 1e-12 * 1e3
    # Xv = cvrath * ρ (gDW/ L)
    rowidx = findfirst(orig_data.id .== "CVv")
    val = orig_data[rowidx, :val] * ρ
    # Expres error relative to the value
    err = abs(val * orig_data[rowidx, "%err"] / 100)
    unit = "gDW/ L"
    push!(conv_data, ["Xv", val, err, unit])

    # This values was already in the desired units
    for ider in [Rd.MED_S; Rd.FEED_C; ["D", "μ"]]
        rowidx = findfirst(orig_data.id .== ider)

        val = orig_data[rowidx, :val]
        err = abs(val * orig_data[rowidx, "%err"] / 100)
        unit = orig_data[rowidx, :unit]
        push!(conv_data, [ider, val, err, unit])
    end

    # Converting fluxes (nmol/ μL hr) to (mmol/ gDW hr)
    # nmol/ μL hr * 1e-6     = mmol/ μL hr          (1 nmol  = 1 mmol * 1e-6)
    # mmol/ μL hr * 1e-9     = mmol/ μm^3 hr        (1 μL = 1e9 μm^3)
    # mmol/ μm^3 hr * (1/ ρ) = mmol/ pgDW hr        (ρ = 0.25 pgDW/ μm^3) Niklas
    # mmol/ pgDW hr * 1e12   = mmol/ gDW hr         (1 pg = 1 g * 1e-12)
    # q = qrath * 1e-6 * 1e-9 * (1/ ρ) * 1e12
    # q = qrath * (1/ ρ) * 1e-3 (mmol/ gDW hr)
    for ider in Rd.RATE_Q
        rowidx = findfirst(orig_data.id .== ider)

        val = (orig_data[rowidx, :val] * 1e-3) / ρ
        err = abs(val * orig_data[rowidx, "%err"] / 100)
        unit = "mmol/ gDW hr"
        push!(conv_data, [ider, val, err, unit])
    end

    # Converting A1AT flux
    rowidx = findfirst(orig_data.id .== "CD")
    CD = orig_data[rowidx, :val] # (μm)
    CV = (4/3)* π *(CD/2)^3 # (μm³)
    CDW = ρ * CV * 1e-12 # (pgDW) -> (gDW)

    a1at_mw = 54; # kDa (kg/mol)
    rowidx = findfirst(orig_data.id .== "qA1AT")
    qa1at = orig_data[rowidx, :val]/ CDW # (pg/cell d) -> (pg/gDW d)
    qa1at = (qa1at/a1at_mw) * 1e-15 # (pg/gDW d) -> (mol/gDW d)
    qa1at = qa1at * 1e3 # (mol/gDW d) -> (mmol/gDW d)
    qa1at = qa1at/24 # (pg/gDW d) -> (mmol/gDW hr)
    err = abs(qa1at * orig_data[rowidx, "%err"] / 100)
    unit = "mmol/ gDW hr"
    push!(conv_data, ["qA1AT", qa1at, err, unit])

    # Saving
    # conv_file_path = Rd.RATH_CONT_CUL_DATA_CONV_FILES[exp]
    conv_file_path = Rd._cont_culture_proc_file(exp)
    CSV.write(conv_file_path, conv_data, delim = "\t")
    println("created $(relpath(conv_file_path))!!!")

end

## ----------------------------------------------------------------------------
# invitro max fluxes

# max_flux_data_orig = CSV.read(Rd.RATH_MAX_FLUX_ORIG_FILE, DataFrame; delim = "\t")
max_flux_data_orig = CSV.read(Rd._max_invitro_fluxs_raw_file(), DataFrame; delim = "\t")
max_flux_data_conv = DataFrame(max_flux_data_orig)

# Converting fluxes (nmol/ μL hr) to (mmol/ gDW hr)
# nmol/ μL hr * 1e-6     = mmol/ μL hr          (1 nmol  = 1 mmol * 1e-6)
# mmol/ μL hr * 1e-9     = mmol/ μm^3 hr        (1 μL = 1e9 μm^3)
# mmol/ μm^3 hr * (1/ ρ) = mmol/ pgDW hr        (ρ = 0.25 pgDW/ μm^3) Niklas
# mmol/ pgDW hr * 1e12   = mmol/ gDW hr         (1 pg = 1 g * 1e-12)
# q = qrath * 1e-6 * 1e-9 * (1/ ρ) * 1e12
# q = qrath * (1/ ρ) * 1e-3 (mmol/ gDW hr)
for (enz_i, enz) in enumerate(max_flux_data_orig.id)
    # flux data
    for data_id in ["cul_1", "cul_2", "mean", "mean_err"]
        max_flux_data_conv[enz_i, data_id] = (max_flux_data_orig[enz_i, data_id] * 1e-3) / ρ
    end
end
# unit
max_flux_data_conv[!, :flux_unit] .= "mmol/ gDW hr";

## ----------------------------------------------------------------------------
# Saving
# CSV.write(Rd.RATH_MAX_FLUX_CONV_FILE, max_flux_data_conv, delim = "\t")
# println("created $(relpath(Rd.RATH_MAX_FLUX_CONV_FILE))!!!")
CSV.write(Rd._max_invitro_fluxs_proc_file(), max_flux_data_conv, delim = "\t")
println("created $(relpath(Rd._max_invitro_fluxs_proc_file()))!!!")


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

# Group examples that are used in the docs (examples.md)
export draw_cooketriplet, draw_schmidtcassegraintelescope, draw_lensconstruction, draw_zoomlenses, draw_HOEfocus, draw_HOEcollimate, draw_multiHOE, draw_stackedbeamsplitters

function draw_cooketriplet(filename::Union{Nothing,AbstractString} = nothing)
    g1, g2 = SCHOTT.N_SK16, SCHOTT.N_SF2
    sys = AxisymmetricOpticalSystem{Float64}(DataFrame(
        SurfaceType  = ["Object", "Standard", "Standard", "Standard", "Stop", "Standard", "Standard", "Image"],
        Radius       = [Inf,      26.777,     66.604,     -35.571,    35.571, 35.571,     -26.777,    Inf    ],
        Thickness    = [Inf,      4.0,        2.0,        4.0,        2.0,    4.0,        44.748,     missing],
        Material     = [Air,      g1,         Air,        g2,         Air,    g1,         Air,        missing],
        SemiDiameter = [Inf,      8.580,      7.513,      7.054,      6.033,  7.003,      7.506,      15.0   ],
    ))

    origins = Origins.Hexapolar(8, 15.0, 15.0)
    directions = Directions.Constant(0.0, 0.0, -1.0)
    s1 = Sources.Source(; origins, directions, sourcenum=1)

    transform = Transform(rotmatd(10, 0, 0), unitZ3())
    s2 = Sources.Source(; transform, origins, directions, sourcenum=2)

    raygenerator = Sources.CompositeSource(Transform(), [s1, s2])

    trackallrays = test = colorbysourcenum = true; resolution = (1000, 700)
    Vis.drawtracerays(sys; raygenerator, trackallrays, test, colorbysourcenum, resolution)
    Vis.make2dy(); Vis.save(filename)
    return sys
end

function draw_zoomlenses(filenames::Vector{<:Union{Nothing,AbstractString}} = repeat([nothing], 3))
    stops = [2.89, 3.99, 4.90]
    zooms = [9.48, 4.48, 2.00]
    dists = [4.46970613, 21.21, 43.81]

    transform = translation(0.0, 0.0, 10.0)
    origins = Origins.Hexapolar(8, 10.0, 10.0)
    directions = Directions.Constant(0.0, 0.0, -1.0)
    raygenerator = Sources.Source(; transform, origins, directions)

    aspherics = [
        ["4" => 1.0386E-04, "6" => 1.4209E-07, "8" => -8.8495E-09, "10" => 1.2477E-10, "12" => -1.0367E-12, "14" => 3.6556E-15],
        ["4" => 4.2721E-05, "6" => 1.2484E-07, "8" => 9.7079E-09, "10" => -1.8444E-10, "12" => 1.8644E-12, "14" => -7.7975E-15],
        ["4" => 1.1339E-04, "6" => 4.8165E-07, "8" => 1.8778E-08, "10" => -5.7571E-10, "12" => 8.9994E-12, "14" => -4.6768E-14],
    ]
    syss = [
        AxisymmetricOpticalSystem{Float64}(DataFrame(
            SurfaceType = ["Object", "Stop", "Standard", "Standard", "Standard", "Aspheric", "Standard", "Standard", "Aspheric", "Aspheric", "Standard", "Standard", "Standard", "Standard", "Standard", "Standard", "Image"],
            Radius = [Inf64, Inf64, -1.6202203499676E+01, -4.8875855327468E+01, 1.5666614444619E+01, -4.2955326460481E+01, 1.0869565217391E+02, 2.3623907394283E+01, -1.6059097478722E+01, -4.2553191489362E+02, -3.5435861091425E+01, -1.4146272457208E+01, -2.5125628140704E+02, -2.2502250225023E+01, -1.0583130489999E+01, -4.4444444444444E+01, Inf64],
            Parameters = [missing, missing, missing, missing, missing, aspherics[1], missing, missing, aspherics[2], aspherics[3], missing, missing, missing, missing, missing, missing, missing],
            Thickness = [Inf64, 0.0, 5.18, 0.10, 4.40, 0.16, 1.0, 4.96, zoom, 4.04, 1.35, 1.0, 2.80, 3.0, 1.22, dist, missing],
            Material = [Air, Air, OHARA.S_LAH66, Air, NIKON.LLF6, Air, OHARA.S_TIH6, OHARA.S_FSL5, Air, OHARA.S_FSL5, Air, OHARA.S_LAL8, OHARA.S_FSL5, Air, OHARA.S_LAH66, Air, missing],
            SemiDiameter = [Inf64, stop, 3.85433218451, 3.85433218451, 4.36304692871, 4.36304692871, 4.72505505439, 4.72505505439, 4.72505505439, 4.45240784026, 4.45240784026, 4.50974054117, 4.50974054117, 4.50974054117, 4.76271114409, 4.76271114409, 15.0]))
        for (stop, zoom, dist) in zip(stops, zooms, dists)]

    for (sys, filename) in zip(syss, filenames)
        Vis.drawtracerays(sys; raygenerator, trackallrays=true, test=true, numdivisions=50, resolution=(1200, 600))
        Vis.make2dy(); Vis.save(filename)
    end
    return syss
end

function draw_schmidtcassegraintelescope(filename::Union{Nothing,AbstractString} = nothing)
    # glass entrance lens on telescope
    topsurf = Plane(
        SVector(0.0, 0.0, 1.0),
        SVector(0.0, 0.0, 0.0),
        interface = FresnelInterface{Float64}(SCHOTT.N_BK7, Air),
        vishalfsizeu = 12.00075,
        vishalfsizev = 12.00075)
    botsurf = AcceleratedParametricSurface(ZernikeSurface(
        12.00075,
        radius = -1.14659768e+4,
        aspherics = [(4, 3.68090959e-7), (6, 2.73643352e-11), (8, 3.20036892e-14)]),
        17,
        interface = FresnelInterface{Float64}(SCHOTT.N_BK7, Air))
    coverlens = Cylinder(12.00075, 1.4) ∩ topsurf ∩ leaf(botsurf, Transform(rotmatd(0, 180, 0), Vec3(0.0, 0.0, -0.65)))

    # big mirror with a hole in it
    frontsurfacereflectance = 1.0
    bigmirror = (
        ConicLens(SCHOTT.N_BK7, -72.65, -95.2773500000134, 0.077235, Inf, 0.0, 0.2, 12.18263; frontsurfacereflectance) -
        leaf(Cylinder(4.0, 0.3, interface = opaqueinterface()), translation(0.0, 0.0, -72.75))
    )

    # small mirror supported on a spider
    backsurfacereflectance = 1.0
    smallmirror = SphericalLens(SCHOTT.N_BK7, -40.65, Inf, -49.6845, 1.13365, 4.3223859; backsurfacereflectance)

    obscuration1 = Circle(4.5, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, -40.649), interface = opaqueinterface())
    obscurations2 = Spider(3, 0.5, 12.0, SVector(0.0, 0.0, -40.65))

    # put it together with the detector
    la = LensAssembly(coverlens(), bigmirror(), smallmirror(), obscuration1, obscurations2...)
    det = Circle(3.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, -92.4542988), interface = opaqueinterface())
    sys = CSGOpticalSystem(la, det)

    # define ray generator
    transform = translation(0.0, 0.0, 10.0)
    origins = Origins.Hexapolar(8, 20.0, 20.0)
    directions = Directions.Constant(0.0, 0.0, -1.0)
    raygenerator = Sources.Source(; transform, origins, directions)

    # draw and output
    Vis.drawtracerays(sys; raygenerator, trackallrays = true, colorbynhits = true, test = true, numdivisions = 100, drawgen = false)
    Vis.save(filename)
    return nothing
end

function draw_lensconstruction(filename::Union{Nothing,AbstractString} = nothing)
    topsurface = leaf(
        AcceleratedParametricSurface(
            QTypeSurface(
                9.0,
                radius = -25.0,
                conic = 0.3,
                αcoeffs = [(1, 0, 0.3), (1, 1, 1.0)],
                βcoeffs = [(1, 0, -0.1), (2, 0, 0.4), (3, 0, -0.6)],
                normradius = 9.5),
            interface = FresnelInterface{Float64}(SCHOTT.N_BK7, Air)),
        translation(0.0, 0.0, 5.0))
    botsurface = Plane(
        SVector(0.0, 0.0, -1.0),
        SVector(0.0, 0.0, -5.0),
        vishalfsizeu = 9.5,
        vishalfsizev = 9.5,
        interface = FresnelInterface{Float64}(SCHOTT.N_BK7, Air))
    barrel = Cylinder(
        9.0, 20.0, interface = FresnelInterface{Float64}(SCHOTT.N_BK7, Air, reflectance=0.0, transmission=0.0)
    )
    lens = (barrel ∩ topsurface ∩ botsurface)(Transform{Float64}(0.0, Float64(π), 0.0, 0.0, 0.0, -5.0))
    detector = Rectangle(15.0, 15.0, [0.0, 0.0, 1.0], [0.0, 0.0, -67.8], interface = opaqueinterface())
    sys = CSGOpticalSystem(LensAssembly(lens), detector)

    Vis.drawtracerays(sys, test = true, trackallrays = true, colorbynhits = true)
    Vis.save(filename)
    return nothing
end

function draw_HOEfocus(filename::Union{Nothing,AbstractString} = nothing)
    rect = Rectangle(5.0, 5.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, 0.0))
    int = HologramInterface(
        SVector(0.0, -3.0, -20.0), ConvergingBeam,
        SVector(0.0, 0.0, -1.0), CollimatedBeam,
        0.55, 9.0, Air, SCHOTT.N_BK7, Air, Air, Air, 0.05, false)
    obj = HologramSurface(rect, int)
    sys = CSGOpticalSystem(
        LensAssembly(obj),
        Rectangle(10.0, 10.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, -25.0),
        interface = opaqueinterface()))

    raygenerator = Sources.Source(;
        transform = translation(0.0, 0.0, 10.0),
        spectrum = Spectrum.DeltaFunction(0.55),
        origins = Origins.RectGrid(3.0, 3.0, 5, 5),
        directions = Directions.Constant(0.0, 0.0, -1.0))

    Vis.drawtracerays(sys; raygenerator, trackallrays = true, rayfilter = nothing, test = true)
    Vis.save(filename)
    return nothing
end

function draw_HOEcollimate(filename::Union{Nothing,AbstractString} = nothing)
    rect = Rectangle(5.0, 5.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, 0.0))
    int = HologramInterface(
        SVector(0.1, -0.05, -1.0), CollimatedBeam,
        SVector(0.0, 0.0, 10), DivergingBeam,
        0.55, 9.0, Air, SCHOTT.N_BK7, Air, Air, Air, 0.05, false)
    obj = HologramSurface(rect, int)
    sys = CSGOpticalSystem(
        LensAssembly(obj),
        Rectangle(10.0, 10.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, -25.0),
        interface = opaqueinterface()))

    raygenerator = Sources.Source(
        transform = Transform(rotmatd(180, 0, 0), Vec3(0.0, 0.0, 10.0)),
        spectrum = Spectrum.DeltaFunction(0.55),
        origins = Origins.Point(),
        directions = Directions.RectGrid(π/4, π/4, 8, 8))

    Vis.drawtracerays(sys; raygenerator, trackallrays = true, rayfilter = nothing, test = true)
    Vis.save(filename)
    return nothing
end

function draw_multiHOE(filename::Union{Nothing,AbstractString} = nothing)
    rect = Rectangle(5.0, 5.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, 0.0))
    int1 = HologramInterface(
        SVector(-5.0, 0.0, -20.0), ConvergingBeam,
        SVector(0.0, -1.0, -1.0), CollimatedBeam,
        0.55, 100.0, Air, SCHOTT.N_BK7, Air, Air, Air, 0.05, false)
    int2 = HologramInterface(
        SVector(5.0, 0.0, -20.0), ConvergingBeam,
        SVector(0.0, 1.0, -1.0), CollimatedBeam,
        0.55, 100.0, Air, SCHOTT.N_BK7, Air, Air, Air, 0.05, false)
    mint = MultiHologramInterface(int1, int2)
    obj = MultiHologramSurface(rect, mint)
    sys = CSGOpticalSystem(
        LensAssembly(obj),
        Rectangle(10.0, 10.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 0.0, -20.0), interface = opaqueinterface()))

    spectrum = Spectrum.DeltaFunction(0.55)
    origins = Origins.RectUniform(3.0, 3.0, 500)
    directions = Directions.Constant(0.0, 0.0, -1.0)
    s1 = Sources.Source(; spectrum, origins, directions, sourcenum = 1,
        transform = Transform(rotmatd(-45, 0, 0), Vec3(0.0, 3.0, 3.0)))
    s2 = Sources.Source(; spectrum, origins, directions, sourcenum = 2,
        transform = Transform(rotmatd(45, 0, 0), Vec3(0.0, -3.0, 3.0)))
    s3 = Sources.Source(; spectrum, origins, directions, sourcenum = 3,
        transform = translation(0.0, 0.0, 3.0))
    raygenerator = Sources.CompositeSource(Transform(), [s1, s2, s3])

    Vis.drawtracerays(sys; raygenerator, trackallrays = true, colorbysourcenum = true, rayfilter = nothing, drawgen = true)
    Vis.save(filename)
    return nothing
end

function draw_stackedbeamsplitters(filenames::Vector{<:Union{Nothing,AbstractString}} = repeat([nothing], 3))
    # ReflectOrTransmit: nondeterministic
    # Transmit: deterministic, all beamsplitters transmissive
    # Reflect: deterministic, all beamsplitters reflective
    interfacemodes = [ReflectOrTransmit, Transmit, Reflect]

    for (interfacemode, filename) in zip(interfacemodes, filenames)
        interface = FresnelInterface{Float64}(SCHOTT.N_BK7, Air; reflectance=0.5, transmission=0.5, interfacemode)
        bs_1 = OpticSim.transform(
                Cuboid(10.0, 20.0, 2.0, interface=interface),
                translation(0.0, 0.0, -30.0-2*sqrt(2))*rotationX(π/4))

        l1 = OpticSim.transform(
            SphericalLens(SCHOTT.N_BK7, -70.0, 30.0, Inf, 5.0, 10.0),
            translation(0.0, -1.34, 0.0))

        bs_2 = OpticSim.transform(
            Cuboid(10.0, 20.0, 2.0, interface=interface),
            translation(0.0, 40.0, -30.0+2*sqrt(2))*rotationX(π/4))

        l2 = OpticSim.transform(
            SphericalLens(SCHOTT.N_BK7, -70.0, 30.0, Inf, 5.0, 10.0),
            translation(0.0, 40.0, 0.0))

        la = LensAssembly(bs_1(), l1(), bs_2(), l2())

        detector = Rectangle(20.0, 40.0, SVector(0.0, 0.0, 1.0), SVector(0.0, 20.0, -130.0); interface = opaqueinterface())
        sys = CSGOpticalSystem(la, detector)

        Vis.drawtracerays(sys; trackallrays=true, rayfilter=nothing, colorbynhits=true)
        Vis.save(filename)
    end
    return nothing
end


text:
#=
Create efficient approximation function

Ignacio Quintero Mächler

t(-_-t)

Created 16 09 2017
Updated 17 01 2019
=#

"""
    linpred(val::Float64, x1::Float64, x2::Float64, y1::Float64, y2::Float64)

Estimate val according to linear interpolation for a range.
"""
linpred(val::Float64, x1::Float64, x2::Float64, y1::Float64, y2::Float64) =
  (y1 + (val - x1)*(y2 - y1)/(x2 - x1))::Float64

"""
    idxrange(x::Array{Float64,1}, val::Float64)

Get indexes in sorted vector `x` corresponding to the range in which
`val` is in using a sort of uniroot algorithm.
"""
function idxrange(x::Array{Float64,1}, val::Float64)

  @inbounds begin

    a::Int64 = 1

    if x[a] > val
      return a, false
    end

    b::Int64 = lastindex(x)

    if x[b] < val
      return b, false
    end

    mid::Int64 = div(b,2)

    while b-a > 1
      val < x[mid] ? b = mid : a = mid
      mid = div(b + a, 2)
    end

    if x[a] == val
      return a, false
    elseif x[b] == val
      return b, false
    else
      return a, true
    end

  end
end

"""
    @generated function approxf_full(t ::Float64,
                                     r ::Array{Float64,1},
                                     x ::Array{Float64,1},
                                     y ::Array{Float64,N},
                                     ::Val{nc}) where {N, nc}
Returns the values of `y` at `t` using an approximation function.
"""
@generated function approxf_full(t ::Float64,
                                 r ::Array{Float64,1},
                                 x ::Array{Float64,1},
                                 y ::Array{Float64,N},
                                 ::Val{nc}) where {N, nc}

  lex1 = quote end
  pop!(lex1.args)

  if N == 1
    push!(lex1.args, :(r[1] = linpred(t, x[a], x[a+1], y[a], y[a+1])::Float64))
  else
    # unroll loop
    for i = Base.OneTo(nc)
      push!(lex1.args, :(r[$i] = linpred(t, xa, xap1, y[a,$i], y[a+1, $i])::Float64))
    end

    # add one assignment
    pushfirst!(lex1.args, :(xap1 = x[a+1]::Float64))
    pushfirst!(lex1.args, :(xa   = x[a]::Float64))
  end

  lex2 = quote end
  pop!(lex2.args)

  if N == 1
    push!(lex2.args, :(r[1] = y[a]::Float64))
  else
    # unroll loop
    for i = Base.OneTo(nc)
      push!(lex2.args, :(r[$i] = y[a,$i]::Float64))
    end
  end

  lex = quote
    a, lp = idxrange(x, t)::Tuple{Int64, Bool}
    if lp
      $lex1
    else
      $lex2
    end
  end

  # aesthetic cleaning
  deleteat!(lex.args,[1,3])

  popfirst!(lex.args[2].args[2].args)
  lex.args[2].args[2] = lex.args[2].args[2].args[1]

  popfirst!(lex.args[2].args[3].args)
  lex.args[2].args[3] = lex.args[2].args[3].args[1]

  return quote
    @inbounds begin
      $lex
    end
    return nothing
  end
end

"""
  make_af(x::Array{Float64,1}, y::Array{Float64,N}, ::Val{ny})

make approximate function closure
"""
function make_af(x::Array{Float64,1}, y::Array{Float64,N}, ::Val{ny}) where {N, ny}

  af! = (t::Float64, r::Array{Float64,1}) ->
    begin
      approxf_full(t::Float64, r::Array{Float64,1}, x::Array{Float64,1}, y::Array{Float64,N}, Val(ny))
      return nothing
    end

  return af!
end

"""
    @generated function approxf_full(t ::Float64,
                                     r ::Array{Float64,1},
                                     x ::Array{Float64,1},
                                     y ::Array{Float64,1},
                                     ::Val{nc}) where {N, nc}

Returns the values of `y` at `t` using an approximation function
when `y` is an array of arrays.
"""
@generated function approxf_full(t ::Float64,
                                 r ::Array{Float64,1},
                                 x ::Array{Float64,1},
                                 y ::Array{Array{Float64,1},1},
                                 ::Val{nc}) where {nc}

  lex1 = quote end
  pop!(lex1.args)

  # unroll loop
  for i = Base.OneTo(nc)
    push!(lex1.args, :(r[$i] = linpred(t, xa, xap1, y[a][$i], y[a+1][$i])::Float64))
  end

  # add one assignment
  pushfirst!(lex1.args, :(xap1 = x[a+1]::Float64))
  pushfirst!(lex1.args, :(xa   = x[a]::Float64))

  lex2 = quote end
  pop!(lex2.args)

  # unroll loop
  for i = Base.OneTo(nc)
    push!(lex2.args, :(r[$i] = y[a][$i]::Float64))
  end

  lex = quote
    a, lp = idxrange(x, t)::Tuple{Int64, Bool}
    if lp
      $lex1
    else
      $lex2
    end
  end

  # aesthetic cleaning
  deleteat!(lex.args,[1,3])

  popfirst!(lex.args[2].args[2].args)
  lex.args[2].args[2] = lex.args[2].args[2].args[1]

  popfirst!(lex.args[2].args[3].args)
  lex.args[2].args[3] = lex.args[2].args[3].args[1]

  return quote
    @inbounds begin
      $lex
    end
    return nothing
  end
end

"""
  make_af(x::Array{Float64,1}, y::Array{Array{Float64,1},1}, ::Val{ny})

make approximate function closure
"""
function make_af(x::Array{Float64,1}, y::Array{Array{Float64,1},1}, ::Val{ny}) where {ny}

  af! = (t::Float64, r::Array{Float64,1}, y::Array{Array{Float64,1},1}) ->
    begin
      approxf_full(t::Float64, r::Array{Float64,1}, x::Array{Float64,1}, y::Array{Array{Float64,1},1}, Val(ny))
      return nothing
    end

  return af!
end

max_stars_repo_name: hmorlon/JPANDA

max_issues_repo_name: hmorlon/JPANDA

max_forks_repo_name: hmorlon/JPANDA


text:
#using Debug
#using NumericExtensions
using Base.LinAlg.copytri!

#import NumericExtensions.evaluate, NumericExtensions.result_type

#include("logging.jl")

function dirichlet_process_mixture_prior(cluster_assignments::Array{Int,1},
                                         alpha::Float64)
    n = length(cluster_assignments)
    d = max(cluster_assignments)
    cluster_counts = zeros(d)
    for i = 1:n
        cluster_counts[cluster_assignments[i]] += 1
    end

    d*log(alpha) + lgamma(alpha) + sum(lfact(cluster_counts-1)) - lgamma(alpha+n)
end

function K_given_W(K, N, W)
    assert(length(W) >= K)
    sum(log(W[1:K])) + (N - K)*log(sum(W[1:K])) + sum(log([N-K+1:N]))
end

function K_given_W_gradient(K, N ,W)
    assert(length(W) >= K)
    sum(log(W[1:K])) ./ W[1:K] + (N - K)/sum(W[1:K])*ones(K)
end

function stick_breaking(W, K, alpha)
    assert(length(W) >= K)

    if sum(W[1:K]) > 1.0
        return -Inf
    end

    betas = W[1:K]
    betas[2:K] ./= 1 - cumsum(W[1:K-1])

    sum(logpdf(Beta(1,alpha), betas))
end

function stick_breaking_gradient(W, K, alpha)
    assert(length(W) >= K)

    if sum(W[1:K]) > 1.0
        return zeros(K)
    end

    betas = W[1:K]
    betas[2:K] ./= 1 - cumsum(W[1:K-1])

    (1-alpha) ./ (1-betas)
end

function GMM_likelihood(data::Array{Float64,2}, # Nxp
                        w::Array{Float64,1}, # not necessarily normalized
                        mu::Array{Float64,2}, # pxK
                        sigma::Array{Array{Float64,2},1})

    K = length(w)
    (N,p) = size(data)
    if K <= 0
        return -Inf
    end

    w_norm = w / sum(w)
    logprobs = zeros(N,K)

    if minimum(w_norm) < 0
        return -Inf
    elseif !isfinite(sum(w_norm))
        lprintln("w_norm: ", w_norm)
    end

    for k = 1:K
        if !isposdef(sigma[k])
            return -Inf
        end
        logprobs[:,k] = normal_logpdf(data, mu[:,k], sigma[k])# + log(w_norm[k])
    end

    logprob = sum(logsumexp(broadcast(+,logprobs, log(w_norm)'), 2) )

#    logprob = 0.0
#    for i = 1:size(data,1)
#        logprob += logsumexp(logprobs[i,:]' + log(w_norm))
#    end

    @assert !isinf(logprob)

    logprob
end

function GMM_likelihood_gradient(data::Array{Float64,2},
                                 w::Array{Float64,1},
                                 mu::Array{Float64,2},
                                 sigma::Array{Array{Float64,2},1})
    K = length(w)
    (N,p) = size(data)

    dmu = zeros(size(mu))
    dw = zeros(size(w))
    dSigma = similar(sigma) #Array( Array{Float64,2}, length(sigma))
    for i = 1:length(sigma)
        dSigma[i] = zeros(size(sigma[i]))
    end

    if K <= 0
        return (dw, dmu, dSigma)
    end

    w_sum = sum(w)
    w_norm = w / w_sum

    if !all(w .>= 0.0)
        return (dw, dmu, dSigma)
    end
    #@assert all(w_norm .> 0.0)

    logprobs = zeros(N,K)
    for k = 1:K
        logprobs[:,k] = normal_logpdf(data, mu[:,k], sigma[k]) # + log(w_norm[k])
    end

    # contribution of this component is zero, so return zero gradients
    if any(logprobs .== -Inf)
        return (dw, dmu, dSigma)
    end

    total_logprobs = zeros(size(data,1))
    for i = 1:N
        total_logprobs[i]= logsumexp(logprobs[i,:]' + log(w_norm))
    end

    for k = 1:K
        sumexp_weights = w_norm[k] * normal_dmu_logpdf(data, mu[:,k], sigma[k])
        sumexp_args = logprobs[:,k] - total_logprobs
        for pind = 1:p
            dmu[pind,k] = sum(sumexp_weights[:,pind] .* exp(sumexp_args))
        end

        sumexp_weights = w_norm[k] * normal_dSigma_logpdf(data, mu[:,k], sigma[k])
        for p1 = 1:p
            for p2 = 1:p
                for i = 1:N
                    dSigma[k][p1,p2] += sumexp_weights[i][p1,p2] .* exp(sumexp_args[i])
                end
            end
        end

        if any(!isfinite(dSigma[k]))
            lprintln("w: ")
            lprintln(w)
            lprintln(w_norm)
            aa = find(!isfinite(logprobs[:,k]))
            lprintln("logprobs: ", logprobs[aa,k])
        end

        @assert issym(dSigma[k])

    end

    for k = 1:K
        # -log(w_sum) because we are using the normalized weights in total_logprobs
        log_dws = logprobs[:,k]-log(w_sum)-total_logprobs
        dw[k] = exp(logsumexp(log_dws)) - N/w_sum
        #dw[k] = sum(exp(logprobs[:,k]-total_logprobs)) - N/w_sum
    end

    (dw, dmu, dSigma)
end

function GMM_L_prior(L, D, K)

    total_prob = 0.0
    for k = 1:K
        for i = 1:size(L[k],1)
            for j = 1:i-1
                total_prob += normal_logpdf(L[k][i,j], 0.0, sqrt(D[i,k]) )
            end
        end
    end
    total_prob
end

function GMM_L_prior_gradient(L, D, K)

    result = similar(L[1:K])
    result_D = similar(D[:,1:K])

    for k = 1:K
        p = size(L[k],1)
        result[k] = zeros(p,p)
        result_D[:,k] = 0
        for i = 1:size(L[k],1)
            for j = 1:i-1
                result[k][i,j] = normal_dx_logpdf(L[k][i,j], 0.0, sqrt(D[i,k]))
                result_D[i,k] += normal_dsigma2_logpdf(L[k][i,j], 0.0, D[i,k])
            end
        end
    end

    (result, result_D)
end

function GMM_D_prior(D, nu, K)
    p = size(D,1)

    total_prob = 0.0
    for k = 1:K
        for i = 1:p
            total_prob += inverse_gamma_logpdf(D[i,k], (nu+i-p)/2, 1/2)
        end
    end

    total_prob
end

function GMM_D_prior_gradient(D, nu, K)
    p = size(D,1)

    result = zeros(p,K)

    for k = 1:K
        for i = 1:p
            result[i,k] = inverse_gamma_dx_logpdf(D[i,k], (nu+i-p)/2, 1/2)
        end
    end

    result
end

function GMM_likelihood(data::Array{Float64,2}, # Nxp
                        w::Array{Float64,1}, # not necessarily normalized
                        mu::Array{Float64,2}, # pxK
                        L::Array{Array{Float64,2},1},
                        D::Array{Float64,2}) # pxK

    K = length(w)
    (N,p) = size(data)
    if K <= 0
        return -Inf
    end

    w_norm = w / sum(w)
    logprobs = zeros(N,K)

    if minimum(w_norm) < 0
        return -Inf
    elseif !isfinite(sum(w_norm))
        lprintln("w_norm: ", w_norm)
    end

    for k = 1:K
        x = broadcast(-, data, mu[:,k]')
        Lk = L[k]
        Dk = D[:,k]
        Lk_x = Lk *  x'
        Lk_x2 = Lk_x .* (diagm(1 ./ Dk) * Lk_x)
        logprobs[:,k] = -sum(log(Dk))/2 .- sum(Lk_x2,1)'/2
    end

    logprob = sum(logsumexp(broadcast(+,logprobs, log(w_norm)'), 2) )

    @assert !isinf(logprob)

    logprob
end

function GMM_likelihood_gradient(data::Array{Float64,2},
                                 w::Array{Float64,1},
                                 mu::Array{Float64,2},
                                 L::Array{Array{Float64,2},1},
                                 D::Array{Float64,2}) # pxK
    K = length(w)
    (N,p) = size(data)

    dmu = zeros(size(mu))
    dw = zeros(size(w))
    dL = similar(L) #Array( Array{Float64,2}, length(L))
    dD = similar(D)

    for i = 1:length(L)
        dL[i] = zeros(size(L[i]))
    end

    if K <= 0
        return (dw, dmu, dL, dD)
    end

    w_sum = sum(w)
    w_norm = w / w_sum

    if !all(w .>= 0.0)
        return (dw, dmu, dL, dD)
    end
    #@assert all(w_norm .> 0.0)

    logprobs = zeros(N,K)
    for k = 1:K
        x = broadcast(-, data, mu[:,k]')
        Lk = L[k]
        Dk = D[:,k]
        Lk_x = Lk *  x'
        Lk_x2 = Lk_x .* (diagm(1 ./ Dk) * Lk_x)
        logprobs[:,k] = -sum(log(Dk))/2 .- sum(Lk_x2,1)/2'
    end

    # contribution of this component is zero, so return zero gradients
    if any(logprobs .== -Inf)
        return (dw, dmu, dL, dD)
    end

    total_logprobs = zeros(size(data,1))
    for i = 1:N
        total_logprobs[i]= logsumexp(logprobs[i,:]' + log(w_norm))
    end

    for k = 1:K
        sumexp_weights = w_norm[k] * normal_dmu_logpdf(data, mu[:,k], L[k], D[:,k])
        sumexp_args = logprobs[:,k] - total_logprobs
        for pind = 1:p
            dmu[pind,k] = sum(sumexp_weights[:,pind] .* exp(sumexp_args))
        end

        sumexp_weights = w_norm[k] * normal_dL_logpdf(data, mu[:,k], L[k], D[:,k])

        for p1 = 1:p
            for p2 = 1:p
                for i = 1:N
                    dL[k][p1,p2] += sumexp_weights[i][p1,p2] .* exp(sumexp_args[i])
                end
            end
        end

        sumexp_weights = w_norm[k] * normal_dD_logpdf(data, mu[:,k], L[k], D[:,k])
        for pind = 1:p
            dD[pind,k] = sum(sumexp_weights[:,pind] .* exp(sumexp_args))
        end

    end

    for k = 1:K
        # -log(w_sum) because we are using the normalized weights in total_logprobs
        log_dws = logprobs[:,k].-log(w_sum).-total_logprobs
        dw[k] = exp(logsumexp(log_dws)) - N/w_sum
        #dw[k] = sum(exp(logprobs[:,k]-total_logprobs)) - N/w_sum
    end

    (dw, dmu, dL, dD)
end

vectorize_L = th -> (inds = tril_inds(th,-1); th[inds])
devectorize_L = (th, y) -> (inds = tril_inds(th,-1); th[inds] = y)

function GMM_cache(data::Array{Float64,2}, # Nxp
                   w::Array{Float64,1}, # not necessarily normalized
                   mu::Array{Float64,2}, # pxK
                   #sigma::Array{Array{Float64,2},1},
                   Lm::Array{Array{Float64,2},1},
                   D::Matrix{Float64},
                   L::Int,
                   l_inds::Range{Int},
                   gradient_name::String,
                   gradient_subind::Int)

    Lmax = L+1

    max_num_L = l_inds[end] <= L ? L-1 : (l_inds[end] == L+1 ? L : L+1)

    @assert length(w) >= max_num_L
    #@assert length(sigma) >= max_num_L
    @assert size(mu, 2) >= max_num_L
    (N,p) = size(data)

    if L < 0
        error("L must be nonnegative")
    end

    total_logprobs = zeros(L+2)
    total_gradients = cell(L+2)
    for l = 1:L+2
        num_L = l<=L ? L-1 : (l==L+1 ? L : L+1)
        total_gradients[l] = cell(4)
        total_gradients[l][1] = zeros(num_L)
        total_gradients[l][2] = zeros(p,num_L)
        total_gradients[l][3] = Array(Array{Float64,2},num_L)
        total_gradients[l][4] = zeros(p,num_L)
        for i = 1:num_L
            total_gradients[l][3][i] = zeros(p,p)
        end
    end

    cache = cell(L+2)

    if minimum(w[1:Lmax]) < 0.0
        for l = 1:L+2
            total_logprobs[l] = -Inf
        end
        for l = 1:L+2
            cache[l] = (total_logprobs[l], total_gradients[l])
        end
        return cache
    end

    logprobs = zeros(N,Lmax)
    Sx = zeros(p,N,Lmax)
    Lk_x = zeros(p,N,Lmax)

    diffs_k = zeros(N,p,Lmax)

    for k = 1:Lmax
        #Sigma = sigma[k]
#        if !isposdef(Sigma)
#            for l = 1:L+2
#                total_logprobs[l] = -Inf
#            end
#            for l = 1:L+2
#                cache[l] = (total_logprobs[l], total_gradients[l])
#            end
#            return cache
#        end

        diffs_k[:,:,k] = broadcast(-, data, mu[:,k]')
        #Sx[:,:,k] = Sigma \ diffs'
        Lk_x[:,:,k] = Lm[k] * diffs_k[:,:,k]'

        Lk_x2 = Lk_x[:,:,k] .* (diagm(1./D[:,k]) * Lk_x[:,:,k])
        #logprobs[:,k] = -0.5*sum(diffs.*Sx[:,:,k]',2) - 0.5p*log(2pi) - 0.5*logdet(Sigma)
        logprobs[:,k] = -0.5*sum(Lk_x2,1) .- 0.5p*log(2pi) .- 0.5*sum(log(D[:,k]))
        #logprobs[:,k] = normal_logpdf(data, mu[:,k], sigma[k])
    end

    l_logprobs = zeros(N,L+2)

    dw = zeros(size(w))

    sumexp_weights_mu = zeros(N,p,Lmax)
    sumexp_weights_Sigma = cell(Lmax)
    sumexp_weights_L = cell(Lmax)
    sumexp_weights_D = zeros(N,p,Lmax)

    # only need the gradient for the currently updating parameter and subparameter
    if gradient_name == "mu"
        k = gradient_subind
        #sumexp_weights_mu[:,:,k] = w[k]*normal_dmu_logpdf(data, mu[:,k], sigma[k], Sx[:,:,k])
        sumexp_weights_mu[:,:,k] = w[k]*normal_dmu_logpdf(data, mu[:,k], Lm[k], D[:,k], Lk_x[:,:,k])
    elseif gradient_name == "Sigma"
        k = gradient_subind
        sumexp_weights_Sigma[k] = w[k]*normal_dSigma_logpdf(data, mu[:,k], sigma[k], Sx[:,:,k])
    elseif gradient_name == "L"
        k = gradient_subind
        sumexp_weights_L[k] = w[k]*normal_dL_logpdf(data, mu[:,k], D[:,k], Lk_x[:,:,k], diffs_k[:,:,k])
    elseif gradient_name == "D"
        k = gradient_subind
        sumexp_weights_D[:,:,k] = w[k]*normal_dD_logpdf(data, mu[:,k], D[:,k], Lk_x[:,:,k])
    end

#    for k = 1:Lmax
#        sumexp_weights_mu[:,:,k] = w[k]*normal_dmu_logpdf(data, mu[:,k], sigma[k], Sx[:,:,k])
#        sumexp_weights_Sigma[k] = w[k]*normal_dSigma_logpdf(data, mu[:,k], sigma[k], Sx[:,:,k])
#    end

    if L == 0
        total_logprobs[1] = -Inf
    end

    for l = l_inds
        if total_logprobs[l] == -Inf
            continue
        end
        k_range = 1:Lmax

        if l == L+1
            k_range = 1:L
        elseif l <= L
            k_range = [ i != l ? i : L for i = 1:L-1]
        end
        num_L = length(k_range)

        dw = total_gradients[l][1]
        dmu = total_gradients[l][2]
        #dSigma = total_gradients[l][3]
        dL = total_gradients[l][3]
        dD = total_gradients[l][4]

        w_sum = sum(w[k_range])

        if num_L > 0
            t_probs = broadcast(+,logprobs[:,k_range], log(w[k_range])')
            l_logprobs[:,l] = logsumexp(t_probs, 2)
        end
#        for i = 1:N
#            l_logprobs[i,l] = logsumexp(logprobs[i,k_range]' + log(w[k_range]))
#        end
        total_logprobs[l] = sum(l_logprobs[:,l]) - N*log(sum(w[k_range]))

        if !isfinite(total_logprobs[l])
            total_logprobs[l] = -Inf
        end

        for k_ind = 1:num_L
            k = k_range[k_ind]
            sumexp_args = logprobs[:,k] - l_logprobs[:,l]

            if gradient_name == "mu" && k == gradient_subind
                for pind = 1:p
                    dmu[pind,k_ind] = sum(sumexp_weights_mu[:,pind,k] .* exp(sumexp_args))
                end
            end

            #dSigma[k_ind] = sum(sumexp_weights_Sigma[k].*exp(sumexp_args) )

            if gradient_name == "Sigma" && k == gradient_subind
                for i = 1:N
                    dSigma[k_ind] += sumexp_weights_Sigma[k][i] .* exp(sumexp_args[i])
                end
            end

            if gradient_name == "L" && k == gradient_subind
                for i = 1:N
                    dL[k_ind] += sumexp_weights_L[k][i] .* exp(sumexp_args[i])
                end
            end

            if gradient_name == "D" && k == gradient_subind
                for pind = 1:p
                    dD[pind,k_ind] = sum(sumexp_weights_D[:,pind,k] .* exp(sumexp_args))
                end
            end

            if gradient_name == "w"
                dw[k_ind] = exp(logsumexp(sumexp_args)) - N/w_sum
            end

#            if any(!isfinite(dSigma[k_ind]))
#                lprintln("k: ", k)
#                lprintln("w: ")
#                lprintln(w)
#                aa = find(!isfinite(logprobs[:,k]))
#                lprintln("logprobs: ", logprobs[aa,k])
#                lprintln(dSigma[k_ind])
#                lprintln(dmu[:,k_ind])
#                lprintln("total_logprobs: ", total_logprobs)
#                aa = find(!isfinite(sumexp_args))
#                lprintln("args: ", sumexp_args[aa])
#
#                lprintln("LL: ", total_logprobs[end] -N*log(sum(w[1:Lmax])))
#                lprintln("LL_actual: ", GMM_likelihood(data, w, mu, sigma))
#                lprintln(GMM_likelihood_gradient(data, w, mu, sigma))
#
#            end

            #@assert issym(dSigma[k_ind])
        end

    end

    for k = 1:Lmax
        if !all(isfinite(total_gradients[k][1]))
            lprintln("k: ", k)
            lprintln(total_gradients[k][1])
        end
    end

    if total_logprobs[1] == Inf
        assert(false)
    end

    for l = 1:L+2
        cache[l] = (total_logprobs[l], total_gradients[l])
    end

    cache
end

function ZWZ_likelihood(Y::Vector{Matrix{Float64}},
                        W::Matrix{Float64},
                        Z::Matrix{Int},
                        A::Vector{Float64},
                        B::Vector{Float64},
                        C::Float64)

    (_,N) = size(Y[1])
    (_,K) = size(Z)

    log_probs = similar(Y, Matrix{Float64})
    if any(isinf(W)) || any(isnan(W))
        for p = 1:length(Y)
            log_probs[p] = -Inf*ones(N,N)
            nan_inds = find( isnan(Y[p]))
            log_probs[p][nan_inds] = 0.0
        end
        return log_probs
    end

    ZWZ = K > 0 ? Z*W*Z' : zeros(N,N)
    ZWZ = broadcast(+,ZWZ, A')
    ZWZ = broadcast(+,ZWZ, B) .+ C

    for p = 1:length(Y)
        log_probs[p] = map(LogLogistic(), ZWZ, Y[p])
        nan_inds = find( isnan(Y[p]))
        log_probs[p][nan_inds] = 0.0
    end

    log_probs
end

function get_nondiagonal_indices(Y::Matrix)
    setdiff(1:length(Y), diagind(Y))
end

function FA_likelihood(data::Matrix{Float64},
                       Z::Matrix{Int},
                       W::Matrix{Float64},
                       sigma_X::Float64)
    ZW = Z*W
    p = length(ZW)

    logprob = sum(-0.5*(ZW - data).^2/sigma_X^2) - 0.5p*log(2pi) - p*log(sigma_X)
#    for i = 1:length(data)
#        logprob += normal_logpdf(ZW[i],data[i], sigma_X)
#    end
    logprob
end

#function dirichlet_process_prior(cluster_counts::Array{Int,1}, alpha::Float64)
#    n = sum(cluster_counts)
#    d = length(cluster_counts)
#    d*log(alpha) + lgamma(alpha) + sum(lfact(cluster_counts-1)) - lgamma(alpha+n)
#end

# Assumes size(x) = Nxp, and size(mu) = p

function normal_logpdf(x::Float64, mu::Float64, sigma::Float64)
    diff = x-mu
    -0.5*diff*diff/(sigma*sigma) - 0.5*log(2pi) - log(sigma)
end
function normal_logpdf(x,mu, sigma::Float64)
    if length(x) == 0
        return 0.0
    end
    p = size(x,2)
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    -0.5*sum(diffs.*diffs,2)/(sigma*sigma) .- 0.5p*log(2pi) .- p*log(sigma)
end

function normal_logpdf(x, mu::Float64, sigma::Float64)
    if length(x) == 0
        return 0.0
    end
    p=1
    diffs = x.-mu
    -0.5*diffs.*diffs/(sigma*sigma) .- 0.5p*log(2pi) .- p*log(sigma)
end
#TODO
# function normal_logpdf(x, mu, L::Matrix{Float64}, D::Vector{Float64})
#     diffs = get_diff(x,mu')
#     diffs * L' * diagm(1./D) * L
# end
function normal_logpdf(x,mu, Sigma::Array{Float64,2})
    p = size(x,2)
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    Sig_inv_diff = Sigma \ diffs'
    #@assert issym(Sigma) cant even pass Symmetric here!
    if !isposdef(Sigma)
        return -Inf
    end

    -0.5*sum(diffs.*Sig_inv_diff',2) .- 0.5p*log(2pi) .- 0.5*logdet(Sigma)
end
sum_normal_logpdf(x, mu, Sigma) = sum(normal_logpdf(x, mu, Sigma))
function normal_dmu_logpdf(x, mu, Sigma::Array{Float64,2})
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    (Sigma \ diffs' )'
    #diffs/(sigma*sigma)
end

function normal_dmu_logpdf(x, mu, Sigma::Matrix{Float64}, Sx::Matrix{Float64})
    Sx'
end

function normal_dmu_logpdf(x, mu, sigma::Float64)
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    diffs/(sigma*sigma)
end

function normal_dmu_logpdf(x, mu, L::Matrix{Float64}, D::Vector{Float64})
    diffs = get_diff(x,mu')
    diffs * L' * diagm(1./D) * L
end

function normal_dmu_logpdf(x, mu, L::Matrix{Float64}, D::Vector{Float64}, Lk_x::Matrix{Float64})
    Lk_x' * diagm(1./D)*L
end

function normal_dx_logpdf(x, mu, Sigma::Array{Float64,2})
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    -(Sigma \ diffs' )'
    #-diffs/(sigma*sigma)
end
function normal_dx_logpdf(x, mu, sigma::Float64)
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    -diffs/(sigma*sigma)
end
function normal_dSigma_logpdf(x, mu, Sigma::Array{Float64,2})
    @assert issym(Sigma)
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    invSig = inv(Sigma)
    ddT = [0.5invSig*diffs[i,:]'diffs[i,:]*invSig - 0.5invSig for i = 1:size(x,1)]
#    if abs(sum(ddT)[end,1] - sum(ddT)[1,end]) >= 10.0^-8 #should be symmetric
#        lprintln(ddT)
#        lprintln(invSig)
#        lprintln(Sigma)
#        @assert false
#    end
    for i = 1:size(x,1)
        copytri!(ddT[i],'U')
    end
    ddT
end

# derivative with respect to sigma^2
function normal_dsigma2_logpdf(x, mu, sigma2::Float64)
    diffs = get_diff(x,mu') #broadcast(-,x,mu')
    0.5*sum(diffs.*diffs)/sigma2^2 .- 0.5/sigma2
end

function normal_dSigma_logpdf(x, mu, Sigma::Array{Array{Float64,2},1})
    [normal_dSigma_logpdf(x, mu, Sigma[k]) for k = 1:length(Sigma)]
end

function normal_dSigma_logpdf(x, mu, Sigma::Array{Float64,2}, Sx::Matrix{Float64})
    @assert issym(Sigma)
    invSig = inv(Sigma)
    N = size(x,1)
    ddT = cell(N)

    for i = 1:N
        ddT[i] = fast_outer_product(Sx[:,i], Sx[:,i],0.5) #0.5Sx[:,i]*Sx[:,i]'
        ddT[i] -= 0.5invSig
        copytri!(ddT[i],'U')
    end

    ddT
end

function normal_dL_logpdf(x, mu, L::Vector{Matrix{Float64}}, D::Matrix{Float64})
    [normal_dL_logpdf(x, mu[:,k], L[k], D[:,k]) for k = 1:length(L)]
end

function normal_dL_logpdf(x, mu, L::Matrix{Float64}, D::Vector{Float64})
    diffs = get_diff(x, mu')
    N = size(x,1)
    ddT = cell(N)
    Ldiffs = diagm(1./D)*L*diffs'
    for i = 1:N
        ddT[i] = fast_outer_product(Ldiffs[:,i], diffs[i,:], -1.0)
        #ddT[i] = diagm(1./D)*L*fast_outer_product(diffs[i,:], diffs[i,:], -1.0)
        ddT[i] -= triu(ddT[i])
    end

    ddT
end

function normal_dL_logpdf(x, mu, D::Vector{Float64}, Lk_x::Matrix{Float64}, diffs::Matrix{Float64})
    N = size(x,1)
    ddT = cell(N)
    Ldiffs = diagm(1./D)*Lk_x
    for i = 1:N
        #ddT[i] = fast_outer_product(Ldiffs[:,i], diffs[i,:], -1.0)
        ddT[i] = fast_outer_product(diffs[i,:], Ldiffs[:,i], -1.0)
        ddT[i] -= triu(ddT[i])
    end

    ddT
end

function normal_dD_logpdf(x, mu, L::Vector{Matrix{Float64}}, D::Matrix{Float64})
    dD = similar(D)

    for k = 1:size(D,2)
        dD[:,k] = normal_dD_logpdf(x, mu[:,k], L[k], D[:,k])
    end
    dD
end

function normal_dD_logpdf(x, mu, L::Matrix{Float64}, D::Vector{Float64})
    diffs = get_diff(x, mu')
    N = size(x,1)
    Ldiffs = L*diffs'
    broadcast(-, (0.5 * D.^-2 .* Ldiffs.^2)', 0.5./D' )
    #(0.5 * D.^-2 .* Ldiffs.^2)'
end

function normal_dD_logpdf(x, mu, D::Vector{Float64}, Lk_x::Matrix{Float64})
    broadcast(-, (0.5 * D'.^-2 * Lk_x.^2)', 0.5./D' )
end

sum_normal_dmu_logpdf(x, mu, sigma) = sum(normal_dmu_logpdf(x, mu, sigma), 1)

standard_normal_logpdf = x -> normal_logpdf(x, zeros(size(x,2)),  eye(size(x,2)))
standard_normal_dmu_logpdf = x -> normal_dmu_logpdf(x, zeros(size(x,2)), eye(size(x,2)))

function inverse_wishart_logpdf(X::Array{Float64,2}, L, nu)
    (p,p) = size(L)
    @assert issym(X)
    if !isposdef(X)
        return -Inf
    end
    -0.5nu * logdet(L) - 0.5p*nu* lpgamma(nu/2,p) - (nu+p+1)/2 * logdet(X) - 0.5*trace(X \ L)
end
function inverse_wishart_logpdf(X::Array{Array{Float64,2},1}, L, nu)
    result = zeros(length(X))
    for i = 1:length(X)
        result[i] = inverse_wishart_logpdf(X[i],L,nu)
    end
    result
end

function inverse_wishart_dX_logpdf(X::Array{Float64,2}, L, nu)
    (p,p) = size(L)

    invX = inv(X)
    -0.5(nu+p+1)*invX' + 0.5invX*L'invX
end
function inverse_wishart_dX_logpdf(X::Array{Array{Float64,2},1}, L, nu)
    result = similar(X)
    for i = 1:length(X)
        result[i] = inverse_wishart_dX_logpdf(X[i],L,nu)
    end
    result
end

function tdist_logpdf(x, nu)
    -0.5(nu+1)*log(1+x^2/nu) + lgamma(0.5(nu+1)) - 0.5log(nu*pi) - lgamma(0.5nu)
end

function tdist_dx_logpdf(x, nu)
    -0.5(nu+1)*2x/(nu+x^2)
end

function tdist_dnu_logpdf(x, nu)
    -0.5log(1+x^2/nu) + 0.5(nu+1)/(nu^2/x^2+nu) + 0.5*digamma(0.5(nu+1)) - 0.5/(nu*pi) - 0.5*digamma(0.5*nu)
end

function inverse_gamma_logpdf(x, shape, scale)
    shape*log(scale) - lgamma(shape) - (shape+1)*log(x) - scale/x
end

function inverse_gamma_dx_logpdf(x, shape, scale)
    -(shape+1)/x + scale/x^2
end

function exponential_logpdf(x, lambda)
    log(lambda) - lambda*x
end
function exponential_dx_logpdf(x, lambda)
    -lambda
end

standard_exponential_logpdf = x -> exponential_logpdf(x, 1.0)
standard_exponential_dx_logpdf = x -> exponential_dx_logpdf(x, 1.0)

function gamma_logpdf(x::Float64, shape, scale)
    if x < 0.0 || isinf(x) || isnan(x)
        return -Inf
    elseif x == 0.0 && shape == 1
        return -x/scale - lgamma(shape) - shape*log(scale)
    else
        return (shape-1)*log(x) - x/scale - lgamma(shape) - shape*log(scale)
    end
end

function gamma_logpdf(x, shape, scale)
    result = zeros(size(x))
    for i in 1:length(x)
        result[i] = gamma_logpdf(x[i], shape, scale)
    end
    result
end
standard_gamma_logpdf(x) = gamma_logpdf(x, 1.0, 1.0)

function gamma_dx_logpdf(x, shape, scale)
    if shape == 1
        return ones(size(x))/scale
    else
        return (shape-1)./x .- 1/scale
    end
end
standard_gamma_dx_logpdf(x) = gamma_dx_logpdf(x, 1.0, 1.0)

function exp_gamma_logpdf(x, shape, scale)
    y = exp(x)
    gamma_logpdf(y, shape, scale)
end

function exp_gamma_dx_logpdf(x, shape, scale)
    y = exp(x)
    gamma_dx_logpdf(y, shape, scale).*y
end

function beta_logpdf(x::Float64, alpha, beta)
    if x < 0.0 || x > 1.0
        return -Inf
    else
        return (alpha-1)log(x)+(beta-1)log(1-x) - lbeta(alpha,beta)
    end
end

function beta_logpdf(x, alpha, beta)
    result = zeros(size(x))
    for i = 1:length(x)
        result[i] = beta_logpdf(x[i], alpha, beta)
    end
    result
end

function beta_dx_logpdf(x, alpha, beta)
    (alpha-1)/x - (beta-1)/(1-x)
end

function bernoulli_logpdf(x, p)
    if p < 0.0 || p > 1.0
        return -Inf*ones(size(x))
    end
    x*log(p) + (1-x)*log(1-p)
end

function bernoulli_dp_logpdf(x, p)
    x/p - (1-x)/(1-p)
end

function beta_binomial_logpdf(k, n, a, b)
    lgamma(n+1) - lgamma(k+1) - lgamma(n-k+1) + lbeta(k+a,n-k+b) - lbeta(a,b)
end

# we need this one for binary vector states (as opposed to counts)
function beta_bernoulli_logpdf(k, n, a, b)
    lbeta(k+a,n-k+b) - lbeta(a,b)
end

function poisson_logpdf(k, lambda)
    k*log(lambda) - lambda - lfact(k)
end

function log_logistic{T<:Real}(effect::T, y::T)
    value = 0.0
    value = -(one(T)-y)*effect - log(one(T)+exp(-effect))
    if isinf(value)
        value = 0.0
    end
    return value
end

function log_logistic_dx(effect, y)
    -(1-y) + exp(-effect)/(1+exp(-effect))
end

#if !isdefined(:BinaryFunctor)
#    typealias BinaryFunctor Functor{2}
#end
#
#type LogLogistic <: BinaryFunctor end
#evaluate(::LogLogistic, x::Number, y::Number) = log_logistic(x,y)
#result_type(::LogLogistic, t1::Type, t2::Type) = promote_type(t1,t2)
#
#function log_predictive(effect)
#    -log(1.+exp(-effect))
#end

# Probability Manipulation
function logsumexp(x)
    max_x = maximum(x)
    xp = x - max_x
    log(sum(exp(xp))) + max_x
end

function logsumexp(x,d)
    max_x = maximum(x,d)
    xp = x .- max_x
    log(sum(exp(xp),d)) + max_x
end

function weighted_logsumexp(x,w)
    if all( x .== -Inf) && all(w .>= 0.0)
        return -Inf
    end

    max_x = max(x)
    xs = x - max_x

    log(sum(w.*exp(xs))) + max_x
end

# xp should be a cell array of dictionaries, so that xp[i][k] gives the gradient
# of parameter k in the ith RTJ mixture component
function logsumexp_d_dx(x,xp)
    #zeros in shape of xp[1]
    xout = Dict{Any,Any}(k => xp[1][k]-xp[1][k] for k = keys(xp[1]))

    max_x = maximum(x)
    xs = x .- max_x

    exp_xs = exp(xs)
    sum_exp_xs = sum(exp_xs)
    for i = 1:length(xp)
        for k = keys(xout)
            xout[k] += exp_xs[i]*xp[i][k]
        end
    end

    for k = keys(xout)
        xout[k] /= sum_exp_xs
    end

    #sum(exp(xs).*xp) / sum(exp(xs))
    xout
end

function logsumexp_d_dx(x,xp::Vector{Vector{Float64}})
    #zeros in shape of xp[1]
    xout = xp[1]-xp[1]

    max_x = maximum(x)
    xs = x .- max_x

    exp_xs = exp(xs)
    sum_exp_xs = sum(exp_xs)
    for i = 1:length(xp)
        xout += exp_xs[i]*xp[i]
    end

    xout /= sum_exp_xs

    #sum(exp(xs).*xp) / sum(exp(xs))
    xout
end

function logsumexp_d_dx(x,xp::Vector{Float64})
    #zeros in shape of xp[1]
    xout = xp[1]-xp[1]

    max_x = maximum(x)
    xs = x .- max_x

    exp_xs = exp(xs)
    sum_exp_xs = sum(exp_xs)
    for i = 1:length(xp)
        xout += exp_xs[i]*xp[i]
    end

    xout /= sum_exp_xs

    #sum(exp(xs).*xp) / sum(exp(xs))
    xout
end

function logsumexp_d_dx(x,xp::Array{Vector{Float64}})
    #zeros in shape of xp[1]
    xout = zeros(length(xp[1]))

    max_x = maximum(x)
    xs = x .- max_x

    exp_xs = exp(xs)
    sum_exp_xs = sum(exp_xs)
    for i = 1:length(xp)
        xout += exp_xs[i]*xp[i]
    end

    xout /= sum_exp_xs

    xout
end

function exp_normalize(x)
    xp = x .- maximum(x)
    exp_x = exp(xp)
    exp_x / sum(exp_x)
end

# Basic Sampling

function randmult(x)
    v = cumsum(x)
    if abs(v[end] - 1.0) >= 10.0^-8
        lprintln(x)
        lprintln(v)
    end
    assert( abs(v[end] - 1.0) < 10.0^-8)
#    if v[end] != 1.0
#        lprintln("v[end]: ", v[end])
#        assert(v[end] == 1.0)
#    end

    u = rand()
    i = 1
    while u > v[i]
        i += 1
    end
    i
end

function randpois(lambda)
    L = exp(-lambda)
    k = 0
    p = 1.0
    while p > L
        k += 1
        p = p * rand()
    end
    k - 1
end

function lpgamma(x, p)
    result = p*(p-1)/2 * log(pi)

    for i = 1:p
        result += lgamma(x + (1-i)/2)
    end
    result
end

# Faster outer product
function fast_outer_product(a,b,w)
    c = zeros(length(a),length(b))
    for i = 1:length(a)
        for j = 1:length(b)
            c[i,j] = w*a[i]*b[j]
        end
    end
    c
end

# Chinese Restaurant Table

function log_CRT(m,r,l::Real)
    P1_cache = calculate_P1_cache(m)
    lcrt = log_CRT_cache(m,r,l,P1_cache)
    (lcrt, P1_cache)
end

function log_CRT(m,r, P1_cache::Array{Float64,2})
    lcrt = zeros(m+1)
    for l = 0:m
        lcrt[l+1] = log_CRT_cache(m,r,l,P1_cache)
    end
    lcrt
end

function log_CRT_cache(m,r,l,P1_cache)
    P1_cache[m+1,l+1] + l*log(r) + lgamma(r) - lgamma(m+r) + lgamma(m+1)
end

function calculate_P1_cache(m)
    P1_cache = zeros(m+1,m+1)

    P1_cache[1,1] = 1

    for n = 1:m
        P1_cache[n+1,2] = log((n-1)/n) + P1_cache[n,2]
        for l = 2:n-1
            P1_cache[n+1,l+1] = log((n-1)/n) + P1_cache[n,l+1] + log(1+exp(P1_cache[n,l]) - P1_cache[n,l+1] - log(n-1))
        end
        P1_cache[n+1,n+1] = P1_cache[n,n] - log(n)
    end

    P1_cache
end

function get_diff(x::Float64,y::Float64)
    x-y
end

function get_diff(x,y::Float64)
    x.-y
end

function get_diff(x::Float64,y)
    x.-y
end

function get_diff(x,y)
    broadcast(-,x,y)
end

function tril_inds(th, k)
    A = ones(size(th))
    A = tril(A,k)
    find(A.==1)
end