text:
module XUnit

using Base: @lock, ReentrantLock
using Distributed #for DistributedTestRunner
using ExceptionUnwrapping: has_wrapped_exception
using EzXML
import Test
using Test: AbstractTestSet, Result, Fail, Broken, Pass, Error, guardseed
using Test: TESTSET_PRINT_ENABLE, get_testset_depth, get_testset, get_test_counts
using Test: _check_testset, parse_testset_args, push_testset, pop_testset
using TestReports
using TestReports: display_reporting_testset
using Random
using Base: Filesystem
using Base.Threads

const Option{T} = Union{Nothing,T}

abstract type TestRunner end
abstract type TestMetrics end

# BEGIN AsyncTestSuite and AsyncTestCase

mutable struct _AsyncTestCase{ASYNC_TEST_SUITE}
    testset_report::AbstractTestSet
    parent_testset::Option{Union{_AsyncTestCase{ASYNC_TEST_SUITE},ASYNC_TEST_SUITE}}
    test_fn::Function
    source::LineNumberNode
    disabled::Bool
    sub_testsets::Vector{ASYNC_TEST_SUITE}
    sub_testcases::Vector{_AsyncTestCase{ASYNC_TEST_SUITE}}
    metrics::Option{TestMetrics}
    runner::TestRunner
    success_handler::Function
    failure_handler::Function
    xml_report::Bool
    html_report::Bool
    modify_lock::ReentrantLock
end

struct AsyncTestSuite
    testset_report::AbstractTestSet
    parent_testset::Option{Union{_AsyncTestCase{AsyncTestSuite},AsyncTestSuite}}
    before_each_hook::Function
    after_each_hook::Function
    source::LineNumberNode
    disabled::Bool
    sub_testsets::Vector{AsyncTestSuite}
    sub_testcases::Vector{_AsyncTestCase{AsyncTestSuite}}
    metrics::Option{TestMetrics}
    runner::TestRunner
    success_handler::Function
    failure_handler::Function
    xml_report::Bool
    html_report::Bool
    modify_lock::ReentrantLock
end

const AsyncTestCase = _AsyncTestCase{AsyncTestSuite}
const AsyncTestSuiteOrTestCase = Union{AsyncTestSuite,AsyncTestCase}

include("rich-reporting-testset.jl")
include("test-runners.jl")
include("test-filter.jl")

function AsyncTestSuite(
    testset_report::AbstractTestSet,
    source::LineNumberNode,
    parent_testset::Option{AsyncTestSuiteOrTestCase}=nothing;
    before_each::Function = () -> nothing,
    after_each::Function = () -> nothing,
    sub_testsets::Vector{AsyncTestSuite} = AsyncTestSuite[],
    sub_testcases::Vector{AsyncTestCase} = AsyncTestCase[],
    disabled::Bool = false,
    metrics = nothing,
    runner::TestRunner = SequentialTestRunner(),
    success_handler::Function = (testset) -> nothing,
    failure_handler::Function = (testset) -> nothing,
    xml_report::Bool = false,
    html_report::Bool = false,
)
    metrics_instance = create_test_metrics(parent_testset, metrics)

    instance = AsyncTestSuite(
        testset_report,
        parent_testset,
        before_each,
        after_each,
        source,
        disabled,
        sub_testsets,
        sub_testcases,
        metrics_instance,
        runner,
        success_handler,
        failure_handler,
        xml_report,
        html_report,
        ReentrantLock(),
    )
    if parent_testset !== nothing
        lock(parent_testset.modify_lock) do
            push!(parent_testset.sub_testsets, instance)
        end
    end
    return instance
end

function AsyncTestSuite(testcase::AsyncTestCase)
    return AsyncTestSuite(
        testcase.testset_report,
        testcase.source;
        disabled = testcase.disabled,
        metrics = testcase.metrics === nothing ? nothing : typeof(testcase.metrics),
        runner = testcase.runner,
        success_handler = testcase.success_handler,
        failure_handler = testcase.failure_handler,
        xml_report = testcase.xml_report,
        html_report = testcase.html_report,
    )
end

function AsyncTestCase(
    test_fn::Function,
    testset_report::AbstractTestSet,
    source::LineNumberNode,
    parent_testset::Option{AsyncTestSuiteOrTestCase};
    disabled::Bool=false,
    metrics = nothing,
    runner::TestRunner = SequentialTestRunner(),
    success_handler::Function = (testset) -> nothing,
    failure_handler::Function = (testset) -> nothing,
    xml_report::Bool = false,
    html_report::Bool = false,
)
    metrics_instance = create_test_metrics(parent_testset, metrics)

    instance = AsyncTestCase(
        testset_report,
        parent_testset,
        test_fn,
        source,
        disabled,
        AsyncTestSuite[],
        AsyncTestCase[],
        metrics_instance,
        runner,
        success_handler,
        failure_handler,
        xml_report,
        html_report,
        ReentrantLock(),
    )
    if parent_testset !== nothing
        lock(parent_testset.modify_lock) do
            push!(parent_testset.sub_testcases, instance)
        end
    end
    return instance
end

include("test-metrics.jl")

function clear_test_reports!(testset::AsyncTestSuite)
    rich_ts = testset.testset_report
    rich_ts.reporting_test_set[] = ReportingTestSet(get_description(rich_ts))

    clear_test_reports!.(testset.sub_testsets)
    clear_test_reports!.(testset.sub_testcases)
end

function clear_test_reports!(testcase::AsyncTestCase)
    rich_ts = testcase.testset_report
    clear_test_reports!.(testset.sub_testsets)
    clear_test_reports!.(testset.sub_testcases)
    rich_ts.reporting_test_set[] = ReportingTestSet(get_description(rich_ts))
end

const TEST_SUITE_PARAMETER_NAMES = (
    Expr(:quote, :before_each),
    Expr(:quote, :after_each),
    Expr(:quote, :metrics),
    Expr(:quote, :runner),
    Expr(:quote, :success_handler),
    Expr(:quote, :failure_handler),
    Expr(:quote, :xml_report),
    Expr(:quote, :html_report),
)

const TEST_CASE_PARAMETER_NAMES = (
    Expr(:quote, :metrics),
    Expr(:quote, :runner),
    Expr(:quote, :success_handler),
    Expr(:quote, :failure_handler),
    Expr(:quote, :xml_report),
    Expr(:quote, :html_report),
)

html_output(testset::AsyncTestSuite) = html_output(testset.testset_report)
html_output(testcase::AsyncTestCase) = html_output(testcase.testset_report)
xml_output(testset::AsyncTestSuite) = xml_output(testset.testset_report)
xml_output(testcase::AsyncTestCase) = xml_output(testcase.testset_report)

function html_report(testset::AsyncTestSuiteOrTestCase)
    return html_report(testset.testset_report)
end

function xml_report(testset::AsyncTestSuiteOrTestCase)
    return xml_report(testset.testset_report)
end

function TestReports.display_reporting_testset(
    testset::AsyncTestSuiteOrTestCase;
    throw_on_error::Bool = true,
)
    TestReports.display_reporting_testset(
        testset.testset_report; throw_on_error=throw_on_error
    )
end

# END AsyncTestSuite and AsyncTestCase

# BEGIN XUnitState

"""
    struct XUnitState

State maintained during a test run, consisting of a stack of strings
for the nested `@testset`s, a maximum depth beyond which we skip `@testset`s,
and a pair of regular expressions over `@testset` nestings used to decide
which `@testset`s should be executed.  We also keep a record of tests run or
skipped so that these can be reported at the end of the test run.
"""
struct XUnitState
    test_suites_stack::Vector{AsyncTestSuiteOrTestCase}
    stack::Vector{String}
    maxdepth::Int
    include::Regex
    exclude::Regex
    seen::Dict{String,Bool}
end

"""
    create_deep_copy(x)

This function works similar to Base.deepcopy.

We saw some failures by applying `Base.deepcopy` on `RichReportingTestSet` and
`ReportingTestSet`. That's why this function is added to have more control over the impl.
"""
function create_deep_copy end
create_deep_copy(x) = deepcopy(x)

