text:
"""
    refine_recursive_images(points, image_points, simplex_inds, maxsize, k)

Refine a triangulation. Consider both simplices and their images.
Continue splitting some fraction of the simplices
with size reduction factor of k until until the maximum simplex size is below
some reference size `refsize` (for example, the maximum simplex size in the
original triangulation).

Returns a tuple containing the refined triangulation and some information about
it, as follows:

(points, image_points, simplex_inds, centroids, centroids_im, radii, radii_im, simplexvolumes, imagevolumes)
"""
function refine_recursive_images(points, image_points, simplex_inds, maxsize, k; niter = 1)
    if niter == 1
        println(size(points, 2), "D triangulation containing ", size(simplex_inds, 1), " simplices.")
        print("Target simplex radius = ", round(maxsize, 4), ". ")
        println("Split simplices into subsimplices with radii reduced by  factor of ", k, ".")
    end

    centroids, radii = centroids_radii2(points, simplex_inds)
    centroids_im, radii_im = centroids_radii2(image_points, simplex_inds)
    # If convergence is reached, return the triangulation and the centroids and radii of
    # the simplices furnishing the triangulation
    maxradius = max(maximum(radii), maximum(radii_im))
    print("Refinement #", niter, "\tMaximum simplex radius = ", floor(maxradius, 4),  ". ")

    if maxradius < maxsize || niter > 100
        if niter > 100
            println("\n\t\tMaximum number of iterations (100) reached. Ending refinement.")
        else
            println("\n\t\tRefinement finished.\n")
        end

        simplexvolumes = simplex_volumes(points, simplex_inds)
        imagevolumes = simplex_volumes(image_points, simplex_inds)
        return points, image_points, simplex_inds, centroids, centroids_im, radii, radii_im,
                simplexvolumes, imagevolumes
    end

    # If no convergence, continue splitting some fraction of the largest simplices.
    if length(radii) == 1
        percentile = 1.0
        split_indices = [1]
    else
        # What portion of the largest simplices to split (1 - percentile)?
        percentile = 0.999
        split_indices = []

        # If the simplices in the triangulation are very regular, there might not be
        # any simplices larger than a given percentile. If so, reduce percentile iteratively
        # until there are simplices larger than the percentile.
        while length(split_indices) == 0
            # Check if all elements are approximately the same size. If so, we must split all
            # of them.
            if all(y-> isapprox(radii[1], y), radii)
                split_indices =  find(radii .> 0)
            elseif all(y-> isapprox(radii[1], y), radii_im)
                split_indices =  find(radii_im .> 0)
            else
                split_indices = find(radii .> quantile(radii, percentile))
                split_indices_im = find(radii_im .> quantile(radii, percentile))
                split_indices = unique([split_indices; split_indices_im])

            end
            percentile = percentile - 0.01
        end
    end

    # The number of simplices to split
    n_split_simplices = length(split_indices)

    # The dimension of the space
    E = size(points, 2)

    # Generic rules for splitting a simplex
    splitting_rules = simplicial_subdivision(k, E)

    # Rules for forming the strictly new vertices of the subtriangulation
    rules = splitting_rules[1]

    # Array where each row represents one of the new simplices in the splitted simplex
    subtriangulation = splitting_rules[2]

    # How many new vertices are created each split?
    n_newvertices_eachsplit = size(rules, 1)

    println(" Splitting ", n_split_simplices, "/", size(simplex_inds, 1), " (", @sprintf("%.4f", 1.0 - percentile), " %)", " simplices and their corresponding images into ", n_split_simplices * n_newvertices_eachsplit - 1, " subsimplices.")

    # We need an array that can accomodate all of them. Each row in this
    # array will be a new vertex. Stacks of n_newvertices_eachsplit * E arrays.
    # We have as many as we have simplices to split.
    new_points = zeros(n_newvertices_eachsplit * n_split_simplices, E)
    new_imagepoints = zeros(n_newvertices_eachsplit * n_split_simplices, E)

    # Fill the array by looping over all simplices that we want to split
    for i = 1:n_split_simplices
        # Figure out what the row indices corresponding to the ith simplex
        # must be. Marks the beginning of each of the simplex stacks in new_points
        ind = n_newvertices_eachsplit * (i - 1)

        # Index of the simplex we need to split
        simplex_idx = split_indices[i]

        # Get the vertices of the simplex currently being splitted. Each of the
        # n_newvertices_eachsplit new vertices will be a linear combination
        # of these vertices. Each row is a vertex.
        vertices = points[simplex_inds[simplex_idx, :], :]
        imagevertices = image_points[simplex_inds[simplex_idx, :], :]

        # Generate the strictly new vertices for each sub
        for j = 1:n_newvertices_eachsplit
            # The index for a particular new vertex of the ith new simplex
            ind_newvertex = ind + j

            # Compute the jth vertex of the subtriangulation of the ith simplex

            # Go the jth new subsimplex of the ith splitted simplex. The entries of this
            # vector
            jth_subsimplex = rules[j, :]

            # Pick the corresponding original vertices with indices contained in rules[j, :]
            original_vertices = vertices[rules[j, :], :]
            original_vertices_image = imagevertices[rules[j, :], :]

            new_points[ind_newvertex, :] = sum(original_vertices, dims=1) ./ k
            new_imagepoints[ind_newvertex, :] = sum(original_vertices_image, dims=1) ./ k

        end
    end

    # Find the unique new vertices
    new_points_noreps = unique(new_points, 1)
    new_imagepoints_noreps = Array{Float64}(size(new_points_noreps, 1), E)

    Ind = Vector{Int}(size(new_points, 1))
    for i = 1:size(new_points_noreps, 1)

        count = 0

        for j = 1:size(new_points, 1) # run over the possibly repeated vertices
            if new_points_noreps[i, :] == new_points[j, :]
                Ind[j] = i

                count = count + 1
                if count == 1
                    new_imagepoints_noreps[i, :] = new_imagepoints[j, :]
                end
            end
        end
    end

    # TEST:  assert all zeros.
    #@show unique(new_imagepoints, 1) - new_imagepoints_noreps

    # Combine old and newly introduced vertices
    num_vertices_beforesplit = size(points, 1)
    points = vcat(points, new_points_noreps)
    image_points = vcat(image_points, new_imagepoints_noreps)

    # Update the Ind array, so that we start at the new vertices
    Ind = Ind + num_vertices_beforesplit
    num_simplices_each_split = size(subtriangulation, 1)

    # The subsimplices formed by the splitting. Each row contains E + 1 indices referencing
    # the vertices furnishing that particular subsimplex (now found in the updated
    # points array).
    newtriangulation = Array{Float64}(num_simplices_each_split * n_split_simplices, E + 1)

    # For each simplex that we need to split,
    for i = 1:n_split_simplices
        # The beginning of the stack we need to fill.
        index = num_simplices_each_split * (i - 1)

        # Figure out what the row indices corresponding to the ith simplex
        # must be. Marks the beginning of each of the simplex stacks in new_points
        ind = n_newvertices_eachsplit * (i - 1)

        # Index of the simplex we need to split
        simplex_idx = split_indices[i]

        # Pick the indices of the original vertices. Should be a column vector.
        inds_original_vertices = simplex_inds[simplex_idx, :]

        # Indices of the new vertices. Should be a column vector
        inds_new_points = Ind[(ind + 1):(ind + n_newvertices_eachsplit)]
        inds_all_vertices = vcat(inds_original_vertices, inds_new_points)

        # Go through each of the subsimplices formed by splitting that simplex
        for j = 1:num_simplices_each_split
            newtriangulation[index + j, :] = inds_all_vertices[subtriangulation[j, :]]
        end
    end

    # Indices of the simplices that are not split
    untouched_indices = complementary(split_indices, size(simplex_inds, 1))

    simplex_inds = round.(Int, vcat(simplex_inds[untouched_indices, :],
                                  newtriangulation))

    refine_recursive_images(points, image_points, simplex_inds, maxsize, k; niter = niter + 1)
end


text:
export move, add_tile, show, initbboard, move_and_reward, score

import Base.show, Base.display, Base.maximum

const ROWMASK = UInt16(2^16 - 1)
const CELLMASK = UInt16(2^4 - 1)
const MASK = UInt8(15)

"""
Bsed on the benchmarks I have done. It's quicker to precompute the results for all rows and make a
set of lookup table that looks up the result of moving left or right.

One of the Rust repositories have taken this approach.
"""
const LEFT, LEFT_REWARD, RIGHT, RIGHT_REWARD = make_row_lookup()

"""
A gameboard is stored on a UInt (64bit)
where each row (TODO shall I make it column major) is stored using 16 bits. So 4 bits for each cell.
So each cell can contain values form 0-15. 2^15 = 32678 and 2^11 is 2048. So it's there's plenty of
room
"""
struct Bitboard
    board::UInt64
end

"""
initialise a bitboard
"""
function initbboard()
    zero(UInt64) |> Bitboard |> add_tile |> add_tile
end

"""
    board2cols(board::UInt64)::Vector{Vector{UInt16}}

Return the columns from a board
"""
function board2cols(bitboard::Bitboard)::Vector{UInt16}
    board = bitboard.board
    # initialise the columns
    cols = zeros(UInt16, 4)

    # take one row at a time
    for row_shift in 48:-16:0
        row = (board >> row_shift) & ROWMASK

        # populate the right cell for each column
        @inbounds for i in 4:-1:1
            cols[i] <<= 4 # doing this here effectively means each col only gets shifted 3 times
            cols[i] |= row & CELLMASK
            row >>= 4
        end
    end
    cols
end

function Base.rotl90(b::Bitboard)
    board = b.board
    # initialise the columns
    # cols = zeros(UInt16, 4)
    # cols = zeros(SVector{4, UInt16})
    col1 = zero(UInt16)
    col2 = zero(UInt16)
    col3 = zero(UInt16)
    col4 = zero(UInt16)

    # take one row at a time
    for row_shift in 48:-16:0
        row = (board >> row_shift) & ROWMASK

        # populate the right cell for each column
        col4 <<= 4 # doing this here effectively means each col only gets shifted 3 times
        col4 |= row & CELLMASK
        row >>= 4

        col3 <<= 4 # doing this here effectively means each col only gets shifted 3 times
        col3 |= row & CELLMASK
        row >>= 4

        col2 <<= 4 # doing this here effectively means each col only gets shifted 3 times
        col2 |= row & CELLMASK
        row >>= 4

        col1 <<= 4 # doing this here effectively means each col only gets shifted 3 times
        col1 |= row & CELLMASK
        row >>= 4
    end

    newboard = zero(UInt64)
    # construct the new board by putting the columns in to rows
    newboard |= col4
    newboard <<= 16
    newboard |= col3
    newboard <<= 16
    newboard |= col2
    newboard <<= 16
    newboard |= col1

    Bitboard(newboard)
end

"""
    move(board::UInt64, dir::Dirs)::UInt64
    move(board::UInt64, LOOKUP::Vector{UInt16}::UInt64

Bitboard game
"""
function move(bitboard::Bitboard, LOOKUP::Vector{UInt16})::Bitboard
    board = bitboard.board

    idx = board >> 48
    new_board = zero(UInt64)
    @inbounds new_board |= LOOKUP[idx+1]

    for i in 32:-16:0
        new_board <<= 16
        idx = (board >> i) & ROWMASK
        @inbounds new_board |= LOOKUP[idx+1]
    end
    Bitboard(new_board)
end

function move_and_reward(bitboard::Bitboard, LOOKUP::Vector{UInt16}, REWARDS)::Bitboard
    board = bitboard.board

    idx = board >> 48
    new_board = zero(UInt64)
    @inbounds new_board |= LOOKUP[idx+1]

    rewards = 0
    rewards += REWARDS[idx+1]

    for i in 32:-16:0
        new_board <<= 16
        idx = (board >> i) & ROWMASK
        @inbounds new_board |= LOOKUP[idx+1]
        rewards += REWARDS[idx+1]
    end
    Bitboard(new_board), rewards
end

function make_column_lookup_up()
    lookup_up = zeros(UInt64, 2^16)
    for u in 1:2^16
        u_moved_up = zero(UInt64)
        u_moved_left = LEFT[u]
        for i in 1:4
            tmp = UInt64((u_moved_left << 4(i - 1)) >> 12)
            u_moved_up |= (tmp << (60 - 16(i - 1)))
        end
        lookup_up[u] = u_moved_up
    end
    lookup_up
end

function make_column_lookup_down()
    lookup_up = zeros(UInt64, 2^16)
    for u in 1:2^16
        u_moved_down = zero(UInt64)
        u_moved_right = RIGHT[u]
        for i in 1:4
            tmp = UInt64((u_moved_right << 4(i - 1)) >> 12)
            u_moved_down |= (tmp << (60 - 16(i - 1)))
        end
        lookup_up[u] = u_moved_down
    end
    lookup_up
end

const UP = make_column_lookup_up()
const DOWN = make_column_lookup_down()

