text:
using Hooking
using Base.Test

addr = cglobal(:jl_)
# Test error return
Hooking.hook(addr) do hook, RC
    error()
end

@test_throws ErrorException ccall(:jl_,Void,(Any,),addr)

Hooking.unhook(addr)

# Should not throw anymore
ccall(:jl_,Void,(Any,),Hooking.hook)

# Now test proper return
didrun = false
Hooking.hook(addr) do hook, RC
    global didrun = true
end
ccall(:jl_,Void,(Any,),Hooking.hook)
@test didrun

bigfib(n) = ((BigInt[1 1; 1 0])^n)[2,1]

Hooking.hook(bigfib, Tuple{Int}) do hook, RC
    for ip in Hooking.rec_backtrace(RC)
        @show (ccall(:jl_lookup_code_address, Any, (Ptr{Void}, Cint),
            reinterpret(Ptr{Void},ip-1), 0))[1]
    end
    println("test")
end
bigfib(20)

@test !Hooking.mem_validate(0,sizeof(Ptr{Void}))
x = Array(UInt8,sizeof(Ptr{Void}))
@test Hooking.mem_validate(pointer(x),sizeof(Ptr{Void}))
@test length(x) == sizeof(Ptr{Void})


text:
using Wordle
using Dates
using TOML
using SHA

package_version = begin
    TOML.parsefile(joinpath(@__DIR__(), "..", "..", "Project.toml"))["version"]
end

data = Dict(
    "version" => package_version,
    "wordle" => Wordle.LATEST_WORDLE_NUMBER,
    "data" => Dict(
        "known-wordles" => join(Wordle.WORDLE_LIST, ',') |> sha256 |> bytes2hex,
        "valid-words" => join(Wordle.VALID_WORD_LIST, ',') |> sha256 |> bytes2hex,
    )
)

open(joinpath(@__DIR__(), "..", "Data.toml"), "w") do io
    TOML.print(io, data)
end

ext: jl


text:
#=
test_collater:
- Julia version: 1.1.0
- Author: bramb
- Date: 2019-03-18
=#

using Test
using HyperCollate,MetaGraphs

include("util.jl")

@testset "collater" begin
    @testset "collating 2 xml texts" begin
        include("util.jl")

        f_xml = """
        <text>
            <s>Hoe zoet moet nochtans zijn dit <subst>werven om<add>trachten naar</add></subst> een vrouw,
                de ongewisheid vóór de liefelijke toestemming!</s>
        </text>
        """
        q_xml = """
        <text>
            <s>Hoe zoet moet nochtans zijn dit <subst>werven om<add>trachten naar</add></subst> een vrouw !
                Die dagen van nerveuze verwachting vóór de liefelijke toestemming.</s>
        </text>
        """
        collation = Collation()
        @test collation.state == needs_witness

        add_witness!(collation,"F",f_xml)
        @test collation.state == needs_witness

        add_witness!(collation,"Q",f_xml)
        @test collation.state == ready_to_collate

        collate!(collation)
        @test collation.state == is_collated
        @debug(collation)
        dot = to_dot(collation.graph)
        _print_dot(dot)
    end

    @testset "ranking" begin
        xml = """
        <text><s><subst>Dit kwam van een<add>De</add></subst> te streng doorgedreven rationalisatie</s></text>
        """

        vwg = to_graph(xml)
        r = ranking(vwg)
        for v in keys(r.by_vertex)
            str = get_prop(vwg,v,:text)
            @debug("$str : $(r.by_vertex[v])")
        end
        for rank in sort(collect(keys(r.by_rank)))
            @debug("$rank : $(r.by_rank[rank])")
        end
    end

end


text:
"""
isntequal(nt1::T1, nt2::T2) where {T1<:NamedTuple, T2<:NamedTuple}

Kind of `isequal()` for `NamedTuple`. Order of the keys doesn't matter.
"""
function isntequal(nt1::T1, nt2::T2) where {T1<:NamedTuple, T2<:NamedTuple}
    # Collect the keys of a `NamedTuple` into an array, then sort it.
    sortkeys(nt) = nt |> keys |> collect |> sort

    sfn1 = sortkeys(nt1)
    sfn2 = sortkeys(nt2)
    fnameseq = sfn1 == sfn2

    # keys are not equal -> they are not equal
    fnameseq || return false
    for key in sfn1
        v1 = getindex(nt1, key)
        v2 = getindex(nt2, key)
        # recursively call
        isntequal(v1, v2) || return false
    end
    return true
end

"""
    isntequal(a, b)

Define for any type, using default `isequal`.
It is used when a value is not `NamedTuple` in recursive calls.
"""
isntequal(a, b) = a == b


text:
using StatsPlots,Dates,JLD2,Statistics,Optim,Parameters,Distributions,DataFrames,CSV
using Plots.PlotMeasures,OrdinaryDiffEq,DiffEqCallbacks
import KenyaCoVSD
include("fitting_methods.jl");
include("plotting_methods.jl");
## Load the collated fits
@load("data/N_kenya.jld2")
@load("data/linelist_data_with_pos_neg_20feb_to_27apr_c_age.jld2")
@load("data/serological_data_with_20feb_to_31dec_age.jld2")
@load("data/serological_data_with_20feb_to_10Mar2021_age.jld2")
@load("data/cleaned_linelist20210521_deaths_c__date_of_lab_confirmation.jld2")
@load("data/p_ID.jld2")
@load("forecasts/condensed_county_forecasts.jld2")
@load("data/rel_sero_detection_after_infection.jld2")
@load("modelfits/Nairobi_model.jld2")
nai_model = model
@load("data/cleaned_linelist20210521_c.jld2")#<--- Positive tests only looking ahead of linelist we fitted to

kenya_owid = DataFrame(CSV.File("data/kenya_owid.csv"))
owid_xs = [(Date(d,DateFormat("dd/mm/yyyy")) - Date(2020,2,20)).value for d in kenya_owid.date]
owid_cases = kenya_owid.new_cases
owid_cases[ismissing.(owid_cases)] .= 0
owid_deaths = kenya_owid.new_deaths
owid_deaths[ismissing.(owid_deaths)] .= 0
june1day = (Date(2021,6,1) - Date(2020,2,24)).value

# gr()
# plotlyjs()
## PCR plot --- Kenya wide

n = size(condensed_county_forecasts[1].pred.mean_PCR_forecast,1)
n_1 = size(linelist_data_with_pos_neg.cases,1) - 14

kenya_pcr_forecast = zeros(n)
var_kenya_pcr_forecast =zeros(n)

for fit in condensed_county_forecasts
     kenya_pcr_forecast .+= fit.pred.mean_PCR_forecast[:]
     var_kenya_pcr_forecast .+= fit.pred.std_PCR_forecast[:].^2
end
# std_kenya_pcr_forecast = sqrt.(std_kenya_pcr_forecast)
kenya_pcr_forecast_mv_av = weekly_mv_av(kenya_pcr_forecast)
kenya_pcr_forecast_mv_av_var = weekly_mv_av(var_kenya_pcr_forecast)
kenya_pos = sum(linelist_data_with_pos_neg.cases[:,:,:,1],dims = [2,3])[:]
kenya_pos_mv_av = weekly_mv_av(kenya_pos[1:(end-14)])

xticktimes = [((Date(2020,2,1) + Month(k))- Date(2020,2,24)).value for k = 1:18 ]
xticklabs = [monthname(k)[1:3]*"/20" for k = 3:12]
xticklabs = vcat(xticklabs,[monthname(k)[1:3]*"/21" for k = 1:8])

gr()
PCR_plt = scatter(kenya_pos[1:(end-14)],
        ms = 4,markerstrokewidth = 0,color = :grey,alpha = 0.5,
        xticks = (xticktimes,xticklabs),
        lab = "Daily cases: Kenyan linelist (used for fitting)",legend = :topleft,
        title = "Kenyan PCR test positives",
        xlims = (-5,june1day),
        ylabel = "Daily PCR-confirmed cases",
        size = (1100,500),dpi = 250,
        legendfont = 13,titlefont = 24,tickfontsize=10,guidefont = 18,
        left_margin = 10mm,right_margin = 7.5mm)

plot!(PCR_plt,(1+3):(length(kenya_pos_mv_av)+3),kenya_pos_mv_av,
        color = :black,lw = 3,
        lab = "Daily cases: 7 day mv-av (used in fitting)")

# plot!(PCR_plt,owid_xs[4:(end-3)],weekly_mv_av(owid_cases),color = :red,lw = 3,lab = "Daily cases: Kenyan MoH (7 day mv-av)")
# plot!(PCR_plt,(n_1+1):(size(linelist_data.cases,1)-3),weekly_mv_av(sum(linelist_data.cases,dims = 2)[:])[(n_1+1):end],color = :red,lw = 3,lab = "Daily cases: Kenyan MoH (7 day mv-av)")
smooth_cases_lookahead = weekly_mv_av(sum(linelist_data.cases[(n_1-2):end,:],dims = 2)[:])
plot!(PCR_plt,(length(kenya_pos_mv_av)+4):(length(kenya_pos_mv_av)+3+length(smooth_cases_lookahead)),smooth_cases_lookahead,
        color = :black,lw = 3,ls = :dash,
        lab = "Daily cases: 7 day mv-av (not used in fitting)")

plot!(PCR_plt,4:(n-3),kenya_pcr_forecast_mv_av,ribbon = 9*sqrt.(kenya_pcr_forecast_mv_av_var),
        color = :red, lw = 3,lab = "Model fit and forecast (7 day mv-av)",
        fillalpha = 0.4)

# savefig(PCR_plt,"plots/kenya_cases.pdf")

## Kenyan deaths

n = size(condensed_county_forecasts[1].pred.mean_PCR_forecast,1)

kenya_deaths_forecast = zeros(n)
var_kenya_deaths_forecast = zeros(n)

for fit in condensed_county_forecasts
     kenya_deaths_forecast .+= fit.pred.mean_deaths[:]
     var_kenya_deaths_forecast .+= fit.pred.std_deaths[:].^2
end

kenya_deaths_forecast_mv_av = weekly_mv_av(kenya_deaths_forecast)
kenya_deaths_forecast_mv_av_var = weekly_mv_av(var_kenya_deaths_forecast)
kenya_deaths = sum(deaths_data.deaths,dims = 2)[1:(end-14)]
kenya_deaths_mv_av = weekly_mv_av(kenya_deaths)

xticktimes = [((Date(2020,2,1) + Month(k))- Date(2020,2,24)).value for k = 1:18 ]
xticklabs = [monthname(k)[1:3]*"/20" for k = 3:12]
xticklabs = vcat(xticklabs,[monthname(k)[1:3]*"/21" for k = 1:8])

deaths_plt = scatter(kenya_deaths,
        ms = 4,markerstrokewidth = 0,color = :grey,alpha = 0.5,
        xticks = (xticktimes,xticklabs),
        lab = "Daily deaths: Kenyan linelist",legend = :topleft,
        title = "Kenyan PCR-confirmed deaths",
        xlims = (-5,june1day),ylims = (0,50),
        ylabel = "Daily PCR-confirmed deaths",
        size = (1100,500),dpi = 250,
        legendfont = 13,titlefont = 24,xtickfontsize=10,ytickfontsize=13,guidefont = 18,
        left_margin = 10mm,right_margin = 7.5mm)

plot!(deaths_plt,(1+3):(length(kenya_deaths_mv_av)+3),kenya_deaths_mv_av,
        color = :black,lw = 3,
        lab = "Daily deaths: 7 day mv-av")

# plot!(deaths_plt,owid_xs[4:(end-3)],weekly_mv_av(owid_deaths),color = :red,lw = 3,lab = "Daily cases: Kenyan MoH (7 day mv-av)")

plot!(deaths_plt,4:(n-3),kenya_deaths_forecast_mv_av,ribbon = 9*sqrt.(kenya_deaths_forecast_mv_av_var),
        color = :green, lw = 5, ls = :dot,lab = "Model fit and forecast (7 day mv-av)")

plot!(deaths_plt,(1+3):(length(kenya_deaths_mv_av)+3),cumsum(kenya_deaths_mv_av),
        xticks = (xticktimes[4:4:end],xticklabs[4:4:end]),
        xlims = (-5,june1day),
        color = :black,lw = 3,
        grid = nothing,
        inset = (1,bbox(0.35, -0.25, 0.25, 0.25, :center)),
        lab="",
        subplot = 2,
        title = "Cumulative confirmed deaths",
        bg_inside = nothing)

# plot!(deaths_plt,owid_xs[4:(end-3)],cumsum(weekly_mv_av(owid_deaths)),
#         color = :red,lw = 3,lab = "",
#         subplot = 2)

plot!(deaths_plt,4:(n-3),cumsum(kenya_deaths_forecast_mv_av),
        xlims = (-5,june1day),color = :green, lw = 5, ls = :dot,lab = "",subplot=2)

# savefig(deaths_plt,"plots/kenya_deaths.pdf")

## Kenya Serology plot
# plotlyjs()
gr()
xticktimes = [((Date(2020,2,1) + Month(k))- Date(2020,2,24)).value for k = 1:18 ]
xticklabs = [monthname(k)[1:3]*"/20" for k = 3:12]
xticklabs = vcat(xticklabs,[monthname(k)[1:3]*"/21" for k = 1:8])

n = size(condensed_county_forecasts[1].pred.mean_PCR_forecast,1)
#Group the serology by week
zeropadtomonday = dayofweek(Date(2020,2,20)) - 1
kenya_sero_pos_rnd1_2 = vcat(zeros(zeropadtomonday),
                        sum(serological_data.serodata[:,:,:,1],dims = [2,3])[:])
kenya_sero_neg_rnd1_2 = vcat(zeros(zeropadtomonday),
                        sum(serological_data.serodata[:,:,:,2],dims = [2,3])[:])
