text:
using Revise
using NNFEM
using PyPlot

NT = 200
T = 1.0
Δt = T/NT

m, n =  15, 15
h = 1/m

# Create a very simple mesh
elements = []
prop = Dict("name"=> "PlaneStrain", "rho"=> 1.0, "E"=> 2.0, "nu"=> 0.35)
coords = zeros((m+1)*(n+1), 2)
for j = 1:n
    for i = 1:m
        idx = (m+1)*(j-1)+i
        elnodes = [idx; idx+1; idx+1+m+1; idx+m+1]
        ngp = 3
        nodes = [
            (i-1)*h (j-1)*h
            i*h (j-1)*h
            i*h j*h
            (i-1)*h j*h
        ]
        coords[elnodes, :] = nodes
        push!(elements, SmallStrainContinuum(nodes, elnodes, prop, ngp))
    end
end

# fixed on the bottom, push on the right
EBC = zeros(Int64, (m+1)*(n+1), 2)
FBC = zeros(Int64, (m+1)*(n+1), 2)
g = zeros((m+1)*(n+1), 2)
f = zeros((m+1)*(n+1), 2)

for i = 2:m
    idx = n*(m+1)+i
    EBC[idx,:] .= -2
    EBC[i,:] .= -2
end

for j = 1:n+1
    idx = (j-1)*(m+1)+1
    EBC[idx,:] .= -2
    idx = (j-1)*(m+1)+m+1
    EBC[idx,:] .= -2
end

ndims = 2
domain = Domain(coords, elements, ndims, EBC, g, FBC, f)

EBC_idx = findall(EBC[:,1] .== -2)
EBC_X = domain.nodes[EBC_idx,:]
function gt_func(t)
    x, y = EBC_X[:,1], EBC_X[:,2]
    u1 = @. (1-y^2)*(x^2+y^2)*exp(-t) * 0.1
    u2 = @. (1-y^2)*(x^2-y^2)*exp(-t) * 0.1
    return [u1;u2], -[u1;u2], [u1;u2]
end

FBC_idx = findall(FBC[:,1] .== -2)
FBC_X = domain.nodes[FBC_idx,:]
# function ft(t)
#     x, y = FBC_X[:,1], FBC_X[:,2]
#     T1 = @. -1.48148148148148*x*exp(-t) - 1.48148148148148*y*exp(-t)
#     T2 = @. -3.45679012345679*x*exp(-t) - 6.41975308641975*y*exp(-t)
#     return [T1;T2] * h * 0.1
# end

function bt(x, y, t)
    f1 = @. 0.987654320987654*x*y*exp(-t) + 0.592592592592593*y^2*exp(-t) + (0.1 - 0.1*y^2)*(x^2 + y^2)*exp(-t) - (0.148148148148148 - 0.148148148148148*y^2)*exp(-t) - (0.641975308641975 - 0.641975308641975*y^2)*exp(-t) - (-0.148148148148148*x^2 - 0.148148148148148*y^2)*exp(-t)
    f2 = @. 0.987654320987654*x*y*exp(-t) - 2.5679012345679*y^2*exp(-t) + (0.1 - 0.1*y^2)*(x^2 - y^2)*exp(-t) - (0.148148148148148 - 0.148148148148148*y^2)*exp(-t) - (-0.641975308641975*x^2 + 0.641975308641975*y^2)*exp(-t) - (0.641975308641975*y^2 - 0.641975308641975)*exp(-t)
    [f1 f2]
end

c = domain.nodes[domain.dof_to_eq]
nn = div(length(c), 2)
x = c[1:nn]
y = c[nn+1:end]

Dstate = zeros(domain.neqs) # d at last step
state = [(@. (1-y^2)*(x^2+y^2)); (@. (1-y^2)*(x^2-y^2))] * 0.1
velo = -[(@. (1-y^2)*(x^2+y^2)); (@. (1-y^2)*(x^2-y^2))] * 0.1
acce = [(@. (1-y^2)*(x^2+y^2)); (@. (1-y^2)*(x^2-y^2))] * 0.1
globdat = GlobalData(state, Dstate, velo, acce, domain.neqs, gt_func, nothing, bt)

assembleMassMatrix!(globdat, domain)
updateTimeDependentEssentialBoundaryCondition!(domain,globdat)
updateStates!(domain, globdat)

# #solve for a0
# SolverInitial!(Δt, globdat, domain)

ω = EigenMode(Δt, globdat, domain)
@show "stable time step is ", 0.8 * 2/ω, " current time step is ", Δt

for i = 1:NT
    @info i
    #global globdat, domain = GeneralizedAlphaSolverStep(globdat, domain, Δt)
    global globdat, domain = ExplicitSolverStep(globdat, domain, Δt)
end
# # # visualize_displacement(domain)
# plot(hcat(domain.history["state"]...)[n*(m+1)+1,:])
# plot(exp.(-LinRange(0,1,NT+1)))

ts = ExplicitSolverTime(Δt, NT)
d_ = hcat(domain.history["state"]...)'

using Random; Random.seed!(233)
for k = 1:5
    i = rand(1:m+1)
    j = rand(1:n+1)
    if k==1
        plot(d_[:,(j-1)*(m+1)+i], color = "C$k", label="Computed")
    else
        plot(d_[:,(j-1)*(m+1)+i], color = "C$k")
    end
    x0 = (i-1)*h
    y0 = (j-1)*h

    if k==1
        plot((@. (1-y0^2)*(x0^2+y0^2)*exp(-ts))*0.1 ,"--", color="C$k", label="Reference")
    else
        plot((@. (1-y0^2)*(x0^2+y0^2)*exp(-ts))*0.1 ,"--", color="C$k")
    end
end
legend()


text:
default_template(paths::String...) = pkgdir(Blueprints, "templates", paths...)

"""
    is_blueprint(x)

Return `true` if `x` is a blueprint type or instance.
"""
is_blueprint(x) = false

macro blueprint(ex)
    return esc(blueprint_m(__module__, ex))
end

macro blueprint(alias::String, ex)
    return esc(blueprint_m(__module__, ex, alias))
end

function blueprint_m(mod::Module, ex, type_alias=nothing)
    ex = macroexpand(mod, ex)
    def = Configurations.JLKwStruct(ex, type_alias)
    return quote
        # we alawys export the blueprint here
        # but this prob should be removed when
        # we generalize this concept
        export $(def.name)
        $(Configurations.codegen_option_type(mod, def))
        $Blueprints.is_blueprint(::$(def.name)) = true
        $Blueprints.is_blueprint(::Type{$(def.name)}) = true
        nothing
    end
end

"""
    @option struct Context
Type for storing the compile context.

# Fields
- `args::Tuple`: the input arguments.
- `kwargs::NamedTuple`: the input keyword arguments.
- `stack::NTuple{N, Blueprint}`: [`Blueprint`](@ref) stack, topmost is the current blueprint under compile.
- `state::Dict{Symbol, Any}`: compile states that each blueprint can pass in.
"""
Base.@kwdef struct Context{Args <: Tuple, Kwargs <: NamedTuple, Blueprint}
    # global inputs that is
    # not from config meta
    args::Args = ()
    kwargs::Kwargs = NamedTuple()
    root::Blueprint # the root blueprint
    # state
    state::Dict{Symbol, Any} = Dict{Symbol, Any}()
end

function assert_blueprint(x)
    is_blueprint(x) || error("expect a blueprint struct type defined via @blueprint")
end

"""
    blueprints(parent)

Return the component blueprints inside a given `blueprint`
"""
function blueprints(x)
    assert_blueprint(x)
    return ()
end

function has_blueprint(ctx::Context, ::Type{T}) where T
    assert_blueprint(T)
    for each in blueprints(ctx.root)
        each isa T && return true
    end
    return false
end

function get_blueprint(ctx::Context, ::Type{T}) where T
    assert_blueprint(T)
    for each in blueprints(ctx.root)
        each isa T && return each
    end
    return
end

abstract type MustacheBlueprint end
function blueprint_view(::MustacheBlueprint, ctx::Context)
    return Dict{String, Any}()
end

function src_file(bp::MustacheBlueprint, ::Context)
    return joinpath(bp.template.root, bp.template.path, bp.template.file)
end

function dst_file(bp::MustacheBlueprint, ctx::Context)
    return joinpath(ctx.kwargs.path, bp.template.path, bp.template.file)
end

function compile(bp::MustacheBlueprint, ctx::Context)
    file = src_file(bp, ctx)::String
    text = render_from_file(file, blueprint_view(bp, ctx))
    dst = dst_file(bp, ctx)::String
    mkpath(dirname(dst))
    write(dst, text)
    return
end

@option struct TemplateFile
    root::String = default_template()
    path::String = "."
    file::String
end

"""
    Badge(hover::AbstractString, image::AbstractString, link::AbstractString)

Container for Markdown badge data.
Each argument can contain placeholders,
which will be filled in with values from [`combined_view`](@ref).

## Arguments
- `hover::AbstractString`: Text to appear when the mouse is hovered over the badge.
- `image::AbstractString`: URL to the image to display.
- `link::AbstractString`: URL to go to upon clicking the badge.
"""
@option struct Badge
    hover::String
    image::String
    link::String
end

Base.string(b::Badge) = "[![$(b.hover)]($(b.image))]($(b.link))"

ext: jl


text:
const SUDOKU_BOX_SIZE = 3
const SUDOKU_GROUP_SIZE = SUDOKU_BOX_SIZE * SUDOKU_BOX_SIZE
const SUDOKU_CANDIDATE_SIZE = SUDOKU_GROUP_SIZE
const SUDOKU_CELL_SIZE = SUDOKU_GROUP_SIZE * SUDOKU_GROUP_SIZE
const SUDOKU_NO_CANDIDATES = 0
const SUDOKU_ALL_CANDIDATES = ((1 << SUDOKU_CANDIDATE_SIZE) - 1)
const SudokuCandidate=UInt8
const SudokuIndex=Int64
const SudokuMap=Array{SudokuCandidate, 2}

function overwrite_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex, candidate_index::SudokuIndex)
    sudoku_map[:, cell_index] .= zero(sudoku_map[1,1])
    sudoku_map[candidate_index, cell_index] = one(sudoku_map[1,1])
end

function string_to_sudokumap(line::String)
    sudoku_map = SudokuMap(ones(SudokuCandidate, SUDOKU_CANDIDATE_SIZE, SUDOKU_CELL_SIZE))
    for (index, c) in enumerate(line)
        if index > SUDOKU_CELL_SIZE
            break
        end

        num = Int(c) - Int('0')
        if num > 0 && num <= SUDOKU_CANDIDATE_SIZE
            overwrite_sudokumap(sudoku_map, index, num)
        end
    end
    sudoku_map
end