function create_deep_copy(ts::XUnitState)::XUnitState
    return XUnitState(
        map(create_deep_copy, ts.test_suites_stack),
        copy(ts.stack),
        ts.maxdepth,
        ts.include,
        ts.exclude,
        copy(ts.seen)
    )
end

function open_testset(rs::XUnitState, name::String)
    push!(rs.stack, name)
    join(rs.stack, "/")
end

function close_testset(rs::XUnitState)
    pop!(rs.stack)
end

const ⊤ = r""       # matches any string
const ⊥ = r"(?!)"   # matches no string

XUnitState() = XUnitState([], [], typemax(Int64), ⊤, ⊥, Dict{String,Bool}())

XUnitState(maxdepth::Int, include::Regex, exclude::Regex) =
   XUnitState([], [], maxdepth, include, exclude, Dict{String,Bool}())

function initialize_xunit_tls_state(tls)
    added_tls = false
    rs = if haskey(tls, :__XUNIT_STATE__)
        tls[:__XUNIT_STATE__]
    else
        added_tls = true
        val = XUnitState()
        tls[:__XUNIT_STATE__] = val
        val
    end
    return added_tls, rs
end

function finalize_xunit_tls_state(tls, added_tls)
    added_tls && delete!(tls, :__XUNIT_STATE__)
end

# END XUnitState

# Enumeration of all possible macro types
@enum SuiteType TestSetType TestCaseType

# BEGIN TestSuite

"""
    @testset "suite name" begin ... end
    @testset [before_each=()->...] [after_each=()->...] [metrics=DefaultTestMetrics]
               [success_handler=(testset)->...] [failure_handler=(testset)->...]
               [xml_report=false] [html_report=false] "suite name" begin ... end

Schedules a Test Suite

Please note that `@testset` and `@testset` macros are very similar.
The only difference is the top-level `@testset` also runs the test-cases, but a top-level
`@testset` does not run its underlying test-cases (and only schedules them).
Then, one needs to explicitly call `run_testset` over the result of this macro.

Also, note that the body of a `@testset` always gets executed at scheduling time, as it
needs to gather possible underlying `@testcase`s. Thus, it's a good practice to put your
tests under a `@testcase` (instead of putting them under a `@testset`), as any tests
defined under a `@testset` are executed sequentially at scheduling time.

## Keyword Arguments

`@testset` takes seven additional parameters:
  - `before_each`: a function to run before each underlying test-case
  - `after_each`: a function to run after each underlying test-case
  - `metrics`: a custom `TestMetrics` type
  - `success_handler`: a function to run after a successful handling of all tests. This
    function accepts the test-suite as an argument. This argument only works for the
    top-most `@testset`.
  - `failure_handler`: a function to run after a failed handling of all tests. This function
    accepts the test-suite as an argument. This argument only works for the top-most
    `@testset`.
  - `xml_report`: whether to produce the XML output file at the end.
    This argument only works for the top-most `@testset`
  - `html_report`: whether to produce the HTML output file at the end.
    This argument only works for the top-most `@testset`
"""
macro testset(args...)
    return testset_handler(args, __source__)
end

function _is_block(e)
    return false
end
function _is_block(e::Expr)
    e.head === :block
end

"""
    get_block_source(e)

A utility function for extracting the source information from a block expression
"""
function get_block_source(e)
    if e.head === :block && !isempty(e.args) && e.args[1] isa LineNumberNode
        return e.args[1]
    end
    return nothing
end

function testset_handler(args, source)
    isempty(args) && error("No arguments to @testset")

    tests = args[end]

    # Determine if a single block or for-loop style
    if !isa(tests, Expr) || (tests.head !== :for && !_is_block(tests))
        error("Expected begin/end block or for loop as argument to @testset")
    end

    if tests.head === :for
        return testset_forloop(args, tests, source)
    else
        return testset_beginend(args, tests, source, TestSetType)
    end
end

"""
    testset_forloop(args, testloop, source)

Generate the code for a `@testset` with a `for` loop argument

Note: It's not supported yet.
"""
function testset_forloop(args, testloop, source)
    error("For loop in `XUnit.@testset` is not supported.")
end

"""
    testset_beginend(args, tests, source, suite_type::SuiteType)

Generate the code for a `@testset` with a `begin`/`end` argument
"""
function testset_beginend(args, tests, source, suite_type::SuiteType)
    is_testcase = suite_type == TestCaseType

    tests_block_location = get_block_source(tests)
    tests_is_block_with_location = tests_block_location !== nothing

    # the location information inside `tests.args[1]` (if available) is more accurate
    # The `source` passed to this function is not correct is this macro is called inside
    # another macro. In that case, source will refer to the upper macro's address, not the
    # place that tests are defined
    source = tests_is_block_with_location ? tests_block_location : source
    desc, testsettype, options = XUnit.parse_testset_args(args[1:end-1])

    # `option` is a tuple creating expression that represents a key-value option
    function filter_hooks_fn(option)
        option.head == :call &&
        option.args[1] == :(=>) &&
        (
            (option.args[2] in TEST_CASE_PARAMETER_NAMES && is_testcase) ||
            (option.args[2] in TEST_SUITE_PARAMETER_NAMES && !is_testcase)
        )
    end

    # separate hook functions from other params
    hook_fn_options = :(Dict{Symbol, Any}())
    append!(hook_fn_options.args, filter(filter_hooks_fn, options.args))

    # keep non-hook options in-place
    filter!(!filter_hooks_fn, options.args)

    if desc === nothing
        desc = "test set"
    end
    # If we're at the top level we'll default to RichReportingTestSet. Otherwise
    # default to the type of the parent testset
    if testsettype === nothing
        testsettype = :(XUnit.get_testset_depth() == 0 ?
                            RichReportingTestSet :
                            typeof(XUnit.get_testset()))
    end

    # Generate a block of code that initializes a new testset, adds
    # it to the task local storage, evaluates the test(s), before
    # finally removes the testset and gives it a chance to take
    # action (such as reporting the results)
    @assert _is_block(tests)
    ex = quote
        # check that `testsettype` is a subtype of `AbstractTestSet`
        # (otherwise, throw an error)
        XUnit._check_testset($testsettype, $(QuoteNode(testsettype.args[1])))
        local ret = nothing
        local ts = $(testsettype)($desc; $options...)
        XUnit.push_testset(ts)
        # we reproduce the logic of guardseed, but this function
        # cannot be used as it changes slightly the semantic of @testset,
        # by wrapping the body in a function
        local RNG = Random.default_rng()
        local oldrng = copy(RNG)
        local ret = nothing
        try
            # RNG is re-seeded with its own seed to ease reproduce a failed test
            Random.seed!(RNG.seed)
            ret = let
                $(checked_testset_expr(
                    desc, tests, source, hook_fn_options; is_testcase = is_testcase
                ))
            end
        finally
            copy!(RNG, oldrng)
            XUnit.pop_testset()
        end

        if ret !== nothing && XUnit.get_testset_depth() == 0
            # if there was no error during the scheduling and it's the topmost `@testset`,
            # `@testset` or `@testcase`, then we want to run the scheduled tests.
            # If the `runner` keyword argument is used, that runner type is going to run
            # the tests. Otherwise, `SequentialTestRunner` is used to keep the same
            # semantics of `@testset` in `Base.Test`
            run_testset(ret)
        end

        ret
    end
    # preserve outer location if possible
    if tests_is_block_with_location
        ex = Expr(:block, tests_block_location, ex)
    end
    return ex
end

function run_testcase_inplace(testcase_obj)
    testcase_obj.disabled && return

    ts = testcase_obj.testset_report

    try
        # a test-case that is under another test-case is treated
        # like a testset and runs immediately
        # Note: `before_each` and `after_each` hooks are already
        # ran for the top-most test-case and won't run again
        # Note: there's no need to gather metrics for a test-case if it's not explicitly
        # marked for statistics collection. These non-top-level test-cases are guaranteed to
        # run sequentially (in a single process/thread) and cannot comprise of smaller
        # execution units that run out-of-order.
        if XUnit.should_report_metric(testcase_obj)
            XUnit.run_and_gather_test_metrics(testcase_obj; run=true)
        else
            testcase_obj.test_fn()
        end
    catch err
        has_wrapped_exception(err, InterruptException) && rethrow()
        # something in the test block threw an error. Count that as an
        # error in this test set
        XUnit.record(ts, XUnit.Error(
            :nontest_error, Expr(:tuple), err, Base.catch_stack(), testcase_obj.source
        ))
    end

    return testcase_obj
