text:
@testset "MIDI pitches" begin
    @testset "types ok and compact" begin
        @test isbitstype(MidiInterval)
        @test isbitstype(MidiIC)
        @test isbitstype(Pitch{MidiInterval})
        @test isbitstype(Pitch{MidiIC})
    end

    @testset "constructors" begin
        @test midi(3) == midi(3)
        @test midip(60) == midip(60)
        @test midic(15) == midic(3)
        @test midipc(63) == midipc(3)
    end

    @testset "macro constructors" begin
        @test (@midi [1,2,3]) == midi.([1,2,3])
        @test (@midip [1,2,3]) == midip.([1,2,3])
        @test (@midic [1,2,3]) == midic.([1,2,3])
        @test (@midipc [1,2,3]) == midipc.([1,2,3])
    end

    @testset "String conversion" begin
        @test string(midi(3)) == "i3"
        @test string(midip(63)) == "p63"
        @test string(midic(15)) == "ic3"
        @test string(midipc(63)) == "pc3"
    end

    @testset "Int conversion" begin
        @test Int(midi(3)) == 3
        @test Int(midip(63)) == 63
        @test Int(midic(15)) == 3
        @test Int(midipc(63)) == 3
    end

    @testset "hashing" begin
        @test hash(midi(3)) == hash(midi(3))
        @test hash(midip(3)) == hash(midip(3))
        @test hash(midic(3)) == hash(midic(3))
        @test hash(midipc(3)) == hash(midipc(3))
    end

    @testset "interval interface" begin
        @test midi(3) + midi(10) == midi(13)
        @test midi(-3) + midi(4) == midi(1)
        @test midi(3) + midi(-4) == midi(-1)
        @test midi(3) - midi(4) == midi(-1)
        @test -midi(5) == midi(-5)
        @test -midi(7) == midi(0) - midi(7)

        @test zero(MidiInterval) == midi(0)
        @test zero(midi(3)) == midi(0)

        @test midi(2) * 7 == midi(14)
        @test midi(-3) * 4 == midi(-12)
        @test midi(4) * -3 == midi(-12)
        @test 6 * midi(2) == midi(12)
        @test 4 * midi(-4) == midi(-16)
        @test 5 * midi(4) == midi(20)

        @test tomidi(midi(42)) == midi(42)
        @test octave(MidiInterval) == midi(12)
        @test sign(midi(3)) == 1
        @test sign(midi(0)) == 0
        @test sign(midi(-3)) == -1
        @test abs(midi(-3)) == midi(3)

        @test ic(midi(15)) == midic(3)
        @test ic(midi(-3)) == midic(9)
        @test embed(midi(42)) == midi(42)
        @test intervaltype(MidiInterval) == MidiInterval
        @test intervalclasstype(MidiInterval) == MidiIC

        for i in -2:2
            @test isstep(midi(i)) == true
        end
        for i in -13:-3
            @test isstep(midi(i)) == false
        end
        for i in 3:13
            @test isstep(midi(i)) == false
        end

        @test chromsemi(MidiInterval) == midi(1)
    end

    @testset "interval class interface" begin
        @test midic(3) + midic(10) == midic(1)
        @test midic(-3) + midic(4) == midic(1)
        @test midic(3) + midic(-4) == midic(11)
        @test midic(3) - midic(4) == midic(11)
        @test -midic(5) == midic(-5)
        @test -midic(7) == midic(0) - midic(7)

        @test zero(MidiIC) == midic(0)
        @test zero(midic(3)) == midic(0)

        @test midic(2) * 7 == midic(14)
        @test midic(-3) * 4 == midic(-12)
        @test midic(4) * -3 == midic(-12)
        @test 6 * midic(2) == midic(12)
        @test 4 * midic(-4) == midic(-16)
        @test 5 * midic(4) == midic(20)

        @test tomidi(midic(42)) == midic(42)
        @test octave(MidiIC) == midic(0)
        @test sign(midic(3)) == 1
        @test sign(midic(0)) == 0
        @test sign(midic(-3)) == -1
        @test abs(midic(-3)) == midic(3)

        @test ic(midic(15)) == midic(3)
        @test ic(midic(-3)) == midic(9)
        @test embed(midic(3)) == midi(3)
        @test embed(midic(-3)) == midi(9)
        @test intervaltype(MidiIC) == MidiInterval
        @test intervalclasstype(MidiIC) == MidiIC

        for i in -2:2
            @test isstep(midic(i)) == true
        end
        for i in -9:-3
            @test isstep(midic(i)) == false
        end
        for i in 3:9
            @test isstep(midic(i)) == false
        end

        @test chromsemi(MidiIC) == midic(1)
    end

    @testset "pitch interface" begin
        @test topitch(midi(3)) == midip(3)
        @test tointerval(midip(42)) == midi(42)

        @test midip(63) + midi(7) == midip(70)
        @test midip(63) + midi(-3) == midip(60)
        @test midip(63) - midi(7) == midip(56)
        @test midip(67) - midip(61) == midi(6)

        @test pc(midip(63)) == midipc(3)
        @test embed(midip(8)) == midip(8)
    end

    @testset "pitch class interface" begin
        @test topitch(midic(3)) == midipc(3)
        @test tointerval(midipc(42)) == midic(42)

        @test midipc(63) + midic(7) == midipc(70)
        @test midipc(63) + midic(-3) == midipc(60)
        @test midipc(63) - midic(7) == midipc(56)
        @test midipc(67) - midipc(61) == midic(6)

        @test pc(midipc(63)) == midipc(3)
        @test embed(midipc(8)) == midip(8)
    end
end


text:
str1 = "arne" ^ 100
str2 = "bsof" ^ 100
str3 = "arnebsof" ^ 100

dm = ncdmatrix([str1, str2, str3])

@test typeof(dm) == Array{Float64,2}

@test dm[1,1] == 0.0
@test dm[2,2] == 0.0
@test dm[3,3] == 0.0

@test dm[1,2] == dm[2,1]
@test dm[1,3] == dm[3,1]
@test dm[2,3] == dm[3,2]

@test dm[1,2] == ncd(str1, str2)
@test dm[1,3] == ncd(str1, str3)
@test dm[2,3] == ncd(str2, str3)

ext: jl


text:
function main(r::Robot)

    path=go_to_right_upper_corner_and_get_path(r)
    i=0
    while !isborder(r,West)
        move!(r,West)
        i=i+1
    end

    n=i
    m=West
    while n==i
        move!(r,Sud)
        n=0
        m=inverse(m)
        while !isborder(r,m)
            move!(r,m)
            n=n+1
        end
    end

    while isborder(r,m)==true
        putmarker!(r)
        move!(r,Sud)
    end

    putmarker!(r)
    move!(r,m)

    while isborder(r,Nord)==true
        putmarker!(r)
        move!(r,m)
    end

    putmarker!(r)
    move!(r,Nord)

    m=inverse(m)

    while isborder(r,m)==true
        putmarker!(r)
        move!(r,Nord)
    end

    putmarker!(r)
    move!(r,m)

    while isborder(r,Sud)==true
        putmarker!(r)
        move!(r,m)
    end

    putmarker!(r)

    go_to_right_upper_corner_and_get_path(r)
    go_back(r,path)
end

inverse(side::HorizonSide) = HorizonSide(mod(Int(side)+2, 4))

function go_to_right_upper_corner_and_get_path(r::Robot) ::Array
    path=[]
    while isborder(r,Nord)==false || isborder(r,Ost)==false
        if isborder(r,Nord)==false
            move!(r,Nord)
            push!(path,Sud)
        end
        if isborder(r,Ost)==false
            move!(r,Ost)
            push!(path,West)
        end
    end
    return path
end

function go_back(r,path::Array)
    n=length(path)
    while n>0
        move!(r,path[n])
        n=n-1
    end
end

max_stars_repo_name: VladislavBukin/mirea-progs

max_issues_repo_name: VladislavBukin/mirea-progs

max_forks_repo_name: VladislavBukin/mirea-progs


text:
#####
##### AlphaZero.jl
##### Jonathan Laurent, Carnegie Mellon University (2019-2020)
#####

module AlphaZero

  import Distributed
  using Formatting
  using Base: @kwdef
  using DataStructures: CircularBuffer
  using Distributions: Categorical, Dirichlet
  using Statistics: mean

  # When running on a CPU, having multiple threads does not play
  # well with BLAS multithreading
  import LinearAlgebra
  LinearAlgebra.BLAS.set_num_threads(1)

  # Internal helper functions
  include("util.jl")
  using .Util
  export Util
  export apply_temperature

  # A generic interface for single-player or zero-sum two-players games.
  include("game.jl")
  using .GameInterface
  const GI = GameInterface
  export GameInterface, GI
  export AbstractGameEnv
  export AbstractGameSpec

  # A standalone, generic MCTS implementation
  include("mcts.jl")
  using .MCTS
  export MCTS

  # A generic network interface
  include("networks/network.jl")
  using .Network
  export Network
  export AbstractNetwork
  export OptimiserSpec
  export CyclicNesterov, Adam

  # Utilities to batch oracle calls
  include("batchifier.jl")
  using .Batchifier
  export Batchifier

  # Schedules
  include("schedule.jl")
  export AbstractSchedule
  export ConstSchedule, PLSchedule, StepSchedule, CyclicSchedule

  # Training params
  include("params.jl")
  export Params
  export MctsParams
  export SimParams
  export SelfPlayParams
  export LearningParams
  export ArenaParams
  export MemAnalysisParams
  export SamplesWeighingPolicy, CONSTANT_WEIGHT, LOG_WEIGHT, LINEAR_WEIGHT

  # Stats about training
  include("report.jl")
  export Report

  # Game traces
  include("trace.jl")
  export Trace
  export total_reward

  # Players and games
  include("play.jl")
  export AbstractPlayer, think, select_move, reset_player!, player_temperature
  export MctsPlayer
  export RandomPlayer
  export NetworkPlayer
  export PlayerWithTemperature
  export EpsilonGreedyPlayer
  export TwoPlayers
  export Human, interactive!
  export play_game

  # Utilities for distributed games simulation
  include("simulations.jl")
  export Simulator, simulate, simulate_distributed
  export record_trace
  export rewards_and_redundancy

  # Memory buffer to hold samples generated during self-play
  include("memory.jl")
  export MemoryBuffer, get_experience

  # Utilities to train the neural network based on collected samples
  include("learning.jl")

  # Main training algorithm
  include("training.jl")
  export Env, train!, initial_report
  export Handlers
  export AlphaZeroPlayer

  # A minmax player to be used as a baseline
  include("minmax.jl")
  using .MinMax
  export MinMax

  # Utilities to write benchmarks
  include("benchmark.jl")
  using .Benchmark
  export Benchmark

  # We provide a library of standard network, both in Knet and Flux.
  # Which backend is used to implement this library is determined during precompilation
  # based on the value of the ALPHAZERO_DEFAULT_DL_FRAMEWORK environment variable.
  const DEFAULT_DL_FRAMEWORK = get(ENV, "ALPHAZERO_DEFAULT_DL_FRAMEWORK", "FLUX")

  # Only Flux works with CUDA,NNlib#master right now.

  if DEFAULT_DL_FRAMEWORK == "FLUX"
    # @info "Using the Flux implementation of AlphaZero.NetLib."
    KNET_ADVICE =
    """
    For optimal performances, we recommend that you configure AlphaZero.jl
    to use Knet as a backend by setting ALPHAZERO_DEFAULT_DL_FRAMEWORK to KNET.
    """
    # Knet does not work with the latest CUDA.jl so we disable this advice for now.
    # @info KNET_ADVICE
    @eval begin
      include("networks/flux.jl")
      const NetLib = FluxLib
    end
  elseif DEFAULT_DL_FRAMEWORK == "KNET"
    # @info "Using the Knet implementation of AlphaZero.NetLib."
    @eval begin
      include("networks/knet.jl")
      const NetLib = KnetLib
    end
  else
    error("Unknown DL framework: $(DEFAULT_DL_FRAMEWORK)")
  end

  using .NetLib
  export NetLib
  export SimpleNet, SimpleNetHP, ResNet, ResNetHP

  # A structure that contains the information necessary to replicate a training session
  include("experiments.jl")
  using .Experiments
  export Experiments
  export Experiment

  # The default user interface is included here for convenience but it could be
  # replaced or separated from the main AlphaZero.jl package (which would also
  # enable dropping some dependencies such as Crayons or JSON3).
  include("ui/ui.jl")
  using .UserInterface
  const UI = UserInterface
  export UserInterface, UI
  export Session, resume!, save
  export explore

  # Bridge with CommonRLInterface.jl
  include("common_rl_intf.jl")
  export CommonRLInterfaceWrapper

  # A small library of standard examples
  include("examples.jl")
  export Examples

  # Scripts
  include("scripts/scripts.jl")
  export Scripts

end


text:
using LightGraphs.SimpleGraphs: SimpleEdge, AbstractSimpleGraph
using LightGraphs.SimpleGraphs: SimpleGraphs, fadj, ne, badj, add_edge!, rem_edge!
using LightGraphs: ne, edgetype, has_edge, is_directed

