text:
using UpdateJulia
using Test
using Suppressor

@testset "report path before" begin
    println(ENV["PATH"])
    if Sys.iswindows()
        for systemwide in (true,false)
            println(strip(open(io -> read(io, String), `powershell.exe -nologo -noprofile -command "[Environment]::GetEnvironmentVariable(\"PATH\"$(systemwide ? "" : ", \"User\""))"`)))
        end
    end
end

@testset "fetch()" begin
    @test UpdateJulia.last_fetched[] == 0
    t0 = time()
    @test nothing === UpdateJulia.fetch()
    t1 = time()
    @test 0 < t0 < UpdateJulia.last_fetched[] <= t1
end

@testset "docstring" begin
    display(@doc update_julia)
    println()
end

@testset "prefer()" begin
    @test UpdateJulia.prefer(v"1.7.0", v"1.3.6")
    @test !UpdateJulia.prefer(v"1.7.0", v"1.7.0")
    @test !UpdateJulia.prefer(v"1.7.0-rc1", v"1.6.4")
    @test !UpdateJulia.prefer(v"17.7.0-rc1", v"1.6.4")
    @test UpdateJulia.prefer(v"17.7.0-rc1", v"1.6.4-DEV")
    @test UpdateJulia.prefer(v"17.7.0-rc1", v"2.6.4-DEV")
    @test !UpdateJulia.prefer(v"17.7.0-DEV", v"2.6.4-rc1")

    @test !UpdateJulia.prefer(missing, v"1.6.4")
    @test !UpdateJulia.prefer(missing, missing)
    @test UpdateJulia.prefer(v"1.7.0-rc1", missing)
end

@testset "latest()" begin
    @test UpdateJulia.latest() >= v"1.7.0"
    @test UpdateJulia.latest().prerelease == ()
    @test UpdateJulia.latest("1.1.") == v"1.1.1"
    @test UpdateJulia.latest("1.5") == v"1.5.4"
    @test UpdateJulia.latest("1.2") == v"1.2.0"
    @test UpdateJulia.latest("1.4.0-") == v"1.4.0-rc2"

    l = UpdateJulia.latest()
    for v in keys(UpdateJulia.versions[])
        @test v === l || UpdateJulia.prefer(l, v)
    end
end

@testset "version_of()" begin
    @test ismissing(UpdateJulia.version_of("this-should-not-be-a-valid-command-gj493fD62l4"))
end

@testset "versions[]" begin
    @test VersionNumber[v"0.1.2", v"0.3.1", v"0.5.0-rc2", v"0.6.0-pre.beta", v"0.7.0-beta",
        v"1.0.0", v"1.1.1", v"1.2.0-rc2", v"1.2.0-rc3", v"1.2.0", v"1.3.0-rc3", v"1.4.1",
        v"1.4.2", v"1.5.1", v"1.5.4", v"1.6.0-rc3", v"1.6.1", v"1.6.4", v"1.7.0-beta2",
        v"1.7.0-beta3"] ⊊ keys(UpdateJulia.versions[])
end

@testset "dry tests" begin
    # Basics
    @test update_julia(dry_run = true) == UpdateJulia.latest()
    @suppress_out @test UpdateJulia.update_julia("nightly", dry_run=true, verbose=false) ==
        UpdateJulia.update_julia(string(UpdateJulia.nightly_version[]), dry_run=true, verbose=false) ==
        UpdateJulia.nightly_version[]

    # Speed
    f() = @test @suppress update_julia("1.7.0-rc1", dry_run=true, verbose=false) == v"1.7.0-rc1"
    @elapsed f()
    @test .05 > @elapsed f()

    # Force fetch
    t_before_fetch = UpdateJulia.last_fetched[]
    @test update_julia("1.7.0-rc3", fetch = true, dry_run=true) == v"1.7.0-rc3"
    t_after_fetch = UpdateJulia.last_fetched[]
    @test 0 < t_before_fetch < t_after_fetch <= time()

    # GUI
    @test update_julia("1.2", dry_run=true, verbose=false) == v"1.2.0"
    @test update_julia("1.6", prefer_gui=true, dry_run=true, verbose=false) >= v"1.6.6"
end

installed_versions = [UpdateJulia.latest(), UpdateJulia.nightly_version[],
    #UpdateJulia.latest("$(VERSION.major).$(VERSION.minor)"),
    v"1.6.3", UpdateJulia.latest("1.6"), v"1.6.2", v"1.5.0-rc2", v"1.5.0", v"1.5.4"]

function random_matrix_test(n)
    versions = vcat(UpdateJulia.nightly_version[], filter!(collect(keys(UpdateJulia.versions[]))) do x
        x >= (Sys.iswindows() ? v"1.5.0-rc2" : v"1.0.0")
    end)
    keywords = [
        # os_str => untested
        # arch => untested
        :prefer_gui => Sys.iswindows() ? [false] : [true, false],
        :fetch => Bool,
        # _v_url => untested
        # v => untested
        :migrate_packages => [true, false, :force],
        # url => untested
        # aliases => untested
        # :systemwide => Bool, unfortunately, we can't do this trivially because userspace installs choose not to overwrite systemwide installs
        :install_location => [mktempdir(), mktempdir()], # Tuple sampling is missing in julia 1.0
        # bin => untested
        :dry_run => Bool,
        :verbose => Bool]

    @testset "install" begin
        for _ in 1:n
            v = rand(versions)
            version = rand(["", "$(v.major)", "$(v.major).$(v.minor)", "$(v.major).$(v.minor).$(v.patch)", "$v"])
            kw = [k=>rand(source) for (k, source) in filter(x->rand(Bool), keywords)]
            try
                v_inst = update_julia(version; kw...)
                @test v_inst === v || UpdateJulia.prefer(v_inst, v)
                if (:dry_run => true) ∉ kw
                    push!(installed_versions, v_inst)
                end
            catch x
                if x isa ProcessFailedException && ((:migrate_packages => :force) ∈ kw || (:migrate_packages => true) ∈ kw)
                    v = UpdateJulia.latest(version)
                    if (v.major, v.minor) < (VERSION.major, VERSION.minor)
                        # Package migratoin requires Pkg at version to be compatible with
                        # VERSION's Project.toml. This is not required by SymVer, and not
                        # the case for 1.0's Pkg with 1.8's Project.toml. So, if we
                        # explicitly ask for package migration in a regression (even without
                        # force) it may fail, and this is okay.
                        @warn "failed backwards package migration to $v"
                    else
                        rethrow()
                    end
                else
                    rethrow()
                end
            end
        end
    end

    commands = unique(vcat("julia", [["julia-$(v.major).$(v.minor)", "julia-$v"]
        for v in vcat(UpdateJulia.nightly_version[], versions)]...))
    @test String["julia", "julia-1.8", "julia-1.5.1", "julia-1.6.2", "julia-1.5.0-rc2",
        "julia-1.7.0-rc1", "julia-1.7.0-beta2"] ⊊ commands
    Sys.iswindows() || @test String["julia-1.0", "julia-1.4",  "julia-1.3.1",
        "julia-1.3.0-rc3", "julia-1.3.0-alpha"] ⊊ commands

    @testset "check" begin
        for c in commands
            cvs = c[7:end]
            installed = filter(installed_versions) do x
                # TODO migrate this to UpdateJulia.match
                xs = string(x)
                xs == cvs && return true # Exact match counts
                # If inexact, and cvs is a complete version string, don't match
                try string(VersionNumber(cvs)) == cvs && return false catch end
                # Otherwise match if prefix
                startswith(xs, cvs)
            end
            actual = UpdateJulia.version_of(c)
            for i in installed
                if UpdateJulia.prefer(i, actual)
                    println("Running: $c")
                    println("Gives result: $actual")
                    println("Installed these compatible versions, in this order: $installed")
                    println("Specifically, $i, which is better than $actual")
                    @test false
                else
                    @test true
                end
            end
            if !ismissing(actual) && actual ∉ installed
                @warn "Unexpectedly found julia version $actual when running `$c -v`. Only installed $installed in testing."
            end
        end
    end
end

if ("CI" => "true") ∈ ENV
    printstyled("WARNING: this will both install out of date julia versions and change what the command `julia` points to.\n", color=Base.warn_color())
    printstyled("If all goes well, it will finish with the latest stable version of julia installed.\n", color=Base.warn_color())

    @testset "curated tests" begin
        # TODO add any failing tests from the random matrix here so that we have a readable,
        # reproducible, and efficient list of tests to start with.

        # default functionality
        @test update_julia() == UpdateJulia.latest()

        # nightly
        @test update_julia("nightly") == UpdateJulia.nightly_version[]

        # Begin PREVIOUSLY FAILED TESTS
        UpdateJulia.migrate_packages(VERSION, true) # Forcibly migrate packages to same directory
        # ERROR: ArgumentError: 'src' and 'dst' refer to the same file/dir. This is not supported.
        # End PREVIOUSLY FAILED TESTS

        # this version
        mm = "$(VERSION.major).$(VERSION.minor)"
        #@test update_julia(mm) == UpdateJulia.latest(mm) TODO if we need this, say why.

        # all these versions have to be at least 1.5.0-rc2 when windows archive became available
        # fallback for prefer_gui when not available
        # explicit systemwide
        # migrate packages
        Sys.iswindows() || println("requesting GUI...")
        @test update_julia("1.6.3", prefer_gui = !Sys.iswindows(), systemwide=false, migrate_packages=true) == v"1.6.3"
        # ensure there is something to overwrite
        @test UpdateJulia.version_of("julia-1.6") ∈ (VERSION == v"1.6.4" ? [v"1.6.3", v"1.6.4"] : [v"1.6.3"])
        # ensure that we actually have packages to migrate (this is not first to test fallback when we don't)
        project_toml = joinpath(first(Base.DEPOT_PATH), "environments", "v$mm", "Project.toml")
        if !isfile(project_toml)
            mkpath(dirname(project_toml))
            open(io -> write(io, "Statistics = \"10745b16-79ce-11e8-11f9-7d13ad32a3b2\"\n"), project_toml, "w")
        end
        # note that the systemwide installation happens after the user installation so that it can overwrite
        @test update_julia("1.6", systemwide=true, migrate_packages=true) == UpdateJulia.latest("1.6") > v"1.6.3"
        # ensure that migration actually happened
        @test isfile(joinpath(first(Base.DEPOT_PATH), "environments", "v1.6", "Project.toml"))
        # do overwrite
        @test UpdateJulia.version_of("julia-1.6") == UpdateJulia.latest("1.6")
        # failed migrate packages without force
        @test update_julia("1.6.2", migrate_packages=true) == v"1.6.2"
        # don't overwrite
        @test UpdateJulia.version_of("julia-1.6") == UpdateJulia.latest("1.6")

        # successfully migrate packages with force
        old_1_5_0 = UpdateJulia.version_of("julia-1.5.0")
        @test update_julia("1.5.0-", systemwide=true, migrate_packages=:force, verbose=true) == v"1.5.0-rc2"
        @test UpdateJulia.version_of("julia-1.5.0") === old_1_5_0
        # Manifest.toml is not neccessarily created until we migrate with force.
        @test isfile(joinpath(first(Base.DEPOT_PATH), "environments", "v1.5", "Manifest.toml"))
        @test update_julia("1.5.0", systemwide=false, migrate_packages=:force, verbose=true) == v"1.5.0"
        @test UpdateJulia.@os(
            # windows doesn't use a bin, and it's only the user path we are reordering.
            UpdateJulia.version_of("julia-1.5") == v"1.5.0",
            # can't overwite the old installation because it was systemwide
            UpdateJulia.version_of("julia-1.5") == v"1.5.0-rc2",
            # system bin may come after user bin, so we have to specify the bin
            UpdateJulia.version_of("/usr/local/bin/julia-1.5") == v"1.5.0-rc2"
        )
        # but still successfully sets julia-1.5.0
        @test UpdateJulia.version_of("julia-1.5.0") == v"1.5.0"
    end

    # Get proper 1.5 installed so we don't trigger false positives in the matrix
    # by masking julia-1.5
    @test update_julia("1.5", systemwide=true) == v"1.5.4"

    random_matrix_test(4) # Small
    random_matrix_test(15) # Big

    # Finish with latest version installed
    @test UpdateJulia.version_of("julia") == UpdateJulia.latest()
else
    @warn "Skipped installation tests!"
end

@testset "report path after" begin
    println(ENV["PATH"])
    if Sys.iswindows()
        for systemwide in (true,false)
            println(strip(open(io -> read(io, String), `powershell.exe -nologo -noprofile -command "[Environment]::GetEnvironmentVariable(\"PATH\"$(systemwide ? "" : ", \"User\""))"`)))
        end
    end
end


text:
import HTTP

# STATIC: Serve index.html, which is the same for every notebook - it's a ⚡🤑🌈 web app
# index.html also contains the CSS and JS

"Attempts to find the MIME pair corresponding to the extension of a filename. Defaults to `text/plain`."
function mime_fromfilename(filename)
    # This bad boy is from: https://developer.mozilla.org/en-US/docs/Web/HTTP/Basics_of_HTTP/MIME_types/Common_types
    mimepairs = Dict(".aac" => "audio/aac", ".bin" => "application/octet-stream", ".bmp" => "image/bmp", ".css" => "text/css", ".csv" => "text/csv", ".eot" => "application/vnd.ms-fontobject", ".gz" => "application/gzip", ".gif" => "image/gif", ".htm" => "text/html", ".html" => "text/html", ".ico" => "image/vnd.microsoft.icon", ".jpeg" => "image/jpeg", ".jpg" => "image/jpeg", ".js" => "text/javascript", ".json" => "application/json", ".jsonld" => "application/ld+json", ".mjs" => "text/javascript", ".mp3" => "audio/mpeg", ".mpeg" => "video/mpeg", ".oga" => "audio/ogg", ".ogv" => "video/ogg", ".ogx" => "application/ogg", ".opus" => "audio/opus", ".otf" => "font/otf", ".png" => "image/png", ".pdf" => "application/pdf", ".rtf" => "application/rtf", ".sh" => "application/x-sh", ".svg" => "image/svg+xml", ".tar" => "application/x-tar", ".tif" => "image/tiff", ".tiff" => "image/tiff", ".ttf" => "font/ttf", ".txt" => "text/plain", ".wav" => "audio/wav", ".weba" => "audio/webm", ".webm" => "video/webm", ".webp" => "image/webp", ".woff" => "font/woff", ".woff2" => "font/woff2", ".xhtml" => "application/xhtml+xml", ".xml" => "application/xml", ".xul" => "application/vnd.mozilla.xul+xml", ".zip" => "application/zip")
    file_extension = getkey(mimepairs, '.' * split(filename, '.')[end], ".txt")
    MIME(mimepairs[file_extension])