end

# XUnit.checked_testsuite_expr consumes hook_fn_options. Callers generate hook_fn_options
# from Test internals, namely the Test.parse_testset_args function. Test.parse_testset_args
# exposes a macro hygiene bug in Julia in its result. Consequently,
# XUnit.checked_testsuite_expr (hygiene-)escapes the result of Test.parse_testset_args
# passed via hook_fn_options to work around the macro hygiene issue. That macro hygiene
# issue was fixed during the Julia 1.6 development cycle, which changed the behavior
# of Test.parse_testset_args such that the hygiene-escaping workaround is no longer correct.
#
# To make Delve happy under both Julia 1.5 and 1.6, we introduce a function that
# applies the workaround only if necessary. This function / its calls should be
# removed once migration to Julia 1.6 is complete.
esc_if_needed(x) = VERSION < v"1.6-" ? esc(x) : x

# This function is the common function for handling the body of all `@testset`,
# and `@testcase` macros. This is how it works:
# - if a `@testset` or `@testcase` is disabled, it skips it.
# - if it's a `@testset` runs its body
# - if it's a `@testcase`:
#   - if it IS NOT enclosed inside another `@testcase`, then it schedules it for running later
#   - if it IS enclosed inside another `@testcase`, runs its body
function checked_testset_expr(
    name::Union{Expr,String}, ts_expr::Expr, source, hook_fn_options; is_testcase::Bool = false
)
    quote
        ts = get_testset()

        tls = task_local_storage()
        added_tls, rs = XUnit.initialize_xunit_tls_state(tls)

        local testset_or_testcase = nothing

        try
            std_io = IOBuffer()
            if XUnit.TESTSET_PRINT_ENABLE[]
                print(std_io, "  "^length(rs.stack))
            end
            path = XUnit.open_testset(rs, $name)
            shouldrun = length(rs.stack) <= rs.maxdepth &&
                        XUnit.pmatch(rs.include, path) != nothing &&
                        XUnit.pmatch(rs.exclude, path) == nothing
            rs.seen[path] = shouldrun
            parent_testset_obj = if isempty(rs.test_suites_stack)
                nothing
            else
                last(rs.test_suites_stack)
            end

            if shouldrun # if it's not disabled
                if XUnit.TESTSET_PRINT_ENABLE[]
                    if $is_testcase && !haskey(tls, :__TESTCASE_IS_RUNNING__)
                        # if it's a `@testcase` NOT enclosed inside another `@testcase`,
                        # then it gets scheduled for running later
                        print(std_io, "Scheduling ")
                    else
                        # otherwise, if it IS enclosed inside another `@testcase`,
                        # we run it immediately
                        print(std_io, "Running ")
                    end
                    printstyled(std_io, path; bold=true)
                    println(std_io, " tests...")
                end
            else # skip disabled `@testset`s and `@testcase`s
                if XUnit.TESTSET_PRINT_ENABLE[]
                    printstyled(
                        std_io, "Skipped Scheduling $path tests...\n"; color=:light_black
                    )
                end
            end

            seekstart(std_io)
            # thread-safe print
            print(read(std_io, String))

            $(
                if !is_testcase #if it's a `@testset`, runs its body
                    quote
                        testset_obj = XUnit.AsyncTestSuite(
                            ts, $(QuoteNode(source)), parent_testset_obj;
                            disabled=!shouldrun, $(esc_if_needed(hook_fn_options))...
                        )

                        push!(rs.test_suites_stack, testset_obj)

                        try
                            if shouldrun # if a test-suite is filtered, we completely ignore it
                                let
                                    $(esc(ts_expr))
                                end
                            end
                        catch err
                            has_wrapped_exception(err, InterruptException) && rethrow()
                            # something in the test block threw an error. Count that as an
                            # error in this test set
                            XUnit.record(ts, XUnit.Error(
                                :nontest_error,
                                Expr(:tuple),
                                err,
                                Base.catch_stack(),
                                testset_obj.source
                            ))
                        finally
                            pop!(rs.test_suites_stack)
                        end

                        testset_or_testcase = testset_obj
                    end
                else
                    quote
                        testset_or_testcase = if shouldrun # if a `@testcase` is not disabled
                            testcase_obj = XUnit.AsyncTestCase(
                                ts, $(QuoteNode(source)), parent_testset_obj;
                                disabled=!shouldrun, $(esc_if_needed(hook_fn_options))...) do
                                    $(esc(ts_expr))
                            end

                            # if it IS enclosed inside another `@testcase`, we run it immediately
                            if haskey(tls, :__TESTCASE_IS_RUNNING__)
                                push!(rs.test_suites_stack, testcase_obj)

                                try
                                    run_testcase_inplace(testcase_obj)
                                finally
                                    pop!(rs.test_suites_stack)
                                end
                            end
                            testcase_obj
                        else  # if a `@testcase` is disabled, skipt it
                            XUnit.AsyncTestCase(
                                ts, $(QuoteNode(source)), parent_testset_obj;
                                disabled=!shouldrun, $(esc_if_needed(hook_fn_options))...) do
                                    nothing
                            end
                        end
                    end
                end
            )
        finally
            XUnit.close_testset(rs)
            if added_tls
                delete!(tls, :__XUNIT_STATE__)
            end
        end

        testset_or_testcase
    end
end

# END TestSuite

# BEGIN TestCase

"""
    @testcase "test-case name" begin ... end
    @testcase [before_each=()->...] [after_each=()->...] [metrics=DefaultTestMetrics]
              [success_handler=(testcase)->...] [failure_handler=(testcase)->...]
              [xml_report=false] [html_report=false] "test-case" begin ... end

Defines a self-contained test-case.

Test-cases are gathered at scheduling time and will get executed using a test-runner.
As a test-runner can run tests in any order (and even on multiple threads/processes), it's
strongly advised that test-cases do not depend on each other.

## Keyword Arguments

`@testcase` takes four additional parameters:
  - `metrics`: a custom `TestMetrics` type
  - `success_handler`: a function to run after a successful handling of all tests. This
    function accepts the test-suite as an argument. This argument only works for the
    top-most `@testcase`.
  - `failure_handler`: a function to run after a failed handling of all tests. This function
    accepts the test-suite as an argument. This argument only works for the top-most
    `@testcase`.
  - `xml_report`: whether to produce the XML output file at the end.
    This argument only works for the top-most `@testcase`
  - `html_report`: whether to produce the HTML output file at the end.
    This argument only works for the top-most `@testcase`
"""
macro testcase(args...)
    return testcase_handler(args, __source__)
end

function testcase_handler(args, source)
    isempty(args) && error("No arguments to @testcase")

    tests = args[end]

    # Determine if a single block or for-loop style
    if !isa(tests, Expr) || !_is_block(tests)
        error("Expected begin/end block or for loop as argument to @testcase")
    end

    return testset_beginend(args, tests, source, TestCaseType)
end

# END TestCase

# BEGIN Overloading Base.Test functions

# We do not want to print all non-relevant stack-trace related to `XUnit` or Julia internals
# This function overload handles this scrubbing and stacktrace cleanup
function Test.scrub_backtrace(bt::Vector)
    do_test_ind = findfirst(ip -> Test.ip_has_file_and_func(
        ip, joinpath(@__DIR__, "rich-reporting-testset.jl"), (:display_reporting_testset,)
    ), bt)
    if do_test_ind !== nothing && length(bt) > do_test_ind
        bt = bt[do_test_ind + 1:end]
    end
    name_ind = findfirst(ip -> Test.ip_has_file_and_func(
        ip, @__FILE__, (Symbol("macro expansion"),)
    ), bt)
    if name_ind !== nothing && length(bt) != 0
        bt = bt[1:name_ind]
    end
    return bt
end