"""
    AbstractWrappedSimpleGraph{T} <: AbstractSimpleGraph{T}

A wrapper type that wraps a SimpleGraph, used to store additional info in it's
type. Mainly used to store information about the spatial embedding (in 1D,2D)
and directions (x,y,z...) of the underlying graph.

Must implement the field `bare_graph::SG` that should store the underlying
graph.
"""
abstract type AbstractWrappedSimpleGraph{T} <: AbstractSimpleGraph{T} end

## imple
bare(g::AbstractWrappedSimpleGraph) = g.bare_graph

@forward AbstractWrappedSimpleGraph.bare_graph SimpleGraphs.fadj, SimpleGraphs.badj, LightGraphs.ne,
    LightGraphs.edgetype, LightGraphs.is_directed, LightGraphs.has_edge, Base.eltype,
        SimpleGraphs.add_edge!, SimpleGraphs.rem_edge!, LightGraphs.nv, LightGraphs.edges

LightGraphs.ne(g::AbstractWrappedSimpleGraph) = ne(bare(g))
LightGraphs.is_directed(g::Type{<:AbstractWrappedSimpleGraph}) = is_directed(fieldtype(g, :bare_graph))


text:
# function testlinearomega()
#   x = uniform(0.0, 1.0)
#   rand(x; ΩT = LinearΩ)
# end

# function x_(ω)
#   a = rand(ω)
#   b = bernoulli(ω, 0.3, Bool)
#   c = normal(ω, 1.0, 2.0)
#   a + b + c
# end

function testlinear()
  x = normal(0.0, 1.0, (1,2,3))
  w = Omega.LinearΩ()
  x(w)
end

# x = ciid(x_)
# ω = LinearΩ()
# x(ω)


text:
using KernelAbstractions: NoneEvent

using Oceananigans.Grids: Flat, Bounded

import Oceananigans.Utils: launch!

#####
##### Implements a "single column model mode" for HydrostaticFreeSurfaceModel
#####

const SingleColumnGrid = AbstractGrid{<:AbstractFloat, <:Flat, <:Flat, <:Bounded}

PressureField(arch, ::SingleColumnGrid) = (pHY′ = nothing,)
FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing
FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, arch, ::SingleColumnGrid) = nothing

validate_momentum_advection(momentum_advection, ::SingleColumnGrid) = nothing
validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, ::SingleColumnGrid) = nothing, NamedTuple()
validate_tracer_advection(tracer_advection::Nothing, ::SingleColumnGrid) = nothing, NamedTuple()

@inline hydrostatic_pressure_y_gradient(i, j, k, grid, ::Nothing) = zero(eltype(grid))

@inline launch!(arch, ::SingleColumnGrid, ::Val{:xy},  args...; kwargs...) = NoneEvent()
@inline launch!(arch, ::SingleColumnGrid, ::Val{dims}, args...; kwargs...) where dims = launch!(arch, grid, dims, args...; kwargs...)


text:
doc = """SoftConfidenceWeighted sample

Usage:
  sample.jl [-c  <conf>]
  sample.jl [--confidence <conf>]
  sample.jl [-a <agr>]
  sample.jl [--aggressiveness <agr>]
  sample.jl -h | --help

Options:
  -h --help                         Show this screen.
  -c <conf>, --confidence=<conf>    Confidence parameter [default: 0.7].
  -a <agr>, --aggressiveness=<agr>  Aggressiveness parameter [default: 0.7].
"""

using DocOpt
using SoftConfidenceWeighted
using JSON

arguments = docopt(doc)
confidence = float(arguments["--confidence"])
aggressiveness = float(arguments["--aggressiveness"])

scw = SCWParameter(confidence, aggressiveness)

for line in EachLine(STDIN)
  line = chomp(line)
  if line == ""
    break
  end

  label, data = JSON.parse(line)
  if int(label) == 0
    label = classify(scw, data)
    println("classify: $label $data")
  else
    update(scw, data, label)
    println("update: $label $data")
  end
end


text:
export Bijection, IdentityTransformation, ClosedTransformation, SquareMinTransformation, ExpMinTransformation, inv

abstract type Bijection end

Base.broadcastable(p::Bijection) = Ref(p)

struct IdentityTransformation{T} <: Bijection end
(f::IdentityTransformation{T})(x::T) where {T} = x
inv(f::IdentityTransformation{T}, x::T) where {T} = x

struct ClosedTransformation{T} <: Bijection
    minValue::T
    maxValue::T
end

function (f::ClosedTransformation{T})(x) where {T}
    return (f.minValue + f.maxValue) / 2 + (f.minValue - f.maxValue) / 2 * cos(x)
end

function inv(f::ClosedTransformation{T}, y) where {T}
    x = (y - f.minValue + (f.minValue - f.maxValue) / 2) / (f.minValue - f.maxValue) * 2
    return acos(x)
end

struct SquareMinTransformation{T} <: Bijection
    minValue::T
end

function (f::SquareMinTransformation{T})(x) where {T}
    return x^2 + f.minValue
end

function inv(f::SquareMinTransformation{T}, y) where {T}
    return sqrt(y - f.minValue)
end

struct ExpMinTransformation{T} <: Bijection
    minValue::T
end

function (f::ExpMinTransformation{T})(x) where {T}
    return exp(x) + f.minValue
end

function inv(f::ExpMinTransformation{T}, y) where {T}
    return log(y - f.minValue)
end

struct MQMinTransformation{T} <: Bijection
    minValue::T
    r::T
end

function (f::MQMinTransformation{T})(x) where {T}
    return sqrt(1+ (f.r*x)^2) + f.minValue - 1
end

function inv(f::MQMinTransformation{T}, y) where {T}
    return sqrt((y + 1 - f.minValue)^2 - 1)/f.r
end


text:
# pilot_testing.jl

# @code_warntype Node(Int16(1), Int16[1], trues(3))
n1 = Node(Int16(1), Int16[1], trues(3))

# @code_warntype Node(Int16(1), Int16[1, 2], trues(3), n1)
n2 = Node(Int16(1), Int16[1, 2], trues(3), n1, 1)
n3 = Node(Int16(1), Int16[1, 3], trues(3), n1, 1)
n4 = Node(Int16(1), Int16[1, 4], trues(3), n1, 1)
n5 = Node(Int16(1), Int16[1, 5], trues(3), n1, 1)
n6 = Node(Int16(1), Int16[1, 6], trues(3), n1, 1)
n7 = Node(Int16(1), Int16[2, 3], trues(3), n1, 1)
n8 = Node(Int16(1), Int16[2, 4], trues(3), n1, 1)
n9 = Node(Int16(1), Int16[2, 5], trues(3), n1, 1)
n10 = Node(Int16(1), Int16[2, 6], trues(3), n1, 1)

push!(n1.children, n2)
push!(n1.children, n3)
push!(n1.children, n4)
push!(n1.children, n5)
push!(n1.children, n6)
push!(n1.children, n7)
push!(n1.children, n8)
push!(n1.children, n9)
push!(n1.children, n10)

@code_warntype has_children(n1)

@code_warntype younger_siblings(n1.children[1])
younger_siblings(n1.children[1])

@code_warntype growtree!(n2, 1, 3, 3)
growtree!(n2, 1, 3, 3)

t = [["a", "b"],
     ["b", "c", "d"],
     ["a", "c"],
     ["e", "b"],
     ["a", "c", "d"],
     ["a", "e"],
     ["a", "b", "c"],
     ["c", "b", "e"]]

@code_warntype unique_items(t);
@time unique_items(t);

unq = unique_items(t)
@code_warntype occurrence(t, unq)
@time occurrence(t, unq)

t = [["a", "b"],
     ["b", "c", "d"],
     ["a", "c"],
     ["e", "b"],
     ["a", "c", "d"],
     ["a", "e"],
     ["a", "b", "c"],
     ["c", "b", "e", "f"]]

unq = unique_items(t)

@code_warntype _frequent(t, unq, 0.01, 3)
xtree1 = _frequent(t, unq, 0.01, 4)

itemlist = randstr(25, 16);

n = 100_000
m = 25             # number of items in transactions
mx_depth = 7
t = [sample(itemlist, m, replace = false) for _ in 1:n];

# @code_warntype _frequent(t, 1)
@time unq2 = unique_items(t);
@time occ2 = occurrence(t, unq2);
@time xtree1 = _frequent(t, unq2, round(Int, 0.01*n), mx_depth);

# Rule Generation Testing
@code_warntype gen_support_dict(xtree1, n)

t1 = [["a", "b"],
     ["b", "c", "d"],
     ["a", "c"],
     ["e", "b"],
     ["a", "c", "d"],
     ["a", "e"],
     ["a", "b", "c"],
     ["c", "b", "e", "f"]]

@code_warntype _frequent(t1, 1, 3)
unq3 = unique_items(t1)
xtree1 = _frequent(t1, unq3, 1, 4);
@code_warntype gen_support_dict(xtree1, length(t1))
xsup = gen_support_dict(xtree1, length(t1))

@code_warntype gen_node_rules(xtree1.children[1].children[1].children[1], xsup, 3, 8)

xrules = gen_node_rules(xtree1.children[1].children[1].children[1], xsup, 3, 8)

rule_arr = Array{Rule, 1}(0)
gen_rules!(rule_arr, xtree1.children[1], xsup, 2, 8)

t1 = [["a", "b"],
     ["b", "c", "d"],
     ["a", "c"],
     ["e", "b"],
     ["a", "c", "d"],
     ["a", "e"],
     ["a", "b", "c"],
     ["c", "b", "e", "f"]]

unq = unique_items(t1)
xtree = _frequent(t1, unq, 1, 4);
shownodes(xtree)

# Comparing with R
a_list = [
    ["a", "b"],
    ["a", "c"],
    ["a", "b", "c"],
    ["a", "b", "d"],
    ["a", "c", "d"],
    ["a", "b", "c", "d"],
    ["a", "b", "c", "e"],
    ["b", "d", "e", "f"],
    ["a", "c", "e", "f"],
    ["b", "c", "d", "e", "f"],
    ["a", "c", "d", "e", "f"],
    ["b", "c", "d", "e", "f"]
]

xunq = unique_items(a_list)
xtree1 = _frequent(a_list, xunq, 1, 6);
xsup = gen_support_dict(xtree1, length(a_list))

xrules = gen_rules(xtree1, xsup, 12)

apriori(a_list, 0.01, 6)

# Example below used to track down bug of unexplored node Int16[4, 10, 14]
transactions = [["milk", "eggs", "bread"],
                ["butter", "milk", "sugar", "flour", "eggs"],
                ["bacon", "eggs", "milk", "beer"],
                ["bread", "ham", "turkey"],
                ["cheese", "ham", "bread", "ketchup"],
                ["mustard", "hot dogs", "buns", "hamburger", "cheese", "beer"],
                ["milk", "sugar", "eggs"],
                ["hamburger", "ketchup", "milk", "beer"],
                ["ham", "cheese", "bacon", "eggs"]]

freq = frequent(transactions, 1, 7)
display(freq)

for x in freq[:itemset]
    println(x)
end

unq = unique_items(transactions)

node = _frequent(transactions, unq, 0.01, 6)
sup = gen_support_dict(node, 9)

for k in keys(sup)
    println(k)
end

rules = apriori(transactions, 0.1, 4)

transactions = [["milk", "eggs", "bread"],
                ["butter", "milk", "sugar", "flour", "eggs"],
                ["bacon", "eggs", "milk", "beer"],
                ["bread", "ham", "turkey"],
                ["cheese", "ham", "bread", "ketchup"],
                ["mustard", "hot dogs", "buns", "hamburger", "cheese", "beer"]]

freq = frequent(transactions, 1, 7)
# display(freq)

unq = unique_items(transactions)
unq[[4, 10, 14]]

node = _frequent(transactions, unq, 0.01, 6)
shownodes(node)

sup = gen_support_dict(node, 9)


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

@doc raw"""HorizontalPodAutoscalerStatus describes the current status of a horizontal pod autoscaler.

    IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus(;
        conditions=nothing,
        currentMetrics=nothing,
        currentReplicas=nothing,
        desiredReplicas=nothing,
        lastScaleTime=nothing,
        observedGeneration=nothing,
    )

    - conditions::Vector{IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerCondition} : conditions is the set of conditions required for this autoscaler to scale its target, and indicates whether or not those conditions are met.
    - currentMetrics::Vector{IoK8sApiAutoscalingV2beta2MetricStatus} : currentMetrics is the last read state of the metrics used by this autoscaler.
    - currentReplicas::Int32 : currentReplicas is current number of replicas of pods managed by this autoscaler, as last seen by the autoscaler.
    - desiredReplicas::Int32 : desiredReplicas is the desired number of replicas of pods managed by this autoscaler, as last calculated by the autoscaler.
    - lastScaleTime::IoK8sApimachineryPkgApisMetaV1Time : lastScaleTime is the last time the HorizontalPodAutoscaler scaled the number of pods, used by the autoscaler to control how often the number of pods is changed.
    - observedGeneration::Int64 : observedGeneration is the most recent generation observed by this autoscaler.
"""
mutable struct IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus <: SwaggerModel
    conditions::Any # spec type: Union{ Nothing, Vector{IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerCondition} } # spec name: conditions
    currentMetrics::Any # spec type: Union{ Nothing, Vector{IoK8sApiAutoscalingV2beta2MetricStatus} } # spec name: currentMetrics
    currentReplicas::Any # spec type: Union{ Nothing, Int32 } # spec name: currentReplicas
    desiredReplicas::Any # spec type: Union{ Nothing, Int32 } # spec name: desiredReplicas
    lastScaleTime::Any # spec type: Union{ Nothing, IoK8sApimachineryPkgApisMetaV1Time } # spec name: lastScaleTime
    observedGeneration::Any # spec type: Union{ Nothing, Int64 } # spec name: observedGeneration

    function IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus(;conditions=nothing, currentMetrics=nothing, currentReplicas=nothing, desiredReplicas=nothing, lastScaleTime=nothing, observedGeneration=nothing)
        o = new()
        validate_property(IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus, Symbol("conditions"), conditions)
        setfield!(o, Symbol("conditions"), conditions)
        validate_property(IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus, Symbol("currentMetrics"), currentMetrics)
        setfield!(o, Symbol("currentMetrics"), currentMetrics)
        validate_property(IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus, Symbol("currentReplicas"), currentReplicas)
        setfield!(o, Symbol("currentReplicas"), currentReplicas)
        validate_property(IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus, Symbol("desiredReplicas"), desiredReplicas)
        setfield!(o, Symbol("desiredReplicas"), desiredReplicas)
        validate_property(IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus, Symbol("lastScaleTime"), lastScaleTime)
        setfield!(o, Symbol("lastScaleTime"), lastScaleTime)
        validate_property(IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus, Symbol("observedGeneration"), observedGeneration)
        setfield!(o, Symbol("observedGeneration"), observedGeneration)
        o
    end