kenya_weekly_sero_pos_rnd1_2 = [sum(grp) for grp in Iterators.partition(kenya_sero_pos_rnd1_2,7)]
kenya_weekly_sero_total_rnd1_2 = [sum(grp) for grp in Iterators.partition(kenya_sero_pos_rnd1_2.+kenya_sero_neg_rnd1_2,7)]

kenya_sero_pos_rnd3 = vcat(zeros(zeropadtomonday),
                        sum(serology_data.sero[:,:,:,1],dims = [2,3])[:])
kenya_sero_neg_rnd3 = vcat(zeros(zeropadtomonday),
                        sum(serology_data.sero[:,:,:,2],dims = [2,3])[:])
kenya_weekly_sero_pos_rnd3 = [sum(grp) for grp in Iterators.partition(kenya_sero_pos_rnd3,7)]
kenya_weekly_sero_total_rnd3 = [sum(grp) for grp in Iterators.partition(kenya_sero_pos_rnd3.+kenya_sero_neg_rnd3,7)]

#Jeffery intervals
seroidxs = kenya_weekly_sero_total_rnd3 .> 0
# rnd_1_2idx =
uerr = [invlogcdf(Beta(pos + 0.5,kenya_weekly_sero_total_rnd3[k] - pos + 0.5),log(0.975)) - pos/kenya_weekly_sero_total_rnd3[k] for (k,pos) in enumerate(kenya_weekly_sero_pos_rnd3) ]
lerr = [pos/kenya_weekly_sero_total_rnd3[k] - invlogcdf(Beta(pos + 0.5,kenya_weekly_sero_total_rnd3[k] - pos + 0.5),log(0.025)) for (k,pos) in enumerate(kenya_weekly_sero_pos_rnd3) ]

total_sero_tests_rnds_1_2 = sum(serological_data.serodata)
total_sero_tests= sum(serology_data.sero)

seroreversionrate = 1/(1.33*365)
seroreversionrate2 = 1/(0.595*365)
sero_array = nai_model.sero_sensitivity.*vcat(rel_sero_array_26days[1:30],[(1-seroreversionrate)^k for k in 1:600])
1/(0.595*365)
sero_array_nw = nai_model.sero_sensitivity.*vcat(rel_sero_array_26days[1:30],[1.0 for k in 1:600])

kenya_serology_forecast = zeros(n)
var_kenya_serology_forecast = zeros(n)
kenya_infections_forecast = zeros(n)
kenya_infections_forecast_weighted = zeros(n)
var_kenya_infections_forecast = zeros(n)
kenya_serology_forecast_nw = zeros(n)

test_weighted_kenya_serology_forecast = zeros(n)
test_weighted_var_kenya_serology_forecast = zeros(n)

for fit in condensed_county_forecasts
     county_sero_pos = KenyaCoVSD.simple_conv(fit.pred.mean_incidence₁ .+ fit.pred.mean_incidence₂,sero_array)
     county_sero_pos_nw = fit.pred.mean_serocoverted₁ .+ fit.pred.mean_serocoverted₂
     kenya_infections_forecast .+= cumsum(fit.pred.mean_incidence₁ .+ fit.pred.mean_incidence₂)
     var_kenya_infections_forecast .+= cumsum((fit.pred.std_incidence₁ .+ fit.pred.std_incidence₂).^2)
     county_test_weight = sum(serological_data.serodata[:,serological_data.areas .== uppercase(fit.name),:,:])/total_sero_tests_rnds_1_2
        # county_test_weight = sum(N_kenya[:,fit.name])/sum(N_kenya)
     test_weighted_kenya_serology_forecast .+= county_sero_pos.*(county_test_weight/sum(N_kenya[:,fit.name]))
     kenya_serology_forecast_nw .+= county_sero_pos_nw.*(county_test_weight/sum(N_kenya[:,fit.name]))
     var_kenya_serology_forecast .+= ((county_test_weight/sum(N_kenya[:,fit.name]))^2 ).*(fit.pred.std_serocoverted₁.^2 .+ fit.pred.std_serocoverted₂.^2)

end

xs_mondays = [-3 + (k-1)*7 for k = 1:length(kenya_weekly_sero_total_rnd3)]
rnd1_2_idxs = xs_mondays .< 300
rnd3_idxs = .~rnd1_2_idxs

plt_sero = scatter(xs_mondays[seroidxs.*rnd1_2_idxs],kenya_weekly_sero_pos_rnd3[seroidxs.*rnd1_2_idxs]./kenya_weekly_sero_total_rnd3[seroidxs.*rnd1_2_idxs],
        lab = "Weekly KNBTS: rounds 1 and 2 (used in fitting)",
        legend = :topleft,
        yerr = (lerr[seroidxs.*rnd1_2_idxs],uerr[seroidxs.*rnd1_2_idxs]),
        xticks = (xticktimes,xticklabs),
        title = "Kenyan overall population exposure",
        size = (1100,500),dpi = 250,
        xlims = xlims = (-5,june1day), ylims = (-0.025,1),
        ylabel = "Proportion of population",
        legendfont = 10,titlefont = 24,xtickfontsize=10,ytickfontsize=13,guidefont = 18,
        left_margin = 10mm,right_margin = 7.5mm)

scatter!(plt_sero,xs_mondays[seroidxs.*rnd3_idxs],kenya_weekly_sero_pos_rnd3[seroidxs.*rnd3_idxs]./kenya_weekly_sero_total_rnd3[seroidxs.*rnd3_idxs],
        yerr = (lerr[seroidxs.*rnd3_idxs],uerr[seroidxs.*rnd3_idxs]),
        lab = "Weekly KNBTS: round 3 (not used in fitting)")
plot!(plt_sero,kenya_serology_forecast_nw,lw = 2,color = :green,
        ribbon = 3*sqrt.(var_kenya_serology_forecast),
        lab = "Model fit: seropositivity (test weighted, no seroreversion)" )
plot!(plt_sero,test_weighted_kenya_serology_forecast,lw = 2,ls = :dash,color = :green,
        lab = "Model fit: seropositivity (test weighted, with seroreversion)" )

plot!(plt_sero,kenya_infections_forecast./sum(N_kenya),
        ribbon = 9*sqrt.(var_kenya_infections_forecast)./sum(N_kenya),
        lab = "Model fit: Overall Kenyan population exposure",
        color = :red)

# savefig(plt_sero,"plots/kenya_sero.pdf")

## Kenyan incidence

kenya_group1_incidence = zeros(n)
var_kenya_group1_incidence = zeros(n)
kenya_group2_incidence = zeros(n)
var_kenya_group2_incidence = zeros(n)

for fit in condensed_county_forecasts
     kenya_group1_incidence .+= fit.pred.mean_incidence₁
     var_kenya_group1_incidence .+= fit.pred.std_incidence₁.^2
     kenya_group2_incidence .+= fit.pred.mean_incidence₂
     var_kenya_group2_incidence .+= fit.pred.std_incidence₂.^2
end

#Remove the callback positions
deleteat!(kenya_group1_incidence,kenya_group1_incidence.==0)
deleteat!(var_kenya_group1_incidence,var_kenya_group1_incidence.==0)
deleteat!(kenya_group2_incidence,kenya_group2_incidence.==0)
deleteat!(var_kenya_group2_incidence,var_kenya_group2_incidence.==0)

plt_inc = plot(kenya_group1_incidence./1e5,
        lab = "Daily incidence: Lower SES",
        legend = :topleft,
        ribbon = 3*sqrt.(var_kenya_group1_incidence)./1e5,
        xticks = (xticktimes,xticklabs),
        title = "Kenyan transmission rates by SES group",
        size = (1100,500),dpi = 250,
        xlims = (-5,june1day),
        ylabel = "Daily infections (100,000s)",
        legendfont = 13,titlefont = 24,xtickfontsize=10,ytickfontsize=13,guidefont = 18,
        left_margin = 10mm,right_margin = 7.5mm)

plot!(plt_inc,kenya_group2_incidence./1e5,
        lab = "Daily incidence: Higher SES",
        legend = :topleft,
        ribbon = 3*sqrt.(var_kenya_group2_incidence)./1e5)

# savefig(plt_inc,"plots/kenya_incidence_SES.pdf")

## Rt plots
@load("modelfits/Nairobi_model.jld2")
nairobi_model = deepcopy(model)
@load("modelfits/Mombasa_model.jld2")
mombasa_model = deepcopy(model)
@load("modelfits/Kiambu_model.jld2")
kiambu_model = deepcopy(model)
@load("modelfits/Mandera_model.jld2")
mandera_model = deepcopy(model)

Rt_plt_nai = plot_Rt_both_SES_groups(nairobi_model,Date(2021,8,1))
Rt_plt_mom = plot_Rt_both_SES_groups(mombasa_model,Date(2021,8,1))
Rt_plt_kiambu= plot_Rt_both_SES_groups(kiambu_model,Date(2021,8,1))
Rt_plt_mandera= plot_Rt_both_SES_groups(mandera_model,Date(2021,8,1))

savefig(Rt_plt_nai,"plots/Rt_nai.pdf")
savefig(Rt_plt_mom,"plots/Rt_mom.pdf")
savefig(Rt_plt_kiambu,"plots/Rt_kiambu.pdf")
savefig(Rt_plt_mandera,"plots/Rt_mandera.pdf")


text:
export AbstractHook,
    EmptyHook,
    StepsPerEpisode,
    RewardsPerEpisode,
    TotalRewardPerEpisode,
    TotalBatchRewardPerEpisode,
    BatchStepsPerEpisode,
    TimePerStep,
    DoEveryNEpisode,
    DoEveryNStep,
    DoOnExit

using UnicodePlots: lineplot, lineplot!
using Statistics: mean, std
using CircularArrayBuffers: CircularArrayBuffer

"""
A hook is called at different stage duiring a [`run`](@ref) to allow users to inject customized runtime logic.
By default, an `AbstractHook` will do nothing. One can customize the behavior by implementing the following methods:

- `(hook::YourHook)(::PreActStage, agent, env)`
- `(hook::YourHook)(::PostActStage, agent, env)`
- `(hook::YourHook)(::PreEpisodeStage, agent, env)`
- `(hook::YourHook)(::PostEpisodeStage, agent, env)`
- `(hook::YourHook)(::PostExperimentStage, agent, env)`

By convention, the `Base.getindex(h::YourHook)` is implemented to extract the metrics we are interested in.
Users can compose different `AbstractHook`s with `+`.
"""
abstract type AbstractHook end

(hook::AbstractHook)(args...) = nothing

struct ComposedHook{H} <: AbstractHook
    hooks::H
end

Base.:(+)(h1::AbstractHook, h2::AbstractHook) = ComposedHook((h1, h2))
Base.:(+)(h1::ComposedHook, h2::AbstractHook) = ComposedHook((h1.hooks..., h2))
Base.:(+)(h1::AbstractHook, h2::ComposedHook) = ComposedHook((h1, h2.hooks...))
Base.:(+)(h1::ComposedHook, h2::ComposedHook) = ComposedHook((h1.hooks..., h2.hooks...))

(h::ComposedHook)(args...) = map(h -> h(args...), h.hooks)

#####
# EmptyHook
#####

"""
Nothing but a placeholder.
"""
struct EmptyHook <: AbstractHook end

const EMPTY_HOOK = EmptyHook()

#####
# StepsPerEpisode
#####

"""
    StepsPerEpisode(; steps = Int[], count = 0)

Store steps of each episode in the field of `steps`.
"""
Base.@kwdef mutable struct StepsPerEpisode <: AbstractHook
    steps::Vector{Int} = Int[]
    count::Int = 0
end

Base.getindex(h::StepsPerEpisode) = h.steps

(hook::StepsPerEpisode)(::PostActStage, args...) = hook.count += 1

function (hook::StepsPerEpisode)(::Union{PostEpisodeStage,PostExperimentStage}, agent, env)
    push!(hook.steps, hook.count)
    hook.count = 0
end

#####
# RewardsPerEpisode
#####

"""
    RewardsPerEpisode(; rewards = Vector{Vector{Float64}}())

Store each reward of each step in every episode in the field of `rewards`.
"""
Base.@kwdef mutable struct RewardsPerEpisode <: AbstractHook
    rewards::Vector{Vector{Float64}} = Vector{Vector{Float64}}()
end

Base.getindex(h::RewardsPerEpisode) = h.rewards

(h::RewardsPerEpisode)(::PreEpisodeStage, agent, env) = push!(h.rewards, [])
(h::RewardsPerEpisode)(::PostActStage, agent, env) = push!(h.rewards[end], reward(env))

#####
# TotalRewardPerEpisode
#####

"""
    TotalRewardPerEpisode(; rewards = Float64[], reward = 0.0, is_display_on_exit = true)

Store the total reward of each episode in the field of `rewards`. If
`is_display_on_exit` is set to `true`, a unicode plot will be shown at the [`PostExperimentStage`](@ref).
"""
Base.@kwdef mutable struct TotalRewardPerEpisode <: AbstractHook
    rewards::Vector{Float64} = Float64[]
    reward::Float64 = 0.0
    is_display_on_exit::Bool = true
end

Base.getindex(h::TotalRewardPerEpisode) = h.rewards

(h::TotalRewardPerEpisode)(::PostActStage, agent, env) = h.reward += reward(env)

function (hook::TotalRewardPerEpisode)(::PostEpisodeStage, agent, env)
    push!(hook.rewards, hook.reward)
    hook.reward = 0
end

function (hook::TotalRewardPerEpisode)(::PostExperimentStage, agent, env)
    if hook.is_display_on_exit
        println(
            lineplot(
                hook.rewards,
                title="Total reward per episode",
                xlabel="Episode",
                ylabel="Score",
            ),
        )
    end
end

