text:
using Plots
using Random
using DataFrames
using UUIDs
using TOML
using Printf
using Dates
import CSV
import BSON
using Statistics
using GridGames
using Serialization
using Base.Threads: @spawn

struct LearnerCallback
    name::String
    event::Symbol
    fn::Function
    period::Int
    warmup::Int
    function LearnerCallback(name, fn; event=:afterstep, period=1, warmup=0)
        new(name, event, fn, period, warmup)
    end
end

tuple(x::LearnerCallback) = x.fn, x.period, x.warmup

mutable struct LearnerStats
    step_count::Counter
    losses::Dict{UInt64,Loss}
    accuracies::Dict{UInt64,PolicyAcc}
    reward_errors::Dict{UInt64,Float16}
    value_errors::Dict{UInt64,Float16}
    times::Dict{UInt64,Float16}
    test_times::Dict{UInt64,Float16}
    policy_test_fails::Union{Nothing,Vector{<:Integer}}
    rec_value_errors::Dict{UInt64,Float16}
    rec_accuracies::Dict{UInt64,PolicyAcc}
    rec_policy_test_fails::Union{Nothing,Vector{<:Integer}}
    games_played::Dict{UInt64,UInt64}
    games_reanalyzed::Dict{UInt64,UInt64}
    games_avg_step::Dict{UInt64,Float16}
end

mutable struct Learner <: AbstractLearner

    uuid::UUID
    dir::String
    training_started::Union{Nothing,DateTime}
    optimizer::Flux.Optimise.AbstractOptimiser
    config::MuZeroConfig
    train::Network
    test::Network
    play::Network
    training_batch::Union{Nothing,TrainingBatch}
    replay_buffers::ReplayBuffers
    callbacks::Vector{LearnerCallback}
    tests::Vector{TrainingTests}
    is_playing::Bool
    is_training::Bool
    play_only::Bool
    stats::LearnerStats

    #@todo - this should be a Function
    valuefn_debug::Union{Nothing,Dict}

    function Learner(
        config,
        train,
        test,
        play,
        replay_buffers;
        tests=[],
        training_batch=nothing,
        optimizer=ADAMW(config.lr_init, (0.9, 0.999), config.weight_decay),
        uuid=uuid4(),
        dir="learner",
        training_started=nothing,
        losses=Dict{UInt64,Float16}(),
        accuracies=Dict{UInt64,Float16}(),
        reward_errors=Dict{UInt64,Float16}(),
        value_errors=Dict{UInt64,Float16}(),
        times=Dict{UInt64,Float16}(),
        test_times=Dict{UInt64,Float16}(),
        policy_test_fails=nothing,
        rec_value_errors=Dict{UInt64,Float16}(),
        rec_accuracies=Dict{UInt64,Float16}(),
        rec_policy_test_fails=nothing,
        callbacks=[],
        play_only=false,
        valuefn_debug=nothing,
        games_played=Dict{UInt64,UInt64}(),
        games_reanalyzed=Dict{UInt64,UInt64}(),
        games_avg_step=Dict{UInt64,Float16}(),
        )
        uuid_object = isnothing(uuid) ? uuid4() : UUID(uuid)
        is_playing = false
        is_training = false
        stats = LearnerStats(
            Counter(0),
            losses,
            accuracies,
            reward_errors,
            value_errors,
            times,
            test_times,
            policy_test_fails,
            rec_value_errors,
            rec_accuracies,
            rec_policy_test_fails,
            games_played,
            games_reanalyzed,
            games_avg_step,
        )
        new(
            uuid_object,
            dir,
            training_started,
            optimizer,
            config,
            train,
            test,
            play,
            training_batch,
            replay_buffers,
            callbacks,
            tests,
            is_playing,
            is_training,
            play_only,
            stats,
            valuefn_debug,
        )
    end
end

function Learner(
    game_type::Type{<:AbstractEnv},
    config::MuZeroConfig;
    kwargs...
    )
    Learner(;
        game_type=game_type,
        config=config,
        kwargs...
    )
end

function Learner(
    game_type::Type{<:AbstractEnv};
    kwargs...
    )
    Learner(;
        game_type=game_type,
        config=getconfig(game_type),
        kwargs...
    )
end

function (x::LearnerCallback)(l::AbstractLearner)
    #println("\nLearnerCallback ($(get_steps(l))): $(x.name), $(x.event), $(x.period), $(x.warmup)")
    if get_steps(l) % x.period == 0 && get_steps(l) >= x.warmup
        #println("$(x.name) running...\n")
        x.fn(l)
    end
end

"Launch a mock training to find a good learning rate and return suggestions based on `suggest_funcs` as a named tuple"
# function lrfind(;
#     start_lr=1e-07,
#     end_lr=10,
#     num_it=100,
#     stop_div=true,
#     show_plot=true,
#     #suggest_funcs=valley
#     )
#     n_epoch = num_it ÷ length(self.dls.train) + 1
#     cb=LRFinder(start_lr=start_lr, end_lr=end_lr, num_it=num_it, stop_div=stop_div)
#     with self.no_logging(): self.fit(n_epoch, cbs=cb)
#     if suggest_funcs is not None:
#         lrs, losses = tensor(self.recorder.lrs[num_it//10:-5]), tensor(self.recorder.losses[num_it//10:-5])
#         nan_idxs = torch.nonzero(torch.isnan(losses.view(-1)))
#         if len(nan_idxs) > 0:
#             drop_idx = min(nan_idxs)
#             lrs = lrs[:drop_idx]
#             losses = losses[:drop_idx]
#         _suggestions, nms = [], []
#         for func in tuplify(suggest_funcs):
#             nms.append(func.__name__ if not isinstance(func, partial) else func.func.__name__) # deal with partials
#             _suggestions.append(func(lrs, losses, num_it))

#         SuggestedLRs = collections.namedtuple('SuggestedLRs', nms)
#         lrs, pnts = [], []
#         for lr, pnt in _suggestions:
#             lrs.append(lr)
#             pnts.append(pnt)
#         if show_plot: self.recorder.plot_lr_find(suggestions=pnts, nms=nms)
#         return SuggestedLRs(*lrs)

#     elif show_plot: self.recorder.plot_lr_find()
# end

function needsreanalyze(l::AbstractLearner, stats)
    stats_path = "$(stats["buffer_dir"])/stats.toml"
    d = TOML.parsefile(stats_path)
    since_last_save = Dates.now() - DateTime(d["saved_at"])
    isabandoned = Dates.value(since_last_save) > 300_000
    if !d["at_capacity"] && !isabandoned
        return false
    end
    curr_step = l.play.step
    buffer_step_min = d["step_min"]
    locked_at = valornothing(d, "locked_at")
    step_diff = curr_step - buffer_step_min
    if step_diff < 1000
        return false
    end
    if isnothing(locked_at)
        return true
    end
    locked_at_diff = Dates.now() - DateTime(locked_at)
    Dates.value(locked_at_diff) > 600_000
end

function getbufferstats(buffer_dir)
    d = TOML.parsefile(
        "$buffer_dir/stats.toml"
    )
    d["buffer_dir"] = buffer_dir
    d
end

function findbuffertoreanalyze(l::AbstractLearner)
    dir = getreplaybuffersdir(l)
    ensuredirexists(dir)
    all_buffer_stats = getbufferstats.(
        readdirexpanduser(dir)
    )
    if isempty(all_buffer_stats)
        return nothing
    end
    d = loadplaysummary(l)
    if isnothing(d)
        return nothing
    end
    avg_step = d["avg_step"]
    avg_and_below = filter(
        x->x["step_mean"] <= avg_step,
        all_buffer_stats
    )
    if isempty(avg_and_below)
        return nothing
    end
    rand(avg_and_below)["buffer_dir"]
end

function reanalyzebuffer(
    l::AbstractLearner,
    buffer_path::String,
    dir=getdir(l),
    )
    buffer = deserializeexpanduser(buffer_path)
    stats_path = getstatspath(
        buffer,
        getreplaybuffersdir(l, dir)
    )
    savereplaybufferstats(
        l,
        buffer,
        stats_path;
        locked_at=string(Dates.now())
    )
    reanalyze!(l, buffer)
    savereplaybuffer!(
        l,
        buffer,
        dir
    )
    println("released lock. ($(buffer.uuid))")
end

function play!(l::AbstractLearner; sleep_period=30)
    l.is_playing = true
    num_actors = getnumactors(l)
    println("play! (num_actors:$num_actors)...")
    responses = []
    if num_actors > 0
        for _ in 1:num_actors
            res = @spawn selfplay!(l)
            push!(responses, res)
        end
    end
    buffer = getwritablereplaybuffer(l)
    total_at_last_save = length(buffer)
    while l.is_playing
        loadplayifnewer!(l)
        if !l.play_only
            loadreplaybuffers!(l)
            saveplaysummary(l)
        end
        sleep(sleep_period)
    end
    responses
end

function loadplaysummary(l::AbstractLearner, dir=getdir(l))
    summary_path = "$dir/play-summary.toml"
    if !isfileexpanduser(summary_path)
        return nothing
    end
    TOML.parsefile(expanduser(summary_path))
end

function saveplaysummary(l::AbstractLearner, dir=getdir(l))
    summary_path = "$dir/play-summary.toml"
    ensure_pathto_exists(summary_path)
    open(expanduser(summary_path), "w") do io
        TOML.print(io, getplaysummary(l))
    end
end

getrrvptests(l::AbstractLearner) = gettests(l, RecurrentRVPTests)
randrrvptest(l::AbstractLearner) = getrrvptests(l) |> rand

getpolicyvaluetests(l::AbstractLearner) = gettests(l, PolicyValueTests)
getpolicytests(l::AbstractLearner) = gettests(l, PolicyTests)

function gettests(l::AbstractLearner, T::Type{<:TrainingTests})
    idx = findfirst(ts->typeof(ts) == T, l.tests)
    isnothing(idx) ? [] : l.tests[idx]
end

#@todo randtest(l::AbstractLearner)

randpolicyvaluetest(l::AbstractLearner) = getpolicyvaluetests(l) |> rand
randpolicytest(l::AbstractLearner) = getpolicytests(l) |> rand
get_steps(l::AbstractLearner) = l.stats.step_count.n
randgame(x::AbstractLearner) = rand(x.replay_buffers)
randtrainingsample(x::AbstractLearner, g::MuZeroGame) = TrainingSample(x.config, g)
randtrainingsample(x::AbstractLearner) = randtrainingsample(x, randgame(x))
getnumactors(x::AbstractLearner) = x.config.num_actors
callbacks(x::AbstractLearner) = x.callbacks
training_network(x::AbstractLearner) = x.train
test_network(x::AbstractLearner) = x.test
play_network(x::AbstractLearner) = x.play
learner_bson_path(l) = learner_bson_path(l, getdir(l))
learner_bson_path(l, dir) = "$dir/$(l.uuid)/learner.bson"
getmodelpath(l, dir=getdir(l)) = "$dir/model.bson"
getmodeloptimizerpath(l, dir=getdir(l)) = "$dir/model-optimizer.bson"
learner_images_dir(l) = "$(getdir(l))/$(l.uuid)/images"
getstatspath(l, dir=getdir(l)) = "$dir/stats.bin"
getreplaybuffersdir(l, dir=getdir(l)) = "$dir/replay_buffers"

function getwritablereplaybufferpath(l, dir=getdir(l))
    getpath(
        getwritablereplaybuffer(l),
        getreplaybuffersdir(l, dir)
    )
end

function getwritablereplaybufferstatspath(l, dir=getdir(l))
    getstatspath(
        getwritablereplaybuffer(l),
        getreplaybuffersdir(l, dir)
    )
end

function Learner(
    config::MuZeroConfig;
    kwargs...
    )
    Learner(;
        game_type=config.game_type,
        config=config,
        kwargs...
    )
end

function Learner(;
    game_type,
    config,
    uuid=nothing,
    dir="./learner",
    tests=[],
    cb=[],
    schedule=nothing,
    play_only=false,
    valuefn_debug=nothing,
    )
    train = new_training_network(config)
    test = new_test_network(config)
    play = new_play_network(config)
    replay_buffers = ReplayBuffers(ReplayBuffer(config))
    #plot_losses_cb = get_plot_losses_cb()
    callbacks = []
    append!(callbacks, cb)

    if !isempty(tests)
        tests_cb = get_tests_cb()
        push!(callbacks,
            LearnerCallback(
                "InitialRunTests",
                tests_cb,
                event=:beforetraining
            )
        )
        push!(callbacks,
            LearnerCallback(
                "UpdateTestNetwork",
                updatetestnet!,
                period=config.checkpoint_interval,
                event=:ontestnetupdate,
            )
        )
        push!(callbacks,
            LearnerCallback(
                "RunTests",
                tests_cb,
                event=:aftertestnetupdate,
            )
        )
    end
    push!(callbacks,
        LearnerCallback(
            "InitialLoss",
            saveloss!,
            event=:beforetraining
        )
    )
    push!(callbacks,
        LearnerCallback(
            "SaveLoss",
            saveloss!,
            event=:afterweightupdate
        )
    )
    push!(callbacks,
        LearnerCallback(
            "SaveLearner",
            save!,
            period=config.save_interval,
            event=:afterstep
        )
    )
    push!(callbacks,
        LearnerCallback(
            "SaveGamePlayStats",
            savegameplaystats!,
            period=1000,
            event=:afterstep
        )
    )

    #accuracy_cb,
    #save_losses_cb,
    #save_accuracy_cb,
    #plot_losses_cb,
    #plot_accuracy_cb, #plot_accuracy_cb = get_plot_accuracy_cb()

    l = Learner(
        config,
        train,
        test,
        play,
        replay_buffers;
        uuid=uuid,
        dir=dir,
        callbacks=callbacks,
        tests=tests,
        play_only=play_only,
        valuefn_debug=valuefn_debug,
    )

    if !isnothing(schedule)
        updateschedule!(l, schedule)
    end

    if !isnothing(uuid)
        if !l.play_only
            loadtrain!(l)
            updatetestnet!(l)
        end

        # @todo - Probably need to load these
        # two only if playing
        loadreplaybuffers!(l)
        loadplay!(l)
        ###

        loadstats!(l)
    else
        l.test = UniformPolicy(l.test)
        l.play = UniformPolicy(l.play)
    end
    l