end # type IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus

const _property_map_IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus = Dict{Symbol,Symbol}(Symbol("conditions")=>Symbol("conditions"), Symbol("currentMetrics")=>Symbol("currentMetrics"), Symbol("currentReplicas")=>Symbol("currentReplicas"), Symbol("desiredReplicas")=>Symbol("desiredReplicas"), Symbol("lastScaleTime")=>Symbol("lastScaleTime"), Symbol("observedGeneration")=>Symbol("observedGeneration"))
const _property_types_IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus = Dict{Symbol,String}(Symbol("conditions")=>"Vector{IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerCondition}", Symbol("currentMetrics")=>"Vector{IoK8sApiAutoscalingV2beta2MetricStatus}", Symbol("currentReplicas")=>"Int32", Symbol("desiredReplicas")=>"Int32", Symbol("lastScaleTime")=>"IoK8sApimachineryPkgApisMetaV1Time", Symbol("observedGeneration")=>"Int64")
Base.propertynames(::Type{ IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus }) = collect(keys(_property_map_IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus))
Swagger.property_type(::Type{ IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus[name]))}
Swagger.field_name(::Type{ IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus }, property_name::Symbol) =  _property_map_IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus[property_name]

function check_required(o::IoK8sApiAutoscalingV2beta2HorizontalPodAutoscalerStatus)
    (getproperty(o, Symbol("conditions")) === nothing) && (return false)
    (getproperty(o, Symbol("currentReplicas")) === nothing) && (return false)
    (getproperty(o, Symbol("desiredReplicas")) === nothing) && (return false)
    true
end

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


text:
module Heuristics

using StaticArrays
using Format
using ThreadPools
import Random: shuffle!

# Required
include("Problem.jl")
include("Options.jl")
include("Optimizers.jl")
include("Results.jl")

# PSO includes
include("./PSO/psoUtil.jl")
include("./PSO/Particle.jl")
include("./PSO/Swarm.jl")
include("./PSO/PSO.jl")
include("./PSO/MS_PSO.jl")

# Exports
export Problem
export Options
export optimize!

export PSO
export MS_PSO

end


text:
"""
Generic ALPS lattice parsed from XML file.
"""
struct ALPSLattice <: AbstractLattice
    sites::Int # n_sites
    dim::Int
    # row = neighbors; col = siteidx (assumption: const. coordination nr.)
    neighs::Matrix{Int}
    n_neighs::Int
    bond_vecs::Matrix{Float64}

    # for generic checkerboard support
    n_bonds::Int
    bonds::Matrix{Int} # src, trg, type
end

# constructors
function ALPSLattice(xmlfile::String)
    sites, dim, n_neighs, bond_vecs, n_bonds, bonds = parse_alpslattice_xml(xmlfile, l)
    neighs = build_neighbortable(ALPSLattice, n_neighs, sites, n_bonds, bonds)
    return ALPSLattice(sites, dim, neighs, n_neighs, bond_vecs, n_bonds, bonds)
end

function build_neighbortable(::Type{ALPSLattice}, n_neighs, sites, n_bonds, bonds)
    neighs = zeros(Int, n_neighs, sites)

    # we store src->trg before trg->src
    try
        nc = fill(1, sites)
        for b in 1:n_bonds
            src = bonds[b,1]
            trg = bonds[b,2]
            neighs[nc[src], src] = trg
            nc[src]+=1
        end
        for b in 1:n_bonds
            src = bonds[b,1]
            trg = bonds[b,2]
            neighs[nc[trg], trg] = src
            nc[trg]+=1
        end
    catch e
        if isa(e, BoundsError)
            warn(
                "Lattice seems to not have a constant coordination number. " *
                "Fields `n_neighs` and `neighs` shouldn't be used."
            )
        else
            throw(e)
        end
    end
    return neighs
end

function parse_alpslattice_xml(filename::String)
    xdoc = parse_file(filename)
    graph = LightXML.root(xdoc)
    sites = 1

    sites = parse(Int, attribute(graph, "vertices"; required=true))
    dim = parse(Int, attribute(graph, "dimension"; required=true))

    edges = get_elements_by_tagname(graph, "EDGE")
    n_bonds = length(edges)

    # bonds & bond vectors
    bonds = zeros(n_bonds, 3)
    bond_vecs = zeros(n_bonds, dim)
    v = Vector{Float64}(dim)
    for (i, edge) in enumerate(edges)
        src = 0
        trg = 0
        src = parse(Int, attribute(edge, "source"; required=true))
        trg = parse(Int, attribute(edge, "target"; required=true))
        typ = parse(Int, attribute(edge, "type"; required=true))
        id = parse(Int, attribute(edge, "id"; required=true))
        v = [parse(Float64, f) for f in split(attribute(edge, "vector"; required=true)," ")]

        if id != i error("Edges in lattice file must be sorted from 1 to N!") end

        bonds[i, 1] = src
        bonds[i, 2] = trg
        bonds[i, 3] = typ
        bond_vecs[i, :] = v

    end

    # neighbors of site 1
    n_neighs = count(x->x==1, bonds[:, 1]) + count(x->x==1, bonds[:, 2])
    return sites, dim, n_neighs, bond_vecs, n_bonds, bonds
end

# Implement AbstractLattice interface: mandatory
@inline Base.length(l::ALPSLattice) = l.sites
function Base.size(l::ALPSLattice)
    @warn "Guessing size of ALPSLattice"
    if l.sites % l.ndims == 0
        return tuple((div(l.sites, l.ndims) for _ in 1:l.ndims)...)
    else
        error("Failed to guess size of ALPSLattice")
    end
end

# Implement AbstractLattice interface: optional
@inline neighbors_lookup_table(l::ALPSLattice) = copy(l.neighs)

# HasNeighborsTable and HasBondsTable traits
has_neighbors_table(::ALPSLattice) = HasNeighborsTable()
has_bonds_table(::ALPSLattice) = HasBondsTable()


text:
using Test
using MechGlueDiffEqBase
using MechanicalUnits: @import_expand, dimension, NoDims
@import_expand(N, s, m)
using DifferentialEquations
@testset "ODE" begin
    f = (y, p, t) -> 0.5y / 3.0  # The derivative, yeah
    u0 = 1.5
    @test zero(u0) == 0 * u0
    @test value(u0) == 1.5
    tspan = (0.0, 1.0)
    prob = ODEProblem(f, u0, tspan)
    integrator = DiffEqBase.__init(prob, Tsit5(), dt = 0.1)
    # Internalnorm's range should be dimensionless
    internalnorm = integrator.opts.internalnorm
    @test dimension(internalnorm(u0, tspan[1])) == NoDims
    sol = solve(prob, Tsit5(), dt=0.1)
    @test sol(0.0) == u0
    @test sol(1.0) ≈ 1.7720406
    @time solve(prob, Tsit5()) # 23 alloc, 4797 KiB
end
@testset "ODE quantity" begin
    f = (y, p, t) -> 0.5y / 3.0s  # The derivative, yeah
    u0 = 1.5N
    @test zero(u0) == 0 * u0
    @test value(u0) == 1.5
    tspan = (0.0s, 1.0s)
    prob = ODEProblem(f, u0, tspan)
    integrator = DiffEqBase.__init(prob, Tsit5(), dt = 0.1s)
    # Internalnorm's range should be dimensionless
    internalnorm = integrator.opts.internalnorm
    @test dimension(internalnorm(u0, tspan[1])) == NoDims
    sol = solve(prob, Tsit5(), dt=0.1s)
    @test sol(0.0s) == u0
    @test sol(1.0s) ≈ 1.7720406N
    @time solve(prob, Tsit5()) # 29 alloc, 5641KiB
end

@testset "ODE quantity algorithms" begin
    f = (y, p, t) -> 0.5y / 3.0s  # The derivative, yeah
    u0 = 1.5N
    @test zero(u0) == 0 * u0
    @test value(u0) == 1.5
    tspan = (0.0s, 1.0s)
    prob = ODEProblem(f, u0, tspan)
    integrator = DiffEqBase.__init(prob, Tsit5(), dt = 0.1s)
    # Internalnorm's range should be dimensionless
    internalnorm = integrator.opts.internalnorm
    @test dimension(internalnorm(u0, tspan[1])) == NoDims
    sol = solve(prob, Tsit5(), dt=0.1s)
    @test sol(0.0s) == u0
    @test sol(1.0s) ≈ 1.7720406N
    @time solve(prob, Tsit5()) # 29 alloc, 5641KiB
end


text:
@testset "Patterns" begin
    items = [
        (:train, :input, :prices, :time),
        (:train, :input, :prices, :id),
        (:train, :input, :prices, :lag),
        (:train, :input, :load, :time),
        (:train, :input, :load, :id),
        (:train, :input, :temperature, :time),
        (:train, :input, :temperature, :id),
        (:train, :output, :prices, :time),
        (:train, :output, :prices, :id),
        (:predict, :input, :prices, :time),
        (:predict, :input, :prices, :id),
        (:predict, :input, :prices, :lag),
        (:predict, :input, :load, :time),
        (:predict, :input, :load, :id),
        (:predict, :input, :temperature, :time),
        (:predict, :input, :temperature, :id),
        (:predict, :output, :prices, :time),
        (:predict, :output, :prices, :id),
    ]

    @testset "Pattern(:_, :_, :_, :time)" begin
        pattern = Pattern(:_, :_, :_, :time)
        @test filter(in(pattern), items) == [
            (:train, :input, :prices, :time),
            (:train, :input, :load, :time),
            (:train, :input, :temperature, :time),
            (:train, :output, :prices, :time),
            (:predict, :input, :prices, :time),
            (:predict, :input, :load, :time),
            (:predict, :input, :temperature, :time),
            (:predict, :output, :prices, :time),
        ]
    end

    @testset "Pattern(:__, :time)" begin
        pattern = Pattern(:__, :time)
        @test filter(in(pattern), items) == [
            (:train, :input, :prices, :time),
            (:train, :input, :load, :time),
            (:train, :input, :temperature, :time),
            (:train, :output, :prices, :time),
            (:predict, :input, :prices, :time),
            (:predict, :input, :load, :time),
            (:predict, :input, :temperature, :time),
            (:predict, :output, :prices, :time),
        ]
    end

    @testset "Pattern(:_, :time)" begin
        pattern = Pattern(:_, :time)
        @test isempty(filter(in(pattern), items))
    end

    @testset "Pattern(:train, :input, :_, :time)" begin
        pattern = Pattern(:train, :input, :_, :time)
        @test filter(in(pattern), items) == [
            (:train, :input, :prices, :time),
            (:train, :input, :load, :time),
            (:train, :input, :temperature, :time),
        ]
    end

    @testset "Pattern(:train, :input, :__)" begin
        pattern = Pattern(:train, :input, :__)
        @test filter(in(pattern), items) == [
            (:train, :input, :prices, :time),
            (:train, :input, :prices, :id),
            (:train, :input, :prices, :lag),
            (:train, :input, :load, :time),
            (:train, :input, :load, :id),
            (:train, :input, :temperature, :time),
            (:train, :input, :temperature, :id),
        ]
    end

    @testset "Reductions" begin
        @testset "Pattern(:__, :time)" begin
            reduced = Pattern(:__, :time)

            patterns = Pattern[
                (:__, :_, :time),
                (:_, :__, :time),
                (:__, :__, :time),
                (:_, :__, :_, :time),
            ]
            @testset "$p" for p in patterns
                @test p == reduced
            end

            @test Pattern(:_, :_, :time) != reduced
            @test Pattern(:_, :__, :time, :_) != reduced
            @test Pattern(:time, :_, :__) != reduced
        end

        @testset "Pattern(:train, :__)" begin
            reduced = Pattern(:train, :__)

            patterns = Pattern[
                (:train, :_, :__),
                (:train, :__, :_),
                (:train, :__, :__),
                (:train, :_, :__, :_),
            ]
            @testset "$p" for p in patterns
                @test p == reduced
            end

            @test Pattern(:train, :_, :_) != reduced
            @test Pattern(:_, :train, :__, :_) != reduced
            @test Pattern(:__, :_, :train) != reduced
        end

        @testset "Pattern(:_, :input, :__)" begin
            reduced = Pattern(:_, :input, :__)

            @test Pattern(:_, :input, :_, :__) == reduced
            @test Pattern(:_, :__, :input, :__) != reduced
        end
    end

    @testset "Mixed types" begin
        t1 = today()
        t2 = t1 + Day(1)
        items = [
            (t1, 1, "prices", :time),
            (t1, 1, "prices", :id),
            (t1, 1, "prices", :lag),
            (t1, 1, "load", :time),
            (t1, 1, "load", :id),
            (t1, 1, "temperature", :time),
            (t1, 1, "temperature", :id),
            (t1, 2, "prices", :time),
            (t1, 2, "prices", :id),
            (t2, 1, "prices", :time),
            (t2, 1, "prices", :id),
            (t2, 1, "prices", :lag),
            (t2, 1, "load", :time),
            (t2, 1, "load", :id),
            (t2, 1, "temperature", :time),
            (t2, 1, "temperature", :id),
            (t2, 2, "prices", :time),
            (t2, 2, "prices", :id),
        ]

        pattern = Pattern(t1, 1, :__)
        @test filter(in(pattern), items) == [
            (t1, 1, "prices", :time),
            (t1, 1, "prices", :id),
            (t1, 1, "prices", :lag),
            (t1, 1, "load", :time),
            (t1, 1, "load", :id),
            (t1, 1, "temperature", :time),
            (t1, 1, "temperature", :id),
        ]

        @testset "Subtype matching" begin
            t1 = Float64
            t2 = Int
            items = [
                (t1, 1, "prices", :time),
                (t1, 1, "prices", :id),
                (t1, 1, "prices", :lag),
                (t1, 1, "load", :time),
                (t1, 1, "load", :id),
                (t1, 1, "temperature", :time),
                (t1, 1, "temperature", :id),
                (t1, 2, "prices", :time),
                (t1, 2, "prices", :id),
                (t2, 1, "prices", :time),
                (t2, 1, "prices", :id),
                (t2, 1, "prices", :lag),
                (t2, 1, "load", :time),
                (t2, 1, "load", :id),
                (t2, 1, "temperature", :time),
                (t2, 1, "temperature", :id),
                (t2, 2, "prices", :time),
                (t2, 2, "prices", :id),
            ]

            pattern = Pattern(AbstractFloat, 1, :__)
            @test filter(in(pattern), items) == [
                (t1, 1, "prices", :time),
                (t1, 1, "prices", :id),
                (t1, 1, "prices", :lag),
                (t1, 1, "load", :time),
                (t1, 1, "load", :id),
                (t1, 1, "temperature", :time),
                (t1, 1, "temperature", :id),
            ]

            pattern = Pattern(Integer, 1, :__)
            @test filter(in(pattern), items) == [
                (t2, 1, "prices", :time),
                (t2, 1, "prices", :id),
                (t2, 1, "prices", :lag),
                (t2, 1, "load", :time),
                (t2, 1, "load", :id),
                (t2, 1, "temperature", :time),
                (t2, 1, "temperature", :id),
            ]

            @test filter(in(Pattern(Real, :__)), items) == items
        end
    end