function to_string_sudokumap(sudoku_map::SudokuMap, one_line::Bool)
    result_str = ""

    for cell_index in 1:SUDOKU_CELL_SIZE
        for candidate in 1:SUDOKU_CANDIDATE_SIZE
            if sudoku_map[candidate, cell_index] != zero(sudoku_map[1,1])
                result_str *= string(candidate)
            end
        end

        if !one_line
            result_str *= ":"
        end

        if cell_index % SUDOKU_GROUP_SIZE == zero(sudoku_map[1,1]) && !one_line
            result_str *= "\n"
        end
    end

    result_str
end

function is_solved_group(sudoku_group)
    all(count -> (count == one(sudoku_group[1,1])), sum(sudoku_group, dims=(1))) && all(count -> (count == one(sudoku_group[1,1])), sum(sudoku_group, dims=(2)))
end

function is_solved_sudoku_rows(sudoku_map::SudokuMap)
    all(row ->
        (index = SUDOKU_GROUP_SIZE * (row - 1);
         is_solved_group(sudoku_map[:, index + 1: index + SUDOKU_GROUP_SIZE])),
        1:SUDOKU_GROUP_SIZE)
end

function is_solved_sudoku_columns(sudoku_map::SudokuMap)
    all(column ->
        (indexes = collect(column: SUDOKU_GROUP_SIZE: column + SUDOKU_GROUP_SIZE * (SUDOKU_GROUP_SIZE - 1));
         is_solved_group(view(sudoku_map, :, indexes))),
        1:SUDOKU_GROUP_SIZE)
end

function is_solved_sudoku_box_group(sudoku_map::SudokuMap, out_box_row::SudokuIndex, out_box_column::SudokuIndex)
    group = zeros(SudokuCandidate, SUDOKU_GROUP_SIZE, SUDOKU_CANDIDATE_SIZE)
    for in_box_index in 1:SUDOKU_BOX_SIZE
        base = (out_box_column - 1) * SUDOKU_BOX_SIZE
        base += (out_box_row - 1) * SUDOKU_GROUP_SIZE * SUDOKU_BOX_SIZE
        base += (in_box_index - 1) * SUDOKU_GROUP_SIZE

        sources = collect(1:SUDOKU_BOX_SIZE)
        sources .+= base
        dest = collect(1:SUDOKU_BOX_SIZE)
        dest .+= (in_box_index - 1) * SUDOKU_BOX_SIZE

        group[:, dest] = view(sudoku_map, :, sources)
    end

    is_solved_group(group)
end

function is_solved_sudoku_boxes(sudoku_map::SudokuMap)
    for out_box_row in 1:SUDOKU_BOX_SIZE
        for out_box_column in 1:SUDOKU_BOX_SIZE
            if !is_solved_sudoku_box_group(sudoku_map, out_box_row, out_box_column)
                return false
            end
        end
    end
    true
end

function is_solved_sudokumap(sudoku_map::SudokuMap)
    is_solved_sudoku_rows(sudoku_map) && is_solved_sudoku_columns(sudoku_map) && is_solved_sudoku_boxes(sudoku_map)
end

function is_consistent_sudokumap(sudoku_map::SudokuMap)
    all(count -> (count > zero(sudoku_map[1,1])), sum(sudoku_map, dims=(1)))
end

function has_unique_candidate_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex)
    sum(sudoku_map[:, cell_index]) == one(sudoku_map[1,1])
end

function set_unused_candidates_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex, other_index::SudokuIndex, candidates::Array{SudokuCandidate,1})
    if other_index != cell_index && has_unique_candidate_sudokumap(sudoku_map, other_index)
        candidates[:] .|= sudoku_map[:, other_index]
    end
end