function Test.record(ts::Test.DefaultTestSet, t::Fail)
    if Distributed.myid() == 1
        printstyled(get_description(ts), ": ", color=:white)
        # don't print for interrupted tests
        if t.test_type !== :test_interrupted
            print(t)
            println()
        end
    end
    push!(ts.results, t)
    t, backtrace()
end

# END Overloading Base.Test functions

function runtests_return_state(fun::Function, depth::Int64=typemax(Int64), args...)
    includes = []
    excludes = ["(?!)"]     # seed with an unsatisfiable regex
    for arg in args
        if startswith(arg, "-") || startswith(arg, "¬")
            push!(excludes, arg[nextind(arg,1):end])
        else
            push!(includes, arg)
        end
    end
    include = partial(join(map(x -> string("(?:", x, ")"), includes), "|"))
    exclude = exact(join(map(x -> string("(?:", x, ")"), excludes), "|"))
    state = XUnitState(depth, include, exclude)
    local fn_res
    task_local_storage(:__XUNIT_STATE__, state) do
        fn_res = fun()
    end
    return (fn_res, state)
end

function runtests(fun::Function, depth::Int64=typemax(Int64), args...)
    (fn_res, state) = runtests_return_state(fun, depth, args...)
    return fn_res
end

"""
    runtests(filepath::String, args...)

Include file `filepath` and execute test sets matching the regular expressions
in `args`.  See alternative form of `runtests` for examples.
"""
function runtests(filepath::String, args...)
    runtests(typemax(Int), args...) do
        if nworkers() == 1
            # Using a single worker, there's no need to do fancy distributed execution
            @eval Main begin
                # Construct a new throw-away module in which to run the tests
                m = @eval Main module $(gensym("XUnitModule")) end  # e.g. Main.##XUnitModule#365
                # Perform the include inside the new module m
                m.include($filepath)
            end
        else
            # make sure to pass the test-state to the worker processes (mostly for test filtering)
            parent_thread_tls = task_local_storage()
            has_xunit_state = haskey(parent_thread_tls, :__XUNIT_STATE__)
            xunit_state = if has_xunit_state
                xs = create_deep_copy(parent_thread_tls[:__XUNIT_STATE__])
                empty!(xs.test_suites_stack)
                empty!(xs.stack)
                empty!(xs.seen)
                xs
            else
                nothing
            end
            Core.eval(Main, Expr(:(=), :GLOBAL_HAS_XUNIT_STATE, has_xunit_state))
            Core.eval(Main, Expr(:(=), :GLOBAL_XUNIT_STATE, xunit_state))
            @passobj 1 workers() GLOBAL_HAS_XUNIT_STATE
            @passobj 1 workers() GLOBAL_XUNIT_STATE

            # When we have several workers available, we'd prepare for `Distributed` execution
            # Even though, still, the underlying test-suite should explicitly request it via
            # the `runner=DistributedTestRunner()` keyword argument.
            XUnitModuleName = replace(string(gensym("XUnitModule")), "#" => "_")
            GLOBAL_TEST_FILENAME = filepath
            # everything is in a single line to correctly report the line numbers in
            # `filepath` back to the user
            GLOBAL_TEST_MOD = string(
                "module ",XUnitModuleName, "; ",
                "import XUnit; ",
                "function __set_tls_xunit_state(); ",
                "    xs = XUnit.create_deep_copy(Main.GLOBAL_XUNIT_STATE); ",
                "    empty!(xs.test_suites_stack); ",
                "    empty!(xs.stack); ",
                "    empty!(xs.seen); ",
                "    tls = task_local_storage(); ",
                "    tls[:__XUNIT_STATE__] = xs; ",
                "end; ",
                "Main.GLOBAL_HAS_XUNIT_STATE && __set_tls_xunit_state(); ",
                read(filepath, String),"\n",
                "end")
            Core.eval(Main, Expr(:(=), :GLOBAL_TEST_MOD, GLOBAL_TEST_MOD))
            Core.eval(Main, Expr(:(=), :GLOBAL_TEST_FILENAME, GLOBAL_TEST_FILENAME))
            @passobj 1 workers() GLOBAL_TEST_FILENAME
            @passobj 1 workers() GLOBAL_TEST_MOD
            @everywhere begin
                tls = task_local_storage()
                has_saved_source_path = haskey(tls, :SOURCE_PATH)
                saved_source_path = has_saved_source_path ? tls[:SOURCE_PATH] : nothing
                try
                    # we are setting the thread-local `:SOURCE_PATH` for Julia's `include`
                    # mechanism to work correctly. Otheriwse, the direct `include`s inside the
                    # test file located at `filepath` won't work.
                    if ispath(Main.GLOBAL_TEST_FILENAME)
                        tls[:SOURCE_PATH] = Main.GLOBAL_TEST_FILENAME
                    end

                    include_string(Main, Main.GLOBAL_TEST_MOD, Main.GLOBAL_TEST_FILENAME)

                finally
                    if has_saved_source_path
                        tls[:SOURCE_PATH] = saved_source_path
                    else
                        delete!(tls, :SOURCE_PATH)
                    end
                end
            end
        end
    end
end

"""
    runtests(args::Vector{String})

Include file `test/runtests.jl` and execute test sets matching the regular
expressions in `args` (where a leading '-' or '¬' indicates that tests
matching the expression should be excluded).

# Examples
```jldoctest
julia> runtests(["t/a/.*"])         # Run all tests under `t/a`

julia> runtests(["t/.*", "¬t/b/2"])  # Run all tests under `t` except `t/b/2`
```
"""
function runtests(args::Vector{String})
    testfile = pwd() * "/test/runtests.jl"
    if !isfile(testfile)
        @error("Could not find test/runtests.jl")
        return
    end
    runtests(testfile, args...)
end

"""
    runtests(depth::Int, args...)

Run test sets up to the provided nesting `depth` and matching the regular
expressions in `args`.
"""
function runtests(depth::Int, args...)
    testfile = pwd() * "/test/runtests.jl"
    if !isfile(testfile)
        @error("Could not find test/runtests.jl")
        return
    end
    runtests(testfile, depth, args...)
end

"""
    get_description(ts::AsyncTestSuiteOrTestCase)
    get_description(ts::AbstractTestSet)
    get_description(ts::ScheduledTest)

Returns a description of the given XUnit object
"""
function get_description end
get_description(ts::AsyncTestSuiteOrTestCase) = get_description(ts.testset_report)
get_description(ts::RichReportingTestSet) = ts.description
get_description(ts::ReportingTestSet) = ts.description
get_description(ts::Test.DefaultTestSet) = ts.description
get_description(ts::ScheduledTest) = get_description(ts.target_testcase)
function get_description(ts::T) where T <: AbstractTestSet
    return hasfield(T, :description) ? ts.description : "$T"
end

export RichReportingTestSet, html_output, html_report, xml_output, xml_report
export clear_test_reports!, test_out_io, test_err_io, test_print, test_println

export @testset, @test_broken
export runtests, showtests

#
# Purely delegated macros and functions
#
using Test: @test, @test_throws, @test_broken, @test_skip,
    @test_warn, @test_nowarn, @test_logs, @test_deprecated
using Test: @inferred
using Test: detect_ambiguities, detect_unbound_args
using Test: GenericString, GenericSet, GenericDict, GenericArray
using Test: TestSetException, LogTestFailure
using Test: get_testset, get_testset_depth
using Test: DefaultTestSet, FallbackTestSet, FallbackTestSetException, record, finish

export @test, @test_throws, @test_broken, @test_skip,
    @test_warn, @test_nowarn, @test_logs, @test_deprecated,
    @testset, @testcase
export @inferred
export detect_ambiguities, detect_unbound_args
export GenericString, GenericSet, GenericDict, GenericArray
export TestSetException
export get_description, get_testset, get_testset_depth, run_testset
export AbstractTestSet, DefaultTestSet, record, finish
export TestRunner
export SequentialTestRunner, ShuffledTestRunner, ParallelTestRunner, DistributedTestRunner
export TestMetrics, DefaultTestMetrics
export display_reporting_testset, gather_test_metrics, combine_test_metrics
export run_and_gather_test_metrics, save_test_metrics

end