# Credit to Nneonneo
# equivalent to rotate left 90
function transpose(bb::Bitboard)
    x = bb.board
    a1 = x & 0xF0F00F0FF0F00F0F
    a2 = x & 0x0000F0F00000F0F0
    a3 = x & 0x0F0F00000F0F0000
    a = a1 | (a2 << 12) | (a3 >> 12)
    b1 = a & 0xFF00FF0000FF00FF
    b2 = a & 0x00FF00FF00000000
    b3 = a & 0x00000000FF00FF00
    return Bitboard(b1 | (b2 >> 24) | (b3 << 24))
end

function move_updown(bitboard::Bitboard, up_or_down)
    tbb = transpose(bitboard)
    moved = zero(UInt64)
    for i in 1:4
        c = (tbb.board << 16(i - 1)) >> 48
        moved |= (up_or_down[c+1] >> 4(i - 1))
    end
    Bitboard(moved)
end

function move(bitboard::Bitboard, dir::Dirs, skip_col_to_row = false)::Bitboard
    if dir == left
        new_bitboard = move(bitboard, LEFT)
    elseif dir == right
        new_bitboard = move(bitboard, RIGHT)
    elseif dir == up
        new_bitboard = move_updown(bitboard, UP)
    else
        new_bitboard = move_updown(bitboard, DOWN)
    end

    new_bitboard
end

function move_and_reward(bitboard::Bitboard, dir::Dirs)
    if dir == left
        new_bitboard = move_and_reward(bitboard, LEFT, LEFT_REWARD)
    elseif dir == right
        new_bitboard = move_and_reward(bitboard, RIGHT, RIGHT_REWARD)
    elseif dir == up
        new_bitboard = move_and_reward(transpose(bitboard), LEFT, LEFT_REWARD)
    else
        new_bitboard = move_and_reward(transpose(bitboard), RIGHT, RIGHT_REWARD)
    end

    new_bitboard
end

"""
Makes moves until not possible
"""
function move(bitboard::Bitboard, dirs::Vector{Dirs})::Bitboard
    for dir in dirs
        new_bitbaord = move(bitboard, dir)
        if new_bitbaord != bitboard
            return new_bitbaord
        end
    end
    return bitboard
end

"""counts how many empty spots there are in the bitboard"""
function count0(bitboard::Bitboard)
    board = bitboard.board
    # firstly count how many empty spots there are
    cnt_empty = 0

    for shift in 0:4:60
        cnt_empty += ((board >> shift) & CELLMASK) == 0
    end

    cnt_empty
end

"""Randomly add a 1-tile or a 2-tile"""
function add_tile(bitboard::Bitboard)::Bitboard
    cnt_empty = count0(bitboard)
    if cnt_empty == 0
        return bitboard
    end

    selected_empty_cell_num = rand(1:cnt_empty)

    add_tile(bitboard, selected_empty_cell_num, rand() < 0.1 ? 2 : 1)
end

"""Add a specified 1-tile or a 2-tile to selected'th empty cell"""
function add_tile(bitboard::Bitboard, selected::Integer, two_or_four)
    board = bitboard.board
    i = 0
    for shift in 0:4:60
        if ((board >> shift) & CELLMASK) == 0
            i += 1
            if i == selected
                board += two_or_four << shift
                return Bitboard(board)
            end
        end
    end
    Bitboard(board)
end

"""Convert a bitboard_to_array"""
function bitboard_to_array(bitboard::Bitboard)::Array{Int8,2}
    board = bitboard.board

    outboard = Array{Int8,2}(undef, 4, 4)

    rowid = 1
    # take one row at a time
    for row_shift in 48:-16:0
        row = (board >> row_shift) & ROWMASK

        # populate the right cell for each column
        for colid in 4:-1:1
            outboard[rowid, colid] = row & CELLMASK
            row >>= 4
        end
        rowid += 1
    end
    outboard
end

function Base.show(io::IO, bitboard::Bitboard)
    show(io, bitboard_to_array(bitboard))
end

function Base.display(bitboard::Bitboard)
    display(bitboard_to_array(bitboard))
end

function Base.maximum(bb::Bitboard)
    maximum(0:4:60) do s
        (bb.board >> s) & MASK
    end
end

function score2(bb::Bitboard)
    mapreduce(+, 0:4:60) do s
        1 << ((bb.board >> s) & MASK)
    end
end

function make_score_lookup()
    u = UInt64.(0:15)
    res = Vector{Int32}(undef, 2^16)
    for (i1, i2, i3, i4) in Iterators.product(u, u, u, u)
        bitboard_row = (i1 << 12) | (i2 << 8) | (i3 << 4) | i4
        bb = Bitboard(bitboard_row)
        @inbounds res [bitboard_row + 1] = score2(bb) - count0(bb)
    end
    res
end

const SCORE = make_score_lookup()

function score(bitboard::Bitboard)
    mapreduce(+, 0:16:48) do s
        idx = (bitboard.board >> s) & ROWMASK
        @inbounds SCORE[idx+1]
    end
end

ext: jl


text:
function calc_rho( Focc, v::Array{Float64,2} )
    Ncols   = size(v)[2]
    Npoints = size(v)[1]
    rho = zeros(Float64,Npoints)
    for ic = 1:Ncols
        rho[:] = rho[:] + v[:,ic] .* v[:,ic] * Focc[ic]
    end
    return rho
end


text:
module QuantumLab

print(" + (ProgressMeter................")
using ProgressMeter
ProgressMeter.printover(STDOUT," + (Reexport...................")
using Reexport
ProgressMeter.printover(STDOUT,"")

print(" + BaseModule..................")
include("SubModules/BaseModule.jl")
@reexport using .BaseModule
println("Done.")

print(" + AtomModule..................")
include("SubModules/AtomModule.jl")
@reexport using .AtomModule
println("Done.")

print(" + GeometryModule..............")
include("SubModules/GeometryModule.jl")
@reexport using .GeometryModule
println("Done.")

print(" + BasisSetModule..............")
include("SubModules/BasisSetModule.jl")
@reexport using .BasisSetModule
println("Done.")

print(" + BasisFunctionsModule........")
include("SubModules/BasisFunctionsModule.jl")
@reexport using .BasisFunctionsModule
println("Done.")

print(" + IntegralsModule.............")
include("SubModules/IntegralsModule.jl")
@reexport using .IntegralsModule
println("Done.")

print(" + BasisModule.................")
include("SubModules/BasisModule.jl")
@reexport using .BasisModule
println("Done.")

print(" + SpecialMatricesModule.......")
include("SubModules/SpecialMatricesModule.jl")
@reexport using .SpecialMatricesModule
println("Done.")

end # module

max_stars_repo_name: tkelman/QuantumLab

max_issues_repo_name: tkelman/QuantumLab

max_forks_repo_name: tkelman/QuantumLab


text:
if Pair != Base.Pair
import Base: Base, IOContext, string, join, sprint
IOContext(io::IO, KV::Pair) = IOContext(io, Base.Pair(KV[1], KV[2]))
length(s::String) = Base.length(s)
^(s::String, i::Int) = Base.:^(s, i)
end
isexpr(e::Expr, s::Symbol) = e.head === s
isexpr(@nospecialize(e), s::Symbol) = false

function Base.show(io::IO, cfg::CFG)
    foreach(pairs(cfg.blocks)) do (idx, block)
        Base.println("$idx\t=>\t", join(block.succs, ", "))
    end
end

print_ssa(io::IO, val::SSAValue, argnames) = Base.print(io, "%$(val.id)")
print_ssa(io::IO, val::Argument, argnames) = Base.print(io, isempty(argnames) ? "%%$(val.n)" : "%%$(argnames[val.n])")
print_ssa(io::IO, val::GlobalRef, argnames) = Base.print(io, val)
print_ssa(io::IO, @nospecialize(val), argnames) = Base.show(io, val)

function print_node(io::IO, idx, stmt, used, argnames, maxsize; color = true, print_typ=true)
    if idx in used
        pad = " "^(maxsize-length(string(idx)))
        Base.print(io, "%$idx $pad= ")
    else
        Base.print(io, " "^(maxsize+4))
    end
    if isa(stmt, PhiNode)
        args = map(1:length(stmt.edges)) do i
            e = stmt.edges[i]
            v = !isassigned(stmt.values, i) ? "#undef" :
                sprint() do io′
                    print_ssa(io′, stmt.values[i], argnames)
                end
            "$e => $v"
        end
        Base.print(io, "φ ", '(', join(args, ", "), ')')
    elseif isa(stmt, PhiCNode)
        Base.print(io, "φᶜ ", '(', join(map(x->sprint(print_ssa, x, argnames), stmt.values), ", "), ')')
    elseif isa(stmt, PiNode)
        Base.print(io, "π (")
        print_ssa(io, stmt.val, argnames)
        Base.print(io, ", ")
        if color
            Base.printstyled(io, stmt.typ, color=:cyan)
        else
            Base.print(io, stmt.typ)
        end
        Base.print(io, ")")
    elseif isa(stmt, UpsilonNode)
        Base.print(io, "ϒ (")
        isdefined(stmt, :val) ?
            print_ssa(io, stmt.val, argnames) :
            Base.print(io, "#undef")
        Base.print(io, ")")
    elseif isa(stmt, ReturnNode)
        if !isdefined(stmt, :val)
            Base.print(io, "unreachable")
        else
            Base.print(io, "return ")
            print_ssa(io, stmt.val, argnames)
        end
    elseif isa(stmt, GotoIfNot)
        Base.print(io, "goto ", stmt.dest, " if not ")
        print_ssa(io, stmt.cond, argnames)
    elseif isexpr(stmt, :call)
        print_ssa(io, stmt.args[1], argnames)
        Base.print(io, "(")
        Base.print(io, join(map(arg->sprint(io->print_ssa(io, arg, argnames)), stmt.args[2:end]), ", "))
        Base.print(io, ")")
        if print_typ && stmt.typ !== Any
            Base.print(io, "::$(stmt.typ)")
        end
    elseif isexpr(stmt, :invoke)
        print(io, "invoke ")
        linfo = stmt.args[1]
        print_ssa(io, stmt.args[2], argnames)
        Base.print(io, "(")
        sig = linfo.specTypes === Tuple ? () : Base.unwrap_unionall(linfo.specTypes).parameters
        print_arg(i) = sprint() do io
            print_ssa(io, stmt.args[2+i], argnames)
            if (i + 1) <= length(sig)
                print(io, "::$(sig[i+1])")
            end
        end
        Base.print(io, join((print_arg(i) for i=1:(length(stmt.args)-2)), ", "))
        Base.print(io, ")")
        if print_typ && stmt.typ !== Any
            Base.print(io, "::$(stmt.typ)")
        end
    elseif isexpr(stmt, :new)
        Base.print(io, "new(")
        Base.print(io, join(String[sprint(io->print_ssa(io, arg, argnames)) for arg in stmt.args], ", "))
        Base.print(io, ")")
    else
        Base.print(io, stmt)
    end
end

function compute_inlining_depth(linetable, iline)
    depth = 0
    while iline != 0
        linetable[iline].inlined_at == 0 && break
        depth += 1
        iline = linetable[iline].inlined_at
    end
    depth
end

function should_print_ssa_type(node)
    if isa(node, Expr)
        return !(node.head in (:gc_preserve_begin, :gc_preserve_end))
    end
    return !isa(node, PiNode)   && !isa(node, GotoIfNot) &&
           !isa(node, GotoNode) && !isa(node, ReturnNode)
end

function default_expr_type_printer(io, typ)
    typ_str = try
        string(typ)
    catch
        "<error_printing>"
    end
    Base.printstyled(io, "::$(typ_str)", color=:cyan)
end

function compute_loc_stack(code, line)
    stack = []
    line === 0 && return stack
    inlined_at = code.linetable[line].inlined_at
    if inlined_at != 0
        push!(stack, inlined_at)
        entry = code.linetable[inlined_at]
        while entry.inlined_at != 0
            push!(stack, entry.inlined_at)
            entry = code.linetable[entry.inlined_at]
        end
        reverse!(stack)
    end
    push!(stack, line)
end