function find_unique_candidates_sudokumap(sudoku_map::SudokuMap)
    for cell_index in 1:SUDOKU_CELL_SIZE
        if has_unique_candidate_sudokumap(sudoku_map, cell_index)
           continue
        end

        row = div(cell_index - 1, SUDOKU_GROUP_SIZE)
        column = (cell_index - 1) % SUDOKU_GROUP_SIZE
        box_base = div(row, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE * SUDOKU_GROUP_SIZE
        box_base += div(column, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE + 1

        candidates = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for index in 1:SUDOKU_GROUP_SIZE
            in_row_index = row * SUDOKU_GROUP_SIZE + index
            in_column_index = column + (index - 1) * SUDOKU_GROUP_SIZE + 1
            in_box_index = div(index - 1, SUDOKU_BOX_SIZE) * SUDOKU_GROUP_SIZE
            in_box_index += (index - 1) % SUDOKU_BOX_SIZE + box_base
            for other_index in [in_row_index, in_column_index, in_box_index]
                set_unused_candidates_sudokumap(sudoku_map, cell_index, other_index, candidates)
            end
        end

        for (index, candidate) in enumerate(candidates)
            if candidate != 0
                sudoku_map[index, cell_index] = zero(sudoku_map[1,1])
            end
        end

        if sum(sudoku_map[:, cell_index]) == 0
            return false
        end
    end

    true
end

function set_unused_candidates_sudokumap(sudoku_map::SudokuMap, cell_index::SudokuIndex, candidates::Array{SudokuCandidate,1})
    if sum(candidates) == SUDOKU_CANDIDATE_SIZE - 1
        index = findfirst(x -> x == 0, candidates)
        overwrite_sudokumap(sudoku_map, cell_index, index)
    end
end

function fill_unused_candidates_sudokumap(sudoku_map::SudokuMap)
    for cell_index in 1:SUDOKU_CELL_SIZE
        if has_unique_candidate_sudokumap(sudoku_map, cell_index)
           continue
        end

        cell_row = div((cell_index - 1), SUDOKU_GROUP_SIZE)
        candidates_row = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for column_index in 1:SUDOKU_GROUP_SIZE
            in_row_index = cell_row * SUDOKU_GROUP_SIZE + column_index
            if in_row_index != cell_index
                candidates_row[:] .|= sudoku_map[:, in_row_index]
            end
        end
        set_unused_candidates_sudokumap(sudoku_map, cell_index, candidates_row)

        cell_column = (cell_index - 1) % SUDOKU_GROUP_SIZE
        candidates_column = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for row_index in 1:SUDOKU_GROUP_SIZE
            in_column_index = cell_column + (row_index - 1) * SUDOKU_GROUP_SIZE + 1
            if in_column_index != cell_index
                candidates_column[:] .|= sudoku_map[:, in_column_index]
            end
        end
        set_unused_candidates_sudokumap(sudoku_map, cell_index, candidates_column)

        box_base = div(cell_row, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE * SUDOKU_GROUP_SIZE
        box_base += div(cell_column, SUDOKU_BOX_SIZE) * SUDOKU_BOX_SIZE + 1
        candidates_box = zeros(SudokuCandidate, SUDOKU_CANDIDATE_SIZE)
        for box_index in 1:SUDOKU_GROUP_SIZE
            in_box_index = div(box_index - 1, SUDOKU_BOX_SIZE) * SUDOKU_GROUP_SIZE
            in_box_index += (box_index - 1) % SUDOKU_BOX_SIZE + box_base
            if in_box_index != cell_index
                candidates_box[:] .|= sudoku_map[:, in_box_index]
            end
        end
        set_unused_candidates_sudokumap(sudoku_map, cell_index, candidates_box)

        if sum(sudoku_map[:, cell_index]) == zero(sudoku_map[1, 1])
            return false
        end
    end

    true
end

function find_backtracking_target(sudoku_map::SudokuMap)
    min_count = SUDOKU_CELL_SIZE
    found = false
    target_index = 0

    candidate_counts = sum(sudoku_map, dims=(1))
    for (index, count) in enumerate(candidate_counts)
        if count <= one(sudoku_map[1,1])
            continue
        end

        if min_count > count
            min_count = count
            found = true
        end
    end

    extract_row = (row -> view(candidate_counts,
                               ((row - 1) * SUDOKU_GROUP_SIZE + 1):(row * SUDOKU_GROUP_SIZE)))
    if found
        row_counts = map(row -> count(n -> (n == min_count), extract_row(row)), 1:SUDOKU_GROUP_SIZE)
        row = argmax(row_counts)
        target_index = findfirst(x -> (x == min_count), extract_row(row)) + (row - 1) * SUDOKU_GROUP_SIZE
    end

    target_index, found
end

function solve_sudokumap(sudoku_map::SudokuMap)
    n_candidates = sum(sudoku_map)
    while true
        if find_unique_candidates_sudokumap(sudoku_map) == false
            return (sudoku_map, false)
        end

        if fill_unused_candidates_sudokumap(sudoku_map) == false
            return (sudoku_map, false)
        end

        new_n_candidates = sum(sudoku_map)
        if n_candidates == new_n_candidates
            break
        end

        n_candidates = new_n_candidates
    end

    if is_solved_sudokumap(sudoku_map)
        return (sudoku_map, true)
    end

    if !is_consistent_sudokumap(sudoku_map)
        return (sudoku_map, false)
    end

    target_index, found = find_backtracking_target(sudoku_map)
    if !found
        return (sudoku_map, false)
    end

    for (index, candidate) in enumerate(sudoku_map[:, target_index])
        if candidate == zero(sudoku_map[1,1])
            continue
        end

        new_map = deepcopy(sudoku_map)
        overwrite_sudokumap(new_map, target_index, index)
        result_map, result = solve_sudokumap(new_map)
        if result
            return (result_map, result)
        end
    end

    return (sudoku_map, false)
end

function main()
    num = 0
    one_line = true
    total_result = true

    if one_line
        println("Solving in Julia")
    end

    for line in eachline(stdin)
        puzzle = string_to_sudokumap(line)
        solution, result = solve_sudokumap(puzzle)
        result = is_solved_sudokumap(solution)
        total_result &= result

        one_line_result = to_string_sudokumap(solution, one_line)
        print(one_line_result)
        if one_line
            println("")
        else
            if result
                println("Solved\n")
            else
                println("Not solved\n")
            end
        end
        num += 1
    end

    if total_result
        println("All ", num, " cases passed.")
    else
        println("Some of ", num, " cases are not solved correctly.")
    end
end

main()

max_stars_repo_name: zettsu-t/sudokusse

max_issues_repo_name: zettsu-t/sudokusse

max_forks_repo_name: zettsu-t/sudokusse


text:
@testset "errors" begin

str = """
[foo]
bar = 3

[foo]
quiz = 3
"""

err = tryparsestring(str)
@test err isa ParserError
@test err.type == Internals.ErrDuplicatedKey

str = """
[[foo.bar]]

[foo]
bar = 2
"""
err = tryparsestring(str)
@test err isa ParserError
@test err.type == Internals.ErrKeyAlreadyHasValue

str = """
[[foo.bar]]

[foo.bar]
q = 3
"""
err = tryparsestring(str)
@test err isa ParserError

end


text:
using Distributions: Beta
using Base.Threads
using Random

mutable struct MCTSNode
    state::GameState
    parent::Union{Nothing, MCTSNode}
    move::Union{Nothing, Move}
    win_counts::Dict{Player, Int}
    num_rollouts::Int
    children::Vector{MCTSNode}
    unvisited_moves::Vector{Move}
end

function MCTSNode(
    state::GameState,
    parent::Union{Nothing, MCTSNode},
    move::Union{Nothing, Move};
    pseudocounts::Int=1
)::MCTSNode
    MCTSNode(
        state, parent, move, Dict(white => pseudocounts, black => pseudocounts),
        pseudocounts*2, MCTSNode[], legal_moves(state))
end
MCTSNode(state::GameState)::MCTSNode = MCTSNode(state, nothing, nothing)

is_terminal(node::MCTSNode)::Bool = is_over(node.state)
can_add_child(node::MCTSNode)::Bool = !isempty(node.unvisited_moves)
winning_pct(node::MCTSNode, player::Player)::Float64 = node.win_counts[player] / node.num_rollouts

function add_random_child!(node::MCTSNode)::MCTSNode
    index = rand(1:length(node.unvisited_moves))
    new_move = node.unvisited_moves[index]
    deleteat!(node.unvisited_moves, index)
    new_state = apply_move(node.state, new_move)
    new_node = MCTSNode(new_state, node, new_move)
    push!(node.children, new_node)
    new_node
end

function record_win!(node::MCTSNode, winner::Player)::Nothing
    node.win_counts[winner] += 1
    node.num_rollouts += 1
    nothing
end

struct MCTSAgent <: AbstractAgent
    num_rounds::Int
    komi::Float64
end

function select_move(agent::MCTSAgent, state::GameState)
    root = MCTSNode(state)

    nodelist = MCTSNode[]
    for i in 1:agent.num_rounds
        node = root
        while !can_add_child(node) && !is_terminal(node)
            node = select_child(agent, node)
        end

        if can_add_child(node)
            node = add_random_child!(node)
        end

        push!(nodelist, node)
    end

    winners = Dict{MCTSNode, Player}()
    @threads for i in 1:length(nodelist)
        node = nodelist[i]
        finalstate = simulate_random_game(node.state)
        result = compute_game_result(finalstate, agent.komi)
        winners[node] = winnercolor(result)
    end

    for (node_, winner) in winners
        node = node_
        while node !== nothing
            record_win!(node, winner)
            node = node.parent
        end
    end

    best_idx = 0
    best_pct = -1.0
    for (i, c) in enumerate(root.children)
        c_pct = winning_pct(c, c.state.next_player)
        if c_pct > best_pct
            best_idx = i
            best_pct = c_pct
        end
    end
    best_move = root.children[best_idx].move

    return best_move
end

function select_child(agent::MCTSAgent, node::MCTSNode)
    best_val = -1.0
    best_idx = 0
    rollouts = sum([c.num_rollouts for c in node.children])
    for (i, c) in enumerate(node.children)
        wins = c.win_counts[c.state.next_player]
        val = rand(Beta(wins, rollouts - wins))
        if val > best_val
            best_idx = i
            best_val = val
        end
    end
    best_child = node.children[best_idx]

    return best_child
end


text:
# Refine elements in the DG solver based on a list of cell_ids that should be refined
function refine!(u_ode::AbstractVector, adaptor, mesh::TreeMesh{3},
                 equations, dg::DGSEM, cache, cells_to_refine)
  # Return early if there is nothing to do
  if isempty(cells_to_refine)
    return
  end

  # Determine for each existing element whether it needs to be refined
  needs_refinement = falses(nelements(dg, cache))

  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_refine = searchsortedfirst.(Ref(cache.elements.cell_ids[1:nelements(dg, cache)]),
                                          cells_to_refine)
  needs_refinement[elements_to_refine] .= true

  # Retain current solution data
  old_n_elements = nelements(dg, cache)
  old_u_ode = copy(u_ode)
  GC.@preserve old_u_ode begin # OBS! If we don't GC.@preserve old_u_ode, it might be GC'ed
    old_u = wrap_array(old_u_ode, mesh, equations, dg, cache)

    # Get new list of leaf cells
    leaf_cell_ids = local_leaf_cells(mesh.tree)

    # re-initialize elements container
    @unpack elements = cache
    resize!(elements, length(leaf_cell_ids))
    init_elements!(elements, leaf_cell_ids, mesh, dg.basis.nodes)
    @assert nelements(dg, cache) > old_n_elements

    resize!(u_ode, nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache))
    u = wrap_array(u_ode, mesh, equations, dg, cache)

    # Loop over all elements in old container and either copy them or refine them
    u_tmp1 = Array{eltype(u), 4}(undef, nvariables(equations), nnodes(dg), nnodes(dg), nnodes(dg))
    u_tmp2 = Array{eltype(u), 4}(undef, nvariables(equations), nnodes(dg), nnodes(dg), nnodes(dg))
    element_id = 1
    for old_element_id in 1:old_n_elements
      if needs_refinement[old_element_id]
        # Refine element and store solution directly in new data structure
        refine_element!(u, element_id, old_u, old_element_id,
                        adaptor, equations, dg, u_tmp1, u_tmp2)
        element_id += 2^ndims(mesh)
      else
        # Copy old element data to new element container
        @views u[:, .., element_id] .= old_u[:, .., old_element_id]
        element_id += 1
      end
    end
    # If everything is correct, we should have processed all elements.
    # Depending on whether the last element processed above had to be refined or not,
    # the counter `element_id` can have two different values at the end.
    @assert element_id == nelements(dg, cache) + 1 || element_id == nelements(dg, cache) + 2^ndims(mesh) "element_id = $element_id, nelements(dg, cache) = $(nelements(dg, cache))"
  end # GC.@preserve old_u_ode

  # re-initialize interfaces container
  @unpack interfaces = cache
  resize!(interfaces, count_required_interfaces(mesh, leaf_cell_ids))
  init_interfaces!(interfaces, elements, mesh)

  # re-initialize boundaries container
  @unpack boundaries = cache
  resize!(boundaries, count_required_boundaries(mesh, leaf_cell_ids))
  init_boundaries!(boundaries, elements, mesh)

  # re-initialize mortars container
  @unpack mortars = cache
  resize!(mortars, count_required_mortars(mesh, leaf_cell_ids))
  init_mortars!(mortars, elements, mesh)

  # Sanity check
  if isperiodic(mesh.tree) && nmortars(mortars) == 0
    @assert ninterfaces(interfaces) == ndims(mesh) * nelements(dg, cache) ("For $(ndims(mesh))D and periodic domains and conforming elements, the number of interfaces must be $(ndims(mesh)) times the number of elements")
  end

  return nothing
end

# TODO: Taal compare performance of different implementations
# Refine solution data u for an element, using L2 projection (interpolation)
function refine_element!(u::AbstractArray{<:Any,5}, element_id,
                         old_u, old_element_id,
                         adaptor::LobattoLegendreAdaptorL2, equations, dg,
                         u_tmp1, u_tmp2)
  @unpack forward_upper, forward_lower = adaptor

  # Store new element ids
  bottom_lower_left_id  = element_id
  bottom_lower_right_id = element_id + 1
  bottom_upper_left_id  = element_id + 2
  bottom_upper_right_id = element_id + 3
  top_lower_left_id     = element_id + 4
  top_lower_right_id    = element_id + 5
  top_upper_left_id     = element_id + 6
  top_upper_right_id    = element_id + 7

  @boundscheck begin
    @assert old_element_id >= 1
    @assert size(old_u, 1) == nvariables(equations)
    @assert size(old_u, 2) == nnodes(dg)
    @assert size(old_u, 3) == nnodes(dg)
    @assert size(old_u, 4) == nnodes(dg)
    @assert size(old_u, 5) >= old_element_id
    @assert     element_id >= 1
    @assert size(    u, 1) == nvariables(equations)
    @assert size(    u, 2) == nnodes(dg)
    @assert size(    u, 3) == nnodes(dg)
    @assert size(    u, 4) == nnodes(dg)
    @assert size(    u, 5) >= element_id + 7
  end

  # Interpolate to bottom lower left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_lower_left_id), forward_lower, forward_lower, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to bottom lower right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_lower_right_id), forward_upper, forward_lower, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to bottom upper left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_upper_left_id), forward_lower, forward_upper, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to bottom upper right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, bottom_upper_right_id), forward_upper, forward_upper, forward_lower,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top lower left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_lower_left_id), forward_lower, forward_lower, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top lower right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_lower_right_id), forward_upper, forward_lower, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top upper left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_upper_left_id), forward_lower, forward_upper, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  # Interpolate to top upper right element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, top_upper_right_id), forward_upper, forward_upper, forward_upper,
    view(old_u, :, :, :, :, old_element_id), u_tmp1, u_tmp2)

  return nothing
end

# Coarsen elements in the DG solver based on a list of cell_ids that should be removed
function coarsen!(u_ode::AbstractVector, adaptor, mesh::TreeMesh{3},
                  equations, dg::DGSEM, cache, child_cells_to_coarsen)
  # Return early if there is nothing to do
  if isempty(child_cells_to_coarsen)
    return
  end

  # Determine for each old element whether it needs to be removed
  to_be_removed = falses(nelements(dg, cache))
  # The "Ref(...)" is such that we can vectorize the search but not the array that is searched
  elements_to_remove = searchsortedfirst.(Ref(cache.elements.cell_ids[1:nelements(dg, cache)]),
                                          child_cells_to_coarsen)
  to_be_removed[elements_to_remove] .= true

  # Retain current solution data
  old_n_elements = nelements(dg, cache)
  old_u_ode = copy(u_ode)
  GC.@preserve old_u_ode begin # OBS! If we don't GC.@preserve old_u_ode, it might be GC'ed
    old_u = wrap_array(old_u_ode, mesh, equations, dg, cache)

    # Get new list of leaf cells
    leaf_cell_ids = local_leaf_cells(mesh.tree)

    # re-initialize elements container
    @unpack elements = cache
    resize!(elements, length(leaf_cell_ids))
    init_elements!(elements, leaf_cell_ids, mesh, dg.basis.nodes)
    @assert nelements(dg, cache) < old_n_elements

    resize!(u_ode, nvariables(equations) * nnodes(dg)^ndims(mesh) * nelements(dg, cache))
    u = wrap_array(u_ode, mesh, equations, dg, cache)

    # Loop over all elements in old container and either copy them or coarsen them
    u_tmp1 = Array{eltype(u), 4}(undef, nvariables(equations), nnodes(dg), nnodes(dg), nnodes(dg))
    u_tmp2 = Array{eltype(u), 4}(undef, nvariables(equations), nnodes(dg), nnodes(dg), nnodes(dg))
    skip = 0
    element_id = 1
    for old_element_id in 1:old_n_elements
      # If skip is non-zero, we just coarsened 2^ndims elements and need to omit the following elements
      if skip > 0
        skip -= 1
        continue
      end

      if to_be_removed[old_element_id]
        # If an element is to be removed, sanity check if the following elements
        # are also marked - otherwise there would be an error in the way the
        # cells/elements are sorted
        @assert all(to_be_removed[old_element_id:(old_element_id+2^ndims(mesh)-1)]) "bad cell/element order"

        # Coarsen elements and store solution directly in new data structure
        coarsen_elements!(u, element_id, old_u, old_element_id,
                          adaptor, equations, dg, u_tmp1, u_tmp2)
        element_id += 1
        skip = 2^ndims(mesh) - 1
      else
        # Copy old element data to new element container
        @views u[:, .., element_id] .= old_u[:, .., old_element_id]
        element_id += 1
      end
    end
    # If everything is correct, we should have processed all elements.
    @assert element_id == nelements(dg, cache) + 1 "element_id = $element_id, nelements(dg, cache) = $(nelements(dg, cache))"
  end # GC.@preserve old_u_ode

  # re-initialize interfaces container
  @unpack interfaces = cache
  resize!(interfaces, count_required_interfaces(mesh, leaf_cell_ids))
  init_interfaces!(interfaces, elements, mesh)

  # re-initialize boundaries container
  @unpack boundaries = cache
  resize!(boundaries, count_required_boundaries(mesh, leaf_cell_ids))
  init_boundaries!(boundaries, elements, mesh)

  # re-initialize mortars container
  @unpack mortars = cache
  resize!(mortars, count_required_mortars(mesh, leaf_cell_ids))
  init_mortars!(mortars, elements, mesh)

  # Sanity check
  if isperiodic(mesh.tree) && nmortars(mortars) == 0
    @assert ninterfaces(interfaces) == ndims(mesh) * nelements(dg, cache) ("For $(ndims(mesh))D and periodic domains and conforming elements, the number of interfaces must be $(ndims(mesh)) times the number of elements")
  end

  return nothing
end

# TODO: Taal compare performance of different implementations
# Coarsen solution data u for four elements, using L2 projection
function coarsen_elements!(u::AbstractArray{<:Any,5}, element_id,
                           old_u, old_element_id,
                           adaptor::LobattoLegendreAdaptorL2, equations, dg,
                           u_tmp1, u_tmp2)
  @unpack reverse_upper, reverse_lower = adaptor

  # Store old element ids
  bottom_lower_left_id  = old_element_id
  bottom_lower_right_id = old_element_id + 1
  bottom_upper_left_id  = old_element_id + 2
  bottom_upper_right_id = old_element_id + 3
  top_lower_left_id     = old_element_id + 4
  top_lower_right_id    = old_element_id + 5
  top_upper_left_id     = old_element_id + 6
  top_upper_right_id    = old_element_id + 7

  @boundscheck begin
    @assert old_element_id >= 1
    @assert size(old_u, 1) == nvariables(equations)
    @assert size(old_u, 2) == nnodes(dg)
    @assert size(old_u, 3) == nnodes(dg)
    @assert size(old_u, 4) == nnodes(dg)
    @assert size(old_u, 5) >= old_element_id + 7
    @assert     element_id >= 1
    @assert size(    u, 1) == nvariables(equations)
    @assert size(    u, 2) == nnodes(dg)
    @assert size(    u, 3) == nnodes(dg)
    @assert size(    u, 4) == nnodes(dg)
    @assert size(    u, 5) >= element_id
  end

  # Project from bottom lower left element
  multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_lower, reverse_lower,
    view(old_u, :, :, :, :, bottom_lower_left_id), u_tmp1, u_tmp2)

  # Project from bottom lower right element_variables
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_lower, reverse_lower,
    view(old_u, :, :, :, :, bottom_lower_right_id), u_tmp1, u_tmp2)

  # Project from bottom upper left element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_upper, reverse_lower,
    view(old_u, :, :, :, :, bottom_upper_left_id), u_tmp1, u_tmp2)

  # Project from bottom upper right element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_upper, reverse_lower,
    view(old_u, :, :, :, :, bottom_upper_right_id), u_tmp1, u_tmp2)

  # Project from top lower left element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_lower, reverse_upper,
    view(old_u, :, :, :, :, top_lower_left_id), u_tmp1, u_tmp2)

  # Project from top lower right element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_lower, reverse_upper,
    view(old_u, :, :, :, :, top_lower_right_id), u_tmp1, u_tmp2)

  # Project from top upper left element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_lower, reverse_upper, reverse_upper,
    view(old_u, :, :, :, :, top_upper_left_id), u_tmp1, u_tmp2)

  # Project from top upper right element
  add_multiply_dimensionwise!(
    view(u,     :, :, :, :, element_id), reverse_upper, reverse_upper, reverse_upper,
    view(old_u, :, :, :, :, top_upper_right_id), u_tmp1, u_tmp2)

  return nothing