text:
module OutlierDetectionTest
    using Test
    using OutlierDetection
    using Random:MersenneTwister

    import OutlierDetectionInterface
    const OD = OutlierDetectionInterface

    import MLJBase
    const MLJ = MLJBase

    export TestData, test_detector, test_meta

    include("generate_data.jl")
    include("test_detector.jl")
    include("test_meta.jl")
end

ext: jl


text:
using ActinRingsMC
using Documenter

DocMeta.setdocmeta!(ActinRingsMC, :DocTestSetup, :(using ActinRingsMC); recursive=true)

makedocs(;
    #modules=[ActinRingsMC],
    authors="Alexander Cumberworth <alex@cumberworth.org>",
    repo="https://github.com/cumberworth/ActinRingsMC.jl/blob/{commit}{path}#{line}",
    sitename="ActinRingsMC.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://cumberworth.github.io/ActinRingsMC.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/cumberworth/ActinRingsMC.jl",
    devbranch="master",
)

max_stars_repo_name: cumberworth/ActinRingsMC

max_issues_repo_name: cumberworth/ActinRingsMC

max_forks_repo_name: cumberworth/ActinRingsMC


text:
using SafeTestsets

@safetestset "testfunctions" begin include("testfunctions.jl") end
@safetestset "testfunctionderivatives" begin include("testfunctionderivatives.jl") end

# @testset "extraSpecialFunctions test" begin
#     # Write your own tests here.
# end


text:
"""
    solve(prob::FODESystem, h, PIEX())

Use explicit Product integration method to solve system of FODE.
"""
struct PIEX <: FractionalDiffEqAlgorithm end

mutable struct M
    an
    bn
    a0
    halpha1
    halpha2
    mu
    mu_tol
    r
    index_fft
    an_fft
    bn_fft
end

function solve(prob::FODESystem, h, ::PIEX)
    @unpack f, α, u0, tspan = prob

    t0 = tspan[1]; T = tspan[2]

    METH = M(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)# Initialization
    α = α[:]
    # Check compatibility size of the problem with number of fractional orders
    alpha_length = length(α)
    problem_size = size(u0, 1)

    m_alpha = ceil.(Int, α)
    m_alpha_factorial = zeros(alpha_length, maximum(m_alpha))
    for i = 1 : alpha_length
        for j = 0 : m_alpha[i]-1
            m_alpha_factorial[i, j+1] = factorial(j)
        end
    end

    f_temp = zeros(length(u0[:, 1]))
    #f_temp = sysf_vectorfield(t0, u0[:, 1], f)
    f(f_temp, u0[:, 1], nothing, t0)

    r::Int = 16
    N::Int = ceil(Int64, (T-t0)/h)
    Nr::Int = ceil(Int64, (N+1)/r)*r
    Qr::Int = ceil(Int64, log2(Nr/r)) - 1
    NNr::Int = 2^(Qr+1)*r

    # Preallocation of some variables
    y = zeros(problem_size, N+1)
    fy = zeros(problem_size, N+1)
    zn = zeros(problem_size, NNr+1)

    # Evaluation of coefficients of the PECE method
    nvett = 0:NNr+1 |> collect
    bn = zeros(alpha_length, NNr+1); an = copy(bn); a0 = copy(bn)
    for i_alpha = 1:alpha_length
        find_alpha = Float64[]
        if α[i_alpha] == α[1:i_alpha-1]
            push!(find_alpha, i_alpha)
        end

        if isempty(find_alpha) == false
            bn[i_alpha, :] = bn[find_alpha[1], :]
        elseif abs(α[i_alpha]-1) < 1e-14
            bn[i_alpha, :] = [1; zeros(NNr)]
        else
            nalpha = nvett.^α[i_alpha]
            bn[i_alpha, :] = nalpha[2:end] - nalpha[1:end-1]
        end
    end
    METH.bn = bn
    METH.halpha1 = h.^α./gamma.(α.+1)

    # Evaluation of FFT of coefficients of the PECE method
    METH.r = r
    if Qr >= 0
        index_fft = zeros(Int, 2, Qr+1)
        for l = 1 : Qr+1
            if l == 1
                index_fft[1, l] = 1; index_fft[2, l] = r*2
            else
                index_fft[1, l] = index_fft[2, l-1]+1; index_fft[2, l] = index_fft[2, l-1]+2^l*r
            end
        end

        bn_fft = zeros(Complex, alpha_length, index_fft[2, end]);# an_fft = copy(bn_fft)
        for l = 1:Qr+1
            coef_end = 2^l*r
            for i_alpha = 1 : alpha_length
                find_alpha = Float64[]
                if α[i_alpha] == α[1:i_alpha-1]
                    push!(find_alpha, i_alpha)
                end
                if isempty(find_alpha) == false
                    bn_fft[i_alpha, index_fft[1, l]:index_fft[2, l]] = bn_fft[find_alpha[1], index_fft[1, l]:index_fft[2, l]]
                else
                    bn_fft[i_alpha, index_fft[1, l]:index_fft[2, l]] = ourfft(METH.bn[i_alpha, 1:coef_end], coef_end)
                end
            end
        end
        METH.bn_fft = bn_fft ; METH.index_fft = index_fft
    end

    # Initializing solution and proces of computation
    t = t0 .+ collect(0:N)*h
    y[:, 1] = u0[:, 1]
    fy[:, 1] = f_temp
    (y, fy) = PIEXTriangolo(1, r-1, t, y, fy, zn, N, METH, problem_size, alpha_length, m_alpha, m_alpha_factorial, u0, t0, f, α)

    # Main process of computation by means of the FFT algorithm
    ff = zeros(1, 2^(Qr+2)); ff[1:2] = [0; 2] ; card_ff = 2
    nx0::Int = 0; ny0::Int = 0
    for qr = 0 : Qr
        L = 2^qr
        (y, fy) = PIDisegnaBlocchi(L, ff, r, Nr, nx0+L*r, ny0, t, y, fy, zn, N, METH, problem_size, alpha_length, m_alpha, m_alpha_factorial, u0, t0, f, α)
        ff[1:2*card_ff] = [ff[1:card_ff]; ff[1:card_ff]]
        card_ff = 2*card_ff
        ff[card_ff] = 4*L
    end

    # Evaluation solution in T when T is not in the mesh
    if T < t[N+1]
        c = (T - t[N])/h
        t[N+1] = T
        y[:, N+1] = (1-c)*y[:, N] + c*y[:, N+1]
    end

    t = t[1:N+1]; y = y[:, 1:N+1]
    return t, y

end

function PIDisegnaBlocchi(L, ff, r, Nr, nx0, ny0, t, y, fy, zn, N, METH, problem_size, alpha_length, m_alpha, m_alpha_factorial, u0, t0, f, alpha)

    nxi::Int = nx0; nxf::Int = nx0 + L*r - 1
    nyi::Int = ny0; nyf::Int = ny0 + L*r - 1
    is::Int = 1
    s_nxi = zeros(N)
    s_nxf = zeros(N)
    s_nyi = zeros(N)
    s_nyf = zeros(N)
    s_nxi[is] = nxi; s_nxf[is] = nxf; s_nyi[is] = nyi; s_nyf[is] = nyf

    i_triangolo = 0; stop = false
    while stop == false

        stop = (nxi+r-1 == nx0+L*r-1) || (nxi+r-1>=Nr-1)

        zn = PIEXQuadrato(nxi, nxf, nyi, nyf, fy, zn, N, METH, problem_size, alpha_length, alpha)

        (y, fy) = PIEXTriangolo(nxi, nxi+r-1, t, y, fy, zn, N, METH, problem_size, alpha_length, m_alpha, m_alpha_factorial, u0, t0, f, alpha)
        i_triangolo = i_triangolo + 1

        if stop == false
            if nxi+r-1 == nxf
                i_Delta = ff[i_triangolo]
                Delta = i_Delta*r
                nxi = s_nxf[is]+1; nxf = s_nxf[is]  + Delta
                nyi = s_nxf[is] - Delta +1; nyf = s_nxf[is]
                s_nxi[is] = nxi; s_nxf[is] = nxf; s_nyi[is] = nyi; s_nyf[is] = nyf
            else
                nxi = nxi + r; nxf = nxi + r - 1; nyi = nyf + 1; nyf = nyf + r
                is = is + 1
                s_nxi[is] = nxi; s_nxf[is] = nxf; s_nyi[is] = nyi; s_nyf[is] = nyf
            end
        end

    end
    return y, fy