end

updatetestnet!(l::AbstractLearner) = copy!(l.train, l.test)

mean10(df::DataFrame) = meanx(df, 10)
mean100(df::DataFrame) = meanx(df, 100)
mean1000(df::DataFrame) = meanx(df, 1000)
mean10000(df::DataFrame) = meanx(df, 10000)
function meanallowmissing(x)
    skipped = skipmissing(x)
    if isempty(collect(skipped))
        return missing
    end
    mean(skipped)
end

function meanx(df::DataFrame, n::Integer)
    df[!, :stepgroup] = df.Step .÷ n
    gdf = groupby(df[2:(end-1), :], :stepgroup)
    colnames = filter(name->!in(name,["stepgroup","Step"]), names(df))
    avgcolnames = map(x->"$(x)Avg", colnames)
    combine(gdf, colnames .=> meanallowmissing .=> avgcolnames)
end

function df_cols(l::AbstractLearner)
    cols = [
        :Step   => (l,step)->step,
        :Loss   => (l,step)->valormissing(l.stats.losses, gettotal, step),
    ]
    if !isnothing(l.stats.accuracies)
        push!(cols, :Acc=>(l,step)->valormissing(l.stats.accuracies, getpct, step))
    end
end

function gettraindf(l::AbstractLearner)
    cols = df_cols(l)
    DataFrame(
        map(step->begin
            colnames = Tuple(map(first, cols))
            colvalues = map(col->col[2](l,step), cols)
            NamedTuple{colnames}(colvalues)
        end, 0:get_steps(l))
    )
end

function showresultsfull(l::AbstractLearner; mean=nothing, max_items=20)
    df = getresultsfull(l, mean=mean)
    len = size(df, 1)
    from = max(1, len-max_items+1)
    to = len
    pretty_table(
       df[from:to, :],
       show_row_number=false,
       nosubheader=true,
       title="Training Results (All columns)",
   )
end

function showresults(l::AbstractLearner; mean=nothing)
    pretty_table(
       getresults(l, mean=mean);
       show_row_number=false,
       nosubheader=true,
       title="Training Results",
   )
end

function getresultsfull(l::AbstractLearner; mean=nothing)
    if isnothing(mean)
        len = max(length(l.stats.losses), 1)
        power = Int(floor(log10(len)))
        mean = 10^power
    end
    meanx(
        gettraindfall(l),
        mean
    )
end

function getresults(l::AbstractLearner; mean=nothing)
    if isnothing(mean)
        power = Int(floor(log10(length(l.stats.losses))))
        mean = 10^power
    end
    meanx(
        select(gettraindfall(l), [:Step, :Loss, :PolicyAcc]),
        mean
    )
end

function showlosses(l::AbstractLearner)
    pretty_table(
        getlosses(l);
        show_row_number=false,
        nosubheader=true,
        title="Losses",
    )
end

function getlosses(l::AbstractLearner)
    select(gettraindfall(l), [:Step, :Loss, :PolicyLoss, :RewardLoss, :ValueLoss])
end

getloss(l::AbstractLearner, step) = valormissing(l.stats.losses, gettotal, step)

function getlosses(l::AbstractLearner)
    map(
        step->getloss(l, step),
        1:get_steps(l)
    )
end

function df_cols_full(l::AbstractLearner)
    cols = [
        :Step       => (l,step)->step,
        :Loss       => (l,step)->valormissing(l.stats.losses, gettotal, step),
        :RewardLoss => (l,step)->valormissing(l.stats.losses, getreward, step),
        :ValueLoss  => (l,step)->valormissing(l.stats.losses, getvalue, step),
        :PolicyLoss => (l,step)->valormissing(l.stats.losses, getpolicy, step),
        :TestTime   => (l,step)->valormissing(l.stats.test_times, step),
    ]
    if !isnothing(l.stats.value_errors)
        push!(cols, :ValueErr=>(l,step)->valormissing(l.stats.value_errors, step))
    end
    if !isnothing(l.stats.accuracies)
        push!(cols, :PolicyPass=>(l,step)->valormissing(l.stats.accuracies, getpass, step))
        push!(cols, :PolicyFail=>(l,step)->valormissing(l.stats.accuracies, getfail, step))
        push!(cols, :PolicyAcc=>(l,step)->valormissing(l.stats.accuracies, getpct, step))
    end
    if !isnothing(l.stats.reward_errors)
        push!(cols, :RewardErr=>(l,step)->valormissing(l.stats.reward_errors, step))
    end
    if !isnothing(l.stats.rec_value_errors)
        push!(cols, :RecValueErr=>(l,step)->valormissing(l.stats.rec_value_errors, step))
    end
    if !isnothing(l.stats.rec_accuracies)
        push!(cols, :RecPolicyPass=>(l,step)->valormissing(l.stats.rec_accuracies, getpass, step))
        push!(cols, :RecPolicyFail=>(l,step)->valormissing(l.stats.rec_accuracies, getfail, step))
        push!(cols, :RecPolicyAcc=>(l,step)->valormissing(l.stats.rec_accuracies, getpct, step))
    end
    push!(cols, :Played=>(l,step)->valormissing(l.stats.games_played, step))
    push!(cols, :Reanalyzed=>(l,step)->valormissing(l.stats.games_reanalyzed, step))
    push!(cols, :GameStep=>(l,step)->valormissing(l.stats.games_avg_step, step))
    #push!(cols, :GameStepQuality=>(l,step)->getgamestepquality(l.replay_buffers, step))
    cols
end

function gettraindfall(l::AbstractLearner)
    cols = df_cols_full(l)
    DataFrame(
        map(step->begin
            colnames = Tuple(map(first, cols))
            colvalues = map(col->col[2](l,step), cols)
            NamedTuple{colnames}(colvalues)
        end, 0:get_steps(l))
    )
end

function appendtraincsv(l::AbstractLearner)
    # Step,TrainLoss,PolicyAcc,ValueErr,RewardErr,PolicyLoss,ValueLoss,RewardLoss,Time,TestTime
    # df = DataFrame(Step=[], PolicyAcc=[])
    # losses = Dict()
    # for (step, loss) in enumerate(l.stats.losses)
    #     losses[step-1] = loss
    # end
    # losses = Dict()
    # for acc in l.stats.accuracies
    #     steps[acc.step] = acc.record
    # end
    # for step in 0:length(keys(steps))
    #     push!(df.Step, step)
    #     push!(df.PolicyAcc, x.record)
    # end
    # path = "$(getdir(l))/$(l.uuid)/csv/train.csv"
    # ensure_pathto_exists(path)
    # @show path, df
    # CSV.write(path, df)
end

function persistent_field_names()
    [
        :training_started,
        :reward_errors,
    ]
end

function getcheckpointdir(l::AbstractLearner, dir=getdir(l))
    "$dir/checkpoints/$(getcheckpointstr(l.train.saved_at))"
end

function ismodelimproved(l::AbstractLearner)
    # @todo
    true
end

function save!(l::AbstractLearner, dir=getdir(l); force=false)
    if force || ismodelimproved(l)
        savemodel(l, dir)
        #savemodeloptimizer(l, dir)
        savestats(l, dir)
        savecheckpoint(l, dir)
    end
end

function savestats(l::AbstractLearner, dir=getdir(l))
    path = getstatspath(l, dir)
    ensure_pathto_exists(path)
    serializeexpanduser(path, l.stats)
end

function loadstats!(l::AbstractLearner, path=getstatspath(l))
    if isfileexpanduser(path)
        l.stats = deserializeexpanduser(path)
    end
end

function loadtrain!(l::AbstractLearner, path=getmodelpath(l))
    if isfileexpanduser(path)
        try
            load!(l.train, path)
        catch e
            println("Training model could not be loaded.")
        end
    end
end

function getlatestsummaryondisk(l::AbstractLearner, dir=getdir(l))
    summary_path = "$dir/summary.toml"
    if !isfileexpanduser(summary_path)
        return nothing
    end
    TOML.parsefile(expanduser(summary_path))
end

function loadplayifnewer!(l::AbstractLearner, dir=getdir(l))
    function isnewer()
        d = getlatestsummaryondisk(l, dir)
        if isnothing(d)
            return false
        end
        saved_at = d["saved_at"]
        if isnothing(saved_at)
            return false
        end
        if isnothing(l.play)
            return true
        end
        isnothing(l.play.saved_at) || saved_at > l.play.saved_at
    end
    if isnewer()
        model_path = getmodelpath(l, dir)
        loadplay!(l, model_path)
    end
end

function loadplay!(l::AbstractLearner, path=getmodelpath(l))
    if isfileexpanduser(path)
        try
            load!(l.play, path)
            if l.play_only
                println("Loaded new play network (step: $(l.play.step))")
            end
        catch e
            @error "Play model could not be loaded." exception=(e, catch_backtrace())
        end
    end
end

# function loadmodeloptimizer!(l::AbstractLearner, path=getmodeloptimizerpath(l))
#     if isfileexpanduser(path)
#         try
#             d = BSON.load(expanduser(path), @__MODULE__)
#             load!(l.train, d)
#             @show d[:opt]
#             #l.optimizer = d[:optimizer]
#         catch e
#             println("Optimizer/model could not be loaded. Try loading only the the model.")
#         end
#     end
# end

function savewritablereplaybuffer(l::AbstractLearner, dir=getdir(l))
    length_ondisk = 0
    stats_path = getwritablereplaybufferstatspath(l, dir)
    if isfileexpanduser(stats_path)
        length_ondisk = TOML.parsefile(expanduser(stats_path))["length"]
    end
    writable = getwritablereplaybuffer(l)
    if length(writable) > length_ondisk
        println("buffer length: $(length(writable)) - $(atcapacity(writable, l.config) ? "" : "not") at capacity (saving...) \n")
        savereplaybuffer!(
            l,
            writable,
            dir
        )
    end
end

function newwriteable!(l::AbstractLearner)
    new_writable = ReplayBuffer(l.config)
    setbuffer!(l.replay_buffers, new_writable)
    l.replay_buffers.writable_uuid = string(new_writable.uuid)
end

function savereplaybuffer!(
    l::AbstractLearner,
    buffer::ReplayBuffer,
    dir=getdir(l)
    )
    path = getpath(
        buffer,
        getreplaybuffersdir(l, dir)
    )
    stats_path = getstatspath(
        buffer,
        getreplaybuffersdir(l, dir)
    )
    buffer.saved_at = string(Dates.now())
    updatenetstep!(buffer)
    ensure_pathto_exists(path)
    serializeexpanduser(path, buffer)
    ensure_pathto_exists(stats_path)
    savereplaybufferstats(
        l,
        buffer,
        stats_path;
        locked_at=nothing
    )
end

function savereplaybufferstats(
    l::AbstractLearner,
    buffer::ReplayBuffer,
    stats_path::String;
    locked_at=nothing,
    )
    replay_data = Dict(
        :uuid => string(buffer.uuid),
        :at_capacity => atcapacity(buffer, l.config),
        :length => length(buffer),
        :play_threads => l.config.num_actors,
    )
    fields = [
        :play_start,
        :saved_at,
        :play_finish,
        :reanalyze_start,
        :reanalyze_finish,
        :step_min,
        :reanalyzed,
        :step_max,
        :step_mean,
        :play_duration,
        :reanalyze_duration,
        :single_play_duration,
        :single_reanalyze_duration,
    ]
    for f in fields
        v = getfield(buffer, f)
        if !isnothing(v)
            replay_data[f] = v
        end
    end
    if !isnothing(locked_at)
        replay_data[:locked_at] = locked_at
        println("lock acquired ($(buffer.uuid)) ...")
    end
    ensure_pathto_exists(expanduser(stats_path))
    open(expanduser(stats_path), "w") do io
        TOML.print(io, replay_data)
    end
end

function loadreplaybuffers!(l::AbstractLearner)
    if l.play_only
        return
    end
    writable_path = getwritablereplaybufferpath(l)
    replay_buffers_dir = getreplaybuffersdir(l)
    if isdirexpanduser(replay_buffers_dir)
        readonly_buffer_dirs = filter(
            x->!occursin(writable_path, x),
            readdirexpanduser(replay_buffers_dir)
        )
        #@show readonly_buffer_dirs
        function isnewbuffer(buffer_dir)
            stats_path = "$buffer_dir/stats.toml"
            buffer_uuid = splitpath(stats_path)[end-1]
            if !hasbuffer(l.replay_buffers, buffer_uuid)
                println("new buffer found ($buffer_uuid)")
                return true
            end
            d = TOML.parsefile(stats_path)
            if !haskey(d, "saved_at")
                @show stats_path
                println("no d.saved_at!")
                return false
            end
            saved_at = d["saved_at"]
            buff = getbuffer(l.replay_buffers, buffer_uuid)
            if saved_at > buff.saved_at
                println("refreshing buffer $buffer_uuid (new games found)...")
                return true
            end
            false
        end
        new_buffer_dirs = filter(
            x->isnewbuffer(x),
            readonly_buffer_dirs
        )
        for new_buffer_dir in new_buffer_dirs
            buffer_path = "$new_buffer_dir/buffer.bin"
            setbuffer!(l.replay_buffers, buffer_path)
        end
    end
    nothing
