text:
#
# Doctest-like library for functional testing.
#

module NarrativeTest

export
    runtests

using Test

# Public interface.

"""
    runtests(files; subs=common_subs(), mod=nothing, quiet=false) :: Bool

Load the specified Markdown files to extract and run the embedded test cases.
When a directory is passed, load all `*.md` files in the directory.  This
function returns `true` if the testing is successful, `false` otherwise.

Specify `subs` to customize substitutions applied to the expected output in
order to convert it to a regular expression.

Specify `mod` to execute tests in the context of the given module.

Set `quiet=true` to suppress all output except for error reports.

    runtests(; default=common_args(), subs=common_subs(), mod=nothing, quiet=false)

In this form, test files are specified as command-line parameters.  When
invoked without parameters, tests are loaded from `*.md` files in the program
directory, which can be overriden using `default` parameter.  This function
terminates the program with code `0` if the testing is successful, `1`
otherwise.

Use this form in `test/runtests.jl`:

```julia
using NarrativeTest
NarrativeTest.runtests()
```
"""
function runtests end

"""
    testset(files = nothing;
            default=common_args(), subs=common_subs(), mod=nothing, quiet=false)

Run NarrativeTest-based tests as a nested test set.  For example, invoke it in
`test/runtests.jl`:

```julia
using Test, NarrativeTest

@testset "MyPackage" begin
    …
    NarrativeTest.testset()
    …
end
```

The given Markdown `files` are analyized to extract and run the embedded test
cases.  When a directory is passed, all nested `*.md` files are loaded.  When
invoked without arguments, all `*.md` files in the program directory are
loaded, which can be overridden using the `default` parameter.

Specify `subs` to customize substitutions applied to the expected output
in order to convert it to a regular expression.

Specify `mod` to execute tests in the context of the given module.

Set `quiet=true` to suppress all output except for error reports.

"""
function testset end

# Position in a file.

struct Location
    file::String
    line::Int
end

Location(file::String) = Location(file, 0)

Base.convert(::Type{Location}, file::String) = Location(file, 0)

function Base.print(io::IO, loc::Location)
    print(io, relpath(loc.file))
    if loc.line > 0
        print(io, ":$(loc.line)")
    end
end

function Base.show(io::IO, ::MIME"text/plain", loc::Location)
    print(io, Base.contractuser(abspath(loc.file)))
    if loc.line > 0
        print(io, ":$(loc.line)")
    end
end

Base.convert(::Type{LineNumberNode}, loc::Location) =
    LineNumberNode(loc.line, Symbol(loc.file))

# Text block with position in a file.

struct TextBlock
    loc::Location
    val::String
end

asexpr(code::TextBlock) =
    Base.parse_input_line("\n" ^ max(0, code.loc.line-1) * code.val,
                          filename=abspath(code.loc.file))

collapse(lines::Vector{TextBlock}) =
    !isempty(lines) ? TextBlock(lines[1].loc, join([line.val for line in lines])) : nothing

# Test case.

abstract type AbstractTestCase end

struct TestCase <: AbstractTestCase
    loc::Location
    code::TextBlock
    pre::Union{TextBlock,Nothing}
    expect::Union{TextBlock,Nothing}
end

location(test::TestCase) = test.loc

function Base.show(io::IO, mime::MIME"text/plain", test::TestCase)
    print(io, "Test case at ")
    show(io, mime, test.loc)
    println(io)
    println(io, indented(test.code))
    if test.pre !== nothing
        println(io, "Precondition:")
        println(io, indented(test.pre))
    end
    if test.expect !== nothing
        println(io, "Expected output:")
        println(io, indented(test.expect))
    end
end

Base.convert(::Type{LineNumberNode}, test::TestCase) =
    convert(LineNumberNode, test.loc)

# Test case that failed to parse.

struct BrokenTestCase <: AbstractTestCase
    loc::Location
    msg::String
end

BrokenTestCase(loc, exc::Exception) =
    BrokenTestCase(loc, sprint(showerror, exc))

location(err::BrokenTestCase) = err.loc

function Base.show(io::IO, mime::MIME"text/plain", err::BrokenTestCase)
    print(io, "Error at ")
    show(io, mime, err.loc)
    println(io)
    println(io, indented(err.msg))
end

Base.convert(::Type{LineNumberNode}, err::BrokenTestCase) =
    convert(LineNumberNode, err.loc)

# Test result types.

abstract type AbstractResult end

# Passed test.

struct Pass <: AbstractResult
    test::TestCase
    output::String
end

function Base.show(io::IO, mime::MIME"text/plain", pass::Pass)
    print(io, "Test passed at ")
    show(io, mime, pass.test.loc)
    println(io)
    println(io, indented(pass.test.code))
    println(io, "Expected output:")
    if pass.test.expect !== nothing
        println(io, indented(pass.test.expect))
    end
    println(io, "Actual output:")
    println(io, indented(pass.output))
end

Base.convert(::Type{Test.Result}, pass::Pass) =
    Test.Pass(:test, asexpr(pass.test.code), nothing, nothing)

# Failed test.

struct Fail <: AbstractResult
    test::TestCase
    output::String
    trace::StackTraces.StackTrace
end

function Base.show(io::IO, mime::MIME"text/plain", fail::Fail)
    print(io, "Test failed at ")
    show(io, mime, fail.test.loc)
    println(io)
    println(io, indented(fail.test.code))
    println(io, "Expected output:")
    if fail.test.expect !== nothing
        println(io, indented(fail.test.expect))
    end
    println(io, "Actual output:")
    println(io, indented(fail.output))
    if !isempty(fail.trace)
        println(io, indented(lstrip(sprint(Base.show_backtrace, fail.trace))))
    end
end

Base.convert(::Type{Test.Result}, fail::Fail) =
    Test.Fail(:test, asexpr(fail.test.code), nothing, false,
              convert(LineNumberNode, fail.test))

# Skipped test.

struct Skip <: AbstractResult
    test::TestCase
end

function Base.show(io::IO, mime::MIME"text/plain", skip::Skip)
    print(io, "Test skipped at ")
    show(io, mime, skip.test.loc)
    println(io)
    println(io, indented(skip.test.code))
    println(io, "Failed precondition:")
    if skip.test.pre !== nothing
        println(io, indented(skip.test.pre))
    end
end

# Ill-formed test.

struct Error <: AbstractResult
    test::BrokenTestCase
end

Base.show(io::IO, mime::MIME"text/plain", error::Error) =
    show(io, mime, error.test)

Base.convert(::Type{Test.Result}, error::Error) =
    Test.Error(:test, error.test.msg, nothing, false,
               convert(LineNumberNode, error.test))

# Summary of the testing results.

struct Summary
    passed::Int
    failed::Int
    skipped::Int
    errors::Int
end

function Base.show(io::IO, mime::MIME"text/plain", sum::Summary)
    if sum.passed > 0
        println(io, "Tests passed: ", sum.passed)
    end
    if sum.failed > 0
        println(io, "Tests failed: ", sum.failed)
    end
    if sum.skipped > 0
        println(io, "Tests skipped: ", sum.skipped)
    end
    if sum.errors > 0
        println(io, "Errors: ", sum.errors)
    end
end

# Output formatting.

struct Esc
    val::String
end

Esc(name::Symbol) = Esc(Base.text_colors[name])

Base.show(io::IO, esc::Esc) = nothing
Base.show(io::Base.TTY, esc::Esc) = print(io, esc.val)

const GOOD = Esc(:green)
const INFO = Esc(:cyan)
const WARN = Esc(:light_red)
const BOLD = Esc(:bold)
const NORM = Esc(:normal)
const CLRL = Esc("\r\e[K")

struct Msg
    args::Tuple
end

Msg(args...) = Msg(args)

Base.show(io::IO, msg::Msg) = print(io, msg.args...)

indicator(loc) =
    Msg(CLRL, BOLD, INFO, Esc("Testing " * string(loc)), NORM)
const MORE =
    Msg(BOLD, INFO, Esc("."), NORM)
const SUCCESS =
    Msg(BOLD, GOOD, "TESTING SUCCESSFUL!", NORM)
const FAILURE =
    Msg(BOLD, WARN, "TESTING UNSUCCESSFUL!", NORM)

const SEPARATOR = "~"^72 * "\n"

const INDENT = " "^4

function indented(str::AbstractString)
    inp = IOBuffer(str)
    out = IOBuffer()
    first = true
    for line in eachline(inp)
        if !first
            println(out)
        end
        first = false
        if !isempty(line)
            print(out, INDENT, line)
        end
    end
    return String(take!(out))
end

indented(text::TextBlock) =
    indented(text.val)

# Implementation of `test/runtests.jl`.

function runtests(; default=common_args(), subs=common_subs(), mod=nothing, quiet=false)
    program_file = !isempty(PROGRAM_FILE) ? PROGRAM_FILE : "runtests.jl"
    files = String[]
    for (i, arg) in enumerate(ARGS)
        if startswith(arg, '-')
            if arg == "--"
                append!(files, ARGS[i+1:end])
                break
            elseif arg == "-h" || arg == "--help"
                println("Usage: $program_file [-q|--quiet] [FILE]...")
                exit(0)
            elseif arg == "-q" || arg == "--quiet"
                quiet = true
            else
                println(stderr, """$program_file: unrecognized option '$arg'
                                   Try '$program_file --help' for more information.""")
                exit(2)
            end
        else
            push!(files, arg)
        end
    end
    if isempty(files)
        files = default
    end
    exit(!runtests(files, subs=subs, mod=mod, quiet=quiet))
end

function runtests(files; subs=common_subs(), mod=nothing, quiet=false)
    files = vcat(String[], findmd.(files)...)
    passed = failed = skipped = errors = 0
    for file in files
        suite = parsemd(file)
        for test in suite
            quiet || print(stderr, indicator(location(test)))
            res = runtest(test, subs=subs, mod=mod)
            if res isa Union{Fail, Error}
                quiet || print(stderr, CLRL)
                print(SEPARATOR)
                show(stdout, MIME"text/plain"(), res)
            end
            passed += res isa Pass
            failed += res isa Fail
            skipped += res isa Skip
            errors += res isa Error
        end
        quiet || print(stderr, CLRL)
    end
    summary = Summary(passed, failed, skipped, errors)
    success = failed == 0 && errors == 0
    if !success
        print(SEPARATOR)
    end
    quiet || show(stdout, MIME"text/plain"(), summary)
    quiet || println(stderr, success ? SUCCESS : FAILURE)
    return success
end

# Compatibility with the Test package.

function testset(files = nothing; default=common_args(), subs=common_subs(), mod=nothing, quiet=false)
    files = vcat(String[], findmd.(something(files, default))...)
    ts = Test.DefaultTestSet("NarrativeTest")
    pass = fail = error = 0
    for file in files
        file_ts = Test.DefaultTestSet(relpath(file))
        push!(ts.results, file_ts)
        suite = parsemd(file)
        for test in suite
            quiet || print(stderr, indicator(location(test)))
            res = runtest(test, subs=subs, mod=mod)
            pass += res isa Pass
            fail += res isa Fail
            error += res isa Error
            if res isa Pass
                file_ts.n_passed += 1
            elseif res isa Union{Fail, Error}
                push!(file_ts.results, convert(Test.Result, res))
                file_ts.anynonpass = true
                ts.anynonpass = true
                quiet || print(stderr, CLRL)
                print(SEPARATOR)
                show(stdout, MIME"text/plain"(), res)
            end
        end
        quiet || print(stderr, CLRL)
    end
    if ts.anynonpass
        print(SEPARATOR)
    end
    if Test.get_testset_depth() > 0
        Test.record(Test.get_testset(), ts)
    elseif ts.anynonpass
        throw(Test.TestSetException(pass, fail, error, 0, []))
    end
    ts
end

# Default parameters.

"""
    common_args() :: Vector{String}

Default test files for use when the test runner has no arguments.
"""
common_args() =
    [relpath(dirname(abspath(PROGRAM_FILE)))]