end


text:
@test floorsum(231248135,143081212,21332,23943) == 3986234580397
@test floorsum(1000000000,1000000000,999999999,999999999) == 499999999500000000
@test floorsum(919229602,311231321,5123123,1231231) == 6954556907962778
@test floorsum(5,5,3,3) == 7


text:
#####
##### CPU channel pressure solver
#####

function ChannelPressureSolver(::CPU, grid, pressure_bcs, planner_flag=FFTW.PATIENT)
    kx², ky², kz² = generate_discrete_eigenvalues(grid, pressure_bcs)
    wavenumbers = (kx² = kx², ky² = ky², kz² = kz²)

    # Storage for RHS and Fourier coefficients is hard-coded to be Float64
    # because of precision issues with Float32.
    # See https://github.com/climate-machine/Oceananigans.jl/issues/55
    storage = zeros(Complex{Float64}, grid.Nx, grid.Ny, grid.Nz)

    @debug "Planning transforms for PressureSolver{Channel, CPU}..."
    x_bc, y_bc, z_bc = pressure_bcs.x.left, pressure_bcs.y.left, pressure_bcs.z.left
    FFTx!   = plan_forward_transform(storage, x_bc, 1, planner_flag)
    DCTyz!  = plan_forward_transform(storage, z_bc, [2, 3], planner_flag)
    IFFTx!  = plan_backward_transform(storage, x_bc, 1, planner_flag)
    IDCTyz! = plan_backward_transform(storage, z_bc, [2, 3], planner_flag)
    @debug "Planning transforms for PressureSolver{Channel, CPU} done!"

    transforms = ( FFTx! =  FFTx!,  DCTyz! =  DCTyz!,
                  IFFTx! = IFFTx!, IDCTyz! = IDCTyz!)

    return PressureSolver(Channel(), CPU(), wavenumbers, storage, transforms, nothing)
end

function solve_poisson_equation!(solver::PressureSolver{Channel, CPU}, grid)
    Nx, Ny, Nz, _ = unpack_grid(grid)
    kx², ky², kz² = solver.wavenumbers.kx², solver.wavenumbers.ky², solver.wavenumbers.kz²

    # We can use the same storage for the RHS and the solution ϕ.
    RHS, ϕ = solver.storage, solver.storage

    solver.transforms.DCTyz! * RHS  # Calculate DCTʸᶻ(f) in place.
    solver.transforms.FFTx!  * RHS  # Calculate FFTˣ(f) in place.

    # Solve the discrete Poisson equation in spectral space. We are computing
    # the Fourier coefficients of the solution from the Fourier coefficients
    # of the RHS.
    @. ϕ = -RHS / (kx² + ky² + kz²)

    # Setting DC component of the solution (the mean) to be zero. This is also
    # necessary because the source term to the Poisson equation has zero mean
    # and so the DC component comes out to be ∞.
    ϕ[1, 1, 1] = 0

    solver.transforms.IFFTx!  * ϕ  # Calculate IFFTˣ(ϕ) in place.
    solver.transforms.IDCTyz! * ϕ  # Calculate IDCTʸᶻ(ϕ) in place.

    # Must normalize by 2N for each dimension transformed via FFTW.REDFT.
    @. ϕ = ϕ / (4Ny*Nz)

    return nothing
end

#####
##### GPU channel pressure solver
#####

function ChannelPressureSolver(::GPU, grid, pressure_bcs, no_args...)
    Nx, Ny, Nz, _ = unpack_grid(grid)

    kx², ky², kz² = generate_discrete_eigenvalues(grid, pressure_bcs) .|> CuArray
    wavenumbers = (kx² = kx², ky² = ky², kz² = kz²)

    ky⁺ = reshape(0:Ny-1,       1, Ny, 1)
    kz⁺ = reshape(0:Nz-1,       1, 1, Nz)
    ky⁻ = reshape(0:-1:-(Ny-1), 1, Ny, 1)
    kz⁻ = reshape(0:-1:-(Nz-1), 1, 1, Nz)

    ω_4Ny⁺ = ω.(4Ny, ky⁺) |> CuArray
    ω_4Nz⁺ = ω.(4Nz, kz⁺) |> CuArray
    ω_4Ny⁻ = ω.(4Ny, ky⁻) |> CuArray
    ω_4Nz⁻ = ω.(4Nz, kz⁻) |> CuArray

    # Indices used when we need views with reverse indexing but index N+1 should
    # return a 0. This can't be enforced using views so we just map N+1 to 1,
    # and use masks M_ky and M_kz to enforce that the value at N+1 is 0.
    r_y_inds = [1, collect(Ny:-1:2)...] |> CuArray
    r_z_inds = [1, collect(Nz:-1:2)...] |> CuArray

    # Masks that are useful for writing broadcast operations involving arrays
    # with reversed indices.
    M_ky = ones(1, Ny, 1) |> CuArray
    M_kz = ones(1, 1, Nz) |> CuArray

    M_ky[1] = 0
    M_kz[1] = 0

    constants = (ω_4Ny⁺ = ω_4Ny⁺, ω_4Nz⁺ = ω_4Nz⁺, ω_4Ny⁻ = ω_4Ny⁻, ω_4Nz⁻ = ω_4Nz⁻,
                 r_y_inds = r_y_inds, r_z_inds = r_z_inds,
                 M_ky = M_ky, M_kz = M_kz)

    # Storage for RHS and Fourier coefficients is hard-coded to be Float64
    # because of precision issues with Float32.
    # See https://github.com/climate-machine/Oceananigans.jl/issues/55
    storage1 = zeros(Complex{Float64}, Nx, Ny, Nz) |> CuArray

    # For solving the Poisson equation with PNN boundary conditions, we need a
    # second storage/buffer array to perform the 2D fast cosine transform. Maybe
    # we can get around this but for now we'll just use up a second array.
    storage2 = zeros(Complex{Float64}, Nx, Ny, Nz) |> CuArray

    storage = (storage1 = storage1, storage2 = storage2)

    @debug "Planning transforms for PressureSolver{Channel, GPU}..."
    x_bc, y_bc, z_bc = pressure_bcs.x.left, pressure_bcs.y.left, pressure_bcs.z.left
    FFTx!   = plan_forward_transform(storage1, x_bc, 1)
    FFTyz!  = plan_forward_transform(storage1, z_bc, [2, 3])
    IFFTx!  = plan_backward_transform(storage1, x_bc, 1)
    IFFTyz! = plan_backward_transform(storage1, z_bc, [2, 3])
    @debug "Planning transforms for PressureSolver{Channel, GPU} done!"

    transforms = ( FFTx! =  FFTx!,  FFTyz! =  FFTyz!,
                  IFFTx! = IFFTx!, IFFTyz! = IFFTyz!)

    return PressureSolver(Channel(), GPU(), wavenumbers, storage, transforms, constants)
end

function solve_poisson_equation!(solver::PressureSolver{Channel, GPU}, grid)
    Nx, Ny, Nz, _ = unpack_grid(grid)

    kx², ky², kz² = solver.wavenumbers.kx², solver.wavenumbers.ky², solver.wavenumbers.kz²
    ω_4Ny⁺, ω_4Ny⁻ = solver.constants.ω_4Ny⁺, solver.constants.ω_4Ny⁻
    ω_4Nz⁺, ω_4Nz⁻ = solver.constants.ω_4Nz⁺, solver.constants.ω_4Nz⁻
    r_y_inds, r_z_inds = solver.constants.r_y_inds, solver.constants.r_z_inds
    M_ky, M_kz = solver.constants.M_ky, solver.constants.M_kz

    # We can use the same storage for the RHS and the solution ϕ.
    RHS =  ϕ = solver.storage.storage1

    B = solver.storage.storage2  # Complex buffer storage.

    # Calculate DCTʸᶻ(f) in place using the FFT.
    solver.transforms.FFTyz! * RHS

    RHS⁻ = view(RHS, 1:Nx, r_y_inds, 1:Nz)
    @. B = 2 * real(ω_4Nz⁺ * (ω_4Ny⁺ * RHS + ω_4Ny⁻ * RHS⁻))

    solver.transforms.FFTx! * B # Calculate FFTˣ(f) in place.

    @. B = -B / (kx² + ky² + kz²)

    B[1, 1, 1] = 0  # Setting DC component of the solution (the mean) to be zero.

    solver.transforms.IFFTx! * B  # Calculate IFFTˣ(ϕ̂) in place.

    # Calculate IDCTʸᶻ(ϕ̂) in place using the FFT.
    B⁻⁺ = view(B, 1:Nx, r_y_inds, 1:Nz)
    B⁺⁻ = view(B, 1:Nx, 1:Ny, r_z_inds)
    B⁻⁻ = view(B, 1:Nx, r_y_inds, r_z_inds)

    @. ϕ = 1/4 *  ω_4Ny⁻ * ω_4Nz⁻ * ((B - M_ky * M_kz * B⁻⁻) - im*(M_kz * B⁺⁻ + M_ky * B⁻⁺))

    solver.transforms.IFFTyz! * ϕ

    return nothing
end


text:
struct GSDTrialParam
    max_ia::Int
    M::Int
    warmup::Int
    nevents::Vector{Vector{Int}}
    max_n::Int
    delta::Float64
    bounds::DataFrame
    prior::AbstractPrior

    function GSDTrialParam(max_ia::Int, M::Int, warmup::Int, nevents::Vector{Vector{Int}}, max_n::Int, delta::Float64, bounds::DataFrame, p::AbstractPrior)
        new(max_ia, M, warmup, nevents, max_n, delta, bounds, p)
    end
end

struct GSDTrial
    model::Model
    param::GSDTrialParam
end

check_bounds(pd1, pd2, u1, u2) = (pd1 > u1) || (pd2 > u2)