#####
# TotalBatchRewardPerEpisode
#####
struct TotalBatchRewardPerEpisode <: AbstractHook
    rewards::Vector{Vector{Float64}}
    reward::Vector{Float64}
    is_display_on_exit::Bool
end

Base.getindex(h::TotalBatchRewardPerEpisode) = h.rewards

"""
    TotalBatchRewardPerEpisode(batch_size::Int; is_display_on_exit=true)

Similar to [`TotalRewardPerEpisode`](@ref), but is specific to environments
which return a `Vector` of rewards (a typical case with `MultiThreadEnv`).
If `is_display_on_exit` is set to `true`, a ribbon plot will be shown to reflect
the mean and std of rewards.
"""
function TotalBatchRewardPerEpisode(batch_size::Int; is_display_on_exit=true)
    TotalBatchRewardPerEpisode(
        [Float64[] for _ in 1:batch_size],
        zeros(batch_size),
        is_display_on_exit,
    )
end

function (hook::TotalBatchRewardPerEpisode)(::PostActStage, agent, env)
    R = reward(env)
    for (i, (t, r)) in enumerate(zip(is_terminated(env), R))
        hook.reward[i] += r
        if t
            push!(hook.rewards[i], hook.reward[i])
            hook.reward[i] = 0.0
        end
    end
end

function (hook::TotalBatchRewardPerEpisode)(::PostExperimentStage, agent, env)
    if hook.is_display_on_exit
        n = minimum(map(length, hook.rewards))
        m = mean([@view(x[1:n]) for x in hook.rewards])
        s = std([@view(x[1:n]) for x in hook.rewards])
        p = lineplot(
            m,
            title="Avg total reward per episode",
            xlabel="Episode",
            ylabel="Score",
        )
        lineplot!(p, m .- s)
        lineplot!(p, m .+ s)
        println(p)
    end
end

#####
# BatchStepsPerEpisode
#####

struct BatchStepsPerEpisode <: AbstractHook
    steps::Vector{Vector{Int}}
    step::Vector{Int}
end

Base.getindex(h::BatchStepsPerEpisode) = h.steps

"""
    BatchStepsPerEpisode(batch_size::Int; tag = "TRAINING")

Similar to [`StepsPerEpisode`](@ref), but is specific to environments
which return a `Vector` of rewards (a typical case with `MultiThreadEnv`).
"""
function BatchStepsPerEpisode(batch_size::Int)
    BatchStepsPerEpisode([Int[] for _ in 1:batch_size], zeros(Int, batch_size))
end

function (hook::BatchStepsPerEpisode)(::PostActStage, agent, env)
    for (i, t) in enumerate(is_terminated(env))
        hook.step[i] += 1
        if t
            push!(hook.steps[i], hook.step[i])
            hook.step[i] = 0
        end
    end
end

#####
# TimePerStep
#####

"""
    TimePerStep(;max_steps=100)
    TimePerStep(times::CircularArrayBuffer{Float64}, t::UInt64)

Store time cost of the latest `max_steps` in the `times` field.
"""
mutable struct TimePerStep <: AbstractHook
    times::CircularArrayBuffer{Float64,1}
    t::UInt64
end

Base.getindex(h::TimePerStep) = h.times

TimePerStep(; max_steps=100) =
    TimePerStep(CircularArrayBuffer{Float64}(max_steps), time_ns())

function (hook::TimePerStep)(::PostActStage, agent, env)
    push!(hook.times, (time_ns() - hook.t) / 1e9)
    hook.t = time_ns()
end

"""
    DoEveryNStep(f; n=1, t=0)

Execute `f(t, agent, env)` every `n` step.
`t` is a counter of steps.
"""
mutable struct DoEveryNStep{F} <: AbstractHook
    f::F
    n::Int
    t::Int
end

DoEveryNStep(f; n=1, t=0) = DoEveryNStep(f, n, t)

function (hook::DoEveryNStep)(::PostActStage, agent, env)
    hook.t += 1
    if hook.t % hook.n == 0
        hook.f(hook.t, agent, env)
    end
end

"""
    DoEveryNEpisode(f; n=1, t=0)

Execute `f(t, agent, env)` every `n` episode.
`t` is a counter of episodes.
"""
mutable struct DoEveryNEpisode{S<:Union{PreEpisodeStage,PostEpisodeStage},F} <: AbstractHook
    f::F
    n::Int
    t::Int
end

DoEveryNEpisode(f::F; n=1, t=0, stage::S=POST_EPISODE_STAGE) where {S,F} =
    DoEveryNEpisode{S,F}(f, n, t)

function (hook::DoEveryNEpisode{S})(::S, agent, env) where {S}
    hook.t += 1
    if hook.t % hook.n == 0
        hook.f(hook.t, agent, env)
    end
end

"""
    DoOnExit(f)

Call the lambda function `f` at the end of an [`Experiment`](@ref).
"""
struct DoOnExit{F} <: AbstractHook
    f::F
end

function (h::DoOnExit)(::PostExperimentStage, agent, env)
    h.f(agent, env)
end


text:
abstract type AbstractObserver end

measure!(o::AbstractObserver; kwargs...) = nothing
checkdone!(o::AbstractObserver; kwargs...) = false

"""
NoObserver is a trivial implementation of an
observer type which can be used as a default
argument for DMRG routines taking an AbstractObserver
"""
struct NoObserver <: AbstractObserver end

"""
A DMRGMeasurement object is an alias for `Vector{Vector{Float64}}`,
in other words an array of arrays of real numbers.

Given a DMRGMeasurement `M`,the result for the
measurement on sweep `n` and site `i` as `M[n][i]`.
"""
const DMRGMeasurement = Vector{Vector{Float64}}

"""
DMRGObserver is an implementation of an
observer object (<:AbstractObserver) which
implements custom measurements and allows
the `dmrg` function to return early if an
energy convergence criterion is met.
"""
struct DMRGObserver <: AbstractObserver
  ops::Vector{String}
  sites::Vector{<:Index}
  measurements::Dict{String,DMRGMeasurement}
  energies::Vector{Float64}
  truncerrs::Vector{Float64}
  etol::Float64
  minsweeps::Int64
end

"""
    DMRGObserver(;energy_tol=0.0,
                  minsweeps=2)

Construct a DMRGObserver by providing the energy
tolerance used for early stopping, and minimum number
of sweeps that must be done.

  - energy_tol: if the energy from one sweep to the
    next no longer changes by more than this amount,
    stop after the current sweep
  - minsweeps: do at least this many sweeps
"""
function DMRGObserver(; energy_tol=0.0, minsweeps=2)
  return DMRGObserver(
    [], Index[], Dict{String,DMRGMeasurement}(), [], [], energy_tol, minsweeps
  )
end

"""
    DMRGObserver(ops::Vector{String},
                 sites::Vector{<:Index};
                 energy_tol=0.0,
                 minsweeps=2)

Construct a DMRGObserver, provide an array
of `ops` of operator names which are strings
recognized by the `op` function. Each of
these operators will be measured on every site
during every step of DMRG and the results
recorded inside the DMRGOberver for later
analysis. The array `sites` is the basis
of sites used to define the MPS and MPO for
the DMRG calculation.

Optionally, one can provide an energy
tolerance used for early stopping, and minimum number
of sweeps that must be done.

  - energy_tol: if the energy from one sweep to the
    next no longer changes by more than this amount,
    stop after the current sweep
  - minsweeps: do at least this many sweeps
"""
function DMRGObserver(
  ops::Vector{String}, sites::Vector{<:Index}; energy_tol=0.0, minsweeps=2
)
  measurements = Dict(o => DMRGMeasurement() for o in ops)
  return DMRGObserver(ops, sites, measurements, [], [], energy_tol, minsweeps)
end

"""
    measurements(o::DMRGObserver)

After using a DMRGObserver object `o` within
a DMRG calculation, retrieve a dictionary
of measurement results, with the keys being
operator names and values being DMRGMeasurement
objects.
"""
measurements(o::DMRGObserver) = o.measurements

"""
    energies(o::DMRGObserver)

After using a DMRGObserver object `o` within
a DMRG calculation, retrieve an array of the
energy after each sweep.
"""
energies(o::DMRGObserver) = o.energies

sites(obs::DMRGObserver) = obs.sites

ops(obs::DMRGObserver) = obs.ops

truncerrors(obs::DMRGObserver) = obs.truncerrs

function measurelocalops!(obs::DMRGObserver, wf::ITensor, i::Int)
  for o in ops(obs)
    m = dot(wf, noprime(op(sites(obs), o, i) * wf))
    imag(m) > 1e-8 && (@warn "encountered finite imaginary part when measuring $o")
    measurements(obs)[o][end][i] = real(m)
  end
end

function measure!(obs::DMRGObserver; kwargs...)
  half_sweep = kwargs[:half_sweep]
  b = kwargs[:bond]
  energy = kwargs[:energy]
  psi = kwargs[:psi]
  truncerr = truncerror(kwargs[:spec])

  if half_sweep == 2
    N = length(psi)

    if b == (N - 1)
      for o in ops(obs)
        push!(measurements(obs)[o], zeros(N))
      end
      push!(truncerrors(obs), 0.0)
    end

    # when sweeping left the orthogonality center is located
    # at site n=b after the bond update.
    # We want to measure at n=b+1 because there the tensor has been
    # already fully updated (by the right and left pass of the sweep).
    wf = psi[b] * psi[b + 1]
    measurelocalops!(obs, wf, b + 1)

    if b == 1
      push!(energies(obs), energy)
      measurelocalops!(obs, wf, b)
    end
    truncerr > truncerrors(obs)[end] && (truncerrors(obs)[end] = truncerr)
  end
end

function checkdone!(o::DMRGObserver; kwargs...)
  outputlevel = get(kwargs, :outputlevel, false)
  if (
    length(energies(o)) > o.minsweeps &&
    abs(energies(o)[end] - energies(o)[end - 1]) < o.etol
  )
    outputlevel > 0 && println("Energy difference less than $(o.etol), stopping DMRG")
    return true
  end
  return false
end


text:
# Iterative Ground State Search
# #+HTML: Source
# #+HTML:

# [[file:~/.julia/dev/MatrixProductStates/README.org::*Iterative%20Ground%20State%20Search][Iterative Ground State Search:1]]
function R_exprs(ψ::MPS{L, T}, H::MPO{L, T}) where {L, T}
    R_exs = Array{T, 3}[]
    R_ex = ones(T, 1, 1, 1)
    for l in L:-1:2
        R_ex = iterate_R_ex(ψ[l], H[l], R_ex)
        push!(R_exs, R_ex)
    end
    reverse(R_exs)
end

# function preallocate_hs(ψ::MPS{L, T}) where {L, T}
#     h_tnsrs = map(ψ.tensors) do M
#         Dˡ⁻¹, Dˡ, d = size(M)
#         Array{T, 6}(undef, d, Dˡ⁻¹, Dˡ, d, Dˡ⁻¹, Dˡ)
#     end
# end

function sweep!(::Right, ψ::MPS{L, T}, H::MPO{L, T}, R_exs) where {L, T}
    L_exs = Array{T, 3}[]
    L_ex  = ones(T, 1, 1, 1)
    E = zero(T)
    for l in 1:(L-1)
        W = H[l]

        E, A, SVp = eigenproblem(right, ψ[l], L_ex, W, R_exs[l])
        ψ.tensors[l] = A

        L_ex = iterate_L_ex(A, W, L_ex)
        push!(L_exs, L_ex)

        Bp1 = ψ.tensors[l+1]
        @tensor Mp1[sⁱ⁻¹, aⁱ, σⁱ] := SVp[sⁱ⁻¹, aⁱ⁻¹] * Bp1[aⁱ⁻¹, aⁱ, σⁱ]
        ψ.tensors[l+1] = Mp1
    end
    L_exs, E
end

function sweep!(::Left, ψ::MPS{L, T}, H::MPO{L, T}, L_exs) where {L, T}
    R_exs = Array{T, 3}[]
    R_ex  = ones(T, 1, 1, 1)
    E = zero(T)
    for l in L:-1:2
        W = H[l]

        E, US, B = eigenproblem(left, ψ[l], L_exs[l-1], W, R_ex)
        ψ.tensors[l] = B

        R_ex = iterate_R_ex(B, W, R_ex)
        push!(R_exs, R_ex)

        Am1 = ψ.tensors[l-1]
        @tensor Mm1[aˡ⁻², sˡ⁻¹, σˡ⁻¹] :=  Am1[aˡ⁻², aˡ⁻¹′, σˡ⁻¹] * US[aˡ⁻¹′, sˡ⁻¹]
        ψ.tensors[l-1] = Mm1
    end
    R_exs, E
end

function h_matrix(L_ex::Array{T,3}, W::Array{T,4}, R_ex::Array{T,3}) where {T}
    @tensor h[σˡ, aˡ⁻¹, aˡ, σˡ′, aˡ⁻¹′, aˡ′] := L_ex[bˡ⁻¹, aˡ⁻¹, aˡ⁻¹′] * W[bˡ⁻¹, bˡ, σˡ, σˡ′] * R_ex[bˡ, aˡ, aˡ′]
    @cast h[(σˡ, aˡ⁻¹, aˡ), (σˡ′, aˡ⁻¹′, aˡ′)] := h[σˡ, aˡ⁻¹, aˡ, σˡ′, aˡ⁻¹′, aˡ′]
end

function eigenproblem(dir::Direction, M::Array{T, 3}, L_ex::Array{T, 3}, W::Array{T, 4}, R_ex::Array{T, 3}) where {T}
    @cast v[(σˡ, aˡ⁻¹, aˡ)] |= M[aˡ⁻¹, aˡ, σˡ]

    h = h_matrix(L_ex, W, R_ex)
    λ, Φ = eigs(h, v0=v, nev=1, which=:SR)
    E  = λ[1]::T
    v⁰ = (Φ[:,1])::Vector{T}

    (E, split_tensor(dir, v⁰, size(M))...)