end

# this method is called when an `ControllerThreeLevel` is constructed
function create_cache(::Type{ControllerThreeLevel}, mesh::TreeMesh{3}, equations, dg::DG, cache)

  controller_value = Vector{Int}(undef, nelements(dg, cache))
  return (; controller_value)
end


text:
function compute_int_K(params, α₁, n₁, α₂, n₂)
    @unpack r, D, dim = params

    m = length(D)
    R = r[m]

    dr = r[m] * 1e-12
    r₁ = [0; r[1:m-1]] .+ dr
    r₂ = r[1:m] .- dr

    K = 0
    if abs(n₁ - n₂) == 1
        λ₁ = α₁^2
        λ₂ = α₂^2
        ν₁ = compute_ν(n₁, dim)
        ν₂ = compute_ν(n₂, dim)

        # Computation at point r₁
        V₁, dV₁ = compute_v(params, α₁, n₁, r₁)
        V₂, dV₂ = compute_v(params, α₂, n₂, r₁)
        I₁a =
            (
                (λ₁ + λ₂) * D .* r₁ .^ (dim - 1) -
                (ν₁ + ν₂ - (dim - 1)) * D .^ 2 .* r₁ .^ (dim - 3)
            ) .* V₁ .* V₂
        I₁b =
            (
                (λ₂ - λ₁) * D .* r₁ .^ dim -
                (ν₁ - ν₂ - (dim - 1)) * D .^ 2 .* r₁ .^ (dim - 2)
            ) .* dV₁ .* V₂
        I₁c =
            (
                (λ₁ - λ₂) * D .* r₁ .^ dim -
                (ν₂ - ν₁ - (dim - 1)) * D .^ 2 .* r₁ .^ (dim - 2)
            ) .* V₁ .* dV₂
        I₁d = 2 * D .^ 2 .* r₁ .^ (dim - 1) .* dV₁ .* dV₂
        I₁ = I₁a + I₁b + I₁c + I₁d

        # Computation at point r₂
        V₁, dV₁ = compute_v(params, α₁, n₁, r₂)
        V₂, dV₂ = compute_v(params, α₂, n₂, r₂)
        I₂a =
            (
                (λ₁ + λ₂) * D .* r₂ .^ (dim - 1) -
                (ν₁ + ν₂ - (dim - 1)) * D .^ 2 .* r₂ .^ (dim - 3)
            ) .* V₁ .* V₂
        I₂b =
            (
                (λ₂ - λ₁) * D .* r₂ .^ dim -
                (ν₁ - ν₂ - (dim - 1)) * D .^ 2 .* r₂ .^ (dim - 2)
            ) .* dV₁ .* V₂
        I₂c =
            (
                (λ₁ - λ₂) * D .* r₂ .^ dim -
                (ν₂ - ν₁ - (dim - 1)) * D .^ 2 .* r₂ .^ (dim - 2)
            ) .* V₁ .* dV₂
        I₂d = 2 * D .^ 2 .* r₂ .^ (dim - 1) .* dV₁ .* dV₂
        I₂ = I₂a + I₂b + I₂c + I₂d
    end

    K = sum(I₂ - I₁) / (λ₁ - λ₂)^2 / R^(dim + 1)
end


text:
### --------------------------------------------------------- ###
### MINLP

function MINLP(dataset::Dataset;
               time_limit::Float64=600.0,
               symm_break_const::Bool=true,
               symm_w_const::Bool=true,
               w_lower::Float64=1e-12,
            )::Tuple{SBM, Matrix{Int}, OptResults}

    A = dataset.A
    K = dataset.n_communities
    n = dataset.n
    m = dataset.m
    k = dataset.k

    # Upper bound on w
    w_upper = getWUpperBound(dataset)
    @show(w_upper);

    # Set time limit
    open("couenne.opt", "w") do f
        write(f, "time_limit $time_limit\n");
        write(f, "allowable_gap 1e-6\n");
        write(f, "allowable_fraction_gap 1e-6\n");
        write(f, "acceptable_tol 1e-6\n");
        write(f, "tol 1e-6\n");
        write(f, "feas_tolerance 1e-7\n");
        write(f, "integer_tolerance 1e-7\n");
        write(f, "constr_viol_tol 1e-7\n");
        write(f, "compl_inf_tol 1e-7\n");
        write(f, "acceptable_constr_viol_tol 1e-7\n");
        write(f, "acceptable_compl_inf_tol 1e-7\n");
    end

    # --- Optimization model -------------------------------------------------
    model = Model(() -> AmplNLWriter.Optimizer("couenne"))

    # Variables
    @variable(model, w_lower <= w[1:K,1:K] <= w_upper);
    @variable(model, z[1:n,1:K], Bin);

    # Objective
    @NLobjective(model, Min, 0.5 * sum(
            (( -A[i,j]*log(w[r,s]) + ((k[i]*k[j])/(2*m))*w[r,s] ) * z[i,r] * z[j,s] )
            for r=1:K, s=1:K, i=1:n, j=1:n if A[i,j] != 0) +
            0.5 * sum(
            (( ((k[i]*k[j])/(2*m))*w[r,s] ) * z[i,r] * z[j,s] )
            for r=1:K, s=1:K, i=1:n, j=1:n if A[i,j] == 0));

    # Constraints
    @constraint(model, con[i = 1:n], sum(z[i,r] for r=1:K) == 1); # assignment constraint

    # Symmetry breaking in numbered clustering - Frank Plastria
    if symm_break_const == true
        # object 1 must be in cluster 1
        @constraint(model, z[1,1] <= 1);
        @constraint(model, z[1,1] >= 1);
        @constraint(model, symBreak[r=2:(K-1), j=r:n], sum(z[i,l] for l=1:(r-1), i=2:(j-1)) - sum(z[j,l] for l=1:r) <= j - 3 )
    end

    # Enforce symmetry on w
    if symm_w_const == true
        @constraint(model, symL[r=1:K, s=1:(r-1)], w[r,s] - w[s,r] <= 0.0)
        @constraint(model, symG[r=1:K, s=1:(r-1)], w[r,s] - w[s,r] >= 0.0)
    end
    # ------------------------------------------------------------------------

    # --- Solve optimization model -------------------------------------------
    start = time();
    println("Start solving")
    optimize!(model)
    status = string(termination_status(model))
    if status == "LOCALLY_SOLVED"
        status = "OPTIMAL"
    elseif status == "OTHER_LIMIT"
        status = "TIME_LIMIT"
    end
    solvetime = time() - start
    obj_lb = Float64(objective_bound(model));
    obj_ub = Float64(objective_value(model));
    obj_lb = -Inf
    if (status == "OPTIMAL")
        obj_lb = obj_ub
    end

    # ------------------------------------------------------------------------

    # --- Recover variable values --------------------------------------------
    w_opt = value.(w)
    z_opt = value.(z)
    z_opt = round.(z_opt[:,:])
    z_opt[map(v -> isnan(v) , z_opt)] .= 0 # replace NaN values with 0
    z_opt = Int.(z_opt)
    # ------------------------------------------------------------------------

    # SBM
    sbm = SBM(w_opt, "poisson") # assumes a poisson distribution

    # Build OptResults
    nodecount = nothing
    lazycount = nothing
    iterations = nothing
    println()
    println("Obj. LB = $obj_lb")
    println("Obj. UB = $obj_ub")
    println("Status: $status")
    println("Solve time: $solvetime")

    # Print results
    @show(w_opt)
    @show(z_opt)

    opt_results = OptResults(obj_lb, obj_ub, status, solvetime, iterations, nodecount, lazycount, nothing)
    displayResults(opt_results)
    return sbm, z_opt, opt_results