"""
    common_subs() :: Vector{Pair{Regex,SubstitutionString{String}}}

Substitutions applied to the expected output in order to convert
it to a regular expression.
"""
common_subs() = [
    r"[\\^$.[|()?*+{]" => s"\\\0",
    r"[\t ]*…[\t ]*" => s".+",
    r"[\t ]*⋮[\t ]*\r?\n?" => s"(.*(\\n|$))+",
    r"\A" => s"\\A",
    r"\z" => s"\\z",
]

# Find `*.md` files in the given directory.

function findmd(base)
    isdir(base) || return [base]
    mdfiles = String[]
    for (root, dirs, files) in walkdir(base, follow_symlinks=true)
        for file in files
            if splitext(file)[2] == ".md"
                push!(mdfiles, joinpath(root, file))
            end
        end
    end
    sort!(mdfiles)
    return mdfiles
end

# Load a file and extract the embedded tests.

"""
    parsemd(file) :: Vector{AbstractTestCase}
    parsemd(name, io) :: Vector{AbstractTestCase}

Parse the specified Markdown file to extract the embedded test suite; return a
list of test cases.
"""
parsemd(args...) =
    loadfile(parsemd!, args...)

"""
    parsejl(file) :: Vector{AbstractTestCase}
    parsejl(name, io) :: Vector{AbstractTestCase}

Load the specified Julia source file and extract the embedded test suite;
return a list of test cases.
"""
parsejl(args...) =
    loadfile(parsejl!, args...)

loadfile(parsefile!::Function, file::String) =
    loadfile(parsefile!, file, file)

function loadfile(parsefile!::Function, name::String, file::Union{String,IO})
    lines =
        try
            readlines(file)
        catch exc
            return AbstractTestCase[BrokenTestCase(Location(name), exc)]
        end
    stack = [TextBlock(Location(name, i), val * "\n") for (i, val) in enumerate(lines)]
    reverse!(stack)
    return parsefile!(stack)
end

# Extract the test suite from Markdown source.

isblank(line) =
    isempty(strip(line.val))
isfence(line) =
    startswith(line.val, "```") || startswith(line.val, "~~~")
isindent(line) =
    startswith(line.val, " "^4)
isadmonition(line) =
    startswith(line.val, "!!!")

function unindent(line)
    val = line.val[5:end]
    TextBlock(line.loc, !isempty(val) ? val : "\n")
end

function parsemd!(stack::Vector{TextBlock})
    # Accumulated test cases.
    suite = AbstractTestCase[]
    # Extract code snippets, either indented or fenced, and pass them to `parsejl!()`.
    while !isempty(stack)
        line = pop!(stack)
        if isfence(line)
            # Extract a fenced block.
            fenceloc = line.loc
            lang = strip(line.val[4:end])
            jlstack = TextBlock[]
            while !isempty(stack) && !isfence(stack[end])
                push!(jlstack, pop!(stack))
            end
            if isempty(stack)
                push!(suite, BrokenTestCase(fenceloc, "incomplete fenced code block"))
            else
                pop!(stack)
                if isempty(lang)
                    reverse!(jlstack)
                    append!(suite, parsejl!(jlstack))
                end
            end
        elseif isindent(line) && !isblank(line)
            # Extract an indented block.
            jlstack = TextBlock[unindent(line)]
            while !isempty(stack) && (isindent(stack[end]) || isblank(stack[end]))
                push!(jlstack, unindent(pop!(stack)))
            end
            reverse!(jlstack)
            append!(suite, parsejl!(jlstack))
        elseif isadmonition(line)
            # Skip an indented admonition block.
            while !isempty(stack) && (isindent(stack[end]) || isblank(stack[end]))
                pop!(stack)
            end
        end
    end
    return suite
end

# Extract the test suite from Julia source.

function parsejl!(stack::Vector{TextBlock})
    # Accumulated test cases.
    suite = AbstractTestCase[]
    # Find and parse test cases.
    while !isempty(stack)
        if isblank(stack[end])
            pop!(stack)
        else
            l = length(stack)
            push!(suite, parsecase!(stack))
            l != length(stack) || pop!(stack)
        end
    end
    return suite
end

# Extract a test case from Julia source.

function parsecase!(stack::Vector{TextBlock})
    loc = stack[end].loc
    pre = TextBlock[]
    code = TextBlock[]
    expect = TextBlock[]
    # Parse precondition.
    if occursin(r"^\s*#\?\s+", stack[end].val)
        m = match(r"^\s*#\?\s+(.*)$", stack[end].val)
        push!(pre, TextBlock(stack[end].loc, rstrip(m[1])))
        pop!(stack)
        while !isempty(stack) && isblank(stack[end])
            pop!(stack)
        end
    end
    # Parse the code block.
    while !isempty(stack) && !occursin(r"^\s*#(\?|->|=>)\s+", stack[end].val)
        line = pop!(stack)
        if occursin(r"\s*#->\s+", line.val)
            # Code and output on the same line.
            m = match(r"^(.+)\s#->\s+(.*)$", line.val)
            push!(code, TextBlock(line.loc, rstrip(m[1])*"\n"))
            push!(expect, TextBlock(line.loc, rstrip(m[2])*"\n"))
            break
        end
        push!(code, line)
    end
    # Skip trailing empty lines.
    while !isempty(code) && isblank(code[end])
        pop!(code)
    end
    if isempty(expect) && !isempty(stack)
        if occursin(r"^\s*#->\s+", stack[end].val)
            # Standalone output line.
            line = pop!(stack)
            m = match(r"^\s*#->\s+(.*)$", line.val)
            push!(expect, TextBlock(line.loc, rstrip(m[1])*"\n"))
        elseif occursin(r"^\s*#=>\s+$", stack[end].val)
            # Multiline output block.
            commentline = pop!(stack)
            while !isempty(stack) && !occursin(r"^\s*=#\s+$", stack[end].val)
                line = pop!(stack)
                push!(expect, line)
            end
            !isempty(stack) || return BrokenTestCase(commentline.loc, "incomplete multiline comment block")
            pop!(stack)
        end
    end
    !isempty(code) || return BrokenTestCase(loc, "missing test code")
    return TestCase(loc, collapse(code), collapse(pre), collapse(expect))
end

# Run a single test case.

"""
    runtest(test::TestCase; subs=common_subs(), mod=nothing) :: AbstractResult
    runtest(loc, code; pre=nothing, expect=nothing, subs=common_subs(), mod=nothing) :: AbstractResult

Run the given test case, return the result.
"""
function runtest(test::TestCase; subs=common_subs(), mod=nothing)
    # Suppress printing of the output value?
    no_output = endswith(test.code.val, ";\n") || test.expect === nothing
    # Generate a module object for running the test code.
    if mod === nothing
        modid = Base.PkgId(module_name(test.loc.file))
        if Base.root_module_exists(modid)
            mod = Base.root_module(modid)
        else
            mod = Module(Symbol(modid.name))
            @eval mod begin
                eval(x) = Core.eval($mod, x)
                include(p) = Base.include($mod, p)
            end
            Base.register_root_module(mod)
        end
    end
    # Replace the standard output/error with a pipe.
    orig_have_color = Base.have_color
    Core.eval(Base, :(have_color = false))
    orig_stdout = stdout
    orig_stderr = stderr
    pipe = Pipe()
    Base.link_pipe!(pipe; reader_supports_async=true, writer_supports_async=false)
    io = IOContext(pipe.in, :limit=>true, :module=>mod)
    redirect_stdout(pipe.in)
    redirect_stderr(pipe.in)
    pushdisplay(TextDisplay(io))
    trace = StackTraces.StackTrace()
    skipped = false
    output = ""
    @sync begin
        @async begin
            # Run the test code and print the result.
            stacktop = length(stacktrace())
            try
                filename = abspath(test.loc.file)
                tls = task_local_storage()
                has_source_path = haskey(tls, :SOURCE_PATH)
                source_path = get(tls, :SOURCE_PATH, nothing)
                tls[:SOURCE_PATH] = filename
                try
                    if test.pre !== nothing
                        pre_body = asexpr(test.pre)
                        pre = Core.eval(mod, pre_body)
                        pre::Bool
                        if !pre
                            skipped = true
                        end
                    end
                    if !skipped
                        body = asexpr(test.code)
                        ans = Core.eval(mod, body)
                        if ans !== nothing && !no_output
                            Base.invokelatest(show, io, ans)
                        end
                    end
                catch exc
                    trace = stacktrace(catch_backtrace())[1:end-stacktop]
                    print(stderr, "ERROR: ")
                    Base.invokelatest(showerror, stderr, exc, trace; backtrace=false)
                end
                println(io)
                if has_source_path
                    tls[:SOURCE_PATH] = source_path
                else
                    delete!(tls, :SOURCE_PATH)
                end
            finally
                # Restore the standard output/error.
                popdisplay()
                redirect_stderr(orig_stderr)
                redirect_stdout(orig_stdout)
                Core.eval(Base, :(have_color = $orig_have_color))
                close(pipe.in)
            end
        end
        @async begin
            # Read the output of the test.
            output = read(pipe.out, String)
            close(pipe.out)
        end
    end
    if skipped
        return Skip(test)
    end
    # Compare the actual output with the expected output and generate the result.
    expect = test.expect !== nothing ? rstrip(test.expect.val) : ""
    actual = rstrip(join(map(rstrip, eachline(IOBuffer(output))), "\n"))
    return expect == actual || occursin(expect2regex(expect, subs), actual) ?
        Pass(test, actual) :
        Fail(test, actual, trace)
end

runtest(test::BrokenTestCase; subs=common_subs(), mod=nothing) =
    Error(test)

runtest(loc, code; pre=nothing, expect=nothing, subs=common_subs(), mod=nothing) =
    runtest(TestCase(loc, TextBlock(loc, code),
                     pre !== nothing ? TextBlock(loc, pre) : nothing,
                     expect !== nothing ? TextBlock(loc, expect) : nothing),
            subs=subs, mod=mod)

# Convert expected output block to a regex.

function expect2regex(pattern, subs)
    for (regex, repl) in subs
        pattern = replace(pattern, regex => repl)
    end
    return Regex(pattern)
end

# Generate a module name from the filename.

function module_name(file)
    file = relpath(abspath(file), abspath(dirname(PROGRAM_FILE)))
    return join(titlecase.(split(file, r"[^0-9A-Za-z]+")))
end

end


text:
# a "dataset" with 10 data points of dimension 6:
x = reshape(1:60, :, 6)
n,d = size(x)

################################################################
# TESTING ILLEGAL INPUT PARAMETER COMBINATIONS
################################################################

@test_throws DomainError structureLearnPoon(x, 0, 0)
@test_throws DomainError structureLearnPoon(x, 0, 1)
@test_throws DomainError structureLearnPoon(x, 2, 2)
@test_throws DomainError structureLearnPoon(x, 3, 1)
@test_throws DomainError structureLearnPoon(x, -3, -2)

@test_throws DomainError structureLearnPoon(x, 3, 2, nsum=0)
@test_throws DomainError structureLearnPoon(x, 3, 2, nleaf=0)
@test_throws DomainError structureLearnPoon(x, 3, 2, nsum=-1)
@test_throws DomainError structureLearnPoon(x, 3, 2, nleaf=-1)
@test_throws DomainError structureLearnPoon(x, 3, 2, baseres=0)
@test_throws DomainError structureLearnPoon(x, 3, 2, baseres=-1)
@test_throws DomainError structureLearnPoon(x, 3, 2, baseres=3)

################################################################
# TEST FUNCTIONS FOR NUMBER OF NODES (FINE REGIONS ONLY)
################################################################

"""
A (m x n) rectangle has m*(m+1)*n*(n+1)/4 subrectangles.
"""
function expectedNumRegions(m::Int, n::Int)
	nsr = m*(m+1)*n*(n+1)/4
	return Int(nsr)
end

function expectedNumDecompositions(m::Int, n::Int)
	nd = 0
	for i in 1:m, j in 1:n
		nd += (m+1-i)*(n+1-j)*(i+j-2)
	end
	return nd
end

function expectedNumLeafNodes(m::Int, n::Int, nleaf::Int)
	return m*n*nleaf
end

function expectedNumSumNodes(m::Int, n::Int, nsum::Int)
	return (expectedNumRegions(m,n)-m*n-1)*nsum + 1