"""
    Compute line number annotations for an IRCode

This functions compute three sets of annotations for each IR line. Take the following
example (taken from `@code_typed sin(1.0)`):

```
    **                                                    ***         **********
    35 6 ── %10  = :(Base.mul_float)(%%2, %%2)::Float64   │╻╷         sin_kernel
       │    %11  = :(Base.mul_float)(%10, %10)::Float64   ││╻          *
```

The three annotations are indicated with `*`. The first one is the line number of the
active function (printed once whenver the outer most line number changes). The second
is the inlining indicator. The number of lines indicate the level of nesting, with a
half-size line (╷) indicating the start of a scope and a full size line (│) indicating
a continuing scope. The last annotation is the most complicated one. It is a heuristic
way to print the name of the entered scope. What it attempts to do is print the outermost
scope that hasn't been printed before. Let's work a number of examples to see the impacts
and tradeoffs involved.

```
f() = leaf_function() # Delibarately not defined to end up in the IR verbatim
g() = f()
h() = g()
top_function() = h()
```

After inlining, we end up with:
```
1 1 ─ %1 = :(Main.leaf_function)()::Any   │╻╷╷ h
  └──      return %1                      │
```

We see that the only function printed is the outermost function. This certainly loses
some information, but the idea is that the outermost function would have the most
semantic meaning (in the context of the function we're looking at).

On the other hand, let's see what happens when we redefine f:
```
function f()
    leaf_function()
    leaf_function()
    leaf_function()
end
```

We get:
```
1 1 ─      :(Main.leaf_function)()::Any   │╻╷╷ h
  │        :(Main.leaf_function)()::Any   ││┃│  g
  │   %3 = :(Main.leaf_function)()::Any   │││┃   f
  └──      return %3                      │
```

Even though we were in the `f` scope since the first statement, it tooks us two statements
to catch up and print the intermediate scopes. Which scope is printed is indicated both
by the indentation of the method name and by an increased thickness of the appropriate
line for the scope.
"""
function compute_ir_line_annotations(code::IRCode)
    loc_annotations = String[]
    loc_methods = String[]
    loc_lineno = String[]
    cur_group = 1
    last_line = 0
    last_lineno = 0
    last_stack = []
    last_printed_depth = 0
    for idx in eachindex(code.stmts)
        buf = IOBuffer()
        line = code.lines[idx]
        depth = compute_inlining_depth(code.linetable, line)
        iline = line
        lineno = 0
        loc_method = ""
        print(buf, "│")
        if line !== 0
            stack = compute_loc_stack(code, line)
            lineno = code.linetable[stack[1]].line
            x = min(length(last_stack), length(stack))
            if length(stack) != 0
                # Compute the last depth that was in common
                first_mismatch = findfirst(i->last_stack[i] != stack[i], 1:x)
                # If the first mismatch is the last stack frame, that might just
                # be a line number mismatch in inner most frame. Ignore those
                if length(last_stack) == length(stack) && first_mismatch == length(stack)
                    last_entry, entry = code.linetable[last_stack[end]], code.linetable[stack[end]]
                    if last_entry.method == entry.method && last_entry.file == entry.file
                        first_mismatch = nothing
                    end
                end
                last_depth = coalesce(first_mismatch, x+1)-1
                if min(depth, last_depth) > last_printed_depth
                    printing_depth = min(depth, last_printed_depth + 1)
                    last_printed_depth = printing_depth
                elseif length(stack) > length(last_stack) || first_mismatch != nothing
                    printing_depth = min(depth, last_depth + 1)
                    last_printed_depth = printing_depth
                else
                    printing_depth = 0
                end
                stole_one = false
                if printing_depth != 0
                    for _ in 1:(printing_depth-1)
                        print(buf, "│")
                    end
                    if printing_depth <= last_depth-1 && first_mismatch === nothing
                        print(buf, "┃")
                        for _ in printing_depth+1:min(depth, last_depth)
                            print(buf, "│")
                        end
                    else
                        stole_one = true
                        print(buf, "╻")
                    end
                else
                    for _ in 1:min(depth, last_depth)
                        print(buf, "│")
                    end
                end
                print(buf, "╷"^max(0,depth-last_depth-stole_one))
                if printing_depth != 0
                    if length(stack) == printing_depth
                        loc_method = String(code.linetable[line].method)
                    else
                        loc_method = String(code.linetable[stack[printing_depth+1]].method)
                    end
                end
                loc_method = string(" "^printing_depth, loc_method)
            end
            last_stack = stack
            entry = code.linetable[line]
        end
        push!(loc_annotations, String(take!(buf)))
        push!(loc_lineno, (lineno != 0 && lineno != last_lineno) ? string(lineno) : "")
        push!(loc_methods, loc_method)
        last_line = line
        (lineno != 0) && (last_lineno = lineno)
    end
    (loc_annotations, loc_methods, loc_lineno)
end

Base.show(io::IO, code::IRCode) = show_ir(io, code)
function show_ir(io::IO, code::IRCode, expr_type_printer=default_expr_type_printer; argnames=Symbol[], verbose_linetable=false)
    (lines, cols) = displaysize(io)
    used = IdSet{Int}()
    foreach(stmt->scan_ssa_use!(push!, used, stmt), code.stmts)
    cfg = code.cfg
    max_bb_idx_size = length(string(length(cfg.blocks)))
    bb_idx = 1
    if any(i->!isassigned(code.new_nodes, i), 1:length(code.new_nodes))
        printstyled(io, :red, "ERROR: New node array has unset entry\n")
    end
    new_nodes = code.new_nodes[filter(i->isassigned(code.new_nodes, i), 1:length(code.new_nodes))]
    foreach(nn -> scan_ssa_use!(push!, used, nn.node), new_nodes)
    perm = sortperm(new_nodes, by = x->x.pos)
    new_nodes_perm = Iterators.Stateful(perm)

    if isempty(used)
        maxsize = 0
    else
        maxused = maximum(used)
        maxsize = length(string(maxused))
    end
    if !verbose_linetable
        (loc_annotations, loc_methods, loc_lineno) = compute_ir_line_annotations(code)
        max_loc_width = maximum(length(str) for str in loc_annotations)
        max_lineno_width = maximum(length(str) for str in loc_lineno)
        max_method_width = maximum(length(str) for str in loc_methods)
    end
    max_depth = maximum(line == 0 ? 1 : compute_inlining_depth(code.linetable, line) for line in code.lines)
    last_stack = []
    for idx in eachindex(code.stmts)
        if !isassigned(code.stmts, idx)
            # This is invalid, but do something useful rather
            # than erroring, to make debugging easier
            printstyled(io, :red, "UNDEF\n")
            continue
        end
        stmt = code.stmts[idx]
        # Compute BB guard rail
        bbrange = cfg.blocks[bb_idx].stmts
        bbrange = bbrange.first:bbrange.last
        bb_pad = max_bb_idx_size - length(string(bb_idx))
        bb_start_str = string("$(bb_idx) ",length(cfg.blocks[bb_idx].preds) <= 1 ? "─" : "┄",  "─"^(bb_pad)," ")
        bb_guard_rail_cont = string("│  "," "^max_bb_idx_size)
        if idx == first(bbrange)
            bb_guard_rail = bb_start_str
        else
            bb_guard_rail = bb_guard_rail_cont
        end
        # Print linetable information
        if verbose_linetable
            stack = compute_loc_stack(code, code.lines[idx])
            # We need to print any stack frames that did not exist in the last stack
            ndepth = max(1, length(stack))
            rail = string(" "^(max_depth+1-ndepth), "│"^ndepth)
            start_column = cols - max_depth - 10
            for (i, x) in enumerate(stack)
                if i > length(last_stack) || last_stack[i] != x
                    entry = code.linetable[x]
                    printstyled(io, "\e[$(start_column)G$(rail)\e[1G", color = :light_black)
                    print(io, bb_guard_rail)
                    ssa_guard = " "^(maxsize+4+(i-1))
                    entry_label = "$(ssa_guard)$(entry.method) at $(entry.file):$(entry.line) "
                    hline = string("─"^(start_column-length(entry_label)-length(bb_guard_rail)+max_depth-i), "┐")
                    printstyled(io, string(entry_label, hline), "\n"; color=:light_black)
                    bb_guard_rail = bb_guard_rail_cont
                end
            end
            printstyled(io, "\e[$(start_column)G$(rail)\e[1G", color = :light_black)
            last_stack = stack
        else
            annotation = loc_annotations[idx]
            loc_method = loc_methods[idx]
            lineno = loc_lineno[idx]
            # Print location information right aligned. If the line below is too long, it'll overwrite this,
            # but that's what we want.
            if get(io, :color, false)
                method_start_column = cols - max_method_width - max_loc_width - 2
                filler = " "^(max_loc_width-length(annotation))
                printstyled(io, "\e[$(method_start_column)G$(annotation)$(filler)$(loc_method)\e[1G", color = :light_black)
            end
            printstyled(io, lineno, " "^(max_lineno_width-length(lineno)+1); color = :light_black)
        end
        idx != last(bbrange) && Base.print(io, bb_guard_rail)
        print_sep = false
        if idx == last(bbrange)
            print_sep = true
        end
        floop = true
        while !isempty(new_nodes_perm) && new_nodes[Iterators.peek(new_nodes_perm)].pos == idx
            node_idx = popfirst!(new_nodes_perm)
            new_node = new_nodes[node_idx]
            node_idx += length(code.stmts)
            if print_sep
                if floop
                    Base.print(io, bb_start_str)
                else
                    Base.print(io, "│  "," "^max_bb_idx_size)
                end
            end
            print_sep = true
            floop = false
            Base.with_output_color(:yellow, io) do io′
                print_node(io′, node_idx, new_node.node, used, argnames, maxsize; color = false, print_typ=false)
            end
            if should_print_ssa_type(new_node.node) && node_idx in used
                expr_type_printer(io, new_node.typ)
            end
            Base.println(io)
        end
        if print_sep
            if idx == first(bbrange) && floop
                Base.print(io, bb_start_str)
            else
                Base.print(io, idx == last(bbrange) ? string("└", "─"^(1+max_bb_idx_size), " ") :
                    string("│  ", " "^max_bb_idx_size))
            end
        end
        if idx == last(bbrange)
            bb_idx += 1
        end
        if !isassigned(code.types, idx)
            # Again, this is an error, but can happen if passes don't update their type information
            printstyled(io, "::UNDEF", color=:red)
            println(io)
            continue
        end
        typ = code.types[idx]
        try
            print_node(io, idx, stmt, used, argnames, maxsize, print_typ=false)
        catch e
            print(io, "<error printing>")
        end
        if should_print_ssa_type(stmt) && idx in used
            expr_type_printer(io, typ)
        end
        println(io)
    end
end

max_stars_repo_name: ninjin/julia

max_issues_repo_name: ninjin/julia

max_forks_repo_name: ninjin/julia


text:
mutable struct BetaDist
    _n_ones::Float64
    _n_zeros::Float64
end

function pdf(beta::BetaDist, mu::Float64)
    a = beta._n_ones
    b = beta._n_zeros
    return gamma(a + b) / (gamma(a) * gamma(b)) * mu^(a - 1.0) * (1.0 - mu)^(b - 1.0)
end

function pdf(beta::BetaDist, mu::AbstractArray{Float64,1})
    a = beta._n_ones
    b = beta._n_zeros
    return gamma(a + b) / (gamma(a) * gamma(b)) * mu .^ (a - 1.0) .*
           (1.0 .- mu) .^ (b - 1.0)
end

max_stars_repo_name: soblin/prml-julia

max_issues_repo_name: soblin/prml-julia

max_forks_repo_name: soblin/prml-julia