end

function assetresponse(path)
    try
        @assert isfile(path)
        response = HTTP.Response(200, read(path, String))
        push!(response.headers, "Content-Type" => string(mime_fromfilename(path)))
        return response
    catch e
        HTTP.Response(404, "Not found!: $(e)")
    end
end

function serveonefile(path)
    return request::HTTP.Request->assetresponse(normpath(path))
end

function serveasset(req::HTTP.Request)
    reqURI = HTTP.URI(req.target)

    filepath = joinpath(PKG_ROOT_DIR, relpath(reqURI.path, "/"))
    assetresponse(filepath)
end

const PLUTOROUTER = HTTP.Router()

function serve_editor(req::HTTP.Request)
    p=HTTP.URI(req.target).path
    b = String(req.body)
    HTTP.Response(200, "Path: $(p) \n\n Body: $(b)")
end

function notebook_redirect(notebook)
    response = HTTP.Response(302, "")
    push!(response.headers, "Location" => "/edit?uuid=" * string(notebook.uuid))
    return response
end

function serve_openfile(req::HTTP.Request)
    uri=HTTP.URI(req.target)
    query = HTTP.URIs.unescapeuri(replace(uri.query, '+' => ' '))

    if length(query) > 5
        path = query[6:end]

        if(isfile(path))
            try
                for nb in values(notebooks)
                    if realpath(nb.path) == realpath(path)
                        return notebook_redirect(nb)
                    end
                end

                nb = load_notebook(path)
                save_notebook(nb)
                notebooks[nb.uuid] = nb
                return notebook_redirect(nb)

            catch e
                return HTTP.Response(500, "Failed to load notebook:\n\n$(e)\n\nGo back")
            end
            # return HTTP.Response(200, """<meta charset="utf-8" />Path: $(path)</bodu>""")
        else
            return HTTP.Response(404, "Can't find a file here.\n\nGo back")
        end
    end
    return HTTP.Response(400, "Bad query.\n\nGo back")
end

function serve_samplefile(req::HTTP.Request)
    nb = samplenotebook()
    save_notebook(nb)
    notebooks[nb.uuid] = nb
    return notebook_redirect(nb)
end

function serve_newfile(req::HTTP.Request)
    nb = emptynotebook()
    save_notebook(nb)
    notebooks[nb.uuid] = nb
    return notebook_redirect(nb)
end

HTTP.@register(PLUTOROUTER, "GET", "/", serveonefile(joinpath(PKG_ROOT_DIR, "assets", "welcome.html")))
HTTP.@register(PLUTOROUTER, "GET", "/index", serveonefile(joinpath(PKG_ROOT_DIR, "assets", "welcome.html")))
HTTP.@register(PLUTOROUTER, "GET", "/index.html", serveonefile(joinpath(PKG_ROOT_DIR, "assets", "welcome.html")))

HTTP.@register(PLUTOROUTER, "GET", "/edit", serveonefile(joinpath(PKG_ROOT_DIR, "assets", "editor.html")))
HTTP.@register(PLUTOROUTER, "GET", "/sample", serve_samplefile)
HTTP.@register(PLUTOROUTER, "GET", "/new", serve_newfile)
HTTP.@register(PLUTOROUTER, "GET", "/open", serve_openfile)

HTTP.@register(PLUTOROUTER, "GET", "/favicon.ico", serveonefile(joinpath(PKG_ROOT_DIR, "assets", "favicon.ico")))
HTTP.@register(PLUTOROUTER, "GET", "/assets/*", serveasset)

HTTP.@register(PLUTOROUTER, "GET", "/ping", r->HTTP.Response(200, JSON.json("OK!")))

ext: jl


text:
@testset "Finite differences" begin
    f = reshape(collect(1.0:25),5,5)
    @testset "Gradients" begin
        outputx = zeros(5,5)
        outputy = zeros(5,5)
        @testset "Simple" begin
            Swalbe.∇f!(outputx, outputy, f)
            @test isa(outputx, Array)
            @test isa(outputy, Array)
            #=
            Analytical results, tested with DiffEqOperators
            Dx = CenteredDifference{1}(1,2,1.0,5)
            Dy = CenteredDifference{2}(1,2,1.0,5)
            fpad = padarray(f, Pad(:circular, 1, 1))
            mul!(solx, Dx, fpad)
            mul!(soly, Dy, fpad)
            =#
            solx = [-1.5 -1.5 -1.5 -1.5 -1.5;
                     1.0 1.0 1.0 1.0 1.0;
                     1.0 1.0 1.0 1.0 1.0;
                     1.0 1.0 1.0 1.0 1.0;
                    -1.5 -1.5 -1.5 -1.5 -1.5]

            soly = [-7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5]
            # Test them
            @test all(outputx .== solx)
            @test all(outputy .== soly)
        end
        @testset "Four Arguments" begin
            a = fill(0.1, 5,5)
            Swalbe.∇f!(outputx, outputy, f, a)
            @test isa(outputx, Array)
            @test isa(outputy, Array)
            # Analytical results
            solx = [-1.5 -1.5 -1.5 -1.5 -1.5;
                     1.0 1.0 1.0 1.0 1.0;
                     1.0 1.0 1.0 1.0 1.0;
                     1.0 1.0 1.0 1.0 1.0;
                    -1.5 -1.5 -1.5 -1.5 -1.5]

            soly = [-7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5]
            # Test them
            @test all(outputx .== 0.1 .* solx)
            @test all(outputy .== 0.1 .* soly)
        end
        @testset "Five Arguments" begin
            a = fill(0.1, 5,5)
            b = zeros(5,5,8)
            Swalbe.∇f!(outputx, outputy, f, b, a)
            @test isa(outputx, Array)
            @test isa(outputy, Array)
            # Analytical results
            solx = [-1.5 -1.5 -1.5 -1.5 -1.5;
                     1.0 1.0 1.0 1.0 1.0;
                     1.0 1.0 1.0 1.0 1.0;
                     1.0 1.0 1.0 1.0 1.0;
                    -1.5 -1.5 -1.5 -1.5 -1.5]

            soly = [-7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5;
                    -7.5 5.0 5.0 5.0 -7.5]
            # Test them
            @test all(outputx .== 0.1 .* solx)
            @test all(outputy .== 0.1 .* soly)
        end
    end

    f1 = collect(1.0:25)
    @testset "Gradients 1D" begin
        output = zeros(25)
        dummy = zeros(25,2)
        @testset "Simple" begin
            Swalbe.∇f!(output, f1, dummy, ones(25))
            @test isa(output, Vector)

            sol = ones(25)
            sol[1] = -11.5
            sol[end] = -11.5
            # Test them
            @test all(output .== sol)

            Swalbe.∇f!(output, f1, dummy)
            @test isa(output, Vector)

            sol = ones(25)
            sol[1] = -11.5
            sol[end] = -11.5
            # Test them
            @test all(output .== sol)
        end

    end

    @testset "Laplacian" begin
        output = zeros(5,5)
        γ = -1.0
        Swalbe.∇²f!(output, f, γ)
        @test isa(output, Array)
        #=
        Analytical results, tested with DiffEqOperators
        Dxx = CenteredDifference{1}(2,2,1.0,5)
        Dyy = CenteredDifference{2}(2,2,1.0,5)
        fpad = padarray(f, Pad(:circular, 1, 1))
        A = Dxx + Dyy
        mul!(sol, A, fpad)
        =#
        sol = [-30.0 -5.0 -5.0 -5.0 20;
               -25.0 0.0 0.0 0.0 25.0;
               -25.0 0.0 0.0 0.0 25.0;
               -25.0 0.0 0.0 0.0 25.0;
               -20.0 5.0 5.0 5.0 30.0];

        @test all(isapprox.(output, sol; atol=1e-10))
    end

    @testset "Laplacian 1D" begin
    output = zeros(25)
    dummy = zeros(25,2)
    input = collect(1:25.0)
    Swalbe.∇²f!(output, input, dummy)
    @test isa(output, Vector)

    sol = zeros(25)
    sol[1] = 25;
    sol[end] = -25;

    @test all(output .== sol)

end

    @testset "Viewneighbors" begin
        f = reshape(collect(1.0:5*5*8),5,5,8)
        f1, f2, f3, f4, f5, f6, f7, f8 = Swalbe.viewneighbors(f)
        allviews = [f1, f2, f3, f4, f5, f6, f7, f8]
        for (index, value) in enumerate(allviews)
            @test all(value .== f[:,:,index])
        end
    end

    @testset "Viewneighbors 1D" begin
        f = reshape(collect(1.0:20),10,2)
        f1, f2 = Swalbe.viewneighbors_1D(f)
        allviews = [f1, f2]
        for (index, value) in enumerate(allviews)
            @test all(value .== f[:,index])
        end
    end
end


text:
# sample sequences of ergms with different parameters' values from R package ergm
# and test the PseudoLikelihoodScoreDrivenERGM filter
using Utilities,AReg,StaticNets,JLD,MLBase,StatsBase,CSV, RCall
using PyCall; pygui(); using PyPlot
# load the required packages in R
R"library(statnet)
    library(ergm)
    library(sna)
    library(coda)
    library(network)
    sessionInfo()"

R"set.seed(0)"

Nsample = 100
T = 50
N = 50
Nterms = 2
# the matrix of parameters values for each t

parDgpT = zeros(Nterms,T)
 parDgpT[1,:] = randSteps(-2.5,-3,2,T)# -3# randSteps(0.05,0.5,2,T) #1.5#.00000000000000001
 parDgpT[2,:] = randSteps(0.05,0.5,2,T)#
 #parDgpT[2,:] = 0.25
 @rput T; @rput parDgpT;@rput N
 #create an empty network, the formula defining ergm, sample the ensemble and store in R
 R"
 net <- network.initialize(N)
 formula_ergm = net ~ edges + gwesp(decay = 0.25,fixed = TRUE)#gwnsp(decay = 0.25,fixed = TRUE)
   sampledMat_T_R =    array(0, dim=c(N,N,T))
   estParSS_T_R =    list()
   changeStats_T_R = list()
   stats_T_R = list()
    for(t in 1:T){
        print(t)
        print(parDgpT[,t])
         net <- simulate(formula_ergm, nsim = 1, seed = sample(1:100000000,1), coef = parDgpT[,t],control = control.simulate.formula(MCMC.burnin = 500000))
         sampledMat_T_R[,,t] <- as.matrix.network( net)
         tmp <- ergm(formula_ergm,estimate = 'MPLE')#)#
         estParSS_T_R[[t]] <- tmp[[1]]
         print(estParSS_T_R[[t]])
         chStat_t <- ergmMPLE(formula_ergm)
         changeStats_T_R[[t]] <- cbind(chStat_t$response, chStat_t$predictor,chStat_t$weights)
         stats_T_R[[t]] <- summary(formula_ergm)
         }"
 # import in julia
 sampledMat_T = BitArray( @rget(sampledMat_T_R))
 estParSS_T = @rget(estParSS_T_R);tmp = zeros(Nterms,T); for t=1:T tmp[:,t] = estParSS_T[t]; end ; estParSS_T = tmp
 changeStats_T = @rget changeStats_T_R; tmp = Array{Array{Float64,2}}(T); for t=1:T tmp[t] =  changeStats_T[t];end;changeStats_T = tmp
 stats_T = @rget(stats_T_R); tmp = zeros(Nterms,T); for t=1:T tmp[:,t] = stats_T[t];end;stats_T = tmp
 close()