end

function savecheckpoint(l, dir)
    saveconfig(l, dir)
    savesummary(l, dir)
    saveresults(l, dir; mean=1000)
    saveresults(l, dir; mean=10000)
end

function saveresults(l::AbstractLearner, dir=getdir(l); mean=1000)
    checkpoint_dir = getcheckpointdir(l, dir)
    results_path = "$checkpoint_dir/results-m$mean.txt"
    ensure_pathto_exists(results_path)
    open(expanduser(results_path), "w") do io
        pretty_table(
            io,
            getresultsfull(l, mean=mean);
            show_row_number=false,
            nosubheader=true,
            title="Training Results (All columns)",
        )
    end
end

function randtrainingresult(x::AbstractLearner)
    if isnothing(x.training_batch)
        println("No training batch defined.")
        return
    end
    sample = x.training_batch |> rand
    initial_output = initial_inference(x.test, sample.image) |> SingleNetworkOutput
    outputs = [initial_output]
    if !isempty(sample.actions)
        # @todo: make this support arbitrarily long rollouts
        recurrent_output = recurrent_inference(x.test, hidden_state(initial_output), sample.actions[1]) |> SingleNetworkOutput
        push!(outputs, recurrent_output)
    end
    TrainingSampleResult(sample, outputs)
end

function getdir(l::AbstractLearner)
    dir = rstrip(l.dir, '/')
    join([dir,string(l.uuid)], "/")
end

function get_plot_losses_cb()
    function fn(l::AbstractLearner)
        losses = l.stats.losses
        if length(losses) <= 100
            plot(losses, fmt = :png)
        else
            plot(avgx(losses, 100), fmt = :png)
        end
        losses_path = "$(learner_images_dir(l))/losses.png"
        copy_path = "$(learner_images_dir(l))/losses-copy.png"
        ensure_pathto_exists(losses_path)
        savefig(copy_path)
        cp(copy_path, losses_path, force=true)
        rm(copy_path)
    end
    #LearnerCallback(fn)
end

function updatepolicytestresults!(l::AbstractLearner, acc, fails)
    l.stats.accuracies[get_steps(l)] = acc
    l.stats.policy_test_fails = fails
end

function updatevaluetestresults!(l::AbstractLearner, mse)
    l.stats.value_errors[get_steps(l)] = mse
end

function updaterewardtestresults!(l::AbstractLearner, mse)
    l.stats.reward_errors[get_steps(l)] = mse
end

function updaterecvaluetestresults!(l::AbstractLearner, mse)
    l.stats.rec_value_errors[get_steps(l)] = mse
end

function updaterecpolicytestresults!(l::AbstractLearner, acc, fails)
    l.stats.rec_accuracies[get_steps(l)] = acc
    l.stats.rec_policy_test_fails = fails
end

function runtests!(l::AbstractLearner, tests::RecurrentRVPTests)
    reward_mse, value_mse, acc, fails = gettestfn(tests)(gettestnet(l))
    updaterewardtestresults!(l, reward_mse)
    updaterecvaluetestresults!(l, value_mse)
    updaterecpolicytestresults!(l, acc, fails)
    nothing
end