## simulate a single one-sided k-stage group-sequential design
function run_trial(trial::GSDTrial, id::Int)
    max_ia = trial.param.max_ia
    M = trial.param.M
    warmup = trial.param.warmup

    model = trial.model

    ## initial data
    data = model.generate_data(model.param, model.fparam) ## BiPHWeibull.generate_data(model.param, model.fparam) ##

    bounds = trial.param.bounds
    z = groupby(bounds, :rule)

    R = length(z)
    stopped = zeros(Bool, R)

    res = DataFrame(trial_id=Int64[], rule=String[], ia=Int64[], time=Float64[], rejectSc=Int64[], rejectS=Int64[],
                    beta1=Float64[], beta2=Float64[],
                    sd1=Float64[], sd2=Float64[], pd1=Float64[], pd2=Float64[], cilo1=Float64[], cihi1=Float64[],
                    cilo2=Float64[], cihi2=Float64[], cov1=Bool[], cov2=Bool[], u1=Float64[], u2=Float64[], n=Int64[], nevents=Int64[])

    true_beta1 = model.param.beta[1]
    true_beta2 = model.param.beta[2]

    for j in 1:max_ia
        ia_time, ia_data = select_ia_data(trial.param.nevents[j], data)

        n = length(ia_data.time)

        x = fit(M, ia_data; warmup=warmup, prior=trial.param.prior)

        beta = get_outcome(x)
        beta1 = beta[:,1]
        beta2 = beta[:,2]

        pd1 = mean(beta1 .<= trial.param.delta)
        pd2 = mean(beta2 .<= trial.param.delta)

        pd1_ci_lo = quantile(beta1, 0.025)
        pd1_ci_hi = quantile(beta1, 0.975)
        pd2_ci_lo = quantile(beta2, 0.025)
        pd2_ci_hi = quantile(beta2, 0.975)

        for r in 1:R
            if !stopped[r]
                u1 = z[r][[:a1, :a2]][j][1]
                u2 = z[r][[:b1, :b2]][j][1]

                rejectSc = pd1 > u1
                rejectS = pd2 > u2
                stopped[r] = rejectSc || rejectS

                push!(res, [id, "$(z[r][:rule][1])", j, ia_time, rejectSc, rejectS,
                            mean(beta1), mean(beta2), sqrt(var(beta1)), sqrt(var(beta2)),
                            pd1, pd2, pd1_ci_lo, pd1_ci_hi, pd2_ci_lo, pd2_ci_hi,
                            (pd1_ci_lo <= true_beta1) & (true_beta1 <= pd1_ci_hi),
                            (pd2_ci_lo <= true_beta2) & (true_beta2 <= pd2_ci_hi),
                            u1, u2, n, sum(ia_data.status)])
            end
        end

        if all(stopped)
            break
        end
    end

    res
end

## find optimal bounds maximizing expected utility for a one-sided two-stage group-sequential design given set of forward simulations
function find_gsd_bounds(type1, type2, fsim, fsim0, fsim1; single_stage=false, ngrid=10, h=0.001, utility=gsd_utility)
    a1 = quantile(fsim0.post_pd[:,1,1], 1-type1)
    a2 = quantile(fsim0.post_pd[:,2,1], 1-type1)##*(1-type1))
    b1 = quantile(fsim0.post_pd[:,1,2], 1-type1)
    b2 = quantile(fsim0.post_pd[:,2,2], 1-type1)##*(1-type1))

    @show a1, a2, b1, b2

    if single_stage
        gamma, best_obj = SA.grid_search(repeat([ngrid], inner=2), (a2, b2), (1.0, 1.0),
                                      x -> objective([1,2], [1.0 1.0; x[1] x[2]], type1, type2, fsim, fsim0, fsim1;
                                                         h=h, utility=utility))
        u = [1.0 1.0; gamma[1] gamma[2]]
    else
        gamma, best_obj = SA.grid_search(repeat([ngrid], inner=4), (a1, a2, b1, b2), (1.0, 1.0, 1.0, 1.0),
                                      x -> objective([1,2], reshape(collect(x), 2, 2), type1, type2, fsim, fsim0, fsim1;
                                                         h=h, utility=utility))
        u = reshape(collect(gamma), 2, 2)
    end

    real_t1 = alpha(fsim0.post_pd, u)
    real_t2 = 1-alpha(fsim1.post_pd, u)

    real_t1_Sc = alpha(fsim0.post_pd[:,:,1:1], u[:,1:1])
    real_t2_Sc = 1-alpha(fsim1.post_pd[:,:,1:1], u[:,1:1])

    real_t1_S = alpha(fsim0.post_pd[:,:,2:2], u[:,2:2])
    real_t2_S = 1-alpha(fsim1.post_pd[:,:,2:2], u[:,2:2])

    u, best_obj, real_t1, real_t2, real_t1_Sc, real_t2_Sc, real_t1_S, real_t2_S
end

## forward simulate one-sided two-stage group-sequential design and find optimal bounds
function get_initial_bounds(B::Int, type1, type2, m::Model, priors, tpar::GSDTrialParam; single_stage=false, h=0.001, fsim_only=false, ngrid=10)
    data = m.generate_data(m.param, m.fparam)
    data0 = select_data(0.0, data)
    x = fit(0, data0; prior=priors[1])
    x0 = fit(0, data0; prior=priors[2])
    x1 = fit(0, data0; prior=priors[3])

    J = tpar.max_ia
    ifrac = map(sum, tpar.nevents) ./ sum(tpar.nevents[end])

    fsim, fsim0, fsim1 = fsim_model(2, B, tpar.M, tpar.warmup, m, data0, priors, tpar.delta, tpar.nevents)

    if fsim_only
        fsim, fsim0, fsim1
    else
        find_gsd_bounds(type1, type2, fsim, fsim0, fsim1; single_stage=single_stage, h=h, ngrid=ngrid)
    end
end


text:
# https://github.com/sisl/PGFPlots.jl

# significant contributions by: @pkofod

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

const _pgfplots_linestyles = KW(
    :solid => "solid",
    :dash => "dashed",
    :dot => "dotted",
    :dashdot => "dashdotted",
    :dashdotdot => "dashdotdotted",
)

const _pgfplots_markers = KW(
    :none => "none",
    :cross => "+",
    :xcross => "x",
    :+ => "+",
    :x => "x",
    :utriangle => "triangle*",
    :dtriangle => "triangle*",
    :circle => "*",
    :rect => "square*",
    :star5 => "star",
    :star6 => "asterisk",
    :diamond => "diamond*",
    :pentagon => "pentagon*",
    :hline => "-",
    :vline => "|",
)

const _pgfplots_legend_pos = KW(
    :bottomleft => "south west",
    :bottomright => "south east",
    :topright => "north east",
    :topleft => "north west",
    :outertopright => "outer north east",
)

const _pgf_series_extrastyle = KW(
    :steppre => "const plot mark right",
    :stepmid => "const plot mark mid",
    :steppost => "const plot",
    :sticks => "ycomb",
    :ysticks => "ycomb",
    :xsticks => "xcomb",
)

# PGFPlots uses the anchors to define orientations for example to align left
# one needs to use the right edge as anchor
const _pgf_annotation_halign = KW(:center => "", :left => "right", :right => "left")

const _pgf_framestyles = [:box, :axes, :origin, :zerolines, :grid, :none]
const _pgf_framestyle_defaults = Dict(:semi => :box)
function pgf_framestyle(style::Symbol)
    if style in _pgf_framestyles
        return style
    else
        default_style = get(_pgf_framestyle_defaults, style, :axes)
        @warn(
            "Framestyle :$style is not (yet) supported by the PGFPlots backend. :$default_style was cosen instead."
        )
        default_style
    end
end

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

# takes in color,alpha, and returns color and alpha appropriate for pgf style
function pgf_color(c::Colorant)
    cstr = @sprintf("{rgb,1:red,%.8f;green,%.8f;blue,%.8f}", red(c), green(c), blue(c))
    cstr, alpha(c)
end

function pgf_color(grad::ColorGradient)
    # Can't handle ColorGradient here, fallback to defaults.
    cstr = @sprintf("{rgb,1:red,%.8f;green,%.8f;blue,%.8f}", 0.0, 0.60560316, 0.97868012)
    cstr, 1
end

# Generates a colormap for pgfplots based on a ColorGradient
function pgf_colormap(grad::ColorGradient)
    join(map(grad.colors) do c
        @sprintf("rgb=(%.8f,%.8f,%.8f)", red(c), green(c), blue(c))
    end, ", ")
end

pgf_thickness_scaling(plt::Plot) = plt[:thickness_scaling]
pgf_thickness_scaling(sp::Subplot) = pgf_thickness_scaling(sp.plt)
pgf_thickness_scaling(series) = pgf_thickness_scaling(series[:subplot])

function pgf_fillstyle(plotattributes, i = 1)
    cstr, a = pgf_color(get_fillcolor(plotattributes, i))
    fa = get_fillalpha(plotattributes, i)
    if fa !== nothing
        a = fa
    end
    "fill = $cstr, fill opacity=$a"
end

function pgf_linestyle(linewidth::Real, color, α = 1, linestyle = "solid")
    cstr, a = pgf_color(plot_color(color, α))
    """
    color = $cstr,
    draw opacity = $a,
    line width = $linewidth,
    $(get(_pgfplots_linestyles, linestyle, "solid"))"""
end

function pgf_linestyle(plotattributes, i = 1)
    lw = pgf_thickness_scaling(plotattributes) * get_linewidth(plotattributes, i)
    lc = get_linecolor(plotattributes, i)
    la = get_linealpha(plotattributes, i)
    ls = get_linestyle(plotattributes, i)
    return pgf_linestyle(lw, lc, la, ls)
end

function pgf_font(fontsize, thickness_scaling = 1, font = "\\selectfont")
    fs = fontsize * thickness_scaling
    return string("{\\fontsize{", fs, " pt}{", 1.3fs, " pt}", font, "}")
end

function pgf_marker(plotattributes, i = 1)
    shape = _cycle(plotattributes[:markershape], i)
    cstr, a = pgf_color(
        plot_color(get_markercolor(plotattributes, i), get_markeralpha(plotattributes, i)),
    )
    cstr_stroke, a_stroke = pgf_color(
        plot_color(
            get_markerstrokecolor(plotattributes, i),
            get_markerstrokealpha(plotattributes, i),
        ),
    )
    return string(
        "mark = $(get(_pgfplots_markers, shape, "*")),\n",
        "mark size = $(pgf_thickness_scaling(plotattributes) * 0.5 * _cycle(plotattributes[:markersize], i)),\n",
        plotattributes[:seriestype] == :scatter ? "only marks,\n" : "",
        "mark options = {
            color = $cstr_stroke, draw opacity = $a_stroke,
            fill = $cstr, fill opacity = $a,
            line width = $(pgf_thickness_scaling(plotattributes) * _cycle(plotattributes[:markerstrokewidth], i)),
            rotate = $(shape == :dtriangle ? 180 : 0),
            $(get(_pgfplots_linestyles, _cycle(plotattributes[:markerstrokestyle], i), "solid"))
        }",
    )
end

function pgf_add_annotation!(o, x, y, val, thickness_scaling = 1)
    # Construct the style string.
    # Currently supports color and orientation
    cstr, a = pgf_color(val.font.color)
    push!(
        o,
        PGFPlots.Plots.Node(
            val.str, # Annotation Text
            x,
            y,
            style = """
              $(get(_pgf_annotation_halign,val.font.halign,"")),
              color=$cstr, draw opacity=$(convert(Float16,a)),
              rotate=$(val.font.rotation),
              font=$(pgf_font(val.font.pointsize, thickness_scaling))
              """,
        ),
    )
end

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

function pgf_series(sp::Subplot, series::Series)
    plotattributes = series.plotattributes
    st = plotattributes[:seriestype]
    series_collection = PGFPlots.Plot[]

    # function args
    args = if st == :contour
        plotattributes[:z].surf, plotattributes[:x], plotattributes[:y]
    elseif RecipesPipeline.is3d(st)
        plotattributes[:x], plotattributes[:y], plotattributes[:z]
    elseif st == :straightline
        straightline_data(series)
    elseif st == :shape
        shape_data(series)
    elseif ispolar(sp)
        theta, r = plotattributes[:x], plotattributes[:y]
        rad2deg.(theta), r
    else
        plotattributes[:x], plotattributes[:y]
    end

    # PGFPlots can't handle non-Vector?
    # args = map(a -> if typeof(a) <: AbstractVector && typeof(a) != Vector
    #         collect(a)
    #     else
    #         a
    #     end, args)

    if st in (:contour, :histogram2d)
        style = []
        kw = KW()
        push!(style, pgf_linestyle(plotattributes))
        push!(style, pgf_marker(plotattributes))
        push!(style, "forget plot")

        kw[:style] = join(style, ',')
        func = if st == :histogram2d
            PGFPlots.Histogram2
        else
            kw[:labels] = series[:contour_labels]
            kw[:levels] = series[:levels]
            PGFPlots.Contour
        end
        push!(series_collection, func(args...; kw...))

    else
        # series segments
        segments = iter_segments(series)
        for (i, rng) in enumerate(segments)
            style = []
            kw = KW()
            push!(style, pgf_linestyle(plotattributes, i))
            push!(style, pgf_marker(plotattributes, i))

            if st == :shape
                push!(style, pgf_fillstyle(plotattributes, i))
            end

            # add to legend?
            if i == 1 && sp[:legend] != :none && should_add_to_legend(series)
                if plotattributes[:fillrange] !== nothing
                    push!(style, "forget plot")
                    push!(series_collection, pgf_fill_legend_hack(plotattributes, args))
                else
                    kw[:legendentry] = plotattributes[:label]
                    if st == :shape # || plotattributes[:fillrange] !== nothing
                        push!(style, "area legend")
                    end
                end
            else
                push!(style, "forget plot")
            end

            seg_args = (arg[rng] for arg in args)

            # include additional style, then add to the kw
            if haskey(_pgf_series_extrastyle, st)
                push!(style, _pgf_series_extrastyle[st])
            end
            kw[:style] = join(style, ',')

            # add fillrange
            if series[:fillrange] !== nothing && st != :shape
                push!(
                    series_collection,
                    pgf_fillrange_series(
                        series,
                        i,
                        _cycle(series[:fillrange], rng),
                        seg_args...,
                    ),
                )
            end

            # build/return the series object
            func = if st == :path3d
                PGFPlots.Linear3
            elseif st == :scatter
                PGFPlots.Scatter
            else
                PGFPlots.Linear
            end
            push!(series_collection, func(seg_args...; kw...))
        end
    end
    series_collection