figure()
 subplot(2,2,1);  plot(1:T,sumSq(sumSq(sampledMat_T,1),1)./(N^2-N))
 subplot(2,2,2);  plot(1:T,stats_T[1,:]  )
 subplot(2,2,3); plot(1:T,parDgpT[1,:],"k",1:T,estParSS_T[1,:],"c")
 subplot(2,2,4); plot(1:T,parDgpT[2,:],"k",1:T,estParSS_T[2,:],"c")
 #cor(stats_T')

 estParStatic,convFlag,UM , ftot_0= estimate(SdErgmDirBinGlobalPseudo(changeStats_T,testGasPar,trues(2),"");UM = [0;0],indTvPar = falses(Nterms))
 constParVec = zeros(sum(!indTvPar)); for i=1:sum(.!indTvPar) constParVec[i] = estParStatic[.!indTvPar][i][1]; end
 grad_T,s_T,I_T = gasScoreSeries(SdErgmDirBinGlobalPseudo(changeStats_T,testGasPar,trues(2),""),constParVec)

 # run one auxiliary regression for each parameters
 using GLM
 for i=1:2
     g_T_t = Float64.(grad_T[i,2:end])
     s_T_tm1 = Float64.(s_T[i,1:end-1])
     g_T_tm1 = Float64.(grad_T[i,1:end-1])
     g_delta_t = Float64.(grad_T[(1:Nterms).!=i,2:end])
     I_T_t = Float64.(I_T[i,2:end])
     #ols = lm(@formula(x1 ~ x2 + x3 + x4 ),data)
     autocor(g_T_t,[1])
     y = ones(T-1)
     ols = lm([  g_delta_t' g_T_t  s_T_tm1.*g_T_t],y)
     yHat = predict(ols)
     yMean =1
     RSS = sum((yHat- yMean).^2)
     ESS = T - RSS
     pval =  ccdf(Chisq(1), ESS)
     scoreAcorr = autocor(Float64.(grad_T[i,:]),[1])[1]
     @show(i,pval,scoreAcorr)
 end

 figure()
 plot(s_T' -mean(s_T)); grid()

## test SDERGM
#  um = 0.5
#   B,A = 0.9, 0.02
#   testGasPar = [[-1.45],[um*(1-B),B ,A]]
#   indTvPar = trues(2);indTvPar[1] = false
#   testModel = DynNets.SdErgmDirBinGlobalPseudo(changeStats_T,testGasPar,trues(1),"")
# #
#   gasFiltPar , pseudolike = DynNets.score_driven_filter_or_dgp(testModel,testGasPar[2],indTvPar;
#     obsT = changeStats_T,vConstPar = testGasPar[1])
#     close()
#     subplot(2,2,1);  plot(1:T,sumSq(sumSq(sampledMat_T,1),1)./(N^2-N))
#     subplot(2,2,2);  plot(1:T,stats_T[1,:]  )
#     subplot(2,2,3); plot(1:T,gasFiltPar[1,:],"k",1:T,estParSS_T[1,:],"c")
#     subplot(2,2,4);  plot(1:T,parDgpT[2,:],"k",1:T,estParSS_T[2,:],"c")
#     subplot(2,2,4); plot(1:T,gasFiltPar[2,:],"k",1:T,estParSS_T[2,:],"c")
#  pseudolike
# #
#
# Npoints = 100
#  testPseudoLike = zeros(Npoints)
#  indTvPar = falses(2);
#  parVals=linspace(-3,3,Npoints)
#  for n=1:Npoints
#      testGasPar = [[parVals[n]],[0.5]]
#     ~,testPseudoLike[n]  = score_driven_filter_or_dgp(testModel,zeros(1),indTvPar;
#        obsT = changeStats_T,vConstPar = [testGasPar[1]; testGasPar[2]])
#  end
#  plot(parVals,testPseudoLike)
um = 0.5
 B,A = 0.9, 0.02
 testGasPar = [[-1.45],[um*(1-B),B ,A]]
 model = SdErgmDirBinGlobalPseudo(changeStats_T,testGasPar,trues(2),"")
 indTvPar = BitArray([true,true])
 estPar,convFlag,UM , ftot_0 = estimate(SdErgmDirBinGlobalPseudo(changeStats_T,testGasPar,trues(2),"");UM = [0;0],indTvPar = indTvPar,
                                                                        indTargPar = indTvPar)#  falses(indTvPar))##
 gasParVec = zeros(sum(indTvPar)*3); for i=0:(sum(indTvPar)-1) gasParVec[1+3i : 3(i+1)] = estPar[indTvPar][i+1]; end
 constParVec = zeros(sum(!indTvPar)); for i=1:sum(.!indTvPar) constParVec[i] = estPar[.!indTvPar][i][1]; end

#gasParVec[1] = -0.101
 #gasParVec[3] = 0.01
gasFiltPar , pseudolike = score_driven_filter_or_dgp(SdErgmDirBinGlobalPseudo(changeStats_T,testGasPar,trues(2),""),gasParVec,indTvPar;
                          vConstPar = constParVec,ftot_0= ftot_0)

 close()
    subplot(2,2,1);  plot(1:T,sumSq(sumSq(sampledMat_T,1),1)./(N^2-N))
    subplot(2,2,2);  plot(1:T,stats_T[1,:]  )
    subplot(2,2,3); plot(1:T,parDgpT[1,:],"k",1:T,gasFiltPar[1,:],"--k",1:T,estParSS_T[1,:],"c")
    subplot(2,2,4); plot(1:T,parDgpT[2,:],"k",1:T,gasFiltPar[2,:],"--k",1:T,estParSS_T[2,:],"c")

###
# close()
#  plt[:hist](estParSS_T[2,1:25])
#
#


text:
## utils.jl - utilities and helpers

# TODO: make recursive like calls_to_bcast
function bcast_to_call(pex::Expr)
    @assert pex.head == :(.)
    return Expr(:call, pex.args[1], pex.args[2].args...)
end

function calls_to_bcast(ex)
    return macroexpand(Main, :(@. $ex)) |> subs_bcast_with_dot
end

deriv_name(z::Symbol, x::Symbol) = Symbol("d$(z)!d$(x)")
split_deriv_name(vname) = Symbol.(split(String(vname), "!"))

function find_related(g::AbstractExGraph, dydx_v::Symbol)
    subderivs = Symbol[]
    i = 1
    name = Symbol("$(dydx_v)__$(i)")
    while haskey(g, name)
        push!(subderivs, name)
        i += 1
        name = Symbol("$(dydx_v)__$(i)")
    end
    return subderivs
end

# # (symbolic) derivative size propagation

const DERIV_NAME_PATTERN = r"(d.+)!(d.+)"

# function propagate_deriv_size!(g::AbstractExGraph, dd_name::Symbol)
#     sizes = @get_or_create(g.ctx, :sizes, Dict())
#     rg = match(DERIV_NAME_PATTERN, String(dd_name))
#     @assert length(rg.captures) == 2
#     str_dnames = rg.captures
#     zname = Symbol(str_dnames[1][2:end])
#     xname = Symbol(split(str_dnames[2][2:end], "__")[1]) # cut down `__$(i)` part if any
#     zsize, xsize = (sizes[zname], sizes[xname])
#     if zsize == :(())
#         # output var is constant
#         sizes[dd_name] = xsize
#     else
#         sizes[dd_name] = :(($zsize..., $xsize...)) |> simplify
#     end
# end

# function propagate_deriv_size!(g::AbstractExGraph)
#     for nd in g.tape
#         vname = varname(nd)
#         if match(DERIV_NAME_PATTERN, String(vname)) != nothing
#             propagate_deriv_size!(g, vname)
#         end
#     end
# end

# # (numeric) derivative size propagation

function infer_deriv_size!(g::AbstractExGraph, dd_name::Symbol)
    rg = match(DERIV_NAME_PATTERN, String(dd_name))
    @assert length(rg.captures) == 2
    str_dnames = rg.captures
    zname = Symbol(str_dnames[1][2:end])
    xname = Symbol(split(str_dnames[2][2:end], "__")[1]) # cut down `__$(i)` part if any
    # in case z or x haven't been evaluated and their size isn't known yet
    evaluate!(g, zname)
    evaluate!(g, xname)
    sizes = @get_or_create(g.ctx, :rsizes, Dict())
    if haskey(sizes, xname)
        # some nodes (e.g. structs) may not have a size
        zsize, xsize = (sizes[zname], sizes[xname])
        sizes[dd_name] = (zsize..., xsize...)
    end
end

function infer_deriv_size!(g::AbstractExGraph)
    for nd in g.tape
        vname = varname(nd)
        if match(DERIV_NAME_PATTERN, String(vname)) != nothing
            infer_deriv_size!(g, vname)
        end
    end
end

# top type

"The top type describing given data"
top_type(x::AbstractArray{T,N}) where {T,N} = AbstractArray{T,N}
top_type(x::Number) = Number
top_type(x::T) where T = T

top_type(::Type{AT}) where {AT <: AbstractArray{T,N}} where {T,N} = AbstractArray{T,N}
top_type(::Type{T}) where {T <: Number} = Number
top_type(::Type{T}) where T = T


text:
module Cards

export NastranCard, GenericCard
export GridCard, ElementCard, PropertyCard, CoordinateSystemCard, ConstraintCard, MaterialCard, SetCard

import Base.convert

abstract type NastranCard{CardType} end
struct CardType{T} end

const GenericCard = Vector{Any}
@generated function NastranCard{U}(u::Type{Val{U}},card)
    return :(convert($U,card))
end

function convert(::Type{NastranCard},card::GenericCard)
    NastranCard(Val{symbol(card[1])},card)
end

export CELAS1,CELAS2,CROD,CBEAM,CBUSH,RBE2,RBE3,CTRIAR,CQUADR,CONM2,CAERO1,SPLINE1,SPLINE2,CHEXA
include("element_cards.jl")

export PBEAM,PSHELL,PROD,PELAS,PAERO1,PSOLID
include("property_cards.jl")

export CORD1R,CORD2R
include("csys_cards.jl")

export GRID,SPOINT
include("grid_cards.jl")

export MPC
include("constraint_cards.jl")

export MAT1
include("material_cards.jl")

export AEFACT
include("set_cards.jl")

end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule openbabel_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("openbabel")
JLLWrappers.@generate_main_file("openbabel", UUID("deeba441-a6a4-5e65-99ae-21494d2aa7fd"))
end  # module openbabel_jll


text:
using ReactionCoordinates
using Documenter

makedocs(;
    modules=[ReactionCoordinates],
    authors="Pablo Zubieta <8410335+pabloferz@users.noreply.github.com> and contributors",
    repo="https://github.com/pabloferz/ReactionCoordinates.jl/blob/{commit}{path}#L{line}",
    sitename="ReactionCoordinates.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)


text:
using LinearAlgebra: getindex
using Distributions: eltype
push!(LOAD_PATH, abspath(@__DIR__,".."))
using Pkg
pkg"activate ."

using Test
using LinearAlgebra,Statistics,StatsBase,Distributions
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs
using BenchmarkTools
using FFTW

function onesparsemat(w::Real)
  mat=Matrix{Float64}(undef,1,1) ; mat[1,1]=w
  return sparse(mat)
end

##

n = 6
myβ = 1.1
mywdistr=Uniform(0.2,0.3)
mywmat = sprand(n,n,0.5,n->rand(mywdistr,n))
mywmat_nonsp = convert(Matrix{Float64},mywmat)
myin = rand(Uniform(0.1,0.5),n)
tfake = NaN
p1 = S.PopulationHawkesExp(n,myβ)
ps1 = S.PSHawkes(p1)
conn1 = S.ConnectionHawkes(ps1,mywmat,ps1)
# rates to test
p1_in = S.PopInputStatic(ps1,myin)
myntw = S.RecurrentNetwork(tfake,(ps1,),(p1_in,),(conn1,) )

#
nspikes = 500_000
# initialize
S.hawkes_initialize!(myntw)

my_act = Vector{Tuple{Int64,Float64}}(undef,nspikes)
for k in 1:nspikes
  S.dynamics_step!(tfake,myntw)
  idx_fire = findfirst(ps1.isfiring)
  t_now = ps1.time_now[1]
  my_act[k] = (idx_fire,t_now)
end

myspikes_all = [S.hawkes_get_spiketimes(i,my_act) for i in 1:n]
myTmax = maximum(x->x[end],myspikes_all)+eps()
@info "Total duration $(round(myTmax;digits=1)) s"
rates_all =  map(x->length(x)/myTmax, myspikes_all )
@info "Rates are $(round.(rates_all;digits=2)) "

# rates with Fourier
# analytic rate from fourier Eq between 6 and  7 in Hawkes
ratefou = let G0 =  mywmat_nonsp .* S.fou_interaction_kernel.(0,p1)
  inv(I-G0)*myin |> real
end

@info "With Fourier, rates are $(round.(ratefou;digits=2)) Hz"

##
# covariance numerical

mydt = 0.2
myτmax = 80.0
mytaus = S.get_times(mydt,myτmax)
ntaus = length(mytaus)
cov_num = S.covariance_density_numerical(myspikes_all,mydt,myτmax)

##

# numerical cov densities
plot(mytaus[2:end],cov_num[1,1,2:end] ; linewidth = 2)
plot!(mytaus[2:end],cov_num[2,2,2:end] ; linewidth = 2)
plot!(mytaus[2:end],cov_num[2,5,2:end] ; linewidth = 2)

## now with Fourier!

# Pick eq 12 from Hawkes
function four_high_res(dt::Real,Tmax::Real)
  k1 = 2
  k2 = 0.2
  n = size(mywmat_nonsp,1)
  myτmax,mydt = Tmax * k1, dt*k2
  mytaus = S.get_times(mydt,myτmax)
  nkeep = div(length(mytaus),k1)
  myfreq = S.get_frequencies_centerzero(mydt,myτmax)
  G_omega = map(mywmat_nonsp) do w
    ifftshift( w .* S.fou_interaction_kernel.(myfreq,p1))
  end
  D = Diagonal(ratefou)
  M = Array{ComplexF64}(undef,n,n,length(myfreq))
  Mt = similar(M,Float64)
  for i in eachindex(myfreq)
    G = getindex.(G_omega,i)
    # M[:,:,i] = (I-G)\(G*D+D*G'-G*D*G')/(I-G')
    M[:,:,i] = (I-G)\D*(G+G'-G*G')/(I-G')
  end
  for i in 1:n,j in 1:n
    Mt[i,j,:] = real.(ifft(M[i,j,:])) ./ mydt
  end
  return mytaus[1:nkeep],Mt[:,:,1:nkeep]
end

taush,Cfou=four_high_res(mydt,myτmax)

# numerical cov densities
plot(mytaus[2:end],cov_num[1,1,2:end] ; linewidth = 2)
plot!(taush[1:end],Cfou[1,1,1:end] ; linewidth=3, linestyle=:dash)

plot(mytaus[2:end],cov_num[2,2,2:end] ; linewidth = 2)
plot!(taush[1:end],Cfou[2,2,1:end] ; linewidth=3, linestyle=:dash)

plot(mytaus[2:end],cov_num[2,5,2:end] ; linewidth = 2)
plot!(taush[3:end],Cfou[2,5,3:end] ; linewidth=3, linestyle=:dash)

plot(mytaus[2:end],cov_num[1,6,2:end] ; linewidth = 2)
plot!(taush[3:end],Cfou[1,6,3:end] ; linewidth=3, linestyle=:dash)


text:
using RoME, IncrementalInference, KernelDensityEstimate, Distributions
using CloudGraphs

export registerGeneralVariableTypes!

# register types of interest (Pose2, etc) in CloudGraphs
# you can register new types at any time (Julia is dynamic)
function registerGeneralVariableTypes!(cloudGraph::CloudGraph)
  # Variable node
  CloudGraphs.registerPackedType!(cloudGraph, VariableNodeData, PackedVariableNodeData, encodingConverter=VNDencoder, decodingConverter=VNDdecoder);
  # factor nodes
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Obsv2}}, PackedFunctionNodeData{PackedObsv2}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Odo}}, PackedFunctionNodeData{PackedOdo}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{GenericMarginal}}, PackedFunctionNodeData{PackedGenericMarginal}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Ranged}}, PackedFunctionNodeData{PackedRanged}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  # Pose2
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{PriorPose2}}, PackedFunctionNodeData{PackedPriorPose2}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Pose2Pose2}}, PackedFunctionNodeData{PackedPose2Pose2}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Pose2DPoint2DBearingRange{Distributions.Normal,Distributions.Normal}}}, PackedFunctionNodeData{PackedPose2DPoint2DBearingRange}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Pose2DPoint2DRange}}, FunctionNodeData{Pose2DPoint2DRange}, encodingConverter=passTypeThrough, decodingConverter=passTypeThrough)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{PriorPoint2D}}, PackedFunctionNodeData{PackedPriorPoint2D}, encodingConverter=FNDencode, decodingConverter=FNDdecode)

  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{PriorPoint2DensityNH}}, PackedFunctionNodeData{PackedPriorPoint2DensityNH}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  #acoustic types
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Pose2DPoint2DRangeDensity}}, PackedFunctionNodeData{PackedPose2DPoint2DRangeDensity}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Pose2DPoint2DBearingRangeDensity}}, PackedFunctionNodeData{PackedPose2DPoint2DBearingRangeDensity}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  # Pose3 stuff
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{PriorPose3}}, PackedFunctionNodeData{PackedPriorPose3}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Pose3Pose3}}, PackedFunctionNodeData{PackedPose3Pose3}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{Pose3Pose3NH}}, PackedFunctionNodeData{PackedPose3Pose3NH}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  # partial constraints
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{PartialPriorRollPitchZ}}, PackedFunctionNodeData{PackedPartialPriorRollPitchZ}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{PartialPose3XYYaw}}, PackedFunctionNodeData{PackedPartialPose3XYYaw}, encodingConverter=FNDencode, decodingConverter=FNDdecode)
  CloudGraphs.registerPackedType!(cloudGraph, FunctionNodeData{GenericWrapParam{PartialPose3XYYawNH}}, PackedFunctionNodeData{PackedPartialPose3XYYawNH}, encodingConverter=FNDencode, decodingConverter=FNDdecode)

  nothing