function runtests!(l::AbstractLearner, tests::PolicyTests)
    acc, fails = gettestfn(tests)(gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    nothing
end

function runpvtests(l::AbstractLearner)
    runtests(
        gettests(l, PolicyValueTests),
        gettestnet(l)
    )
end
function runrrvptests(l::AbstractLearner)
    runtests(
        gettests(l, RecurrentRVPTests),
        gettestnet(l)
    )
end

function runtests!(l::AbstractLearner, tests::ValueTests)
    mse = gettestfn(tests)(gettestnet(l))
    updatevaluetestresults!(l, mse)
end

function runtests!(l::AbstractLearner, tests::RewardTests)
    mse = gettestfn(tests)(gettestnet(l))
    updaterewardtestresults!(l, mse)
end

function runtests!(l::AbstractLearner, tests::PolicyValueTests)
    acc, fails, value_mse = runtests(tests, gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    updatevaluetestresults!(l, value_mse)
end

function runtests!(l::AbstractLearner, tests::PolicyRewardTests)
    prtestfn = gettestfn(tests)
    acc, fails, reward_mse = prtestfn(gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    updaterewardtestresults!(l, reward_mse)
end

function runtests!(l::AbstractLearner, tests::ValueRewardTests)
    vrtestfn = gettestfn(tests)
    value_mse, reward_mse = vrtestfn(gettestnet(l))
    updatevaluetestresults!(l, value_mse)
    updaterewardtestresults!(l, reward_mse)
end

function runtests!(l::AbstractLearner, tests::PolicyValueRewardTests)
    rvptestfn = gettestfn(tests)
    acc, fails, value_mse, reward_mse = rvptestfn(gettestnet(l))
    updatepolicytestresults!(l, acc, fails)
    updatevaluetestresults!(l, value_mse)
    updaterewardtestresults!(l, reward_mse)
end

function get_tests_cb()
    function fn(l::AbstractLearner)
        if isempty(l.tests)
            return nothing
        end
        t = @elapsed begin
            for testset in l.tests
                if get_steps(l) % testset.period == 0 && get_steps(l) >= testset.warmup
                    runtests!(l,testset)
                end
            end
        end
        l.stats.test_times[get_steps(l)] = round(t, digits=2)
    end
end

function get_plot_accuracy_cb()
    function fn(l::AbstractLearner)
        accuracies = get_record.(l.stats.accuracies)
        if length(accuracies) <= 100
            plot(accuracies, fmt = :png)
        else
            plot(avgx(accuracies, 100), fmt = :png)
        end
        accuracy_path = "$(learner_images_dir(l))/accuracy.png"
        copy_path = "$(learner_images_dir(l))/accuracy-copy.png"
        ensure_pathto_exists(accuracy_path)
        savefig(copy_path)
        cp(copy_path, accuracy_path, force=true)
        rm(copy_path)
    end
    #LearnerCallback(fn)
end

function clearreplaybuffer!(l::AbstractLearner)
    @assert false "punting on this for now"
    # stopplaying!(l)
    # setwritablereplaybuffer!(l, ReplayBuffer(l.config))
    # nothing
end

stopplaying!(l::AbstractLearner) = l.is_playing = false
getwritablereplaybuffer(l::AbstractLearner) = getwritablebuffer(l.replay_buffers)

function selfplay!(l::AbstractLearner)
    println("selfplay! started on thread", threadid(), " ...")
    while l.is_playing
        game_limit_reached = false
        d = loadplaysummary(l)
        if !isnothing(d)
            if d["num_games"] >= d["max_games"]
                game_limit_reached = true
            end
        end
        if timetoreanalyze(l)
            if l.play_only
                println("reanalyzing on ", threadid(), " ...")
            end
            reanalyze!(l)
        elseif !game_limit_reached
            if l.play_only
                println("playing on ", threadid(), " ...")
            end
            savegames!(l, playgames(l))
        else
            sleep(30)
        end
    end
end

function timetoreanalyze(l::AbstractLearner)
    if isnothing(findbuffertoreanalyze(l))
        return false
    end
    d = loadplaysummary(l)
    if isnothing(d)
        return false
    end
    l.play.step > d["avg_step"] && d["step_quality"] < 1.0
end

function reanalyze!(l::AbstractLearner)
    dir = findbuffertoreanalyze(l)
    if !isnothing(dir)
        reanalyzebuffer(l, "$dir/buffer.bin")
    end
end

function atreplaycapacity(l::AbstractLearner)
    atcapacity(getwritablereplaybuffer(l), l.config)
end

function finishplay!(l::AbstractLearner)
    rbs = getwritablereplaybuffer(l)
    newwriteable!(l)
    finishplay!(rbs)
    savereplaybuffer!(l, rbs)
end

function savegames!(l::AbstractLearner, games::Vector{MuZeroGame})
    replay = getwritablereplaybuffer(l)
    replay.games = games
    # @todo - remove mzgame.net_step
    # (should only be at buffer level)
    replay.step_min = games[1].net_step
    newwriteable!(l)
    finishplay!(replay)
    savereplaybuffer!(l, replay)
end

struct LearnerLossesPNG
    learner::AbstractLearner
end

losses_png(x::AbstractLearner) = x |> LearnerLossesPNG

function Base.show(io::IO, m::MIME"image/png", g::LearnerLossesPNG)
    print(io, read("losses.png", String))
end

randtrainingbatch(x::AbstractLearner) = randtrainingbatch(x.replay_buffers, x.config)

function new_training_batch!(x::AbstractLearner)
    set_training_batch!(x, randtrainingbatch(x))
end

training_batch(x::AbstractLearner) = x.training_batch
training_tensor(x::AbstractLearner) = x |> training_batch |> tensor

# @todo - this should be at the LearnerStats level
function push_time!(l::AbstractLearner, field_name::Symbol, t)
    getfield(l.stats, field_name)[get_steps(l)] = round(t, digits=2)
end

stoptraining!(l::AbstractLearner) = l.is_training = false

function fit!(l::AbstractLearner; steps=nothing)
    l.is_training = true
    println("fit! started on thread ", threadid(), " ...")
    updatetrainingstarted!(l)
    #learning_rate = config.lr_init * config.lr_decay_rate^(get_steps(l) / config.lr_decay_steps)
    starting_step = get_steps(l)
    maxsteps() = isnothing(steps) ? l.config.training_steps : (starting_step + steps)
    while l.is_training && get_steps(l) < maxsteps()
        new_training_batch!(l)
        if get_steps(l) == 0
            runcallbacks(l, :beforetraining)
        end
        runcallbacks(l, :beforeweightupdate)
        t = @elapsed update_weights!(l)
        push_time!(l, :times, t)
        runcallbacks(l, :afterweightupdate)
        runcallbacks(l, :ontestnetupdate)
        runcallbacks(l, :aftertestnetupdate)
        runcallbacks(l, :afterstep)
    end
    save!(l)
    nothing
end

function saveloss!(l::AbstractLearner)
    l.stats.losses[get_steps(l)] = getloss(
        l.train,
        training_tensor(l),
    )
end

getavggamestep(l::AbstractLearner) = getavgstep(l.replay_buffers)

function savegameplaystats!(l::AbstractLearner)
    l.stats.games_played[get_steps(l)] = length(l.replay_buffers)
    l.stats.games_reanalyzed[get_steps(l)] = calcnumreanalyzed(l.replay_buffers)
    l.stats.games_avg_step[get_steps(l)] = getavggamestep(l)
end

function getcallbacks(l::AbstractLearner, event)
    filter(cb->cb.event == event, l.callbacks)
end

function runcallbacks(l::AbstractLearner, event)
    foreach(cb->cb(l), getcallbacks(l, event))
end

function update_weights!(l::AbstractLearner)
    network = l.train
    batch = training_tensor(l)
    weight_decay = l.config.weight_decay
    local loss
    ps = get_weights(network)
    gs = Zygote.gradient(ps) do
        loss = getloss(network, batch)
        gettotal(loss)
    end
    Flux.update!(l.optimizer, ps, gs)
    increment!(l.stats.step_count)
    l.train.step = l.stats.step_count.n
    loss
end

function export_csv(l::AbstractLearner)
    dir = getdir(l)
    train_path = "$dir/csv/train.csv"
    ensure_pathto_exists(train_path)
    df = DataFrame(
        "Step"=>[1],
        "Loss"=>[1],
        "PolicyAcc"=>[1],
        "ValueErr"=>[1],
        "RewardErr"=>[1],
        "PolicyLoss"=>[1],
        "ValueLoss"=>[1],
        "RewardLoss"=>[1],
        "TestTime"=>[1],
    )
    CSV.write(expanduser(train_path), df)
end

function savemodel(l::AbstractLearner, dir)
    if !l.config.skip_save
        save!(l.train, getmodelpath(l, dir))
    end
end

# function savemodeloptimizer(l::AbstractLearner, dir)
#     if !l.config.skip_save
#         save!(
#             l.train,
#             l.optimizer,
#             getmodeloptimizerpath(l, dir)
#         )
#     end
# end

updatetrainingstarted!(l::AbstractLearner) = l.training_started = now()

function getgamestepquality(l::AbstractLearner)
    d = getlatestsummaryondisk(l)
    if isnothing(d) || !haskey(d,"step_count")
        return 0
    end
    getgamestepquality(
        l.replay_buffers,
        d["step_count"],
    )
end

function getplaysummary(l::AbstractLearner)
    Dict(
        "num_games" => length(l.replay_buffers),
        "max_games" => l.config.max_games,
        "step_quality" => getgamestepquality(l),
        "avg_step" => getavggamestep(l),
    )
end

function get_summary(l::AbstractLearner)
    d = Dict(
        "uuid" => string(l.uuid),
        "training_started" => string(l.training_started),
        "saved_at" => l.train.saved_at,
        "step_count" => get_steps(l),
        "avg_game_step" => getavggamestep(l),
    )
    loss = get_loss(l)
    if !isnothing(loss)
        d["loss"] = round(loss, digits=6)
    end
    acc = get_accuracy(l)
    if !isnothing(acc)
        d["accuracy"] = round(getpct(acc), digits=6)
    end
    d
end

get_reward_error(l::AbstractLearner) = lastbykeyornothing(l.stats.reward_errors)
get_accuracy(l::AbstractLearner) = lastbykeyornothing(l.stats.accuracies)
get_loss(l::AbstractLearner) = lastbykeyornothing(l.stats.losses, gettotal)
function getconfigpath(l::AbstractLearner, dir=getdir(l))
     "$(getcheckpointdir(l, dir))/config.txt"
end

function saveconfig(l::AbstractLearner, dir=getdir(l))
    save(l.config, getconfigpath(l, dir))
end

#loadconfig!(l::AbstractLearner) = l.config = load(MuZeroConfig, getconfigpath(l))

function savesummary(l::AbstractLearner, dir)
    summary_path = "$dir/summary.toml"
    ensure_pathto_exists(summary_path)
    open(expanduser(summary_path), "w") do io
        TOML.print(io, get_summary(l))
    end
end

function policy_test_result(x::AbstractLearner)
    sample = x.training_batch |> rand
    initial_output = initial_inference(x.test, sample.image) |> SingleNetworkOutput
    recurrent_output = recurrent_inference(x.test, hidden_state(initial_output), sample.actions[1]) |> SingleNetworkOutput
    PolicyTestResult(sample, [initial_output, recurrent_output])
end

randpolicytestresult(x::AbstractLearner) = x |> getpolicytestresults |> rand
showpolicytestresults(x::AbstractLearner) = display(getpolicytestresults(x))
showpolicyvaluetestresults(x::AbstractLearner) = display(getpolicyvaluetestresults(x))
showrrvptestresults(x::AbstractLearner) = display(getrrvptestresults(x))

function getpvtestresult(l::AbstractLearner, idx)
    getpolicyvaluetestresults(l)[idx]
end

getpolicyvaluetestresult(l::AbstractLearner, idx) = getpolicyvaluetestresults(l)[idx]

function getpolicyvaluetestresults(x::AbstractLearner)
    tests = getpolicyvaluetests(x)
    states = [t.state for t in tests]
    images = tensorize(map(get_image, states))
    output = initial_inference(x.test, images) |> NetworkOutput
    policies = collect(eachcol(policy(output)))
    results = PolicyValueTestResult.(tests, policies)
    PolicyValueTestResults(results)
end

function getrrvptestresults(x::AbstractLearner)
    tests = getrrvptests(x)
    output = recurrentinferencefn(tests)(
        x.test
    )
    RecurrentRVPTestResults(
        RecurrentRVPTestResult.(
            tests,
            first.(eachcol(reward(output))),
            first.(eachcol(value(output))),
            collect(eachcol(policy(output))),
        )
    )
end

function getpolicyvaluetestresults(x::AbstractLearner)
    tests = getpolicyvaluetests(x)
    states = [t.state for t in tests]
    images = tensorize(map(get_image, states))
    output = initial_inference(x.test, images) |> NetworkOutput
    policies = collect(eachcol(policy(output)))
    results = PolicyValueTestResult.(tests, policies)
    PolicyValueTestResults(results)
end

function getpolicy(l::AbstractLearner, s::AbstractEnvState)
    initial_inference(l.test, s) |> SingleNetworkOutput |> policy
end

function getvalue(l::AbstractLearner, s::AbstractEnvState)
    initial_inference(l.test, s) |> SingleNetworkOutput |> value
end

function getvalue(l::AbstractLearner, s::AbstractEnvState, a::Integer)
    hs = initial_inference(l.test, s) |> SingleNetworkOutput |> hidden_state
    recurrent_inference(l.test, hs, a) |> SingleNetworkOutput |> value
end

function findmuzerogames(l::AbstractLearner, target_state::AbstractEnvState)
    mzgames = []
    for (idx,mzgame) in enumerate(l.replay_buffers)
        for s in mzgame.game.states
            if s == target_state
                println("found... $(idx)")
                push!(mzgames, mzgame)
                break
            end
        end
    end
    mzgames
end

function getpolicytestresults(x::AbstractLearner)
    tests = getpolicytests(x)
    states = [t.state for t in tests]
    images = tensorize(map(get_image, states))
    output = initial_inference(x.test, images) |> NetworkOutput
    policies = collect(eachcol(policy(output)))
    results = PolicyTestResult.(tests, policies)
    PolicyTestResults(results)
end

function cycleschedule(; activatation_step=0, min_lr=1e-6, max_lr=1e-2, cycle_steps=10_000)
    (l::AbstractLearner)->begin
        new_lr = trianglelr(
            get_steps(l);
            activatation_step=activatation_step,
            min_lr=min_lr,
            max_lr=max_lr,
            cycle_steps=cycle_steps,
        )
        #println("Change learning rate to: $new_lr")
        l.optimizer[3].eta = new_lr
    end
end

function newcycleschedule(
    l::AbstractLearner;
    activatation_step=get_steps(l),
    min_lr=1e-6,
    max_lr=1e-2,
    cycle_steps=10_000,
    )
    cycleschedule(;
        activatation_step=activatation_step,
        min_lr=min_lr,
        max_lr=max_lr,
        cycle_steps=cycle_steps,
    )
end

function updateschedule!(l::AbstractLearner, schedule)
    updatecallback!(
        l,
        LearnerCallback(
            "UpdateLearningRate",
            schedule,
            event=:beforeweightupdate
        )
    )
end

function updatecallback!(l::AbstractLearner, cb::LearnerCallback)
    idx = findfirst(x->x.name == cb.name, l.callbacks)
    if !isnothing(idx)
        deleteat!(l.callbacks, idx)
    end
    push!(l.callbacks, cb)
end

function trianglelr(step; activatation_step, min_lr=1e-6, max_lr=1e-2, cycle_steps=10_000)
    if step < activatation_step
        return min_lr
    end
    calibrated_step = step - activatation_step
    curr_step = mod1(calibrated_step, cycle_steps)
    halfcycle_steps = cycle_steps ÷ 2
    delta = (max_lr - min_lr) / halfcycle_steps
    if curr_step <= halfcycle_steps
        return min_lr + curr_step*delta
    end
    min_lr + (cycle_steps - curr_step)*delta
end

set_training_batch!(l::AbstractLearner, b::TrainingBatch) = l.training_batch = b
show_game(l::AbstractLearner) = randgame(l) |> display
show_tests(l::AbstractLearner) = show_policy_tests(l)

function show_policy_tests(x::AbstractLearner, n=10)
    indices = 1:n
    len = length(x.policy_tests)
    title_paren = "all $len"
    if len > n
        title_paren = "random batch of $n"
        indices = randperm(length(x.policy_tests))[1:n]
    end
    d = Dict(
        "Test#"=>[],
        "Action(s)"=>[],
    )
    for idx in indices
        t = x.policy_tests[idx]
        push!(d["Test#"], idx)
        push!(d["Action(s)"], join(t.actions, ","))
    end
    df = DataFrame(
        "Test Key"=>d["Test#"],
        "Action(s)"=>d["Action(s)"],
    )
    println()
    println("Policy Tests ($title_paren)\n")
    show(df, show_row_number=false, nosubheader=true, title="")
    println()
end

show_pv_tests(x::AbstractLearner, n=10) = show_policyvalue_tests(x, n)

function show_policyvalue_tests(x::AbstractLearner, n=10)
    indices = 1:n
    len = length(x.policyvalue_tests)
    title_paren = "all $len"
    if len > n
        title_paren = "random batch of $n"
        indices = randperm(length(x.policyvalue_tests))[1:n]
    end
    d = Dict(
        "Test Key"=>[],
        "Value"=>[],
        "Action(s)"=>[],
    )
    for idx in indices
        t = x.policyvalue_tests[idx]
        push!(d["Test Key"], idx)
        push!(d["Value"], t.value)
        push!(d["Action(s)"], join(t.actions, ","))
    end
    df = DataFrame(
        "Test Key"=>d["Test Key"],
        "Value Test"=>d["Value"],
        "Action Test"=>d["Action(s)"],
    )
    println()
    println("Policy-Value Tests ($title_paren)\n")
    show(df, show_row_number=false, nosubheader=true, title="")
    println()
end

function show_reward_tests(x::AbstractLearner, n=10)
    indices = 1:n
    len = length(x.reward_tests)
    title_paren = "all $len"
    if len > n
        title_paren = "random batch of $n"
        indices = randperm(length(x.reward_tests))[1:n]
    end
    d = Dict(
        "Idx"=>[],
        "Action"=>[],
        "Reward Test"=>[],
    )
    for idx in indices
        t = x.reward_tests[idx]
        push!(d["Idx"], idx)
        push!(d["Action"], t.action)
        push!(d["Reward Test"], t.reward)
    end
    df = DataFrame(
        "Test Key"=>d["Idx"],
        "Action"=>d["Action"],
        "Reward Test"=>d["Reward Test"],
    )
    println()
    println("Reward Tests ($title_paren)\n")
    show(df, show_row_number=false, nosubheader=true, title="")
    println()
end

gettestnet(l::AbstractLearner) = l.test
getplaynet(l::AbstractLearner) = l.play

function play_game(l::AbstractLearner, action_idxs::Vector{<:Integer})
    action_list = getactions(l.config.game_type)
    actions = map(idx->action_list[idx], action_idxs)
    play_game(l, plan=actions)
end

function newgames(l::AbstractLearner, i)
    step = getplaynet(l).step
    [new_game(l.config, step) for _ in 1:i]
end

movesmaxedout(config::MuZeroConfig, game) = movesmaxedout(game, config.max_moves)
isactive(config::MuZeroConfig, game) = !is_terminal(game) && !movesmaxedout(config, game)

function playgames(l::AbstractLearner; n=100, plans=[])
    config = l.config
    net = getplaynet(l)
    games = newgames(l, n)
    active_games = view(games, :)
    active_plans = view(plans, :)
    while !isempty(active_games)
        roots = Node.(0, get_state.(active_games))
        net_output = initial_inference(net, active_games) |> NetworkOutput
        expand_node!(roots, legal_actions.(active_games), net_output, zeros(Float32, length(roots)))
        add_exploration_noise(config, roots)
        run_mcts(net, config, roots, action_history.(active_games))
        store_search_statistics.(active_games, roots)
        next_actions = getnextactions(
            roots,
            active_plans;
            num_moves=length.(history.(active_games)),
            num_steps=net.step,
            temp_fn=config.visit_softmax_temperature_fn
        )
        apply!.(active_games, next_actions)
        active_idxs = isactive.(config, active_games)
        active_games = view(active_games, active_idxs)
        if !isempty(active_plans)
            active_plans = view(active_plans, active_idxs)
        end
    end
    games
end

function play_game(l::AbstractLearner; show_root=false, plan=[])
    net = getplaynet(l)
    game = new_game(l.config, net.step)
    while !is_terminal(game) && !movesmaxedout(l.config, game)
        root = Node(0, get_state(game))
        net_output = initial_inference(net, game) |> SingleNetworkOutput
        expand_node!(root, legal_actions(game), net_output, Float32(0))
        add_exploration_noise(l.config, root)
        run_mcts(net, l.config, root, action_history(game))
        store_search_statistics(game, root)
        action = select_action(
            root;
            num_moves=length(history(game)),
            num_steps=net.step,
            temp_fn=l.config.visit_softmax_temperature_fn,
            plan=plan
        )
        apply!(game, action)
    end
    game
end

function reanalyze!(l::AbstractLearner, b::ReplayBuffer)
    try
        b.reanalyze_start = string(now())
        b.games = playgames(
            l;
            n=length(b),
            plans=gethistory.(getgame.(b))
        )
        b.reanalyze_finish = string(now())
        b.reanalyze_duration = Dates.value(
            now() - DateTime(b.reanalyze_start)
        )
        b.single_reanalyze_duration = Int32(
            round(b.reanalyze_duration / length(b))
        )
        incrementreanalyzed!(b)
        updatenetstep!(b)
    catch e
        @error "reanalyze errored out." exception=(e, catch_backtrace())
    end
end

function reanalyze(l::AbstractLearner, mzg::MuZeroGame)
    play_game(l, plan=gethistory(mzg.game))
end

function run_mcts(
    net::Network,
    config::MuZeroConfig,
    root::Node,
    action_history::ActionHistory;
    valuefn_debug=nothing
    )
    min_max_stats = MinMaxStats(config.known_bounds)
    logging = false #length(action_history) == 3
    for sim_idx in 1:config.num_simulations
        if logging
            println()
            @show sim_idx
        end
        history = clone(action_history)
        node = root
        search_path = [node]
        while expanded(node)
            action, node = select_child!(
                node,
                min_max_stats,
                config,
            )
            add_action!(history, action)
            push!(search_path, node)
        end
        parent = search_path[end-1]
        inference = recurrent_inference(net, parent.hidden_state, last_action(history)) |> SingleNetworkOutput
        r = reward(inference)
        expand_node!(node, legal_actions(config, history), inference, r)
        v = value(inference)
        if !isnothing(valuefn_debug)
            v = valuefn_debug[getnextstate(parent.state, last_action(history))]
        end
        backpropagate!(
            search_path,
            v,
            to_play(history),
            config.discount,
            min_max_stats
        )
    end
end

function run_mcts(
    net::Network,
    config::MuZeroConfig,
    roots::Vector{Node},
    action_histories::Vector{ActionHistory}
    )
    min_max_stats_set = MinMaxStats(config.known_bounds, length(roots))
    logging = false

    for sim_idx in 1:config.num_simulations
        if logging
            println()
            @show sim_idx
        end
        histories = clone.(action_histories)
        nodes = roots
        search_paths = [[node] for node in nodes]
        nodes = searchexpanded!.(
            search_paths,
            nodes,
            histories,
            min_max_stats_set,
            config,
        )
        lens = length.(search_paths)
        parents = getindex.(search_paths, lens .- 1)

        inference = recurrent_inference(
            net,
            gethiddenstate.(parents),
            last_action.(histories)
        ) |> NetworkOutput

        r = reward(inference)
        v = value(inference)

        expand_node!(
            nodes,
            legal_actions.(config.game_type, histories),
            inference,
            r[1, :]
        )

        backpropagate!.(
            search_paths,
            r[1, :],
            to_play.(histories),
            config.discount,
            min_max_stats_set
        )
    end
end

function searchexpanded!(
    search_path,
    node,
    history,
    min_max_stats,
    config,
    )
    while expanded(node)
        action, node = select_child!(
            node,
            min_max_stats,
            config,
            )
        add_action!(history, action)
        push!(search_path, node)
    end
    node
end


text:
using PolyChaos, Test
include("../src/sparse_pce.jl")

# Univariate test models and their corresponsing ideal bases
μ, σ = 0, 1
# model_1(x) = exp(μ + σ * x)
# model_2(x) = x^7 - 21 * x^5 + 105 * x^3 - 105 * x + x^3 - 3 * x + 1 # 0th, 3th and 7th Hermite polynomial
# model_3(x) = x^10 - 45 * x^8 + 630 * x^6 - 3150 * x^4 + 4725 * x^2 - 945 # He_10
# model_4(x) = x^10 - 45 * x^8 + 630 * x^6 - 3150 * x^4 + 4725 * x^2 - 945 + x^3 - 3* x # He_10 + He 3
model_5(x) = 0.0
model_6(x) = 1.0
model_7(x) = 50 + 5/2 * (5 * x^3) + 13/16 * (231 * x^6 - 315 * x^4 + 105 * x^2 - 5)

models = [
        #   model_1,
        #   model_2,
        #   model_3,
        #   model_4,
          model_5,
          model_6,
          model_7
         ]
fullBases =
    [
        # GaussOrthoPoly(10),
        # GaussOrthoPoly(10),
        # GaussOrthoPoly(10),
        # GaussOrthoPoly(10),
        GaussOrthoPoly(5),
        GaussOrthoPoly(5),
        LegendreOrthoPoly(10)
    ]
resultBases =  [
        #   [0,1,2,3,4,5,6],
        #   [0, 3, 7],
        #   [10],
        #   [3, 10],
          [0],
          [0],
          [0, 3, 6]
         ]
# PCEs =   [[1.6483774649263525, 1.6490473953964484, 0.8219974064089732, 0.2752047708599591, 0.06663165539581808, 0.013745562538423796, 0.0017708687329313392, 0.00028892409055291307],
        #   [1,1,1],[3,1]]

# tol = 1e-7
# Parameters for sparse algorithm
pMax = 10
Q²tgt = .99999

# Gaussian mixture sampling
function sampleFromGaussianMixture(n::Int,μ::Vector{},σ::Vector{},w::Vector{})
    X = Float64[]
    for i in 1:n
        k = findfirst(x -> x > rand(), cumsum(w))
        push!(X, μ[k] + σ[k]*randn())
    end
    return X
end

@testset "Sparse PCE, univariate bases" begin
    for i in 1:length(models)
        model = models[i]
        op = fullBases[i]
        basis_ref = resultBases[i]
        # pce_ref = PCEs[i]
        pce, basis, maxDeg, R², Q² = sparsePCE(op, model, Q²tgt=Q²tgt, pMax=pMax, jMax=1)

        @test Q² >= Q²tgt
        @test basis == basis_ref
        # @test all(isapprox.(pce, pce_ref; atol = tol))
    end
    @testset "Custom sampling method"

end

# TODO: Test ill-conditioning of matrix with restarting for larger sample set

# TODO: Test multivariate basis
@testset "Sparse PCE, multivariate bases" begin

end

ext: jl

max_stars_repo_name: adriangrupp/sparsePolyChaos

max_issues_repo_name: adriangrupp/sparsePolyChaos

max_forks_repo_name: adriangrupp/sparsePolyChaos


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ a73297b8-ed15-11ea-0915-63edf6769bb1
begin
	using Plots
	using PlutoUI
	using PyPlot
	using DIVAnd
end

# ╔═╡ 51c4cf54-ed23-11ea-3116-e3341d2f2190
md"""
## Set bathymetry parameters
### Bathymetry file
Indicate the path of the directory that stores the bathymetry:
"""

# ╔═╡ 385e2f9a-ed20-11ea-24d7-cd7a536c78df
@bind bathdir TextField()

# ╔═╡ 1710e054-ed24-11ea-30cb-03a941b3b1c3
md"""
Chose among the 3 available files:
"""

# ╔═╡ 53e5bd40-ed17-11ea-1135-036c66a58649
@bind bathfile Select(["gebco_30sec_16.nc", "gebco_30sec_8.nc", "gebco_30sec_4.nc"])

# ╔═╡ 75d3e690-ed24-11ea-2961-c7a76c2572d2
md"""
Check the box if the bathymetry is global
"""

# ╔═╡ 75166930-ed24-11ea-33bc-ab751079962b
@bind bathisglobal CheckBox()

# ╔═╡ 953e6faa-ed24-11ea-077d-7fbbbc4afe85
bathisglobal

# ╔═╡ 77af82a4-ed23-11ea-3424-9368c3f93d71
md"""
### Domain extension
Use the sliders to set the domain
"""

# ╔═╡ e2102e0e-ed21-11ea-3684-77dd737dd550
begin
	aaa = @bind lonmin Slider(20:1:30, default=27);
	bbb = @bind lonmax Slider(40:1:45, default=42);
	ccc = @bind latmin Slider(35:1:45, default=40);
	ddd = @bind latmax Slider(45:1:55, default=47);

	md"""

	lonmin: $(aaa)
	lonmax: $(bbb)

	latmin: $(ccc)
	latmax: $(ddd)

	"""
end

# ╔═╡ e3d08c90-ed22-11ea-2b1c-ffbcc9af6c92
"lon min: $(lonmin), lon max: $(lonmax), lat min: $(latmin), lat max: $(latmax)"

# ╔═╡ a792096a-ed23-11ea-23de-bd35056d7a66
md"""
### Make the plot
"""

# ╔═╡ 3301da4c-f82e-11ea-18bb-adec63e6c73b
function with_pyplot(f::Function)
    f()
    fig = gcf()
    close(fig)
    return fig
end

# ╔═╡ 33785f66-f82c-11ea-13ca-53c5bf06589b
with_pyplot() do
    lonr = lonmin:0.2:lonmax;
	latr = latmin:0.2:latmax;
	bathname = joinpath(bathdir, bathfile);
	bx, by, b = extract_bath(bathname, bathisglobal, lonr, latr);
	llon, llat = ndgrid(bx, by);
	PyPlot.title("Domain bathymetry")
	PyPlot.pcolor(llon, llat, b)
end

# ╔═╡ 87420de6-ed27-11ea-098b-7f0afe067ec8
@info("Interpolation finished")

# ╔═╡ Cell order:
# ╠═a73297b8-ed15-11ea-0915-63edf6769bb1
# ╟─51c4cf54-ed23-11ea-3116-e3341d2f2190
# ╟─385e2f9a-ed20-11ea-24d7-cd7a536c78df
# ╟─1710e054-ed24-11ea-30cb-03a941b3b1c3
# ╟─53e5bd40-ed17-11ea-1135-036c66a58649
# ╟─75d3e690-ed24-11ea-2961-c7a76c2572d2
# ╟─75166930-ed24-11ea-33bc-ab751079962b
# ╠═953e6faa-ed24-11ea-077d-7fbbbc4afe85
# ╟─77af82a4-ed23-11ea-3424-9368c3f93d71
# ╟─e2102e0e-ed21-11ea-3684-77dd737dd550
# ╟─e3d08c90-ed22-11ea-2b1c-ffbcc9af6c92
# ╟─a792096a-ed23-11ea-23de-bd35056d7a66
# ╟─3301da4c-f82e-11ea-18bb-adec63e6c73b
# ╟─33785f66-f82c-11ea-13ca-53c5bf06589b
# ╠═87420de6-ed27-11ea-098b-7f0afe067ec8

max_stars_repo_name: ctroupin/PresentationsHTML

max_issues_repo_name: ctroupin/PresentationsHTML

max_forks_repo_name: ctroupin/PresentationsHTML


text:
__precompile__()
module PSD

import Base: mean

using SymPy
# using DataFrames

# %% General distribution
"""
Abstract type for all distributions
"""
abstract Distribution

abstract UnboundDistribution <: Distribution
abstract BoundDistribution <: Distribution

"""
Probability density function
"""
function pdf(d :: Distribution, x)
	z = symbols("z")
    subs(diff(cdf(d, z), z), z => x)
end

"""
Integrate `f(x) p(x)` expression
"""
function integrate_distribution(f, d :: Distribution; numerical = false, xmin = -oo, xmax = oo, assumption = false)
	if numerical
		quadgk(z -> f(z)*pdf(d, z), xmin, xmax) |> first
	else
		x = assumption ? symbols("x"; (assumption, true)) : symbols("x")
		integrate(f(x)*pdf(d,x), (x, xmin, xmax))
	end
end

"""
Mathematical expectation, mean value, of the distribution

Keyword arguements:
- `numerical` (`false`) use numerical integration?
- `xmin = -oo` - minimum value of x
- `xmax = oo` - maximum value of x
- `assumption = false` - add any assumptions about x (e.g., `:positive`)
"""
function mean(d :: Distribution; kv...)
	integrate_distribution(x->x, d; kv...)
end

"""
Expectation of ξ². Sometimes it is more efficient to calculate just this
without calculating the full variance (that requires also mean of ξ)
"""
function mean_of_square(d :: Distribution; kv...)
	integrate_distribution(x->x^2, d; kv...)
end

"""
Variance of the distribution
"""
function variance(d :: Distribution; kv...)
	x2 = mean_of_square(d; kv...)
    me = mean(d; kv...)
    x2 - me^2
end

function relvar(d :: Distribution; kv...)
	x2 = mean_of_square(d; kv...)
	me = mean(d; kv...)
	x2 / me^2 - (me/me)
end

"""
Returns coefficient of determination R² defined as

R² = 1 - SS(res)/SS(tot)

where

SS(tot) = ∑ (y(i)-E[y])²
SS(res) = ∑(f(i)-y(i))²
"""
function r2{T}(data :: Vector{T}, approximation :: Vector{T})
    data_mean = mean(data)
    ss_tot = sum((y - data_mean)^2 for y ∈ data)
    ss_res = sum((data[i]-approximation[i])^2 for i=1:length(data))
    one(T) - ss_res/ss_tot
end

function r2{T}(y :: Vector{T}, x :: Vector{T}, f :: Function)
    y_mean = mean(y)
    ss_tot = sum((η - y_mean)^2 for η ∈ y)
    ss_res = sum((y[i]-f(x[i]))^2 for i=1:length(y))
    one(T) - ss_res/ss_tot
end

# %% GGS

"""
Gates-Gaudin-Schuhmann distribution

P(x) = (x/xmax)^α
"""
type GGS <: BoundDistribution
    α
    xmax
end

# All properties can be expressed directly:
cdf(d :: GGS, x) = (x / d.xmax)^d.α
pdf(d :: GGS, x) = d.α / d.xmax * (x / d.xmax)^(d.α-1)
mean(d :: GGS) = d.α / (d.α + 1) * d.xmax
variance(d :: GGS) = d.α/((d.α+1)^2 * (d.α+2))*d.xmax^2
relvar(d :: GGS) = 1/(d.α*(d.α+2))

"""
Fits data representing a CFD to a particular
distribution. Returns a pair (D, R2) where
D is a distribution object and R2 is R² coefficient
of determinition of the fit.
"""
function fit_cdf(::Type{Val{GGS}}, x, y)
    xmax = x[end]
    xw = log.(x / xmax)
    yw = log.(y)

    a, λ = linreg(xw,yw)
    f = z -> λ*z
    (GGS(λ,xmax), r2(yw,xw,f))
end

# %%
# %% Harris
"""
Harris distribution
"""
type Harris <: PSD.BoundDistribution
    n
    s
    xmax
end

cdf(d :: Harris, x) = 1-(1-(x/d.xmax)^d.s)^d.n

function pdf(d :: Harris, x)
	ξ = x / d.xmax
	ξs = ξ^(d.s-1)
	d.s*d.n/d.xmax*ξs*(1-ξs*ξ)^(d.n-1)
end

mean(d ::Harris) = d.n * d.xmax / (1+d.n*d.s)*beta(1/d.s, d.n)

function variance(d :: Harris)
	me = mean(d)
	(2*d.n*d.xmax^2)/(2+d.n*d.s)*beta(2/d.s,d.n)-me^2
end

function relvar(d :: Harris)
	(2*(1+d.n*d.s)^2)/(d.n*(2+d.n*d.s)) * beta(2/d.s,d.n)/(beta(1/d.s,d.n))^2-1
end

function fit_cdf(::Type{Val{Harris}}, x, y)
    xmax = x[end]
    ξ = x / xmax
    # Stage 1: small ξ ⇒ get s
    ξ_small = filter(u -> u < 0.1, ξ)
    lnξ_small = log.(ξ_small)
    yw_small = [log(log(1.0/(1.0-y[i]))) for i=1:length(ξ_small)]
    (lnn, s) = linreg(lnξ_small, yw_small)
    r2s = r2(yw_small, lnξ_small, z -> s*z + lnn)
    # Stage 2: get n from all the data points
    # except for last where y=1
    arg = log.(1.0 .- ξ[1:end-1] .^ s)
    val = log.(1.0 .- y[1:end-1])
    (_,n) = linreg(arg,val)
    r2n = r2(val, arg, z -> n*z)
    (Harris(n,s,xmax), (r2n,r2s))
end

# %% Log-normal
"""
Log-normal distribution
"""
type LogNormal <: UnboundDistribution
	μ
	σ
end

cdf(d :: LogNormal, x) = 0.5*(1+erf((log(x)-d.μ)/(√2*d.σ)))
pdf(d :: LogNormal, x) = 1/(x * d.σ*√(2π))*exp(-0.5*((log(x)-d.μ)/d.σ)^2)
mean(d :: LogNormal) = exp(d.μ+(d.σ)^2/2)
variance(d :: LogNormal) = (exp((d.σ)^2)-1)*(mean(d))^2
relvar(d :: LogNormal) = exp((d.σ)^2)-1

function fit_cdf(::Type{Val{LogNormal}}, x, y)
    h = erfinv.(2y-1)
    lnx = log.(x)
    (l,k) = linreg(h,lnx)
    σ = k / sqrt(2.0)
    μ = l
    (LogNormal(μ,σ), r2(lnx, h, z -> k*z+l))
end

# %% Rosin-Rammler

"""
Rosin-Rammler distribution
"""
type RosinRammler <: UnboundDistribution
    α
    x632
end

cdf(d :: RosinRammler, x) = 1 - exp(-(x/d.x632)^d.α)
pdf(d :: RosinRammler, x) = d.α/x * (x/d.x632)^d.α * exp(-(x/d.x632)^d.α)
mean(d :: RosinRammler) = d.x632 * gamma(1+1/d.α)
variance(d :: RosinRammler) = d.x632^2*(gamma(1+2/d.α)-(gamma(1+1/d.α))^2)
relvar(d :: RosinRammler) = gamma(1+2/d.α)/(gamma(1+1/d.α))^2-1

function fit_cdf(::Type{Val{RosinRammler}}, x, y)
    xw = log.(x)
    yw = log.(log.(1.0 ./ (1.0 .- y)))
    (l,k) = linreg(yw, xw)
    α = 1.0/k
    x632 = exp(l)
    (RosinRammler(α,x632), r2(xw, yw, z -> k*z+l))
end

# %% Truncation
"""
Truncate: forward transformation: from 0 < x < ∞ to 0 < y < ymax
"""
truncf(ymax, x) = ymax*x/(ymax+x)
"""
Truncate: backward transformation: from 0 < y < ymax to 0 < x < ∞
"""
truncb(ymax, y) = ymax*y /(ymax - y)

"""
Derivative of truncate backward
"""
dtruncb(ymax, y) = (ymax/(y-ymax))^2

abstract AbstractTruncatedDistribution <: BoundDistribution

type TruncatedDistribution{T <: UnboundDistribution} <: AbstractTruncatedDistribution
    d :: T
    xmax
end

cdf{T <: UnboundDistribution}(d::TruncatedDistribution{T}, x) =
	cdf(d.d, truncb(d.xmax, x))
pdf{T <: UnboundDistribution}(d :: TruncatedDistribution{T}, x) =
	pdf(d.d, truncb(d.xmax, x))*dtruncb(d.xmax, x)

function integrate_distribution(f, d :: TruncatedDistribution; numerical = true, xmin = 0.0, xmax = d.xmax, assumotions = false)
	integrate_distribution(f, d.d, xmin = xmin, xmax = xmax, numerical = numerical, assumptions = assumptions)
end

function mean(d :: TruncatedDistribution; kv...)
	integrate_distribution(x->x, d; kv...)
end

function mean_of_square(d :: TruncatedDistribution; kv...)
	integrate_distribution(x->x^2, d; kv...)
end

function variance(d :: TruncatedDistribution; kv...)
	E = mean(d; kv...)
	E2 = mean_of_square(d; kv...)
	E2 - E^2
end

function relvar(d :: TruncatedDistribution; kv...)
	E = mean(d; kv...)
	E2 = mean_of_square(d; kv...)
	E2/E^2 - 1.0
end

function fit_cdf{T <: UnboundDistribution}(::Type{Val{TruncatedDistribution{T}}}, x, y)
    xmax = x[end]
    xw = truncb.(xmax, x[1:end-1])
    d, r2d = fit_cdf(Val{T}, xw, y[1:end-1])
    (TruncatedDistribution(d, xmax), r2d)
end

# %% Truncated Rosin-Rammler

# Need to create the type to specify p() and P() functions explicitly
"""
Truncated Rosin-Rammler
"""
type TRR <: BoundDistribution
    d :: TruncatedDistribution{RosinRammler}
    TRR(α, x632, xmax) =
		new(TruncatedDistribution(RosinRammler(α, truncb(xmax, x632)), xmax))
end

cdf(d :: TRR, x) = 1 - exp(-(x*d.d.xmax / (d.d.d.x632*(d.d.xmax - x)))^d.d.d.α)

function pdf(d :: TRR, x)
    td = d.d # truncated distribution
    rr = td.d
    α = rr.α
    x632 = rr.x632
    xmax = td.xmax
    tmp = (x*xmax/(x632*(xmax-x)))^α
    α*xmax*tmp / (x*(xmax-x)) * exp(-tmp)
end

# These only work numerically
mean(d :: TRR) = quadgk(x -> x * pdf(d, x), 0.0, d.d.xmax) |> first
variance(d :: TRR) = (quadgk(x->x^2 * pdf(d, x), 0.0, d.d.xmax) |> first) - (mean(d))^2

function fit_cdf(::Type{Val{TRR}}, x, y)
    d, r2d = fit_cdf(Val{TruncatedDistribution{RosinRammler}}, x, y)
    (TRR(d.d.α, truncf(d.xmax, d.d.x632), d.xmax), r2d)
end

# %% Jumped (discontinuous) distribution
type JumpedDistribution{T <: Distribution} <: Distribution
	d :: T
	xmin
end

cdf(d :: JumpedDistribution, x) = cdf(d.d,x)*Heaviside(x-d.xmin)

function pdf(d :: JumpedDistribution, x)
	# No δ(x)-part as it cannot be hanled numerically
	pdf(d.d,x)*Heaviside(x-d.xmin)+cdf(d.d,x)*DiracDelta(x-d.xmin)
end

function integrate_distribution(f, d :: JumpedDistribution; xmin = -oo, kv...)
	if xmin > d.xmin
		integrate_distribution(f, d.d; xmin = xmin, kv...)
	else
		integrate_distribution(f, d.d; xmin = d.xmin, kv...) + f(d.xmin)*cdf(d.d, d.xmin)
	end
end

# %% Characteristics of the distribution

function particles_number(φ, d :: Distribution; numerical = false, kv...)
	(1-φ) * 6 / (numerical? pi : PI) *
	integrate_distribution(x->x^(-3), d; numerical = numerical, kv...)
end

function particles_area(φ, d :: Distribution; kv...)
	6(1-φ) * integrate_distribution(x -> x^(-1), d; kv...)
end

# %%

end # module


text:
function generateU(ul::Vector{Complex128},xn::Vector{Complex128})
	L = length(ul) #Num of basis functions
	N = length(xn) #signal length
	p = 1 #calculate U matrices up to this num (1) for U0,U1
	K = ifloor((N-1-p)/2)
    NPOW = 3

	ul_inv = [1/ul[i] for i=1:L]
	ul_invK = [Base.power_by_squaring(ul_inv[i],K) for i=1:L]

	#initialize memory
	g0 = zeros(Complex128,L)
	g0_K = zeros(Complex128,L)
    g0_NK = zeros(Complex128,L)
	D0 = zeros(Complex128,L)
	U0 = zeros(Complex128,(L,L))
	U1 = zeros(Complex128,(L,L))
    ul_invk = ones(Complex128,L)
	for i=0:(K)
		for j=0:(L-1)
			g0[j+1] += ul_invk[j+1]*xn[i+1]
			g0_K[j+1] += ul_invk[j+1]*xn[K+1+i+1]
            g0_NK[j+1] += ul_invk[j+1]*xn[N-K+i]
            D0[j+1] += (i + 1)*xn[i+1]*ul_invk[j+1] + (K - i)*xn[i + K + 1+1]*ul_invk[j+1]*ul_inv[j+1]*ul_invK[j+1]
            ul_invk[j+1] = ((i%NPOW)==(NPOW-1)) ? Base.power_by_squaring(ul_inv[j+1],i+1) : ul_invk[j+1]*ul_inv[j+1]
		end
	end

	for i=1:L
		for j=1:i
			U0[i,j] = (1/(ul[i]-ul[j]))*(ul[i]*g0[j] - ul[j]*g0[i]
					+ ul_invK[j]*g0_K[i] - ul_invK[i]*g0_K[j])
			U0[j,i] = U0[i,j]
			U1[i,j] = (1/(ul[i]-ul[j]))*(ul[i]*ul[j]*g0[j] - ul[j]*ul[i]*g0[i]
					+ ul_invK[j]*ul[i]*g0_K[i] - ul_invK[i]*ul[j]*g0_K[j])
			U1[j,i] = U1[i,j]
		end
		U0[i,i] = D0[i]
		U1[i,i] = D0[i]*ul[i]-ul[i]*g0[i]+ul_invK[i]*g0_K[i]
	end
	return (U0,U1,g0,g0_K, g0_NK)
end

function xft(xn::Vector{Complex128},ul::Vector{Complex128},Z::Vector{Complex128},q::Float64)

    N = length(xn)
    if length(Z)<N
        error("Length of Z must be at least the length of N")
    end
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC2 = genRdagG0(R,g0_NK*Z[k]^(-N+M))
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0_K*Z[k]^(-M))*(\((RR+q^2),RC2)))[1]
    end
    return Spec+fft(vcat(xn,zeros(length(Z)-N)))
end

function rrt(xn::Vector{Complex128},ul::Vector{Complex128},Z::Vector{Complex128},q::Float64)
    N = length(xn)
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC = genRdagG0(R,g0)
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0)*(\((RR+q^2),RC)))[1]
    end
    return Spec