end

### --------------------------------------------------------- ###
### MILP

function MILP(dataset::Dataset;                          # dataset
              time_limit::Float64=600.0,                 # time limit (in seconds)
              symm_break_const::Bool=true,               # whether to add symmetry-breaking constraints in the model
              symm_w_const::Bool=true,                   # whether to enforce symmetry on w
              balanced_communities::Bool=false,          # whether to enforce balanced communities
              TOL::Float64=1e-6,                         # tolerance
              w_lower::Float64=1e-12,                    # lower bound on w
              n_breakpoints::Int64=50)                   # number of breakpoints
    """
    Solves the Mixed Integer Linear Programming formulation.

    Parameters
    ----------
    dataset : Dataset
        Dataset representing an observed graph.

    Returns
    -------
    sbm : SBM
        Stochastic Block Model
    z : Array{Int64,2}
        Matrix of assignments of nodes to groups
    opt_results : OptResults
        Results of the optimization process
    """

    # --- Retrieve instance data --------------------------------------------
    A = dataset.A
    K = dataset.n_communities
    n = dataset.n
    m = dataset.m
    k = dataset.k

    println("K = $K groups\nn = $n nodes\nm = $m edges");
    println("Breakpoints: $n_breakpoints")

    # Calculate constant term
    Cnst = getObjectiveConstant(dataset)

    # Upper bound on w
    rho = 0
    for i=1:n, j=1:n
        if A[i,j] != 0
            rho = maximum([rho, (A[i,j] / (k[i]*k[j]))])
        end
    end
    w_upper = 2 * m * rho
    @show(w_upper)
    # ------------------------------------------------------------------------

    # --- Initialize all break-points ----------------------------------------
    # Initialize wp
    wp = collect(range(w_lower, stop=w_upper, length=n_breakpoints))
    # ------------------------------------------------------------------------

    # --- Linear pieces coefficients------------------------------------------
    # Calculating slope and intercept of linear approximations
    # Linear pieces coefficients
    # Calculating slope and intercept of linear approximations
    a = Array{Array{Float64},2}(undef, n, n) # Slope of each line
    b = Array{Array{Float64},2}(undef, n, n) # Intercept of each line
    for i=1:n, j=1:n
        if A[i,j] != 0
            # slope: a = f'(xp) ( derivative evaluated at the point wp[p] )
            a[i,j] = -A[i,j]./wp .+ (k[i]*k[j])/(2*m)
            # intercept: b = f(xp) - f'(xp)*xp ( derive it from yp = m * xp + b -> b = yp - m * xp )
            b[i,j] = A[i,j] .* (1 .- log.(wp))
        else
            a[i,j] = [(k[i]*k[j])/(2*m)]
            b[i,j] = [0]
        end
    end
    # ------------------------------------------------------------------------

    # --- Defining bounds on variables ---------------------------------------
    M_lower = zeros(n,n,K,K);
    M_upper = zeros(n,n,K,K);
    for i=1:n, j=1:n, r=1:K, s=1:K
        if A[i,j] == 0
            M_lower[i,j,r,s] = 0 # Lower bound
            M_upper[i,j,r,s] = ((k[i]*k[j])/(2*m))*w_upper  # Upper bound
        else
            # Lower bound
            M_lower[i,j,r,s] = (A[i,j] * (1 - log(A[i,j]) + log((k[i]*k[j])/(2*m))))

            # Upper bound
            M_upper[i,j,r,s] = maximum([-A[i,j]*log(w_lower) + ((k[i]*k[j])/(2*m))*w_lower,
                                        -A[i,j]*log(w_upper) + ((k[i]*k[j])/(2*m))*w_upper])
        end
    end
    M_LOWER = minimum(M_lower)
    M_UPPER = maximum(M_upper)
    # ------------------------------------------------------------------------

    # --- Optimization model -------------------------------------------------
    solver = CPLEX.Optimizer
    model = Model(solver)
    set_optimizer_attribute(model, "CPXPARAM_MIP_Strategy_CallbackReducedLP", 0) # Access node information of the original model in the callback: CPX_OFF
    set_optimizer_attribute(model, "CPXPARAM_ScreenOutput", 1)              # Switching ON the display: CPX_ON
    set_optimizer_attribute(model, "CPXPARAM_Read_DataCheck", 1)            # Print warnings: CPX_DATACHECK_WARN
    set_optimizer_attribute(model, "CPX_PARAM_THREADS", 1)                  # Number of threads
    set_optimizer_attribute(model, "CPX_PARAM_TILIM", time_limit) # Time limit for the solver
    set_optimizer_attribute(model, "CPX_PARAM_EPGAP", TOL)                  # Relative MIP gap tolerance
    set_optimizer_attribute(model, "CPX_PARAM_PRELINEAR", 0)                # CPX_OFF
    set_optimizer_attribute(model, "CPX_PARAM_REDUCE", 0)                   # CPX_PREREDUCE_NOPRIMALORDUAL
    set_optimizer_attribute(model, "CPX_PARAM_EPINT", 1e-9)                 # Integrality tolerance
    set_optimizer_attribute(model, "CPX_PARAM_EPRHS", 1e-9)                 # Feasibility tolerance
    set_optimizer_attribute(model, "CPX_PARAM_EPOPT", 1e-9)                 # Optimality tolerance
    set_optimizer_attribute(model, "CPXPARAM_Preprocessing_Reduce", 1)      # Primal and dual preprocessing reduction type: Only primal reductions

    # Variables
    @variable(model, w_lower <= w[1:K,1:K] <= w_upper);
    @variable(model, z[1:n,1:K], Bin)
    @variable(model, 0 <= y[1:n,1:n,1:K,1:K] <= 1);
    @variable(model, M_LOWER <= x[1:n,1:n,1:K,1:K] <= M_UPPER);

    # Objective
    @objective(model, Min, - Cnst + 0.5 * sum( x[i,j,r,s] for r=1:K, s=1:K, i=1:n, j=1:n) );

    # Constraints
    # PWL Constraints
    @constraint(model, bigM1[i=1:n, j=1:n, r=1:K, s=1:K, p=1:length(a[i,j])], x[i,j,r,s] >= a[i,j][p]*w[r,s] + b[i,j][p] - M_upper[i,j,r,s]*(1 - y[i,j,r,s] )); # Big M-1 constraint

    # Big M constraints
    @constraint(model, bigM2[i=1:n, j=1:n, r=1:K, s=1:K], x[i,j,r,s] <= M_upper[i,j,r,s] * y[i,j,r,s] ); # Big M-2 constraint
    @constraint(model, bigM3[i=1:n, j=1:n, r=1:K, s=1:K], x[i,j,r,s] >= M_lower[i,j,r,s] * y[i,j,r,s] ); # Big M-3 constraint

    # Indicator variables constraints
    @constraint(model, bigM4[i=1:n, j=1:n, r=1:K, s=1:K], z[i,r] - y[i,j,r,s] >= 0 ); # Big M-4 constraint
    @constraint(model, bigM5[i=1:n, j=1:n, r=1:K, s=1:K], z[j,s] - y[i,j,r,s] >= 0 ); # Big M-5 constraint
    @constraint(model, bigM6[i=1:n, j=1:n, r=1:K, s=1:K], 1 - z[i,r] - z[j,s] + y[i,j,r,s] >= 0 ); # Big M-6 constraint

    # Assignment constraints
    @constraint(model, assign[i=1:n], sum(z[i,r] for r=1:K) == 1); # assignment constraint

    # Balanced communities
    if balanced_communities == true
        @constraint(model, bal[r=1:K], sum(z[i,r] for i=1:n) == n/K)
    end

    # Symmetry breaking constraints
    if symm_break_const == true
        @constraint(model, conObj1, z[1,1] == 1); # object 1 must be in cluster 1
        @constraint(model, symBreak[r=2:(K-1), j=r:n], sum(z[i,l] for l=1:(r-1), i=2:(j-1)) - sum(z[j,l] for l=1:r) <= j - 3 )
    end

    # Enforce symmetry on w
    if symm_w_const == true
        @constraint(model, symL[r=1:K, s=1:(r-1)], w[r,s] - w[s,r] <= 0.0)
        @constraint(model, symG[r=1:K, s=1:(r-1)], w[r,s] - w[s,r] >= 0.0)
    end

    # Get number of constraints of type GenericAffExpr
    num_cons = sum([num_constraints(model, func, set)
                    for (func, set) in list_of_constraint_types(model)
                    if func == GenericAffExpr{Float64, VariableRef}])
    println("Number of constraints before start: $num_cons")

    # Get the number of constraints of type VariableRef
    num_cons_vars = sum([num_constraints(model, func, set)
                         for (func, set) in list_of_constraint_types(model)
                         if func == VariableRef])
    println("Number of constraints on variables before start: $num_cons_vars")
    # ------------------------------------------------------------------------

    # --- Lazy constraint callback -------------------------------------------
    # Lazy constraints counter
    lazycount = 0

    # Defining lazy constraint callback
    function lazy_callback_function(cb_data)
        x_val = zeros(size(x))
        y_val = zeros(size(y))
        w_val = zeros(size(w))
        z_val = zeros(size(z))
        for i=1:n, j=1:n, r=1:K, s=1:K
            x_val[i,j,r,s] = callback_value(cb_data, x[i,j,r,s])
        end
        for i=1:n, j=1:n, r=1:K, s=1:K
            y_val[i,j,r,s] = callback_value(cb_data, y[i,j,r,s])
        end
        for r=1:K, s=1:K
            w_val[r,s] = callback_value(cb_data, w[r,s])
            if (w_val[r,s] <= 0)
                w_val[r,s] = 1e-50; # TODO: use a tolerance value here
            end
        end
        for i=1:n, r=1:K
            z_val[i,r] = callback_value(cb_data, z[i,r])
        end

        z_val = round.(z_val[:,:])
        z_val[map(v -> isnan(v) , z_val)] .= 0 # replace NaN values with 0
        z_val = Int.(z_val)

        # Calculate objective function and approximation difference
        pwl_obj = - Cnst + 0.5 * sum( x_val[i,j,r,s] for r=1:K, s=1:K, i=1:n, j=1:n)
        nl_obj = calculateObjective(dataset, w_val, z_val)

        # Check that the linearized objective value is less than or equal to the real objective value
        @assert nl_obj >= pwl_obj "Something weird inside the lazy constraint. PWL objective should be < NL obj. Got PWL obj: $pwl_obj and NL obj: $nl_obj"
        obj_relative_diff = (nl_obj - pwl_obj) / pwl_obj
        println("Lazy count $lazycount")
        println("  PWL obj: $(@sprintf("%.10f", pwl_obj))")
        println("  NL obj: $(@sprintf("%.10f", nl_obj))")
        println("  obj relative diff: $(@sprintf("%.10e", obj_relative_diff))")
        println("  z_val : $z_val")
        println("  w_val : $w_val")

        # Allow for some impreciseness in the solution
        if abs(obj_relative_diff) >= TOL
            lazycount += 1;
            # Calculate new approximation
            w_cut = w_val[0 .< w_val .< Inf] # Only finite values
            size_w = length(w_cut)

            for i=1:n, j=1:n, p=1:size_w
                if (A[i,j] != 0)
                    a = -A[i,j]/w_cut[p] + ((k[i]*k[j])/(2*m)) # Slope
                    a = max(-1e8, min(1e8, a)) # Clip values for the slope
                    b = A[i,j]*(1 - log(w_cut[p])) # Intercept of each line

                    # Submit constraints
                    for r=1:K, s=1:K
                        rhs = (M_upper[i,j,r,s] - b)
                        lhs = -x_val[i,j,r,s] + M_upper[i,j,r,s]*y_val[i,j,r,s] + a*w_val[r,s]

                        if (!( lhs <= rhs ))
                            lazycon = @build_constraint(a*w[r,s] + b - M_upper[i,j,r,s]*(1 - y[i,j,r,s]) <= x[i,j,r,s] ); # Lazy constraint
                            MOI.submit(model, MOI.LazyConstraint(cb_data), lazycon)
                        end
                    end
                end
            end
        end
    end # End of callback function

    # Tell JuMP/Gurobi/CPLEX to use our callback function
    MOI.set(model, MOI.LazyConstraintCallback(), lazy_callback_function)
    # ------------------------------------------------------------------------

    # --- Solve optimization model -------------------------------------------
    optimize!(model) # Solve
    status      = string(termination_status(model))
    solvetime   = solve_time(model)
    obj_lb      = objective_bound(model);
    has_vals    = has_values(model)

    # Get objective value if there is a solution
    if has_vals
        obj_ub = objective_value(model);
    else
        obj_ub = nothing
    end

    # Get node count
    nodecount = Int64(MOI.get(model, MOI.NodeCount()))
    # ------------------------------------------------------------------------

    # --- Recover variable values --------------------------------------------
    if has_vals
        w_opt = value.(w);
        x_opt = value.(x);
        z_opt = value.(z);
    else
        throw("No solution found. Revise the model output")
    end

    # Convert z to int
    z_opt = round.(z_opt[:,:])
    z_opt[map(v -> isnan(v) , z_opt)] .= 0 # replace NaN values with 0
    z_opt = Int.(z_opt)

    # Check the tolerance in the objective value relative difference
    nonlinear_obj = calculateObjective(dataset, w_opt, z_opt)
    obj_relative_diff = (nonlinear_obj - obj_ub) / obj_ub
    if status == "OPTIMAL"
        @assert (obj_relative_diff <= TOL) "Tolerance of the objective value violated. Obj. value: $(@sprintf("%.10f", obj_ub)); NL obj: $(@sprintf("%.10f", nonlinear_obj)); Obj. relative diff: $(@sprintf("%.10e", obj_relative_diff))"
    end
    # ------------------------------------------------------------------------
    # SBM
    sbm = SBM(w_opt, "poisson") # assumes a poisson distribution

    # Print results
    println()
    println("w_opt =")
    display(w_opt)
    println()
    println("z_opt =")
    display(z_opt)
    println()
    println("Solve time:                $solvetime")
    println("Status:                    $status")
    println("Objective value:           $obj_ub")
    println("Objective lower bound:     $obj_lb")
    println("Non-linear objective:      $(@sprintf("%.10f", nonlinear_obj))")
    println("Obj relative diff:         $(@sprintf("%.10e", obj_relative_diff))")
    println("Nodes processed:           $nodecount")
    println("Lazy constraints:          $lazycount")

    # Build OptResults
    iterations = nothing
    opt_results = OptResults(obj_lb, obj_ub, status, solvetime, iterations, nodecount, lazycount, num_cons)
    displayResults(opt_results)
    return sbm, z_opt, opt_results