end

function pgf_fillrange_series(series, i, fillrange, args...)
    st = series[:seriestype]
    style = []
    kw = KW()
    push!(style, "line width = 0")
    push!(style, "draw opacity = 0")
    push!(style, pgf_fillstyle(series, i))
    push!(style, pgf_marker(series, i))
    push!(style, "forget plot")
    if haskey(_pgf_series_extrastyle, st)
        push!(style, _pgf_series_extrastyle[st])
    end
    kw[:style] = join(style, ',')
    func = RecipesPipeline.is3d(series) ? PGFPlots.Linear3 : PGFPlots.Linear
    return func(pgf_fillrange_args(fillrange, args...)...; kw...)
end

function pgf_fillrange_args(fillrange, x, y)
    n = length(x)
    x_fill = [x; x[n:-1:1]; x[1]]
    y_fill = [y; _cycle(fillrange, n:-1:1); y[1]]
    return x_fill, y_fill
end

function pgf_fillrange_args(fillrange, x, y, z)
    n = length(x)
    x_fill = [x; x[n:-1:1]; x[1]]
    y_fill = [y; y[n:-1:1]; x[1]]
    z_fill = [z; _cycle(fillrange, n:-1:1); z[1]]
    return x_fill, y_fill, z_fill
end

function pgf_fill_legend_hack(plotattributes, args)
    style = []
    kw = KW()
    push!(style, pgf_linestyle(plotattributes, 1))
    push!(style, pgf_marker(plotattributes, 1))
    push!(style, pgf_fillstyle(plotattributes, 1))
    push!(style, "area legend")
    kw[:legendentry] = plotattributes[:label]
    kw[:style] = join(style, ',')
    st = plotattributes[:seriestype]
    func = if st == :path3d
        PGFPlots.Linear3
    elseif st == :scatter
        PGFPlots.Scatter
    else
        PGFPlots.Linear
    end
    return func(([arg[1]] for arg in args)...; kw...)
end

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

function pgf_axis(sp::Subplot, letter)
    axis = sp[Symbol(letter, :axis)]
    style = []
    kw = KW()

    # turn off scaled ticks
    push!(style, "scaled $(letter) ticks = false")

    # set to supported framestyle
    framestyle = pgf_framestyle(sp[:framestyle])

    # axis guide
    kw[Symbol(letter, :label)] = axis[:guide]

    # axis label position
    labelpos = ""
    if letter == :x && axis[:guide_position] == :top
        labelpos = "at={(0.5,1)},above,"
    elseif letter == :y && axis[:guide_position] == :right
        labelpos = "at={(1,0.5)},below,"
    end

    # Add label font
    cstr, α = pgf_color(plot_color(axis[:guidefontcolor]))
    push!(
        style,
        string(
            letter,
            "label style = {",
            labelpos,
            "font = ",
            pgf_font(axis[:guidefontsize], pgf_thickness_scaling(sp)),
            ", color = ",
            cstr,
            ", draw opacity = ",
            α,
            ", rotate = ",
            axis[:guidefontrotation],
            "}",
        ),
    )

    # flip/reverse?
    axis[:flip] && push!(style, "$letter dir=reverse")

    # scale
    scale = axis[:scale]
    if scale in (:log2, :ln, :log10)
        kw[Symbol(letter, :mode)] = "log"
        scale == :ln || push!(style, "log basis $letter=$(scale == :log2 ? 2 : 10)")
    end

    # ticks on or off
    if axis[:ticks] in (nothing, false, :none) || framestyle == :none
        push!(style, "$(letter)majorticks=false")
    end

    # grid on or off
    if axis[:grid] && framestyle != :none
        push!(style, "$(letter)majorgrids = true")
    else
        push!(style, "$(letter)majorgrids = false")
    end

    # limits
    # TODO: support zlims
    if letter != :z
        lims =
            ispolar(sp) && letter == :x ? rad2deg.(axis_limits(sp, :x)) :
            axis_limits(sp, letter)
        kw[Symbol(letter, :min)] = lims[1]
        kw[Symbol(letter, :max)] = lims[2]
    end

    if !(axis[:ticks] in (nothing, false, :none, :native)) && framestyle != :none
        ticks = get_ticks(sp, axis)
        #pgf plot ignores ticks with angle below 90 when xmin = 90 so shift values
        tick_values =
            ispolar(sp) && letter == :x ? [rad2deg.(ticks[1])[3:end]..., 360, 405] :
            ticks[1]
        push!(style, string(letter, "tick = {", join(tick_values, ","), "}"))
        if axis[:showaxis] && axis[:scale] in (:ln, :log2, :log10) && axis[:ticks] == :auto
            # wrap the power part of label with }
            tick_labels = Vector{String}(undef, length(ticks[2]))
            for (i, label) in enumerate(ticks[2])
                base, power = split(label, "^")
                power = string("{", power, "}")
                tick_labels[i] = string(base, "^", power)
            end
            push!(
                style,
                string(letter, "ticklabels = {\$", join(tick_labels, "\$,\$"), "\$}"),
            )
        elseif axis[:showaxis]
            tick_labels =
                ispolar(sp) && letter == :x ? [ticks[2][3:end]..., "0", "45"] : ticks[2]
            if axis[:formatter] in (:scientific, :auto)
                tick_labels = string.("\$", convert_sci_unicode.(tick_labels), "\$")
                tick_labels = replace.(tick_labels, Ref("×" => "\\times"))
            end
            push!(style, string(letter, "ticklabels = {", join(tick_labels, ","), "}"))
        else
            push!(style, string(letter, "ticklabels = {}"))
        end
        push!(
            style,
            string(
                letter,
                "tick align = ",
                (axis[:tick_direction] == :out ? "outside" : "inside"),
            ),
        )
        cstr, α = pgf_color(plot_color(axis[:tickfontcolor]))
        push!(
            style,
            string(
                letter,
                "ticklabel style = {font = ",
                pgf_font(axis[:tickfontsize], pgf_thickness_scaling(sp)),
                ", color = ",
                cstr,
                ", draw opacity = ",
                α,
                ", rotate = ",
                axis[:tickfontrotation],
                "}",
            ),
        )
        push!(
            style,
            string(
                letter,
                " grid style = {",
                pgf_linestyle(
                    pgf_thickness_scaling(sp) * axis[:gridlinewidth],
                    axis[:foreground_color_grid],
                    axis[:gridalpha],
                    axis[:gridstyle],
                ),
                "}",
            ),
        )
    end

    # framestyle
    if framestyle in (:axes, :origin)
        axispos = framestyle == :axes ? "left" : "middle"
        if axis[:draw_arrow]
            push!(style, string("axis ", letter, " line = ", axispos))
        else
            # the * after line disables the arrow at the axis
            push!(style, string("axis ", letter, " line* = ", axispos))
        end
    end

    if framestyle == :zerolines
        push!(style, string("extra ", letter, " ticks = 0"))
        push!(style, string("extra ", letter, " tick labels = "))
        push!(
            style,
            string(
                "extra ",
                letter,
                " tick style = {grid = major, major grid style = {",
                pgf_linestyle(
                    pgf_thickness_scaling(sp),
                    axis[:foreground_color_border],
                    1.0,
                ),
                "}}",
            ),
        )
    end

    if !axis[:showaxis]
        push!(style, "separate axis lines")
    end
    if !axis[:showaxis] || framestyle in (:zerolines, :grid, :none)
        push!(style, string(letter, " axis line style = {draw opacity = 0}"))
    else
        push!(
            style,
            string(
                letter,
                " axis line style = {",
                pgf_linestyle(
                    pgf_thickness_scaling(sp),
                    axis[:foreground_color_border],
                    1.0,
                ),
                "}",
            ),
        )
    end

    # return the style list and KW args
    style, kw
end

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

function _update_plot_object(plt::Plot{PGFPlotsBackend})
    plt.o = PGFPlots.Axis[]
    # Obtain the total height of the plot by extracting the maximal bottom
    # coordinate from the bounding box.
    total_height = bottom(bbox(plt.layout))

    for sp in plt.subplots
        # first build the PGFPlots.Axis object
        style = ["unbounded coords=jump"]
        kw = KW()

        # add to style/kw for each axis
        for letter in (:x, :y, :z)
            if letter != :z || RecipesPipeline.is3d(sp)
                axisstyle, axiskw = pgf_axis(sp, letter)
                append!(style, axisstyle)
                merge!(kw, axiskw)
            end
        end

        # bounding box values are in mm
        # note: bb origin is top-left, pgf is bottom-left
        # A round on 2 decimal places should be enough precision for 300 dpi
        # plots.
        bb = bbox(sp)
        push!(
            style,
            """
    xshift = $(left(bb).value)mm,
    yshift = $(round((total_height - (bottom(bb))).value, digits=2))mm,
    axis background/.style={fill=$(pgf_color(sp[:background_color_inside])[1])}
""",
        )
        kw[:width] = "$(width(bb).value)mm"
        kw[:height] = "$(height(bb).value)mm"

        if sp[:title] != ""
            kw[:title] = "$(sp[:title])"
            cstr, α = pgf_color(plot_color(sp[:titlefontcolor]))
            push!(
                style,
                string(
                    "title style = {font = ",
                    pgf_font(sp[:titlefontsize], pgf_thickness_scaling(sp)),
                    ", color = ",
                    cstr,
                    ", draw opacity = ",
                    α,
                    ", rotate = ",
                    sp[:titlefontrotation],
                    "}",
                ),
            )
        end

        if get_aspect_ratio(sp) in (1, :equal)
            kw[:axisEqual] = "true"
        end

        legpos = sp[:legend]
        if haskey(_pgfplots_legend_pos, legpos)
            kw[:legendPos] = _pgfplots_legend_pos[legpos]
        end
        cstr, bg_alpha = pgf_color(plot_color(sp[:background_color_legend]))
        fg_alpha = alpha(plot_color(sp[:foreground_color_legend]))

        push!(
            style,
            string(
                "legend style = {",
                pgf_linestyle(
                    pgf_thickness_scaling(sp),
                    sp[:foreground_color_legend],
                    fg_alpha,
                    "solid",
                ),
                ",",
                "fill = $cstr,",
                "fill opacity = $bg_alpha,",
                "text opacity = $(alpha(plot_color(sp[:legendfontcolor]))),",
                "font = ",
                pgf_font(sp[:legendfontsize], pgf_thickness_scaling(sp)),
                "}",
            ),
        )

        if any(s[:seriestype] == :contour for s in series_list(sp))
            kw[:view] = "{0}{90}"
            kw[:colorbar] = !(sp[:colorbar] in (:none, :off, :hide, false))
        elseif RecipesPipeline.is3d(sp)
            azim, elev = sp[:camera]
            kw[:view] = "{$(azim)}{$(elev)}"
        end

        axisf = PGFPlots.Axis
        if sp[:projection] == :polar
            axisf = PGFPlots.PolarAxis
            #make radial axis vertical
            kw[:xmin] = 90
            kw[:xmax] = 450
        end

        # Search series for any gradient. In case one series uses a gradient set
        # the colorbar and colomap.
        # The reasoning behind doing this on the axis level is that pgfplots
        # colorbar seems to only works on axis level and needs the proper colormap for
        # correctly displaying it.
        # It's also possible to assign the colormap to the series itself but
        # then the colormap needs to be added twice, once for the axis and once for the
        # series.
        # As it is likely that all series within the same axis use the same
        # colormap this should not cause any problem.
        for series in series_list(sp)
            for col in (:markercolor, :fillcolor, :linecolor)
                if typeof(series.plotattributes[col]) == ColorGradient
                    push!(
                        style,
                        "colormap={plots}{$(pgf_colormap(series.plotattributes[col]))}",
                    )

                    if sp[:colorbar] == :none
                        kw[:colorbar] = "false"
                    else
                        kw[:colorbar] = "true"
                    end
                    # goto is needed to break out of col and series for
                    @goto colorbar_end
                end
            end
        end
        @label colorbar_end

        push!(style, "colorbar style={title=$(sp[:colorbar_title])}")
        o = axisf(; style = join(style, ","), kw...)

        # add the series object to the PGFPlots.Axis
        for series in series_list(sp)
            push!.(Ref(o), pgf_series(sp, series))

            # add series annotations
            anns = series[:series_annotations]
            for (xi, yi, str, fnt) in EachAnn(anns, series[:x], series[:y])
                pgf_add_annotation!(
                    o,
                    xi,
                    yi,
                    PlotText(str, fnt),
                    pgf_thickness_scaling(series),
                )
            end
        end

        # add the annotations
        for ann in sp[:annotations]
            pgf_add_annotation!(
                o,
                locate_annotation(sp, ann...)...,
                pgf_thickness_scaling(sp),
            )
        end

        # add the PGFPlots.Axis to the list
        push!(plt.o, o)
    end