text:
"""
Single variable bound constraints
"""
# MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{AbstractVector{<:MOI.AbstractScalarSet}}) = true
MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{MOI.LessThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{MOI.GreaterThan{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{MOI.EqualTo{Float64}}) = true
MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{MOI.Interval{Float64}}) = true

"""
Binary/Integer variable support
"""
MOI.supports_constraint(::Optimizer, ::Type{SVF}, ::Type{<:VAR_TYPES}) = true

"""
    check_var_bounds(model)

Checks whether all variables are integral and bounded on both sides. If not an error is throws
"""
function check_var_bounds(model::Optimizer)
    for var in model.variable_info
        if !var.has_lower_bound || !var.has_upper_bound || !var.is_integer
            throw(ErrorException("Each variable must be an integer and bounded. Currently the variable index $(var.idx) doesn't fulfill this requirements."))
        end
    end
end

"""
addupd_var_in_inner_model(model::Optimizer, var_index::Int)

Adds all variables to model.inner.search_space
"""
function addupd_var_in_inner_model(model::Optimizer, var_index::Int)
    if length(model.inner.search_space) < var_index
        push!(model.inner.search_space, model.variable_info[var_index])
    else
        model.inner.search_space[var_index] = model.variable_info[var_index]
    end
end

"""
    MOI.add_variable(model::Optimizer)
"""
function MOI.add_variable(model::Optimizer)
    index = length(model.variable_info)+1
    push!(model.variable_info, Variable(index))
    changes = changes = Vector{Vector{Tuple{Symbol,Int64,Int64,Int64}}}()
    push!(changes, Vector{Tuple{Symbol,Int64,Int64,Int64}}())
    model.variable_info[index].changes = changes
    push!(model.inner.subscription, Int[])
    push!(model.inner.bt_infeasible, 0)
    addupd_var_in_inner_model(model, index)
    return MOI.VariableIndex(index)
end

MOI.add_variables(model::Optimizer, n::Int) = [MOI.add_variable(model) for i in 1:n]

function check_inbounds(model::Optimizer, vi::VI)
	num_variables = length(model.variable_info)
	if !(1 <= vi.value <= num_variables)
	    @error "Invalid variable index $vi. ($num_variables variables in the model.)"
	end
	return
end

check_inbounds(model::Optimizer, var::SVF) = check_inbounds(model, var.variable)

has_upper_bound(model::Optimizer, vi::VI) =
    model.variable_info[vi.value].has_upper_bound

has_lower_bound(model::Optimizer, vi::VI) =
    model.variable_info[vi.value].has_lower_bound

is_fixed(model::Optimizer, vi::VI) =
    model.variable_info[vi.value].is_fixed

function MOI.add_constraint(model::Optimizer, v::SVF, t::MOI.Integer)
    vi = v.variable
    model.variable_info[vi.value].is_integer = true

    cindex = length(model.var_constraints)+1
    push!(model.var_constraints, (vi.value, :int, typemin(Int64), typemax(Int64)))

    addupd_var_in_inner_model(model, vi.value)
    return MOI.ConstraintIndex{SVF, MOI.Integer}(cindex)
end

function MOI.add_constraint(model::Optimizer, v::SVF, t::MOI.ZeroOne)
    vi = v.variable
    model.variable_info[vi.value].is_integer = true

    model.variable_info[vi.value].upper_bound = 1
    model.variable_info[vi.value].max = 1
    model.variable_info[vi.value].has_upper_bound = true
    model.variable_info[vi.value].lower_bound = 0
    model.variable_info[vi.value].min = 0
    model.variable_info[vi.value].has_lower_bound = true
    model.variable_info[vi.value].values =  [0,1]
    model.variable_info[vi.value].offset =  1
    model.variable_info[vi.value].indices =  1:2
    model.variable_info[vi.value].first_ptr = 1
    model.variable_info[vi.value].last_ptr = 2
    addupd_var_in_inner_model(model, vi.value)

    cindex = length(model.var_constraints)+1
    push!(model.var_constraints, (vi.value, :bin, typemin(Int64), typemax(Int64)))
    return MOI.ConstraintIndex{SVF, MOI.ZeroOne}(cindex)
end

#=
Populating Variable bounds
=#
function MOI.add_constraint(model::Optimizer, v::SVF, interval::MOI.Interval{Float64})
    vi = v.variable
    check_inbounds(model, vi)
    isnan(interval.upper) && throw(ErrorException("The interval bounds can not contain NaN and must be an Integer. Currently it has an upper bound of $(interval.upper)"))
    has_upper_bound(model, vi) && @error "Upper bound on variable $vi exists already."
    isnan(interval.lower) && throw(ErrorException("The interval bounds can not contain NaN and must be an Integer. Currently it has a lower bound of $(interval.lower)"))
    has_lower_bound(model, vi) && @error "Lower bound on variable $vi exists already."
    is_fixed(model, vi) && @error "Variable $vi is fixed. Cannot also set upper bound."

    model.variable_info[vi.value].upper_bound = interval.upper
    model.variable_info[vi.value].max = interval.upper
    model.variable_info[vi.value].has_upper_bound = true
    model.variable_info[vi.value].lower_bound = interval.lower
    model.variable_info[vi.value].offset =  1-interval.lower
    model.variable_info[vi.value].min = interval.lower
    model.variable_info[vi.value].has_lower_bound = true

    model.variable_info[vi.value].values =  model.variable_info[vi.value].min:model.variable_info[vi.value].max
    num_vals = model.variable_info[vi.value].max-model.variable_info[vi.value].min+1
    model.variable_info[vi.value].indices =  1:num_vals
    model.variable_info[vi.value].first_ptr = 1
    model.variable_info[vi.value].last_ptr = num_vals

    addupd_var_in_inner_model(model, vi.value)

    cindex = length(model.var_constraints)+1
    push!(model.var_constraints, (vi.value, :interval, interval.lower, interval.upper))
    return MOI.ConstraintIndex{SVF, MOI.Interval{Float64}}(cindex)
end

function MOI.add_constraint(model::Optimizer, v::SVF, lt::MOI.LessThan{Float64})
    vi = v.variable
    check_inbounds(model, vi)
    isnan(lt.upper) && throw(ErrorException("The variable bounds can not contain NaN and must be an Integer. Currently it has an upper bound of $(lt.upper)"))
    has_upper_bound(model, vi) && @error "Upper bound on variable $vi already exists."
    is_fixed(model, vi) && @error "Variable $vi is fixed. Cannot also set upper bound."

    model.variable_info[vi.value].upper_bound = lt.upper
    model.variable_info[vi.value].max = lt.upper
    model.variable_info[vi.value].has_upper_bound = true

    if has_lower_bound(model, vi)
        model.variable_info[vi.value].values =  model.variable_info[vi.value].min:model.variable_info[vi.value].max
        num_vals = model.variable_info[vi.value].max-model.variable_info[vi.value].min+1
        model.variable_info[vi.value].indices =  1:num_vals
        model.variable_info[vi.value].first_ptr = 1
        model.variable_info[vi.value].last_ptr = num_vals
    end

    addupd_var_in_inner_model(model, vi.value)

    cindex = length(model.var_constraints)+1
    push!(model.var_constraints, (vi.value, :leq, typemin(Int64), lt.upper))
    return MOI.ConstraintIndex{SVF, MOI.LessThan{Float64}}(cindex)
end

function MOI.add_constraint(model::Optimizer, v::SVF, gt::MOI.GreaterThan{Float64})
    vi = v.variable
    check_inbounds(model, vi)
    isnan(gt.lower) && throw(ErrorException("The variable bounds can not contain NaN and must be an Integer. Currently it has an upper bound of $(gt.upper)"))
    has_lower_bound(model, vi) && @error "Lower bound on variable $vi already exists."
    is_fixed(model, vi) && @error "Variable $vi is fixed. Cannot also set lower bound."

    model.variable_info[vi.value].lower_bound = gt.lower
    model.variable_info[vi.value].min = gt.lower
    model.variable_info[vi.value].offset =  1-gt.lower
    model.variable_info[vi.value].has_lower_bound = true

    if has_upper_bound(model, vi)
        model.variable_info[vi.value].values =  model.variable_info[vi.value].min:model.variable_info[vi.value].max
        num_vals = model.variable_info[vi.value].max-model.variable_info[vi.value].min+1
        model.variable_info[vi.value].indices =  1:num_vals
        model.variable_info[vi.value].first_ptr = 1
        model.variable_info[vi.value].last_ptr = num_vals
    end
    addupd_var_in_inner_model(model, vi.value)

    cindex = length(model.var_constraints)+1
    push!(model.var_constraints, (vi.value, :geq, gt.lower, typemax(Int64)))
    return MOI.ConstraintIndex{SVF, MOI.GreaterThan{Float64}}(cindex)
end

function MOI.add_constraint(model::Optimizer, v::SVF, eq::MOI.EqualTo{Float64})
    vi = v.variable
    check_inbounds(model, vi)
    isnan(eq.value) && @error "Invalid fixed value $(eq.value)."
    has_lower_bound(model, vi) && @error "Variable $vi has a lower bound. Cannot be fixed."
    has_upper_bound(model, vi) && @error "Variable $vi has an upper bound. Cannot be fixed."
    is_fixed(model, vi) && @error "Variable $vi is already fixed."
    model.variable_info[vi.value].lower_bound = eq.value
    model.variable_info[vi.value].upper_bound = eq.value
    model.variable_info[vi.value].is_fixed = true
    model.variable_info[vi.value].has_lower_bound = true
    model.variable_info[vi.value].has_upper_bound = true
    model.variable_info[vi.value].min = eq.value
    model.variable_info[vi.value].max = eq.value
    model.variable_info[vi.value].values = [eq.value]
    model.variable_info[vi.value].indices = [1]
    model.variable_info[vi.value].first_ptr = 1
    model.variable_info[vi.value].last_ptr = 1
    addupd_var_in_inner_model(model, vi.value)

    cindex = length(model.var_constraints)+1
    push!(model.var_constraints, (vi.value, :eq, eq.value, eq.value))
    return MOI.ConstraintIndex{SVF, MOI.EqualTo{Float64}}(cindex)
end


text:
import Pkg; Pkg.activate(".")
using Distributed
@time @everywhere include("src/required.jl")

using Latexify

## EBD data preparation

# Load data from CSV files (from file cut with terminal)
@time df = CSV.read("data/raw/ebd_warblers_cut.csv", header=true, delim="\t")

# Prepare data (arrange values & columns)
df = prepare_ebd_data(df)

# Remove duplicates (BUT not ok for counts, so-so for dates, see group-observation.jl in src/test/)
df_nogroups = filter(x -> ismissing(x[:groupIdentifier]), df)
df_groups = dropmissing(df, :groupIdentifier)
df_groups_unique = unique(df_groups, [:species, :groupIdentifier])
newdf = vcat(df_nogroups, df_groups_unique)

# Select subset with specific columns
select!(newdf, [:species, :year, :latitude, :longitude, :groupIdentifier, :countryCode, :samplingEventIdentifier])
# Remove 1 Aleutian Islands observation with positive longitude
filter!(x -> x.longitude < 0, newdf)

# Clean workspace
df = newdf
df_nogroups = nothing
df_groups = nothing
df_groups_unique = nothing

# Data for table
table1 = by(df, :countryCode, [:countryCode, :species, :samplingEventIdentifier] =>
        x -> (n_obs = length(x.countryCode),
              n_checklist = length(unique(x.samplingEventIdentifier)),
              n_sp = length(unique(x.species))
              ))
tmp = by(df, [:countryCode, :samplingEventIdentifier], nrow)
table2 = by(tmp, :countryCode, [:x1] => x -> (n_sp_moy = mean(x.x1),
                                              n_sp_med = median(x.x1),
                                              n_sp_max = maximum(x.x1)
                                              ))
tables = join(table1, table2, on = :countryCode)

n_obs_tot = nrow(df)
n_checklist_tot = length(unique(df.samplingEventIdentifier))
n_sp_tot = length(unique(df.species))
tmp2 = by(df, :samplingEventIdentifier, nrow)
n_sp_moy_tot = mean(tmp[:x1])
n_sp_med_tot = median(tmp[:x1])
n_sp_max_tot = maximum(tmp[:x1])
table_tot = DataFrame(countryCode = "Total", n_obs = n_obs_tot, n_checklist = n_checklist_tot, n_sp = n_sp_tot,
                        n_sp_moy = n_sp_moy_tot, n_sp_med = n_sp_med_tot, n_sp_max = n_sp_max_tot)

tables = vcat(tables, table_tot)

show(md(tables))
latexify(tables) |> print

max_stars_repo_name: gabrieldansereau/betadiversity-hotspots

max_issues_repo_name: gabrieldansereau/betadiversity-hotspots

max_forks_repo_name: gabrieldansereau/betadiversity-hotspots


text:
#=
Copyright 2021 BlackRock, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
=#

@doc raw"""
    ProxCache

Store the proximal operators corresponding to each element of a vector of `PiecewiseQuadratic`s

# Note
The proximal operator of ``f``, ``\rho``, denoted
```math
\text{prox}_{f,\rho}(u) =
\argmin_{x \in \text{dom}(f)} f(x) + \frac{\rho}{2}  \|x - u \|_2^2
```
"""
struct ProxCache
    f_plus_quad::Vector{PiecewiseQuadratic}
end

"""
    ProxCache(params::AdmmParams, settings::Settings)

Construct a ProxCache given AdmmParams `params` and ADMM Settings `settings`
"""
function ProxCache(params::AdmmParams, settings::Settings)
    f_plus_quad = Vector{PiecewiseQuadratic}(undef, length(params.g))
    work = PwqSumWorkspace(2)
    for i in 1:length(params.g)
        σi = settings.σ[i]
        cm = _sum([params.g[i], PiecewiseQuadratic(BoundedQuadratic(σi / 2.0, 0.0, 0.0))],
                  work)
        env = simplify(envelope(cm))
        f_plus_quad[i] = env
    end

    return ProxCache(f_plus_quad)
end

"""
    prox_step(pc::ProxCache, params::AdmmParams, settings::Settings, u::Vector{Float64})

Return the next `x` by using the proximal operators of the separable function pieces
"""
function prox_step(pc::ProxCache, params::AdmmParams, settings::Settings,
                   u::Vector{Float64})
    x = zeros(length(u))
    for i in 1:length(u)
        x[i] = prox(pc.f_plus_quad[i], u[i], settings.σ[i]; use_quadratic=false)
    end
    return x
end


text:
include("twihes.jl")
#include("pxxx.jl") <- Globaly defines constant parameters. Useful for when running in REPL
using BenchmarkTools

#Alice' first part of the key-exchange
function ephimeral_A(s::T, Pa::TwiHesPoint{T}, Qa::TwiHesPoint{T}, Pb::TwiHesPoint{T}, Qb::TwiHesPoint{T}, a0::Fp2Elem{T}, ea::U, strat_A::Array{Int64}) where {T <: Integer, U <:Integer}
    G = add(Pa, mul(s, Qa, a0), a0)
    P_B, Q_B, a_new = iso2pow_opt(G, Pb, Qb, a0, ea, strat_A)
    d_new = recover_d(P_B, a_new) #Ignores d-parameter during calculation, so must recover it at the end
    return P_B, Q_B, a_new, d_new
end

#Bob's first part of the key-exchange
function ephimeral_B(s::T, Pb::TwiHesPoint{T}, Qb::TwiHesPoint{T}, Pa::TwiHesPoint{T}, Qa::TwiHesPoint{T}, a0::Fp2Elem{T}, d0::Fp2Elem{T}, eb::U, strat_B::Array{Int64}, w::Fp2Elem{T}) where {T <: Integer, U <:Integer}
    G = add(Pb, mul(s, Qb, a0), a0)
    P_A, Q_A, a_new, d_new = iso3pow_opt(G, Pa, Qa, a0, d0, w, eb, strat_B)
    return P_A, Q_A, a_new, d_new
end

#Alice' second part of the key-exchange
function secretAgreement_A(s::Integer, P_A::TwiHesPoint{T}, Q_A::TwiHesPoint{T}, a_B::Fp2Elem{T}, ea::U, strat_A::Array{Int64}) where {T <: Integer, U <:Integer}
    G = add(P_A, mul(s, Q_A, a_B), a_B)
    a_new, d_new = iso2pow_opt(G, a_B, ea, strat_A) #Recovering d is build in the function, see twihes.jl
    return jInvariant(a_new, d_new)
end

#Bob's second part of the key-exchange
function secretAgreement_B(s::Integer, P_B::TwiHesPoint{T}, Q_B::TwiHesPoint{T}, a_A::Fp2Elem{T}, d_A::Fp2Elem{T}, eb::U, strat_B::Array{Int64}, w::Fp2Elem{T}) where {T <: Integer, U <:Integer}
    G = add(P_B, mul(s, Q_B, a_A), a_A)
    a_new, d_new = iso3pow_opt(G, a_A, d_A, w, eb, strat_B)
    return jInvariant(a_new, d_new)
end

#Full key-exhange example
function example()
    sA = rand(1:la^ea)
    sB = rand(1:lb^eb)
    println("\nAlice' secret: $sA")
    println("\nBob's secret: $sB")

    P_B, Q_B, a_A, d_A = ephimeral_A(sA, Pa, Qa, Pb, Qb, a0, ea, strat_A)
    P_A, Q_A, a_B, d_B = ephimeral_B(sB, Pb, Qb, Pa, Qa, a0, d0, eb, strat_B, w)

    println("\nAlice sends:")
    println("phi_A(P_B) = $P_B")
    println("phi_A(Q_B) = $Q_B")
    println("E_A : ($(a_A))X^3 + Y^3 + Z^3 = ($(d_A))XYZ\n")
    println("Bob sends:")
    println("phi_B(P_A) = $P_A")
    println("phi_B(Q_A) = $Q_A")
    println("E_B : ($(a_B))X^3 + Y^3 + Z^3 = ($(d_B))XYZ\n")
    println("--------------------------------\n")

    j_A = secretAgreement_A(sA, P_A, Q_A, a_B, ea, strat_A)
    j_B = secretAgreement_B(sB, P_B, Q_B, a_A, d_A, eb, strat_B, w)

    println("Results:")
    println("key_A: $j_A")
    println("key_B: $j_B \n \n")
end

#Timing the second part of Bob's key-exhange to compare with OptimizedExperimental
function benchmarkBob()
    sA = rand(1:la^ea)
    sB = rand(1:lb^eb)
    P_B, Q_B, a_A, d_A = ephimeral_A(sA, Pa, Qa, Pb, Qb, a0, ea, strat_A)
    P_A, Q_A, a_B, d_B = ephimeral_B(sB, Pb, Qb, Pa, Qa, a0, d0, eb, strat_B, w)
    j_A = secretAgreement_A(sA, P_A, Q_A, a_B, ea, strat_A)
    j_B = secretAgreement_B(sB, P_B, Q_B, a_A, d_A, eb, strat_B, w)
    if j_A != j_B
        error("Shared secret was not the same")
    end

    return @benchmark secretAgreement_B($sB, $P_B, $Q_B, $a_A, $d_A, $eb, $strat_B, $w)
end

#Timing the second part of Alice' key-exhange to compare with OptimizedExperimental
function benchmarkAlice()
    sA = rand(1:la^ea)
    sB = rand(1:lb^eb)
    P_B, Q_B, a_A, d_A = ephimeral_A(sA, Pa, Qa, Pb, Qb, a0, ea, strat_A)
    P_A, Q_A, a_B, d_B = ephimeral_B(sB, Pb, Qb, Pa, Qa, a0, d0, eb, strat_B, w)
    j_A = secretAgreement_A(sA, P_A, Q_A, a_B, ea, strat_A)
    j_B = secretAgreement_B(sB, P_B, Q_B, a_A, d_A, eb, strat_B, w)
    if j_A != j_B
        error("Shared secret was not the same")
    end

    return @benchmark secretAgreement_A($sA, $P_A, $Q_A, $a_B, $eb, $strat_B)
end

function main()
    choice = "0"
    while !(choice in ["1","2","3","4","5"])
        println("Which parameter-sets do you want?")
        println("[1] p132, [2] p434, [3] p503, [4] p610, [5] p751")
        choice = readline()
        if choice == "1"
            include("p132.jl")
        elseif choice == "2"
            include("p434.jl")
        elseif choice == "3"
            include("p503.jl")
        elseif choice == "4"
            include("p610.jl")
        elseif choice == "5"
            include("p751.jl")
        else
            println("pick a number between 1 and 5")
            println("----------------------------------------------------------------")
        end
    end
    choice = "0"
    while choice != "4"
        println("What do you want to do?")
        println("[1] Run as example, [2] Benchmark Alice, [3] Benchmark Bob, [4] Close")
        choice = readline()
        if choice == "1"
            example()
        elseif choice == "2"
            println("\n")
            show(stdout,MIME"text/plain"(),benchmarkAlice())
            println("\n")
        elseif choice == "3"
            println("\n")
            show(stdout,MIME"text/plain"(),benchmarkBob())
            println("\n")
        elseif choice == "4"
            return
        else
            println("pick a number between 1 and 4")
            println("----------------------------------------------------------------")
        end
    end
end

main()


text:
### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ 2697620e-5e5e-11eb-3568-ab2248561ebd
begin
	import Pkg
	Pkg.activate(Base.current_project())
	using ReinforcementLearning
	using Flux
	using Statistics
	using Plots
end

# ╔═╡ fb3c6816-5e5d-11eb-06dd-b5ba442b05ed
md"""
# Chapter 5 The Blackjack Environment

In this notebook, we'll study monte carlo based methods to play the **Blackjack** game.
"""

# ╔═╡ 32125940-5e5e-11eb-3a93-a30c1afca000
md"""
As usual, let's define the environment first. The implementation of the **Blackjack** environment is mainly taken from [openai/gym](https://github.com/openai/gym/blob/c4d0af393ef9fba641bd3ebbbf1f60d291c8475d/gym/envs/toy_text/blackjack.py) with some necessary modifications for our following up experiments.
"""

# ╔═╡ 9e0bc938-5e5e-11eb-1665-c331d08ac768
begin
	# 1 = Ace, 2-10 = Number cards, Jack/Queen/King = 10
	DECK = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 10, 10, 10]

	mutable struct BlackjackEnv <: AbstractEnv
		dealer_hand::Vector{Int}
		player_hand::Vector{Int}
		done::Bool
		reward::Int
		init::Union{Tuple{Vector{Int}, Vector{Int}}, Nothing}
	end

	function BlackjackEnv(;init=nothing)
		env = BlackjackEnv([], [], false, 0., init)
		reset!(env)
		env
	end

	function RLBase.reset!(env::BlackjackEnv)
		empty!(env.dealer_hand)
		empty!(env.player_hand)
		if isnothing(env.init)
			push!(env.dealer_hand, rand(DECK))
			push!(env.dealer_hand, rand(DECK))
			while sum_hand(env.player_hand) < 12
				push!(env.player_hand, rand(DECK))
			end
		else
			append!(env.player_hand, env.init[1])
			append!(env.dealer_hand, env.init[2])
		end
		env.done=false
		env.reward = 0.
	end

	RLBase.state_space(env::BlackjackEnv) = Space([Base.OneTo(31), Base.OneTo(10), Base.OneTo(2)])
	RLBase.action_space(env::BlackjackEnv) = Base.OneTo(2)

	usable_ace(hand) = (1 in hand) && (sum(hand) + 10 <= 21)
	sum_hand(hand) = usable_ace(hand) ? sum(hand) + 10 : sum(hand)
	is_bust(hand) = sum_hand(hand) > 21
	score(hand) = is_bust(hand) ? 0 : sum_hand(hand)

	RLBase.state(env::BlackjackEnv) = (sum_hand(env.player_hand), env.dealer_hand[1], usable_ace(env.player_hand)+1)
	RLBase.reward(env::BlackjackEnv) = env.reward
	RLBase.is_terminated(env::BlackjackEnv) = env.done

	function (env::BlackjackEnv)(action)
		if action == 1
			push!(env.player_hand, rand(DECK))
			if is_bust(env.player_hand)
				env.done = true
				env.reward = -1
			else
				env.done = false
				env.reward = 0
			end
		elseif action == 2
			env.done = true
			while sum_hand(env.dealer_hand) < 17
				push!(env.dealer_hand, rand(DECK))
			end
			env.reward = cmp(score(env.player_hand), score(env.dealer_hand))
		else
			@error "unknown action"
		end
	end