end

function PIEXQuadrato(nxi, nxf, nyi, nyf, fy, zn, N, METH, problem_size, alpha_length, alpha)
    coef_end::Int = nxf-nyi+1
    i_fft::Int = log2(coef_end/METH.r)
    funz_beg::Int = nyi+1; funz_end::Int = nyf+1
    Nnxf::Int = min(N, nxf)

    # Evaluation convolution segment for the predictor
    vett_funz = fy[:, funz_beg:funz_end]
    vett_funz_fft = rowfft(vett_funz, coef_end)
    zzn = zeros(problem_size, coef_end)
    for i = 1 : problem_size
        i_alpha = min(alpha_length, i)
        if abs(alpha[i_alpha]-1)>1e-14
        Z = METH.bn_fft[i_alpha, METH.index_fft[1, i_fft]:METH.index_fft[2, i_fft]].*vett_funz_fft[i, :]
        zzn[i, :] = real.(ourifft(Z, coef_end))
        end
    end
    zzn = zzn[:, nxf-nyf:end-1]
    zn[:, nxi+1:Nnxf+1] = zn[:, nxi+1:Nnxf+1] + zzn[:, 1:Nnxf-nxi+1]
    return zn
end

function PIEXTriangolo(nxi, nxf, t, y, fy, zn, N, METH, problem_size, alpha_length, m_alpha, m_alpha_factorial, u0, t0, f, alpha)

    for n = nxi:min(N, nxf)
        St = StartingTerm(t[n+1], u0, m_alpha, t0, m_alpha_factorial)
        # Evaluation of the predictor
        Phi = zeros(problem_size, 1)
        if nxi == 1 # Case of the first triangle
            j_beg::Int = 0
        else # Case of any triangle but not the first
            j_beg = nxi
        end
        for j = j_beg:n-1
            Phi = Phi + METH.bn[1:alpha_length,n-j].*fy[:, j+1]
        end

        i_alpha_1 = findall(alpha -> abs(alpha - 1) < 1e-14, alpha)
        Phi[i_alpha_1] = fy[i_alpha_1, n]
        St[i_alpha_1] = y[i_alpha_1, n]

        y[:, n+1] = St + METH.halpha1.*(zn[:, n+1] + Phi)
        temp = zeros(length(y[:, n+1]))
        f(temp, y[:, n+1], nothing, t[n+1])
        #fy[:, n+1] = sysf_vectorfield(t[n+1], y[:, n+1], f)
        fy[:, n+1] = temp

    end
    return y, fy
end

sysf_vectorfield(t, y, f_fun) = f_fun(t, y)

function  StartingTerm(t, u0, m_alpha, t0, m_alpha_factorial)
    ys = zeros(size(u0, 1), 1)
    for k = 1 : maximum(m_alpha)
        if length(m_alpha) == 1
            ys = ys + (t-t0)^(k-1)/m_alpha_factorial[k]*u0[:, k]
        else
            i_alpha = findall(x -> x>=k, m_alpha)
            ys[i_alpha, 1] = ys[i_alpha, 1] + (t-t0)^(k-1)*u0[i_alpha, k]./m_alpha_factorial[i_alpha, k]
        end
    end
    return ys
end


text:
# returns a function that can repeatedly and efficiently create new seeds
function getseedfunction(seeds, pqueue, visited, weights, wrapped, keyargs)
    seedqueue = getseedqueue(weights)
    notvisited(i) = checkbounds(Bool, visited, i) && (visited[i] == 0)
    stridelist = getdimoffsets(wrapped)
    function addseed!()
        seed = findseed!(seedqueue, weights, visited)
        if seed == 0
            return 255
        end
        for i in 1:6 # 6 directions
            e = getnewedge(seed, notvisited, stridelist, i)
            if e != 0 && weights[e] > 0
                enqueue!(pqueue, e, weights[e])
            end
        end
        seedcorrection!(wrapped, seed, keyargs)
        push!(seeds, seed)
        visited[seed] = length(seeds)
        # new seed thresh
        seed_weights = weights[getedgeindex.(seed, 1:3)]
        new_seed_thresh = NBINS - div(NBINS - sum(seed_weights)/3, 2)
        return new_seed_thresh
    end
    return addseed!
end

function getseedqueue(weights)
    queue = PQueue{Int}(3NBINS)
    for (i, w) in enumerate(sum([w == 0 ? UInt8(255) : w for w in weights]; dims=1))
        enqueue!(queue, i, w)
    end
    return queue
end

function findseed!(queue::PQueue, weights, visited)
    while !isempty(queue)
        ind = dequeue!(queue)
        if visited[ind] == 0
            return ind
        end
    end
    return 0
end

function seedcorrection!(wrapped, vox, keyargs)
    if haskey(keyargs, :phase2) && haskey(keyargs, :TEs) # requires multiecho
        phase2 = keyargs[:phase2]
        TEs = keyargs[:TEs]
        best = Inf
        offset = 0
        for off1 in -2:2, off2 in -1:1
            diff = abs((wrapped[vox] + 2π*off1) / TEs[1] - (phase2[vox] + 2π*off2) / TEs[2])
            diff += (abs(off1) + abs(off2)) / 100 # small panelty for wraps (if TE1 == 2*TE2 wrong value is chosen otherwise)
            if diff < best
                best = diff
                offset = off1
            end
        end
        wrapped[vox] += 2π * offset
    else
        wrapped[vox] = rem2pi(wrapped[vox], RoundNearest)
    end
end


text:
export every, fps, fpswhen, throttle

"""
    throttle(dt, input, f=(acc,x)->x, init=value(input), reinit=x->x)

Throttle a signal to update at most once every dt seconds. By default, the throttled signal holds the last update in the time window.

This behavior can be changed by the `f`, `init` and `reinit` arguments. The `init` and `f` functions are similar to `init` and `f` in `foldp`. `reinit` is called when a new throttle time window opens to reinitialize the initial value for accumulation, it gets one argument, the previous accumulated value.

For example
    y = throttle(0.2, x, push!, Int[], _->Int[])
will create vectors of updates to the integer signal `x` which occur within 0.2 second time windows.

"""
function throttle{T}(dt, node::Signal{T}, f=(acc,x)->x, init=value(node), reinit=x->x; typ=typeof(init))
    output = Signal(typ, init, (node,))
    throttle_connect(dt, output, node, f, init, reinit)
    output
end

# Aggregate a signal producing an update at most once in dt seconds
function throttle_connect(dt, output, input, f, init, reinit)
    let collected = init, timer = Timer(x->x, 0)
        add_action!(input, output) do output, timestep
            collected = f(collected,  value(input))
            close(timer)
            timer = Timer(x -> begin push!(output, collected); collected=reinit(collected) end, dt)
        end
    end
end

"""
    every(dt)

A signal that updates every `dt` seconds to the current timestamp. Consider using `fpswhen` or `fps` before using `every`.
"""
function every(dt)
    n = Signal(time(), ())
    every_connect(dt, n)
    n
end

function every_connect(dt, output)
    outputref = WeakRef(output)
    timer = Timer(x -> _push!(outputref, time(), ()->close(timer)), dt, dt)
    finalizer(output, _->close(timer))
    output
end

"""
    fpswhen(switch, rate)

returns a signal which when `switch` signal is true, updates `rate` times every second. If `rate` is not possible to attain because of slowness in computing dependent signal values, the signal will self adjust to provide the best possible rate.
"""
function fpswhen(switch, rate)
    switch_ons = filter(x->x, false, switch) # only turn-ons
    n = Signal(Float64, 0.0, (switch, switch_ons,))
    fpswhen_connect(rate, switch, switch_ons, n)
    n
end

function setup_next_tick(outputref, switchref, dt, wait_dt)
    if value(switchref.value)
        Timer(t -> if value(switchref.value)
                       _push!(outputref, dt)
                   end, wait_dt)
    end
end