end

function _show(io::IO, mime::MIME"image/svg+xml", plt::Plot{PGFPlotsBackend})
    show(io, mime, plt.o)
end

function _show(io::IO, mime::MIME"application/pdf", plt::Plot{PGFPlotsBackend})
    # prepare the object
    pgfplt = PGFPlots.plot(plt.o)

    # save a pdf
    fn = tempname() * ".pdf"
    PGFPlots.save(PGFPlots.PDF(fn), pgfplt)

    # read it into io
    write(io, read(open(fn), String))

    # cleanup
    PGFPlots.cleanup(plt.o)
end

function _show(io::IO, mime::MIME"application/x-tex", plt::Plot{PGFPlotsBackend})
    fn = tempname() * ".tex"
    PGFPlots.save(
        fn,
        backend_object(plt),
        include_preamble = plt.attr[:tex_output_standalone],
    )
    write(io, read(open(fn), String))
end

function _display(plt::Plot{PGFPlotsBackend})
    # prepare the object
    pgfplt = PGFPlots.plot(plt.o)

    # save an svg
    fn = string(tempname(), ".svg")
    PGFPlots.save(PGFPlots.SVG(fn), pgfplt)

    # show it
    open_browser_window(fn)

    # cleanup
    PGFPlots.cleanup(plt.o)
end


text:
using SparseDiffTools
using Test

@testset "Contraction Algorithm" begin include("test_contraction_algo.jl") end
@testset "Greedy D1 Algorithm" begin include("test_greedy_d1.jl") end


text:
using Documenter, GeoPhyInv

#makedocs(modules=[GeoPhyInv],
#	         doctest=true)
#makedocs()

makedocs(
	 format = Documenter.HTML(
		  prettyurls = get(ENV, "CI", nothing) == "true"),
   sitename = "Toolbox for Geophysical Inversion",
   pages = ["Home" => "index.md",
	    "SeisForwExpt" => Any[
				  "Introduction" => "Fdtd/intro.md",
	#			  "Basic usage" => "Fdtd/reuse_expt.md",
				  "Generate snaps" => "Fdtd/create_snaps.md",
				  ],
	    "PoissonExpt" => Any[
				  "Introduction" => "Poisson/intro.md",
				 "Record data" => "Poisson/forw.md",
				 "Born map" => "Poisson/test_born.md",
				 ],
	    ]
#	    "Seismic Born Modeling" => "FWI/born_map.md",
#	    "Seismic Full Waveform Inversion" => "FWI/gradient_accuracy.md",
#    modules = []
)

deploydocs(
    repo   = "github.com/pawbz/GeoPhyInv.jl.git",
)
#=
deploydocs(
    deps   = Deps.pip("mkdocs", "python-markdown-math"),
    repo   = "github.com/pawbz/GeoPhyInv.jl.git",
    julia  = "1.0",
    osname = "linux",
    target = "build",
    make   = nothing
)
=#


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct Bus <: Topology
        number::Int64
        name::String
        bustype::Union{Nothing, BusTypes.BusType}
        angle::Union{Nothing, Float64}
        voltage::Union{Nothing, Float64}
        voltagelimits::Union{Nothing, Min_Max}
        basevoltage::Union{Nothing, Float64}
        area::Union{Nothing, Area}
        load_zone::Union{Nothing, LoadZone}
        ext::Dict{String, Any}
        internal::InfrastructureSystemsInternal
    end

A power-system bus.

# Arguments
- `number::Int64`: number associated with the bus
- `name::String`: the name of the bus
- `bustype::Union{Nothing, BusTypes.BusType}`: bus type
- `angle::Union{Nothing, Float64}`: angle of the bus in radians, validation range: (-1.571, 1.571), action if invalid: error
- `voltage::Union{Nothing, Float64}`: voltage as a multiple of basevoltage, validation range: voltagelimits, action if invalid: warn
- `voltagelimits::Union{Nothing, Min_Max}`: limits on the voltage variation as multiples of basevoltage
- `basevoltage::Union{Nothing, Float64}`: the base voltage in kV, validation range: (0, nothing), action if invalid: error
- `area::Union{Nothing, Area}`: the area containing the bus
- `load_zone::Union{Nothing, LoadZone}`: the load zone containing the bus
- `ext::Dict{String, Any}`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct Bus <: Topology
    "number associated with the bus"
    number::Int64
    "the name of the bus"
    name::String
    "bus type"
    bustype::Union{Nothing, BusTypes.BusType}
    "angle of the bus in radians"
    angle::Union{Nothing, Float64}
    "voltage as a multiple of basevoltage"
    voltage::Union{Nothing, Float64}
    "limits on the voltage variation as multiples of basevoltage"
    voltagelimits::Union{Nothing, Min_Max}
    "the base voltage in kV"
    basevoltage::Union{Nothing, Float64}
    "the area containing the bus"
    area::Union{Nothing, Area}
    "the load zone containing the bus"
    load_zone::Union{Nothing, LoadZone}
    ext::Dict{String, Any}
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal

    function Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, internal, )
        (number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, internal, ) = CheckBusParams(
            number,
            name,
            bustype,
            angle,
            voltage,
            voltagelimits,
            basevoltage,
            area,
            load_zone,
            ext,
            internal,
        )
        new(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, internal, )
    end
end

function Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext=Dict{String, Any}(), )
    Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, InfrastructureSystemsInternal(), )
end

function Bus(; number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext=Dict{String, Any}(), )
    Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, )
end

# Constructor for demo purposes; non-functional.
function Bus(::Nothing)
    Bus(;
        number=0,
        name="init",
        bustype=nothing,
        angle=0.0,
        voltage=0.0,
        voltagelimits=(min=0.0, max=0.0),
        basevoltage=nothing,
        area=Area(nothing),
        load_zone=LoadZone(nothing),
        ext=Dict{String, Any}(),
    )
end

"""Get Bus number."""
get_number(value::Bus) = value.number

InfrastructureSystems.get_name(value::Bus) = value.name
"""Get Bus bustype."""
get_bustype(value::Bus) = value.bustype
"""Get Bus angle."""
get_angle(value::Bus) = value.angle
"""Get Bus voltage."""
get_voltage(value::Bus) = value.voltage
"""Get Bus voltagelimits."""
get_voltagelimits(value::Bus) = value.voltagelimits
"""Get Bus basevoltage."""
get_basevoltage(value::Bus) = value.basevoltage
"""Get Bus area."""
get_area(value::Bus) = value.area
"""Get Bus load_zone."""
get_load_zone(value::Bus) = value.load_zone
"""Get Bus ext."""
get_ext(value::Bus) = value.ext
"""Get Bus internal."""
get_internal(value::Bus) = value.internal


text:
"""
    tasklearner(task, traindata, validdata[; callbacks=[], kwargs...]) -> Learner
    tasklearner(task, data; pctgval = 0.2, kwargs...)

Create a [`Learner`](#) to train a model for learning task `task` using
`data`.

## Keyword arguments

- `callbacks = []`: [`Callback`](#)s to use during training.
- `batchsize = 16`: Batch size used for the training data loader.
- `backbone = nothing`: Backbone model to construct task-specific model from using
   [`taskmodel`](#)`(task, backbone)`.
- `model = nothing`: Complete model to use. If given, the `backbone` argument is ignored.
- `optimizer = ADAM()`: Optimizer passed to `Learner`.
- `lossfn = `[`tasklossfn`](#)`(task)`: Loss function passed to `Learner`.

Any other keyword arguments will be passed to [`taskdataloaders`](#).

## Examples

Full example:

```julia
data, blocks = loaddataset("imagenette2-160", (Image, Label))
task = ImageClassificationSingle(blocks)
learner = tasklearner(task, data)
fitonecycle!(learner, 10)
```

Custom training and validation split:

```julia
learner = tasklearner(task, traindata, validdata)
```

Using callbacks:

```julia
learner = tasklearner(task, data; callbacks=[
    ToGPU(), Checkpointer(), LogMetrics(TensorboardBackend())
])
```
"""
function tasklearner(
        task::LearningTask,
        traindata,
        validdata;
        backbone=nothing,
        model=nothing,
        callbacks=[],
        pctgval=0.2,
        batchsize=16,
        optimizer=ADAM(),
        lossfn=tasklossfn(task),
        kwargs...,
    )
    if isnothing(model)
        model = isnothing(backbone) ? taskmodel(task) : taskmodel(task, backbone)
    end
    dls = taskdataloaders(traindata, validdata, task, batchsize; kwargs...)
    return Learner(model, dls, optimizer, lossfn, callbacks...)
end

function tasklearner(task, data; pctgval=0.2, kwargs...)
    traindata, validdata = splitobs(data, at=1 - pctgval)
    return tasklearner(task, traindata, validdata; kwargs...)
end

"""
    getbatch(learner[; validation = false, n = nothing])

Get a batch of data from `learner`. Take a batch of training data by default
or validation data if `validation = true`. If `n` take only the first
`n` samples from the batch.

"""
function getbatch(learner; context = Training(), n = nothing)
    dl = context == Validation() ? learner.data.validation : learner.data.training
    batch = first(dl)
    b = min(isnothing(n) ? Inf : n, Datasets.batchsize(batch))
    batch = MLUtils.batch([s for (s, _) in zip(Datasets.unbatch(batch), 1:b)])
    return batch
end

# ## Tests

@testset "getbatch" begin
    batch = rand(1, 10), rand(1, 10)
    learner = Learner(identity, ([batch], [batch]), nothing, nothing)
    @test size.(getbatch(learner)) == ((1,10), (1, 10))
    @test size.(getbatch(learner, n=4)) == ((1,4), (1,4))
end

@testset "tasklearner" begin
    task = SupervisedTask((Label(1:2), Label(1:2)), (OneHot(),))
    data = (rand(1:2, 1000), rand(1:2, 1000))
    @test_nowarn learner = tasklearner(task, data, model=identity)

    @testset "batch sizes" begin
        learner = tasklearner(task, data, model=identity, batchsize=100)
        @test length(learner.data.training) == 8
        @test length(learner.data.validation) == 1

        learner = tasklearner(task, data, model=identity, pctgval=0.4, batchsize=100)
        @test length(learner.data.training) == 6
        @test length(learner.data.validation) == 2

        learner = tasklearner(task, data, model=identity, batchsize=100, validbsfactor=1)
        @test length(learner.data.training) == 8
        @test length(learner.data.validation) == 2
    end

    @testset "callbacks" begin
        learner = tasklearner(
            task, data, model=identity,
            callbacks=[ToGPU(), Checkpointer(mktempdir())])
        @test !isnothing(FluxTraining.getcallback(learner, Checkpointer))

    end
end


text:
using Raylib
using StaticArrays

struct BallGame <: Game
  pos::MVector{2,Float32}
end

function Raylib.update!(g::BallGame)
  if is_key_down(Keys.RIGHT)
    g.pos[1] += 2.0
  elseif is_key_down(Keys.LEFT)
    g.pos[1] -= 2.0
  elseif is_key_down(Keys.UP)
    g.pos[2] -= 2.0
  elseif is_key_down(Keys.DOWN)
    g.pos[2] += 2.0
  end
end

function Raylib.draw(g::BallGame)
  clear_background(WHITE)
  draw_text("move the ball with arrow keys", 10, 10, 20, LIGHTGREY)
  draw_circle(g.pos, 50, RayColor(0,0,0, 255))
end

run(BallGame(SA_F32[400, 225]))


text:
"""
```julia
vstd(A; dims=:, mean=nothing, corrected=true, multithreaded=:auto)
```
Compute the variance of all elements in `A`, optionally over dimensions specified by `dims`.
As `Statistics.var`, but vectorized and (optionally) multithreaded.

A precomputed `mean` may optionally be provided, which results in a somewhat faster
calculation. If `corrected` is `true`, then _Bessel's correction_ is applied, such
that the sum is divided by `n-1` rather than `n`.

## Examples
```julia
julia> using VectorizedStatistics

julia> A = [1 2; 3 4]
2×2 Matrix{Int64}:
 1  2
 3  4

julia> vstd(A, dims=1)
1×2 Matrix{Float64}:
 1.41421  1.41421

julia> vstd(A, dims=2)
2×1 Matrix{Float64}:
 0.7071067811865476
 0.7071067811865476
```
"""
function vstd(A; dims=:, mean=nothing, corrected=true, multithreaded=:auto)
    if (multithreaded===:auto && length(A) > 4095) || multithreaded===true
        _vtstd(mean, corrected, A, dims)
    else
        _vstd(mean, corrected, A, dims)
    end
end
export vstd

sqrt!(x::Number) = sqrt(x)
function sqrt!(A::AbstractArray)
    @turbo for i ∈ eachindex(A)
        A[i] = sqrt(A[i])
    end
    return A
end