end


text:
# TODO: DEPRECATE THIS!
function fit_fs_pt!(init::StateFS, cfs::ConstantsFS, dfs::DataFS, tfs::TunersFS;
                    use_rand_inits=false,
                    tempers::Vector{Float64}, ncores::Int,
                    nmcmc::Int, nburn::Int,
                    swap_freq::Int=10,
                    remove_current_workers::Bool=true,
                    monitors=[monitor1, monitor2],
                    fix::Vector{Symbol}=Symbol[],
                    thins::Vector{Int}=[2, nsamps_to_thin(10, nmcmc)],
                    thin_dden::Int=1,
                    printFreq::Int=0,
                    randpair=0.0,
                    computeDIC::Bool=false, computeLPML::Bool=false,
                    computedden::Bool=false,
                    sb_ibp::Bool=false,
                    use_repulsive::Bool=true, Z_marg_lamgam::Bool=true,
                    verbose::Int=1,
                    time_updates=[false for _ in tempers],
                    Z_thin::Int=1, seed::Int=-1)

  # Number of temperatures
  num_tempers = length(tempers)
  @assert mod(num_tempers, 2) == 0

  # Create distributed arrays for:
  # - loglike, ConstantsFS, StatesFS, TunersFS
  lls = [Float64[0.0] for _ in tempers]
  cs = [let
          ct = deepcopy(cfs)
          ct.constants.temper = tau
          ct
        end for tau in tempers]
  states = if use_rand_inits
    println("Using random inits!")
    [let
       s = genInitialState(cfs.constants, dfs.data)
       StateFS{Float64}(s, dfs)
     end for _ in tempers]
  else
    [deepcopy(init) for _ in tempers]
  end
  tuners = [deepcopy(tfs) for _ in tempers]
  swapcounts = zeros(Int, num_tempers, num_tempers)
  paircounts = zeros(Int, num_tempers, num_tempers)

  println("About to start parallel chains ...")
  for iter in 1:(nburn+nmcmc)
    # Perform updates in parallel
    # @time out = pmap(states, cs, tuners, lls, time_updates) do s, c, t, ll, tu
    #   @time for _ in 1:swap_freq
    out = pmap(states, cs, tuners, lls, time_updates) do s, c, t, ll, tu
      for _ in 1:swap_freq
        update_state_feature_select!(s, c, dfs, t,
                                     ll=ll, fix=fix,
                                     use_repulsive=use_repulsive,
                                     # TODO: make this more explicit, instead
                                     # of random. See `update_Z_v2`.
                                     Z_marg_lamgam=Z_marg_lamgam,
                                     sb_ibp=sb_ibp, time_updates=tu)
        # println(ll)
      end
      return Dict(:s => s, :c => c, :t => t, :ll => ll)
    end
    println("Done with iter $iter")

    # replace stuff
    states = [o[:s] for o in out]
    cs = [o[:c] for o in out]
    tuners = [o[:t] for o in out]
    lls = [o[:ll] for o in out]

    if iter > nburn  # && iter % swap_freq == 0
      llf(s, t) = compute_marg_loglike(s, cfs.constants, dfs.data, t)
      swapchains!(states, llf, tempers,
                  paircounts=paircounts, swapcounts=swapcounts,
                  randpair=randpair,
                  verbose=verbose)
    end
  end
  println("After running parallel chains ...")

  out = Dict(:lls => lls,
             :state => states[1],
             :c => cfs,
             :d => dfs,
             :t => tfs,
             :paircounts => paircounts,
             :swapcounts => swapcounts)
  return out
end

#= Possible implementation detail
@assert mod(nmcmc + nburn, swap_freq) == 0

for iter in 1:div(nmcmc + nburn, swap_freq)
  # Update states in parallel
  for _ in 1:swap_freq
    states = update(states)
  end

  # SWAP STATES
end
=#


text:
VERSION < v"0.7.0-beta2.199" && __precompile__()

"""
The LaTeXStrings module exists mainly to make LaTeX equations easier to type as
literal strings, and so that the resulting strings display as formatted equations
in supporting environments like IJulia.

See in particular the `LaTeXString` type and the `L"..."` constructor macro.
"""
module LaTeXStrings
export LaTeXString, latexstring, @L_str, @L_mstr
using Compat

# IJulia supports LaTeX output for any object with a text/latex
# writemime method, but these are annoying to type as string literals
# in Julia because of all the escaping required, e.g. "\$\\alpha +
# \\beta\$".  To simplify this, we add a new string type with a macro
# constructor, so that one can simply do L"$\alpha + \beta$".

@doc raw"""
A `LaTeXString` is a string type whose contents represent a fragment of LaTeX code,
typically containing an equation (`$...$`).   In certain environments (e.g. IJulia)
this will display with LaTeX-like formatting.   For the most part, you can use
a `LaTeXString` object in any context that expects an `AbstractString` object.

The `L"..."` macro is convenient for constructing `LaTeXString` objects, because
it eliminates the need to escape backslashes and dollar signs, and implicitly inserts
dollar signs around the string if none are present.  For example, `L"$\alpha$"`, `L"\alpha"`,
and `LaTeXString("\$\\alpha\$")` are all equivalent.
"""
struct LaTeXString <: AbstractString
    s::String
end

"""
    latexstring(args...)

Similar to `string(args...)`, but generates a `LaTeXString` instead of a `String`.
"""
latexstring(args...) = latexstring(string(args...))
function latexstring(s::String)
    # the only point of using LaTeXString to represent equations, since
    # IJulia doesn't support LaTeX output other than equations, so add $'s
    # around the string if they aren't there (ignoring \$)
    return occursin(r"[^\\]\$|^\$", s) ? LaTeXString(s) :  LaTeXString(string('\$', s, '\$'))
end
latexstring(s::AbstractString) = latexstring(String(s))

macro L_str(s, flags...) latexstring(s) end
macro L_mstr(s, flags...) latexstring(s) end

Base.write(io::IO, s::LaTeXString) = write(io, s.s)
Base.show(io::IO, ::MIME"application/x-latex", s::LaTeXString) = write(io, s)
Base.show(io::IO, ::MIME"text/latex", s::LaTeXString) = write(io, s)
function Base.show(io::IO, s::LaTeXString)
    print(io, "L")
    Base.print_quoted_literal(io, s.s)
end

Compat.firstindex(s::LaTeXString) = Compat.firstindex(s.s)
Compat.lastindex(s::LaTeXString) = Compat.lastindex(s.s)
@static if isdefined(Base, :iterate)
    Base.iterate(s::LaTeXString, i::Int) = iterate(s.s, i)
else
    Base.start(s::LaTeXString) = start(s.s)
    Base.next(s::LaTeXString, i) = next(s.s, i)
    Base.done(s::LaTeXString, i) = done(s.s, i)
end
Base.nextind(s::LaTeXString, i::Int) = nextind(s.s, i)
Base.prevind(s::LaTeXString, i::Int) = prevind(s.s, i)
Base.eachindex(s::LaTeXString) = eachindex(s.s)
Base.length(s::LaTeXString) = length(s.s)
Base.getindex(s::LaTeXString, i::Integer) = getindex(s.s, i)
Base.getindex(s::LaTeXString, i::Int) = getindex(s.s, i) # for method ambig in Julia 0.6
Base.getindex(s::LaTeXString, i::UnitRange{Int}) = getindex(s.s, i)
Base.getindex(s::LaTeXString, i::UnitRange{<:Integer}) = getindex(s.s, i)
Base.getindex(s::LaTeXString, i::AbstractVector{<:Integer}) = getindex(s.s, i)
Compat.codeunit(s::LaTeXString, i::Integer) = codeunit(s.s, i)
Compat.codeunit(s::LaTeXString) = codeunit(s.s)
Compat.ncodeunits(s::LaTeXString) = ncodeunits(s.s)
Compat.codeunits(s::LaTeXString) = codeunits(s.s)
Base.sizeof(s::LaTeXString) = sizeof(s.s)
Base.isvalid(s::LaTeXString, i::Integer) = isvalid(s.s, i)
Base.pointer(s::LaTeXString) = pointer(s.s)
Base.IOBuffer(s::LaTeXString) = IOBuffer(s.s)
Base.unsafe_convert(T::Union{Type{Ptr{UInt8}},Type{Ptr{Int8}},Cstring}, s::LaTeXString) = Base.unsafe_convert(T, s.s)

end # module


text:
export SAGA, SVRG, LSVRG, SVAG, SAG, QSAGA, ILSVRG, ASVAG, SG

export primiterates, dualiterates

################################################################################
# Index Samplings
################################################################################
abstract type IndexSampling end

struct UniformSampling{RNG} <: IndexSampling
	rng::RNG
	n::Int
end
(s::UniformSampling)() = (rand(s.rng, 1:s.n), 1/s.n)

struct WeightedSampling{RNG, CP} <: IndexSampling
	rng::RNG
	p::CP
	cp::CP
	function WeightedSampling(rng, w)
		p = w./sum(w)
		cp = cumsum(p)
		new{typeof(rng), typeof(cp)}(rng, p, cp)
	end
end
function (s::WeightedSampling)()
	u = rand(s.rng)
	i = searchsortedfirst(s.cp, u)
	return i, s.p[i]
end

################################################################################
# Variance Reduced Algorithms
################################################################################
abstract type VRAlgorithm end

update!(alg::VRAlgorithm) = nothing
update!(alg::VRAlgorithm, iter) = update!(alg::VRAlgorithm)
update!(alg::VRAlgorithm, iter, stage) = update!(alg::VRAlgorithm, iter)

stageupdate!(alg::VRAlgorithm) =
	error("stageupdate! need to return the stage length")
stageupdate!(alg::VRAlgorithm, stage) = stageupdate!(alg::VRAlgorithm)
stageupdate!(alg::VRAlgorithm, iter, stage) =
	stageupdate!(alg::VRAlgorithm, stage)

initialize!(alg::VRAlgorithm) = nothing

primiterates(alg::VRAlgorithm) = nothing
dualiterates(alg::VRAlgorithm) = nothing

expectedstagelen(alg::VRAlgorithm) = nothing

##############################
# SAGA
function SAGA(vrg, x0, stepsize, rng; weights=nothing, prox_f=IndFree())
	innoweight = length(vrg)
	SVAG(
		vrg, x0, stepsize, innoweight, rng, weights=weights, prox_f=prox_f)
end

##############################
# SAG
function SAG(vrg, x0, stepsize, rng; weights=nothing, prox_f=IndFree())
	innoweight = 1
	SVAG(
		vrg, x0, stepsize, innoweight, rng, weights=weights, prox_f=prox_f)
end

##############################
# SVAG
struct SVAG{X,VRG,PF,T,S} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	innoweight::T
	sampling::S
end
# Constructor
function SVAG(
		vrg, x0, stepsize, innoweight, rng; weights=nothing, prox_f=IndFree())

	T = eltype(x0)
	stepsize = T(stepsize)
	innoweight = T(innoweight)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)

	SVAG(x,x_tmp,x0,vrg,prox_f,stepsize,innoweight,sampling)
end
# Interface
function initialize!(alg::SVAG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::SVAG, iter, stage)
	(i,p) = alg.sampling()

	n = length(alg.vrg)
	innoweight_i = alg.innoweight/(n*n*p)
	vrgrad_store!(alg.x_tmp, alg.vrg, alg.x, i, innoweight_i)
	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::SVAG) = 1