function fpswhen_connect(rate, switch, switch_ons, output)
    let prev_time = time()
        dt = 1.0/rate
        outputref = WeakRef(output)
        switchref = WeakRef(switch)

        for inp in [output, switch_ons]
            add_action!(inp, output) do output, timestep
                start_time = time()
                setup_next_tick(outputref, switchref, start_time-prev_time, dt)
                prev_time = start_time
            end
        end

        setup_next_tick(outputref, switchref, dt, dt)
    end
end

"""
    fps(rate)

Same as `fpswhen(Input(true), rate)`
"""
fps(rate) = fpswhen(Signal(Bool, true, ()), rate)


text:
# internal model structure consisting of prior and log-likelihood function

struct ESSModel{P,L} <: AbstractMCMC.AbstractModel
    "Gaussian prior."
    prior::P
    "Log likelihood function."
    loglikelihood::L

    function ESSModel{P,L}(prior::P, loglikelihood::L) where {P,L}
        isgaussian(P) || error("prior distribution has to be a Gaussian distribution")
        return new{P,L}(prior, loglikelihood)
    end
end

function ESSModel(prior, loglikelihood)
    return ESSModel{typeof(prior),typeof(loglikelihood)}(prior, loglikelihood)
end

# obtain prior
prior(model::ESSModel) = model.prior

# evaluate the loglikelihood of a sample
Distributions.loglikelihood(model::ESSModel, f) = model.loglikelihood(f)


text:
mutable struct ImageClassifier{B,F,O,L} <: MLJModelInterface.Probabilistic
    builder::B
    finaliser::F
    optimiser::O    # mutable struct from Flux/src/optimise/optimisers.jl
    loss::L         # can be called as in `loss(yhat, y)`
    epochs::Int          # number of epochs
    batch_size::Int # size of a batch
    lambda::Float64 # regularization strength
    alpha::Float64  # regularizaton mix (0 for all l2, 1 for all l1)
    optimiser_changes_trigger_retraining::Bool
    acceleration::AbstractResource
end

function ImageClassifier(; builder::B   = Short()
                         , finaliser::F = Flux.softmax
                         , optimiser::O = Flux.Optimise.ADAM()
                         , loss::L      = Flux.crossentropy
                         , epochs       = 10
                         , batch_size   = 1
                         , lambda       = 0
                         , alpha        = 0
                         , optimiser_changes_trigger_retraining = false
                         , acceleration = CPU1()
                         ) where {B,F,O,L}

    model = ImageClassifier{B,F,O,L}(builder
                                     , finaliser
                                     , optimiser
                                     , loss
                                     , epochs
                                     , batch_size
                                     , lambda
                                     , alpha
                                     , optimiser_changes_trigger_retraining
                                     , acceleration
                                     )

   message = clean!(model)
   isempty(message) || @warn message

    return model
end

function MLJModelInterface.fit(model::ImageClassifier,
                               verbosity::Int,
                               X_,
                               y_)

    data = collate(model, X_, y_)

    levels = MLJModelInterface.classes(y_[1])
    n_output = length(levels)
    n_input = size(X_[1])

    if scitype(first(X_)) <: GrayImage{A, B} where A where B
        n_channels = 1      # 1-D image
    else
        n_channels = 3      # 3-D color image
    end

    chain0 = build(model.builder, n_input, n_output, n_channels)
    chain = Flux.Chain(chain0, model.finaliser)

    optimiser = deepcopy(model.optimiser)

    chain, history = fit!(chain,
                          optimiser,
                          model.loss,
                          model.epochs,
                          model.lambda,
                          model.alpha,
                          verbosity,
                          model.acceleration,
                          data[1],
                          data[2])

    # `optimiser` is now mutated

    cache = (deepcopy(model), data, history, n_input, n_output, optimiser)
    fitresult = (chain, levels)

    report = (training_losses=history, )

    return fitresult, cache, report
end

function MLJModelInterface.predict(model::ImageClassifier, fitresult, Xnew)
    chain, levels = fitresult
    X = reformat(Xnew)
    probs = vcat([chain(X[:,:,:,idx:idx])'
                  for idx in 1:length(Xnew)]...)
    return MLJModelInterface.UnivariateFinite(levels, probs)
end

function MLJModelInterface.update(model::ImageClassifier,
                                  verbosity::Int,
                                  old_fitresult,
                                  old_cache,
                                  X,
                                  y)

    old_model, data, old_history, n_input, n_output, optimiser = old_cache
    old_chain, levels = old_fitresult

    optimiser_flag = model.optimiser_changes_trigger_retraining &&
        model.optimiser != old_model.optimiser

    keep_chain = !optimiser_flag && model.epochs >= old_model.epochs &&
        MLJModelInterface.is_same_except(model, old_model, :optimiser, :epochs)

    if keep_chain
        chain = old_chain
        epochs = model.epochs - old_model.epochs
    else
        if scitype(first(X)) <: GrayImage{A, B} where A where B
            n_channels = 1      # 1-D image
        else
            n_channels = 3      # 3-D color image
        end
        chain = Flux.Chain(build(model.builder, n_input, n_output, n_channels),
                           model.finaliser)
        data = collate(model, X, y)
        epochs = model.epochs
    end

    # we only get to keep the optimiser "state" carried over from
    # previous training if we're doing a warm restart and the user has not
    # changed the optimiser hyper-parameter:
    if !keep_chain ||
        !MLJModelInterface._equal_to_depth_one(model.optimiser,
                                              old_model.optimiser)
        optimiser = deepcopy(model.optimiser)
    end

    chain, history = fit!(chain,
                          optimiser,
                          model.loss,
                          epochs,
                          model.lambda,
                          model.alpha,
                          verbosity,
                          model.acceleration,
                          data[1],
                          data[2])
    if keep_chain
        # note: history[1] = old_history[end]
        history = vcat(old_history[1:end-1], history)
    end

    fitresult = (chain, levels)
    cache = (deepcopy(model), data, history, n_input, n_output, optimiser)
    report = (training_losses=history, )

    return fitresult, cache, report

end

MLJModelInterface.fitted_params(::ImageClassifier, fitresult) =
    (chain=fitresult[1],)