end

function xczt(xn::Vector{Complex128},A::Complex128,W::Complex128,K::Int,q::Float64)
    N = length(xn)
    Q = N*sum(abs(xn))
    q = Q*q
    Nul = minimum((N,100))
    ul = [exp(im*2*pi/Nul)^(i-1) for i=1:Nul] #fourier basis for expanding Krylov operators U0,U1
    Z = [A*W^(-(i-1)) for i=1:K] #this is the spiral the CZT is to be evaluated on
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC2 = genRdagG0(R,g0_NK*Z[k]^(-N+M))
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0_K*Z[k]^(-M))*(\((RR+q^2),RC2)))[1]
    end
    return Spec+czt(A,W,K,xn)
end

function rczt(xn::Vector{Complex128},A::Complex128,W::Complex128,K::Int,q::Float64)
    N = length(xn)
    Q = N*sum(abs(xn))
    q = Q*q
    Nul = minimum((N,100))
    ul = [exp(im*2*pi/Nul)^(i-1) for i=1:Nul] #fourier basis for expanding Krylov operators U0,U1
    Z = [A*W^(-(i-1)) for i=1:K] #this is the spiral the CZT is to be evaluated on
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC = genRdagG0(R,g0)
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0)*(\((RR+q^2),RC)))[1]
    end
    return Spec