end

function expectedNumProdNodes(m::Int, n::Int, nsum::Int, nleaf::Int)
	# compute p1:
	p1 = 2*m*n - m - n

	# compute p2:
	f(x,y) = Int(x*(y*(y+1)/2 - 2y + 1))
	p2 = 2*(f(m,n) + f(n,m))

	# compute p3:
	p3 = 0
	for i in 2:m, j in 2:n
		p3 += (m+1-i)*(n+1-j)*(i+j-2)
	end
	for j in 3:n
		p3 += m*(n+1-j)*(j-3)
	end
	for i in 3:m
		p3 += n*(m+1-i)*(i-3)
	end

	return p1*nleaf^2 + p2*nleaf*nsum + p3*nsum^2
end

function expectedNumNodes(m::Int, n::Int, nsum::Int, nleaf::Int)
	return expectedNumLeafNodes(m,n,nleaf) + expectedNumSumNodes(m,n,nsum) + expectedNumProdNodes(m,n,nsum,nleaf)
end

################################################################
# TESTS FOR THE FUNCTIONS ABOVE
################################################################

@test expectedNumRegions(3,2) == 18
@test expectedNumDecompositions(3,2) == 18
@test expectedNumLeafNodes(3,2,4) == 24
@test expectedNumSumNodes(3,2,2) == 23
@test expectedNumProdNodes(3,2,2,4) == 172
@test expectedNumNodes(3,2,2,4) == 219

################################################################
# TESTING POON ARCHITECTURE SPN WITH BASERES=1
################################################################

h = 3
w = 2
nsum = 2
nleaf = 4
baseres = 1

spn = structureLearnPoon(x, h, w, nsum=nsum, nleaf=nleaf, baseres=1)

@test length(spn) == expectedNumNodes(h,w,nsum,nleaf)
@test numNodes(spn) == expectedNumNodes(h,w,nsum,nleaf)
@test numSumNodes(spn) == expectedNumSumNodes(h,w,nsum)
@test numProdNodes(spn) == expectedNumProdNodes(h,w,nsum,nleaf)
@test numLeafNodes(spn) == expectedNumLeafNodes(h,w,nleaf)

################################################################
# TEST FUNCTIONS FOR NUMBER OF NODES (INCL. COARSE REGIONS)
################################################################

function expectedNumRegions(m::Int, n::Int, baseres::Int)
	# coarse regions:
	cm = m ÷ baseres
	cn = n ÷ baseres
	# don't count coarse unit regions:
	numCoarse = expectedNumRegions(cm, cn) - cm*cn
	# fine regions:
	numFine = cm*cn * expectedNumRegions(baseres, baseres)
	return numCoarse + numFine
end

function expectedNumDecompositions(m::Int, n::Int, baseres::Int)
	cm = m ÷ baseres
	cn = n ÷ baseres
	numCoarse = expectedNumDecompositions(cm, cn)
	numFine = cm * cn * expectedNumDecompositions(baseres, baseres)
	return numCoarse + numFine
end

function expectedNumLeafNodes(m::Int, n::Int, nleaf::Int, baseres::Int)
	return m*n*nleaf
end

function expectedNumSumNodes(m::Int, n::Int, nsum::Int, baseres::Int)
	cm = m ÷ baseres
	cn = n ÷ baseres
	numCoarse = expectedNumSumNodes(cm, cn, nsum)
	# the (+numsum-1) part takes into account that for
	# the fine decomposition the root region is represented
	# by nsum sum nodes instead of by a single sum node.
	numFine = cm * cn * (expectedNumSumNodes(baseres, baseres, nsum) + nsum - 1)
	return numCoarse + numFine
end

function expectedNumProdNodes(m::Int, n::Int, nsum::Int, nleaf::Int, baseres::Int)
	cm = m ÷ baseres
	cn = n ÷ baseres
	numCoarse = expectedNumDecompositions(cm, cn)*nsum^2
	numFine = cm * cn * expectedNumProdNodes(baseres, baseres, nsum, nleaf)
	return numCoarse + numFine
end

function expectedNumNodes(m::Int, n::Int, nsum::Int, nleaf::Int, baseres::Int)
	return expectedNumLeafNodes(m,n,nleaf,baseres) + expectedNumSumNodes(m,n,nsum,baseres) + expectedNumProdNodes(m,n,nsum,nleaf,baseres)
end

################################################################
# TESTS FOR THE FUNCTIONS ABOVE
################################################################

h=4
w=6
b=2
s=2
l=4

@test expectedNumRegions(h, w, b) == (18-6) + 6*9
@test expectedNumDecompositions(h, w, b) == 18 + 6*6
@test expectedNumLeafNodes(h, w, l, b) == 96
@test expectedNumSumNodes(h, w, s, b) == (1 + 11*2) + 6*(5*2)
@test expectedNumProdNodes(h, w, s, l, b) == 18*s^2 + 6*(2*s^2+4*l^2)
@test expectedNumNodes(h, w, s, l, b) == 96 + 23 + 60 + 72 + 432

################################################################
# TESTING POON ARCHITECTURE SPN WITH BASERES=2
################################################################

# a "dataset" with 10 data points of dimension 24:
x = reshape(1:240, :, 24)
n,d = size(x)

h = 4
w = 6
nsum = 2
nleaf = 4
b = 2

spn = structureLearnPoon(x, h, w, nsum=nsum, nleaf=nleaf, baseres=b)

@test length(spn) == expectedNumNodes(h,w,nsum,nleaf,b)
@test numNodes(spn) == expectedNumNodes(h,w,nsum,nleaf,b)
@test numSumNodes(spn) == expectedNumSumNodes(h,w,nsum,b)
@test numProdNodes(spn) == expectedNumProdNodes(h,w,nsum,nleaf,b)
@test numLeafNodes(spn) == expectedNumLeafNodes(h,w,nleaf,b)

ext: jl


text:
function deltaJl!(
    R1 :: Matrix, R2 :: Matrix, R3 :: Matrix, R4 :: Matrix, R5 :: Matrix, R6 :: Matrix,
    Out :: Matrix,
    r1 :: Real, r2 :: Real, r3 :: Real , r4 :: Real, r5 :: Real, r6 :: Real)

    d1, d2 = size(Out)
    for j in 1:d2
      for i in 1:d1
        Out[i,j] = r1 * R1[i,j] + r2 * R2[i,j] + r3 * R3[i,j] + r4 * R4[i,j] + r5 * R5[i,j] + r6 * R6[i,j]
      end
    end
end


text:
@enum(HttpMethod,
HttpGet,
HttpPost,
HttpHead,
HttpPut,
HttpDelete)

@enum(HttpStatus,
    # 2xx: success
    HttpOk             = 200, # Most common code returned when operation was successful
    HttpCreated        = 201, # The resource has successfully been created
    HttpAccepted       = 202, # The request has been accepted, but will be processed later by the server
    HttpNoContent      = 204, # Sent when the server didn't send any data in return
    HttpResetContent   = 205, # The server informs the client that it should clear the view (form) that caused the request to be sent
    HttpPartialContent = 206, # The server has sent a part of the resource, as a response to a partial GET request

    # 3xx: redirection
    HttpMultipleChoices  = 300, # The requested page can be accessed from several locations
    HttpMovedPermanently = 301, # The requested page has permanently moved to a new location
    HttpMovedTemporarily = 302, # The requested page has temporarily moved to a new location
    HttpNotModified      = 304, # For conditional requests, means the requested page hasn't changed and doesn't need to be refreshed

    # 4xx: client error
    HttpBadRequest          = 400, # The server couldn't understand the request (syntax error)
    HttpUnauthorized        = 401, # The requested page needs an authentication to be accessed
    HttpForbidden           = 403, # The requested page cannot be accessed at all, even with authentication
    HttpNotFound            = 404, # The requested page doesn't exist
    HttpRangeNotSatisfiable = 407, # The server can't satisfy the partial GET request (with a "Range" header field)

    # 5xx: server error
    HttpInternalServerError = 500, # The server encountered an unexpected error
    HttpNotImplemented      = 501, # The server doesn't implement a requested feature
    HttpBadGateway          = 502, # The gateway server has received an error from the source server
    HttpServiceNotAvailable = 503, # The server is temporarily unavailable (overloaded, in maintenance, ...)
    HttpGatewayTimeout      = 504, # The gateway server couldn't receive a response from the source server
    HttpVersionNotSupported = 505, # The server doesn't support the requested HTTP version

    # 10xx: SFML custom codes
    HttpInvalidResponse  = 1000, # Response is not a valid HTTP one
    HttpConnectionFailed = 1001  # Connection with server failed
)

type HttpRequest
    ptr::Ptr{Void}

    function HttpRequest(ptr::Ptr{Void})
        self = new(ptr)
        finalizer(self, destroy)
        self
    end
end

function HttpRequest()
    HttpRequest(ccall((:sfHttpRequest_create, libcsfml_network), Ptr{Void}, ()))
end

function destroy(request::HttpRequest)
    ccall((:sfHttpRequest_destroy, libcsfml_network), Void, (Ptr{Void},), request.ptr)
end

function set_field(request::HttpRequest, field::AbstractString, value::AbstractString)
    ccall((:sfHttpRequest_setField, libcsfml_network), Void, (Ptr{Void}, Ptr{Cchar}, Ptr{Cchar},), request.ptr, field, value)
end

function set_method(request::HttpRequest, method::HttpMethod)
    ccall((:sfHttpRequest_setMethod, libcsfml_network), Void, (Ptr{Void}, Int32,), request.ptr, Int32(method))
end

function set_uri(request::HttpRequest, uri::AbstractString)
    ccall((:sfHttpRequest_setUri, libcsfml_network), Void, (Ptr{Void}, Ptr{Cchar},), request.ptr, uri)
end

function set_http_version(request::HttpRequest, major::Integer, minor::Integer)
    ccall((:sfHttpRequest_setHttpVersion, libcsfml_network), Void, (Ptr{Void}, Int32, Int32,), request.ptr, major, minor)
end

function set_body(request::HttpRequest, body::AbstractString)
    ccall((:sfHttpRequest_setBody, libcsfml_network), Void, (Ptr{Void}, Ptr{Cchar},), request.ptr, body)
end

type HttpResponse
    ptr::Ptr{Void}

    function HttpResponse(ptr::Ptr{Void})
        self = new(ptr)
        finalizer(self, destroy)
        self
    end
end

function destroy(response::HttpResponse)
    ccall((:sfHttpResponse_destroy, libcsfml_network), Void, (Ptr{Void},), response.ptr)
end

function get_field(response::HttpResponse, field::AbstractString)
    bytestring(ccall((:sfHttpResponse_getField, libcsfml_network), Ptr{Cchar}, (Ptr{Void}, Ptr{Cchar},), response.ptr, field))
end

function get_status(response::HttpResponse)
    HttpStatus(ccall((:sfHttpResponse_getStatus, libcsfml_network), Int32, (Ptr{Void},), response.ptr))
end

function get_major_version(response::HttpResponse)
    ccall((:sfHttpResponse_getMajorVersion, libcsfml_network), UInt32, (Ptr{Void},), response.ptr)
end

function get_minor_version(response::HttpResponse)
    ccall((:sfHttpResponse_getMinorVersion, libcsfml_network), UInt32, (Ptr{Void},), response.ptr)
end

function get_body(response::HttpResponse)
    bytestring(ccall((:sfHttpResponse_getBody, libcsfml_network), Ptr{Cchar}, (Ptr{Void},), response.ptr))
end

type Http
    ptr::Ptr{Void}

    function Http(ptr::Ptr{Void})
        self = new(ptr)
        finalizer(self, destroy)
        self
    end
end

function Http()
    Http(ccall((:sfHttp_create, libcsfml_network), Ptr{Void}, ()))
end

function destroy(http::Http)
    ccall((:sfHttp_destroy, libcsfml_network), Void, (Ptr{Void},), http.ptr)
end

function set_host(http::Http, host::AbstractString, port::Integer=0)
    ccall((:sfHttp_setHost, libcsfml_network), Void, (Ptr{Void}, Ptr{Cchar}, UInt16), http.ptr, host, port)