end

function split_tensor(::Right, v⁰::Vector, (Dˡ⁻¹, Dˡ, d))
    @cast Mm[(σˡ, aˡ⁻¹), aˡ] := v⁰[(σˡ, aˡ⁻¹, aˡ)] (aˡ⁻¹:Dˡ⁻¹, aˡ:Dˡ, σˡ:d)
    U, S, V = svd(Mm)
    @cast A[aˡ⁻¹, aˡ, σˡ] |= U[(σˡ, aˡ⁻¹), aˡ] (σˡ:d, aˡ⁻¹:Dˡ⁻¹, aˡ:Dˡ)
    A, Diagonal(S)*V'
end

function split_tensor(::Left, v⁰::Vector, (Dˡ⁻¹, Dˡ, d))
    @cast Mm[aˡ⁻¹, (σˡ, aˡ)] |= v⁰[(σˡ, aˡ⁻¹, aˡ)] (aˡ⁻¹:Dˡ⁻¹, aˡ:Dˡ, σˡ:d)
    U, S, V = svd(Mm)
    @cast B[aˡ⁻¹, aˡ, σˡ] |= V'[aˡ⁻¹, (σˡ, aˡ)] (σˡ:d)
    U*Diagonal(S), B
end

function iterate_R_ex(B, W, R_ex) where {T}
    @tensoropt R_ex′[bⁱ⁻¹, aⁱ⁻¹, aⁱ⁻¹′] := (conj.(B))[aⁱ⁻¹,aⁱ,σⁱ] * W[bⁱ⁻¹,bⁱ,σⁱ,σⁱ′] * B[aⁱ⁻¹′,aⁱ′,σⁱ′] * R_ex[bⁱ,aⁱ,aⁱ′]
end

function iterate_L_ex(A, W, L_ex) where {T}
    @tensoropt L_ex′[bˡ, aˡ, aˡ′] := L_ex[bˡ⁻¹,aˡ⁻¹,aˡ⁻¹′] * (conj.(A))[aˡ⁻¹,aˡ,σˡ] * W[bˡ⁻¹,bˡ,σˡ,σˡ′] * A[aˡ⁻¹′,aˡ′,σˡ′]
end

"""
    ground_state(ψ::MPS{L, T}, H::MPO{L, T}; maxiter=10, quiet=false, ϵ=1e-8) where {L, T}

Perform the finite system density matrix renormalization group
algorithm. First this will build up the R expressions, then do right
and left sweeps until either
 1) The state converges to an eigenstate `ϕ` such that
    ϕ' * H * H * ϕ ≈ (ϕ' * H * ϕ)
to the requested tolerance `ϵ`
 2) The energy eigenvalue stops changing (possible signaling the algorithm is
stuck in a local minimum)
 3) The number of full (right and left) sweeps exceeds `maxiter`.

Setting `quiet=true` will suppress notifications about the algorithm's
progress but *not* warnings due to non-convergence.
"""
function ground_state(ψ::MPS{L, T}, H::MPO{L, T}; maxiter=10, quiet=false, ϵ=1e-8) where {L, T}
    ϕ = ψ |> copy

    quiet || println("Computing R expressions")
    R_exs = R_exprs(ψ, H)

    converged = false
    count     = 0
    E₀ = zero(T)
    enable_cache(maxsize=5*10^9)
    while not(converged)
        quiet || println("Performing right sweep")
        L_exs, E₀′ = sweep!(right, ϕ, H, R_exs)

        quiet || println("Performing left sweep")
        R_exs, E₀  = sweep!(left,  ϕ, H, L_exs)

        count += 1
        if iseigenstate(ϕ, H, ϵ=ϵ)
            quiet || println("Converged in $count iterations")
            converged = true
        elseif count > 1 && E₀ ≈ E₀′
                @warn """
Energy eigenvalue converged but state is not an eigenstate.
Consider either lowering your requested tolerance or
implementing a warm-up algorithm to avoid local minima.
"""
            break
        elseif count >= maxiter
            @warn "Did not converge in $maxiter iterations"
            break
        end
    end
    clear_cache()
    ϕ, E₀
end

function iseigenstate(ψ::MPS, H::MPO; ϵ=1e-8)
    ϕ = rightcanonical(ψ)
    isapprox(ϕ' * (H * H * ϕ), (ϕ' * (H * ϕ))^2, rtol=ϵ)
end
# Iterative Ground State Search:1 ends here


text:
### Functions to Compute Pose Graph Error ###
### Ted Steiner ###
### September 2014 ###

# Compute pose position RMS error
function posegraphError( posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1} )
    alignPoses!(posesA,posesB)

    err = poseError(posesA,posesB)

    rmse = sqrt( sum(err.*err) / length(posesA) )
    return rmse
end

function poseError( posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1} )
    err = zeros(length(posesA),3)

    for k = 1:length(posesA)
        err[k,:] = poseError(posesA[k],posesB[k])
    end

    return err
end

function poseError( poseA::Pose2D, poseB::Pose2D )
    err = zeros(3)
    err[1] = poseA.x - poseB.x
    err[2] = poseA.y - poseB.y
    err[3] = 0.0 #poseA.theta - poseB.theta
    return err
end

function posegraphError( truth_filename::String, result_filename::String )
    posesTruth = readiSAMOutput(truth_filename)
    posesResult = readiSAMOutput(result_filename)
    return posegraphError(posesTruth,posesResult)
end

function poseErrorCovTrace( cov::Array{Float64,2} )
    err = zeros(size(cov,1))

    for k = 1:length(err)
        err[k] = sqrt(cov[k,1]) + sqrt(cov[k,2]) + sqrt(cov[k,3])
    end

    return err
end

# Pose Sigma (Pose position error standard deviation)
function poseSigma( cov::Array{Float64,2} )
    return sqrt( cov[1,1] + cov[1,2] )
end

function poseSigmas( covs::Array{Float64,2} )
    err = zeros(size(covs,1))

    for k = 1:length(err)
        err[k] = poseSigma( covs[k,:] )
    end

    return err
end

function epsilonTrajectory( covs::Array{Float64,2} )
    return mean(poseSigmas(covs))
end

# Trajectory Error Epsilon
function errorRatio( mean_sigmas::Array, rounds::Int=Inf )
    if rounds > size(mean_sigmas,1)
        rounds = size(mean_sigmas,1)
    end

    return mean(errorRatios(mean_sigmas,rounds),1)
end

function errorRatios( mean_sigmas::Array, rounds::Int=Inf )
    if rounds > size(mean_sigmas,1)
        rounds = size(mean_sigmas,1)
    end

    anchors = size(mean_sigmas,2)-1
    e_ratios = zeros(rounds,anchors)

    for r = 1:rounds
        e_full = mean_sigmas[r,end]

        for a = 1:anchors
            e_ratios[r,a] = (mean_sigmas[r,a] - e_full) / e_full
        end
    end

    return e_ratios
end

######################
### Pose Alignment ###
######################