end


text:
using ..CartesianDomains: RectangularDomain

#####
##### RectangularElement
#####

struct RectangularElement{D, X, Y, Z, I}
    data::D
    x::X
    y::Y
    z::Z
    index::I # gives state.realdata[:, variable, index]
    # M::Mⁱʲ
end

"""
    RectangularElement(domain, grid, realdata, element_index)

Returns a Cartesian view into the `realdata` and nodes (located in `grid`)
associated with `element_index`.
"""
function RectangularElement(
    domain::RectangularDomain,
    grid,
    realdata,
    element_index,
)
    volume_geometry = grid.vgeo

    Np = domain.Np
    Te = prod(domain.Ne) # total number of elements

    # Extract views of Cartesian coordinates
    x = view(volume_geometry, :, grid.x1id, :)
    y = view(volume_geometry, :, grid.x2id, :)
    z = view(volume_geometry, :, grid.x3id, :)

    # Reshape x, y, z to (xnode, ynode, znode, element)
    x = reshape(x, Np + 1, Np + 1, Np + 1, Te)
    y = reshape(y, Np + 1, Np + 1, Np + 1, Te)
    z = reshape(z, Np + 1, Np + 1, Np + 1, Te)

    # Reshape realdata as coordinate arrays
    reshaped_realdata = reshape(realdata, Np + 1, Np + 1, Np + 1, Te)

    # Build views to realdata and x, y, z
    data = view(reshaped_realdata, :, :, :, element_index)
    x = view(x, :, :, :, element_index)
    y = view(y, :, :, :, element_index)
    z = view(z, :, :, :, element_index)

    return RectangularElement(data, x, y, z, element_index)
end

# Constructor for index-less elements
RectangularElement(data, x, y, z) = RectangularElement(data, x, y, z, nothing)