end

# ╔═╡ f0454136-5e5e-11eb-01ad-2fda9f6d0a78
game = BlackjackEnv()

# ╔═╡ f6d2fbe0-5e5e-11eb-0cae-7b9f5d343ae3
state_space(game)

# ╔═╡ dfebc2b8-5e5e-11eb-2030-4766a533b521
md"""
As you can see, the `state_space` of the **Blackjack** environment has 3 discrete features. To reuse the tabular algorithms in `ReinforcementLearning.jl`, we need to flatten the state and wrap it in a `StateTransformedEnv`.
"""

# ╔═╡ 651a9d74-5e5f-11eb-251d-1bab0d8d3384
STATE_MAPPING = s ->  LinearIndices((31, 10, 2))[CartesianIndex(s)]

# ╔═╡ 6f95e1a0-5e5f-11eb-1203-21668fdb4d1d
world = StateTransformedEnv(
    BlackjackEnv(),
    state_mapping=STATE_MAPPING
)

# ╔═╡ 82391e12-5e5f-11eb-1eb9-f1239a781589
RLBase.state_space(x::typeof(world)) = Base.OneTo(31* 10*2)

# ╔═╡ 7e18ee70-5e5f-11eb-00aa-877362e303b3
NS = state_space(world)

# ╔═╡ cffe4ec4-5e5f-11eb-2f7d-03deb8eb3680
md"""
## Figure 5.1
"""

# ╔═╡ e56fbcca-5e5f-11eb-0c04-d3b2aa8a634e
agent = Agent(
    policy = VBasedPolicy(
        learner=MonteCarloLearner(;
            approximator=TabularVApproximator(;n_state=NS, opt=InvDecay(1.0)),
            γ = 1.0
            ),
        mapping= (env, V) -> sum_hand(env.env.player_hand) in (20, 21) ? 2 : 1
        ),
    trajectory=VectorSARTTrajectory()
)

# ╔═╡ f032cbf2-5e5f-11eb-0bca-457ebd9e2f01
run(agent, world, StopAfterEpisode(10_000))

# ╔═╡ fd16d76e-5e5f-11eb-06f3-ebfe8c142e13
VT = agent.policy.learner.approximator.table

# ╔═╡ 0b4702e4-5e60-11eb-1d2a-7528302acbfe
X, Y = 12:21, 1:10

# ╔═╡ 129d507c-5e60-11eb-05de-75e719976129
plot(X, Y, [VT[STATE_MAPPING((i,j,1))] for i in X, j in Y],linetype=:wireframe)

# ╔═╡ 1d37cfda-5e60-11eb-0923-3d5f8a53521a
plot(X, Y, [VT[STATE_MAPPING((i,j,2))] for i in X, j in Y],linetype=:wireframe)

# ╔═╡ 2e18d632-5e60-11eb-3988-8de51cb1706a
# now run more simulations
run(agent, world, StopAfterEpisode(500_000))

# ╔═╡ 3b0f8098-5e60-11eb-035c-7339e9c07157
plot(X, Y, [VT[STATE_MAPPING((i,j,1))] for i in X, j in Y],linetype=:wireframe)

# ╔═╡ 2b67b004-5e60-11eb-0c57-9b94ffc1799e
plot(X, Y, [VT[STATE_MAPPING((i,j,2))] for i in X, j in Y],linetype=:wireframe)

# ╔═╡ 475e9956-5e60-11eb-1108-13deb8642ac0
md"""
## Figure 5.2
"""