primiterates(alg::SVAG) = alg.x
dualiterates(alg::SVAG) = alg.vrg
expectedstagelen(alg::SVAG) = 1

##############################
# Adaptive SVAG
struct ASVAG{X,VRG,PF,T,S,F} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	innoweight::Base.RefValue{T}
	innov::X
	expinnov::X
	decay::T
	biascorr::Base.RefValue{T}
	sampling::S
	stepcorrection::F
end
# Constructor
function ASVAG(
		vrg, x0, stepsize, rng;
		weights=nothing, prox_f=IndFree(), enforce_step=false)

	T = eltype(x0)
	stepsize = T(stepsize)
	innoweight = T(length(vrg))

	x = similar(x0)
	x_tmp = similar(x0)

	innov = similar(x0)
	expinnov = zero(x0)
	biascorr = zero(T)
	decay = T(1/10)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	if enforce_step
		stepcorrection = (n,theta) -> max(
			2/(2 + abs(n-theta)),
			2/(2 + (n-theta)*(theta-1)/n*((theta-1)/n-1+sign(theta-1)*sqrt(2)) )
			)
	else
		stepcorrection = (n,theta) -> 1
	end
	ASVAG(
		x, x_tmp, x0, vrg, prox_f, stepsize, Ref(innoweight),
		innov, expinnov, decay, Ref(biascorr),
		sampling, stepcorrection
		)
end
# Interface
function initialize!(alg::ASVAG)
	reset!(alg.vrg)
	alg.x .= alg.x0
	alg.innoweight[] = length(alg.vrg)
	alg.expinnov .= zero(eltype(alg.expinnov))
	alg.biascorr[] = zero(alg.biascorr[])
end
function update!(alg::ASVAG, iter, stage)
	(i,p) = alg.sampling()
	n = length(alg.vrg)

	# Update expected innovation
	state = innov!(alg.innov, alg.vrg, alg.x, i)
	alg.expinnov .= alg.innov .+ (1-alg.decay).*alg.expinnov
	alg.biascorr[] = 1 + (1-alg.decay)*alg.biascorr[]

	# Calc innovation weight and make stepsize correction
	alg.innoweight[] =
		n*dot(alg.expinnov, alg.innov) /
		(alg.biascorr[]*norm(alg.innov)^2 + 1e-8)
	lim = n
	alg.innoweight[] = min(lim, alg.innoweight[])
	alg.innoweight[] = max(-lim, alg.innoweight[])

	innoweight_i = alg.innoweight[] / (n*n*p)
	stepsize = alg.stepsize * alg.stepcorrection(n,alg.innoweight[])

	# Take VRSG step
	grad_approx!(alg.x_tmp, alg.vrg)
	alg.x_tmp .= alg.x .- stepsize.*(innoweight_i.*alg.innov .+ alg.x_tmp)
	prox!(alg.x, alg.prox_f, alg.x_tmp, stepsize)
	store_from_innov!(alg.vrg, alg.innov, state)
end
stageupdate!(alg::ASVAG) = 1
primiterates(alg::ASVAG) = alg.x
dualiterates(alg::ASVAG) = alg.vrg
expectedstagelen(alg::ASVAG) = 1

##############################
# SVRG
struct SVRG{X,VRG,PF,T,I,S} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	stagelen::I
	sampling::S
end
# Constructor
function SVRG(vrg, x0, stepsize, stagelen, rng;weights=nothing,prox_f=IndFree())
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	SVRG(x,x_tmp,x0,vrg,prox_f,stepsize,stagelen,sampling)
end
# Interface
function initialize!(alg::SVRG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::SVRG, iter, stage)
	(i,p) = alg.sampling()
	innoweight = 1/(length(alg.vrg)*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)
	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
function stageupdate!(alg::SVRG)
	store_from_point!(alg.vrg, alg.x)
	return alg.stagelen
end
primiterates(alg::SVRG) = alg.x
dualiterates(alg::SVRG) = alg.vrg
expectedstagelen(alg::SVRG) = alg.stagelen

##############################
# Loopless SVRG
struct LSVRG{X,VRG,PF,T,Q,RNG,S} <:VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	q::Q
	rng::RNG
	sampling::S
end
# Constructor
function LSVRG(vrg, x0, stepsize, q, rng; weights=nothing,prox_f=IndFree())
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	LSVRG(x,x_tmp,x0,vrg,prox_f,stepsize,q,rng,sampling)
end
# Interface
function initialize!(alg::LSVRG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::LSVRG, iter, stage)
	(i,p) = alg.sampling()

	innoweight = 1/(length(alg.vrg)*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)

	rand(alg.rng) < alg.q && store_from_point!(alg.vrg, alg.x)

	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::LSVRG) = 1
primiterates(alg::LSVRG) = alg.x
dualiterates(alg::LSVRG) = alg.vrg
expectedstagelen(alg::LSVRG) = 1

##############################
# Incoherent Loopless SVRG
struct ILSVRG{X,VRG,PF,T,Q,RNG,S} <:VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	q::Q
	rng::RNG
	sampling::S
end
# Constructor
function ILSVRG(vrg, x0, stepsize, q, rng; weights=nothing,prox_f=IndFree())
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	ILSVRG(x,x_tmp,x0,vrg,prox_f,stepsize,q,rng,sampling)
end
# Interface
function initialize!(alg::ILSVRG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::ILSVRG, iter, stage)
	(i,p) = alg.sampling()
	n = length(alg.vrg)

	innoweight = 1/(n*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)

	for j = 1:n
		rand(alg.rng) < alg.q && store_from_point!(alg.vrg, alg.x, j)
	end

	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::ILSVRG) = 1
primiterates(alg::ILSVRG) = alg.x
dualiterates(alg::ILSVRG) = alg.vrg
expectedstagelen(alg::ILSVRG) = 1

##############################
# qSAGA with or without replacement in the dual update
struct QSAGA{X,VRG,PF,T,RNG,S} <:VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize::T
	dualindex::Vector{Int}
	replace::Bool
	rng::RNG
	sampling::S
end
# Constructor
function QSAGA(
		vrg, x0, stepsize, q, rng;
		weights=nothing, prox_f=IndFree(), replace=true)
	T = eltype(x0)
	stepsize = T(stepsize)

	x = similar(x0)
	x_tmp = similar(x0)
	dualindex = Vector{Int}(undef,q)

	sampling = (weights == nothing) ?	UniformSampling(rng, length(vrg)) :
										WeightedSampling(rng, weights)
	QSAGA(x,x_tmp,x0,vrg,prox_f,stepsize,dualindex,replace,rng,sampling)
end
# Interface
function initialize!(alg::QSAGA)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::QSAGA, iter, stage)
	n = length(alg.vrg)
	(i,p) = alg.sampling()

	innoweight = 1/(n*p)
	vrgrad!(alg.x_tmp, alg.vrg, alg.x, i, innoweight)

	sample!(alg.rng, 1:n, alg.dualindex, replace=alg.replace)
	store_from_point!(alg.vrg, alg.x, alg.dualindex)

	alg.x_tmp .= alg.x .- alg.stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, alg.stepsize)
end
stageupdate!(alg::QSAGA) = 1
primiterates(alg::QSAGA) = alg.x
dualiterates(alg::QSAGA) = alg.vrg
expectedstagelen(alg::QSAGA) = 1

##############################
# SG, ordinary stoachastic gradient, no variance reduction
struct SG{X,T,VRG,PF,SF,RNG,S} <: VRAlgorithm
	x::X
	x_tmp::X
	x0::X
	vrg::VRG
	prox_f::PF
	stepsize_f::SF
	iter::Base.RefValue{T}
	rng::RNG
	sampling::S
end
# Constructor
function SG(vrg, x0, stepsize_f, rng; weights=nothing, prox_f=IndFree())
	iter = zero(eltype(x0))

	x = similar(x0)
	x_tmp = similar(x0)

	sampling = (weights == nothing) ?
		UniformSampling(rng, length(vrg)) :
		WeightedSampling(rng, weights)

	SG(x,x_tmp,x0,vrg,prox_f,stepsize_f,Ref(iter),rng,sampling)
end
# Interface
function initialize!(alg::SG)
	reset!(alg.vrg)
	alg.x .= alg.x0
end
function update!(alg::SG, iter, stage)
	(i,p) = alg.sampling()

	stepsize = alg.stepsize_f(alg.iter[])
	alg.iter[] += 1

	grad!(alg.x_tmp, alg.vrg, alg.x, i)
	alg.x_tmp .= alg.x .- stepsize.*alg.x_tmp
	prox!(alg.x, alg.prox_f, alg.x_tmp, stepsize)
end
stageupdate!(alg::SG) = 1
primiterates(alg::SG) = alg.x
dualiterates(alg::SG) = alg.vrg
expectedstagelen(alg::SG) = 1


text:
#---------------------------------------------------------------------#
#This function computes the Norm of two arrays
#Written by F.X. Giraldo on April 19, 2019
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function compute_norms(qn,qe,DFloat)

    norms=zeros(DFloat,2,1)
    norms[1]=norm(qn[:,1]-qe[:,1],2)/norm(qe[:,1],2)
    norms[2]=norm(qn[:,2]-qe[:,2],2)/norm(qe[:,2],2)
    return (norms)

end


text:
import GraphNeuralNetworks: compute_message, update_node, update_edge, propagate