end

max_stars_repo_name: vidalt/Exact-SBM

max_issues_repo_name: vidalt/Exact-SBM

max_forks_repo_name: vidalt/Exact-SBM


text:
(+)(td::TimeDate) = td
(+)(tdz::TimeDateZone) = tdz

function (-)(td1::TimeDate, td2::TimeDate)
    dtm1 = DateTime(td1)
    dtm2 = DateTime(td2)
    dtm0 = canonical(Microsecond(td1) - Microsecond(td2) + Nanosecond(td1) - Nanosecond(td2))
    dtm = dtm1 - dtm2
    dtm += dtm0
    return sum(fldmod(dtm))
end

(-)(td::TimeDate, dt::DateTime) = td - TimeDate(dt)
(-)(dt::DateTime, td::TimeDate) = TimeDate(dt) - td
(-)(td::TimeDate, dt::Date) = td - TimeDate(dt)
(-)(dt::Date, td::TimeDate) = TimeDate(dt) - td

function (-)(td::TimeDate, tm::Time)
    tm_td, dt_td = at_time(td), on_date(td)
    if tm_td < tm
        dt_td -= Day(1)
        tm_td = tm_td - tm + Hour(24)
    else
        tm_td = tm_td - tm
    end
    tm_td = canonical(tm_td)
    tm = Time(tm_td)
    return TimeDate(tm, dt_td)
end

for P in (:Year, :Month, :Day)
    @eval begin
        function (+)(td::TimeDate, p::$P)
            tm, dt = at_time(td), on_date(td)
            dt = dt + p
            return TimeDate(tm, dt)
        end
        function (-)(td::TimeDate, p::$P)
            p = -p
            return td + p
        end
    end
end

function (+)(td::TimeDate, p::Union{Hour, Minute, Second, Millisecond, Microsecond, Nanosecond})
    tm, dt = at_time(td), on_date(td)
    extradays,extratime=divrem(p,oftype(p,oneunit(Day)))
    newtm = sum(promote(extratime,tm.instant))
    if newtm >= Nanosecond(oneunit(Day))
        newtm -= Nanosecond(oneunit(Day))
        extradays += 1
    elseif newtm < zero(Nanosecond)
        newtm += Nanosecond(oneunit(Day))
        extradays -= 1
    end
    return TimeDate(Time(newtm),dt+Day(extradays))
end

function (-)(td::TimeDate, p::Union{Hour, Minute, Second, Millisecond, Microsecond, Nanosecond})
    td + (-p)
end

@inline nonempty(x) = !isempty(x) ? x : Nanosecond(0)

function (-)(x::TimeDateZone, y::TimeDateZone)
    xx = Microsecond(x) + Nanosecond(x)
    xx = nonempty(xx)
    yy = Microsecond(y) + Nanosecond(y)
    yy = nonempty(yy)
    xy = xx - yy
    xy = nonempty(xy)
    zx = ZonedDateTime(x)
    zy = ZonedDateTime(y)
    zxy = zx - zy
    zxy = nonempty(zxy)
    result = zxy + xy
    result = nonempty(result)

    return canonical(result)
end

(-)(x::TimeDateZone, y::ZonedDateTime) = x - TimeDateZone(y)
(-)(x::ZonedDateTime, y::TimeDateZone) = TimeDateZone(y) - y

function(+)(tdz::TimeDateZone, p1::Period, p2::Period)
    p = p1 + p2
    return tdz + p
end

function(-)(tdz::TimeDateZone, p1::Period, p2::Period)
    p = p1 + p2
    return tdz - p
end

function(+)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period)
    p = p1 + p2 + p3
    return tdz + p
end

function(-)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period)
    p = p1 + p2 + p3
    return tdz - p
end

function(+)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period, p4::Period)
    p = p1 + p2 + p3 + p4
    return tdz + p
end

function(-)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period, p4::Period)
    p = p1 + p2 + p3 + p4
    return tdz - p
end

function(+)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period, p4::Period, p5::Period)
    p = p1 + p2 + p3 + p4 + p5
    return tdz + p
end

function(-)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period, p4::Period, p5::Period)
    p = p1 + p2 + p3 + p4 + p5
    return tdz - p
end

function(+)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period, p4::Period, p5::Period, p6::Period)
    p = p1 + p2 + p3 + p4 + p5 + p6
    return tdz + p
end

function(-)(tdz::TimeDateZone, p1::Period, p2::Period, p3::Period, p4::Period, p5::Period, p6::Period)
    p = p1 + p2 + p3 + p4 + p5 + p6
    return tdz - p
end


text:
### implements the coveriance matrices as described in Patchwork Kriging for Large-scale Gaussian Process Regression
# https://jmlr.org/papers/v19/17-042.html