# ╔═╡ 54ac6d18-5e60-11eb-030c-190885e37497
md"""
In Chapter 5.3, a **Monte Carlo Exploring Start** method is used to solve the **Blackjack** game. Although several variants of monte carlo methods are supported in `ReinforcementLearning.jl` package, they do not support the *exploring start*. Nevertheless, we can define it very easily.
"""

# ╔═╡ d33e736a-5e60-11eb-3765-69a54bbf6049
begin
	Base.@kwdef mutable struct ExploringStartPolicy{P} <: AbstractPolicy
		policy::P
		is_start::Bool = true
	end

	function (p::ExploringStartPolicy)(env::AbstractEnv)
		if p.is_start
			p.is_start = false
			rand(action_space(env))
		else
			p.policy(env)
		end
	end

	(p::ExploringStartPolicy)(s::AbstractStage, env::AbstractEnv) = p.policy(s, env)

	function (p::ExploringStartPolicy)(s::PreEpisodeStage, env::AbstractEnv)
		p.is_start = true
		p.policy(s, env)
	end

	function (p::ExploringStartPolicy)(s::PreActStage, env::AbstractEnv, action)
		p.policy(s, env, action)
	end
end

# ╔═╡ d9a5a368-5e60-11eb-0a0d-313dc4697d0d
solver = Agent(
    policy = QBasedPolicy(
        learner=MonteCarloLearner(;
            approximator=TabularQApproximator(
				;n_state=NS,
				n_action=2,
				opt=InvDecay(1.0)),
            γ = 1.0,
            ),
        explorer=GreedyExplorer()
        ),
    trajectory=VectorSARTTrajectory()
)

# ╔═╡ e904b632-5e60-11eb-05f3-abca0adeac69
run(ExploringStartPolicy(policy=solver), world, StopAfterEpisode(10_000_000))

# ╔═╡ 10750292-5e61-11eb-0a2f-51c6de1737b4
QT = solver.policy.learner.approximator.table

# ╔═╡ 114b9d7e-5e61-11eb-3579-f9fd7d0a4bd3
heatmap([argmax(QT[:,STATE_MAPPING((i,j,1))]) for i in 11:21, j in Y])

# ╔═╡ 1aec9c8c-5e61-11eb-2b12-6b9dbf4b61db
heatmap([argmax(QT[:,STATE_MAPPING((i,j,2))]) for i in 11:21, j in Y])

# ╔═╡ 15f4a8fa-5e61-11eb-2313-81caa4ed00e6
V_agent = Agent(
    policy = VBasedPolicy(
        learner=MonteCarloLearner(;
            approximator=TabularVApproximator(;n_state=NS, opt=InvDecay(1.0)),
            γ = 1.0
            ),
        mapping=(env, V) -> solver.policy(env)
        ),
    trajectory=VectorSARTTrajectory()
)

# ╔═╡ 8c366ee0-5e61-11eb-31b9-6722b4dfc9a3
run(V_agent, world, StopAfterEpisode(500_000))

# ╔═╡ 937068f0-5e61-11eb-0326-4f2e7e313b83
V_agent_T = V_agent.policy.learner.approximator.table;

# ╔═╡ 9b0617b0-5e61-11eb-27be-15bcf24d11f5
plot(X, Y, [V_agent_T[STATE_MAPPING((i,j,1))] for i in X, j in Y],linetype=:wireframe)

# ╔═╡ a0a52486-5e61-11eb-212a-45ca32bfb987
plot(X, Y, [V_agent_T[STATE_MAPPING((i,j,2))] for i in X, j in Y],linetype=:wireframe)

# ╔═╡ aa68a5ec-5e61-11eb-07c8-9103a0ada51c
md"""
## Figure 5.3
"""

# ╔═╡ d874fc22-5e61-11eb-2fee-37742f910182
static_env = StateTransformedEnv(
    BlackjackEnv(;init=([1,2], [2])),
    state_mapping=STATE_MAPPING
);

# ╔═╡ e31a1da6-5e61-11eb-09ac-1f29247434ca
INIT_STATE = state(static_env)

# ╔═╡ f1437b16-5e61-11eb-2f14-976483f8d924
GOLD_VAL = -0.27726

# ╔═╡ ecade5b4-5e61-11eb-0fd5-d1b149adbece
begin
	Base.@kwdef struct StoreMSE <: AbstractHook
		mse::Vector{Float64} = []
	end
	(f::StoreMSE)(::PostEpisodeStage, agent, env) = push!(f.mse, (GOLD_VAL - agent.policy.π_target.learner.approximator[1](INIT_STATE))^2)
end

# ╔═╡ 4f8a0a80-5e64-11eb-13db-3b117cdd35b6
target_policy_mapping = (env, V) -> sum_hand(env.env.player_hand) in (20, 21) ? 2 : 1

# ╔═╡ 5f5473ba-5e64-11eb-1b37-7706eb629b45
function RLBase.prob(
	p::VBasedPolicy{<:Any,typeof(target_policy_mapping)},
	env::AbstractEnv,
	a
)
	s = sum_hand(env.env.player_hand)
    if s in (20, 21)
        Int(a == 2)
    else
        Int(a == 1)
    end
end

# ╔═╡ bceb8f5a-5e63-11eb-25fb-1d73dc919323
function ordinary_mse()
	agent = Agent(
		policy=OffPolicy(
			π_target = VBasedPolicy(
				learner=MonteCarloLearner(
					approximator=(
						TabularVApproximator(;n_state=NS, opt=Descent(1.0)),
						TabularVApproximator(;n_state=NS, opt=InvDecay(1.0))
						),
					kind=FIRST_VISIT,
					sampling=ORDINARY_IMPORTANCE_SAMPLING
					),
				mapping=target_policy_mapping
			),
			π_behavior = RandomPolicy(Base.OneTo(2))
			),
		trajectory=VectorWSARTTrajectory()
	)
	h = StoreMSE()
	run(agent, static_env, StopAfterEpisode(10_000), h)
	h.mse
end

# ╔═╡ 0d5ef608-5e65-11eb-2881-d72c5478e1ce
function weighted_mse()
	agent = Agent(
		policy=OffPolicy(
			π_target = VBasedPolicy(
				learner=MonteCarloLearner(
					approximator=(
						TabularVApproximator(;n_state=NS, opt=Descent(1.0)),
						TabularVApproximator(;n_state=NS, opt=InvDecay(1.0)),
						TabularVApproximator(;n_state=NS, opt=InvDecay(1.0))),
					kind=FIRST_VISIT,
					sampling=WEIGHTED_IMPORTANCE_SAMPLING
					),
				mapping=target_policy_mapping
			),
			π_behavior = RandomPolicy(Base.OneTo(2))
			),
		trajectory=VectorWSARTTrajectory()
	)
	h = StoreMSE()
	run(agent, static_env, StopAfterEpisode(10_000), h)
	h.mse
end

# ╔═╡ d054b474-5e64-11eb-281d-c75d62dcad55
begin
	fig_5_3 = plot()
	plot!(fig_5_3, mean(ordinary_mse() for _ in 1:100), xscale=:log10,label="Ordinary Importance Sampling")
	plot!(fig_5_3, mean(weighted_mse() for _ in 1:100), xscale=:log10,label="Weighted Importance Sampling")
	fig_5_3
end

# ╔═╡ Cell order:
# ╟─fb3c6816-5e5d-11eb-06dd-b5ba442b05ed
# ╠═2697620e-5e5e-11eb-3568-ab2248561ebd
# ╟─32125940-5e5e-11eb-3a93-a30c1afca000
# ╠═9e0bc938-5e5e-11eb-1665-c331d08ac768
# ╠═f0454136-5e5e-11eb-01ad-2fda9f6d0a78
# ╠═f6d2fbe0-5e5e-11eb-0cae-7b9f5d343ae3
# ╟─dfebc2b8-5e5e-11eb-2030-4766a533b521
# ╠═651a9d74-5e5f-11eb-251d-1bab0d8d3384
# ╠═6f95e1a0-5e5f-11eb-1203-21668fdb4d1d
# ╠═82391e12-5e5f-11eb-1eb9-f1239a781589
# ╠═7e18ee70-5e5f-11eb-00aa-877362e303b3
# ╟─cffe4ec4-5e5f-11eb-2f7d-03deb8eb3680
# ╠═e56fbcca-5e5f-11eb-0c04-d3b2aa8a634e
# ╠═f032cbf2-5e5f-11eb-0bca-457ebd9e2f01
# ╠═fd16d76e-5e5f-11eb-06f3-ebfe8c142e13
# ╠═0b4702e4-5e60-11eb-1d2a-7528302acbfe
# ╠═129d507c-5e60-11eb-05de-75e719976129
# ╠═1d37cfda-5e60-11eb-0923-3d5f8a53521a
# ╠═2e18d632-5e60-11eb-3988-8de51cb1706a
# ╠═3b0f8098-5e60-11eb-035c-7339e9c07157
# ╠═2b67b004-5e60-11eb-0c57-9b94ffc1799e
# ╟─475e9956-5e60-11eb-1108-13deb8642ac0
# ╟─54ac6d18-5e60-11eb-030c-190885e37497
# ╠═d33e736a-5e60-11eb-3765-69a54bbf6049
# ╠═d9a5a368-5e60-11eb-0a0d-313dc4697d0d
# ╠═e904b632-5e60-11eb-05f3-abca0adeac69
# ╠═10750292-5e61-11eb-0a2f-51c6de1737b4
# ╠═114b9d7e-5e61-11eb-3579-f9fd7d0a4bd3
# ╠═1aec9c8c-5e61-11eb-2b12-6b9dbf4b61db
# ╠═15f4a8fa-5e61-11eb-2313-81caa4ed00e6
# ╠═8c366ee0-5e61-11eb-31b9-6722b4dfc9a3
# ╠═937068f0-5e61-11eb-0326-4f2e7e313b83
# ╠═9b0617b0-5e61-11eb-27be-15bcf24d11f5
# ╠═a0a52486-5e61-11eb-212a-45ca32bfb987
# ╟─aa68a5ec-5e61-11eb-07c8-9103a0ada51c
# ╠═d874fc22-5e61-11eb-2fee-37742f910182
# ╠═e31a1da6-5e61-11eb-09ac-1f29247434ca
# ╠═f1437b16-5e61-11eb-2f14-976483f8d924
# ╠═ecade5b4-5e61-11eb-0fd5-d1b149adbece
# ╠═4f8a0a80-5e64-11eb-13db-3b117cdd35b6
# ╠═5f5473ba-5e64-11eb-1b37-7706eb629b45
# ╠═bceb8f5a-5e63-11eb-25fb-1d73dc919323
# ╠═0d5ef608-5e65-11eb-2881-d72c5478e1ce
# ╠═d054b474-5e64-11eb-281d-c75d62dcad55


text:
#using Modia
@static if VERSION < v"0.7.0-DEV.2005"
  using Base.Test
else
  using Test
end

@testset "RunExamples" begin

Modia.ModiaLogging.resetTestStatus()

global figure = 1
include("CurrentController.jl")
include("Rectifier.jl")
include("CauerLowPassFilter.jl")
include("LinearSystems.jl")
include("SynchronousExamples.jl")
include("ElectricalVehicleAndCharger.jl") # Problem in 1.0
include("CollidingBalls.jl")
include("HeatTransfer2D.jl")

Modia.ModiaLogging.printTestStatus()

end # testset
nothing


text:
module LazyLieManifolds

using LinearAlgebra
using Manifolds
using StaticArrays
using DocStringExtensions

const LLM = LazyLieManifolds

export LLM
##==============================================================================
## Abstract Lie Groups and Manifolds
##==============================================================================

abstract type AbstractLieGroup end
abstract type AbstractLieAlgebra end

Base.getindex(p::AbstractLieGroup) = p.value
Base.getindex(X::AbstractLieAlgebra) = X.value

# extend Manifolds.jl
Manifolds.affine_matrix(p::AbstractLieGroup) = affine_matrix(manifold(p), p[])
Manifolds.screw_matrix(X::AbstractLieAlgebra) = affine_matrix(manifold(X), X[])

"""
    $(SIGNATURES)
"""
function manifold end

"""
    $(SIGNATURES)
"""
function _identity end

##------------------------------------------------------------------------------
## Extending Manifolds.jl
##------------------------------------------------------------------------------

function Manifolds.vee(X::AbstractLieAlgebra)
    M = manifold(X)
    return vee(M, make_identity(M, X[]), X[]) #FIXME
end

# we must use
# seVec = Manifolds.hat(M, p, coords)