function alignPoses!(posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1})
    R, t = kabsch(posesA,posesB)
    transformPoses!(posesB,R',-t)
    return nothing
end

# Kabsch Algorithm: Computes rotation and translation between two sets of points
function kabsch( posesA::Array{Pose2D,1}, posesB::Array{Pose2D,1} )
    A = poses2matrix(posesA)
    B = poses2matrix(posesB)

    return kabsch(A,B)
end

function kabsch( A::Array{Float64,2}, B::Array{Float64,2} )
    assert(size(A,1) == size(B,1))

    centroid_A = mean(A,1)
    centroid_B = mean(B,1)

    H = (A - repmat(centroid_A, size(A,1), 1))' * (B - repmat(centroid_B, size(B,1), 1))
    U,S,V = svd(H)
    R = V*U';

    if det(R) < 0
        println("Reflection detected.")
        V[:,3] *= -1
        R = V*U'
    end

    t = -R*centroid_A' + centroid_B'

    return R, t
end

function transformPoses!( poses::Array{Pose2D,1}, R, t )
    R_theta = atan2(R[2,1],R[1,1])

    for pose in poses
        pose.x += t[1]
        pose.y += t[2]
        pose.theta += R_theta
    end

    return nothing
end

function poses2matrix( poses::Array{Pose2D,1} )
    A = zeros(length(poses),3)

    for k = 1:length(poses)
        A[k,1] = poses[k].x
        A[k,2] = poses[k].y
    end

    return A
end


text:
function generate_animation(genp::GeneratorPyramid, α, β, amplifier, z_rec)
    z = build_zero_pyramid(z_rec, genp.noise_shapes)
    st = Base.length(genp.noise_shapes)

    z_rand = amplifier * randn_like(z_rec, size(z_rec))
    z_prev1 = @. 0.95f0 * z_rec + 0.05f0
    z_prev2 = z_rec

    for i in 1:100
        z_rand = amplifier * randn_like(z_rec, size(z_rec))
        diff_curr = @. β * (z_prev1 - z_prev2) + (1 - β) * z_rand
        z_curr = @. α * z_rec + (1 - α) * (z_prev1 + diff_curr)
        z_prev1 = z_curr
        z_prev2 = z_prev1

        z[1] = z_curr
        img = genp(z, st, false)
        save_array_as_image(animation_savepath(i), view(img, :, :, :, 1))
    end
end

max_stars_repo_name: matsueushi/SinGAN

max_issues_repo_name: matsueushi/SinGAN

max_forks_repo_name: matsueushi/SinGAN


text:
"""
    Base.rand(pp::PoissonProcess, intervallength::Float64) -> Vector{Float64}

Simulate `pp` over the time interval `[0, intervallength]`, and return a vector of the
arrival times.

##### Complexity
- Time complexity: `O(E[N(t)]log(E[N(t)]))`, in expectation, where E[N(t)] denotes the
expected number of arrivals over the time interval `[0, intervallength]`.
"""
function Base.rand(pp::PoissonProcess, intervallength::AbstractFloat)
    # instantiate a Poisson random variable with distribution F1
    F1 = Poisson(nexpectedarrivals(pp, intervallength))
    # sample from F1
    narrivals = rand(F1)
    # no arrivals; return an empty Vector
    if narrivals == 0
        return []
    end

    # instantiate a Uniform random variable with distribution F2
    F2 = Uniform(0, intervallength)
    # allocate memory
    arrivaltimes = Vector{Float64}(undef, narrivals)
    for i in 1:narrivals
        arrivaltimes[i] = rand(F2)
    end
    return sort!(arrivaltimes)
end

"""
    Base.rand(pp::PoissonProcess, intervallength::Integer) -> Vector{Float64}

Allow the user to supply an integer `intervallength`, and work-around method ambiguity
arising from `Random`.
"""
function Base.rand(pp::PoissonProcess, intervallength::Integer)
    return Base.rand(pp, Float64(intervallength))
end

max_stars_repo_name: LoganDGraham/StochasticProcesses

max_issues_repo_name: LoganDGraham/StochasticProcesses

max_forks_repo_name: LoganDGraham/StochasticProcesses


text:
"""
    function single_excitonic_basis(clusters, fspace::FockConfig{N}; R=1, Nk=2, T=Float64) where {N}
"""
function single_excitonic_basis(clusters, fspace::FockConfig{N}; R=1, Nk=2, T=Float64) where {N}
    ci_vector = ClusteredState(clusters, fspace, R=R)
    config = ones(Int,N)
    for ci in 1:N
        config_i = deepcopy(config)
        for k in 1:Nk
            config_i[ci] = k
            ci_vector[fspace][ClusterConfig(config_i)] = zeros(Int,R)
        end
    end
    set_vector!(ci_vector, Matrix{T}(I,size(ci_vector)...))
    return ci_vector
end

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
using Test
using GeoParams

@testset "DislocationCreepLaws" begin

    # This tests the MaterialParameters structure
    CharUnits_GEO = GEO_units(; viscosity=1e19, length=1000km)

    # Define a linear viscous creep law ---------------------------------
    x1 = DislocationCreep()
    @test isbits(x1)
    @test x1.n.val == 1.0
    @test x1.A.val == 1.5

    x2 = DislocationCreep(; n=3)
    @test x2.A.val == 1.5

    # perform a computation with the dislocation creep laws
    # Calculate EpsII, using a set of pre-defined values
    CharDim = GEO_units()
    EpsII = GeoUnit(0s^-1)
    EpsII = nondimensionalize(EpsII, CharDim)
    TauII = GeoUnit(100MPa)
    TauII = nondimensionalize(TauII, CharDim)
    P = GeoUnit(100MPa)
    P = nondimensionalize(P, CharDim)
    T = GeoUnit(500C)
    T = nondimensionalize(T, CharDim)
    f = GeoUnit(50MPa)
    f = nondimensionalize(f, CharDim)
    p = CreepLawVariables(; P=P, T=T, f=f)
    Phase = SetMaterialParams(;
        Name="Viscous Matrix",
        Phase=2,
        Density=ConstantDensity(),
        CreepLaws=DislocationCreep(; n=3NoUnits, r=1NoUnits),
        CharDim=CharDim,
    )
#    εII = computeCreepLaw_EpsII(TauII, Phase.CreepLaws[1], p)
#    @test εII ≈ 2.1263214994323903e-11 rtol = 1e-8

 #   @test εII == computeCreepLaw_EpsII(TauII, Phase.CreepLaws[1], P.val, T.val, f.val)

    # Check that once inverted, we get back the TauII that we used to calculate EpsII
 #   NewTau = computeCreepLaw_TauII(εII, Phase.CreepLaws[1], p)
 #   @test NewTau ≈ TauII.val
 #   @test NewTau == computeCreepLaw_TauII(εII, Phase.CreepLaws[1], P.val, T.val, f.val)

    # Given strainrate
    #@test computeCreepLaw_EpsII(1e-13/s, x1, CreepLawParams())==1e18*2*1e-13Pa       # dimensional input
    # -------------------------------------------------------------------

end


text:
struct DistributionFreeSKCETest{E<:SKCE,B,V} <: HypothesisTests.HypothesisTest
    """Calibration estimator."""
    estimator::E
    """Uniform upper bound of the terms of the estimator."""
    bound::B
    """Number of observations."""
    n::Int
    """Calibration error estimate."""
    estimate::V
end

function DistributionFreeSKCETest(
    estimator::SKCE,
    predictions::AbstractVector,
    targets::AbstractVector;
    bound=uniformbound(estimator),
)
    estimate = estimator(predictions, targets)
    return DistributionFreeSKCETest(estimator, bound, length(predictions), estimate)
end

# HypothesisTests interface

HypothesisTests.default_tail(::DistributionFreeSKCETest) = :right

function HypothesisTests.pvalue(test::DistributionFreeSKCETest{<:BiasedSKCE})
    s = sqrt(test.n * test.estimate / test.bound) - 1
    p = exp(-max(s, zero(s))^2 / 2)
    return p
end

function HypothesisTests.pvalue(
    test::DistributionFreeSKCETest{<:Union{UnbiasedSKCE,BlockUnbiasedSKCE}}
)
    return exp(-div(test.n, 2) * (test.estimate / test.bound)^2 / 2)
end

HypothesisTests.testname(::DistributionFreeSKCETest) = "Distribution-free SKCE test"

# parameter of interest: name, value under H0, point estimate
function HypothesisTests.population_param_of_interest(test::DistributionFreeSKCETest)
    return nameof(typeof(test.estimator)), zero(test.estimate), test.estimate
end

function HypothesisTests.show_params(io::IO, test::DistributionFreeSKCETest, ident="")
    println(io, ident, "number of observations: $(test.n)")
    return println(io, ident, "uniform bound of the terms of the estimator: $(test.bound)")
end

# uniform bound `B_{p;q}` of the absolute value of the terms in the estimators
uniformbound(kce::SKCE) = 2 * uniformbound(kce.kernel)

# uniform bounds of the norm of base kernels
uniformbound(kernel::ExponentialKernel) = 1
uniformbound(kernel::SqExponentialKernel) = 1
uniformbound(kernel::WhiteKernel) = 1

# uniform bound of the norm of a scaled kernel
uniformbound(kernel::ScaledKernel) = first(kernel.σ²) * uniformbound(kernel.kernel)

# uniform bound of the norm of a kernel with input transformations
# assume transform is bijective (i.e., transform does not affect the bound) as default
uniformbound(kernel::TransformedKernel) = uniformbound(kernel.kernel)

# uniform bounds of the norm of tensor product kernels
uniformbound(kernel::KernelTensorProduct) = prod(uniformbound, kernel.kernels)


text:
# Imagenet training script based on https://github.com/pytorch/examples/blob/main/imagenet/main.py

using ArgParse                                          # Parse Arguments from Commandline
using Augmentor                                         # Image Augmentation
using CUDA                                              # GPUs <3
using DataLoaders                                       # Pytorch like DataLoaders
using Dates                                             # Printing current time
using Lux                                               # Neural Network Framework
using FluxMPI                                           # Distibuted Training
using Formatting                                        # Pretty Printing
using Functors                                          # Parameter Manipulation
using Images                                            # Image Processing
using Metalhead                                         # Image Classification Models
using MLDataUtils                                       # Shuffling and Splitting Data
using NNlib                                             # Neural Network Backend
using Optimisers                                        # Collection of Gradient Based Optimisers
using ParameterSchedulers                               # Collection of Schedulers for Parameter Updates
using Random                                            # Make things less Random
using Serialization                                     # Serialize Models
using Setfield                                          # Easy Parameter Manipulation
using Statistics                                        # Statistics
using Zygote                                            # Our AD Engine

import Flux: OneHotArray, onecold, onehot, onehotbatch  # Only being used for OneHotArrays
import DataLoaders: LearnBase                           # Extending Datasets
import MLUtils

# Distributed Training
FluxMPI.Init(; verbose=true)
CUDA.allowscalar(false)

# unsafe_free OneHotArrays
CUDA.unsafe_free!(x::OneHotArray) = CUDA.unsafe_free!(x.indices)

# Image Classification Models
VGG11_BN(args...; kwargs...) = VGG11(args...; batchnorm=true, kwargs...)
VGG13_BN(args...; kwargs...) = VGG13(args...; batchnorm=true, kwargs...)
VGG16_BN(args...; kwargs...) = VGG16(args...; batchnorm=true, kwargs...)
VGG19_BN(args...; kwargs...) = VGG19(args...; batchnorm=true, kwargs...)
MobileNetv3_small(args...; kwargs...) = MobileNetv3(:small, args...; kwargs...)
MobileNetv3_large(args...; kwargs...) = MobileNetv3(:large, args...; kwargs...)
ResNeXt50(args...; kwargs...) = ResNeXt(50, args...; kwargs...)
ResNeXt101(args...; kwargs...) = ResNeXt(101, args...; kwargs...)
ResNeXt152(args...; kwargs...) = ResNeXt(152, args...; kwargs...)

AVAILABLE_IMAGENET_MODELS = [
    AlexNet,
    VGG11,
    VGG13,
    VGG16,
    VGG19,
    VGG11_BN,
    VGG13_BN,
    VGG16_BN,
    VGG19_BN,
    ResNet18,
    ResNet34,
    ResNet50,
    ResNet101,
    ResNet152,
    ResNeXt50,
    ResNeXt101,
    ResNeXt152,
    GoogLeNet,
    DenseNet121,
    DenseNet161,
    DenseNet169,
    DenseNet201,
    MobileNetv1,
    MobileNetv2,
    MobileNetv3_small,
    MobileNetv3_large,
    ConvMixer,
]

IMAGENET_MODELS_DICT = Dict(string(model) => model for model in AVAILABLE_IMAGENET_MODELS)

function get_model(model_name::String, models_dict::Dict, rng, args...; warmup=true,
                   kwargs...)
    model = Lux.transform(models_dict[model_name](args...; kwargs...).layers)
    ps, st = Lux.setup(rng, model) .|> gpu
    if warmup
        # Warmup for compilation
        x__ = randn(rng, Float32, 224, 224, 3, 1) |> gpu
        y__ = onehotbatch([1], 1:1000) |> gpu
        should_log() && println("$(now()) ==> staring `$model_name` warmup...")
        model(x__, ps, st)
        should_log() && println("$(now()) ==> forward pass warmup completed")
        (l, _, _), back = Zygote.pullback(p -> logitcrossentropyloss(x__, y__, model, p, st),
                                          ps)
        back((one(l), nothing, nothing))
        should_log() && println("$(now()) ==> backward pass warmup completed")
    end

    if is_distributed()
        ps = FluxMPI.synchronize!(ps; root_rank=0)
        st = FluxMPI.synchronize!(st; root_rank=0)
        should_log() && println("$(now()) ===> models synced across all ranks")
    end

    return model, ps, st
end

# Parse Training Arguments
function parse_commandline_arguments()
    parse_settings = ArgParseSettings("Lux ImageNet Training")
    @add_arg_table! parse_settings begin
        "--arch"
        default = "ResNet18"
        range_tester = x -> x ∈ keys(IMAGENET_MODELS_DICT)
        help = "model architectures: " * join(keys(IMAGENET_MODELS_DICT), ", ", " or ")
        "--epochs"
        help = "number of total epochs to run"
        arg_type = Int
        default = 90
        "--start-epoch"
        help = "manual epoch number (useful on restarts)"
        arg_type = Int
        default = 0
        "--batch-size"
        help = "mini-batch size, this is the total batch size across all GPUs"
        arg_type = Int
        default = 256
        "--learning-rate"
        help = "initial learning rate"
        arg_type = Float32
        default = 0.1f0
        "--momentum"
        help = "momentum"
        arg_type = Float32
        default = 0.9f0
        "--weight-decay"
        help = "weight decay"
        arg_type = Float32
        default = 1.0f-4
        "--print-freq"
        help = "print frequency"
        arg_type = Int
        default = 10
        "--resume"
        help = "resume from checkpoint"
        arg_type = String
        default = ""
        "--evaluate"
        help = "evaluate model on validation set"
        action = :store_true
        "--pretrained"
        help = "use pre-trained model"
        action = :store_true
        "--seed"
        help = "seed for initializing training. "
        arg_type = Int
        default = 0
        "data"
        help = "path to dataset"
        required = true
    end

    return parse_args(parse_settings)
end

# Loss Function
logitcrossentropyloss(ŷ, y) = mean(-sum(y .* logsoftmax(ŷ; dims=1); dims=1))

function logitcrossentropyloss(x, y, model, ps, st)
    ŷ, st_ = model(x, ps, st)
    return logitcrossentropyloss(ŷ, y), ŷ, st_
end

# Optimisers / Parameter Schedulers
function update_lr(st::ST, eta) where {ST}
    if hasfield(ST, :eta)
        @set! st.eta = eta
    end
    return st
end
update_lr(st::Optimisers.OptimiserChain, eta) = update_lr.(st.opts, eta)
function update_lr(st::Optimisers.Leaf, eta)
    @set! st.rule = update_lr(st.rule, eta)
end
update_lr(st_opt::NamedTuple, eta) = fmap(l -> update_lr(l, eta), st_opt)

# Accuracy
function accuracy(ŷ, y, topk=(1,))
    maxk = maximum(topk)

    pred_labels = partialsortperm.(eachcol(ŷ), (1:maxk,), rev=true)
    true_labels = onecold(y)

    accuracies = Vector{Float32}(undef, length(topk))

    for (i, k) in enumerate(topk)
        accuracies[i] = sum(map((a, b) -> sum(view(a, 1:k) .== b), pred_labels,
                                true_labels))
    end

    return accuracies .* 100 ./ size(y, ndims(y))
end

# Distributed Utils
is_distributed() = FluxMPI.Initialized() && total_workers() > 1
should_log() = !FluxMPI.Initialized() || local_rank() == 0

# Checkpointing
function save_checkpoint(state, is_best, filename="checkpoint.pth.tar")
    if should_log()
        serialize(filename, state)
        if is_best
            cp(filename, "model_best.pth.tar"; force=true)
        end
    end
end

# DataLoading
struct ImageDataset
    image_files::Any
    labels::Any
    mapping::Any
    augmentation_pipeline::Any
    normalization_parameters::Any
end

function ImageDataset(folder::String, augmentation_pipeline, normalization_parameters)
    ulabels = readdir(folder)
    label_dirs = joinpath.((folder,), ulabels)
    @assert length(label_dirs)==1000 "There should be 1000 subdirectories in $folder"

    classes = readlines(joinpath(@__DIR__, "synsets.txt"))
    mapping = Dict(z => i for (i, z) in enumerate(ulabels))

    istrain = endswith(folder, r"train|train/")

    if istrain
        image_files = vcat(map((x, y) -> joinpath.((x,), y), label_dirs,
                               readdir.(label_dirs))...)

        remove_files = [
            "n01739381_1309.JPEG",
            "n02077923_14822.JPEG",
            "n02447366_23489.JPEG",
            "n02492035_15739.JPEG",
            "n02747177_10752.JPEG",
            "n03018349_4028.JPEG",
            "n03062245_4620.JPEG",
            "n03347037_9675.JPEG",
            "n03467068_12171.JPEG",
            "n03529860_11437.JPEG",
            "n03544143_17228.JPEG",
            "n03633091_5218.JPEG",
            "n03710637_5125.JPEG",
            "n03961711_5286.JPEG",
            "n04033995_2932.JPEG",
            "n04258138_17003.JPEG",
            "n04264628_27969.JPEG",
            "n04336792_7448.JPEG",
            "n04371774_5854.JPEG",
            "n04596742_4225.JPEG",
            "n07583066_647.JPEG",
            "n13037406_4650.JPEG",
            "n02105855_2933.JPEG",
        ]
        remove_files = joinpath.((folder,),
                                 joinpath.(first.(rsplit.(remove_files, "_", limit=2)),
                                           remove_files))

        image_files = [setdiff(Set(image_files), Set(remove_files))...]

        labels = [mapping[x] for x in map(x -> x[2], rsplit.(image_files, "/", limit=3))]
    else
        vallist = hcat(split.(readlines(joinpath(@__DIR__, "val_list.txt")))...)
        labels = parse.(Int, vallist[2, :]) .+ 1
        filenames = [joinpath(classes[l], vallist[1, i]) for (i, l) in enumerate(labels)]
        image_files = joinpath.((folder,), filenames)
        idxs = findall(isfile, image_files)
        image_files = image_files[idxs]
        labels = labels[idxs]
    end

    return ImageDataset(image_files, labels, mapping, augmentation_pipeline,
                        normalization_parameters)
end

LearnBase.nobs(data::ImageDataset) = length(data.image_files)

function LearnBase.getobs(data::ImageDataset, i::Int)
    img = Images.load(data.image_files[i])
    img = augment(img, data.augmentation_pipeline)
    cimg = channelview(img)
    if ndims(cimg) == 2
        cimg = reshape(cimg, 1, size(cimg, 1), size(cimg, 2))
        cimg = vcat(cimg, cimg, cimg)
    end
    img = Float32.(permutedims(cimg, (3, 2, 1)))
    img = (img .- data.normalization_parameters.mean) ./ data.normalization_parameters.std
    return img, onehot(data.labels[i], 1:1000)
end

MLUtils.numobs(data::ImageDataset) = length(data.image_files)

MLUtils.getobs(data::ImageDataset, i::Int) = LearnBase.getobs(data, i)

## DataLoaders doesn't yet work with MLUtils
LearnBase.nobs(data::DistributedDataContainer) = MLUtils.numobs(data)

LearnBase.getobs(data::DistributedDataContainer, i::Int) = MLUtils.getobs(data, i)

# Tracking
Base.@kwdef mutable struct AverageMeter
    fmtstr::Any
    val::Float64 = 0.0
    sum::Float64 = 0.0
    count::Int = 0
    average::Float64 = 0
end

function AverageMeter(name::String, fmt::String)
    fmtstr = FormatExpr("$name {1:$fmt} ({2:$fmt})")
    return AverageMeter(; fmtstr=fmtstr)
end

function update!(meter::AverageMeter, val, n::Int)
    meter.val = val
    meter.sum += val * n
    meter.count += n
    meter.average = meter.sum / meter.count
    return meter.average
end

print_meter(meter::AverageMeter) = printfmt(meter.fmtstr, meter.val, meter.average)

struct ProgressMeter{N}
    batch_fmtstr::Any
    meters::NTuple{N, AverageMeter}
end

function ProgressMeter(num_batches::Int, meters::NTuple{N}, prefix::String="") where {N}
    fmt = "%" * string(length(string(num_batches))) * "d"
    prefix = prefix != "" ? endswith(prefix, " ") ? prefix : prefix * " " : ""
    batch_fmtstr = generate_formatter("$prefix[$fmt/" * sprintf1(fmt, num_batches) * "]")
    return ProgressMeter{N}(batch_fmtstr, meters)
end

function print_meter(meter::ProgressMeter, batch::Int)
    base_str = meter.batch_fmtstr(batch)
    print(base_str)
    foreach(x -> (print("\t"); print_meter(x)), meter.meters[1:end])
    return println()
end

# Validation
function validate(val_loader, model, ps, st, args)
    batch_time = AverageMeter("Batch Time", "6.3f")
    data_time = AverageMeter("Data Time", "6.3f")
    forward_time = AverageMeter("Forward Pass Time", "6.3f")
    losses = AverageMeter("Loss", ".4f")
    top1 = AverageMeter("Acc@1", "6.2f")
    top5 = AverageMeter("Acc@5", "6.2f")

    progress = ProgressMeter(length(val_loader),
                             (batch_time, data_time, forward_time, losses, top1, top5),
                             "Val:")

    st_ = Lux.testmode(st)
    t = time()
    for (i, (x, y)) in enumerate(CuIterator(val_loader))
        t_data, t = time() - t, time()

        # Compute Output
        ŷ, st_ = model(x, ps, st_)
        loss = logitcrossentropyloss(ŷ, y)
        t_forward = time() - t

        # Metrics
        acc1, acc5 = accuracy(cpu(ŷ), cpu(y), (1, 5))
        update!(top1, acc1, size(x, ndims(x)))
        update!(top5, acc5, size(x, ndims(x)))
        update!(losses, loss, size(x, ndims(x)))

        # Measure Elapsed Time
        update!(data_time, t_data, size(x, ndims(x)))
        update!(forward_time, t_forward, size(x, ndims(x)))
        update!(batch_time, t_data + t_forward, size(x, ndims(x)))

        # Print Progress
        if i % args["print-freq"] == 0 || i == length(val_loader)
            should_log() && print_meter(progress, i)
        end

        t = time()
    end

    return top1.average, top5.average, losses.average
end

# Training
function train(train_loader, model, ps, st, optimiser_state, epoch, args)
    batch_time = AverageMeter("Batch Time", "6.3f")
    data_time = AverageMeter("Data Time", "6.3f")
    forward_time = AverageMeter("Forward Pass Time", "6.3f")
    backward_time = AverageMeter("Backward Pass Time", "6.3f")
    optimize_time = AverageMeter("Optimize Time", "6.3f")
    losses = AverageMeter("Loss", ".4e")
    top1 = AverageMeter("Acc@1", "6.2f")
    top5 = AverageMeter("Acc@5", "6.2f")
    progress = ProgressMeter(length(train_loader),
                             (batch_time, data_time, forward_time, backward_time,
                              optimize_time, losses, top1, top5), "Epoch: [$epoch]")

    st = Lux.trainmode(st)

    t = time()
    for (i, (x, y)) in enumerate(CuIterator(train_loader))
        t_data, t = time() - t, time()

        # Gradients and Update
        (loss, ŷ, st), back = Zygote.pullback(p -> logitcrossentropyloss(x, y, model, p,
                                                                          st), ps)
        t_forward, t = time() - t, time()
        gs = back((one(loss) / total_workers(), nothing, nothing))[1]
        t_backward, t = time() - t, time()
        if is_distributed()
            gs = FluxMPI.allreduce_gradients(gs)
        end
        optimiser_state, ps = Optimisers.update!(optimiser_state, ps, gs)
        t_opt = time() - t

        # Metrics
        acc1, acc5 = accuracy(cpu(ŷ), cpu(y), (1, 5))
        update!(top1, acc1, size(x, ndims(x)))
        update!(top5, acc5, size(x, ndims(x)))
        update!(losses, loss, size(x, ndims(x)))

        # Measure Elapsed Time
        update!(data_time, t_data, size(x, ndims(x)))
        update!(forward_time, t_forward, size(x, ndims(x)))
        update!(backward_time, t_backward, size(x, ndims(x)))
        update!(optimize_time, t_opt, size(x, ndims(x)))
        update!(batch_time, t_data + t_forward + t_backward + t_opt, 1)

        # Print Progress
        if i % args["print-freq"] == 0 || i == length(train_loader)
            should_log() && print_meter(progress, i)
        end

        t = time()
    end

    return ps, st, optimiser_state, (top1.average, top5.average, losses.average)
end

# Main Function
function main(args)
    best_acc1 = 0

    # Seeding
    rng = Random.default_rng()
    Random.seed!(rng, args["seed"])

    # Model Construction
    if should_log()
        if args["pretrained"]
            println("$(now()) => using pre-trained model `$(args["arch"])`")
        else
            println("$(now()) => creating model `$(args["arch"])`")
        end
    end
    model, ps, st = get_model(args["arch"], IMAGENET_MODELS_DICT, rng; warmup=true,
                              pretrain=args["pretrained"])

    normalization_parameters = (mean=reshape([0.485f0, 0.456f0, 0.406f0], 1, 1, 3),
                                std=reshape([0.229f0, 0.224f0, 0.225f0], 1, 1, 3))
    train_data_augmentation = Resize(256, 256) |> FlipX(0.5) |> RCropSize(224, 224)
    val_data_augmentation = Resize(256, 256) |> CropSize(224, 224)
    train_dataset = ImageDataset(joinpath(args["data"], "train"),
                                 train_data_augmentation,
                                 normalization_parameters)
    val_dataset = ImageDataset(joinpath(args["data"], "val"),
                               val_data_augmentation,
                               normalization_parameters)
    if is_distributed()
        train_dataset = DistributedDataContainer(train_dataset)
        val_dataset = DistributedDataContainer(val_dataset)
    end

    train_loader = DataLoader(shuffleobs(train_dataset),
                              args["batch-size"] ÷ total_workers())
    val_loader = DataLoader(val_dataset, args["batch-size"] ÷ total_workers())

    # Optimizer and Scheduler
    should_log() && println("$(now()) => creating optimiser")
    optimiser = Optimisers.OptimiserChain(Optimisers.Momentum(args["learning-rate"],
                                                              args["momentum"]),
                                          Optimisers.WeightDecay(args["weight-decay"]))
    optimiser_state = Optimisers.setup(optimiser, ps)
    if is_distributed()
        optimiser_state = FluxMPI.synchronize!(optimiser_state)
        should_log() && println("$(now()) ==> synced optimiser state across all ranks")
    end
    scheduler = Step(λ=args["learning-rate"], γ=0.1f0, step_sizes=30)

    if args["resume"] != ""
        if isfile(args["resume"])
            checkpoint = deserialize(args["resume"])
            args["start-epoch"] = checkpoint["epoch"]
            optimiser_state = checkpoint["optimiser_state"] |> gpu
            ps = checkpoint["model_parameters"] |> gpu
            st = checkpoint["model_states"] |> gpu
            should_log() &&
                println("$(now()) => loaded checkpoint `$(args["resume"])` (epoch $(args["start-epoch"]))")
        else
            should_log() &&
                println("$(now()) => no checkpoint found at `$(args["resume"])`")
        end
    end

    if args["evaluate"]
        @assert !is_distributed() "We are not syncing statistics. For evaluation run on 1 process"
        validate(val_loader, model, ps, st, args)
        return
    end

    GC.gc(true)
    CUDA.reclaim()

    for epoch in args["start-epoch"]:args["epochs"]
        # Train for 1 epoch
        ps, st, optimiser_state, _ = train(train_loader, model, ps, st, optimiser_state,
                                           epoch, args)

        # Some Housekeeping
        GC.gc(true)
        CUDA.reclaim()

        # Evaluate on validation set
        acc1, _, _ = validate(val_loader, model, ps, st, args)

        # ParameterSchedulers
        eta_new = scheduler(epoch)
        optimiser_state = update_lr(optimiser_state, eta_new)

        # Some Housekeeping
        GC.gc(true)
        CUDA.reclaim()

        # Remember Best Accuracy and Save Checkpoint
        is_best = acc1 > best_acc1
        best_acc1 = max(acc1, best_acc1)

        save_state = Dict("epoch" => epoch,
                          "arch" => args["arch"],
                          "model_states" => st |> cpu,
                          "model_parameters" => ps |> cpu,
                          "optimiser_state" => optimiser_state |> cpu)
        save_checkpoint(save_state, is_best)
    end
end

main(parse_commandline_arguments())


text:
using OrdinaryDiffEq, ModelingToolkit, Test, SparseArrays

N = 3
xyd_brusselator = range(0, stop = 1, length = N)
brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0
limit(a, N) = ModelingToolkit.ifelse(a == N + 1, 1, ModelingToolkit.ifelse(a == 0, N, a))
function brusselator_2d_loop(du, u, p, t)
    A, B, alpha, dx = p
    alpha = alpha / dx^2
    @inbounds for I in CartesianIndices((N, N))
        i, j = Tuple(I)
        x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
        ip1, im1, jp1, jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N),
                             limit(j - 1, N)
        du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] -
                       4u[i, j, 1]) +
                      B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] +
                      brusselator_f(x, y, t)
        du[i, j, 2] = alpha * (u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] -
                       4u[i, j, 2]) +
                      A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2]
    end