end

function send_request(http::Http, request::HttpRequest, timeout::Time=Time(0))
    HttpResponse(ccall((:sfHttp_sendRequest, libcsfml_network), Ptr{Void}, (Ptr{Void}, Ptr{Void}, Time), http.ptr, request.ptr, timeout))
end


text:
using FastAI

# https://github.com/FluxML/FastAI.jl/blob/master/src/models/tabularmodel.jl
function tabular_embedding_backbone(embedding_sizes, dropout_rate=0.)
    embedslist = [Flux.Embedding(ni, nf) for (nf, ni) in embedding_sizes]
    emb_drop = iszero(dropout_rate) ? identity : Dropout(dropout_rate)
    Chain(
        Parallel(vcat, embedslist),
        emb_drop
    )
    Parallel(vcat, embedslist)
end

# FROM: https://github.com/FluxML/FastAI.jl/blob/master/src/models/blocks.jl
# function was not importing, so I copied it over, for now.

function linbndrop(h_in, h_out; use_bn=true, p=0., act=identity, lin_first=false)
    bn = BatchNorm(lin_first ? h_out : h_in)
    dropout = p == 0 ? identity : Dropout(p)
    dense = Dense(h_in, h_out, act; bias=!use_bn)
    if lin_first
        return Chain(dense, bn, dropout)
    else
        return Chain(bn, dropout, dense)
    end
end


text:
for i in 1:2:10
  println(i^3)
end

arr = [x^2 for x in 1:10]


text:
# generic implementations

macro reroute(f, g)
    quote
        @inline function Cassette.overdub(ctx::HessianSparsityContext,
                                          f::typeof($(esc(f))),
                                          args...)
            println("rerouted")
            Cassette.overdub(
                ctx,
                invoke,
                $(esc(g.args[1])),
                $(esc(:(Tuple{$(g.args[2:end]...)}))),
                args...)
        end
    end
end

@reroute LinearAlgebra.BLAS.dot LinearAlgebra.dot(Any, Any)
@reroute LinearAlgebra.BLAS.axpy! LinearAlgebra.axpy!(Any,
                                                      AbstractArray,
                                                      AbstractArray)


text:
# generate examples
import Literate

EXAMPLE_BL_DIR = joinpath(@__DIR__, "..", "examples", "DGmethods_old")
GENERATED_BL_DIR = joinpath(@__DIR__, "src", "examples", "DGmethods_old", "generated")
for example in readdir(EXAMPLE_BL_DIR)
    endswith(example, ".jl") || continue
    input = abspath(joinpath(EXAMPLE_BL_DIR, example))
    script = Literate.script(input, GENERATED_BL_DIR)
    code = strip(read(script, String))
    mdpost(str) = replace(str, "@__CODE__" => code)
    Literate.markdown(input, GENERATED_BL_DIR, postprocess = mdpost)
    Literate.notebook(input, GENERATED_BL_DIR, execute = true)
end

# remove any .vtu files in the generated dir (should not be deployed)
cd(GENERATED_BL_DIR) do
    foreach(file -> endswith(file, ".vtu") && rm(file), readdir())
end

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
using RungeKuttaFehlberg
using Base.Test

config = RKFConfig((x, y) -> vecnorm(x .- y, 1), 0.0001, 1.0)

# the zero function should always return 0
for x0 in [0, 1, 2]
    @test_approx_eq rkf45_step((x, t)->0, x0, 0.0, 0.1, config)[1] 0
end

# test linear function in x
@test_approx_eq rkf45_step((x, t)->x, 0, 0.0, 0.1, config)[1] 0
@test rkf45_step((x, t)->x,  1, 0.0, 0.1, config)[1] > 0
@test rkf45_step((x, t)->x, -1, 0.0, 0.1, config)[1] < 0

# test linear function in t
@test rkf45_step((x, t)->t,  0, -1.0, 0.1, config)[1] < 0
@test rkf45_step((x, t)->t,  0,  1.0, 0.1, config)[1] > 0
@test rkf45_step((x, t)->t, -1, -1.0, 0.1, config)[1] < 0
@test rkf45_step((x, t)->t, -1,  1.0, 0.1, config)[1] > 0

# test vector input
function test_vector()
    x0 = float([1, 2, 3])
    res = rkf45_step((x, t)->reverse(x), x0, 0.0, 0.1, config)
    return typeof(res[1]) == typeof(x0)
end
@test test_vector()


text:
# Try including hello.jl

include("hello.jl");

max_stars_repo_name: neilpanchal/Julia-Algorithms

max_issues_repo_name: neilpanchal/Julia-Algorithms

max_forks_repo_name: neilpanchal/Julia-Algorithms


text:
using ConstLab
using Test

include("test_examples.jl")


text:
module PolyharmonicInterpolation
using LinearAlgebra

export PolyharmonicInterpolator,polyharmonicSpline

struct PolyharmonicInterpolator
	knots::AbstractArray
	values::Vector{Float64}
	nPts::Int
	dim::Int
	order::Int
	w::Vector{Float64}
	v::Vector{Float64}
end

PolyharmonicInterpolator(knots::AbstractArray{T},values::Vector{T},
order::Int=2) where T <: Real = begin

	# Check to make sure knots is N x d and values is N x 1
	if size(knots)[1] != length(values)
		error("knots and values length must agree")
	end

	if length(size(knots)) == 1
		nPts = length(knots)
		dim = 1
	else
		nPts,dim = size(knots)
	end

	# initialize w and v
	w = Vector{Float64}(undef,nPts)
	v = Vector{Float64}(undef,dim+1)

	# Fill npts and dim using the dimension of knots
	computeWeights(PolyharmonicInterpolator(knots,values,nPts,dim,order,w,v))

end

function polyharmonicRBF(r::Float64,k::Int)

	if r >= 1
		val = (k % 2 == 0) ? r^k*log(r) : r^k
	else
		val = (k % 2 == 0) ? r^(k-1)*log(r^r) : r^k
	end

	return val

end

function createBlockMatrix(p::PolyharmonicInterpolator)

	# Create the matrix of differences between the knots
	dist = [norm(p.knots[i,:] - p.knots[j,:]) for i in 1:p.nPts, j in 1:p.nPts]

	# Create the matrix A
	A = map(r->polyharmonicRBF(r,p.order),dist)

	# Create the matrix B
	B = [ones(p.nPts) p.knots];

	return A,B

end

function computeWeights(p::PolyharmonicInterpolator)

	# Create block matrices used to solve for weights
	A,B = createBlockMatrix(p)

	# Make block matrix for the symmetric linear system
	blockMat = [A B; transpose(B) zeros(p.dim+1,p.dim+1)]

	# Get weights nPts and number of dimensions
	weights = blockMat \ [p.values; zeros(p.dim+1,1)]

	p.w .= weights[1:p.nPts]
	p.v .= weights[p.nPts+1:end]

	return p

end