Base.eltype(::RectangularElement) = eltype(elem.data)

Base.size(element::RectangularElement) = size(element.data)

Base.@propagate_inbounds Base.getindex(elem::RectangularElement, i, j, k) =
    elem.data[i, j, k]

Base.maximum(f, element::RectangularElement) = maximum(f, element.data)
Base.minimum(f, element::RectangularElement) = minimum(f, element.data)

Base.maximum(element::RectangularElement) = maximum(element.data)
Base.minimum(element::RectangularElement) = minimum(element.data)

function Base.show(io::IO, elem::RectangularElement{D}) where {D}
    intro = "RectangularElement{$(D.name.wrapper)} with "
    data = @sprintf(
        "data ∈ [%.2e, %.2e]\n",
        minimum(elem.data),
        maximum(elem.data)
    )
    x = @sprintf("    x ∈ [%.2e, %.2e]\n", minimum(elem.x), maximum(elem.x))
    y = @sprintf("    y ∈ [%.2e, %.2e]", minimum(elem.y), maximum(elem.y))
    z = @sprintf("    z ∈ [%.2e, %.2e]", minimum(elem.z), maximum(elem.z))

    return print(io, intro, data, x, y, z)
end

#####
##### ⟨⟨ Assemble! ⟩⟩
#####

""" Assemble an array along the first dimension. """
function assemble(::Val{1}, west::AbstractArray, east::AbstractArray)
    contact = @. (west[end:end, :, :] + east[1:1, :, :]) / 2

    east = east[2:end, :, :]
    west = west[1:(end - 1), :, :]

    assembled = cat(west, contact, east, dims = 1)

    return assembled
end

""" Assemble an array along the second dimension. """
function assemble(::Val{2}, south::AbstractArray, north::AbstractArray)
    contact = @. (south[:, end:end, :] + north[:, 1:1, :]) / 2

    north = north[:, 2:end, :]
    south = south[:, 1:(end - 1), :]

    assembled = cat(south, contact, north, dims = 2)

    return assembled
end

""" Assemble an array along the third dimension. """
function assemble(::Val{3}, bottom::AbstractArray, top::AbstractArray)
    contact = @. (bottom[:, :, end:end] + top[:, :, 1:1]) / 2

    top = top[:, :, 2:end]
    bottom = bottom[:, :, 1:(end - 1)]

    assembled = cat(bottom, contact, top, dims = 3)

    return assembled
end

""" Assemble elements along `dim`ension. """
function assemble(dim::Val, left::RectangularElement, right::RectangularElement)
    data = assemble(dim, left.data, right.data)
    x = assemble(dim, left.x, right.x)
    y = assemble(dim, left.y, right.y)
    z = assemble(dim, left.z, right.z)

    return RectangularElement(data, x, y, z, nothing)
end

assemble(dim::Val, e1, e2, e3...) = assemble(dim, e1, assemble(dim, e2, e3...))
assemble(dim::Val, elem) = elem

data(elem::RectangularElement) = elem.data

"""
    assemble(elements::Array{<:RectangularElement, 3})

Assemble the three-dimensional data in `elements` into a single `Array`,
averaging data on shared nodes.
"""
function assemble(
    elements::Array{T, 3},
) where {T <: Union{RectangularElement, AbstractArray}}

    Nx, Ny, Nz = size(elements)

    pencils = [assemble(Val(1), elements[:, j, k]...) for j in 1:Ny, k in 1:Nz]

    slabs = [assemble(Val(2), pencils[:, k]...) for k in 1:Nz]

    volume = assemble(Val(3), slabs...)

    return volume
end


text:
import Statistics

"""
node :
Node object for constructing tree
"""
mutable struct Node
    id              :: Int64
    feature         :: Int64
    threshold       :: Number
    l               :: Node
    r               :: Node
    region          :: Union{Array{Int64,1},Nothing}
    depth           :: Int64
    is_leaf         :: Bool
    value           :: Any
    n               :: Int64

    function Node(region,depth)
        node = new()
        node.depth = depth
        node.region = region
        return(node)
    end
end

"""
CausalTree :
Causal tree object
"""
struct Tree
    tree            :: Node
    n_features      :: Int64
end

"""
predict :
"""
function predict(
    tree            :: Tree,
    X               :: Array{Float64,2}
    )

    tau_ =  Array{Float64,1}()
    n_samples, n_features = size(X)
    tree = tree.tree
    for i in 1:n_samples
        node = tree
        while node.is_leaf == false
            if X[i,node.feature] < node.threshold
                node = node.l
            else
                node = node.r
            end
        end
        push!(tau_,node.value)
    end
    return(tau_)
end

"""
CausalForest :
Causal forest object
"""
struct Forest
    trees           :: Array{Tree,1}
    feature_index   :: Any
    n_trees         :: Int64
    n_features      :: Int64
end

"""
predict :
"""
function predict(
    forest          :: Forest,
    X               :: Array{Float64,2}
    )

    # initialise prediction arrays
    n_samples, n_features = size(X)
    tau_m_ = Array{Float64,2}(undef,(n_samples,forest.n_trees))
    tau_ = Array{Float64,1}(undef,n_samples)

    # predictions from individual trees
    for t in 1:forest.n_trees
        tau_m_[:,t] = predict(forest.trees[t],X[:,forest.feature_index[t,:]])
    end

    # average across trees (columns) for individual predictions
    for i in 1:n_samples
        tau_[i] = Statistics.mean(tau_m_[i,:])
    end

    return(tau_)
end

"""
summary :
Summary of a tree
"""
# function summary(tree)
# end


text:
# # Tracking transient deformation of a cantilever beam

# Source code: [`bending_wave_Ray_tut.jl`](bending_wave_Ray_tut.jl)

# ## Description

# A cantilever beam is given an initial velocity and then at time 0.0 it is
# suddenly stopped by fixing one of its ends. This sends a wave down the beam.

# The beam is modeled as a solid. Consistent mass matrix is used.
# Trapezoidal rule is used to integrate the
# equations of motion in time. Rayleigh mass-proportional damping is
# incorporated. The dynamic stiffness is factorized for efficiency.

# ## Goals

# - Show how to create the discrete model.
# - Demonstrate  trapezoidal-rule time stepping.

##
# ## Definitions

# Basic imports.
using LinearAlgebra
using Arpack

# This is the finite element toolkit itself.
using FinEtools

# The linear stress analysis application is implemented in this package.
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule

# Input parameters
E = 205000*phun("MPa");# Young's modulus
nu = 0.3;# Poisson ratio
rho = 7850*phun("KG*M^-3");# mass density
loss_tangent = 0.0001;
frequency = 1/0.0058;
Rayleigh_mass = 2*loss_tangent*(2*pi*frequency);
L = 200*phun("mm");
W = 4*phun("mm");
H = 8*phun("mm");
tolerance = W/500;
vmag = 0.1*phun("m")/phun("SEC");
tend = 0.013*phun("SEC");

##
# ## Create the discrete model

MR = DeforModelRed3D
fens,fes  = H8block(L,W,H, 50,1,4)

geom = NodalField(fens.xyz)
u = NodalField(zeros(size(fens.xyz,1),3)) # displacement field

nl = selectnode(fens, box=[L L -Inf Inf -Inf Inf], inflate=tolerance)
setebc!(u, nl, true, 1)
setebc!(u, nl, true, 2)
setebc!(u, nl, true, 3)
applyebc!(u)
numberdofs!(u)

corner = selectnode(fens, nearestto=[0 0 0])
cornerzdof = u.dofnums[corner[1], 3]

material = MatDeforElastIso(MR, rho, E, nu, 0.0)

femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3,2)), material)
femm = associategeometry!(femm, geom)
K = stiffness(femm, geom, u)
femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3,3)), material)
M = mass(femm, geom, u)

# Find the relationship of the sum of all the elements of the
# mass matrix and the total mass of the structure.
@show sum(sum(M))
@show L*W*H*rho

# Form the damping matrix.
C = Rayleigh_mass * M

# Figure out the highest frequency in the model, and use a time step that is
# considerably larger than the period of that highest frequency.
evals, evecs = eigs(K, M; nev=1, which=:LM);
@show dt = 350 * 2/sqrt(evals[1]);

# The time stepping loop is protected by `let end` to avoid unpleasant surprises
# with variables getting clobbered by globals.
ts, corneruzs = let dt = dt
    # Initial displacement, velocity, and acceleration.
    U0 = gathersysvec(u)
    v = deepcopy(u)
    v.values[:, 3] .= vmag
    V0 = gathersysvec(v)
    F0 = fill(0.0, length(V0))
    U1 = fill(0.0, length(V0))
    V1 = fill(0.0, length(V0))
    F1 = fill(0.0, length(V0))
    R  = fill(0.0, length(V0))

    # Factorize the dynamic stiffness
    DSF = cholesky((M + (dt/2)*C + ((dt/2)^2)*K))

    # The times and displacements of the corner will be collected into two vectors
    ts = Float64[]
    corneruzs = Float64[]
    # Let us begin the time integration loop:
    t = 0.0;
    step = 0;
    while t < tend
        push!(ts, t)
        push!(corneruzs, U0[cornerzdof])
        t = t+dt;
        step = step + 1;
        (mod(step,25)==0) && println("Step $(step): $(t)")
        # Zero out the load
        fill!(F1, 0.0);
        # Compute the out of balance force.
        R = (M*V0 - C*(dt/2*V0) - K*((dt/2)^2*V0 + dt*U0) + (dt/2)*(F0+F1));
        # Calculate the new velocities.
        V1 = DSF\R;
        # Update the velocities.
        U1 = U0 + (dt/2)*(V0+V1);
        # Switch the temporary vectors for the next step.
        U0, U1 = U1, U0;
        V0, V1 = V1, V0;
        F0, F1 = F1, F0;
        if (t == tend) # Are we done yet?
            break;
        end
        if (t+dt > tend) # Adjust the last time step so that we exactly reach tend
            dt = tend-t;
        end
    end
    ts, corneruzs # return the collected results
end

##
# ## Plot the results

using Gnuplot

@gp "set terminal windows 0 " :-
@gp  :- ts corneruzs./phun("mm") "lw 2 lc rgb 'red' with lines title 'Displacement of the corner' "
@gp  :- "set xlabel 'Time [s]'"
@gp  :- "set ylabel 'Displacement [mm]'"

# The end.
true


text:
@testset "42.trapping-rain-water.jl" begin
    @test trap_rain([4, 2, 0, 3, 2, 5]) == 9
    @test trap_rain([0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]) == 6
end


text:
using BioformatsLoader
using JavaCall
using Test
using Downloads
using AxisArrays

@test BioformatsLoader.init() == nothing

oxr = OMEXMLReader()
@test !JavaCall.isnull(oxr.reader)
@test !JavaCall.isnull(oxr.meta)

ome_imgs_url = "https://downloads.openmicroscopy.org/images"

ome_paths =
    Dict(
        "ND2" => [
            "maxime/BF007.nd2",
            "aryeh/MeOh_high_fluo_003.nd2"
        ],
    )

test_slice = (Axis{:T}(1), Axis{:Z}(1), Axis{:C}(1), Axis{:X}(1), Axis{:Y}(:))
for (fmt, paths) in ome_paths
    for p in paths
        url = "$(ome_imgs_url)/$(fmt)/$(p)"
        imgs = bf_import(url)
        @test length(imgs) > 0

        filename = basename(p)

        mktempdir() do path
            imgpath = joinpath(path, filename)
            Downloads.download(url, imgpath)
            for (order, import_order) in [("TZYXC", (:T, :Z, :Y, :X, :C)),
                                          ("XYCZT", (:X, :Y, :C, :Z, :T))]
                img = bf_import(imgpath; order)[1]
                @test img.ImportOrder == import_order
                @test arraydata(img[test_slice...]) == arraydata(imgs[1][test_slice...])
            end
        end
    end
end


text:
# ---
# title: 867. Transpose Matrix
# id: problem867
# author: Qling
# date: 2021-02-25
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/transpose-matrix/description/>
# hidden: true
# ---
#
# Given a matrix `A`, return the transpose of `A`.
#
# The transpose of a matrix is the matrix flipped over it's main diagonal,
# switching the row and column indices of the matrix.
#
#
# ![](https://assets.leetcode.com/uploads/2019/10/20/hint_transpose.png)
#
#
#
# **Example 1:**
#
#
#
#     Input: [[1,2,3],[4,5,6],[7,8,9]]
#     Output: [[1,4,7],[2,5,8],[3,6,9]]
#
#
# **Example 2:**
#
#
#
#     Input: [[1,2,3],[4,5,6]]
#     Output: [[1,4],[2,5],[3,6]]
#
#
#
#
# **Note:**
#
#   1. `1 <= A.length <= 1000`
#   2. `1 <= A[0].length <= 1000`
#
#
## @lc code=start
using LeetCode

function transpose_matrix(matrix::Vector{Vector{Int}})::Vector{Vector{Int}}
    m, n = length(matrix), length(matrix[1])

    transpose_mat = [[0 for _ = 1:m] for _ = 1:n]

    for i = 1:m, j = 1:n
        transpose_mat[j][i] = matrix[i][j]
    end

    return transpose_mat
end
## @lc code=end


text:
include("AWSCorePrototype.jl")

using .AWSCorePrototype.Services: s3, auto_scaling, glacier
using SymDict