function Base.:∘(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(compose(M, p[], q[]))
end

function Base.inv(p::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(inv(M, p[]))
end

function Base.:\(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(compose(M, inv(M, p[]), q[]))
end

function Base.:/(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    typeof(p)(compose(M,  p[], inv(M, q[])))
end

function Base.exp(X::AbstractLieAlgebra) #look into retract with RightAction
    M = manifold(X)
    e = identity_element(M, X[]) #TODO this might not be completely correct
    G = convert(AbstractLieGroup, typeof(X))
    return G(exp(M, e, X[]))
end

function Base.log(p::AbstractLieGroup) #look into inverse_retract with RightAction
    M = manifold(p)
    e = identity_element(M, p[])
    g = convert(AbstractLieAlgebra, typeof(p))
    return g(log(M, e, p[]))
end

function Base.log(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    g = convert(AbstractLieAlgebra, typeof(p))
    return g(log(M, p[], q[]))
end

function Manifolds.group_exp(X::AbstractLieAlgebra) #look into retract with RightAction
    M = manifold(X)
    G = convert(AbstractLieGroup, typeof(X))
    return G(group_exp(M, X[]))
end

function Manifolds.distance(p::AbstractLieGroup, q::AbstractLieGroup)
    M = manifold(p)
    return distance(M, p[], q[])
end

Manifolds.manifold_dimension(p::AbstractLieGroup) = manifold_dimension(manifold(p))

##==============================================================================
## SpecialEuclidean SE{N} and se{N}
##==============================================================================
export se, SE

"""
    $TYPEDEF

Lie Group representation for `Manifolds.SpecialEuclidean(N)`

Related

[`se`](@ref)
"""
struct SE{N, T} <: AbstractLieGroup#TODO {N, T}
  value::T # Group
end

#TODO assert N and/or SE(p::ProductRepr) = ...
SE{N}(p::ProductRepr) where N = SE{N, typeof(p)}(p)

SE{N}() where N = _identity(SE{N})

function SE(R::Matrix{T}, t::Vector{T}) where T<:Real
    N = length(t)
    @assert size(R,1) == size(R,2) == N

    _t = SVector{N}(t)
    _R = SMatrix{N,N}(R)

    p = ProductRepr(_t, _R)

    return SE{N, typeof(p)}(p)
end

manifold(::SE{N,T}) where {N,T} = SpecialEuclidean(N)

getNparameter(::SE{N}) where N = N

#Default SE{N} representation
function _identity(::Type{SE{N}}, ::Type{T}=Float64) where {N, T}
    t = zeros(SVector{N, T})
    R = SMatrix{N,N,T}(one(T)I)
    SE{N}(ProductRepr(t, R))
end

"""
    $TYPEDEF

Lie algebra representation for `Manifolds.SpecialEuclidean(N)`

Related

[`SE`](@ref)
"""
struct se{N, T} <: AbstractLieAlgebra
  value::T # Tangent Vector (Lie algebra, skew simetric matrix)
end

#TODO assert N and/or se(X::ProductRepr) = ...
se{N}(X::ProductRepr) where N = se{N, typeof(X)}(X)

function se{N}(X::AbstractVector) where N
    l = length(X)
    @assert manifold_dimension(SpecialEuclidean(N)) == l "X dimension, $l, does not match manifold dimension, $N"

    M = SpecialEuclidean(N)

    e = _identity(SE{N}, eltype(X))[]
    _X = hat(M, e, X)

    return se{N}(_X)
end

manifold(::se{N,T}) where {N,T} = SpecialEuclidean(N)

Base.convert(::Type{AbstractLieGroup}, ::Type{se{N, T}}) where {N,T} = SE{N}
Base.convert(::Type{AbstractLieAlgebra}, ::Type{SE{N, T}}) where {N,T} = se{N}

# function default_point_type(M::SpecialOrthogonal, ::Type{ET}=Float64, ::Type{AT}=SMatrix) where {ET, AT}
#     s = representation_size(M)
#     return AT{s..., ET, prod(s)}
# end

# function default_point_type(M::SpecialEuclidean, ::Type{ET}=Float64, ::Type{AT}=SArray) where {ET, AT}

#     s1 = representation_size(M.manifold[1])
#     s2 = representation_size(M.manifold[2])
#     return ProductRepr{Tuple{AT{Tuple{s1[1]}, ET}, AT{Tuple{s2[1],s2[2]}, ET}}}
# end

##------------------------------------------------------------------------------
## other operators
##------------------------------------------------------------------------------
#TODO abstract
function ⊕(p::SE{N}, X::se{N}) where N
    M = manifold(p)
    e = identity_element(M, p[])
    typeof(p)(compose(M, p[], exp(M, e, X[])))
end

function ⊕(X::se{N}, p::SE{N}) where N
    M = manifold(p)
    e = identity_element(M, p[])
    typeof(p)(compose(M, exp(M, e, X[]), p[]))
end

# ⊖ - right minus
# X = log(M, e, inv(x) ∘ y)
function ⊖(y::SE{N}, x::SE{N}) where N
    M = manifold(y)
    e = identity_element(M, y[])
    return log(M, e, inv(x) ∘ y )
end

##==============================================================================
## SpecialOrthogonal SO{N} and so{N}
##==============================================================================
export so, SO

"""
    $TYPEDEF

Lie Group representation for `Manifolds.SpecialOrthogonal(N)`

Related

[`so`](@ref)
"""
struct SO{N, T} <: AbstractLieGroup
  value::T # Group
end

#TODO assert N and/or SO(p::ProductRepr) = ...
SO{N}(p::StaticMatrix{N,N,T}) where {N, T} = SO{N, typeof(p)}(p)

SO{N}() where N = _identity(SO{N})

function SO(R::Matrix{T}) where T<:Real
    N = size(R,1)
    @assert size(R,1) == size(R,2)

    return SO{N}(SMatrix{N,N}(R))
end

manifold(::SO{N,T}) where {N,T} = SpecialOrthogonal(N)

#Default SO{N} representation
function _identity(::Type{SO{N}}, ::Type{T}=Float64) where {N, T}
    R = SMatrix{N,N,T}(one(T)I)
    SO{N}(R)
end

"""
    $TYPEDEF

Lie algebra representation for `Manifolds.SpecialOrthogonal(N)`

Related

[`SO`](@ref)
"""
struct so{N, T} <: AbstractLieAlgebra
  value::T # Tangent Vector (Lie algebra, skew simetric matrix)
end

#TODO assert N and/or so(X::ProductRepr) = ...
so{N}(X::StaticMatrix{N,N,T}) where {N,T} = so{N, typeof(X)}(X)

function so{N}(X::AbstractVector) where N
    l = length(X)
    @assert manifold_dimension(SpecialOrthogonal(N)) == l "X dimension, $l, does not match manifold dimension, $N"

    M = SpecialOrthogonal(N)

    e = _identity(SO{N}, eltype(X))[]
    _X = hat(M, e, X)

    return so{N}(_X)
end

manifold(::so{N,T}) where {N,T} = SpecialOrthogonal(N)

Base.convert(::Type{AbstractLieGroup}, ::Type{so{N, T}}) where {N,T} = SO{N}
Base.convert(::Type{AbstractLieAlgebra}, ::Type{SO{N, T}}) where {N,T} = so{N}

##==============================================================================
## TranslationGroup 𝕋{N} and 𝕥{N} feels a bit redundant, but adding for completeness
##==============================================================================
export 𝕥, 𝕋

"""
    $TYPEDEF

Lie Group representation for `Manifolds.TranslationGroup(N)`

Related

[`𝕥`](@ref)
"""
struct 𝕋{N, T} <: AbstractLieGroup
  value::T # Group
end

#TODO assert N and/or 𝕋(p::ProductRepr) = ...
𝕋{N}(p::SVector{N,T}) where {N, T} = 𝕋{N, typeof(p)}(p)

𝕋{N}() where N = _identity(𝕋{N})

function 𝕋(t::Vector{T}) where T<:Real
    N = size(t,1)

    return 𝕋{N}(SVector{N, T}(t))
end

manifold(::𝕋{N,T}) where {N,T} = TranslationGroup(N)

#Default 𝕋{N} representation
function _identity(::Type{𝕋{N}}, ::Type{T}=Float64) where {N, T}
    t = zeros(SVector{N, T})
    𝕋{N}(t)
end

"""
    $TYPEDEF

Lie algebra representation for `Manifolds.TranslationGroup(N)`

Related

[`𝕋`](@ref)
"""
struct 𝕥{N, T} <: AbstractLieAlgebra
  value::T # Tangent Vector (Lie algebra)
end

#TODO assert N and/or 𝕥(X::ProductRepr) = ...
𝕥{N}(X::SVector{N,T}) where {N,T} = 𝕥{N, typeof(X)}(X)

function 𝕥{N}(X::AbstractVector) where N
    l = length(X)
    @assert manifold_dimension(TranslationGroup(N)) == l "X dimension, $l, does not match manifold dimension, $N"

    M = TranslationGroup(N)

    e = _identity(𝕋{N}, eltype(X))[]
    _X = hat(M, e, X)

    return 𝕥{N}(_X)
end

manifold(::𝕥{N,T}) where {N,T} = TranslationGroup(N)

Base.convert(::Type{AbstractLieGroup}, ::Type{𝕥{N, T}}) where {N,T} = 𝕋{N}
Base.convert(::Type{AbstractLieAlgebra}, ::Type{𝕋{N, T}}) where {N,T} = 𝕥{N}

##==============================================================================
## TODO Quaternion Group https://github.com/JuliaManifolds/Manifolds.jl/issues/382
##==============================================================================

##==============================================================================
## RoME like functions sandbox
##==============================================================================
export PosePose, PosePose_x, PosePose_d, Prior_d

# this is in the "local frame" sequential, a similar operation can be done with the group_log (but together)
function PosePose(ⁱmⱼ::SE{N}, ʷxᵢ::SE{N}, ʷxⱼ::SE{N}) where N

    ʷT̂ⱼ = ʷxᵢ ∘ ⁱmⱼ

    ʲT̂ⱼ = ʷxⱼ \ ʷT̂ⱼ

    X = log(ʲT̂ⱼ)

    return vee(X)
end

# this is in the "global frame" e component wise on submanifolds
function PosePose_x(ⁱmⱼ::SE{N}, ʷxᵢ::SE{N}, ʷxⱼ::SE{N}) where N
    M = manifold(ʷxᵢ)

    ʷT̂ⱼ = ʷxᵢ ∘ ⁱmⱼ
    X = se{N}(log(M, ʷxⱼ[], ʷT̂ⱼ[])) # TODO check sign, this (currently) is the vector at ʷT̂ⱼ that points to ʷxⱼ

    return vee(X)
end

# this is in the "global frame" e component wise on submanifolds, but uses distance
function PosePose_d(m, p, q)

    q̂ = p ∘ m
    return distance(q, q̂)

end

#manifolds point to manifold point factor
function PointPoint_d(m, p, q)

    q̂ = p ∘ m
    return distance(q, q̂)

end

function PointPoint_d(X::AbstractLieAlgebra, p, q)
    q̂ = exp(p, X)
    # q̂ = p ⊕ X
    return distance(q, q̂)
end

# manifold prior factor
function Prior_d(meas, p)
#
    return distance(meas, p)
end

function PointPoint_velocity_d(X, dt, p, q)
    q̂ = p ∘ group_exp(X*dt)
    return distance(q, q̂)
end

function Pose2Pose2(meas, wxi, wxj)
    wTjhat = SE2(wxi)*SE2(meas)
    jTjhat = SE2(wxj) \ wTjhat
    return se2vee(jTjhat)
end

end


text:
"""
    BalanceLaw

An abstract type representing a PDE balance law of the form

elements for balance laws of the form

```math
q_{,t} + Σ_{i=1,...d} F_{i,i} = s
```

Subtypes `L` should define the following methods:
- `vars_state_conservative(::L)`: a tuple of symbols containing the state variables
- `vars_state_auxiliary(::L)`: a tuple of symbols containing the auxiliary variables
- `vars_state_gradient(::L)`: a tuple of symbols containing the transformed variables of which gradients are computed
- `vars_state_gradient_flux(::L)`: a tuple of symbols containing the diffusive variables
- `vars_gradient_laplacian(::L)`: a tuple of symbols containing the transformed variables of which gradients of laplacian are computed, they have to be a subset of `vars_state_gradient`
- `vars_hyperdiffusive(::L)`: a tuple of symbols containing the hyperdiffusive variables
- `flux_first_order!(::L, flux::Grad, state_conservative::Vars, state_auxiliary::Vars, t::Real)`
- `flux_second_order!(::L, flux::Grad, state_conservative::Vars, state_gradient_flux::Vars, Qhypervisc_div::Vars, state_auxiliary::Vars, t::Real)`
- `compute_gradient_argument!(::L, transformstate::Vars, state_conservative::Vars, state_auxiliary::Vars, t::Real)`: transformation of state variables to variables of which gradients are computed
- `compute_gradient_flux!(::L, state_gradient_flux::Vars, ∇transformstate::Grad, state_auxiliary::Vars, t::Real)`: transformation of gradients to the diffusive variables
- `transform_post_gradient_laplacian!(::L, Qhypervisc_div::Vars, ∇Δtransformstate::Grad, state_auxiliary::Vars, t::Real)`: transformation of laplacian gradients to the hyperdiffusive variables
- `source!(::L, source::Vars, state_conservative::Vars, diffusive::Vars, state_auxiliary::Vars, t::Real)`
- `wavespeed(::L, n⁻, state_conservative::Vars, state_auxiliary::Vars, t::Real)`
- `boundary_state!(::NumericalFluxGradient, ::L, state_conservative⁺::Vars, state_auxiliary⁺::Vars, normal⁻, state_conservative⁻::Vars, state_auxiliary⁻::Vars, bctype, t)`
- `boundary_state!(::NumericalFluxFirstOrder, ::L, state_conservative⁺::Vars, state_auxiliary⁺::Vars, normal⁻, state_conservative⁻::Vars, state_auxiliary⁻::Vars, bctype, t)`
- `boundary_state!(::NumericalFluxSecondOrder, ::L, state_conservative⁺::Vars, state_gradient_flux⁺::Vars, state_auxiliary⁺::Vars, normal⁻, state_conservative⁻::Vars, state_gradient_flux⁻::Vars, state_auxiliary⁻::Vars, bctype, t)`
- `init_state_auxiliary!(::L, state_auxiliary::Vars, coords, args...)`
- `init_state_conservative!(::L, state_conservative::Vars, state_auxiliary::Vars, coords, args...)`

"""
abstract type BalanceLaw end # PDE part

# function stubs
function vars_state_conservative end
function vars_state_auxiliary end

function vars_state_gradient end
function vars_state_gradient_flux end

vars_gradient_laplacian(::BalanceLaw, FT) = @vars()
vars_hyperdiffusive(::BalanceLaw, FT) = @vars()

vars_integrals(::BalanceLaw, FT) = @vars()
vars_reverse_integrals(::BalanceLaw, FT) = @vars()

number_state_conservative(m::BalanceLaw, FT) =
    varsize(vars_state_conservative(m, FT))
number_state_auxiliary(m::BalanceLaw, FT) = varsize(vars_state_auxiliary(m, FT))

number_state_gradient(m::BalanceLaw, FT) = varsize(vars_state_gradient(m, FT))
number_state_gradient_flux(m::BalanceLaw, FT) =
    varsize(vars_state_gradient_flux(m, FT))

num_gradient_laplacian(m::BalanceLaw, FT) =
    varsize(vars_gradient_laplacian(m, FT))
num_hyperdiffusive(m::BalanceLaw, FT) = varsize(vars_hyperdiffusive(m, FT))

num_integrals(m::BalanceLaw, FT) = varsize(vars_integrals(m, FT))
num_reverse_integrals(m::BalanceLaw, FT) =
    varsize(vars_reverse_integrals(m, FT))

function init_state_conservative! end
function init_state_auxiliary! end

function flux_first_order! end
function flux_second_order! end
function source! end

compute_gradient_argument!(::BalanceLaw, args...) = nothing
compute_gradient_flux!(::BalanceLaw, args...) = nothing
function transform_post_gradient_laplacian! end

function wavespeed end
function boundary_state! end

function update_auxiliary_state! end
function update_auxiliary_state_gradient! end

function integral_load_auxiliary_state! end
function integral_set_auxiliary_state! end
function reverse_integral_load_auxiliary_state! end
function reverse_integral_set_auxiliary_state! end

using ..Courant
"""
    calculate_dt(dg, model, Q, Courant_number, direction, t)

For a given model, compute a time step satisying the nondiffusive Courant number
`Courant_number`
"""
function calculate_dt(dg, model, Q, Courant_number, t, direction)
    Δt = one(eltype(Q))
    CFL = courant(nondiffusive_courant, dg, model, Q, Δt, t, direction)
    return Courant_number / CFL
end

function create_conservative_state(balance_law::BalanceLaw, grid)
    topology = grid.topology
    # FIXME: Remove after updating CUDA
    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    Np = dofs_per_element(grid)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    V = vars_state_conservative(balance_law, FT)
    state_conservative = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        number_state_conservative(balance_law, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )
    return state_conservative
end

function create_auxiliary_state(balance_law, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    V = vars_state_auxiliary(balance_law, FT)
    state_auxiliary = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        number_state_auxiliary(balance_law, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    dim = dimensionality(grid)
    polyorder = polynomialorder(grid)
    vgeo = grid.vgeo
    device = typeof(state_auxiliary.data) <: Array ? CPU() : CUDA()
    nrealelem = length(topology.realelems)
    event = Event(device)
    event = kernel_init_state_auxiliary!(device, min(Np, 1024), Np * nrealelem)(
        balance_law,
        Val(dim),
        Val(polyorder),
        state_auxiliary.data,
        vgeo,
        topology.realelems,
        dependencies = (event,),
    )
    event = MPIStateArrays.begin_ghost_exchange!(
        state_auxiliary;
        dependencies = event,
    )
    event = MPIStateArrays.end_ghost_exchange!(
        state_auxiliary;
        dependencies = event,
    )
    wait(device, event)

    return state_auxiliary
end

function create_gradient_state(balance_law, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    # TODO: Clean up this MPIStateArray interface...
    V = vars_state_gradient_flux(balance_law, FT)
    state_gradient_flux = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        number_state_gradient_flux(balance_law, FT),
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    return state_gradient_flux
end

function create_higher_order_states(balance_law, grid)
    topology = grid.topology
    Np = dofs_per_element(grid)

    h_vgeo = Array(grid.vgeo)
    FT = eltype(h_vgeo)
    DA = arraytype(grid)

    weights = view(h_vgeo, :, grid.Mid, :)
    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))

    ngradlapstate = num_gradient_laplacian(balance_law, FT)
    # TODO: Clean up this MPIStateArray interface...
    V = vars_gradient_laplacian(balance_law, FT)
    Qhypervisc_grad = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        3ngradlapstate,
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )

    V = vars_hyperdiffusive(balance_law, FT)
    Qhypervisc_div = MPIStateArray{FT, V}(
        topology.mpicomm,
        DA,
        Np,
        ngradlapstate,
        length(topology.elems),
        realelems = topology.realelems,
        ghostelems = topology.ghostelems,
        vmaprecv = grid.vmaprecv,
        vmapsend = grid.vmapsend,
        nabrtorank = topology.nabrtorank,
        nabrtovmaprecv = grid.nabrtovmaprecv,
        nabrtovmapsend = grid.nabrtovmapsend,
        weights = weights,
    )
    return Qhypervisc_grad, Qhypervisc_div
end


text:
if basename(pwd()) == "aoc"
    cd("2021/20")
end
using Base.Iterators: cycle, partition, Stateful

function loadstartingpositions(filename::AbstractString)
    [parse(Int, last(line)) for line in eachline(filename)]
end

part1(filename::AbstractString) = part1(loadstartingpositions(filename))

examplepositions = [4, 8]
inputpositions = [6, 7]

function part1(positions::Vector{Int})
    positions = copy(positions)
    scores = [0, 0]
    rolls = 0
    index = 0
    dice = Stateful(sum(p) for p in partition(cycle(1:100), 3))
    while all(<(1000), scores)
        index = offsetmod(index + 1, 2)
        positions[index] = offsetmod(positions[index] + popfirst!(dice), 10)
        scores[index] += positions[index]
        rolls += 3
    end
    rolls * minimum(scores)
end

offsetmod(x, n) = ((x - 1) % n) + 1

max_stars_repo_name: stellartux/advent-of-code

max_issues_repo_name: stellartux/advent-of-code

max_forks_repo_name: stellartux/advent-of-code


text:
# __precompile__()

module JuliaRC

using Compat, OhMyREPL

source_files = [
    "constants.jl",
    "environment.jl",
    "aliases.jl",
    "helpers.jl",
    "greedy_banner.jl",
    "macros.jl",
    "emerge.jl",
    "repl.jl",
    "setup.jl",
]

foreach(include, source_files)

end # module


text:
using LinearAlgebra

import GLPK

# Needed until gets merged in GLPK.jl
MOI.get(::GLPK.Optimizer, ::MOI.BarrierIterations) = 0
MOI.get(::GLPK.Optimizer, ::MOI.SimplexIterations) = 0

# Dummy oracle: minimize c'x over the hypercube
mutable struct DummyOracle <: Linda.Oracle
    id::Int
    dim::Int

    c::Vector{Float64}
    x::Vector{Float64}

    farkas::Bool
    π::Vector{Float64}
    σ::Float64

    DummyOracle(id::Int, dim::Int, c) = new(id, dim, c, zeros(dim), false, zeros(dim), 0.0)

end

function Linda.update!(o::DummyOracle, farkas, π, σ)
    o.farkas = farkas
    o.π .= π
    o.σ = σ
    return nothing
end

Linda.optimize!(::DummyOracle) = nothing

function Linda.get_columns(o::DummyOracle)
    o.x .= 0.0
    rc = -o.σ

    if o.farkas
        # Farkas pricing
        for i in 1:o.dim
            if -o.π[i] <= 0.0
                o.x[i] = 1.0
                rc -= o.π[i]
            end
        end
    else
        # Regular pricing
        for i in 1:o.dim
            if o.c[i] - o.π[i] < 0.0
                o.x[i] = 1.0
                rc += o.c[i] - o.π[i]
            end
        end
    end

    col = Linda.Column(dot(o.x, o.c), o.id, collect(1:o.dim), copy(o.x), true)
    return [(col, rc)]
end

Linda.get_dual_bound(o::DummyOracle) = o.farkas ? -Inf : (dot(o.x, o.c) - dot(o.x, o.π) - o.σ)

Linda.get_objective_value(o::DummyOracle) = o.farkas*dot(o.x, o.c) - dot(o.x, o.π) - o.σ

# Create initial RMP
import GLPK
rmp = GLPK.Optimizer()
# Disable RMP solver's output
MOI.set(rmp, MOI.Silent(), true)
# Set objective sense to minimization, otherwise all the duals are wrong
MOI.set(rmp, MOI.ObjectiveSense(), MOI.MIN_SENSE)

m = 4  # number of linking constraints
R = 1  # number of sub-problems
b = [0.2, 0.4, 0.6, 0.8]

# Create the convexity constraints
con_cvx = MOI.ConstraintIndex[]
for r in 1:R
    cidx = MOI.add_constraint(
        rmp,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.0),
        MOI.EqualTo(1.0)
    )
    push!(con_cvx, cidx)
end
# Create linking constraints
con_link = MOI.ConstraintIndex[]
for i in 1:m
    cidx = MOI.add_constraint(
        rmp,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.0),
        MOI.LessThan(b[i])
    )
    push!(con_link, cidx)
end

mp = Linda.Master(rmp, con_cvx, con_link, b, MOI.VariableIndex[], MOI.VariableIndex[])

# Add a single column
# Linda.add_column!(mp,
#     Linda.Column(
#         0.0, 1, Int[], Float64[], true
# ))

# Add a single column
# Linda.add_column!(mp,
#     Linda.Column(
#         -4.0, 1, [1], [1.0], true
# ))
# Linda.add_column!(mp,
#     Linda.Column(
#         -8.0, 1, [2], [1.0], true
# ))
# Linda.add_column!(mp,
#     Linda.Column(
#         -12.0, 1, [3], [1.0], true
# ))
# Linda.add_column!(mp,
#     Linda.Column(
#         -16.0, 1, [4], [1.0], true
# ))
# The following three give an optimal solution
# Linda.add_column!(mp,
#     Linda.Column(
#         -16.0, 1, [1, 3], [1.0, 1.0], true
# ))
# Linda.add_column!(mp,
#     Linda.Column(
#         -40.0, 1, [1, 2, 3, 4], [1.0, 1.0, 1.0, 1.0], true
# ))

# Setup the sub-problems
oracles = [DummyOracle(1, m, [-5.0, -4.0, -3.0, -2.0])]

# CG parameters
env = Linda.LindaEnv()
env.num_cgiter_max.val = 10
env.verbose.val = 1

# Solve the problem with Column Generation
Linda.solve_colgen!(env, mp, oracles)


text:
"""
    idx(I)

Handle different indexing types (ALL, scalar, range-based, and vector). Returns the
proper format for GraphBLAS indexing. Should *not* be used for functions that take a single
scalar index like [`extractElement`].
"""
function idx(I)
    if I == ALL
        return I, 0 #ni doesn't matter if I=ALL
    elseif I isa UnitRange
        return [I.start, I.stop], libgb.GxB_RANGE #Simple ranges
    elseif I isa StepRange
        #The step must survive tozerobased(I), so we add 1 to it.
        if I.step > 0
            return [I.start, I.stop, I.step + 1], libgb.GxB_STRIDE #Normal stried ranges
        elseif I.step < 0
            #Strided ranges with negative increment
            return [I.start, I.stop, -I.step + 1], libgb.GxB_BACKWARDS
        end
    elseif I isa Vector
        return Vector{libgb.GrB_Index}(I), length(I) #Assume ni = length(I) otherwise
    elseif I isa Integer
        return [UInt64(I)], 1
    elseif I isa CartesianIndex{1}
        return [UInt64(I[1])], 1
    else
        throw(TypeError(:idx,
            Union{UnitRange, StepRange, Vector, Integer}, typeof(I)))
    end
end