end

# Test with tuple parameters
p = (3.4, 1.0, 10.0, step(xyd_brusselator))

function init_brusselator_2d(xyd)
    N = length(xyd)
    u = zeros(N, N, 2)
    for I in CartesianIndices((N, N))
        x = xyd[I[1]]
        y = xyd[I[2]]
        u[I, 1] = 22 * (y * (1 - y))^(3 / 2)
        u[I, 2] = 27 * (x * (1 - x))^(3 / 2)
    end
    u
end

u0 = init_brusselator_2d(xyd_brusselator)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
                                     u0, (0.0, 11.5), p)
sys = modelingtoolkitize(prob_ode_brusselator_2d)

# test sparse jacobian pattern only.
prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = false)
JP = prob.f.jac_prototype
@test findnz(Symbolics.jacobian_sparsity(map(x -> x.rhs, equations(sys)), states(sys)))[1:2] ==
      findnz(JP)[1:2]

# test sparse jacobian
prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
@test findnz(calculate_jacobian(sys, sparse = true))[1:2] ==
      findnz(prob.f.jac_prototype)[1:2]

# test when not sparse
prob = ODEProblem(sys, u0, (0, 11.5), sparse = false, jac = true)
@test prob.f.jac_prototype == nothing

prob = ODEProblem(sys, u0, (0, 11.5), sparse = false, jac = false)
@test prob.f.jac_prototype == nothing

# test when u0 is nothing
f = DiffEqBase.ODEFunction(sys, u0 = nothing, sparse = true, jac = true)
@test findnz(f.jac_prototype)[1:2] == findnz(JP)[1:2]
@test eltype(f.jac_prototype) == Float64

f = DiffEqBase.ODEFunction(sys, u0 = nothing, sparse = true, jac = false)
@test findnz(f.jac_prototype)[1:2] == findnz(JP)[1:2]
@test eltype(f.jac_prototype) == Float64

# test when u0 is not Float64
u0 = similar(init_brusselator_2d(xyd_brusselator), Float32)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,
                                     u0, (0.0, 11.5), p)
sys = modelingtoolkitize(prob_ode_brusselator_2d)

prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = false)
@test eltype(prob.f.jac_prototype) == Float32

prob = ODEProblem(sys, u0, (0, 11.5), sparse = true, jac = true)
@test eltype(prob.f.jac_prototype) == Float32


text:
# see examples.jl first

# Extractor
xml = parsexml("""
<?xml version="1.0" encoding="UTF-8"?>
<university university-name="Julia University">
  <person id="1">
    <age>24</age>
    <study-field>Mechanical Engineering</study-field>
    <GPA>4.5</GPA>
    <taken-courses>Artificial Intelligence</taken-courses>
    <taken-courses>Robotics</taken-courses>
    He is a genius
  </person>
  <person id="2">
    <age>18</age>
    <study-field>Computer Engineering</study-field>
    <GPA>4.2</GPA>
    <taken-courses>Julia</taken-courses>
  </person>
</university>
""")

# extract University
U = University(xml) # StructName(xml) extracts the data and stores them in proper format

# Now you can access all of the data by calling the fieldnames

U.name # Julia University

# extract Person

P1 = U.people[1]

P1.comment
P1.age # 24
P1.field # Mechanical Engineering
P1.GPA # 4.5
P1.courses # ["Artificial Intelligence", "Robotics"]
P1.id # 1


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Julia 1.5.2
#     language: julia
#     name: julia-1.5
# ---

using WebIO
using JSExpr

# +
s = Scope(
      imports = [
        "./vars.js",
        "./mystyle.css",
    　],
)