function setuppatchGP!(X_parts::Vector{Vector{Vector{T}}},
    θ::KT) where {T,KT}

    N_parts = length(X_parts)

    ### K_XX.
    K_XX_set = Vector{Matrix{T}}(undef, N_parts)

    for n = 1:N_parts

        K_XX_set[n] = constructkernelmatrix(X_parts[n], θ)
    end

    boundary_labels, bb_positive_list, bb_negative_list,
    Xb_positive_list, Xb_negative_list = setupboundaryquantities(X_parts)

    ### K_bb.

    ### K_Xb.

    ### think about fast serach for queryX and queryb.

    return boundary_labels
end

function setupboundaryquantities(X_parts::Vector{Vector{Vector{T}}}) where T

    N_parts = length(X_parts)

    # get boundary labels.
    M = round(Int, N_parts*(N_parts-1)/2)
    boundary_labels = Vector{Tuple{Int,Int}}(undef, M)
    a = 0

    for j = 2:N_parts
        for i = 1:j-1
            a += 1
            boundary_labels[a] = (i,j)
        end
    end
    N_boundaries = length(boundary_labels)

    ### mark part of the non-zero entries of K_bb.
    bb_positive_list = Vector{Tuple{Int,Int}}(undef,0)
    bb_negative_list = Vector{Tuple{Int,Int}}(undef,0)

    #for j = 2:M
        #for i = 1:j-1
    for m1 = 1:N_boundaries
        for m2 = 1:N_boundaries

            k, l = boundary_labels[m1]
            u, v = boundary_labels[m2]

            if (k == u && l != v) || (k != u && l == v)

                push!(bb_positive_list, (m1,m2))

            elseif (k==v && l != u) || (k!=v && l == u)

                push!(bb_negative_list, (m1,m2))
            end

        end
    end

    ### get non-zero entries of K_Xb.

    Xb_positive_list = Vector{Tuple{Int,Int}}(undef,0)
    Xb_negative_list = Vector{Tuple{Int,Int}}(undef,0)

    # number of training data inputs, i.e., excluding patchwork/boundary inputs.
    N_X = sum( length(X_parts[n]) for n = 1:N_parts )

    for m = 1:N_boundaries
        for j = 1:N_parts

            #part_ind = div(i, N_parts) +1
            k, l = boundary_labels[m]

            if k == j
                push!(Xb_positive_list, (m,j))

            elseif l == j
                push!(Xb_positive_list, (m,j))
            end
        end
    end

    return boundary_labels, bb_positive_list, bb_negative_list,
        Xb_positive_list, Xb_negative_list
end
# ####### patchworkGP: kernel matrix contruction.
# boundary_labels, bb_positive_list, bb_negative_list,
# Xb_positive_list, Xb_negative_list = PatchMixtureKriging.setupboundaryquantities(X_parts)

# pause development since the GP inducing-point styled approx. inference could be very slow.
# instead, work on the mixture of GP idea.

max_stars_repo_name: RoyCCWang/PatchMixtureKriging

max_issues_repo_name: RoyCCWang/PatchMixtureKriging


text:
outer_(x) = x*x'
outer_(x::Number) = x*x'*I
outer_(x::AbstractVector) = Diagonal(x.*x)

construct_a(σ::Number,P) = outer_(σ)*similar_type(P, Size(size(P,1),size(P,1)))(I)
construct_a(σ::Number,P::Number) = σ*σ'
construct_a(σ,P) = Matrix(outer_(σ))

get_dt(ts::AbstractRange) = step(ts)
get_dt(ts::AbstractVector) = ts[2] - ts[1]
get_tspan(ts) = (first(ts),last(ts))

timechange(s, tstart=first(s), tend=last(s)) = tstart .+ (s .- tstart).*(2.0 .- (s .- tstart)/(tend-tstart))

isinplace(R::AbstractRegression{inplace}) where {inplace} = inplace

# linear approximation
(a::AffineMap)(u,p,t) = a.B*u .+ a.β
(a::ConstantMap)(u,p,t) = a.x

function convert_message(message, F1::InformationFilter, F2::CovarianceFilter)
  @unpack ktilde, ts, soldis, sol = message
  # P = inv(H)

  # ν = inv(H)F

  return Message(ktilde, ts, soldis, nothing)
end

function convert_message(message, F1::CovarianceFilter, F2::InformationFilter)
  @unpack ktilde, ts, soldis, sol = message
  # H = inv(P)

  # F = Hν

  return Message(ktilde, ts, soldis, nothing)
end

function construct_discrete_sol(sol::SciMLBase.ODESolution)
  soldis = reverse(Array(sol), dims=2)
  ν, P, c = myunpack(sol.u[1])
  d = length(ν)
  #TODO: only store views, make an option for Cholesky(P), qr(P), etc.
  return [( view(cols,1:d), reshape(view(cols, d+1:d+d*d),d,d), cols[end]) for cols in eachcol(soldis)]
end

function construct_discrete_sol(sol::AbstractVector)
  ν, P, c = myunpack(sol[1])
  d = length(ν)
  #TODO: only store views, make an option for Cholesky(P), qr(P), etc.
  return [( view(cols,1:length(ν)), reshape(view(cols, d+1:d+d*d),d,d), cols[end]) for cols in sol]
end

# handle noise conversion between solvers
function compute_Z(::Nothing, ::Nothing, trange, u0)
    n = length(trange)

    return collect(zip(1:n, trange, cumsum([[zero(u0)];[sqrt(trange[i+1]-ti)*randn(size(u0))
            for (i,ti) in enumerate(trange[1:end-1])]])))
end

function compute_Z(::Nothing, noise_rate_prototype, trange, u0)
    n = length(trange)

    return collect(zip(1:n, trange, cumsum([[zeros(size(noise_rate_prototype,2))];[sqrt(trange[i+1]-ti)*randn(size(noise_rate_prototype,2))
            for (i,ti) in enumerate(trange[1:end-1])]])))
end

function compute_Z(Z::DiffEqNoiseProcess.NoiseGrid, noise_rate_prototype, trange, u0)
    n = length(trange)

    return collect(zip(1:n, trange, Z.W))
end

function compute_Z(Z::DiffEqNoiseProcess.NoiseWrapper, noise_rate_prototype, trange, u0)
    n = length(trange)

    return collect(zip(1:n, trange, Z.source.W))
end

function compute_Z(Z, noise_rate_prototype, trange, u0)
    n = length(trange)

    return collect(zip(1:n, trange, Z))
end

function AugmentedNoiseGrid(ts,Ws,Zs=nothing)
  Waug = [vcat(W,zero(eltype(W))) for W in Ws]
  if Zs!==nothing
    Zaug = [vcat(Z,zero(eltype(Z))) for Z in Zs]
  else
    Zaug = nothing
  end
  NGaug = NoiseGrid(ts,Waug,Zaug)
end

"""
  innov(t)
t:  either an array of Numbers or StepRangeLen{Number}
returns Brownian values
"""
function innov(t, Z)
    dt = diff(t)
    w = [sqrt(dt[i])*convert(typeof(Z[i]),randn(size(Z[i]))) for i in 1:length(t)-1]
    brownian_values = cumsum(pushfirst!(w, zero(Z[1])))
end

"""
  pCN(Z, ρ)
update noise values Z by pCN-step with parameter ρ. Taking ρ=1 just returns a
NoiseGrid with the noise values Z
"""
function pCN(t,Z,ρ)
    Znew = innov(t, Z)
    a = cumsum(pushfirst!(ρ * diff(Z) + sqrt(1. - ρ^2) * diff(Znew), zero(Z[1])))
    NoiseGrid(t, a)
end


text:
#*************************************************************************#
#
# FeedforwardNetwork Type
#
#*************************************************************************#

# Constructor I
"""
    FeedforwardNetwork(ninputs, activfunc)

Initialize `ParametricModel` representing feedforward neural network.
`ninputs` is the network number of inputs and `activfunc`
is a vector of `AparametricModel` containg activation
functions for each one of the layers.
"""
function FeedforwardNetwork{T<:AparametricModel}(ninputs::Int,
                                                 activfunc::Vector{T})
    nlayers = length(activfunc)
    mdls = Vector{StaticModel}(3*nlayers)
    node_noutputs = ninputs
    k = 1
    for i = 1:nlayers
        node_ninputs = node_noutputs
        node_noutputs = activfunc[i].ninputs
        mdls[k] = Linear(node_ninputs, node_noutputs)
        k += 1
        mdls[k] = Bias(node_noutputs)
        k += 1
        mdls[k] = activfunc[i]
        k += 1
    end

    SeriesModel(mdls)
end

# Constructor II
"""
    FeedforwardNetwork(ninputs, noutputs, nhidden)

Initialize `ParametricModel` representing feedforward neural network.
`ninputs` and `outputs` are the network number of inputs and outputs
and `nhidden` is a vector containg the number of nodes in each hidden
layer.
"""
function FeedforwardNetwork(ninputs::Int,
                            noutputs::Int,
                            nhidden::Vector{Int})
    nlayers = length(nhidden)+1

    activfunc = Vector{ElementwiseOperator}(nlayers)
    for k = 1:nlayers
        activfunc[k] = (k==nlayers) ? Identity(noutputs) : HyperbolicTangent(nhidden[k])
    end

    FeedforwardNetwork(ninputs, activfunc)
end


text:
export Variable, @variables, @parameters

const IndexMap = Dict{Char,Char}(
            '0' => '₀',
            '1' => '₁',
            '2' => '₂',
            '3' => '₃',
            '4' => '₄',
            '5' => '₅',
            '6' => '₆',
            '7' => '₇',
            '8' => '₈',
            '9' => '₉')
function map_subscripts(indices)
    str = string(indices)
    join(IndexMap[c] for c in str)
end

"""
$(TYPEDEF)

A named variable which represents a numerical value. The variable's value may
be known (parameters, independent variables) or unknown (dependent variables).

# Fields
$(FIELDS)
"""
struct Variable <: Function
    """The variable's unique name."""
    name::Symbol
    """
    Whether the variable's value is known.
    """
    known::Bool
    Variable(name; known = false) = new(name, known)
end
function Variable(name, indices...; known = false)
    var_name = Symbol("$(name)$(join(map_subscripts.(indices), "̒"))")
    Variable(var_name; known=known)
end

(x::Variable)(args...) = Operation(x, collect(Expression, args))

Base.isequal(x::Variable, y::Variable) = (x.name, x.known) == (y.name, y.known)
Base.print(io::IO, x::Variable) = show(io, x)
Base.show(io::IO, x::Variable) = print(io, x.name)
function Base.show(io::IO, ::MIME"text/plain", x::Variable)
    known = x.known ? "known" : "unknown"
    print(io, x.name, " (callable ", known, " variable)")
end