_vstd(mean, corrected, A, dims) = sqrt!(_vvar(mean, corrected, A, dims))

sqrtt!(x::Number) = sqrt(x)
function sqrtt!(A::AbstractArray)
    @tturbo for i ∈ eachindex(A)
        A[i] = sqrt(A[i])
    end
    return A
end

_vtstd(mean, corrected, A, dims) = sqrtt!(_vtvar(mean, corrected, A, dims))


text:
"""
    dist_extrema(values::T...) where {T <: Number}
Computes the distance between extrema in an ordered set.
"""
function o_dist_extrema(values)
    m, M = extrema(values)
    return Float64(M - m)
end


text:
module TestGraphs
using Test

import LightGraphs
using Catlab.CategoricalAlgebra.Graphs

# Graphs
########

g = Graph()
@test keys(g.indices) == (:src,:tgt)
add_vertex!(g)
add_vertices!(g, 2)
@test nv(g) == 3
@test ne(g) == 0

add_edge!(g, 1, 2)
add_edge!(g, 2, 3)
@test ne(g) == 2
@test ne(g, 1, 2) == 1
@test has_edge(g, 1, 2)
@test !has_edge(g, 1, 3)
@test outneighbors(g, 2) == [3]
@test inneighbors(g, 2) == [1]
@test collect(all_neighbors(g, 2)) == [1,3]

add_edge!(g, 1, 2)
@test ne(g) == 3
@test ne(g, 1, 2) == 2
@test collect(edges(g, 1, 2)) == [1,3]
@test outneighbors(g, 1) == [2,2]
@test inneighbors(g, 1) == []
@test LightGraphs.DiGraph(g) == LightGraphs.path_digraph(3)

g = Graph()
add_vertices!(g, 4)
add_edges!(g, [1,2,3], [2,3,4])
@test LightGraphs.DiGraph(g) == LightGraphs.path_digraph(4)

# Symmetric graphs
##################

g = SymmetricGraph()
@test keys(g.indices) == (:src,)

add_vertices!(g, 3)
@test nv(g) == 3
@test ne(g) == 0

add_edge!(g, 1, 2)
add_edge!(g, 2, 3)
@test ne(g) == 4
@test collect(edges(g, 1, 2)) == [1]
@test neighbors(g, 1) == [2]
@test neighbors(g, 2) == [1,3]
@test neighbors(g, 3) == [2]
@test LightGraphs.Graph(g) == LightGraphs.path_graph(3)

g = SymmetricGraph()
add_vertices!(g, 4)
add_edges!(g, [1,2,3], [2,3,4])
lg = LightGraphs.DiGraph(4)
map((src, tgt) -> add_edge!(lg, src, tgt), [1,2,3,2,3,4], [2,3,4,1,2,3])
@test LightGraphs.DiGraph(g) == lg

# Property graphs
#################

g = PropertyGraph{String}()
add_vertex!(g, a="foo", b="bar")
@test nv(g) == 1
@test vprops(g, 1) == Dict(:a => "foo", :b => "bar")
@test get_vprop(g, 1, :b) == "bar"
set_vprop!(g, 1, :b, "baz")
@test get_vprop(g, 1, :b) == "baz"
add_vertices!(g, 2)
@test nv(g) == 3
@test vprops(g, 3) == Dict()

add_edge!(g, 1, 2, c="car")
@test ne(g) == 1
@test src(g, 1) == 1
@test tgt(g, 1) == 2
@test eprops(g, 1) == Dict(:c => "car")
@test get_eprop(g, 1, :c) == "car"
set_eprop!(g, 1, :c, "cat")
@test get_eprop(g, 1, :c) == "cat"
add_edges!(g, [2,2], [3,3])
set_eprop!(g, 2, :d, "dog")
set_eprop!(g, 3, :d, "dig")
@test src(g, [2,3]) == [2,2]
@test tgt(g, [2,3]) == [3,3]
@test (get_eprop(g, 2, :d), get_eprop(g, 3, :d)) == ("dog", "dig")

# Symmetric property graphs
###########################

g = SymmetricPropertyGraph{String}()
add_vertex!(g, a="aardvark")
@test vprops(g, 1) == Dict(:a => "aardvark")
add_vertex!(g)
@test vprops(g, 2) == Dict()

add_edge!(g, 1, 2, c="car")
@test ne(g) == 2
@test eprops(g, 1) == Dict(:c => "car")
@test eprops(g, 1) === eprops(g, 2)

end


text:
using Documenter, PackageCompiler

makedocs(
    format = Documenter.HTML(
        prettyurls = "deploy" in ARGS,
    ),
    sitename = "PackageCompiler",
    pages = [
        "Home" => "index.md",

        "Manual" => [
            "sysimages.md"
            "libs.md"
            "apps.md"
        ],

        "Examples" => [
            "examples/ohmyrepl.md",
            "examples/plots.md",
        ],

        "PackageCompiler - the manual way" => [
            "devdocs/intro.md",
            "devdocs/sysimages_part_1.md",
            "devdocs/binaries_part_2.md",
            "devdocs/relocatable_part_3.md",
        ],

        "References" => "refs.md",
        "Upgrade notes" => "upgrade.md",
    ]
)

deploydocs(
    repo = "github.com/JuliaLang/PackageCompiler.jl.git",
    push_preview = true,
)


text:
using MLDatasets: MNIST

struct MNISTDataset{T, D} <: ImageDataset{T, D}
    X
    y
    Y
    Xt
    yt
    Yt
end

n_display(::MNISTDataset) = 100

function MNISTDataset(
    n_data::Int=60_000;
    seed::Int=1,
    test_ratio=1/6,
    n_test::Int=ratio2num(n_data, test_ratio),
    is_flatten::Bool=false,
    alpha::T=1f-6,
    is_link::Bool=true,
    rng=MersenneTwister(seed),
) where {T<:AbstractFloat}
    if n_data > 60_000
        @warn "Attempts to access more than 60,000 train points in MNIST; clamped to 60,000."
        n_data = 60_000
    end
    if n_test > 10_000
        @warn "Attempts to access more than 10,000 test points in MNIST; clamped to 10,000."
        n_test = 10_000
    end
    X, y, Y, Xt, yt, Yt = get_image_data(
        MNIST, n_data, n_test, is_flatten, alpha, is_link, 10, (2, 1, 3); rng=rng
    )
    @info "Oh you just get the MNIST dataset" n_data=n_data n_test=n_test is_flatten=is_flatten alpha=alpha is_link=is_link
    return MNISTDataset{Val{is_link}, 784}(X, y, Y, Xt, yt, Yt)
end


text:
@testset "Choi with $(factory.constructor)" for factory in sdp_factories
    @polyvar x y z

    m = SOSModel(factory)

    C = [x^2+2y^2 -x*y -x*z;
         -x*y y^2+2z^2 -y*z;
         -x*z -y*z z^2+2x^2]

    @SDconstraint m C >= 0

    JuMP.optimize!(m)
    @test JuMP.dual_status(m) == MOI.InfeasibilityCertificate
end


text:
using Distributed
@everywhere using Pkg
@everywhere Pkg.activate(pwd())
Pkg.instantiate()
using MyMCMCSampler
using StatsBase
using MCMCDiagnostics
using Mamba
using StanSample
using ArgParse
using JLD2, FileIO
using ProgressMeter
ProjDir = @__DIR__
tmpdir= ProjDir * "/tmp"

sa = ArgParseSettings()
@add_arg_table sa begin
    "--outdir"
        help = "folder to store dataset file"
        arg_type = String
        default = "data/tmp/"
    "--n"
        help = "number of variable"
        arg_type = Int
        default = 10
    "--ndata"
        help = "number of different graphs per hyper-parameter"
        arg_type = Int
        default = 1
    "--beta"
        help = "inverse temperature"
        arg_type = Float64
        default = 0.3
    "--alpha"
        help = "relative strength of the 4th order term"
        arg_type = Float64
        default = 1.0
    "--scale"
        help = "scaling of singleton and interaction coefficients"
        arg_type = Float64
        default = 1.0
end

args = parse_args(ARGS, sa)
println(args)
n = args["n"]
ndata = args["ndata"]
outdir = args["outdir"]

stan_program = "
data {
    int N; // number of variables
    int M2;  // number of pariwise potentials
    int M3; // number of third-order potentials
    vector[N] s[3];  // singleton potentials, up to third order
    vector[M2] p; // pairwise potentials
    vector[M3] t;  //  third-order potentials
    int pid[2,M2]; // indices for pairwise potentials
    int tid[3,M3]; // indices for third-order potentials
    real alpha;  // coeff for 4th order base measure
    real beta;  // inverse temp
}
parameters {
    vector[N] x;  // variables
}

model{
    // singleton potentials
    target += - beta * ( s[1].*x + s[2].*x.*x + s[3].*x.*x.*x ) ;
    // pairwise potentials
    target += - beta * (p .* x[pid[1]] .* x[pid[2]]);
    // third-order potentials
    target += - beta * (t .* x[tid[1]] .* x[tid[2]] .* x[tid[3]]);
    // base measure
    // target += - beta * (alpha * pow(sum(x.*x),2));
    target += - beta * alpha * sum(pow(x,4));
}

";

# ks = range(2.0,n,length=30)
# p0s = range(0.0,1.0,length=30)
k2range = Uniform(2.0,n-1)
k3range = Uniform(2.0, floor((n-1)*(n-2)/6))
prange = Uniform(0.0, 1.0)

# alpha =  0.2
# beta = 0.05
# beta = 0.3
beta0 = args["beta"]
alpha0 = args["alpha"]
scale = args["scale"]
beta = beta0
alpha = alpha0
gamma = 1/ scale
gamma1 = 1.0

counter = 0
prog = Progress(ndata)
# for k in ks,p0 in p0s, i in 1:ndata
for i in 1:ndata
    k2 = rand(k2range)
    k3 = rand(k3range)
    p0 = rand(prange)
    global counter
    counter += 1
    valid = false
    sm = nothing
    sm0 = nothing
    data = nothing
    data0 = nothing
    g = nothing
    while !valid
        # g = random_ws_factor_graph(n,k,p0;β=3.0,degrees=[2,3])
        # g = random_ws_factor_graph(n,k2,k3,p0;β=1.0, degrees=[2,3],singleton_dist=:randn)  # old, used for training set
        g = random_ws_factor_graph(n,k2,k3,p0;γ=gamma,γ1=gamma1, degrees=[2,3],singleton_dist=:randn )
        idx2 = findall(x->x==2,g.edge_type)
        idx3 = findall(x->x==3,g.edge_type)
        M2 = length(idx2)
        M3 = length(idx3)
        # s = permutedims(hcat(g.node_attr...),[2,1])
        s = hcat(g.node_attr...)  # (3,N)
        # s = [s[i,:] for i in 1:size(s,1)]
        p = vcat(g.edge_attr[idx2]...) |> Array{Float64}
        t = vcat(g.edge_attr[idx3]...) |> Array{Float64}
        pid = hcat(g.edge_index[idx2]...) |> Array{Int}
        tid = hcat(g.edge_index[idx3]...) |> Array{Int}
        stan_data = Dict(
            "N" => n,
            "M2" =>M2,
            "M3" =>M3,
            "s" => s,
            "p" =>p,
            "t"=>t,
            "pid"=>pid,
            "tid"=>tid,
            "alpha"=>alpha,
            "beta"=>beta,
        )
        # baseline model without all the interactions

        stan_data_singleton = Dict(
            "N" => n,
            "M2" =>M2,
            "M3" =>M3,
            "s" => s,
            "p" =>zero(p),
            "t"=>zero(t),
            "pid"=>pid,
            "tid"=>tid,
            "alpha"=>alpha,
            "beta"=>beta,
        )
        sm = SampleModel("test", stan_program, [8],
          method=StanSample.Sample(
                save_warmup=false, num_warmup=10000,
                num_samples=10000,
                thin=4,
                adapt=StanSample.Adapt(delta=0.85),
                algorithm=StanSample.Hmc(StanSample.Nuts(max_depth=30))
                ),
          tmpdir=tmpdir
        );

        sm0 = SampleModel("test0", stan_program, [8],
          method=StanSample.Sample(save_warmup=false, num_warmup=5000,
          num_samples=5000, thin=4, adapt=StanSample.Adapt(delta=0.9)),
          tmpdir=tmpdir
        );

        try
            rc = stan_sample(sm, data=stan_data);
            rc0 = stan_sample(sm0, data=stan_data_singleton);
            chains = read_samples(sm; output_format=:mambachains,mpsrf=true, transform=true)
            gel = gelmandiag(chains)
            if all(gel.value[:,2,1] .< 1.3)
                valid=true
                data = chains.value
                chains0 = read_samples(sm0; output_format=:mambachains,mpsrf=true, transform=true)
                data0 = chains0.value
            end
        catch e
            if isa(e,LoadError)
                println("Process failed, skipping")
            end
            valid = false
        end
    end
    # save graph to file
    outdict = Dict(
        "graph" => g,
        "samples" => data,
        "samples0" => data0,
        "p" => p0,
        "k2" => k2,
        "k3" => k3,
        "alpha" => alpha,
        "beta" => beta,
        "gamma" =>gamma,
        "gamma1" =>gamma1,
        "n" => n
    )
    outfilepath = "$outdir/$(counter).jld2"
    save(outfilepath, outdict)
    next!(prog)
end