function polyharmonicSpline(p::PolyharmonicInterpolator,interpPoints::AbstractArray)

	# Check to see the dimensions of the interpPoints array agree
	let
		nInterpDims = nothing
		try
			_, nInterpDims = size(interpPoints)
		catch
			nInterpDims = 1
		end

		if p.dim != nInterpDims
			error("dimension of knots and points to be interpolated must agree")
		end
	end

	# Initialize array
	interpValues = zeros(size(interpPoints)[1])

	for (i,point) in enumerate(eachrow(interpPoints))

		tmppoint =[1.0;point]

		# Get distance to all scattered points
		dist = sqrt.(sum((p.knots .- point').^2,dims=2))

		# reduce to normal size (Julia returns a 2 dimensional row array for some reason...)
		dist = reduce(vcat,dist)

		# Get distance between basis and point of interest and compute the
		# polyharmonic basis
		phi = map(r->polyharmonicRBF(r,p.order),dist)

		interpValues[i] = p.w'phi + p.v'tmppoint

	end

	return interpValues

end

end

max_stars_repo_name: MathOnco/ploidyEvolution

max_issues_repo_name: MathOnco/ploidyEvolution

max_forks_repo_name: MathOnco/ploidyEvolution


text:
module GraphOrdering

struct Result
    perm :: Vector{Int}
    invperm :: Vector{Int}
    degrees :: Vector{Int}
    edge :: Vector{Int}
    dist :: Vector{Int}
end

"""
    bandwidth(G)

Calculate the bandwidth of the graph G.
"""

function bandwidth(G)
    bw = -1
    for (v, adj) in G
        for w in adj
            bw = max(bw, abs(v-w))
        end
    end
    return 2*bw + 1
end

"""
    symrcm(G, v)

Sparse reverse Cuthill-McKee ordering. `G` is graph presented as adjacency list
and `v` is starting vertex. Returns `Result` type, which contains bandwidth
minimizing permutation, inverse of permutation and some other results.
"""
function symrcm(G, v)
    n = length(G)
    permutation = zeros(Int, n)
    permutation[1] = v
    visited = zeros(Bool, n)
    visited[v] = true
    degrees = [length(G[i]) for i in 1:n]
    order = Base.Order.By(j -> degrees[j])
    connected = zeros(Int, n)
    edge = zeros(Int, n)
    dist = zeros(Int, n)
    nwrk = 0
    wrk = zeros(Int, nwrk)
    idx = 2

    for i=1:n

        v = permutation[i]
        adj = G[v]
        nadj = length(adj)

        aux = adj
        # aux is adj or wrk, depending is adj in order. If adj is not sorted,
        # copy degrees from adj to aux and make in-place sort
        if (!issorted(aux, order))
            aux = wrk
            resize!(aux, max(nadj, length(aux)))
            copyto!(aux, adj)
            sort!(aux, 1, nadj, InsertionSort, order)
        end

        for wi = 1:nadj
            w = aux[wi]
            visited[w] && continue
            visited[w] = true
            permutation[idx] = w
            edge[w] = v
            dist[w] = dist[v] + 1
            idx += 1
        end

    end

    perm = reverse(permutation)
    return Result(perm, invperm(perm), degrees, edge, dist)
end

"""
    reorder(G, result)

Given result of `symrcm`, reorder nodes of graph `G`. Returns
new graph with nodes ordered corresponding the new permuation.
"""
function reorder(G, res)
    H = empty(G)
    for (v, adj) in G
        H[res.invperm[v]] = [res.invperm[w] for w in adj]
    end
    return H
end

end


text:
function normalization(x)
    return (x .- minimum(x))./(maximum(x).-minimum(x))
end


text:
# code inside graph_functions.jl
function sum_values(x,y)
  return x+y
end

max_stars_repo_name: jpbrasile/VegaGraphs

max_issues_repo_name: jpbrasile/VegaGraphs

max_forks_repo_name: jpbrasile/VegaGraphs


text:
function maybe_default_show(io, x)
    if !get(io, :limit, false)
        invoke(show, Tuple{IO, Any}, io, x)
        return true
    end
    return false
end

function shortsummary(x; context=devnull)
    context = IOContext(
        context,
        :compact => true,
        :limit => true,
    )
    return first(sort(
        [
            sprint(summary, x; context=context),
            sprint(print, x; context=context),
        ],
        by = length,
    ))
end

macro nograd(funs...)
    exprs = map(funs) do f
        f = esc(f)
        args = esc(gensym("args"))
        return macroexpand(
            __module__,
            :(@adjoint $f($args...) = $f($args...), _ -> nothing),
        )
        # Without `macroexpand`, I get `syntax:
        # "SteadyStateFit.__context__" is not a valid function
        # argument name`.
    end
    return Expr(:block, __source__, exprs...)
end

# https://github.com/FluxML/Zygote.jl/pull/351
@nograd Base.gc_num Base.time_ns

nograd(f) = f()
@nograd nograd


text:
using AES, Test

a=AES.AESBlock(rand(UInt8, 100), 3)
ac=copy(a)
@test ac == a
@test ac !== a
@test ac isa AES.AESBlock


text:
using PyCall

struc = get_structure("2HHB.pdb")

# Extremely rudimentary test, basically do we get anything back from our PyCall?
@test isa(struc,PyObject)

chains = get_chains(struc)

# Are the chains there? We expect two alpha and two beta chains
@test length(chains) == 4

# Are there enough c_alpha atoms? With enough coordinates?
@test size(chains[1]) == (141,3)
@test size(chains[2]) == (146,3)
@test size(chains[3]) == (141,3)
@test size(chains[4]) == (146,3)

# PDB export test
P = [31.132 16.439 58.160;
	 6.870 17.784 4.702
]

f = open("pdb_export_test.pdb")
pdb_expected = readlines(f)

pdb_handler(r::Test.Success) = rm("pdb_exported_test.pdb")

Test.with_handler(pdb_handler) do
	@test pdb_expected == export_to_pdb("VAL", "A", P, "pdb_exported_test.pdb")
end

# remote pdb test

# flush cache
if isreadable(joinpath(dirname(@__FILE__),"..", ".cache"))
	rm(joinpath(dirname(@__FILE__),"..", ".cache"), recursive=true)
end
remote_pdb_handler(r::Test.Success) = rm(joinpath(dirname(@__FILE__),"..", ".cache", "2HHB"))
f = open("2HHB.pdb", "r")
reference = readall(f)

# hit downloading first
@test reference == readall(open(get_remote_pdb("2HHB"), "r"))

# then check caching code path
Test.with_handler(remote_pdb_handler) do
	@test reference == readall(open(get_remote_pdb("2HHB"), "r"))
end


text:
struct Rename{oldnames, newnames}
end
@inline Rename(oldnames::Tuple{Vararg{Symbol}}, newnames::Tuple{Vararg{Symbol}}) = Rename{oldnames, newnames}()
@inline Rename(oldname::Symbol, newname::Symbol) = Rename{(oldname,), (newname,)}()

@inline rename(x, oldname::Symbol, newname::Symbol) = Rename(oldname, newname)(x)
@inline rename(x, oldnames::Tuple{Vararg{Symbol}}, newnames::Tuple{Vararg{Symbol}}) = Rename(oldnames, newnames)(x)

@inline function (::Rename{oldnames, newnames})(nt::NamedTuple{names}) where {oldnames, newnames, names}
    return NamedTuple{_rename(Val(oldnames), Val(newnames), Val(names))}(Tuple(nt))
end

# Tried @pure here but didn't seem to do what I expected...
@generated function _rename(::Val{oldnames}, ::Val{newnames}, ::Val{names}) where {oldnames, newnames, names}
    if length(unique(oldnames::Tuple{Vararg{Symbol}})) != length(unique(oldnames))
        error("The old column names $oldnames are not unique")
    end

    if length(unique(newnames::Tuple{Vararg{Symbol}})) != length(unique(newnames))
        error("The new column names $newnames are not unique")
    end

    if length(oldnames) != length(newnames)
        error("The number of new column names does not match the number of old column names")
    end

    out = map(names::Tuple{Vararg{Symbol}}) do n
        i = findfirst(==(n), oldnames)
        return i === nothing ? n : newnames[i]
    end

    if length(unique(out)) != length(unique(names))
        error("The output column names $out are not unique")
    end

    return out
end


text:
##########################################################################################
#
#    MCMC specific derivation rules
#
##########################################################################################

####### creates multiple rules at once for logpdf(Distrib, x)
# FIXME : using undocumented dprefix function of ReverseDiffSource (should be replaced)
macro dlogpdfd(dist::Symbol, rule)
	sig = :( logpdf($(Expr(:(::), :d, dist)), x::Real) )
	ReverseDiffSource.deriv_rule( sig, :d, rule )

	sig = :( logpdf($(Expr(:(::), :d, dist)), x::AbstractArray) )
	rule2 = ReverseDiffSource.substSymbols(rule, {:x => :(x[i]), :ds => :(ds[i])})
	ReverseDiffSource.deriv_rule( sig, :d, :(for i in 1:length(x) ; $rule2 ; end))

	sig = :( logpdf($(Expr(:(::), :d, Expr(:curly, :Array, dist))), x::AbstractArray) )
	rule2 = ReverseDiffSource.substSymbols(rule, {:dd1 => :(dd1[i]), :dd2 => :(dd2[i]), :dd3 => :(dd3[i]),
		:x => :(x[i]), :ds => :(ds[i]), :d => :(d[i]) })
	ReverseDiffSource.deriv_rule(sig, :d, :(for i in 1:length(x) ; $rule2 ; end))
end

macro dlogpdfx(dist::Symbol, rule)
	sig = :( logpdf($(Expr(:(::), :d, dist)), x::Real) )
	ReverseDiffSource.deriv_rule( sig, :x, rule )

	sig = :( logpdf($(Expr(:(::), :d, dist)), x::AbstractArray) )
	rule2 = ReverseDiffSource.substSymbols(rule, {:dx => :(dx[i]), :x => :(x[i]), :ds => :(ds[i])})
	ReverseDiffSource.deriv_rule( sig, :x, :(for i in 1:length(x) ; $rule2 ; end))

	sig = :( logpdf($(Expr(:(::), :d, Expr(:curly, :Array, dist))), x::AbstractArray) )
	rule3 = ReverseDiffSource.substSymbols(rule2, {:d => :(d[i])})
	ReverseDiffSource.deriv_rule( sig, :x, :(for i in 1:length(x) ; $rule3 ; end))
end

####### derivation for Distribution types constructors
ReverseDiffSource.declareType(Distribution, :Distribution)

for d in [:Bernoulli, :TDist, :Exponential, :Poisson]
	ReverseDiffSource.declareType(eval(d), d)

	ReverseDiffSource.deriv_rule(:( ($d)(p::Real) ), :p, :( dp = ds1 ))
	ReverseDiffSource.deriv_rule(:( ($d)(p::AbstractArray) ), :p, :( copy!(dp, ds1) ))
end

for d in [ :Normal, :Uniform, :Weibull, :Gamma, :Cauchy, :LogNormal, :Binomial, :Beta, :Laplace]
	ReverseDiffSource.declareType(eval(d), d)

	ReverseDiffSource.deriv_rule(:( ($d)(p1::Real, p2::Real) ),   :p1, :( dp1 = ds1 ) )
	ReverseDiffSource.deriv_rule(:( ($d)(p1::Real, p2::Real) ),   :p2, :( dp2 = ds2 ) )
	ReverseDiffSource.deriv_rule(:( ($d)(p1::AbstractArray, p2::AbstractArray) ), :p1, :( copy!(dp1, ds1) ) )
	ReverseDiffSource.deriv_rule(:( ($d)(p1::AbstractArray, p2::AbstractArray) ), :p2, :( copy!(dp2, ds2) ) )
end

#######   Normal distribution
@dlogpdfx Normal dx += (d.μ - x) / (d.σ * d.σ) * ds
@dlogpdfd Normal ( 	dd1 += (x - d.μ) / (d.σ*d.σ) * ds;
					dd2 += ((x - d.μ)*(x - d.μ) / (d.σ*d.σ) - 1.) / d.σ * ds )

## Uniform distribution
@dlogpdfx Uniform dx += 0.
@dlogpdfd Uniform ( dd1 += (d.a <= x <= d.b) / (d.b - d.a) * ds ;
					dd2 += (d.a <= x <= d.b) / (d.a - d.b) * ds )

## Weibull distribution
@dlogpdfd Weibull   ( 	dd1 += ((1. - (x/d.scale)^d.shape) * log(x/d.scale) + 1./d.shape) * ds ;
						dd2 += ((x/d.scale)^d.shape - 1.) * d.shape/d.scale * ds )
@dlogpdfx Weibull   dx += ((1. - (x/d.scale)^d.shape) * d.shape - 1.) / x * ds

## Beta distribution
@dlogpdfd Beta   ( dd1 += (digamma(d.alpha+d.beta) - digamma(d.alpha) + log(x)) * ds ;
				    dd2 += (digamma(d.alpha+d.beta) - digamma(d.beta) + log(1-x)) * ds )
@dlogpdfx Beta   dx += ((d.alpha-1) / x - (d.beta-1)/(1-x)) * ds

## TDist distribution
@dlogpdfd TDist   dd1 += ((x*x-1)/(x*x + d.df)+log(d.df/(x*x+d.df))+digamma((d.df+1)/2)-digamma(d.df/2))/2 * ds
@dlogpdfx TDist   dx += (-(d.df+1)*x / (d.df+x*x)) * ds

## Exponential distribution
@dlogpdfd Exponential   dd1 += (x-d.scale) / (d.scale*d.scale) * ds
@dlogpdfx Exponential   dx -= ds / d.scale

## Gamma distribution
@dlogpdfd Gamma   ( dd1 += (log(x) - log(d.scale) - digamma(d.shape)) * ds ;
					dd2 += ((x - d.scale*d.shape) / (d.scale*d.scale)) * ds )
@dlogpdfx Gamma   dx += (-( d.scale + x - d.shape*d.scale)/(d.scale*x)) * ds

## Cauchy distribution
@dlogpdfd Cauchy   ( dd1 += (2(x-d.location) / (d.scale*d.scale + (x-d.location)*(x-d.location))) * ds ;
					 dd2 += (((x-d.location)*(x-d.location) - d.scale*d.scale) / (d.scale*(d.scale*d.scale + (x-d.location)*(x-d.location)))) * ds )
@dlogpdfx Cauchy   dx += (2(d.location-x) / (d.scale*d.scale + (x-d.location)*(x-d.location))) * ds

## Log-normal distribution
@dlogpdfd LogNormal   ( dd1 += (log(x) - d.meanlog) / (d.sdlog*d.sdlog) * ds ;
					 	dd2 += (d.meanlog*d.meanlog - d.sdlog*d.sdlog - log(x)*(2d.meanlog-log(x))) / (d.sdlog*d.sdlog*d.sdlog) * ds )
@dlogpdfx LogNormal   dx += (d.meanlog - d.sdlog*d.sdlog - log(x)) / (d.sdlog*d.sdlog*x) * ds

## Laplace distribution
@dlogpdfd Laplace   ( dd1 += (x > d.location ? 1 : -1) / d.scale * ds ;  # location
	                  dd2 += ( abs(x-d.location)/d.scale - 2. ) /d.scale ) # scale
@dlogpdfx Laplace   dx += (x > d.location ? -1 : 1) / d.scale * ds

# Note : vectorization will not be easily possible for multi variate distribs (Dirichlet, Categorical)
# TODO : add other continuous distribs ? : Pareto, Rayleigh, Logistic, Levy, Dirichlet, FDist
# TODO : add other discrete distribs ? : NegativeBinomial, DiscreteUniform, HyperGeometric, Geometric, Categorical

## Bernoulli distribution (Note : no derivation on x parameter as it is an integer)
@dlogpdfd Bernoulli     dd1 += 1. / (d.p1 - 1. + x) * ds

## Binomial distribution (Note : no derivation on x and n parameters as they are integers)
@dlogpdfd Binomial      dd2 += (x / d.prob - (d.size-x) / (1 - d.prob)) * ds

## Poisson distribution (Note : no derivation on x parameter as it is an integer)
@dlogpdfd Poisson       dd1 += (x / d.lambda - 1) * ds


text: abstract type AbstractSubject <: AbstractObject end;

max_stars_repo_name: daviddoret/julia-cybersecurity

max_issues_repo_name: daviddoret/julia-cybersecurity

max_forks_repo_name: daviddoret/julia-cybersecurity


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -3.611664788975792 <= q <= 3.6444325166036777)

                     add_NL_constraint(m, :(1/(1 + exp(-(-0.4383440965768486 + -0.2075365332338559*$(x[1]) + 0.46530061880514095*$(x[2])))) + 1/(1 + exp(-(0.22314644246023452 + -0.6314968564820513*$(x[1]) + 0.5512122876813113*$(x[2])))) + 1/(1 + exp(-(0.30053153463966176 + -0.15357242883325783*$(x[1]) + 0.254162988138138*$(x[2])))) + 1/(1 + exp(-(-0.06895001670910483 + 0.7534492827978192*$(x[1]) + -0.6113176568181604*$(x[2])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
# Comparison

function (==)(a::MultiValue,b::MultiValue)
  a.array == b.array
end

function (≈)(a::MultiValue,b::MultiValue)
  a.array ≈ b.array
end

function (≈)(
  a::AbstractArray{<:MultiValue}, b::AbstractArray{<:MultiValue})
  if size(a) != size(b); return false; end
  for (ai,bi) in zip(a,b)
    if !(ai≈bi); return false; end
  end
  true
end

# Addition / subtraction

for op in (:+,:-)
  @eval begin

    function ($op)(a::MultiValue{S}) where S
      r = $op(a.array)
      MultiValue(r)
    end

    function ($op)(a::MultiValue{S},b::MultiValue{S}) where S
      r = $op(a.array, b.array)
      MultiValue(r)
    end

  end
end

# Matrix Division

function (\)(a::TensorValue, b::MultiValue)
  r = a.array \ b.array
  MultiValue(r)
end

# Operations with other numbers

for op in (:+,:-,:*)
  @eval begin
    ($op)(a::MultiValue,b::Number) = MultiValue($op(a.array,b))
    ($op)(a::Number,b::MultiValue) = MultiValue($op(a,b.array))
  end
end

(/)(a::MultiValue,b::Number) = MultiValue(a.array/b)

# Dot product (simple contraction)

(*)(a::VectorValue{D}, b::VectorValue{D}) where D = inner(a,b)

function (*)(a::MultiValue,b::MultiValue)
  r = a.array * b.array
  MultiValue(r)
end

@inline dot(u::VectorValue,v::VectorValue) = inner(u,v)

# Inner product (full contraction)

inner(a::Real,b::Real) = a*b

@generated function inner(a::MultiValue{S,T,N,L}, b::MultiValue{S,W,N,L}) where {S,T,N,L,W}
  str = join([" a.array.data[$i]*b.array.data[$i] +" for i in 1:L ])
  Meta.parse(str[1:(end-1)])
end

# Reductions

for op in (:sum,:maximum,:minimum)
  @eval begin
    $op(a::MultiValue) = $op(a.array)
  end
end

# Outer product (aka dyadic product)

outer(a::Real,b::Real) = a*b

outer(a::MultiValue,b::Real) = a*b

outer(a::Real,b::MultiValue) = a*b

@generated function outer(a::VectorValue{D},b::VectorValue{Z}) where {D,Z}
  str = join(["a.array[$i]*b.array[$j], " for j in 1:Z for i in 1:D])
  Meta.parse("MultiValue(SMatrix{$D,$Z}($str))")
end

# Linear Algebra

det(a::TensorValue) = det(a.array)

inv(a::TensorValue) = MultiValue(inv(a.array))

# Measure

meas(a::VectorValue) = sqrt(inner(a,a))

meas(a::TensorValue) = abs(det(a))

@inline norm(u::VectorValue) = sqrt(inner(u,u))

# conj

conj(a::MultiValue) = MultiValue(conj(a.array))

# Trace

@generated function trace(v::TensorValue{D}) where D
  str = join([" v.array.data[$i+$((i-1)*D)] +" for i in 1:D ])
  Meta.parse(str[1:(end-1)])
end

@inline tr(v::TensorValue) = trace(v)

# Adjoint

function adjoint(v::TensorValue)
  t = adjoint(v.array)
  TensorValue(t)
end

# Symmetric part

@generated function symmetic_part(v::TensorValue{D}) where D
  str = "("
  for j in 1:D
    for i in 1:D
      str *= "0.5*v.array.data[$i+$((j-1)*D)] + 0.5*v.array.data[$j+$((i-1)*D)], "
    end
  end
  str *= ")"
  Meta.parse("TensorValue($str)")
end


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

mutable struct ApplicationGatewayBackendAddressPool2 <: SwaggerModel
    id::Any # spec type: Union{ Nothing, String } # spec name: id

    function ApplicationGatewayBackendAddressPool2(;id=nothing)
        o = new()
        validate_property(ApplicationGatewayBackendAddressPool2, Symbol("id"), id)
        setfield!(o, Symbol("id"), id)
        o
    end
end # type ApplicationGatewayBackendAddressPool2

const _property_map_ApplicationGatewayBackendAddressPool2 = Dict{Symbol,Symbol}(Symbol("id")=>Symbol("id"))
const _property_types_ApplicationGatewayBackendAddressPool2 = Dict{Symbol,String}(Symbol("id")=>"String")
Base.propertynames(::Type{ ApplicationGatewayBackendAddressPool2 }) = collect(keys(_property_map_ApplicationGatewayBackendAddressPool2))
Swagger.property_type(::Type{ ApplicationGatewayBackendAddressPool2 }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ApplicationGatewayBackendAddressPool2[name]))}
Swagger.field_name(::Type{ ApplicationGatewayBackendAddressPool2 }, property_name::Symbol) =  _property_map_ApplicationGatewayBackendAddressPool2[property_name]

function check_required(o::ApplicationGatewayBackendAddressPool2)
    true
end

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


text:
function write_scopf_summary(scenario_id, network, objective; branch_flow_cuts=0, objective_lb=-Inf, load_time=-1.0, solve_time=-1.0, filter_time=-1.0, total_time=load_time+solve_time+filter_time)
    println("")

    data = [
        "----",
        "scenario id",
        "num bus",
        "num branch",
        "num gen",
        "num load",
        "num shunt",
        "num gen cont",
        "num branch cont",
        "active gen cont",
        "active branch cont",
        "branch flow cuts",
        "objective ub",
        "objective lb",
        "load time (sec.)",
        "solve time (sec.)",
        "filter time (sec.)",
        "total time (sec.)",
    ]
    println(join(data, ", "))

    data = [
        "DATA",
        scenario_id,
        length(network["bus"]),
        length(network["branch"]),
        length(network["gen"]),
        length(network["load"]),
        length(network["shunt"]),
        length(network["gen_contingencies"]),
        length(network["branch_contingencies"]),
        length(network["gen_contingencies_active"]),
        length(network["branch_contingencies_active"]),
        branch_flow_cuts,
        objective,
        objective_lb,
        load_time,
        solve_time,
        filter_time,
        total_time,
    ]
    println(join(data, ", "))

end

function write_power_balance_summary(scenario_id, p_delta_abs_max, q_delta_abs_max, p_delta_abs_mean, q_delta_abs_mean)
    println("")

    data = [
        "----",
        "scenario id",
        "p_delta abs max",
        "q_delta abs max",
        "p_delta abs mean",
        "q_delta abs mean",
    ]
    println(join(data, ", "))

    data = [
        "DATA_PB",
        scenario_id,
        p_delta_abs_max,
        q_delta_abs_max,
        p_delta_abs_mean,
        q_delta_abs_mean,
    ]
    println(join(data, ", "))

end


text:
# TODO: seed, copy

"""
A box in R^n, i.e., each coordinate is bounded.

Two kinds of valid input:
    Box(-1.0, 1.0, (3,4)) # low and high are scalars, and shape is provided
    Box([-1.0,-2.0], [2.0,4.0]) # low and high are arrays of the same shape
"""
mutable struct Box <: AbstractSpace
    low::Array
    high::Array
    shape::Tuple
end

function Box(low::Number, high::Number, shape::Union{Tuple, Array{Int64, 1}}, dtype::Union{DataType, Nothing}=nothing)
    if isnothing(dtype)
        dtype = high == 255 ? UInt8 : Float32
        @warn "dtype was autodetected as $(dtype). Please provide explicit data type."
    end

    if low > high
        @warn "low  > high. Swapping values to preserve sanity"
        (low, high) = (high, low)  # Preserves sanity if low > high
    end

    if dtype <: Integer
        if !isa(low, Integer) || !isa(high, Integer)
            @warn "dtype is an Integer, but the values are floating points. Using ceiling of lower bound and floor of upper bound"
        end
        low = ceil(dtype, low)
        high = floor(dtype, high)
    end

    Low = dtype(low) .+ zeros(dtype, shape)
    High = dtype(high) .+ zeros(dtype, shape)
    return Box(Low, High, shape)
end

function Box(low::Array, high::Array, dtype::Union{DataType, Nothing}=nothing)
    @assert size(low) == size(high) "Dimension mismatch between low and high arrays."
    shape = size(low)
    @assert all(low .< high) "elements of low must be lesser than their respective counterparts in high"

    if isnothing(dtype)
        dtype = all(high .== 255) ? UInt8 : Float32
        @warn "dtype was autodetected as $(dtype). Please provide explicit data type."
    end
    if dtype <: Integer
        if !all(isa.(low, Integer)) || !all(isa(high, Integer))
            @warn "dtype is an Integer, but the values are floating points. Using ceiling of lower bound and floor of upper bound"
        end
        low = ceil.(dtype, low)
        high = floor.(dtype, high)
    else
        low = dtype.(low)
        high = dtype.(high)
    end
    return Box(low, high, shape)
end
#=
function seed!(box_obj::Box, seed::Int)
    box_obj.seed = seed
end
=#

Base.:(==)(box_obj::Box, other::Box) = checkvalidtypes(box_obj, other) && isapprox(box_obj.low, other.low) && isapprox(box_obj.high, other.high)

function sample(box_obj::Box)
    dtype = eltype(box_obj.low)
    dtype <: AbstractFloat ?
        rand(dtype, size(box_obj)) .* (box_obj.high .- box_obj.low) .+ box_obj.low :
        rand.(UnitRange.(box_obj.low, box_obj.high))
end

function contains(x::Union{Real, AbstractArray, NTuple}, box_obj::Box)
    isa(x, Number) && size(box_obj.low) == (1,) && (x = [x])
    size(x) == size(box_obj) && all(box_obj.low .<= x .<= box_obj.high)
end

function checkvalidtypes(box_obj1::Box, box_obj2::Box)
    dtype1, dtype2 = eltype(box_obj1.low), eltype(box_obj2.low)
    dtype1 == dtype2 ||                            # If the dtypes of both boxes are not the same...
            (dtype1 <: Unsigned && dtype2 <: Unsigned) || (dtype1 <: Signed && dtype2 <: Signed)  # then check if they're both signed or both unsigned.
end


text:
using Pkg

pkg"add https://github.com/faf0/AES.jl"
#pkg"rm AES"
pkg"add https://github.com/oxinabox/MD5.jl"
pkg"add JSON"


text:
module TestStep
using Debug, Debug.Flow
import Debug.BPNode, Debug.Scope

state = DBState()
ip    = 1

function trap(node, scope::Scope)
    global firstline, ip
    if isa(node, BPNode);  firstline = node.loc.line;  end
    if !Flow.pretrap(state, node, scope); return; end

    line = node.loc.line - firstline + 1
    l, f = instructions[ip]
    ip += 1
    @assert line == l

    if f === stepout!; stepout!(state, node, scope)
    else               f(state)
    end
end

macro test_step(ex)
    Debug.code_debug(Flow.instrument(trap, ex))
end

const instructions = Any[
    (1, stepover!),
    (2, singlestep!), (3, singlestep!), (4, singlestep!),
    (6, stepover!), (9, stepover!),
    (10, singlestep!), (11, stepout!),
    (14, stepover!),
    (17, singlestep!), (18, stepout!),
    (20, stepover!),
    (23, singlestep!), (24, stepout!),
    (26, stepover!),
    (29, singlestep!), (30, stepout!),
    (32, stepover!),
    (35, singlestep!), (36, stepout!),
    (38, singlestep!),
    (42, stepover!),
    (43, singlestep!), (39, stepout!),
    (44, continue!),
]

@test_step let
    @bp         #  1
    let         #  2
        x = 1   #  3
        y = 2   #  4
    end
    let         #  6
        x = 1   #  7
    end
    x = 5       #  9
    let         # 10
        x = 1   # 11
        y = 2   # 12
    end
    for x=1:3   # 14
        z=x^2   # 15
    end
    for x=1:3   # 17
        z=x^2   # 18
    end
    while x>0   # 20
        x-=1    # 21
    end
    while x<5   # 23
        x+=1    # 24
    end
    try         # 26
        x       # 27
    end
    try         # 29
        x       # 30
    end
    [begin      # 32
        x^2     # 33
    end for x=1:3]
    [begin      # 35
        x^2     # 36
    end for x=1:3]
    function f(x) # 38
         x -= 1 # 39
         2x     # 40
    end
    f(3)        # 42
    f(4)        # 43
    y = 7       # 44
    y = 8       # 45
end

@assert ip == length(instructions)+1

end  # module


text:
using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots
using DifferentialEquations
using DiffEqSensitivity
using Zygote
using ForwardDiff
using LinearAlgebra
using Random
using Statistics
using ProgressBars, Printf
using Flux.Optimise: update!
using Flux.Losses: mae, mse
using BSON: @save, @load

is_restart = false;
n_epoch = 1000000;
n_plot = 10;

opt = ADAMW(0.001, (0.9, 0.999), 1.f-6);
datasize = 40;
n_exp_train = 20;
n_exp_val = 10;
n_exp = n_exp_train + n_exp_val;
noise = 1.f-4;
ns = 3;
nr = 3;

grad_max = 10 ^ (4);
maxiters = 10000;

alg = AutoTsit5(Rosenbrock23(autodiff=false));
# alg = Rosenbrock23(autodiff=false);
atol = 1e-5;
rtol = 1e-3;
lb = 1e-5;
ub = 1.f1;

# Generate datasets
u0_list = rand(Float32, (n_exp, ns)) .+ 0.5;
u0_list[:, 2:2] .= 0 .+ lb;

tsteps = 10 .^ range(-2, 5, length=datasize);
tspan = Float32[0, tsteps[end]];
t_end = tsteps[end]

k = [4.f-2, 3.f7, 1.f4];
ode_data_list = zeros(Float32, (n_exp, ns, datasize));
yscale_list = [];

function rober!(dydt, y, k, t)
    r1 = k[1] * y[1]
    r2 = k[2] * y[2] * y[2]
    r3 = k[3] * y[2] * y[3]
    dydt[1] = -r1 + r3
    dydt[2] = r1 - r2 - r3
    dydt[3] = r2
end

u0 = u0_list[1, :];
prob_trueode = ODEProblem(rober!, u0, tspan, k);

function max_min(ode_data)
    return maximum(ode_data, dims=2) .- minimum(ode_data, dims=2)
end

for i = 1:n_exp
    u0 = u0_list[i, :]
    prob_trueode = ODEProblem(rober!, u0, tspan, k)
    ode_data = Array(solve(prob_trueode, alg, saveat=tsteps, atol=atol, rtol=rtol))
    ode_data += randn(size(ode_data)) .* ode_data .* noise
    ode_data_list[i, :, :] = ode_data
    push!(yscale_list, max_min(ode_data))
end

yscale = maximum(hcat(yscale_list...), dims=2);
dydt_scale = yscale[:, 1] ./ t_end
show(stdout, "text/plain", round.(yscale', digits=8))

np = nr * (2 * ns + 1) + 1;
pcrnn = (rand(Float32, np) .- 0.5) * 2 * sqrt(6 / (ns + nr));
pcrnn[end] = 1.e-1;

function p2vec(p)
    slope = abs(p[np]) * 10
    w_b = @view(p[1:nr]) .* slope

    w_in = reshape(@view(p[nr * (ns + 1) + 1:nr * (2 * ns + 1)]), ns, nr)

    w_out = reshape(@view(p[nr + 1:nr * (ns + 1)]), ns, nr)
    w_out = @. -w_in * abs(w_out)

    w_in = clamp.(w_in, 0, 2.5)

    return w_in, w_b, w_out
end

function display_p(p)
    w_in, w_b, w_out = p2vec(p)
    println("species (column) reaction (row)")
    println("w_in | w_b | w_out")
    display(hcat(w_in', w_b, w_out'))

    println("w_out_scale")
    w_out_ = (w_out)' .* exp.(w_b)
    # display(w_out_)
    # display(maximum(abs.(w_out_), dims=2)')
    display(w_out_ ./ maximum(abs.(w_out_), dims=2))
    println("slope = $(p[end])")
end
display_p(pcrnn)

node = 4;
dudt2 = Chain(x -> x,
              Dense(2, node, gelu),
              Dense(node, node, gelu),
              Dense(node, node, gelu),
            #   Dense(node, node, gelu),
              Dense(node, 1, exp));
pnn, re = Flux.destructure(dudt2);
rep = re(pnn);
p = vcat(pcrnn, pnn);

function crnn(du, u, p, t)
    u_ = vcat(@views(u[1]), rep(@views(u[[1, 3]])), @views(u[3]))
    w_in_x = w_in' * @. log(clamp(u_, lb, ub))
    du .= w_out * (@. exp(w_in_x + w_b))
end

u0 = u0_list[1, :]
prob = ODEProblem(crnn, u0, tspan, saveat=tsteps, atol=atol, rtol=rtol)

sense = BacksolveAdjoint(checkpointing=true; autojacvec=ZygoteVJP());
function predict_neuralode(u0, p)
    pcrnn = p[1:np]
    pnn = p[np + 1:end]
    global w_in, w_b, w_out = p2vec(pcrnn)
    global rep = re(pnn)
    sol = solve(prob, alg, u0=u0, p=pcrnn, sensalg=sense, verbose=false, maxiters=maxiters)
    pred = Array(sol)
    pred[2:2, :] .= rep(pred[[1, 3], :])

    if sol.retcode == :Success
        nothing
    else
        println("ode solver failed")
    end
    return pred
end
pred = predict_neuralode(u0, p);
# using BenchmarkTools
# @benchmark predict_neuralode(u0, p)

function loss_neuralode(p, i_exp)
    pred = predict_neuralode(u0_list[i_exp, :], p)[[1, 3], :]
    ode_data = ode_data_list[i_exp, [1, 3], 1:size(pred)[2]]
    loss = mae(ode_data, pred)
    return loss
end
loss_neuralode(p, 1)
# @benchmark ForwardDiff.gradient(x -> loss_neuralode(x, 1), p)

cbi = function (p, i_exp)
    ode_data = ode_data_list[i_exp, :, :]
    pred = predict_neuralode(u0_list[i_exp, :], p)
    l_plt = []
    for i = 1:ns
        plt = scatter(tsteps, ode_data[i, :], xscale=:log10,
                      markercolor=:transparent, label=string("data_", i))
        plot!(plt, tsteps[1:size(pred)[2]], pred[i, :], xscale=:log10, label=string("pred_", i))
        push!(l_plt, plt)
    end
    plt_all = plot(l_plt..., legend=false)
    png(plt_all, string("figs/i_exp_", i_exp))

    return false
end

l_loss_train = []
l_loss_val = []
l_grad = []
iter = 1
cb = function (p, loss_train, loss_val, g_norm)
    global l_loss_train, l_loss_val, l_grad, iter
    push!(l_loss_train, loss_train)
    push!(l_loss_val, loss_val)
    push!(l_grad, g_norm)

    if iter % n_plot == 0
        display_p(p)

        l_exp = randperm(n_exp)[1:1]
        println("update plot for ", l_exp)
        for i_exp in l_exp
            cbi(p, i_exp)
        end

        plt_loss = plot(l_loss_train, xscale=:identity, yscale=:log10, label="train")
        plot!(plt_loss, l_loss_val, xscale=:identity, yscale=:log10, label="val")
        plt_grad = plot(l_grad, xscale=:identity, yscale=:log10, label="grad_norm")
        xlabel!(plt_loss, "Epoch")
        xlabel!(plt_grad, "Epoch")
        ylabel!(plt_loss, "Loss")
        ylabel!(plt_grad, "Grad Norm")
        plt_all = plot([plt_loss, plt_grad]..., legend=:top)
        png(plt_all, "figs/loss_grad")

        @save "./checkpoint/mymodel.bson" p opt l_loss_train l_loss_val l_grad iter
    end
    iter += 1
end

if is_restart
    @load "./checkpoint/mymodel.bson" p opt l_loss_train l_loss_val l_grad iter
    iter += 1
    # opt = ADAMW(0.001, (0.9, 0.999), 1.f-6)
end

epochs = ProgressBar(iter:n_epoch);
loss_epoch = zeros(Float32, n_exp);
grad_norm = zeros(Float32, n_exp_train);
for epoch in epochs
    global p
    for i_exp in randperm(n_exp_train)
        grad = ForwardDiff.gradient(x -> loss_neuralode(x, i_exp), p)
        grad_norm[i_exp] = norm(grad, 2)
        if grad_norm[i_exp] > grad_max
            grad = grad ./ grad_norm[i_exp] .* grad_max
        end
        update!(opt, p, grad)
    end
    for i_exp in 1:n_exp
        loss_epoch[i_exp] = loss_neuralode(p, i_exp)
    end
    loss_train = mean(loss_epoch[1:n_exp_train]);
    loss_val = mean(loss_epoch[n_exp_train + 1:end]);
    g_norm = mean(grad_norm)
    set_description(epochs, string(@sprintf("Loss train %.4e val %.4e gnorm %.4e", loss_train, loss_val, g_norm)))
    cb(p, loss_train, loss_val, g_norm);
end

@printf("min loss train %.4e val %.4e\n", minimum(l_loss_train), minimum(l_loss_val))

for i_exp in 1:n_exp
    cbi(p, i_exp)
end


text:
using FileIO
using FactCheck
facts("FileIO") do
	include("query.jl")
	include("loadsave.jl")
	include("error_handling.jl")
end

# make Travis fail when tests fail:
FactCheck.exitstatus()


text:
        # real_to_parametric
        p₁ = Point2D{T}(0, 0)
        p₂ = Point2D{T}(2, 0)
        p₃ = Point2D{T}(2, 2)
        p₄ = Point2D{T}(3//2, 1//4)
        p₅ = Point2D{T}(3, 1)
        p₆ = Point2D{T}(1, 1)
        tri6 = QuadraticTriangle(p₁, p₂, p₃, p₄, p₅, p₆)
        @test tri6(real_to_parametric(p₁, tri6)) ≈ p₁
        @test tri6(real_to_parametric(p₂, tri6)) ≈ p₂
        @test tri6(real_to_parametric(p₃, tri6)) ≈ p₃
        @test tri6(real_to_parametric(p₄, tri6)) ≈ p₄
        @test tri6(real_to_parametric(p₅, tri6)) ≈ p₅
        @test tri6(real_to_parametric(p₆, tri6)) ≈ p₆

        # real_to_parametric
        p₁ = Point2D{T}(0, 0)
        p₂ = Point2D{T}(2, 0)
        p₃ = Point2D{T}(2, 3)
        p₄ = Point2D{T}(0, 3)
        p₅ = Point2D{T}(3//2, 1//2)
        p₆ = Point2D{T}(5//2, 3//2)
        p₇ = Point2D{T}(3//2, 5//2)
        p₈ = Point2D{T}(0, 3//2)
        quad8 = QuadraticQuadrilateral(p₁, p₂, p₃, p₄, p₅, p₆, p₇, p₈)
        @test quad8(real_to_parametric(p₁, quad8)) ≈ p₁
        @test quad8(real_to_parametric(p₂, quad8)) ≈ p₂
        @test quad8(real_to_parametric(p₃, quad8)) ≈ p₃
        @test quad8(real_to_parametric(p₄, quad8)) ≈ p₄
        @test quad8(real_to_parametric(p₅, quad8)) ≈ p₅
        @test quad8(real_to_parametric(p₆, quad8)) ≈ p₆

max_stars_repo_name: khurrumsaleem/MOCNeutronTransport

max_issues_repo_name: khurrumsaleem/MOCNeutronTransport

max_forks_repo_name: khurrumsaleem/MOCNeutronTransport


text:
struct SOSPolynomialBridge{
    T, F <: MOI.AbstractVectorFunction,
    DT <: SemialgebraicSets.AbstractSemialgebraicSet,
    UMCT <: MOI.ConstraintIndex{MOI.VectorOfVariables},
    UMST, MCT,
    BT <: PolyJuMP.AbstractPolynomialBasis,
    CT <: SOS.Certificate.AbstractIdealCertificate,
    MT <: MP.AbstractMonomial,
    MVT <: AbstractVector{MT}} <: MOIB.Constraint.AbstractBridge

    Q::Vector{MOI.VariableIndex} # TODO the type will be different for sparse SOS
    cQ::UMCT # TODO the type will be different for sparse SOS
    certificate_monomials::MVT
    zero_constraint::MOI.ConstraintIndex{F, PolyJuMP.ZeroPolynomialSet{DT, BT, MT, MVT}}
    domain::DT
    monomials::MVT
    certificate::CT
end

function MOI.Bridges.Constraint.bridge_constraint(
    ::Type{SOSPolynomialBridge{T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT}},
    model::MOI.ModelLike, f::MOI.AbstractVectorFunction,
    s::SOS.SOSPolynomialSet{<:SemialgebraicSets.AbstractAlgebraicSet}) where {T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT}

    @assert MOI.output_dimension(f) == length(s.monomials)
    p = MP.polynomial(collect(MOIU.eachscalar(f)), s.monomials)
    # As `*(::MOI.ScalarAffineFunction{T}, ::S)` is only defined if `S == T`, we
    # need to call `changecoefficienttype`. This is critical since `T` is
    # `Float64` when used with JuMP and the coefficient type is often `Int` if
    # `set.domain.V` is `FullSpace` or `FixedPolynomialsSet`.
    # FIXME convert needed because the coefficient type of `r` is `Any` otherwise if `domain` is `AlgebraicSet`
    r = SOS.Certificate.get(s.certificate, SOS.Certificate.ReducedPolynomial(), p, MP.changecoefficienttype(s.domain, T))
    X = SOS.Certificate.get(s.certificate, SOS.Certificate.GramBasis(), r)
    g, Q, cQ = SOS.add_gram_matrix(model, MCT, X)
    q = r - g
    set = PolyJuMP.ZeroPolynomialSet(s.domain, SOS.Certificate.zero_basis(s.certificate), MP.monomials(q))
    coefs = MOIU.vectorize(MP.coefficients(q))
    zero_constraint = MOI.add_constraint(model, coefs, set)
    return SOSPolynomialBridge{T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT}(
        Q, cQ, X, zero_constraint, s.domain, s.monomials, s.certificate)
end

function MOI.supports_constraint(::Type{SOSPolynomialBridge{T}},
                                 ::Type{<:MOI.AbstractVectorFunction},
                                 ::Type{<:SOS.SOSPolynomialSet{<:SemialgebraicSets.AbstractAlgebraicSet}}) where T
    return true
end
function MOIB.added_constrained_variable_types(::Type{<:SOSPolynomialBridge{T, F, DT, UMCT, UMST, MCT}}) where {T, F, DT, UMCT, UMST, MCT}
    return constrained_variable_types(MCT)
end
function MOIB.added_constraint_types(::Type{SOSPolynomialBridge{T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT}}) where {T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT}
    return [(F, PolyJuMP.ZeroPolynomialSet{DT, BT, MT, MVT})]
end
function MOIB.Constraint.concrete_bridge_type(
    ::Type{<:SOSPolynomialBridge{T}},
    F::Type{<:MOI.AbstractVectorFunction},
    ::Type{<:SOS.SOSPolynomialSet{DT, MT, MVT, CT}}) where {T, DT<:SemialgebraicSets.AbstractAlgebraicSet, MT, MVT, CT}
    # promotes VectorOfVariables into VectorAffineFunction, it should be enough
    # for most use cases
    G = MOIU.promote_operation(-, T, F, MOI.VectorOfVariables)
    MCT = SOS.matrix_cone_type(CT)
    UMCT = union_constraint_types(MCT)
    UMST = union_set_types(MCT)
    BT = SOS.Certificate.zero_basis_type(CT)
    return SOSPolynomialBridge{T, G, DT, UMCT, UMST, MCT, BT, CT, MT, MVT}
end

# Attributes, Bridge acting as an model
function MOI.get(bridge::SOSPolynomialBridge, ::MOI.NumberOfVariables)
    return length(bridge.Q)
end
function MOI.get(bridge::SOSPolynomialBridge, ::MOI.ListOfVariableIndices)
    return bridge.Q
end
function MOI.get(bridge::SOSPolynomialBridge{T, F, DT, UMCT, UMST},
                 ::MOI.NumberOfConstraints{MOI.VectorOfVariables, S}) where {T, F, DT, UMCT, UMST, S<:UMST}
    return bridge.cQ isa MOI.ConstraintIndex{MOI.VectorOfVariables, S} ? 1 : 0
end
function MOI.get(bridge::SOSPolynomialBridge{T, F, DT, UMCT, UMST},
                 ::MOI.ListOfConstraintIndices{MOI.VectorOfVariables, S}) where {T, F, DT, UMCT, UMST, S<:UMST}
    if bridge.cQ isa MOI.ConstraintIndex{MOI.VectorOfVariables, S}
        return [bridge.cQ]
    else
        return MOI.ConstraintIndex{MOI.VectorOfVariables, S}[]
    end
end
function MOI.get(::SOSPolynomialBridge{T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT},
                 ::MOI.NumberOfConstraints{F, PolyJuMP.ZeroPolynomialSet{DT, BT, MT, MVT}}) where {
        T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT
    }
    return 1
end
function MOI.get(b::SOSPolynomialBridge{T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT},
                 ::MOI.ListOfConstraintIndices{F, PolyJuMP.ZeroPolynomialSet{DT, BT, MT, MVT}}) where {
        T, F, DT, UMCT, UMST, MCT, BT, CT, MT, MVT
    }
    return [b.zero_constraint]
end

# Indices
function MOI.delete(model::MOI.ModelLike, bridge::SOSPolynomialBridge)
    # First delete the constraints in which the Gram matrix appears
    MOI.delete(model, bridge.zero_constraint)
    # Now we delete the Gram matrix
    MOI.delete(model, bridge.Q)
end

# Attributes, Bridge acting as a constraint
function MOI.get(model::MOI.ModelLike,
                 ::MOI.ConstraintSet,
                 bridge::SOSPolynomialBridge)
    return SOS.SOSPolynomialSet(bridge.domain, bridge.monomials, bridge.certificate)
end
function MOI.get(::MOI.ModelLike,
                 ::MOI.ConstraintPrimal,
                 ::SOSPolynomialBridge)
    throw(SOS.ValueNotSupported())
end

function MOI.get(model::MOI.ModelLike, attr::MOI.ConstraintDual,
                 bridge::SOSPolynomialBridge)
    dual = MOI.get(model, attr, bridge.zero_constraint)
    set = MOI.get(model, MOI.ConstraintSet(), bridge.zero_constraint)
    μ = MultivariateMoments.measure(dual, set.monomials)
    I = SemialgebraicSets.ideal(bridge.domain)
    return [dot(rem(mono, I), μ) for mono in bridge.monomials]
end
function MOI.get(model::MOI.ModelLike,
                 attr::MOI.ConstraintDual,
                 bridge::SOSPolynomialBridge{T, <:MOI.AbstractVectorFunction, SemialgebraicSets.FullSpace}) where {T}
    return MOI.get(model, attr, bridge.zero_constraint)
end
function MOI.get(model::MOI.ModelLike, attr::PolyJuMP.MomentsAttribute,
                 bridge::SOSPolynomialBridge)
    return MOI.get(model, attr, bridge.zero_constraint)
end

function MOI.get(::MOI.ModelLike, ::SOS.CertificateMonomials,
                 bridge::SOSPolynomialBridge)
    return bridge.certificate_monomials
end
function MOI.get(model::MOI.ModelLike,
                 attr::SOS.GramMatrixAttribute,
                 bridge::SOSPolynomialBridge)
    return SOS.build_gram_matrix(MOI.get(model, MOI.VariablePrimal(attr.N), bridge.Q),
                                 bridge.certificate_monomials)
end
function MOI.get(model::MOI.ModelLike,
                 attr::SOS.MomentMatrixAttribute,
                 bridge::SOSPolynomialBridge)
    return SOS.build_moment_matrix(MOI.get(model, MOI.ConstraintDual(attr.N),
                                           bridge.cQ),
                                   bridge.certificate_monomials)
end


text:
using Surrogates
using Stheno
using Zygote
using Test

@testset "1D - 1D" begin

    lb = 0.0
    ub = 10.0
    f = x -> log(x)*exp(x)
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)
    my_k = SthenoKriging(x,y)
    x_new = 4.0
    y_new = f.(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [5.0, 6.0]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        Zygote.gradient(x -> my_k(x), 1.0)
    end

    @testset "Optimization" begin
        objective_function = x -> 2*x+1
        x = [2.0,4.0,6.0]
        y = [5.0,9.0,13.0]
        p = 2
        a = 2
        b = 6
        my_k_EI1 = SthenoKriging(x,y)
        surrogate_optimize(objective_function,EI(),a,b,my_k_EI1,UniformSample(),maxiters=200,num_new_samples=155)
    end
end

@testset "ND - 1D" begin

    lb = [0.0; 0.0]
    ub = [5.0; 10.0]
    f = x -> log(x[1])*exp(x[2])
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)
    my_k = SthenoKriging(x,y)
    x_new = (2.0, 1.0)
    y_new = f(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [(1.0, 2.0), (0.5, 2.5)]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        g1 = Zygote.gradient(x -> my_k(x), (1.0, 1.0))[1]
        g2 = Zygote.gradient(x -> my_k(x), [1.0, 1.0])[1]
        @test [g1...] ≈ g2
    end
    @testset "Optimization" begin
        objective_function_ND = z -> 3*hypot(z...)+1
        x = [(1.2,3.0),(3.0,3.5),(5.2,5.7)]
        y = objective_function_ND.(x)
        theta = [2.0,2.0]
        lb = [1.0,1.0]
        ub = [6.0,6.0]

        my_k_E1N = SthenoKriging(x,y)
        surrogate_optimize(objective_function_ND,EI(),lb,ub,my_k_E1N,UniformSample())
    end
end

@testset "1D - ND" begin

    lb = 0.0
    ub = 10.0
    f = x -> [x, x^2]
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)

    my_k = SthenoKriging(x,y,Stheno.GP(Stheno.EQ(), Stheno.GPC()))
    x_new = 4.0
    y_new = f.(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [5.0, 6.0]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    # separate Stheno models
    gpc = Stheno.GPC()
    gp1 = Stheno.GP(Stheno.EQ(), gpc)
    gp2 = Stheno.GP(Stheno.Exp(), gpc)
    my_k = SthenoKriging(x, y, (gp1, gp2))
    x_new = 0.2
    y_new = f.(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [5.5, 6.5]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        Zygote.gradient(x -> sum(my_k(x)), 1.0)
    end
end

@testset "ND - ND" begin

    lb = [0.0; 0.0]
    ub = [5.0; 10.0]
    f = x -> [x[1]*x[2], x[2]^2]
    x = sample(5,lb,ub,SobolSample())
    y = f.(x)

    my_k = SthenoKriging(x,y,Stheno.GP(Stheno.EQ(), Stheno.GPC()))
    x_new = (1.0, 1.0)
    y_new = f(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [(1.0, 2.0), (0.5, 2.5)]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    # separate Stheno models
    gpc = Stheno.GPC()
    gp1 = Stheno.GP(Stheno.EQ(), gpc)
    gp2 = Stheno.GP(Stheno.Exp(), gpc)
    my_k = SthenoKriging(x, y, (gp1, gp2))
    x_new = (0.2, 0.1)
    y_new = f(x_new)
    add_point!(my_k, x_new, y_new)
    xs_new = [(5.5, 0.5), (6.5, 0.5)]
    ys_new = f.(xs_new)
    add_point!(my_k, xs_new, ys_new)
    y_pred = f(x_new)
    pred = my_k(x_new)
    @test pred ≈ y_pred
    pred_std = std_error_at_point(my_k, x_new)
    @test pred_std ≈ zero(y_pred) atol=1e-6

    @testset "AD" begin
        g1 = Zygote.gradient(x -> sum(my_k(x)), (1.1, 1.1))[1]
        g2 = Zygote.gradient(x -> sum(my_k(x)), [1.1, 1.1])[1]
        @test [g1...] ≈ g2
    end
end