@testset "message passing" begin
    in_channel = 10
    out_channel = 5
    num_V = 6
    num_E = 14
    T = Float32

    adj =  [0 1 0 0 0 0
            1 0 0 1 1 1
            0 0 0 0 0 1
            0 1 0 0 1 0
            0 1 0 1 0 1
            0 1 1 0 1 0]

    X = rand(T, in_channel, num_V)
    E = rand(T, in_channel, num_E)
    U = rand(T, in_channel)

    @testset "no aggregation" begin
        struct NewLayer1{G} end

        (l::NewLayer1{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, nothing)

        l = NewLayer1{GRAPH_T}()
        g = GNNGraph(adj, ndata=X, graph_type=GRAPH_T)
        g_ = l(g)

        @test adjacency_matrix(g_) == adj
        @test node_features(g_) === nothing
        @test edge_features(g_)  === nothing
        @test graph_features(g_) === nothing
    end

    @testset "neighbor aggregation (+)" begin
        struct NewLayer2{G} end

        (l::NewLayer2{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)

        l = NewLayer2{GRAPH_T}()
        g = GNNGraph(adj, ndata=X, edata=E, gdata=U, graph_type=GRAPH_T)
        g_ = l(g)

        @test adjacency_matrix(g_) == adj
        @test size(node_features(g_)) == (in_channel, num_V)
        @test edge_features(g_) ≈ E
        @test graph_features(g_) ≈ U
    end

    @testset "custom message and neighbor aggregation" begin
        struct NewLayer3{G} end

        GraphNeuralNetworks.compute_message(l::NewLayer3{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
        (l::NewLayer3{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)

        l = NewLayer3{GRAPH_T}()
        g = GNNGraph(adj, ndata=X, edata=E, gdata=U, graph_type=GRAPH_T)
        g_ = l(g)

        @test adjacency_matrix(g_) == adj
        @test size(node_features(g_)) == (out_channel, num_V)
        @test edge_features(g_) ≈ edge_features(g)
        @test graph_features(g_) ≈ graph_features(g)
    end

    @testset "update_edge" begin
        struct NewLayer4{G} end

        GraphNeuralNetworks.update_edge(l::NewLayer4{GRAPH_T}, m, e) = m
        GraphNeuralNetworks.compute_message(l::NewLayer4{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
        (l::NewLayer4{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)

        l = NewLayer4{GRAPH_T}()
        g = GNNGraph(adj, ndata=X, edata=E, gdata=U, graph_type=GRAPH_T)
        g_ = l(g)

        @test adjacency_matrix(g_) == adj
        @test size(node_features(g_)) == (out_channel, num_V)
        @test size(edge_features(g_)) == (out_channel, num_E)
        @test graph_features(g_) ≈ graph_features(g)
    end

    @testset "update edge/vertex" begin
        struct NewLayer5{G} end

        GraphNeuralNetworks.update_node(l::NewLayer5{GRAPH_T}, m̄, xi) = rand(T, 2*out_channel, size(xi, 2))
        GraphNeuralNetworks.update_edge(l::NewLayer5{GRAPH_T}, m, e) = m
        GraphNeuralNetworks.compute_message(l::NewLayer5{GRAPH_T}, xi, xj, e) = ones(T, out_channel, size(e,2))
        (l::NewLayer5{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)

        l = NewLayer5{GRAPH_T}()
        g = GNNGraph(adj, ndata=X, edata=E, gdata=U, graph_type=GRAPH_T)
        g_ = l(g)

        @test all(adjacency_matrix(g_) .== adj)
        @test size(node_features(g_)) == (2*out_channel, num_V)
        @test size(edge_features(g_)) == (out_channel, num_E)
        @test size(graph_features(g_)) == (in_channel, 1)
    end

    @testset "message and update with weights" begin

        struct NewLayerW{G}
            weight
        end

        NewLayerW(in, out) = NewLayerW{GRAPH_T}(randn(T, out, in))

        GraphNeuralNetworks.compute_message(l::NewLayerW{GRAPH_T}, x_i, x_j, e_ij) = l.weight * x_j
        GraphNeuralNetworks.update_node(l::NewLayerW{GRAPH_T}, m, x) = l.weight * x + m

        (l::NewLayerW{GRAPH_T})(g) = GraphNeuralNetworks.propagate(l, g, +)

        l = NewLayerW(in_channel, out_channel)
        g = GNNGraph(adj, ndata=X, edata=E, gdata=U, graph_type=GRAPH_T)
        g_ = l(g)

        @test adjacency_matrix(g_) == adj
        @test size(node_features(g_)) == (out_channel, num_V)
        @test edge_features(g_) === E
        @test vec(graph_features(g_)) ≈ U
    end

    @testset "NamedTuples" begin
        struct NewLayerNT{G}
            W
        end

        NewLayerNT(in, out) = NewLayerNT{GRAPH_T}(randn(T, out, in))

        function GraphNeuralNetworks.compute_message(l::NewLayerNT{GRAPH_T}, di, dj, dij)
            a = l.W * (di.x .+ dj.x .+ dij.e)
            b = l.W * di.x
            return (; a, b)
        end
        function GraphNeuralNetworks.update_node(l::NewLayerNT{GRAPH_T}, m, d)
            return (α=l.W * d.x + m.a + m.b, β=m)
        end
        function GraphNeuralNetworks.update_edge(l::NewLayerNT{GRAPH_T}, m, e)
            return m.a
        end

        function (::NewLayerNT{GRAPH_T})(g, x, e)
            x, e = propagate(l, g, mean, (; x), (; e))
            return x.α .+ x.β.a, e
        end

        l = NewLayerNT(in_channel, out_channel)
        g = GNNGraph(adj, graph_type=GRAPH_T)
        X′, E′ = l(g, X, E)

        @test size(X′) == (out_channel, num_V)
        @test size(E′) == (out_channel, num_E)
    end
end


text:
module MOMDPs

include("reexport.jl")
@reexport using POMDPs

export
    MOMDP,
    create_fully_obs_transition,
    create_partially_obs_transition,
    fully_obs_space,
    part_obs_space

abstract MOMDP <: POMDP

create_fully_obs_transition(pomdp::MOMDP) = error("$(typeof(pomdp)) does not implement create_fully_obs_transition")
create_partially_obs_transition(pomdp::MOMDP) = error("$(typeof(pomdp)) does not implement create_part_obs_transition")

fully_obs_space(pomdp::POMDP) = error("$(typeof(pomdp)) does not implement fully_obs_space")
part_obs_space(pomdp::POMDP) = error("$(typeof(pomdp)) does not implement part_obs_space")

end # module


text:
const Idx = Union{Real,Colon,AbstractArray{Int}}

using Base: ViewIndex, @propagate_inbounds, tail

abstract type Value{T} end

struct TolValue{T} <: Value{T}
    val::T
    tol::T
end

TolValue(x, tol=Base.rtoldefault(typeof(x))*abs(x)) = TolValue(promote(x,tol)...)

struct ExactValue{T} <: Value{T}
    val::T
end

atvalue(x::Number; rtol=Base.rtoldefault(typeof(x)), atol=zero(x)) = TolValue(x, atol+rtol*abs(x))
atvalue(x) = ExactValue(x)

const Values = AbstractArray{<:Value}

# For throwing a BoundsError with a Value index, we need to define the following
# (note that we could inherit them for free, were Value <: Number)
Base.iterate(x::Value, state = false) = state ? nothing : (x, true)

# Values have the indexing trait of their wrapped type
_axistrait_el(::Type{<:Value{T}}) where {T} = _axistrait_el(T)

# How to show Value objects (e.g. in a BoundsError)
Base.show(io::IO, v::TolValue) =
    print(io, string("TolValue(", v.val, ", tol=", v.tol, ")"))
Base.show(io::IO, v::ExactValue) = print(io, string("ExactValue(", v.val, ")"))

# Defer IndexStyle to the wrapped array
Base.IndexStyle(::Type{AxisArray{T,N,D,Ax}}) where {T,N,D,Ax} = IndexStyle(D)

# Simple scalar indexing where we just set or return scalars
@propagate_inbounds Base.getindex(A::AxisArray, idxs::Int...) = A.data[idxs...]
@propagate_inbounds Base.setindex!(A::AxisArray, v, idxs::Int...) = (A.data[idxs...] = v)

# Cartesian iteration
Base.eachindex(A::AxisArray) = eachindex(A.data)

# Avoid an ambiguity -- IntervalSets@0.4 takes .. from EllipsisNotation,
# which defines A[..] for any AbstractArray, like this:
Base.getindex(A::AxisArray, ::typeof(..)) = A

"""
    reaxis(A::AxisArray, I...)

This internal function determines the new set of axes that are constructed upon
indexing with I.
"""
reaxis(A::AxisArray, I::Idx...) = _reaxis(make_axes_match(axes(A), I), I)
# Linear indexing
reaxis(A::AxisArray{<:Any,1}, I::AbstractArray{Int}) = _new_axes(A.axes[1], I)
reaxis(A::AxisArray, I::AbstractArray{Int}) = default_axes(I)
reaxis(A::AxisArray{<:Any,1}, I::Real) = ()
reaxis(A::AxisArray, I::Real) = ()
reaxis(A::AxisArray{<:Any,1}, I::Colon) = _new_axes(A.axes[1], Base.axes(A, 1))
reaxis(A::AxisArray, I::Colon) = default_axes(Base.OneTo(length(A)))
reaxis(A::AxisArray{<:Any,1}, I::AbstractArray{Bool}) = _new_axes(A.axes[1], findall(I))
reaxis(A::AxisArray, I::AbstractArray{Bool}) = default_axes(findall(I))

# Ensure the number of axes matches the number of indexing dimensions
@inline function make_axes_match(axs, idxs)
    nidxs = Base.index_ndims(idxs...)
    ntuple(i->(Base.@_inline_meta; _default_axis(i > length(axs) ? Base.OneTo(1) : axs[i], i)), length(nidxs))
end

# Now we can reaxis without worrying about mismatched axes/indices
@inline _reaxis(axs::Tuple{}, idxs::Tuple{}) = ()
# Scalars are dropped
const ScalarIndex = Union{Real, AbstractArray{<:Any, 0}}
@inline _reaxis(axs::Tuple, idxs::Tuple{ScalarIndex, Vararg{Any}}) = _reaxis(tail(axs), tail(idxs))
# Colon passes straight through
@inline _reaxis(axs::Tuple, idxs::Tuple{Colon, Vararg{Any}}) = (axs[1], _reaxis(tail(axs), tail(idxs))...)
# But arrays can add or change dimensions and accompanying axis names
@inline _reaxis(axs::Tuple, idxs::Tuple{AbstractArray, Vararg{Any}}) =
    (_new_axes(axs[1], idxs[1])..., _reaxis(tail(axs), tail(idxs))...)

# Vectors simply create new axes with the same name; just subsetted by their value
@inline _new_axes(ax::Axis{name}, idx::AbstractVector) where {name} = (Axis{name}(ax.val[idx]),)
# Arrays create multiple axes with _N appended to the axis name containing their indices
@generated function _new_axes(ax::Axis{name}, idx::AbstractArray{<:Any,N}) where {name,N}
    newaxes = Expr(:tuple)
    for i=1:N
        push!(newaxes.args, :($(Axis{Symbol(name, "_", i)})(Base.axes(idx, $i))))
    end
    newaxes
end
# And indexing with an AxisArray joins the name and overrides the values
@generated function _new_axes(ax::Axis{name}, idx::AxisArray{<:Any, N}) where {name,N}
    newaxes = Expr(:tuple)
    idxnames = axisnames(idx)
    for i=1:N
        push!(newaxes.args, :($(Axis{Symbol(name, "_", idxnames[i])})(idx.axes[$i].val)))
    end
    newaxes
end

@propagate_inbounds function Base.getindex(A::AxisArray, idxs::Idx...)
    AxisArray(A.data[idxs...], reaxis(A, idxs...))
end

# To resolve ambiguities, we need several definitions
using Base: AbstractCartesianIndex
@propagate_inbounds Base.view(A::AxisArray, idxs::Idx...) = AxisArray(view(A.data, idxs...), reaxis(A, idxs...))

# Setindex is so much simpler. Just assign it to the data:
@propagate_inbounds Base.setindex!(A::AxisArray, v, idxs::Idx...) = (A.data[idxs...] = v)

# Logical indexing
@propagate_inbounds function Base.getindex(A::AxisArray, idx::AbstractArray{Bool})
    AxisArray(A.data[idx], reaxis(A, idx))
end
@propagate_inbounds Base.setindex!(A::AxisArray, v, idx::AbstractArray{Bool}) = (A.data[idx] = v)

### Fancier indexing capabilities provided only by AxisArrays ###
# To avoid StackOverflowErrors on indexes that we don't know how to convert, we
# give AxisArrays once chance to convert into known format and then defer to the parent
@propagate_inbounds Base.getindex(A::AxisArray, idxs...) = getindex_converted(A, to_index(A,idxs...)...)
@propagate_inbounds Base.setindex!(A::AxisArray, v, idxs...) = setindex!_converted(A, v, to_index(A,idxs...)...)
# Deal with lots of ambiguities here
@propagate_inbounds Base.view(A::AxisArray, idxs::ViewIndex...) = view(A, to_index(A,idxs...)...)
@propagate_inbounds Base.view(A::AxisArray, idxs::Union{ViewIndex,AbstractCartesianIndex}...) = view(A, to_index(A,Base.IteratorsMD.flatten(idxs)...)...)
@propagate_inbounds Base.view(A::AxisArray, idxs...) = view(A, to_index(A,idxs...)...)

@propagate_inbounds getindex_converted(A, idxs::Idx...) = A[idxs...]
@propagate_inbounds setindex!_converted(A, v, idxs::Idx...) = (A[idxs...] = v)

@propagate_inbounds getindex_converted(A, idxs...) = A.data[idxs...]
@propagate_inbounds setindex!_converted(A, v, idxs...) = (A.data[idxs...] = v)

# First is indexing by named axis. We simply sort the axes and re-dispatch.
# When indexing by named axis the shapes of omitted dimensions are preserved
# TODO: should we handle multidimensional Axis indexes? It could be interpreted
#       as adding dimensions in the middle of an AxisArray.
# TODO: should we allow repeated axes? As a union of indices of the duplicates?
@generated function to_index(A::AxisArray{T,N,D,Ax}, I::Axis...) where {T,N,D,Ax}
    dims = Int[axisdim(A, ax) for ax in I]
    idxs = Expr[:(Colon()) for d = 1:N]
    names = axisnames(A)
    for i=1:length(dims)
        idxs[dims[i]] == :(Colon()) ||
            return :(throw(ArgumentError(string("multiple indices provided ",
                "on axis ", $(string(names[dims[i]]))))))
        idxs[dims[i]] = :(I[$i].val)
    end

    meta = Expr(:meta, :inline)
    return :($meta; to_index(A, $(idxs...)))
end

function Base.reshape(A::AxisArray, ::Val{N}) where N
    axN, _ = Base.IteratorsMD.split(axes(A), Val(N))
    AxisArray(reshape(A.data, Val(N)), Base.front(axN))
end

# Keyword indexing, reconstructs the Axis{}() objects
@propagate_inbounds Base.view(A::AxisArray; kw...) =
    view(A, kw_to_axes(parent(A), kw.data)...)
@propagate_inbounds Base.getindex(A::AxisArray; kw...) =
    getindex(A, kw_to_axes(parent(A), kw.data)...)
@propagate_inbounds Base.setindex!(A::AxisArray, val; kw...) =
    setindex!(A, val, kw_to_axes(parent(A), kw.data)...)

function kw_to_axes(A::AbstractArray, nt::NamedTuple)
    length(nt) == 0 && throw(BoundsError(A, ())) # Trivial case A[] lands here
    nt_to_axes(nt)
end

### Indexing along values of the axes ###

# Default axes indexing throws an error
"""
    axisindexes(ax::Axis, axis_idx) -> array_idx
    axisindexes(::Type{<:AxisTrait}, axis_values, axis_idx) -> array_idx

Translate an index into an axis into an index into the underlying array.
Users can add additional indexing behaviours for custom axes or custom indices by adding
methods to this function.

## Examples

Add a method for indexing into an `Axis{name, SortedSet}`:

```julia
AxisArrays.axisindexes(::Type{Categorical}, ax::SortedSet, idx::AbstractVector) = findin(collect(ax), idx)
```

Add a method for indexing into a `Categorical` axis with a `SortedSet`:

```julia
AxisArrays.axisindexes(::Type{Categorical}, ax::AbstractVector, idx::SortedSet) = findin(ax, idx)
```
"""
axisindexes(ax, idx) = axisindexes(axistrait(ax.val), ax.val, idx)
axisindexes(::Type{Unsupported}, ax, idx) = error("elementwise indexing is not supported for axes of type $(typeof(ax))")
axisindexes(t, ax, idx) = error("cannot index $(typeof(ax)) with $(typeof(idx)); expected $(eltype(ax)) axis value or Integer index")

# Dimensional axes may be indexed directly by their elements if Non-Real and unique
# Maybe extend error message to all <: Numbers if Base allows it?
axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::Real) =
    throw(ArgumentError("invalid index: $idx. Use `atvalue` when indexing by value."))
function axisindexes(::Type{Dimensional}, ax::AbstractVector{T}, idx::T) where T
    idxs = searchsorted(ax, ClosedInterval(idx,idx))
    length(idxs) > 1 && error("more than one datapoint lies on axis value $idx; use an interval to return all values")
    if length(idxs) == 1
        idxs[1]
    else
        throw(BoundsError(ax, idx))
    end
end
function axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::Axis)
    idxs = searchsorted(ax, idx.val)
    length(idxs) > 1 && error("more than one datapoint lies on axis value $idx; use an interval to return all values")
    if length(idxs) == 1
        idxs[1]
    else
        throw(BoundsError(ax, idx))
    end
end
# Dimensional axes may always be indexed by value if in a Value type wrapper.
function axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::TolValue)
    idxs = searchsorted(ax, ClosedInterval(idx.val,idx.val))
    length(idxs) > 1 && error("more than one datapoint lies on axis value $idx; use an interval to return all values")
    if length(idxs) == 1
        idxs[1]
    else # it's zero
        last(idxs) > 0 && abs(ax[last(idxs)] - idx.val) < idx.tol && return last(idxs)
        first(idxs) <= length(ax) && abs(ax[first(idxs)] - idx.val) < idx.tol && return first(idxs)
        throw(BoundsError(ax, idx))
    end
end
function axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::ExactValue)
    idxs = searchsorted(ax, ClosedInterval(idx.val,idx.val))
    length(idxs) > 1 && error("more than one datapoint lies on axis value $idx; use an interval to return all values")
    if length(idxs) == 1
        idxs[1]
    else # it's zero
        throw(BoundsError(ax, idx))
    end
end
# For index types that AxisArrays doesn't know about
axisindexes(::Type{Dimensional}, ax::AbstractVector, idx) = idx

# Dimensional axes may be indexed by intervals to select a range
axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::ClosedInterval) = searchsorted(ax, idx)

# Or repeated intervals, which only work if the axis is a range since otherwise
# there will be a non-constant number of indices in each repetition.
#
# There are a number of challenges here:
#   * This operation adds a dimension to the result; rows represent the interval
#     (or subset) and columns are offsets (or repetition). A RepeatedRangeMatrix
#     represents the resulting matrix of indices very nicely.
#   * We also want the returned matrix to keep track of its axes; the axis
#     subset (ax_sub) is the relative location of the interval with respect to
#     each offset, and the repetitions (ax_rep) is the array of offsets.
#   * We are interested in the resulting *addition* of the interval against the
#     offsets. Either the offsets or the interval may independently be out of
#     bounds prior to this addition. Even worse: the interval may have different
#     units than the axis (e.g., `(Day(-1)..Day(1)) + dates` for a three-day
#     span around dates of interest over a Date axis).
#   * It is possible (and likely!) that neither the interval endpoints nor the
#     offsets fall exactly upon an axis value. Or even worse: the some offsets
#     when added to the interval could span more elements than others (the
#     fencepost problem). As such, we need to be careful about how and when we
#     snap the provided intervals and offsets to exact axis values (and indices).
#
# To avoid the fencepost problems and to define the axes, we convert the
# interval to a UnitRange of relative indices and the array of offsets to an
# array of absolute indices (independently of each other). Exactly how we do so
# must be carefully considered.
#
# Note that this is fundamentally different than indexing by a single interval;
# whereas those intervals are specified in the same units as the elements of the
# axis itself, these intervals are specified in terms of _offsets_. At the same
# time, we want `A[interval] == vec(A[interval + [0]])`. To make these
# computations as similar as possible, we use a phony range of the form
# `step(ax):step(ax):step(ax)` in order to search for the interval.
phony_range(r::AbstractRange) = step(r):step(r):step(r)
phony_range(r::AbstractUnitRange) = step(r):step(r)
phony_range(r::StepRangeLen) = StepRangeLen(r.step, r.step, 1)
function relativewindow(r::AbstractRange, x::ClosedInterval)
    pr = phony_range(r)
    idxs = Extrapolated.searchsorted(pr, x)
    vals = Extrapolated.getindex(pr, idxs)
    return (idxs, vals)
end

axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::RepeatedInterval) = error("repeated intervals might select a varying number of elements for non-range axes; use a repeated Range of indices instead")
function axisindexes(::Type{Dimensional}, ax::AbstractRange, idx::RepeatedInterval)
    idxs, vals = relativewindow(ax, idx.window)
    offsets = [Extrapolated.searchsortednearest(ax, offset) for offset in idx.offsets]
    AxisArray(RepeatedRangeMatrix(idxs, offsets), Axis{:sub}(vals), Axis{:rep}(Extrapolated.getindex(ax, offsets)))
end

# We also have special datatypes to represent intervals about indices
axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::IntervalAtIndex) = searchsorted(ax, idx.window + ax[idx.index])
function axisindexes(::Type{Dimensional}, ax::AbstractRange, idx::IntervalAtIndex)
    idxs, vals = relativewindow(ax, idx.window)
    AxisArray(idxs .+ idx.index, Axis{:sub}(vals))
end
axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::RepeatedIntervalAtIndexes) = error("repeated intervals might select a varying number of elements for non-range axes; use a repeated Range of indices instead")
function axisindexes(::Type{Dimensional}, ax::AbstractRange,
                     idx::RepeatedIntervalAtIndexes)
    idxs, vals = relativewindow(ax, idx.window)
    AxisArray(RepeatedRangeMatrix(idxs, idx.indexes), Axis{:sub}(vals), Axis{:rep}(ax[idx.indexes]))
end

# Categorical axes may be indexed by their elements
function axisindexes(::Type{Categorical}, ax::AbstractVector, idx)
    i = findfirst(isequal(idx), ax)
    i === nothing && throw(ArgumentError("index $idx not found"))
    i
end
function axisindexes(::Type{Categorical}, ax::AbstractVector, idx::Value)
    val = idx.val
    i = findfirst(isequal(val), ax)
    i === nothing && throw(ArgumentError("index $val not found"))
    i
end
# Categorical axes may be indexed by a vector of their elements
function axisindexes(::Type{Categorical}, ax::AbstractVector, idx::AbstractVector)
    res = findall(in(idx), ax)
    length(res) == length(idx) || throw(ArgumentError("index $(setdiff(idx,ax)) not found"))
    res
end

# Creates *instances* of axis traits for a set of axes.
# TODO: Transition axistrait() to return trait instances in line with common
# practice in Base and other packages.
#
# This function is a utility tool to ensure that `axistrait` is only called
# from outside the generated function below. (If not, we can get world age
# errors.)
_axistraits(ax1, rest...) = (axistrait(ax1)(), _axistraits(rest...)...)
_axistraits() = ()

# This catch-all method attempts to convert any axis-specific non-standard
# indexing types to their integer or integer range equivalents using axisindexes
# It is separate from the `Base.getindex` function to allow reuse between
# set- and get- index.
@inline to_index(A::AxisArray, I...) = _to_index(A, _axistraits(I...), I...)

@generated function _to_index(A::AxisArray{T,N,D,Ax}, axtraits, I...) where {T,N,D,Ax}
    ex = Expr(:tuple)
    n = 0
    axtrait_types = axtraits.parameters
    for i=1:length(I)
        if axtrait_types[i] <: Categorical && i <= length(Ax.parameters)
            if I[i] <: Axis
                push!(ex.args, :(axisindexes(A.axes[$i], I[$i].val)))
            else
                push!(ex.args, :(axisindexes(A.axes[$i], I[$i])))
            end
            n += 1

            continue
        end

        if I[i] <: Idx
            push!(ex.args, :(I[$i]))
            n += 1
        elseif I[i] <: AbstractArray{Bool}
            push!(ex.args, :(findall(I[$i])))
            n += 1
        elseif I[i] <: Values
            push!(ex.args, :(axisindexes.(Ref(A.axes[$i]), I[$i])))
            n += 1
        elseif I[i] <: CartesianIndex
            for j = 1:length(I[i])
                push!(ex.args, :(I[$i][$j]))
            end
            n += length(I[i])
        elseif i <= length(Ax.parameters)
            if I[i] <: Axis
                push!(ex.args, :(axisindexes(A.axes[$i], I[$i].val)))
            else
                push!(ex.args, :(axisindexes(A.axes[$i], I[$i])))
            end
            n += 1
        else
            push!(ex.args, :(error("dimension ", $i, " does not have an axis to index")))
        end
    end
    for _=n+1:N
        push!(ex.args, :(Colon()))
    end
    meta = Expr(:meta, :inline)
    return :($meta; $ex)
end

## Extracting the full axis (name + values) from the Axis{:name} type
@inline Base.getindex(A::AxisArray, ::Type{Ax}) where {Ax<:Axis} = getaxis(Ax, axes(A)...)
@inline getaxis(::Type{Ax}, ax::Ax, axs...) where {Ax<:Axis} = ax
@inline getaxis(::Type{Ax}, ax::Axis, axs...) where {Ax<:Axis} = getaxis(Ax, axs...)
@noinline getaxis(::Type{Ax}) where {Ax<:Axis} = throw(ArgumentError("no axis of type $Ax was found"))

# Boundschecking specialization: defer to the data array.
# Note that we could unwrap AxisArrays when they are used as indices into other
# arrays within Base's to_index/to_indices methods, but that requires a bigger
# refactor to merge our to_index method with Base's.
@inline Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, A::AxisArray) = Base.checkindex(Bool, inds, A.data)