"""
$(TYPEDEF)

An expression which wraps a constant numerical value.

The value of the constant can be extracted with [`Base.get`](@ref).
"""
struct Constant <: Expression
    """The constant's numerical value"""
    value::Number
end

"""
$(TYPEDSIGNATURES)

Get the value of a [`ModelingToolkit.Constant`](@ref).
"""
Base.get(c::Constant) = c.value

Base.iszero(ex::Expression) = isa(ex, Constant) && iszero(ex.value)
Base.isone(ex::Expression)  = isa(ex, Constant) && isone(ex.value)

# Variables use isequal for equality since == is an Operation
Base.isequal(c::Constant, n::Number) = c.value == n
Base.isequal(n::Number, c::Constant) = c.value == n
Base.isequal(a::Constant, b::Constant) = a.value == b.value

Base.convert(::Type{Expr}, c::Constant) = c.value

# Build variables more easily
function _parse_vars(macroname, known, x)
    ex = Expr(:block)
    var_names = Symbol[]
    # if parsing things in the form of
    # begin
    #     x
    #     y
    #     z
    # end
    x = x isa Tuple && first(x) isa Expr && first(x).head == :tuple ? first(x).args : x # tuple handling
    x = flatten_expr!(x)
    for _var in x
        iscall = isa(_var, Expr) && _var.head == :call
        isarray = isa(_var, Expr) && _var.head == :ref
        issym  = _var isa Symbol
        @assert iscall || isarray || issym "@$macroname expects a tuple of expressions or an expression of a tuple (`@$macroname x y z(t) v[1:3] w[1:2,1:4]` or `@$macroname x, y, z(t) v[1:3] w[1:2,1:4]`)"

        if iscall
            var_name, expr = _construct_vars(_var.args[1], known, _var.args[2:end])
        else
            var_name, expr = _construct_vars(_var, known, nothing)
        end
        push!(var_names, var_name)
        push!(ex.args, expr)
    end
    push!(ex.args, build_expr(:tuple, var_names))
    return ex
end

function _construct_vars(_var, known, call_args)
    issym  = _var isa Symbol
    isarray = isa(_var, Expr) && _var.head == :ref
    if isarray
        var_name = _var.args[1]
        indices = _var.args[2:end]
        expr = _construct_array_vars(var_name, known, call_args, indices...)
    else
        # Implicit 0-args call
        var_name = _var
        expr = _construct_var(var_name, known, call_args)
    end
    var_name, :($var_name = $expr)
end

function _construct_var(var_name, known, call_args)
    if call_args === nothing
        :(Variable($(Meta.quot(var_name)); known = $known)())
    elseif !isempty(call_args) && call_args[end] == :..
        :(Variable($(Meta.quot(var_name)); known = $known))
    else
        :(Variable($(Meta.quot(var_name)); known = $known)($(call_args...)))
    end
end

function _construct_var(var_name, known, call_args, ind)
    if call_args === nothing
        :(Variable($(Meta.quot(var_name)), $ind...; known = $known)())
    elseif !isempty(call_args) && call_args[end] == :..
        :(Variable($(Meta.quot(var_name)), $ind...; known = $known))
    else
        :(Variable($(Meta.quot(var_name)), $ind...; known = $known)($(call_args...)))
    end
end

function _construct_array_vars(var_name, known, call_args, indices...)
    :(map(Iterators.product($(indices...))) do ind
        $(_construct_var(var_name, known, call_args, :ind))
    end)
end

"""
$(SIGNATURES)

Define one or more unknown variables.
"""
macro variables(xs...)
    esc(_parse_vars(:variables, false, xs))
end

"""
$(SIGNATURES)

Define one or more known variables.
"""
macro parameters(xs...)
    esc(_parse_vars(:parameters, true, xs))
end


text:
N = 10
Ngroups = N
groupsInds = distributeAinVecN(Vector(1:Ngroups),N)
tmpDic = countmap(groupsInds)
#number of nodes per group
NgroupsMemb = sortrows([[i for i in keys(tmpDic)] [  tmpDic[i] for i in keys(tmpDic) ] ])

deltaN = 3
groupsPar =  ones(Ngroups)
degGroups =  DegSeq2graphDegSeq(fooErgmBin1, Array{Real,1}(linspace(1+deltaN,N - round(0.1*N)-deltaN,N)))
#tmpdegs = [8, 4, 7, 8, 5, 5, 6, 5, 6, 0]
#sort([7, 5, 8, 9, 8, 5, 4, 2, 8, 4],rev = true)# DegSeq2graphDegSeq(fooErgmBin1, sort([9, 8, 4, 7, 3, 6, 4, 7, 4, 2],rev = true))
Mod = ErgmBin1(degGroups,groupsPar,groupsInds)
groupsPar, nIterm,estMod = estimate(Mod)
firstOrderCond(Mod;deg = tmpdegs,par = groupsPar, groupsInds = groupsInds )
expMatrix(Mod,groupsPar[groupsInds])

## test sampling
groupsPar[10] -= 10
Nsample = 1000
sam = sampl(estMod,Nsample)
meanMatSam = squeeze(mean(sam,3),3)
degsSam = squeeze(sum(meanMatSam,2),2)
degGroupsSam = zeros(Ngroups) ; for i =1:N degGroupsSam[groupsInds[i]] += degsSam[i] end
relErrSam = maximum(abs.(( degGroupsSam - degGroups )./degGroups ))
## Test estimate umbiasedness
estParSam = zeros(N,Nsample)
@time for n=1:Nsample
    degs_n = expValStatsFromMat(fooErgmBin1,sam[:,:,n] )
    estParSam[:,n], ~,~ = estimate(Mod; deg= degs_n)
end

##
using Plots, AverageShiftedHistograms, Distributions, StatPlots
plotly()
listPlots=""
    for n=1:N
        err = ((estParSam .- groupsPar))[n,:]
        rng_space = -3:.1:3
        n==10 ? (rng_space = -300:.1:300) : ()
        tmpAsh = ash(err;rng = rng_space)
        pA1 = plot(tmpAsh)
        expr = "PW$(n) = plot!(Normal(0,1),linewidth = 5) "
        eval(parse(expr))
        listPlots *= "PW$(n) , "
    end

exprLast = "plot(" * listPlots  *" titlefont = font(12) ,titleloc=:center,size=(1350,600),legend = :none )"
 eval(parse(exprLast))

##

#


text:
@testset "Select" begin

    @testset "Select parameters" begin
        import NaiveNASflux: select
        mat = reshape(collect(1:3*4), 3, 4)

        @test select(mat, 2 => [2,-1,-1,4]) == [4 0 0 10;5 0 0 11;6 0 0 12]
        @test select(mat, 1 => [-1,1,3]) == [0 0 0 0;1 4 7 10;3 6 9 12]

        @test select(mat, 1 => [2,-1,3,-1], 2 => [-1,1,-1,4]) == [0 2 0 11;0 0 0 0;0 3 0 12;0 0 0 0]
    end

    @testset "KernelSizeAligned" begin
        import NaiveNASflux: selectfilters
        pars = reshape(1:2*3*4*5, 5,4,3,2)
        ps = selectfilters(KernelSizeAligned(2, -1), pars)

        @test ps[1] == (1 => [-1, 1, 2, 3, 4, 5, -1])
        @test ps[2] == (2 => [2,3,4])
    end

    @testset "KernelSizeAligned MutableLayer $convtype" for convtype in (Conv, ConvTranspose, DepthwiseConv)
        import NaiveNASflux: weights
        m = mutable(convtype((3,4), 5=>5, pad=(1,1,1,2)), inputvertex("in", 5), layerfun=identity)
        indata = ones(Float32, 5,5,5,2)

        @test size(weights(layer(m)))[1:2] == (3,4)
        @test size(m(indata)) == size(indata)

        mutate_weights(m, KernelSizeAligned((-1, 2), (1,0,3,2)))

        @test size(weights(layer(m)))[1:2] == (2,6)
        @test size(m(indata)) == size(indata)
    end

    @testset "KernelSizeAligned LazyMutable" begin
        import NaiveNASflux: weights
        m = mutable(Conv((3,4), 5=>6, pad=(1,1,1,2)), inputvertex("in", 5), layerfun=LazyMutable)
        indata = ones(Float32, 5,5,5,2)

        @test size(weights(layer(m))) == (3,4,5,6)
        @test size(m(indata))[1:2] == size(indata)[1:2]

        mutate_weights(m, KernelSizeAligned((-1, 2), (1,0,3,2)))

        @test size(m(indata)) == (size(indata,1), size(indata,2), 6, size(indata,4))
        @test size(weights(layer(m))) == (2,6,5,6)

        mutate_weights(m, KernelSizeAligned((2, -1), (1,2,2,2)))
        mutate_outputs(m, [1,3,5])

        @test size(m(indata)) == (size(indata,1), size(indata,2), 3, size(indata,4))
        @test size(weights(layer(m))) == (4,5,5,3)
    end

    @testset "KernelSizeAligned Dense is Noop with layerfun $lfun" for lfun in (identity, LazyMutable)
        m = mutable(Dense(3,4), inputvertex("in", 3), layerfun = lfun)
        indata = ones(Float32, nin(m)[], 2)

        @test size(m(indata)) == (4,2)
        mutate_weights(m, KernelSizeAligned((-1,-1), (1,1)))
        @test size(m(indata)) == (4,2)
    end
end


text:
## Memory allocation strategies

abstract type AbstractMemoryPool end

mutable struct SimplePool <: AbstractMemoryPool
end

alloc(mp::SimplePool, T, sz) = T(undef, sz...)
free(mp::SimplePool, a) = ()

mutable struct CachingPool
    cache::Dict{Any, Any}    # (T, sz) => [buffer1, buffer2, ...]
end
CachingPool() = CachingPool(Dict())

function alloc(mp::CachingPool, T, sz)
    key = (T, sz)
    if haskey(mp.cache, key)
        buffers = mp.cache[key]
        arr = pop!(buffers)
        isempty(buffers) &&  delete!(mp.cache, key)
        return arr
    else
        try
            return T(undef, sz)
        catch e
            if e isa OutOfMemoryError
                # worst case - release all cached memory, run GC and repeat attempt
                mp.cache = Dict()
                GC.gc()
                return T(undef, sz)
            end
        end
    end
end

function free(mp::CachingPool, arr)
    key = (typeof(arr), size(arr))
    if !haskey(mp.cache, key)
        mp.cache[key] = Any[]
    end
    push!(mp.cache[key], arr)
end

using LinearAlgebra

function usage()
    a = rand(5, 4) |> cu
    b = rand(4, 10) |> cu

    mp = CachingPool()
    # alloc x
    x = alloc(mp, CuArray{Float32, 2}, (5, 10))
    mul!(x, a, b)
    # free x when it's not needed anymore
    free(mp, x)
    x = nothing
end