function rest_xml()
    s3("PUT", "/mattbr-test-bucket-ayyy-lmao"; content=
    """
       <CreateBucketConfiguration
           <LocationConstraint>us-east-2</LocationConstraint>
       </CreateBucketConfiguration>"""
    )

    buckets = s3("GET", "/")
    println(buckets)
end

function query()
    autoscaling_groups = auto_scaling("DescribeAutoScalingGroups")
    println(autoscaling_groups)
end

function rest_json()
    query = @SymDict(
        "351396300852",
        "mattbr-test-vault-name"
    )
    vault_name = "mattbr-test-vault-name"
    vault = glacier("PUT", "/accountId/vaults/vaultName", query=query)
end

rest_json()
# rest_xml()
# query()

max_stars_repo_name: mattBrzezinski/Julia-AWS-Status

max_issues_repo_name: mattBrzezinski/Julia-AWS-Status

max_forks_repo_name: mattBrzezinski/Julia-AWS-Status


text:
include("prime.jl")

function compute()
	index=1
	count=0
	while true
		index+=1
		if is_prime(index)
			count+=1
			if( count==10001)
				return index
			end
		end
	end
end

println(compute())


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

using Markdown
using InteractiveUtils

# ╔═╡ 42cb734c-2b47-11eb-12de-6d2181461d88
begin
	using LinearAlgebra
	using BioSequences
	using Plots
	using Random
	using Statistics
end

# ╔═╡ 52fdcf9e-2b47-11eb-3710-6f1ef98c4e51
const KERNEL = let
	equations = []

	# Constrain one: All scaled frequencies sum to zero
	push!(equations, ones(Int, 256))

	# Constrain two: sum(xABC) == sum(ABCx)
	for base in 0:63
		equation = zeros(Int, 256)
		for affix in 0:3
			equation[base << 2 + affix + 1] += 1
			equation[affix << 6 + base + 1] -= 1
		end
		push!(equations, equation)
	end

	# Return the nullspace
	nullspace(Matrix(transpose(reduce(hcat, equations))))
end;

# ╔═╡ fea5f9c0-2b47-11eb-268b-03414fc4887e
"Use the kernel to project kmer freq matrix to orthonormal matrix"
function project_kmer_freq(freq::AbstractMatrix{<:Number})
	mean = sum(freq, dims=2) / size(freq, 2)
	shifted = freq .- mean
	return shifted * KERNEL
end;

# ╔═╡ 7e06120e-2b48-11eb-10fa-5765d37c680a
"Generate kmer spectrum from random nucleotides"
function random_kmer_spectrum(len::Integer, N::Integer)
	seq = LongSequence{DNAAlphabet{2}}(len)
	factor = 1 / (length(seq) - 3)
	counts = zeros(Int, 256)
	result = zeros(Float64, (N, 256))

	for row in 1:N
		rand!(seq.data)
		fill!(counts, zero(eltype(counts)))

		@inbounds for mer in each(DNAMer{4}, seq)
			index = reinterpret(Int, mer.fw) + 1
			counts[index] += 1
		end

		@inbounds for i in eachindex(counts)
			result[row, i] = counts[i] * factor
		end
	end

	return result
end

# ╔═╡ 9609e20c-2b4a-11eb-0f01-016bc09728d0
"Obs x features matrix"
function auto_correlation(matrix::AbstractMatrix)
	correlation = cor(matrix, dims=1)

	# Zero out correlation with a feature and itself, so as not
	# to saturate the heatplots below needlessly
	for i in 1:size(correlation, 1)
		correlation[i, i] = 0
	end
	correlation
end

# ╔═╡ 605618d0-2b4d-11eb-1bd3-ad2a83c52183
begin
	spectrum = random_kmer_spectrum(2_000, 10_000)
	orthonormal = project_kmer_freq(spectrum)
	corr_spectrum = auto_correlation(spectrum)
	corr_orthonormal = auto_correlation(orthonormal)

	mn = min(minimum(corr_spectrum), minimum(corr_orthonormal))
	mx = max(maximum(corr_spectrum), maximum(corr_orthonormal))
	clim = (mn, mx)
end;

# ╔═╡ d424aef0-2b4a-11eb-1e20-d93067c6038a
heatmap(corr_spectrum, clim=clim)

# ╔═╡ 70b600e0-2b49-11eb-1229-5fc87c84fc1e
heatmap(corr_orthonormal, clim=clim)

# ╔═╡ Cell order:
# ╠═42cb734c-2b47-11eb-12de-6d2181461d88
# ╠═52fdcf9e-2b47-11eb-3710-6f1ef98c4e51
# ╠═fea5f9c0-2b47-11eb-268b-03414fc4887e
# ╠═7e06120e-2b48-11eb-10fa-5765d37c680a
# ╠═9609e20c-2b4a-11eb-0f01-016bc09728d0
# ╠═605618d0-2b4d-11eb-1bd3-ad2a83c52183
# ╠═d424aef0-2b4a-11eb-1e20-d93067c6038a
# ╠═70b600e0-2b49-11eb-1229-5fc87c84fc1e

max_stars_repo_name: jakobnissen/play

max_issues_repo_name: jakobnissen/play

max_forks_repo_name: jakobnissen/play


text:
module MimiDICE2010

using Mimi
using XLSX: readxlsx

include("parameters.jl")
include("marginaldamage.jl")

include("components/grosseconomy_component.jl")
include("components/emissions_component.jl")
include("components/co2cycle_component.jl")
include("components/radiativeforcing_component.jl")
include("components/climatedynamics_component.jl")
include("components/sealevelrise_component.jl")
include("components/damages_component.jl")
include("components/neteconomy_component.jl")
include("components/welfare_component.jl")

export construct_dice

# Allow these to be accessed by, e.g., EPA DICE model
const model_years = 2005:10:2595

function get_model(params=nothing)
    params_dict = params == nothing ? dice2010_excel_parameters() : params

    m = Model()
    set_dimension!(m, :time, model_years)

    # --------------------------------------------------------------------------
    # Add components in order
    # --------------------------------------------------------------------------

    add_comp!(m, grosseconomy)
    add_comp!(m, emissions)
    add_comp!(m, co2cycle)
    add_comp!(m, radiativeforcing)
    add_comp!(m, climatedynamics)
    add_comp!(m, sealevelrise)
    add_comp!(m, damages)
    add_comp!(m, neteconomy)
    add_comp!(m, welfare)

    # --------------------------------------------------------------------------
    # Make internal parameter connections
    # --------------------------------------------------------------------------

    # Socioeconomics
    connect_param!(m, :grosseconomy, :I, :neteconomy, :I)
    connect_param!(m, :emissions, :YGROSS, :grosseconomy, :YGROSS)

    # Climate
    connect_param!(m, :co2cycle, :E, :emissions, :E)
    connect_param!(m, :radiativeforcing, :MAT, :co2cycle, :MAT)
    connect_param!(m, :radiativeforcing, :MAT_final, :co2cycle, :MAT_final)
    connect_param!(m, :climatedynamics, :FORC, :radiativeforcing, :FORC)
    connect_param!(m, :sealevelrise, :TATM, :climatedynamics, :TATM)

    # Damages
    connect_param!(m, :damages, :TATM, :climatedynamics, :TATM)
    connect_param!(m, :damages, :YGROSS, :grosseconomy, :YGROSS)
    connect_param!(m, :damages, :TotSLR, :sealevelrise, :TotSLR)
    connect_param!(m, :neteconomy, :YGROSS, :grosseconomy, :YGROSS)
    connect_param!(m, :neteconomy, :DAMFRAC, :damages, :DAMFRAC)
    connect_param!(m, :welfare, :CPC, :neteconomy, :CPC)

    # --------------------------------------------------------------------------
    # Set external parameter values
    # --------------------------------------------------------------------------
    for (name, value) in params_dict
        set_param!(m, name, value)
    end

    return m
end

construct_dice = get_model # still export the old version of the function name

end # module


text:
module Mocha

include("compatibility.jl")
include("logging.jl")
include("config.jl")

if Config.use_native_extension
  include("native.jl")
end

include("macros.jl")
include("base.jl")
include("exception.jl")

include("utils/blas.jl")
include("utils/math.jl")
include("utils/io.jl")
include("utils/tensor.jl")
include("utils/ref-count.jl")

if Config.use_native_extension
  include("utils/im2col-native.jl")
else
  include("utils/im2col.jl")
end

include("backend.jl")
include("blob.jl")

if Config.use_cuda
  include("cuda/cuda.jl")
  include("cuda/cublas.jl")
  include("cuda/cudnn.jl")
  include("cuda/backend.jl")
  include("cuda/blob.jl")
  include("cuda/utils/math.jl")
  include("cuda/utils/padded-copy.jl")
  include("cuda/utils/shifted-copy.jl")
end

export DefaultBackend
const DefaultBackend = Config.use_cuda ? GPUBackend : CPUBackend
const default_backend_type = Config.use_cuda ? "gpu" : "cpu"
@show DefaultBackend

include("initializers.jl")
include("regularizers.jl")
include("constraints.jl")
include("neurons.jl")

if Config.use_cuda
  include("cuda/regularizers.jl")
  include("cuda/constraints.jl")
  include("cuda/neurons.jl")
end

include("pooling-functions.jl")
include("parameter.jl")

include("data-transformers.jl")
if Config.use_cuda
  include("cuda/data-transformers.jl")
end

include("layers.jl")
if Config.use_cuda
  include("cuda/layers.jl")
end

if Config.use_native_extension
  include("layers/pooling/native-impl.jl")
else
  include("layers/pooling/julia-impl.jl")
end
include("layers/pooling/channel-pooling.jl")

include("net.jl")

include("solvers.jl")
include("coffee-break.jl")
include("solvers/sgd-common.jl") # for SGD and Nesterov
include("solvers/policies.jl")
include("solvers/sgd.jl")
include("solvers/nesterov.jl")
include("solvers/adam.jl")
include("solvers/adagrad.jl")
include("solvers/adadelta.jl")

if Config.use_cuda
  include("cuda/solvers.jl")
end

include("utils/gradient-checking.jl")
include("utils/graphviz.jl")

end # module


text:
using Documenter, TSML

makedocs(
  source = "src",
  build = "build",
  modules = [TSML],
  clean = true,
  sitename = "TSML Documentation",
  doctest = false,
  pages = Any[
    "HOME" => "index.md",
    "Tutorial" => Any[
       "tutorial/aggregators.md",
       "tutorial/pipeline.md",
       "tutorial/statistics.md",
       "tutorial/monotonic_plotting.md",
       "tutorial/tsclassifier.md"
    ],
    "Manual" => Any[
      "Value PreProcessing" => "man/valueproc.md",
      "Date PreProcessing" => "man/dateproc.md",
      "Aggregation" => "man/aggregation.md",
      "Imputation" => "man/imputation.md",
    ],
    "ML Library" => Any[
      "Types and Functions" => "lib/functions.md"
    ]
  ],
  format = Documenter.HTML(
     prettyurls = get(ENV, "CI", nothing) == "true"
  )
)

deploydocs(
    repo   = "github.com/IBM/TSML.jl.git",
)


text:
module DynamicSelectApp

using Matte

const title = "Matte Example :: Dynamic Selection"

function ui()
    sidebar_layout(
        side_panel(
            selector("first_select", "Select First Option:", "['Group #1', 'Group #2']"),
            selector("main_select", "Choose a subgroup", "dynamic_items")
        ),
        main_panel(
            h1("Chosen Sub-group"),
            text_output("subgroup")
        )
    )
end

module Server

function dynamic_items(first_select)
  if (first_select == "Group #1")
    ["Subgroup 1.$i" for i in 1:10]
  elseif (first_select == "Group #2")
    ["Subgroup 2.$i" for i in 1:10]
  else
    ["No group selected"]
  end
end

function subgroup(main_select)
    main_select
end

end

end


text:
##########################################################################

#Using Stopping, the idea is to create a buffer function
"""
solveIpopt: Use Ipopt to solve the problem. If stp is an MPCCStoping the MPCC
is set as a NLP using NLMPCC.

`solveIpopt(:: Union{NLPStopping,MPCCStopping}; print_level :: Int = 0)`
"""
function solveIpopt(stp :: Union{NLPStopping,MPCCStopping}; print_level :: Int = 0)

 if typeof(stp.pb) <: AbstractMPCCModel
     nlp = NLMPCC(stp.pb)
 else #typeof(stp.pb) == NLPModels
     nlp = stp.pb
 end
 #xk = solveIpopt(stop.pb, stop.current_state.x)
 stats = ipopt(nlp, print_level     = print_level,
                    tol             = stp.meta.rtol,
                    x0              = stp.current_state.x, #temporary fix
                    max_iter        = stp.meta.max_iter,
                    max_cpu_time    = stp.meta.max_time,
                    dual_inf_tol    = stp.meta.atol,
                    constr_viol_tol = stp.meta.atol,
                    compl_inf_tol   = stp.meta.atol)

 stp.meta.nb_of_stop = stats.iter
 #stats.elapsed_time

 x = stats.solution

 #Not mandatory, but in case some entries of the State are used to stop
  fill_in!(stp, x)
  stop!(stp)

  #Update the meta boolean with the output message
  if stats.status == :first_order stp.meta.suboptimal      = true end
  if stats.status == :acceptable  stp.meta.suboptimal      = true end
  if stats.status == :infeasible  stp.meta.infeasible      = true end
  if stats.status == :small_step  stp.meta.stalled         = true end
  if stats.status == :max_iter    stp.meta.iteration_limit = true end
  if stats.status == :max_time    stp.meta.tired           = true end

 return stp
end