text:
module Hills

using Flux
using Juno
using Plots
import Base: push!, append!
import Flux: train!

struct Container
    n::Int64
    data::Vector{Tuple{Vector{Float64}, Vector{Float64}}}
    function Container(N::Int64)
        return new(N, Vector{Tuple{Vector{Float64}, Vector{Float64}}}(0))
    end
end

function append!(this::Container,
                orig::Matrix{Float64},
                mask::Matrix{Float64}, can_invert::Bool = true)

    if can_invert
        if mean(orig) > 0.5
            orig = 1.0 - orig
        end
    end

    for k = 1 : size(orig, 1)
        n = min(this.n, size(orig, 2))
        sx = zeros(this.n)
        sy = zeros(this.n)

        sx = orig[k, 1 : n]
        sy = mask[k, 1 : n]

        push!(this, sx, sy)
    end

    for k = 1 : size(orig, 2)
        n = min(this.n, size(orig, 1))
        sx = zeros(this.n)
        sy = zeros(this.n)

        sx = orig[1 : n, k]
        sy = mask[1 : n, k]

        push!(this, sx, sy)
    end

    this
end

function push!(this::Container, x::Vector{Float64}, y::Vector{Float64})
    xn = maximum(x) == 0.0 ? x : x / maximum(x)
    yn = maximum(y) == 0.0 ? y : y / maximum(y)

    push!(this.data, (xn, yn))
end

function loss(x::Vector{Float64}, y::Vector{Float64}, mdl::Flux.Chain)
    return Flux.mse(mdl(x), y)
end

function loss(this::Tuple{Vector{Float64}, Vector{Float64}}, mdl::Flux.Chain)
    return loss(this[1], this[2], mdl)
end

function train(this::Container; mdl = get_model(this.n))
    numepochs = 100
    opt = Flux.Nesterov(params(mdl), 0.1)

    ntrain = Int64(round(length(this.data)*0.5))
    jtrain = unique(rand(1:length(this.data), ntrain))
    jtest = setdiff(1:length(this.data), jtrain)

    train_losses = Vector{Tuple{Float64, Float64}}(0)
    test_losses  = Vector{Tuple{Float64, Float64}}(0)

    Juno.progress(name="Optimizing") do prog
        testdata  = this.data[jtest]
        traindata = this.data[jtrain]

        for k = 1 : numepochs
            Flux.train!((x, y) -> loss(x, y, mdl), traindata, opt)
            this_train_losses = [x.tracker.data for x in loss.(traindata, mdl)]
            this_test_losses = [x.tracker.data for x in loss.(testdata, mdl)]

            msg = @sprintf("%2.4f | %2.4f", mean(this_test_losses), std(this_test_losses))
            Juno.progress(prog, k / numepochs)
            Juno.right_text(prog, msg)

            push!(train_losses, (mean(this_train_losses), std(this_train_losses)))
            push!(test_losses, (mean(this_test_losses), std(this_test_losses)))

            tr_m = [x[1] for x in train_losses]
            tr_s = [x[2] for x in train_losses]
            p1 = scatter([tr_m tr_s], label = ["mean", "std"], title="train")

            ts_m = [x[1] for x in test_losses]
            ts_s = [x[2] for x in test_losses]
            p2 = scatter([ts_m ts_s], label = ["mean", "std"], title="test")

            display(plot(p1, p2, layout = grid(2, 1)))
        end
    end

    return mdl, train_losses, test_losses
end

function get_model(n::Int64)
    return Chain(Dense(n, n, Flux.σ),
                Dense(n, n, Flux.σ),
                Dense(n, n, Flux.σ),
                Dense(n, n, Flux.σ))
end

function filter_image(img::Matrix{Float64}, mdl::Chain, N::Int64)
    outx = zeros(size(img))
    outy = zeros(size(img))

    for k = 1 : size(img, 1)
        n = min(size(img, 2), N)
        outx[k, 1 : n] = mdl(img[k, 1 : n]).data
    end

    for k = 1 : size(img, 2)
        n = min(size(img, 1), N)
        outy[1 : n, k] = mdl(img[1 : n, k]).data
    end

    (outx, outy)
end

end


text:
function fba(lp_model::JuMP.Model, obj_idx::Integer;
        sense = MAX_SENSE, drop_LPsol = true
    )

    optimize!(lp_model, obj_idx; sense)

    return FBAOut(lp_model, obj_idx; drop_LPsol)
end

function fba(S, b, lb, ub, obj_idx::Integer;
        sense = MAX_SENSE,
        drop_LPsol = true,
        lp_model = nothing,
        up_lb_con = false,
        up_ub_con = false,
        up_stoi_con = false,
        solver = Clp.Optimizer
    )

    # build model
    if isnothing(lp_model)
        lp_model = build_lp_model(S, b, lb, ub, solver)
    else
        # update cons
        up_stoi_con && set_stoi_con!(lp_model, S, b)
        up_lb_con && set_lb_con!(lp_model, lb)
        up_ub_con && set_ub_con!(lp_model, ub)
    end

    fba(lp_model, obj_idx; sense, drop_LPsol)
end

function fba(S, b, lb, ub, idx1::Integer, idx2::Integer;
        sense1 = MAX_SENSE,
        sense2 = MIN_SENSE,
        btol = 0.0, # tol of the fixation
        drop_LPsol = true,
        lp_model = nothing,
        up_lb_con = false,
        up_ub_con = false,
        up_stoi_con = false,
        solver = Clp.Optimizer
    )

    # build model
    if isnothing(lp_model)
        lp_model = build_lp_model(S, b, lb, ub, solver)
    else
        # update cons
        up_stoi_con && set_stoi_con!(lp_model, S, b)
        up_lb_con && set_lb_con!(lp_model, lb)
        up_ub_con && set_ub_con!(lp_model, ub)
    end

    # optimization idx1
    fbaout1 = fba(lp_model, idx1; sense = sense1, drop_LPsol)
    isempty(fbaout1) && return fbaout1
    val1 = objval(fbaout1)

    # fix obj1
    set_lb_con!(lp_model, idx1, val1 - btol)
    set_ub_con!(lp_model, idx1, val1 + btol)

    # optimization idx2
    fbaout2 = fba(lp_model, idx2; sense = sense2, drop_LPsol)

    # set back bounds
    set_lb_con!(lp_model, idx1, lb[idx1])
    set_ub_con!(lp_model, idx1, ub[idx1])

    return fbaout2