MLJModelInterface.metadata_model(ImageClassifier,
               input=AbstractVector{<:MLJModelInterface.Image},
               target=AbstractVector{<:Multiclass},
               path="MLJFlux.ImageClassifier",
               descr="A neural network model for making probabilistic predictions of a `GrayImage` target,
                given a table of `Continuous` features. ")


text:
using Base.Test
using LinearResponseVariationalBayes.SubModels
using JuMP

println("Testing SubModels.")

function test_sub_models()
  m = Model()
  @defVar(m, x)
  @defVar(m, y)
  @defVar(m, z)

  setValue(x, 1.0)
  setValue(y, 2.0)
  setValue(z, 3.0)

  @defNLExpr(obj, x * y + z ^ 2)
  @setNLObjective(m, Min, obj)

  getValue(obj)

  jo_full = JuMPObjective(m)

  z_col = z.col
  sub_vars = Bool[ false for i=1:length(m.colVal)]
  sub_vars[z_col] = true
  jo_sub = JuMPObjective(m, "sub", sub_vars)

  param_val = Float64[5, 6, 7]
  @test_approx_eq(get_objective_val(param_val, jo_full),
                  get_objective_val([param_val[z_col]], jo_sub))
  @test_approx_eq(get_objective_val(param_val, jo_full), getValue(obj))
  SubModels.get_full_objective_deriv!(param_val, jo_full)
  SubModels.get_full_objective_deriv!([param_val[z_col]], jo_sub)

  @test_approx_eq(jo_full.grad, jo_sub.grad)

  sub_grad = Float64[0.0]
  SubModels.get_objective_deriv!([param_val[z_col]], jo_sub, sub_grad)
  @test_approx_eq(jo_full.grad[z_col], sub_grad)

  @test_approx_eq(
    SubModels.get_objective_hess(param_val, jo_full)[z_col, z_col],
    SubModels.get_objective_hess([param_val[z_col]], jo_sub))

  sub_hess = zeros(1, 1)
  full_hess = zeros(3, 3)

  SubModels.get_objective_hess!(param_val, jo_full, full_hess)
  SubModels.get_objective_hess!([param_val[z_col]], jo_sub, sub_hess)
  @test_approx_eq(full_hess[z_col, z_col], sub_hess)

  optimize_subobjective(param_val, jo_sub)
  @test_approx_eq(jo_sub.colVal[z_col], 0.0)
end

function test_sub_model_optimization()
  m = Model()
  @defVar(m, x)
  @defVar(m, y)
  @defVar(m, z)

  setValue(x, 1.0)
  setValue(y, 2.0)
  setValue(z, 3.0)

  @defNLExpr(obj, x^2 * y - x * z + z ^ 2)
  @setNLObjective(m, Min, obj)

  getValue(obj)

  jo_full = JuMPObjective(m)

  z_col = z.col
  x_col = x.col
  sub_vars = Bool[ false for i=1:length(m.colVal)]
  sub_vars[[z_col, x_col]] = true
  jo_sub = JuMPObjective(m, "sub", sub_vars)

  param_val = Float64[5, 6, 7]
  z_final, val_final, optim_res =
    optimize_subobjective(param_val, jo_sub, method=:Optim)
  @test_approx_eq_eps z_final [0., 0.] 1e-6
  @test_approx_eq_eps val_final 0. 1e-6

  param_val = Float64[5, 6, 7]
  z_final, val_final, nlsolve_res =
    optimize_subobjective(param_val, jo_sub, method=:NLsolve)
  @test_approx_eq_eps z_final [0., 0.] 1e-6
  @test_approx_eq_eps val_final 0. 1e-6

end

test_sub_models()


text:
export bachelierImpliedVolatility
const alpha4 = 0.15
const aLFK4n = 7
const aLFK4d = 5

const aLFK4 = [
    0.06155371425063157,
    2.723711658728403,
    10.83806891491789,
    301.0827907126612,
    1082.864564205999,
    790.7079667603721,
    109.330638190985,
    0.1515726686825187,
    1.436062756519326,
    118.6674859663193,
    441.1914221318738,
    313.4771127147156,
    40.90187645954703,
]

const cLFK4n = 9
const cLFK4d = 7

const bLFK4 = [
    0.6409168551974357,
    788.5769356915809,
    445231.8217873989,
    149904950.4316367,
    32696572166.83277,
    4679633190389.852,
    420159669603232.9,
    2.053009222143781e+16,
    3.434507977627372e+17,
    2.012931197707014e+16,
    644.3895239520736,
    211503.4461395385,
    42017301.42101825,
    5311468782.258145,
    411727826816.0715,
    17013504968737.03,
    247411313213747.3,
]
const cLFK4 = [
    0.6421106629595358,
    654.5620600001645,
    291531.4455893533,
    69009535.38571493,
    9248876215.120627,
    479057753706.175,
    9209341680288.471,
    61502442378981.76,
    107544991866857.5,
    63146430757.94501,
    437.9924136164148,
    90735.89146171122,
    9217405.224889684,
    400973228.1961834,
    7020390994.356452,
    44654661587.93606,
    76248508709.85633,
]
const dLFK4 = [
    0.936024443848096,
    328.5399326371301,
    177612.3643595535,
    8192571.038267588,
    110475347.0617102,
    545792367.0681282,
    1033254933.287134,
    695066365.5403566,
    123629089.1036043,
    756.3653755877336,
    173.9755977685531,
    6591.71234898389,
    82796.56941455391,
    396398.9698566103,
    739196.7396982114,
    493626.035952601,
    87510.31231623856,
]

const betaStart4 = -log(alpha4)
const betaEnd = -log(eps(0.0))

@inline function eta(u::AbstractFloat)::AbstractFloat
    if u < 0.01
        return 1 -
               u * (
            0.5 +
            u *
            (1.0 / 12 + u * (1.0 / 24 + u * (19.0 / 720 + u * (3.0 / 160 + u * (863.0 / 60480 + u * 275.0 / 24192)))))
        )
    end
    return -u / log1p(-u)
end
@inline function zofy(z::AbstractFloat, betaStart::AbstractFloat, betaEnd::AbstractFloat)::AbstractFloat
    return -(log(z) + betaStart) / (betaEnd - betaStart)
end

#
# @inline function ratval(cofs::Vector{Float64}, nn::Int, dd::Int, x::Float64)::Float64
# 	sumn = cofs[nn]
# 	sumd = 0.0
# 	for j = nn - 1:-1:1
# 		sumn = sumn*x + cofs[j]
# 	end
# 	for j = nn + dd - 1:-1:nn+1
# 		sumd = (sumd + cofs[j]) * x
# 	end
# 	return sumn / (1.0 + sumd)
# end

function bachelierImpliedVolatility(
    price::Number,
    isCall::Bool,
    strike::Number,
    tte::Number,
    forward::Number,
    discountDf::Number,
)::Number
    sign = 1
    if !isCall
        sign = -1
    end
    x = (forward - strike) * sign
    if abs(x) < eps(forward)
        return price / discountDf * sqrt(2 * pi / tte)
    end
    z = if x > 0
        (price / discountDf - x) / x
    else
        -price / (discountDf * x)
    end

    if z < alpha4 #highly OTM options up to 1e-100
        u = zofy(z, betaStart4, betaEnd)
        hz = if u < 0.0091
            num =
                bLFK4[1] +
                u * (
                    bLFK4[2] +
                    u * (
                        bLFK4[3] +
                        u * (
                            bLFK4[4] +
                            u * (
                                bLFK4[5] +
                                u * (bLFK4[6] + u * (bLFK4[7] + u * (bLFK4[8] + u * (bLFK4[9] + u * bLFK4[10]))))
                            )
                        )
                    )
                )
            den =
                1.0 +
                u * (
                    bLFK4[11] +
                    u * (
                        bLFK4[12] +
                        u * (bLFK4[13] + u * (bLFK4[14] + u * (bLFK4[15] + u * (bLFK4[16] + u * (bLFK4[17])))))
                    )
                )
            num / den
        elseif u < 0.088
            num =
                cLFK4[1] +
                u * (
                    cLFK4[2] +
                    u * (
                        cLFK4[3] +
                        u * (
                            cLFK4[4] +
                            u * (
                                cLFK4[5] +
                                u * (cLFK4[6] + u * (cLFK4[7] + u * (cLFK4[8] + u * (cLFK4[9] + u * cLFK4[10]))))
                            )
                        )
                    )
                )
            den =
                1.0 +
                u * (
                    cLFK4[11] +
                    u * (
                        cLFK4[12] +
                        u * (cLFK4[13] + u * (cLFK4[14] + u * (cLFK4[15] + u * (cLFK4[16] + u * (cLFK4[17])))))
                    )
                )
            num / den
        else
            num =
                dLFK4[1] +
                u * (
                    dLFK4[2] +
                    u * (
                        dLFK4[3] +
                        u * (
                            dLFK4[4] +
                            u * (
                                dLFK4[5] +
                                u * (dLFK4[6] + u * (dLFK4[7] + u * (dLFK4[8] + u * (dLFK4[9] + u * dLFK4[10]))))
                            )
                        )
                    )
                )
            den =
                1.0 +
                u * (
                    dLFK4[11] +
                    u * (
                        dLFK4[12] +
                        u * (dLFK4[13] + u * (dLFK4[14] + u * (dLFK4[15] + u * (dLFK4[16] + u * (dLFK4[17])))))
                    )
                )
            num / den
        end
        return abs(x) / sqrt(hz * tte)
    else
        putPrice = if x <= 0
            price - x
        else
            price
        end
        z = abs(x) / putPrice
        u = eta(z)
        num =
            aLFK4[1] +
            u * (
                aLFK4[2] +
                u * (aLFK4[3] + u * (aLFK4[4] + u * (aLFK4[5] + u * (aLFK4[6] + u * (aLFK4[7] + u * (aLFK4[8]))))))
            )
        den = 1.0 + u * (aLFK4[9] + u * (aLFK4[10] + u * (aLFK4[11] + u * (aLFK4[12] + u * (aLFK4[13])))))
        hz = num / den

        return putPrice * hz / sqrt(tte)
    end
end