b = Node(
    :button,
    "press",
    className="button",
    events=Dict(
        "click" => js"""
        function () {
            const old = document.querySelector(".button")
                                 .textContent;
            if(old == "pressed"){
                document.querySelector(".button")
                        .textContent = "press";
            } else {
                document.querySelector(".button")
                        .textContent = "pressed";
            }
            console.log(myVar);
            console.log(f(3));
        }
        """
    )
)

d = Node(
    :div,
    "A",
    className="rectangle",
)

s(Node(:div, b, d))
# -

max_stars_repo_name: terasakisatoshi/jsexample

max_issues_repo_name: terasakisatoshi/jsexample

max_forks_repo_name: terasakisatoshi/jsexample


text:
module TwoElectronTDDFT

using LinearAlgebra

include("abstract_types.jl")
include("operator.jl")
include("Potential.jl")
include("Hamiltonian.jl")

end


text:
# Arvind Prasadan
# 2019 January
# DMF Package
# Generates a realization of an ARMA process

"""
	gen_arma_sequence(n = 100, ar_comp = [], ma_comp = [], arma_std = 1.0)

Generates a realization of an ARMA process

# Arguments
- `n`: Length of process
- `ar_comp`: List of AR Coefficients; one of this and ma_comp must be non-empty
- `ma_comp`: List of MA Coefficients; one of this and ar_comp must be non-empty
- `arma_std`: Noise Standard Deviation

# Outputs
- `x`: vector of ARMA process realization
"""
function gen_arma_sequence(n = 100, ar_comp = [], ma_comp = [], arma_std = 1.0)

    n = round(abs(n[1])) # Samples
    n = (1 <= n) ? n : 1

    arma_std = abs(arma_std[1]) # Standard Deviation

    if 0 == length(ar_comp) + length(ma_comp)
	x = arma_std * randn(n)
    else
    	# Define ARMA object, simulate from this set of parameters
    	arma = ARMA(ar_comp, ma_comp, arma_std)
    	x = simulation(arma, ts_length = n)
    end

    return x
end


text:
"""
    mpifit_hsvgp(get_data, n_parts, n_dims, bounds_low, bounds_high;
    ni_low     = 5,
    ni_high    = 50,
    batch_size = 20,
    n_iters    = 200,
    n_subcycle_ELBO_local  = 20,
    n_subcycle_ELBO_global = 20)

Fit HSVGP to data for distributed computation

    get_data:                function to define how a local partition should read in its data
    ni_low:                  number of local inducing points
    ni_high:                 number of global inducing points
    batch_size:              size of batches for stochastic optimizer
    n_iters:                 number of optimization steps
    n_subcycle_ELBO_local:   subcycles for local  ELBO before MPI communication
    n_subcycle_ELBO_global:  subcycles for global ELBO before MPI communication
    n_batch:                 number of partitions in a local to global communication batch

WARNING: Currently writes out several jld2 files containing local and global GPss
"""
function mpifit_hsvgp(get_data, n_parts, n_dims, bounds_low, bounds_high;
    ni_low     = 5,               # Number of local inducing points
    ni_high    = 50,              # Number of global inducing points
    batch_size = 20,              # Number of points used in updating local ELBO
    n_iters    = 200,             # Total updates in whole cycle
    n_subcycle_ELBO_local  = 20,  # subcycles for local ELBO before comm.
    n_subcycle_ELBO_global = 20,  # subcycles for global ELBO before comm.
    n_batch    = 5)               # Number of partitions selected at each iteration

    #TODO:
    #     Need to clean up saving out of local and global objects at the end.
    #     It's messy now and will get worse with many MPI processes.
    #
    #     Would be nice to write a function for prediction to identify which
    #     partition a new value of X would be in, then direct the prediction to
    #     that local svgp
    #
    #     Need to generalize the interface here. Should have a way to pass in
    #     the structure of data to be worked with. Maybe pass in a get_data
    #     function that returns the x,y data for each partition.
    #
    #     Would like to clean up some of those arguments also. I dont think
    #     we should need to have to pass all that stuff in, but not sure the
    #     best way to give global bounds information without simply specifying
    #     it for cases where all the data cannot be stored on one node.
    #

    grad_clip = 1.e-3

    MPI.Init()
    comm = MPI.COMM_WORLD
    myrank = MPI.Comm_rank(comm)
    mysize = MPI.Comm_size(comm)

    root = 0 # global node

    # Fixing this for Gaussian case for now.
    #     Need to generalize this for the future but that will come in part of a large generalization of the code
    # n_latent = 1

    total_inducing = n_parts * ni_low

    part_split = [[ii] for ii in 1:(mysize-1)]

    if n_parts > mysize - 1
        for ii = mysize:n_parts
            append!(part_split[1 + ((ii-1) % (mysize - 1))], ii)
        end
    end

    ######################## DATA SPLIT ########################################
    # This part need not be done separately. Data will already be stored on nodes

    # Root process is the global node and rest of the processes are local processes
    # datarow goes from 2 to 48601
    # numrowscsv = 48600
    # numrowseachrank = Int(numrowscsv/(mysize-1))

    if myrank != root

        data_array = [get_data(part) for part in part_split[myrank]]
        len        = length(part_split[myrank])

    end # rank != root

    # Initializing "trace" array to store ELBO values
    if myrank == root
        trace_h       = zeros(n_iters)          # To collect the global ELBO value
        store_trace_l = zeros(n_iters,n_parts) # Store trace_l from all proesses
    end
    if myrank != root
        trace_l = [zeros(n_iters) for ii in 1:len] # To collect the local ELBO value
    end

    # Now that each local node has seen the data, it will send xi_low, cov_low,
    # mean_low to global node.
    # Once the global node gets data from local nodes, it will evaluate xi_high,
    # cov_high, mean_high and theta_high
    # Following that, global node will calculate prior_mean and prior_cov and
    # send it back to local nodes.

    ###################### INITIALIZATION ######################################

    if myrank != root
        # Step 1
        # Build local SVGPs ###################
        local_svgps = [SVGP_obj(data_array[ii][1], float(data_array[ii][2]), ni_low) for ii in 1:len]

        for ii in 1:len
            # Step 2
            # Sending mean_factor, xi_low, cov_low, mean_low and N to global #######
            send_inducing_points_to_global(local_svgps[ii].params.inducing_locs, local_svgps[ii].params.inducing_mean, root, comm)
        end

        # Step 3
        # Initialize local optimizers
        lopt_cmean = [Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for ii in 1:len]
        lopt_lrho  = [Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for ii in 1:len]
        lopt_lkap  = [Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for ii in 1:len]
        lopt_lsig  = [Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075)) for ii in 1:len]
        loptm      = [Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.01)  ) for ii in 1:len]
        loptS      = [Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.002) ) for ii in 1:len]
        loptx      = [Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.01)  ) for ii in 1:len]

    end # myrank != root

    if myrank == root
        wts = StatsBase.AnalyticWeights(ones(n_parts));

        # Step 3
            # Initialize global svgp and optimizers
            # Root process will receive inducing point data
            # Root process will update global parameters
            # Root process will calculate prior_mean prior_cov for local nodes
            # Root process will send prior_mean prior_cov back to local processes

        # Initialize global optimizers
        opt_cmean = Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075))
        opt_lrho  = Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075))
        opt_lkap  = Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075))
        opt_lsig  = Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.0075))
        optm      = Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.01)  )
        optS      = Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.002) )
        optx      = Flux.Optimiser(Flux.ClipValue(grad_clip), ADAM(1. * 0.01)  )

        # Initialize global svgp
        # TODO:
        #     Need bounds on input space and some idea of scale for output
        #         to set initial GP variance parameter and
        #         Latin Hypercube for inducing point locations
        #
        #     Assuming we will work in the case that we can't send all inducing point
        #         data to global. Instead should probably get bounds as arguments to original
        #         function or otherwise.
        fake_data_to_init = SVGP_data(
            vcat(bounds_low', bounds_high'),
            randn(4)*10 # sampling 4 points with something like the variance I was to initialize at
        )
        global_svgp = SVGP_obj(ni_high, SVGP_data("global"), SVGP_params(fake_data_to_init, ni_high), "gaussian")

        for part = (root+1):n_parts
            recv_id = 1 + (part - 1) % (mysize - 1)
            # TODO: Would prefer to not do this sequentially by rank number but as communications arrive
            #

            # println("In Process(1): ",recv_id,"\n")

            # Receiving xi_low and yi_low #
            xi_low, yi_low = receive_inducing_points_in_global(ni_low, n_dims, recv_id, comm) ### TODO: IMPORTANT! Need to communicate which partition is being recieved

            # Update global parameters from local data #
            grads  = gradient(gps -> svgp_elbo_global(
                    xi_low,
                    yi_low,
                    gps,
                    global_svgp
                    ), global_svgp.params)[1][]

            Flux.Optimise.update!(opt_cmean, global_svgp.params.const_mean,    -grads.const_mean)
            Flux.Optimise.update!(opt_lrho,  global_svgp.params.log_rho,       -grads.log_rho)
            Flux.Optimise.update!(opt_lkap,  global_svgp.params.log_kappa,     -grads.log_kappa)
            Flux.Optimise.update!(opt_lsig,  global_svgp.params.log_sigma,     -grads.log_sigma)
            Flux.Optimise.update!(optm,      global_svgp.params.inducing_mean, -grads.inducing_mean)
            Flux.Optimise.update!(optS,      global_svgp.params.inducing_L,    -LowerTriangular(grads.inducing_L) )
            Flux.Optimise.update!(optx,      global_svgp.params.inducing_locs, -grads.inducing_locs)
            # Sending prior_mean, prior_cov to local process ###################
            send_prior_to_local( pred_vgp(xi_low, global_svgp.params), recv_id, comm)

        end # for loop
    end

    MPI.Barrier(comm)
    if myrank != root
        #Get prior information from globa
        prior_mean, prior_sd = receive_prior_in_local(
                                                        ni_low,
                                                        root,
                                                        comm)
    end # myrank == root

    ###################### END OF INITIALIZATION ###############################

    ##################### OUTER LOOP ###########################################

    for iter= 1:n_iters
        if mod(iter, 2000) == 0
            print(iter, " ")
        end

        if myrank != root

            for ii in 1:len ## TODO: Clean this up with Kelin to process non-sequentially
                            ## TODO: ENSURE THE MPI PROPERLY ASSOCIATED DATA COMMUNICATED TO CORRECT PARTITION
                # Step 4
                # Each local process will receive prior_mean and prior_cov

                # Receiving prior_mean, prior_cov to local process #################
                # prior_mean, prior_sd = receive_prior_in_local(
                #                                        ni_low,
                #                                        root,
                #                                        comm)
                # println("prior_mean, prior_cov received in local")

                # Step 5
                # Update local svgp using prior_mean prior_cov from global
                nsteps_curr  = rand(Distributions.Poisson(n_subcycle_ELBO_local)) + 3;
                for sub_iter = 1:nsteps_curr
                    inds   = rand(1:size(local_svgps[ii].data.y)[1], batch_size)

                    # Update global parameters from local data #
                    grads  = gradient(gps -> svgp_elbo_local(
                            local_svgps[ii].data.x[inds,:],
                            local_svgps[ii].data.y[inds],
                            gps,
                            local_svgps[ii],
                            (prior_mean, prior_sd)
                            ), local_svgps[ii].params)[1][]

                    Flux.Optimise.update!(lopt_cmean[ii], local_svgps[ii].params.const_mean,    -grads.const_mean)
                    Flux.Optimise.update!(lopt_lrho[ii],  local_svgps[ii].params.log_rho,       -grads.log_rho)
                    Flux.Optimise.update!(lopt_lkap[ii],  local_svgps[ii].params.log_kappa,     -grads.log_kappa)
                    Flux.Optimise.update!(lopt_lsig[ii],  local_svgps[ii].params.log_sigma,     -grads.log_sigma)
                    Flux.Optimise.update!(loptm[ii],      local_svgps[ii].params.inducing_mean, -grads.inducing_mean)
                    Flux.Optimise.update!(loptS[ii],      local_svgps[ii].params.inducing_L,    -LowerTriangular(grads.inducing_L) )
                    Flux.Optimise.update!(loptx[ii],      local_svgps[ii].params.inducing_locs, -grads.inducing_locs)
                    # # Sending prior_mean, prior_cov to local process ###################
                    # send_prior_to_local( pred_vgp(xi_low, global_svgp.params[jj]), recv_id, comm)

                    if sub_iter == nsteps_curr
                        # Collecting local ELBO in an array ###########################
                        trace_l[ii][iter] = svgp_elbo_local(
                            local_svgps[ii].data.x[inds,:],
                            local_svgps[ii].data.y[inds],
                            local_svgps[ii].params,
                            local_svgps[ii],
                            (prior_mean, prior_sd)
                        )
                    end # if sub_iter
                end # for sub_iter
                # println("local ELBO optimized")

                # THIS IS NOW DONE OUTSIDE THE ii LOOP
                # Step 6
                # Sending inducing point data to global
                # for jj in 1:n_latent
                #     send_inducing_points_to_global(local_svgps[ii].params[jj].inducing_locs, local_svgps[ii].params[jj].inducing_mean, root, comm)
                # end
            end # for ii

            # Kelin block for sending inducing point data to global only if you were selected
            is_message, status = MPI.Iprobe(0, 8, comm)
            is_message, status = MPI.Iprobe(0, 8, comm) # Kelin: Silly solution, but try this for now

            msg  = [1];
            msg2 = ones(1);

            if is_message
                MPI.Recv!(msg, 0, 8, comm);
                global msg = convert(Vector{Int64}, msg)
                global msg2 = convert(Vector{Int64}, ones(msg[1]));
                #println("process ", myrank, "recieved first message ", msg);
                MPI.Recv!(msg2, 0, 10+myrank, comm);
                global msg2 = msg2;
                #println("process ", myrank, "recieved second message ", msg2);

                indx = findall(in(part_split[myrank]), msg2)
                for ii in indx
                    send_inducing_points_to_global(local_svgps[ii].params.inducing_locs, local_svgps[ii].params.inducing_mean, root, comm)
                end
               #println("process ", myrank, "sent inducing points")

               #Get prior information from global
               prior_mean, prior_sd = receive_prior_in_local(
                                                        ni_low,
                                                        root,
                                                        comm)
            end # if is_message
        end # if myrank != root

        if myrank == root
            #Select random local partition for batch
            aa = collect(1:n_parts);
            selection = convert(Vector{Int64}, zeros(n_batch));
            StatsBase.efraimidis_ares_wsample_norep!(aa, wts, selection)
            part_split_curr = [intersect(part_split[ii], selection) for ii in 1:(mysize-1)];
            wts[selection] .= 0;
            wts .+= 1;

            for ii in 1:(mysize-1)
               msg_to_send = part_split_curr[ii]
               if length(msg_to_send) > 0
                  MPI.Isend([length(msg_to_send)], ii, 8, comm)
                  MPI.Isend(msg_to_send, ii, 10+ii, comm);
                  #println("sending message ", msg_to_send, "to process ", ii)
               end
            end

            step_elbo = 0.
            for part = 1:n_batch
                status  = MPI.Probe(MPI.MPI_ANY_SOURCE, MPI.MPI_ANY_TAG, comm)
                xi_low  = zeros(ni_low, n_dims)

                rmsg    = MPI.Recv!(xi_low, status.source, 0, comm)
                yi_low  = zeros(ni_low)

                rmsg    = MPI.Recv!(yi_low, status.source, 1, comm)
                recv_id = status.source

                # Update global parameters from local data #
                for sub_iter = 1:n_subcycle_ELBO_global
                    grads  = gradient(gps -> svgp_elbo_global(
                            xi_low,
                            yi_low,
                            gps,
                            global_svgp
                            ), global_svgp.params)[1][]

                    Flux.Optimise.update!(opt_cmean, global_svgp.params.const_mean,    -grads.const_mean)
                    Flux.Optimise.update!(opt_lrho,  global_svgp.params.log_rho,       -grads.log_rho)
                    Flux.Optimise.update!(opt_lkap,  global_svgp.params.log_kappa,     -grads.log_kappa)
                    Flux.Optimise.update!(opt_lsig,  global_svgp.params.log_sigma,     -grads.log_sigma)
                    Flux.Optimise.update!(optm,      global_svgp.params.inducing_mean, -grads.inducing_mean)
                    Flux.Optimise.update!(optS,      global_svgp.params.inducing_L,    -LowerTriangular(grads.inducing_L) )
                    Flux.Optimise.update!(optx,      global_svgp.params.inducing_locs, -grads.inducing_locs)
                end # for sub_iter
                # println("global ELBO optimized")

                # Update step_elbo
                step_elbo += svgp_elbo_global(xi_low, yi_low, global_svgp.params, global_svgp)

                # Sending prior_mean, prior_cov to local process ###################
                send_prior_to_local( pred_vgp(xi_low, global_svgp.params), recv_id, comm) # MAKE SURE COMM. INDICATES WHICH PARTITION IN A PROCESS THIS COMM IS FOR
            end # for part

            trace_h[iter] = step_elbo
        end # myrank == root

        barrier_req = MPI.Ibarrier(comm);
        barrier_flag, barrier_status = MPI.Test!(barrier_req);
        while !barrier_flag
            is_message, status = MPI.Iprobe(0, 8, comm)
            is_message, status = MPI.Iprobe(0, 8, comm) #Silly solution, but try this for now
            msg  = [1];
            msg2 = ones(1);
            #println("process ", myrank, "message recieved? ", is_message)
            if is_message
                MPI.Recv!(msg, 0, 8, comm);
                global msg = convert(Vector{Int64}, msg);
                global msg2 = convert(Vector{Int64}, ones(msg[1]));
                #println("process ", myrank, "recieved first message ", msg);
                MPI.Recv!(msg2, 0, 10+myrank, comm);
                global msg2 = msg2;
                #println("process ", myrank, "recieved second message ", msg2);

                indx = findall(in(part_split[myrank]), msg2)
                #println("process ", myrank, "indx = ", indx);
                for ii in indx
                    send_inducing_points_to_global(local_svgps[ii].params.inducing_locs, local_svgps[ii].params.inducing_mean, root, comm)
                end
                #println("process ", myrank, "sent inducing points")

                #Get prior information from global
                prior_mean, prior_sd = receive_prior_in_local(ni_low, root, comm)
            end
            barrier_flag, barrier_status = MPI.Test!(barrier_req);
        end # while !barrier_flag
    end # outeriter

    ############################################################################
    # Returning trace_h and trace_l to main function
    if myrank != root

        for ii in 1:len
            send_trace(trace_l[ii], root ,comm)
        end

    end #myrank != root

    if myrank == root

        for part = (root+1):n_parts
            recv_id = 1 + (part - 1) % (mysize - 1)
            trace_l = receive_trace(n_iters, recv_id, comm)
            store_trace_l[:,part] = trace_l
        end # recv_id

        return_trace = hcat(trace_h, store_trace_l)
        # println(return_trace)

        # Write trace to a csv file
        df_return_trace = DataFrame(return_trace)
        CSV.write("./return_trace.csv", df_return_trace)

    end # myrank == root
    ############################################################################

    ############################################################################
    # Save svgp_params
    if myrank != root

        # send_svgp_params(local_svgp.params, root ,comm)
        for ii in 1:len
            save(string("local_fitted_parameters",part_split[myrank][ii],"_",mysize-1, ".jld2"), string("local_svgp",part_split[myrank][ii]), local_svgps[ii])
        end

    end #myrank != root

    if myrank == root

        # local_svgp_params = [receive_svgp_params(recv_id, comm) for recv_id = 1:mysize-1]

        save(string("global_fitted_svgp_",mysize-1,".jld2"), "global_svgp", global_svgp)

    end # myrank == root
    ############################################################################

    MPI.Finalize()