end

function fba(model::MetNet, obj_ider::IDER_TYPE; kwargs...)
    obj_idx = rxnindex(model, obj_ider)
    model_fields = _extract_dense(model, [:S, :b, :lb, :ub])
    return fba(model_fields..., obj_idx; kwargs...)
end

function fba(model::MetNet, ider1::IDER_TYPE, ider2::IDER_TYPE; kwargs...)
    idx1 = rxnindex(model, ider1)
    idx2 = rxnindex(model, ider2)
    model_fields = _extract_dense(model, [:S, :b, :lb, :ub])
    return fba(model_fields..., idx1, idx2; kwargs...)
end


text:
using Test, OffsetArrays
using VanFoFy.Ellipticals: Theta, theta

@testset "θ-functions" begin
    ω1 = complex(1.0)
    ω3 = exp(1im)
    tol_check = 1e-15
    tol_compute = 1e-20
    θ = Theta(ω1, ω3, tol_compute)

    z = 0.63256929128615824176 + 0.47122375149242207160im

    th = OffsetArray{ComplexF64}(undef, 1:4, 0:4)

    for n in 0:4
        for k in 1:4
            th[k, n] = theta(θ, th_k=k, d_n=n, z=z)
        end
    end

    #####

    th_ma = [ 0.46214180304899833678 + 0.65513824162092381551im,
               0.96467023166692638429 + 0.11945667130003175272im,
               1.13835201868809736642 + 0.08123099419582718546im,
               0.86149287144921489500 - 0.08153295492698743884im]

    @test th[:, 0] ≈ th_ma atol=tol_check

    #####

    d1_th_ma = [ 0.97727706225681589769 + 0.08049915546539055744im,
             -0.42045304182205492196 - 0.63824580596766239552im,
              0.14241824316243509356 - 0.38564309515712238239im,
             -0.14363258751912619862 + 0.38620176200689986572im]

    @test th[:, 1] ≈ d1_th_ma atol=tol_check

    #####

    d2_th_ma = [-0.54549882845031761232 - 0.68893198268462155795im,
             -0.93946558497506382639 - 0.19739257822094024267im,
             -0.55247742739635853599 - 0.32311219018650595886im,
              0.55495918519936412129 + 0.32794356188507020189im]

    @test th[:, 2] ≈ d2_th_ma atol=tol_check

    #####

    d3_th_ma = [-1.05294508933080125267 + 0.15318331373277377347im,
              0.17025898335751795341 + 0.53691780124847588798im,
             -0.5623867730174600491 + 1.5392204494203911042im,
              0.5818162827245192532 - 1.5481591190168449923im]

    @test th[:, 3] ≈ d3_th_ma atol=tol_check

    #####

    d4_th_ma = [1.2954660922120757234 + 0.9931820896285708616im,
             0.71273193831970297909 + 0.89906624330941167631im,
             2.1950195882908139481 + 1.2634597310003502748im,
            -2.2347277131390164682 - 1.3407616781773902616im]

    @test th[:, 4] ≈ d4_th_ma atol=tol_check

    #####

    th1, th2, dth1, dth2 = theta(θ, z)

    @test th1  ≈ theta(θ, th_k=1, d_n=0, z=z) atol=tol_check
    @test th2  ≈ theta(θ, th_k=2, d_n=0, z=z) atol=tol_check
    @test dth1 ≈ theta(θ, th_k=1, d_n=1, z=z) atol=tol_check
    @test dth2 ≈ theta(θ, th_k=2, d_n=1, z=z) atol=tol_check
end


text:
function _precompile_()
    ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
    precompile(OhMyREPL.test_passes, (Base.AbstractIOBuffer{Array{UInt8, 1}}, OhMyREPL.PassHandler, String, Int64, Bool,))
    precompile(OhMyREPL.add_pass!, (OhMyREPL.PassHandler, String, OhMyREPL.Passes.SyntaxHighlighter.SyntaxHighlighterSettings, Bool,))
    precompile(OhMyREPL._find_pass, (OhMyREPL.PassHandler, String,))
    precompile(OhMyREPL.add_pass!, (OhMyREPL.PassHandler, String, OhMyREPL.Passes.BracketHighlighter.BracketHighlighterSettings, Bool,))
    precompile(OhMyREPL.Passes.BracketHighlighter.bracket_match, (Array{Tokenize.Tokens.Token, 1}, Int64,))
    precompile(OhMyREPL.test_pass, (Base.AbstractIOBuffer{Array{UInt8, 1}}, Function, String,))
    precompile(OhMyREPL.apply_passes!, (OhMyREPL.PassHandler, Array{Tokenize.Tokens.Token, 1}, Int64, Bool,))
    precompile(OhMyREPL.add_pass!, (OhMyREPL.PassHandler, String, Function, Bool,))
    precompile(OhMyREPL.test_pass, (Base.AbstractIOBuffer{Array{UInt8, 1}}, OhMyREPL.Passes.BracketHighlighter.BracketHighlighterSettings, String, Int64, Bool,))
    precompile(OhMyREPL.test_pass, (Base.AbstractIOBuffer{Array{UInt8, 1}}, OhMyREPL.Passes.SyntaxHighlighter.SyntaxHighlighterSettings, String, Int64, Bool,))
    precompile(OhMyREPL._check_pass_name, (OhMyREPL.PassHandler, String, Bool,))
    precompile(OhMyREPL.test_pass, (Base.AbstractIOBuffer{Array{UInt8, 1}}, Function, String, Int64, Bool,))
    precompile(OhMyREPL.test_passes, (Base.AbstractIOBuffer{Array{UInt8, 1}}, OhMyREPL.PassHandler, String, Int64,))
    precompile(OhMyREPL.test_pass, (Base.AbstractIOBuffer{Array{UInt8, 1}}, OhMyREPL.Passes.SyntaxHighlighter.SyntaxHighlighterSettings, String,))
    precompile(OhMyREPL.test_pass, (Base.AbstractIOBuffer{Array{UInt8, 1}}, OhMyREPL.Passes.BracketHighlighter.BracketHighlighterSettings, String, Int64,))
    precompile(OhMyREPL.Prompt.rewrite_with_ANSI, (Base.LineEdit.PrefixSearchState, Bool))
    precompile(OhMyREPL.Prompt.rewrite_with_ANSI, (Base.LineEdit.MIState, Bool))
    precompile(OhMyREPL.Prompt.rewrite_with_ANSI, (Base.LineEdit.PromptState, Bool))
    precompile(OhMyREPL.Prompt.refresh_multi_line, (Base.Terminals.TerminalBuffer, Base.Terminals.TTYTerminal, Base.AbstractIOBuffer{Array{UInt8, 1}}, Base.LineEdit.InputAreaState, Int64,))
    precompile(OhMyREPL.Prompt.get_prompt_length, (Base.LineEdit.MIState,))
    precompile(OhMyREPL.ANSICodes.maybe_print, (Base.AbstractIOBuffer{Array{UInt8, 1}}, Bool, OhMyREPL.ANSICodes.ANSIValue, Bool, Bool,))
    precompile(OhMyREPL.ANSICodes.merge, (Crayons.Crayon, Crayons.Crayon,))
    precompile(OhMyREPL.input_prompt!, (String, Symbol,))
    precompile(OhMyREPL._find_pass, (OhMyREPL.PassHandler, String,))
    precompile(OhMyREPL.Prompt.insert_keybindings, ())
    precompile(OhMyREPL.Passes.BracketHighlighter.bracket_match, (Array{Tokenize.Tokens.Token, 1}, Int64,))
    precompile(OhMyREPL.Prompt._commit_line, (Base.LineEdit.MIState, Base.REPL.LineEditREPL, String,))
    precompile(OhMyREPL.Prompt.rewrite_with_ANSI, (Base.LineEdit.MIState, Bool,))
    precompile(OhMyREPL.output_prompt!, (String, Symbol,))
    precompile(OhMyREPL.untokenize_with_ANSI, (Base.AbstractIOBuffer{Array{UInt8, 1}}, Array{Crayons.Crayon, 1}, Array{Tokenize.Tokens.Token, 1}, Int64,))
    precompile(OhMyREPL.ANSICodes.Type, (Type{OhMyREPL.ANSICodes.ANSIValue}, Int64,))
    precompile(OhMyREPL.__init__, ())
    precompile(OhMyREPL.Passes.BracketHighlighter.BracketHighlighterSettings, (OhMyREPL.Passes.BracketHighlighter.BracketHighlighterSettings, Array{Crayons.Crayon, 1}, Array{Tokenize.Tokens.Token, 1}, Int64,))
    precompile(OhMyREPL._check_pass_name, (OhMyREPL.PassHandler, String, Bool,))
    precompile(OhMyREPL.ANSICodes.merge!, (Array{Crayons.Crayon, 1}, Array{Crayons.Crayon, 1},))
    precompile(OhMyREPL.ANSICodes._print, (Base.AbstractIOBuffer{Array{UInt8, 1}}, Crayons.Crayon, Bool,))
    precompile(OhMyREPL.apply_passes!, (OhMyREPL.PassHandler, Array{Tokenize.Tokens.Token, 1}, Int64, Bool,))
    end


text:
## #----------------------------------------------------------------------------------------------#
## #                                  Concrete Type Definitions                                   #
## #----------------------------------------------------------------------------------------------#
##
## #······························································································#
## #                    Concrete Conceptual States -- from Maxwell's Relations                    #
## #······························································································#
##
## # Helmholtz's
## """
## `struct TvPair{𝘁<:AbstractFloat} <: ConceptualState{𝘁}`\n
## The \$(T, v)\$ conceptual Helmholtz state, or unsubstantial intensive property pair.\n
## Conceptual states are just a pair (or a triad) of predefined intensive properties, without
## regard of whether those properties are independent or not, since that analysis would require
## knowledge of substance model.
## """
## struct TvPair{𝘁<:AbstractFloat} <: ConceptualState{𝘁}
##     T::system_T{𝘁}      # The state temperature
##     v::intensive_V{𝘁}   # The state specific volume
##     function TvPair(T::system_T{𝘁}, v::intensive_V{𝘁}) where 𝘁<:AbstractFloat
##         @assert minT(𝘁) <= T            "T < minimum T"
##         @assert minv(𝘁, BO(v)) <= v     "v < minimum v"
##         @assert T <= maxT(𝘁)            "T > maximum T"
##         @assert v <= maxv(𝘁, BO(v))     "v > maximum v"
##         new{𝘁}(T, v)
##     end
## end
##
## export TvPair
##
## "`T(st::TvPair{𝘁})::system_T{𝘁} where 𝘁`"
## T(st::TvPair) = st.T
## "`v(st::TvPair{𝘁})::intensive_V{𝘁} where 𝘁`"
## v(st::TvPair) = st.v
## Tuple(st::TvPair) = (st.T, st.v)
## TvPair(t::Tuple{system_T{𝘁},intensive_V{𝘁}}) where 𝘁 = TvPair(t...)
##
##
## # Gibbs'
## """
## `struct TPPair{𝘁<:AbstractFloat} <: ConceptualState{𝘁}`\n
## The \$(T, P)\$ conceptual Gibbs state, or unsubstantial intensive property pair.\n
## Conceptual states are just a pair (or a triad) of predefined intensive properties, without
## regard of whether those properties are independent or not, since that analysis would require
## knowledge of substance model.
## """
## struct TPPair{𝘁<:AbstractFloat} <: ConceptualState{𝘁}
##     T::system_T{𝘁}      # The state temperature
##     P::system_P{𝘁}      # The state pressure
##     function TPPair(T::system_T{𝘁}, P::system_P{𝘁}) where 𝘁<:AbstractFloat
##         @assert minT(𝘁) <= T            "T < minimum T"
##         @assert minP(𝘁) <= P            "P < minimum P"
##         @assert T <= maxT(𝘁)            "T > maximum T"
##         @assert P <= maxP(𝘁)            "P > maximum P"
##         new{𝘁}(T, P)
##     end
## end
##
## export TPPair
##
## "`T(st::TPPair{𝘁})::system_T{𝘁} where 𝘁`"
## T(st::TPPair) = st.T
## "`P(st::TPPair{𝘁})::system_P{𝘁} where 𝘁`"
## P(st::TPPair) = st.P
## Tuple(st::TPPair) = (st.T, st.P)
## TPPair(t::Tuple{system_T{𝘁},system_P{𝘁}}) where 𝘁 = TPPair(t...)
##
##
## #----------------------------------------------------------------------------------------------#
## #                                      State Type Unions                                       #
## #----------------------------------------------------------------------------------------------#
##
## """
## `anyT{𝘁} = Union{system_T{𝘁},TvPair{𝘁},TPPair{𝘁}} where 𝘁<:AbstractFloat`\n
## A `st::anyT` is any `T`-function-bearing type that retrieves an _actual_ temperature, instead of
## creating one. This _includes_ all types declared in its definition, since an actual temperature
## can be retrieved from those types; however _excludes_ plain `Number`s and
## `Unitful.Quantity`(ies) that have dimensions of temperature, despite the fact that a `T()` call
## with those argument types successfully builds a `system_T`, since a conversion takes place—the
## `Number` wasn't a temperature before the conversion, and `Unitful.Quantity`(ies) conventions are
## broader than this package `AbstractAmount` subtypes.
## """
## anyT{𝘁} = Union{system_T{𝘁},TvPair{𝘁},TPPair{𝘁}} where 𝘁<:AbstractFloat
##
## export anyT
##
##