end     # mpifit_HSVGP

max_stars_repo_name: ayanbiswas/PRISM

max_issues_repo_name: ayanbiswas/PRISM

max_forks_repo_name: ayanbiswas/PRISM


text:
Base.angle(d1, d2) = acos(min(dot(normalize(d1), normalize(d2)), 1))
Base.angle(a::T, b::T) where {T<:Union{CalibratedHit, PMT, MCTrack}} = Base.angle(a.dir, b.dir)
Base.angle(a, b::Union{CalibratedHit, PMT, MCTrack}) = Base.angle(a, b.dir)
Base.angle(a::Union{CalibratedHit, PMT, MCTrack}, b) = Base.angle(a.dir, b)

"""
$(SIGNATURES)

Calculate the azimuth angle for a given direction.
"""
azimuth(d) = atan(d[2], d[1])

"""
$(SIGNATURES)

Calculate the zenith angle for a given direction.
"""
zenith(d) = acos(-d[3]/norm(d))

"""
$(SIGNATURES)

Calculate the cartesian coordinates for given `ϕ`, `θ` and radius `r`.
"""
function cartesian(ϕ, θ; r=1.0)
    return Position(r*sin(θ) * cos(ϕ), r*sin(θ) * sin(ϕ), cos(θ))
end

"""
$(SIGNATURES)

Calculate the distance between a point (p1) and a line (given by p2 and d2).
"""
function pld3(p1, p2, d2)
    norm(cross(d2, (p2 - p1))) / norm(d2)
end

"""
$(SIGNATURES)

Calculate the distance between two lines.
"""
function lld3(P, u, Q, v)
    R = Q - P
    n = cross(u, v)
    return norm(R⋅n) / norm(n)
end

"""
$(SIGNATURES)

Calculate the distance between two tracks.
"""
function lld3(t₁::Track, t₂::Track)
    return lld3(t₁.pos, t₁.dir, t₂.pos, t₂.dir)
end

"""
$(SIGNATURES)

Projects a point to a track.
"""
function project(P, t::Track)
    A = t.pos
    B = A + t.dir
    project(P, A, B)
end

"""
$(SIGNATURES)

Project P onto a line spanned by A and B.
"""
function project(P, A, B)
    Position(A + ((P-A)⋅(B-A))/((B-A)⋅(B-A)) * (B-A))
end

"""
$(SIGNATURES)

The function to fit time residual plot distributions.
"""
function langauss(x, LA, Lμ, Lσ, GA, Gμ, Gσ, offset)
    LA * pdf(Landau(Lμ, Lσ), x) + GA * pdf(Normal(Gμ, Gσ), x) + offset
end


text:
ENV["JULIA_CXX_RTTI"]=1
using Cxx
push!(LOAD_PATH,"../src/")
using Documenter, ROS

makedocs(
    modules = [ROS],
    authors = "George Stavrinos",
    sitename = "ROS.jl",
    format = Documenter.HTML(prettyurls = false, footer = nothing),
    pages = [
        "Home" => "index.md",
        "Features" => "features.md",
        "Examples" => "examples.md",
        "API" => [
            "roscpp.md",
            "nodehandle.md",
            "subscriber.md",
            "publisher.md",
            "serviceServer.md",
            "serviceClient.md",
            "time.md",
            "package.md",
            "actionlib.md",
            "tf.md",
            "typegen.md"
        ],
    ]
)

deploydocs(
    repo = "github.com/gstavrinos/ROS.jl.git"
)


text:
# from https://github.com/vitskvara/PaperUtils.jl

"""
    eol(s)
Replaces the "& " at the end of s with a tabular end of line.
"""
function eol(s)
    return string(s[1:end-2], " \\\\ \n")
end

"""
    wspad(s, n)
Pads s with n white spaces.
"""
function wspad(s, n)
    return string(s, repeat(" ", n))
end
function ranks2tikzcd(ranks, algnames, c, caption = ""; scale=1.0,
    vpad=0.3, label = "", pos = "h", levelwidth=0.06, vdist=0.1)
    n = length(ranks)
    @assert n == length(algnames)
    ranks = reshape(ranks, n)

    # header
    s = "\\begin{tikzpicture}[scale=$(scale)] \n"
    s = wspad(s, 2)
    #axis
    s = string(s, drawaxis(ranks))
    # nodes
    s = string(s, drawnodes(ranks, algnames; vpad=vpad))
    # levels
    s = string(s, drawlevels(ranks, c; levelwidth=levelwidth, vdist=vdist))
    # end of figure
    s = wspad(s, 1)
    s = string(s, "\\end{tikzpicture} \n")

    return s
end

function drawaxis(ranks)
    mx = ceil(maximum(ranks))
    mn = floor(minimum(ranks))

    s = "\\draw ($(mn),0) -- ($(mx),0); \n"
    s = wspad(s, 2)
    s = string(s, "\\foreach \\x in {$(Int(mn)),...,$(Int(mx))} \\draw (\\x,0.10) -- (\\x,-0.10) node[anchor=north]{\$\\x\$}; \n")
    return s
end

function drawnodes(ranks, algnames; vpad=0.3)
    isort = sortperm(ranks)
    ranks = ranks[isort]
    algnames = algnames[isort]

    mx = ceil(ranks[end])
    mn = floor(ranks[1])
    nor = length(ranks) # number of ranks
    nos = Int(ceil(nor/2)) # number of ranks on one side (start with left, there may be one more if nor is odd)

    s = ""
    for (i, r) in enumerate(ranks)
        s = wspad(s,2)
        if i<=nos
            s = string(s, "\\draw ($(r),0) -- ($r,$(i*vpad-0.1)) -- ($(mn-0.1), $(i*vpad-0.1)) node[anchor=east] {$(algnames[i])}; \n")
        else
            s = string(s, "\\draw ($(r),0) -- ($r,$((nor-i)*vpad+0.2)) -- ($(mx+0.1), $((nor-i)*vpad+0.2)) node[anchor=west] {$(algnames[i])}; \n")
        end
    end

    return s
end

function drawlevels(ranks, cv; levelwidth=0.06, vdist=0.1)
    ranks = sort(ranks)

    mx = ceil(ranks[end])
    mn = floor(ranks[1])
    nor = length(ranks) # number of ranks

    # do this differently
    ranks = sort(ranks)

    mx = ceil(ranks[end])
    mn = floor(ranks[1])
    nor = length(ranks) # number of ranks

    # do a matrix that says which ranks are connected
    con_mat = zeros(nor, nor)
    for i in 1:nor
        for j in 1:nor
            con_mat[i,j] = abs(ranks[j] - ranks[i]) < cv
        end
    end

    # now
    levels = [] # each element is a tuple containing the (start, end, height) values
    lastlasti = 0
    for i in 1:nor
        if length(levels) == 0
            global levelh = vdist
        elseif lastlasti < i
            global levelh = vdist
        else
            lastlevelh = levels[end][end]
            if lastlevelh == vdist
                global levelh = lastlevelh + vdist
            else
                global levelh = lastlevelh + vdist
            end
        end
        startl = ranks[i]
        lasti = findlast(con_mat[i,:] .== 1)
        if lasti != nothing && lasti != lastlasti && lasti != i
            endl = ranks[lasti]
            push!(levels, (startl, endl, levelh))
        end
        lastlasti = lasti
    end

    # draw them
    s = ""
    for level in levels
        s = wspad(s,2)
        s = string(s, "\\draw[line width=$(levelwidth)cm,color=black,draw opacity=1.0] ($(level[1]-0.03),$(level[3])) -- ($(level[2]+0.03),$(level[3])); \n")
    end

    return s
end

"""
    string2file(f, s)
Save string s to file f.
"""
function string2file(f, s)
    open(f, "w") do _f
        write(_f, s)
    end
end