text:
using Test
using OpenQuantumSystems
using LinearAlgebra
using SparseArrays

# https://github.com/JuliaIO/Suppressor.jl/blob/master/src/Suppressor.jl
macro suppress(block)
    quote
        if ccall(:jl_generating_output, Cint, ()) == 0
            original_stdout = stdout
            out_rd, out_wr = redirect_stdout()
            out_reader = @async read(out_rd, String)

            original_stderr = stderr
            err_rd, err_wr = redirect_stderr()
            err_reader = @async read(err_rd, String)

            # approach adapted from https://github.com/JuliaLang/IJulia.jl/pull/667/files
            logstate = Base.CoreLogging._global_logstate
            logger = logstate.logger
            if logger.stream == original_stderr
                new_logstate =
                    Base.CoreLogging.LogState(typeof(logger)(err_wr, logger.min_level))
                Core.eval(Base.CoreLogging, Expr(:(=), :(_global_logstate), new_logstate))
            end
        end

        try
            $(esc(block))
        finally
            if ccall(:jl_generating_output, Cint, ()) == 0
                redirect_stdout(original_stdout)
                close(out_wr)

                redirect_stderr(original_stderr)
                close(err_wr)

                if logger.stream == stderr
                    Core.eval(Base.CoreLogging, Expr(:(=), :(_global_logstate), logstate))
                end
            end
        end
    end
end

@testset "aggregate" begin

    D(op1::Array, op2::Array) = abs(norm(op1 - op2))

    mode1 = Mode(0.2, 1.0)
    mode2 = Mode(0.4, 1.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1, mode2], 3, Energy)
    mol2 = Molecule([mode1, mode2], 3, Energy)
    agg = Aggregate([mol1, mol2])

    @test getNvib(agg) == [[3, 3], [3, 3]]
    @test electronicIndices(agg) == [[1, 1], [2, 1], [1, 2]]

    vibInds = [
        [[1, 1], [1, 1]],
        [[1, 1], [1, 2]],
        [[1, 1], [1, 3]],
        [[1, 1], [2, 1]],
        [[1, 1], [2, 2]],
        [[1, 1], [2, 3]],
        [[1, 1], [3, 1]],
        [[1, 1], [3, 2]],
        [[1, 1], [3, 3]],
        [[1, 2], [1, 1]],
        [[1, 2], [1, 2]],
        [[1, 2], [1, 3]],
        [[1, 2], [2, 1]],
        [[1, 2], [2, 2]],
        [[1, 2], [2, 3]],
        [[1, 2], [3, 1]],
        [[1, 2], [3, 2]],
        [[1, 2], [3, 3]],
        [[1, 3], [1, 1]],
        [[1, 3], [1, 2]],
        [[1, 3], [1, 3]],
        [[1, 3], [2, 1]],
        [[1, 3], [2, 2]],
        [[1, 3], [2, 3]],
        [[1, 3], [3, 1]],
        [[1, 3], [3, 2]],
        [[1, 3], [3, 3]],
        [[2, 1], [1, 1]],
        [[2, 1], [1, 2]],
        [[2, 1], [1, 3]],
        [[2, 1], [2, 1]],
        [[2, 1], [2, 2]],
        [[2, 1], [2, 3]],
        [[2, 1], [3, 1]],
        [[2, 1], [3, 2]],
        [[2, 1], [3, 3]],
        [[2, 2], [1, 1]],
        [[2, 2], [1, 2]],
        [[2, 2], [1, 3]],
        [[2, 2], [2, 1]],
        [[2, 2], [2, 2]],
        [[2, 2], [2, 3]],
        [[2, 2], [3, 1]],
        [[2, 2], [3, 2]],
        [[2, 2], [3, 3]],
        [[2, 3], [1, 1]],
        [[2, 3], [1, 2]],
        [[2, 3], [1, 3]],
        [[2, 3], [2, 1]],
        [[2, 3], [2, 2]],
        [[2, 3], [2, 3]],
        [[2, 3], [3, 1]],
        [[2, 3], [3, 2]],
        [[2, 3], [3, 3]],
        [[3, 1], [1, 1]],
        [[3, 1], [1, 2]],
        [[3, 1], [1, 3]],
        [[3, 1], [2, 1]],
        [[3, 1], [2, 2]],
        [[3, 1], [2, 3]],
        [[3, 1], [3, 1]],
        [[3, 1], [3, 2]],
        [[3, 1], [3, 3]],
        [[3, 2], [1, 1]],
        [[3, 2], [1, 2]],
        [[3, 2], [1, 3]],
        [[3, 2], [2, 1]],
        [[3, 2], [2, 2]],
        [[3, 2], [2, 3]],
        [[3, 2], [3, 1]],
        [[3, 2], [3, 2]],
        [[3, 2], [3, 3]],
        [[3, 3], [1, 1]],
        [[3, 3], [1, 2]],
        [[3, 3], [1, 3]],
        [[3, 3], [2, 1]],
        [[3, 3], [2, 2]],
        [[3, 3], [2, 3]],
        [[3, 3], [3, 1]],
        [[3, 3], [3, 2]],
        [[3, 3], [3, 3]],
    ]
    @test vibrationalIndices(agg) == vibInds

    aggInds = [
        [[1, 1], [[1, 1], [1, 1]]],
        [[1, 1], [[1, 1], [1, 2]]],
        [[1, 1], [[1, 1], [1, 3]]],
        [[1, 1], [[1, 1], [2, 1]]],
        [[1, 1], [[1, 1], [2, 2]]],
        [[1, 1], [[1, 1], [2, 3]]],
        [[1, 1], [[1, 1], [3, 1]]],
        [[1, 1], [[1, 1], [3, 2]]],
        [[1, 1], [[1, 1], [3, 3]]],
        [[1, 1], [[1, 2], [1, 1]]],
        [[1, 1], [[1, 2], [1, 2]]],
        [[1, 1], [[1, 2], [1, 3]]],
        [[1, 1], [[1, 2], [2, 1]]],
        [[1, 1], [[1, 2], [2, 2]]],
        [[1, 1], [[1, 2], [2, 3]]],
        [[1, 1], [[1, 2], [3, 1]]],
        [[1, 1], [[1, 2], [3, 2]]],
        [[1, 1], [[1, 2], [3, 3]]],
        [[1, 1], [[1, 3], [1, 1]]],
        [[1, 1], [[1, 3], [1, 2]]],
        [[1, 1], [[1, 3], [1, 3]]],
        [[1, 1], [[1, 3], [2, 1]]],
        [[1, 1], [[1, 3], [2, 2]]],
        [[1, 1], [[1, 3], [2, 3]]],
        [[1, 1], [[1, 3], [3, 1]]],
        [[1, 1], [[1, 3], [3, 2]]],
        [[1, 1], [[1, 3], [3, 3]]],
        [[1, 1], [[2, 1], [1, 1]]],
        [[1, 1], [[2, 1], [1, 2]]],
        [[1, 1], [[2, 1], [1, 3]]],
        [[1, 1], [[2, 1], [2, 1]]],
        [[1, 1], [[2, 1], [2, 2]]],
        [[1, 1], [[2, 1], [2, 3]]],
        [[1, 1], [[2, 1], [3, 1]]],
        [[1, 1], [[2, 1], [3, 2]]],
        [[1, 1], [[2, 1], [3, 3]]],
        [[1, 1], [[2, 2], [1, 1]]],
        [[1, 1], [[2, 2], [1, 2]]],
        [[1, 1], [[2, 2], [1, 3]]],
        [[1, 1], [[2, 2], [2, 1]]],
        [[1, 1], [[2, 2], [2, 2]]],
        [[1, 1], [[2, 2], [2, 3]]],
        [[1, 1], [[2, 2], [3, 1]]],
        [[1, 1], [[2, 2], [3, 2]]],
        [[1, 1], [[2, 2], [3, 3]]],
        [[1, 1], [[2, 3], [1, 1]]],
        [[1, 1], [[2, 3], [1, 2]]],
        [[1, 1], [[2, 3], [1, 3]]],
        [[1, 1], [[2, 3], [2, 1]]],
        [[1, 1], [[2, 3], [2, 2]]],
        [[1, 1], [[2, 3], [2, 3]]],
        [[1, 1], [[2, 3], [3, 1]]],
        [[1, 1], [[2, 3], [3, 2]]],
        [[1, 1], [[2, 3], [3, 3]]],
        [[1, 1], [[3, 1], [1, 1]]],
        [[1, 1], [[3, 1], [1, 2]]],
        [[1, 1], [[3, 1], [1, 3]]],
        [[1, 1], [[3, 1], [2, 1]]],
        [[1, 1], [[3, 1], [2, 2]]],
        [[1, 1], [[3, 1], [2, 3]]],
        [[1, 1], [[3, 1], [3, 1]]],
        [[1, 1], [[3, 1], [3, 2]]],
        [[1, 1], [[3, 1], [3, 3]]],
        [[1, 1], [[3, 2], [1, 1]]],
        [[1, 1], [[3, 2], [1, 2]]],
        [[1, 1], [[3, 2], [1, 3]]],
        [[1, 1], [[3, 2], [2, 1]]],
        [[1, 1], [[3, 2], [2, 2]]],
        [[1, 1], [[3, 2], [2, 3]]],
        [[1, 1], [[3, 2], [3, 1]]],
        [[1, 1], [[3, 2], [3, 2]]],
        [[1, 1], [[3, 2], [3, 3]]],
        [[1, 1], [[3, 3], [1, 1]]],
        [[1, 1], [[3, 3], [1, 2]]],
        [[1, 1], [[3, 3], [1, 3]]],
        [[1, 1], [[3, 3], [2, 1]]],
        [[1, 1], [[3, 3], [2, 2]]],
        [[1, 1], [[3, 3], [2, 3]]],
        [[1, 1], [[3, 3], [3, 1]]],
        [[1, 1], [[3, 3], [3, 2]]],
        [[1, 1], [[3, 3], [3, 3]]],
        [[2, 1], [[1, 1], [1, 1]]],
        [[2, 1], [[1, 1], [1, 2]]],
        [[2, 1], [[1, 1], [1, 3]]],
        [[2, 1], [[1, 1], [2, 1]]],
        [[2, 1], [[1, 1], [2, 2]]],
        [[2, 1], [[1, 1], [2, 3]]],
        [[2, 1], [[1, 1], [3, 1]]],
        [[2, 1], [[1, 1], [3, 2]]],
        [[2, 1], [[1, 1], [3, 3]]],
        [[2, 1], [[1, 2], [1, 1]]],
        [[2, 1], [[1, 2], [1, 2]]],
        [[2, 1], [[1, 2], [1, 3]]],
        [[2, 1], [[1, 2], [2, 1]]],
        [[2, 1], [[1, 2], [2, 2]]],
        [[2, 1], [[1, 2], [2, 3]]],
        [[2, 1], [[1, 2], [3, 1]]],
        [[2, 1], [[1, 2], [3, 2]]],
        [[2, 1], [[1, 2], [3, 3]]],
        [[2, 1], [[1, 3], [1, 1]]],
        [[2, 1], [[1, 3], [1, 2]]],
        [[2, 1], [[1, 3], [1, 3]]],
        [[2, 1], [[1, 3], [2, 1]]],
        [[2, 1], [[1, 3], [2, 2]]],
        [[2, 1], [[1, 3], [2, 3]]],
        [[2, 1], [[1, 3], [3, 1]]],
        [[2, 1], [[1, 3], [3, 2]]],
        [[2, 1], [[1, 3], [3, 3]]],
        [[2, 1], [[2, 1], [1, 1]]],
        [[2, 1], [[2, 1], [1, 2]]],
        [[2, 1], [[2, 1], [1, 3]]],
        [[2, 1], [[2, 1], [2, 1]]],
        [[2, 1], [[2, 1], [2, 2]]],
        [[2, 1], [[2, 1], [2, 3]]],
        [[2, 1], [[2, 1], [3, 1]]],
        [[2, 1], [[2, 1], [3, 2]]],
        [[2, 1], [[2, 1], [3, 3]]],
        [[2, 1], [[2, 2], [1, 1]]],
        [[2, 1], [[2, 2], [1, 2]]],
        [[2, 1], [[2, 2], [1, 3]]],
        [[2, 1], [[2, 2], [2, 1]]],
        [[2, 1], [[2, 2], [2, 2]]],
        [[2, 1], [[2, 2], [2, 3]]],
        [[2, 1], [[2, 2], [3, 1]]],
        [[2, 1], [[2, 2], [3, 2]]],
        [[2, 1], [[2, 2], [3, 3]]],
        [[2, 1], [[2, 3], [1, 1]]],
        [[2, 1], [[2, 3], [1, 2]]],
        [[2, 1], [[2, 3], [1, 3]]],
        [[2, 1], [[2, 3], [2, 1]]],
        [[2, 1], [[2, 3], [2, 2]]],
        [[2, 1], [[2, 3], [2, 3]]],
        [[2, 1], [[2, 3], [3, 1]]],
        [[2, 1], [[2, 3], [3, 2]]],
        [[2, 1], [[2, 3], [3, 3]]],
        [[2, 1], [[3, 1], [1, 1]]],
        [[2, 1], [[3, 1], [1, 2]]],
        [[2, 1], [[3, 1], [1, 3]]],
        [[2, 1], [[3, 1], [2, 1]]],
        [[2, 1], [[3, 1], [2, 2]]],
        [[2, 1], [[3, 1], [2, 3]]],
        [[2, 1], [[3, 1], [3, 1]]],
        [[2, 1], [[3, 1], [3, 2]]],
        [[2, 1], [[3, 1], [3, 3]]],
        [[2, 1], [[3, 2], [1, 1]]],
        [[2, 1], [[3, 2], [1, 2]]],
        [[2, 1], [[3, 2], [1, 3]]],
        [[2, 1], [[3, 2], [2, 1]]],
        [[2, 1], [[3, 2], [2, 2]]],
        [[2, 1], [[3, 2], [2, 3]]],
        [[2, 1], [[3, 2], [3, 1]]],
        [[2, 1], [[3, 2], [3, 2]]],
        [[2, 1], [[3, 2], [3, 3]]],
        [[2, 1], [[3, 3], [1, 1]]],
        [[2, 1], [[3, 3], [1, 2]]],
        [[2, 1], [[3, 3], [1, 3]]],
        [[2, 1], [[3, 3], [2, 1]]],
        [[2, 1], [[3, 3], [2, 2]]],
        [[2, 1], [[3, 3], [2, 3]]],
        [[2, 1], [[3, 3], [3, 1]]],
        [[2, 1], [[3, 3], [3, 2]]],
        [[2, 1], [[3, 3], [3, 3]]],
        [[1, 2], [[1, 1], [1, 1]]],
        [[1, 2], [[1, 1], [1, 2]]],
        [[1, 2], [[1, 1], [1, 3]]],
        [[1, 2], [[1, 1], [2, 1]]],
        [[1, 2], [[1, 1], [2, 2]]],
        [[1, 2], [[1, 1], [2, 3]]],
        [[1, 2], [[1, 1], [3, 1]]],
        [[1, 2], [[1, 1], [3, 2]]],
        [[1, 2], [[1, 1], [3, 3]]],
        [[1, 2], [[1, 2], [1, 1]]],
        [[1, 2], [[1, 2], [1, 2]]],
        [[1, 2], [[1, 2], [1, 3]]],
        [[1, 2], [[1, 2], [2, 1]]],
        [[1, 2], [[1, 2], [2, 2]]],
        [[1, 2], [[1, 2], [2, 3]]],
        [[1, 2], [[1, 2], [3, 1]]],
        [[1, 2], [[1, 2], [3, 2]]],
        [[1, 2], [[1, 2], [3, 3]]],
        [[1, 2], [[1, 3], [1, 1]]],
        [[1, 2], [[1, 3], [1, 2]]],
        [[1, 2], [[1, 3], [1, 3]]],
        [[1, 2], [[1, 3], [2, 1]]],
        [[1, 2], [[1, 3], [2, 2]]],
        [[1, 2], [[1, 3], [2, 3]]],
        [[1, 2], [[1, 3], [3, 1]]],
        [[1, 2], [[1, 3], [3, 2]]],
        [[1, 2], [[1, 3], [3, 3]]],
        [[1, 2], [[2, 1], [1, 1]]],
        [[1, 2], [[2, 1], [1, 2]]],
        [[1, 2], [[2, 1], [1, 3]]],
        [[1, 2], [[2, 1], [2, 1]]],
        [[1, 2], [[2, 1], [2, 2]]],
        [[1, 2], [[2, 1], [2, 3]]],
        [[1, 2], [[2, 1], [3, 1]]],
        [[1, 2], [[2, 1], [3, 2]]],
        [[1, 2], [[2, 1], [3, 3]]],
        [[1, 2], [[2, 2], [1, 1]]],
        [[1, 2], [[2, 2], [1, 2]]],
        [[1, 2], [[2, 2], [1, 3]]],
        [[1, 2], [[2, 2], [2, 1]]],
        [[1, 2], [[2, 2], [2, 2]]],
        [[1, 2], [[2, 2], [2, 3]]],
        [[1, 2], [[2, 2], [3, 1]]],
        [[1, 2], [[2, 2], [3, 2]]],
        [[1, 2], [[2, 2], [3, 3]]],
        [[1, 2], [[2, 3], [1, 1]]],
        [[1, 2], [[2, 3], [1, 2]]],
        [[1, 2], [[2, 3], [1, 3]]],
        [[1, 2], [[2, 3], [2, 1]]],
        [[1, 2], [[2, 3], [2, 2]]],
        [[1, 2], [[2, 3], [2, 3]]],
        [[1, 2], [[2, 3], [3, 1]]],
        [[1, 2], [[2, 3], [3, 2]]],
        [[1, 2], [[2, 3], [3, 3]]],
        [[1, 2], [[3, 1], [1, 1]]],
        [[1, 2], [[3, 1], [1, 2]]],
        [[1, 2], [[3, 1], [1, 3]]],
        [[1, 2], [[3, 1], [2, 1]]],
        [[1, 2], [[3, 1], [2, 2]]],
        [[1, 2], [[3, 1], [2, 3]]],
        [[1, 2], [[3, 1], [3, 1]]],
        [[1, 2], [[3, 1], [3, 2]]],
        [[1, 2], [[3, 1], [3, 3]]],
        [[1, 2], [[3, 2], [1, 1]]],
        [[1, 2], [[3, 2], [1, 2]]],
        [[1, 2], [[3, 2], [1, 3]]],
        [[1, 2], [[3, 2], [2, 1]]],
        [[1, 2], [[3, 2], [2, 2]]],
        [[1, 2], [[3, 2], [2, 3]]],
        [[1, 2], [[3, 2], [3, 1]]],
        [[1, 2], [[3, 2], [3, 2]]],
        [[1, 2], [[3, 2], [3, 3]]],
        [[1, 2], [[3, 3], [1, 1]]],
        [[1, 2], [[3, 3], [1, 2]]],
        [[1, 2], [[3, 3], [1, 3]]],
        [[1, 2], [[3, 3], [2, 1]]],
        [[1, 2], [[3, 3], [2, 2]]],
        [[1, 2], [[3, 3], [2, 3]]],
        [[1, 2], [[3, 3], [3, 1]]],
        [[1, 2], [[3, 3], [3, 2]]],
        [[1, 2], [[3, 3], [3, 3]]],
    ]
    @test getIndices(agg) == aggInds

    mode1 = Mode(0.2, 1.0)
    mode2 = Mode(0.3, 2.0)
    mode3 = Mode(0.4, 3.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1, mode2, mode3], 2, Energy)
    mol2 = Molecule([mode1, mode2], 2, Energy)
    mol3 = Molecule([mode1], 2, Energy)
    agg = Aggregate([mol1, mol2, mol3])

    @test getShifts(agg) == [[1.0, 2.0, 3.0], [1.0, 2.0], [1.0]]
    @test getFrequencies(agg) == [[0.2, 0.3, 0.4], [0.2, 0.3], [0.2]]

    mode1 = Mode(0.2, 1.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 2, Energy)
    mol2 = Molecule([mode1], 2, Energy)
    agg = Aggregate([mol1, mol2])
    aggInds = getIndices(agg)
    FC1 = getFranckCondonFactors(agg)
    FC2 = getFranckCondonFactors(agg, aggInds)

    @test size(FC1) == (12, 12)
    @test FC1 == FC2
    FC = [
        1.0 0.0 0.0 0.0 0.7788007830714049 0.0 0.5506953149031838 0.0 0.7788007830714049 0.5506953149031838 0.0 0.0;
        0.0 1.0 0.0 0.0 0.0 0.7788007830714049 0.0 0.5506953149031838 -0.5506953149031837 0.3894003915357024 -0.0 0.0;
        0.0 0.0 1.0 0.0 -0.5506953149031837 -0.0 0.3894003915357024 0.0 0.0 0.0 0.7788007830714049 0.5506953149031838;
        0.0 0.0 0.0 1.0 -0.0 -0.5506953149031837 0.0 0.3894003915357024 -0.0 0.0 -0.5506953149031837 0.3894003915357024;
        0.7788007830714049 0.0 -0.5506953149031837 -0.0 1.0 0.0 0.0 0.0 0.6065306597126334 0.42888194248035344 -0.42888194248035333 -0.3032653298563167;
        0.0 0.7788007830714049 -0.0 -0.5506953149031837 0.0 1.0 0.0 0.0 -0.42888194248035333 0.30326532985631666 0.30326532985631666 -0.21444097124017664;
        0.5506953149031838 0.0 0.3894003915357024 0.0 0.0 0.0 1.0 0.0 0.42888194248035344 0.30326532985631677 0.30326532985631666 0.2144409712401767;
        0.0 0.5506953149031838 0.0 0.3894003915357024 0.0 0.0 0.0 1.0 -0.3032653298563167 0.2144409712401767 -0.21444097124017664 0.15163266492815833;
        0.7788007830714049 -0.5506953149031837 0.0 -0.0 0.6065306597126334 -0.42888194248035333 0.42888194248035344 -0.3032653298563167 1.0 0.0 0.0 0.0;
        0.5506953149031838 0.3894003915357024 0.0 0.0 0.42888194248035344 0.30326532985631666 0.30326532985631677 0.2144409712401767 0.0 1.0 0.0 0.0;
        0.0 -0.0 0.7788007830714049 -0.5506953149031837 -0.42888194248035333 0.30326532985631666 0.30326532985631666 -0.21444097124017664 0.0 0.0 1.0 0.0;
        0.0 0.0 0.5506953149031838 0.3894003915357024 -0.3032653298563167 -0.21444097124017664 0.2144409712401767 0.15163266492815833 0.0 0.0 0.0 1.0
    ]
    @test 1e-12 > D(FC1, FC)

    @test getAggStateEnergy(agg, [1, 1], [[1], [2]]) == 0.4

    @test getAggStateEnergy(agg, [1, 1], [[1], [2]]) == 0.4

    @test OpenQuantumSystems.elIndOrder([1, 1, 1]) == 1
    @test OpenQuantumSystems.elIndOrder([1, 1, 2]) == 4
    @test OpenQuantumSystems.elIndOrder([2, 1, 1]) == 2

    agg.coupling[2, 3] = 200
    agg.coupling[3, 2] = 200

    FCSparse = getFranckCondonFactorsSparse(agg, aggInds)
    @test 1e-12 > D(FC, Matrix(FCSparse))

    FCSparse = getFranckCondonFactorsSparse(agg)
    @test 1e-12 > D(FC, Matrix(FCSparse))

    #=
    HamSparse1 = getAggHamiltonianSparse(agg, aggInds, FCSparse)
    HamSparse2 = getAggHamiltonianSparse(agg, aggInds)
    HamSparse3 = getAggHamiltonianSparse(agg)
    @test 1e-12 > D(Ham_ref, Matrix(HamSparse1.data))
    @test 1e-12 > D(Ham_ref, Matrix(HamSparse2.data))
    @test 1e-12 > D(Ham_ref, Matrix(HamSparse3.data))
    =#

    agg = Aggregate([mol1, mol2])
    agg.coupling[2, 3] = 200
    agg.coupling[3, 2] = 200
    agg2 = Aggregate([mol1, mol2], [0.0 0.0 0.0; 0.0 0.0 200.0; 0.0 200.0 0.0])
    @test agg2.coupling == agg.coupling

    mode1 = Mode(0.2, 1.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 2, [2.0, 200.0])
    mol2 = Molecule([mode1], 2, [3.0, 300.0])
    agg = Aggregate([mol1, mol2])
    aggInds = getIndices(agg)
    vibIndices = getVibIndices(agg, aggInds)
    FCFact = getFranckCondonFactors(agg, aggInds)

    agg.coupling[2, 3] = 50
    agg.coupling[3, 2] = 50

    @test vibIndices == [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]

    Ham_S_ref = [
        5.0 0.0 0.0;
        0.0 203.5 50.0;
        0.0 50.0 302.5
    ]
    Ham_S = getAggHamSystemSmall(agg; groundEnergy=true)
    @test 1e-12 > D(Ham_S.data, Ham_S_ref)

    Ham_S_ref = [5.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 5.0 0.0 0.0 0.0 0.0 0.0 0.0 -0.0 0.0 -0.0 0.0; 0.0 0.0 5.0 0.0 -0.0 -0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 5.0 -0.0 -0.0 0.0 0.0 -0.0 0.0 -0.0 0.0; 0.0 0.0 -0.0 -0.0 203.5 0.0 0.0 0.0 30.326532985631673 21.444097124017674 -21.444097124017667 -15.163266492815836; 0.0 0.0 -0.0 -0.0 0.0 203.5 0.0 0.0 -21.444097124017667 15.163266492815833 15.163266492815833 -10.722048562008831; 0.0 0.0 0.0 0.0 0.0 0.0 203.5 0.0 21.444097124017674 15.163266492815838 15.163266492815833 10.722048562008835; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 203.5 -15.163266492815836 10.722048562008835 -10.722048562008831 7.581633246407916; 0.0 -0.0 0.0 -0.0 30.326532985631673 -21.444097124017667 21.444097124017674 -15.163266492815836 302.5 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 21.444097124017674 15.163266492815833 15.163266492815838 10.722048562008835 0.0 302.5 0.0 0.0; 0.0 -0.0 0.0 -0.0 -21.444097124017667 15.163266492815833 15.163266492815833 -10.722048562008831 0.0 0.0 302.5 0.0; 0.0 0.0 0.0 0.0 -15.163266492815836 -10.722048562008831 10.722048562008835 7.581633246407916 0.0 0.0 0.0 302.5]
    Ham_S = getAggHamSystemBig(agg, aggInds, FCFact; groundEnergy=true)
    @test 1e-12 > D(Ham_S.data, Ham_S_ref)

    Ham_B_ref = [
        0.2  0.0  0.0  0.0
        0.0  0.4  0.0  0.0
        0.0  0.0  0.4  0.0
        0.0  0.0  0.0  0.6
    ]
    Ham_B = getAggHamBathSmall(agg; groundEnergy=true)
    @test 1e-12 > D(Ham_B.data, Ham_B_ref)

    Ham_B_ref = [
        0.2  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
        0.0  0.4  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.4  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.6  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.2  0.0  0.0  0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0  0.4  0.0  0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0  0.0  0.4  0.0  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.6  0.0  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.2  0.0  0.0  0.0
        0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.4  0.0  0.0
        0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.4  0.0
        0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.6
    ]
    Ham_B = getAggHamBathBig(agg; groundEnergy=true)
    @test 1e-12 > D(Ham_B.data, Ham_B_ref)

    Ham_0_ref = [5.2 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 5.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 5.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 5.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 203.7 0.0 0.0 0.0 30.326532985631673 21.444097124017674 -21.444097124017667 -15.163266492815836; 0.0 0.0 0.0 0.0 0.0 203.9 0.0 0.0 -21.444097124017667 15.163266492815833 15.163266492815833 -10.722048562008831; 0.0 0.0 0.0 0.0 0.0 0.0 203.9 0.0 21.444097124017674 15.163266492815838 15.163266492815833 10.722048562008835; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 204.1 -15.163266492815836 10.722048562008835 -10.722048562008831 7.581633246407916; 0.0 0.0 0.0 0.0 30.326532985631673 -21.444097124017667 21.444097124017674 -15.163266492815836 302.7 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 21.444097124017674 15.163266492815833 15.163266492815838 10.722048562008835 0.0 302.9 0.0 0.0; 0.0 0.0 0.0 0.0 -21.444097124017667 15.163266492815833 15.163266492815833 -10.722048562008831 0.0 0.0 302.9 0.0; 0.0 0.0 0.0 0.0 -15.163266492815836 -10.722048562008831 10.722048562008835 7.581633246407916 0.0 0.0 0.0 303.1]
    Ham_0 = getAggHamSystemBath(agg, aggInds, FCFact; groundEnergy=true)
    @test 1e-12 > D(Ham_0.data, Ham_0_ref)

    Ham_I_ref = [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0]
    Ham_I = getAggHamInteraction(agg, aggInds, FCFact)
    @test 1e-12 > D(Ham_I.data, Ham_I_ref)

    Ham_ref = [5.2 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 5.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 5.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 5.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 203.7 0.0 -0.7071067811865476 0.0 30.326532985631673 21.444097124017674 -21.444097124017667 -15.163266492815836; 0.0 0.0 0.0 0.0 0.0 203.9 0.0 -0.7071067811865476 -21.444097124017667 15.163266492815833 15.163266492815833 -10.722048562008831; 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 203.9 0.0 21.444097124017674 15.163266492815838 15.163266492815833 10.722048562008835; 0.0 0.0 0.0 0.0 0.0 -0.7071067811865476 0.0 204.1 -15.163266492815836 10.722048562008835 -10.722048562008831 7.581633246407916; 0.0 0.0 0.0 0.0 30.326532985631673 -21.444097124017667 21.444097124017674 -15.163266492815836 302.7 -0.7071067811865476 0.0 0.0; 0.0 0.0 0.0 0.0 21.444097124017674 15.163266492815833 15.163266492815838 10.722048562008835 -0.7071067811865476 302.9 0.0 0.0; 0.0 0.0 0.0 0.0 -21.444097124017667 15.163266492815833 15.163266492815833 -10.722048562008831 0.0 0.0 302.9 -0.7071067811865476; 0.0 0.0 0.0 0.0 -15.163266492815836 -10.722048562008831 10.722048562008835 7.581633246407916 0.0 0.0 -0.7071067811865476 303.1]
    Ham = getAggHamiltonian(agg, aggInds, FCFact; groundEnergy=true)
    @test 1e-12 > D(Ham.data, Ham_ref)

    #=
    aggInds_ref = getIndices(agg)
    vibindices_ref = getVibIndices(agg, aggInds_ref)
    aggIndLen_ref = length(aggInds_ref)
    basis_ref = GenericBasis([aggIndLen_ref])
    FCFact_ref = getFranckCondonFactors(agg, aggInds_ref)
    FCProd_ref =
        getFCProd(agg, FCFact_ref, aggInds_ref, vibindices_ref)
    Ham_ref = getAggHamiltonian(
        agg,
        aggInds_ref,
        FCFact_ref;
        groundEnergy = true,
    )
    Ham_0_ref = getAggHamSysBath(agg, aggInds_ref; groundEnergy = true)
    Ham_I_ref = Ham_ref - Ham_0_ref

    aggInds, vibindices, aggIndLen, basis, FCFact, FCProd, Ham, Ham_0, Ham_I =
        setupAggregate(agg;  groundEnergy = true, verbose = false)
    @test aggInds_ref == aggInds
    @test vibindices_ref == vibindices
    @test aggIndLen_ref == aggIndLen
    @test basis_ref == basis
    @test FCFact_ref == FCFact
    @test FCProd_ref == FCProd
    @test Ham_ref == Ham
    @test Ham_0_ref == Ham_0
    @test Ham_I_ref == Ham_I

    @suppress aggInds, vibindices, aggIndLen, basis, FCFact, FCProd, Ham, Ham_0, Ham_I =
        setupAggregate(agg; groundEnergy = true, verbose = true)
    @test aggInds_ref == aggInds
    @test vibindices_ref == vibindices
    @test aggIndLen_ref == aggIndLen
    @test basis_ref == basis
    @test FCFact_ref == FCFact
    @test FCProd_ref == FCProd
    @test Ham_ref == Ham
    @test Ham_0_ref == Ham_0
    @test Ham_I_ref == Ham_I
    =#
end


text:
using Images, Colors, GeometryTypes
using Reactive, FileIO, GLVisualize
using GLAbstraction, GeometryTypes, GLWindow, ImageFiltering

if !isdefined(:runtests)
    window = glscreen()
end
description = """
Simple slider example.
You can drag the number to change it interactively.
"""

# loadasset loads data from the GLVisualize asset folder and is defined as
# FileIO.load(assetpath(name))
racoon = loadasset("racoon.png")
# Convert to RGBA{Float32}. Float for filtering and 32 because it fits the GPU better
img = convert(Matrix{RGBA{Float32}}, racoon)

# create a slider that goes from 1-20 in 0.1 steps
slider, slider_s = widget(Signal(1f0), range = 1f0:0.1f0:20f0, window)

if !isdefined(:clamp01)
    function clamp01(color)
        mapc(color) do c
            clamp(c, 0, 1)
        end
    end
end
"""
Applies a gaussian filter to `img` and converts it to RGBA{N0f8}
"""
function myfilter(img, sigma)
    img = imfilter(img, KernelFactors.IIRGaussian((sigma, sigma)))
    # map color compononts and clamp them
    clamp01.(img)
end

startvalue = myfilter(img, value(slider_s))
# Use Reactive.async_map, to filter the image without blocking the main process
task, imgsig = async_map(myfilter, startvalue, Signal(img), slider_s)
# visualize the image signal
image_renderable = visualize(
    imgsig,
    model = translationmatrix(Vec3f0(50,100,0))
)

w = widths(value(boundingbox(slider)))
h = round(Int, w[2])

slider_screen = Screen(window, area = map(window.area) do a
    SimpleRectangle(0, 0, a.w, h)
end)
image_screen = Screen(window, area = map(window.area) do a
    SimpleRectangle(0, h, a.w, a.h-h)
end)
_view(image_renderable, image_screen, camera = :orthographic_pixel)
center!(image_screen, :orthographic_pixel)
_view(slider, slider_screen, camera = :fixed_pixel)

if !isdefined(:runtests)
    renderloop(window)
end

ext: jl


text:
module EnsembleNest

using ..Acor
using ..EnsembleSampler
using ..Stats

using HDF5
using Printf

import Base:
    write

using Statistics

export NestState, ndim, nlive, ndead, retire!, logZ, logdZ, postsample, run!

"""
Nested sampling state object.
"""
mutable struct NestState
    likelihood
    prior
    nmcmc::Int
    nmcmc_exact::Float64
    livepts::Array{Float64, 2}
    livelogps::Array{Float64, 1}
    livelogls::Array{Float64, 1}
    deadpts::Array{Float64, 2}
    deadlogls::Array{Float64,1}
    deadlogwts::Array{Float64,1}
    logx::Float64
    loglthresh::Float64
end

"""
    NestState(loglike, logprior, init, nmcmc)

Return a `NestState` object initialised from the like points in the 2D
array `init` using the given log-likelihood and log-prior functions.

`nmcmc` will be used initially to control the number of steps taken
before generating a live point to replace the low-likelihood point
that is being retired.  This parameter will be adjusted in a control
loop that tries to ensure that successive live points are
uncorrelated, as described below in the docs for `retire!`.  A setting
of `32` is a reasonable default.
"""
function NestState(logl, logp, pts::Array{Float64, 2}, nmcmc)
    npts = size(pts, 2)
    ndim = size(pts, 1)

    livelogps = zeros(npts)
    livelogls = zeros(npts)
    for i in 1:npts
        livelogps[i] = logp(pts[:,i])
        livelogls[i] = logl(pts[:,i])
    end

    NestState(logl, logp, nmcmc, nmcmc, copy(pts), livelogps, livelogls, zeros((ndim,0)), zeros(0), zeros(0), 0.0, -Inf)
end

"""Reading and writing NestState objects from HDF5"""
function NestState(f::Union{HDF5File,HDF5Group}; logl=nothing, logp=nothing)
    NestState(logl, logp, read(f, "nmcmc"), read(f, "nmcmc_exact"),
              read(f, "livepts"), read(f, "livelogps"), read(f, "livelogls"),
              read(f, "deadpts"), read(f, "deadlogls"), read(f, "deadlogwts"),
              read(f, "logx"), read(f, "loglthresh"))
end

function write(f::Union{HDF5File, HDF5Group}, ns::NestState)
    f["nmcmc"] = ns.nmcmc
    f["nmcmc_exact"] = ns.nmcmc_exact
    f["livepts", "compress", 3, "shuffle", ()] = ns.livepts
    f["livelogps", "compress", 3, "shuffle", ()] = ns.livelogps
    f["livelogls", "compress", 3, "shuffle", ()] = ns.livelogls
    f["deadpts", "compress", 3, "shuffle", ()] = ns.deadpts
    f["deadlogls", "compress", 3, "shuffle", ()] = ns.deadlogls
    f["deadlogwts", "compress", 3, "shuffle", ()] = ns.deadlogwts
    f["logx"] = ns.logx
    f["loglthresh"] = ns.loglthresh
end

"""Return the dimension of the problem in `n`."""
function ndim(n::NestState)
    size(n.livepts, 1)
end

"""Return the number of live points."""
function nlive(n::NestState)
    size(n.livepts, 2)
end

"""Return the number of dead (retired) points."""
function ndead(n::NestState)
    size(n.deadpts, 2)
end

function retire!(n::NestState)
    retire!(n::NestState, true)
end

"""
    retire!(nstate, verbose=true)

Retire the lowest-likelihood live point, using the stretch move in an MCMC to
produce its replacement.  If `verbose == true` then print information on the
retired point and the number of MCMC steps used.  In particular, print an
estimate of the autocorrelation length based on the current point's sample
chain; a '*' is added to the estimate if internal diagnostics suggest that it
may be unreliable due to a short chain.

The number of MCMC steps is adjusted so that it approaches `10*(2/accept_rate -
1)` exponentially with a rate that is `1/nlive(nstate)`.  If each accepted
stretch move generated a truly independent point, this would correspond to
running for 10 autocorrelation lengths of the resulting series to produce the
replacement live point.  The factor 10 is a "safety factor", and the exponential
approach with rate constant ensures that the internal MCMC adapts to the "local"
conditions of the likelihood vs. prior curve.  """
function retire!(n::NestState, verbose)
    nd = ndim(n)
    nl = nlive(n)

    # Find the point we are killing
    imin = argmin(n.livelogls)

    # Update the dead points, logl, logwt
    n.deadpts = cat(n.deadpts, reshape(n.livepts[:,imin], (nd, 1)), dims=2)
    push!(n.deadlogls, n.livelogls[imin])
    push!(n.deadlogwts, n.logx - log(nl))

    # Update prior fraction
    n.logx += log1p(-1.0/nl)

    # Update threshold
    n.loglthresh = n.livelogls[imin]

    # Now update the retired point
    i = rand(1:nl)
    pt = n.livepts[:,i]
    ll = n.livelogls[i]
    lp = n.livelogps[i]

    nacc = 0

    pts = zeros(nd, n.nmcmc)

    for i in 1:n.nmcmc
        q = n.livepts[:,rand(1:nl)]
        z = exp(log(0.5) + rand()*(log(2.0)-log(0.5)))
        newpt = q + z*(pt - q)

        newlp = n.prior(newpt)

        logpacc = newlp - lp + nd*log(z)

        if log(rand()) < logpacc
            newll = n.likelihood(newpt)
            if newll > n.loglthresh
                nacc += 1
                pt = newpt
                ll = newll
                lp = newlp
            end
        end

        pts[:,i] = pt
    end

    facc = float(nacc)/float(n.nmcmc)

    acl = -Inf
    star_warn = "*"
    for j in 1:nd
        ac = Acor.acf(vec(pts[j,:]))

        cac = 2.0*cumsum(ac) .- 1.0

        a = cac[round(Int, size(cac,1)/2)]
        broken = false
        for i in 1:round(Int, size(cac,1)/2) # Sum up to half the chain
            if 2*cac[i] < i # Ensure we have summed over at least 2 ACLs, then we can stop
                a = cac[i]
                broken = true
                break
            end
        end

        if a > acl
            acl = a
            if broken
                star_warn = " "
            else
                star_warn = "*"
            end
        end
    end

    # Run to at least 2 ideal ACLs; if no steps were accepted, then double current run length
    if facc == 0
        n.nmcmc_exact = (1.0 + 1.0/nl)*n.nmcmc_exact
    else
        nrun = 2*(2/facc - 1)
        n.nmcmc_exact = (1.0-1.0/nl)*n.nmcmc_exact + 1.0/nl*nrun
    end
    n.nmcmc = round(Int, n.nmcmc_exact)

    n.livepts[:,imin] = pt
    n.livelogls[imin] = ll
    n.livelogps[imin] = lp

    if verbose
        println(@sprintf("Retired point with ll = %.4f; accept = %.4f; ACL = %.1f%s; next nmcmc = %d", n.deadlogls[end], facc, acl, star_warn, n.nmcmc))
    end

    n
end

""" Return the (natural) log of the evidence and an estimate of its
uncertainty.  """
function logZ(n::NestState)
    nl = nlive(n)

    loglivewt = n.logx - log(nl)

    logZdead = logsumexp(n.deadlogwts .+ n.deadlogls)
    logZlive = logsumexp(n.livelogls .+ loglivewt)
    logZlive_big = maximum(n.livelogls) + n.logx
    logZlive_small = minimum(n.livelogls) + n.logx

    logZ = logsumexp(logZdead, logZlive)
    logZ_big = logsumexp(logZdead, logZlive_big)
    logZ_small = logsumexp(logZdead, logZlive_small)

    logZ, (logZ_big - logZ_small)
end

""" Return the pair `(samples, lnlikes)` resulting from a posterior
sampling of the given nested state."""
function postsample(n::NestState)
    nd = ndim(n)
    nl = nlive(n)

    loglivewt = n.logx - log(nl)

    pts = cat(n.deadpts, n.livepts, dims=2)
    lls = cat(n.deadlogls, n.livelogls, dims=1)
    logwts = cat(n.deadlogls .+ n.deadlogwts, n.livelogls .+ loglivewt, dims=1)

    logwtmax = maximum(logwts)
    post = zeros((nd, 0))
    logls = Float64[]

    for i in 1:size(pts, 2)
        if logwtmax + log(rand()) < logwts[i]
            post = cat(post, reshape(pts[:,i], (nd, 1)), dims=2)
            logls = push!(logls, lls[i])
        end
    end

    post, logls
end

"""
    run!(nstate, dZStop, verbose=true, ckpt_file=nothing)

Run `retire!` on the live points in `nstate` until the uncertainty in
the log evidence calculation is smaller than `dZStop`.

`verbose` is as in `retire!`.

`ckpt_file` is a filename in which to store intermediate, serialised
states of the computation.  Deserialising one of these states and
calling `run!` on it again will continue the computation"""
function run!(n::NestState, dZStop; verbose=true, ckpt_file=nothing)
    while true
        for i in 1:nlive(n)
            retire!(n, verbose)
        end

        lZ, dlZ = logZ(n)

        if verbose
            println(@sprintf("Now evolved for %d steps, dlog(Z) = %.4f", ndead(n), dlZ))
        end

        if ckpt_file != nothing
            tmpfile = "$(ckpt_file).temp"
            h5open(f -> write(f, n), tmpfile, "w")
            mv(tmpfile, ckpt_file, force=true)
        end

        if dlZ < dZStop
            break
        end
    end
end

function dic(lls::Array{Float64, 1})
    -2.0*(mean(lls) - var(lls))
end

"""
    dic(nstate)

Return the DIC for the given nested state.

The DIC is defined as ``\\mathrm{DIC} \\equiv -2\\left( \\langle \\log L
\\rangle - \\var \\log L \\right)``.
"""
function dic(ns::NestState)
    _, lls = postsample(ns)

    dic(lls)
end

end


text:
# Note. `ModelProxy` is the type of a model's metadata entry (a named
# tuple). So, `info("PCA")` has this type, for example.

## BASIC IDEA

if false

    matching(model::MLJModels.ModelProxy, X) =
        !(model.is_supervised) && scitype(X) <: model.input_scitype

    matching(model::MLJModels.ModelProxy, X, y) =
        model.is_supervised &&
        scitype(X) <: model.input_scitype &&
        scitype(y) <: model.target_scitype

    matching(model::MLJModels.ModelProxy, X, y, w::AbstractVector{<:Real}) =
        model.is_supervised &&
        model.supports_weights &&
        scitype(X) <: model.input_scitype &&
        scitype(y) <: model.target_scitype

end

## IMPLEMENTATION

struct ModelChecker{is_supervised,
                  supports_weights,
                  input_scitype,
                    target_scitype} end

function Base.getproperty(::ModelChecker{is_supervised,
                                         supports_weights,
                                         input_scitype,
                                         target_scitype},
                          field::Symbol) where {is_supervised,
                                                supports_weights,
                                                input_scitype,
                                                target_scitype}
    if field === :is_supervised
        return is_supervised
    elseif field === :supports_weights
        return supports_weights
    elseif field === :input_scitype
        return input_scitype
    elseif field === :target_scitype
        return target_scitype
    else
        throw(ArgumentError("Unsupported property. "))
    end
end

Base.propertynames(::ModelChecker) =
    (:is_supervised, :supports_weights, :input_scitype, :target_scitype)

function _as_named_tuple(s::ModelChecker)
    names = propertynames(s)
    NamedTuple{names}(Tuple(getproperty(s, p) for p in names))
end

# function Base.show(io::IO, ::MIME"text/plain", S::ModelChecker)
#     show(io, MIME("text/plain"), _as_named_tuple(S))
# end

"""
   matching(model, X, y)

Returns `true` exactly when the registry metadata entry `model` is
supervised and admits inputs and targets with the scientific types of
`X` and `y`, respectively.

   matching(model, X)

Returns `true` exactly when `model` is unsupervised and admits inputs
with the scientific types of `X`.

    matching(model), matching(X, y), matching(X)

Curried versions of the preceding methods, i.e., `Bool`-valued
callable objects satisfying `matching(X, y)(model) = matching(model,
X, y)`, etc.

### Example

    models(matching(X))

Finds all unsupervised models compatible with input data `X`.

    models() do model
        matching(model, X, y) && model.prediction_type == :probabilistic
    end

Finds all supervised models compatible with input data `X` and target
data `y` and making probabilistic predictions.

See also [`models`](@ref)

"""
matching(X)       = ModelChecker{false,false,scitype(X),missing}()
matching(X, y)    = ModelChecker{true,false,scitype(X),scitype(y)}()
matching(X, y, w) = ModelChecker{true,true,scitype(X),scitype(y)}()

(f::ModelChecker{false,false,XS,missing})(model::MLJModels.ModelProxy) where XS =
    !(model.is_supervised) &&
    XS <: model.input_scitype

(f::ModelChecker{true,false,XS,yS})(model::MLJModels.ModelProxy) where {XS,yS} =
    model.is_supervised &&
    XS <: model.input_scitype &&
    yS <: model.target_scitype

(f::ModelChecker{true,true,XS,yS})(model::MLJModels.ModelProxy) where {XS,yS} =
    model.is_supervised &&
    model.supports_weights &&
    XS <: model.input_scitype &&
    yS <: model.target_scitype

(f::ModelChecker)(name::String; pkg=nothing) = f(info(name, pkg=pkg))
(f::ModelChecker)(realmodel::Model) = f(info(realmodel))

matching(model::MLJModels.ModelProxy, args...) = matching(args...)(model)
matching(name::String, args...; pkg=nothing) =
    matching(info(name, pkg=pkg), args...)
matching(realmodel::Model, args...) = matching(info(realmodel), args...)

## DUAL NOTION

struct DataChecker
    model::MLJModels.ModelProxy
end

matching(model::MLJModels.ModelProxy) = DataChecker(model)
matching(name::String; pkg=nothing) = DataChecker(info(name, pkg=pkg))
matching(realmodel::Model) = matching(info(realmodel))

(f::DataChecker)(args...) = matching(f.model, args...)


text:
using Random, Plots, Distributions, Statistics

n = 20
k = n + 1
b = 10  # Poisson upper bound
γ = 0.9
θ = 0.1
V = zeros(Float64, k, k)
PI = zeros(Int64, k, k)

# day rental probility at location A and B for one day
A_rp =
    [(j - i, i, pdf(Poisson(3), i) * pdf(Poisson(4), j)) for i = 0:b, j = 0:b]
B_rp =
    [(j - i, i, pdf(Poisson(3), i) * pdf(Poisson(2), j)) for i = 0:b, j = 0:b]

# day rental at the end of day with a given number of cars
# return element: (car number delta, number of cars rent out, probility)
function day_rental(rp, start)
    broadcast((x, y) -> (clamp(y + x[1], 0, n), min(y, x[2]), x[3]), rp, start)
end

# value evaluation
function value_eval(A, B, orig_act)
    global k, γ, V
    v = 0.0

    for a in A, b in B
        # calculate actual action
        act = orig_act >= 0 ? min(a[1], orig_act, n - b[1]) :
            -min(b[1], abs(orig_act), n - a[1])

        # next state
        i2 = a[1] - act + 1
        j2 = b[1] + act + 1

        # extra parking cost for car number > 10
        extra = 0.0
        if a[1] - act > 10
            extra += 4
        end
        if b[1] + act > 10
            extra += 4
        end

        # free bus transport 1 car from A to B
        if act > 1
            act -= 1
        end

        v +=
            a[3] *
            b[3] *
            ((a[2] + b[2]) * 10 - abs(act) * 2 - extra + γ * V[i2, j2])
    end

    v
end

function policy_evaluation()
    global n, k, γ, θ, V, PI, A_rp, B_rp

    while true
        delta = 0.0
        for i = 1:k, j = 1:k
            A = day_rental(A_rp, i - 1)
            B = day_rental(B_rp, j - 1)

            v = value_eval(A, B, PI[i, j])

            delta = max(delta, abs(v - V[i, j]))
            V[i, j] = v
        end
        @show delta
        if delta < θ
            break
        end
    end
end

function policy_improvement()
    global n, k, γ, V, PI, A_rp, B_rp
    stable = true

    for i = 1:k, j = 1:k
        A = day_rental(A_rp, i - 1)
        B = day_rental(B_rp, j - 1)
        max_v, max_act = 0.0, 0

        for act = -5:5
            v = value_eval(A, B, act)

            if v >= max_v
                max_act = act
                max_v = v
            end
        end

        if PI[i, j] != max_act
            stable = false
        end
        PI[i, j] = max_act
    end

    stable
end

# policy iteration
policy_evaluation()
while !policy_improvement()
    policy_evaluation()
end

# draw plot for V and π
heatmap(V)
heatmap(PI)

max_stars_repo_path: cheatsheet/Solutions/Chapter 4/Ex4.7.jl

max_issues_repo_path: cheatsheet/Solutions/Chapter 4/Ex4.7.jl

max_forks_repo_path: cheatsheet/Solutions/Chapter 4/Ex4.7.jl


text:
module PowerNumbers

using Base, DualNumbers, DomainSets, LinearAlgebra, SingularIntegralEquations, SingularIntegralEquations.HypergeometricFunctions

import Base: convert, *, +, -, ==, <, <=, >, |, !, !=, >=, /, ^, \, isinf, in, isapprox
import Base: exp, atanh, log1p, abs, max, min, log, inv, real, imag, conj, sqrt, sin, cos

import DualNumbers: Dual, realpart, epsilon, dual
import DomainSets: UnionDomain, TypedEndpointsInterval
import SingularIntegralEquations: Directed, undirected
import SingularIntegralEquations.HypergeometricFunctions: speciallog

export PowerNumber, LogNumber

include("LogNumber.jl")

struct PowerNumber <: Number
    realpart::ComplexF64
    epsilon::ComplexF64
    alpha::Float64

    PowerNumber(realpart,epsilon,alpha) =
    alpha >= 2.0 ? new(realpart,0,1) :
    alpha <= -2.0 ? new(Inf,0,1) :
    alpha ≈ 0 ? new(realpart+epsilon,0,1) :
    new(realpart,epsilon,alpha)
end

realpart(r::PowerNumber) = r.realpart
epsilon(r::PowerNumber) = r.epsilon
alpha(r::PowerNumber) = r.alpha

Base.promote_rule(::Type{PowerNumber}, ::Type{<:Number}) = PowerNumber
Base.convert(::Type{PowerNumber}, z::PowerNumber) = z
Base.convert(::Type{PowerNumber}, z::Number) = PowerNumber(0, z, 0)

undirected(r::PowerNumber) = alpha(r) < 0.0 ? Inf : undirected(realpart(r))
isinf(r::PowerNumber) = isinf(realpart(r)) || alpha(r) < 0.0

PowerNumber(x::Dual) = PowerNumber(realpart(x), epsilon(x), 1)
Dual(x::PowerNumber) = alpha(x) == 1 ? Dual(realpart(x), epsilon(x)) : throw("α must equal 1 to convert to dual.")
dual(x::PowerNumber) = Dual(x)

function (x::PowerNumber)(ε)
    a,b,α = realpart(x), epsilon(x), alpha(x)
    isinf(a) && return b/(ε^α)
    realpart(x) + epsilon(x)*ε^α
end

in(x::PowerNumber, d::Domain) = realpart(x) ∈ d
in(x::PowerNumber, d::TypedEndpointsInterval{:closed,:closed}) = realpart(x) ∈ d

function +(x::PowerNumber, y::PowerNumber)
    a, b, α = realpart(x), epsilon(x), alpha(x)
    c, d, β = realpart(y), epsilon(y), alpha(y)
    coeff = [b, d]
    alphs = [α, β]
    for γ in sort(alphs)
        tot = sum(coeff .* (alphs .≈ γ))
        if tot != 0
            return PowerNumber(a+c,tot,γ)
        end
    end
    return PowerNumber(a+c,0,1)
end

+(x::PowerNumber, y::Number) = PowerNumber(realpart(x)+y,epsilon(x), alpha(x))
+(y::Number, x::PowerNumber) = PowerNumber(realpart(x)+y,epsilon(x), alpha(x))

function *(x::PowerNumber, y::PowerNumber)
    a, b, α = realpart(x), epsilon(x), alpha(x)
    c, d, β = realpart(y), epsilon(y), alpha(y)
    if α+β ≈ 0
        return PowerNumber(a*c,b*c,α) + PowerNumber(b*d,a*d,β)
    else
        return PowerNumber(a*c,b*c,α) + PowerNumber(0,a*d,β) + PowerNumber(0,b*d,α+β)
    end
end

for Typ in (:Bool, :Number)
    @eval begin
        *(x::PowerNumber, y::$Typ) = PowerNumber(realpart(x)*y, epsilon(x)*y, alpha(x))
        *(y::$Typ, x::PowerNumber) = PowerNumber(realpart(x)*y, epsilon(x)*y, alpha(x))
    end
end

-(x::PowerNumber) = PowerNumber(-realpart(x), -epsilon(x), alpha(x))
-(x::PowerNumber, y::PowerNumber) = +(x, -y)

function inv(z::PowerNumber)
    x, y, α = realpart(z), epsilon(z), alpha(z)
    (α > 0.0 && x!=0) ? PowerNumber(1/x,-y/(x^2),α) : PowerNumber(0,1/y,-α)
end

/(z::PowerNumber, x::PowerNumber) = z*inv(x)
/(z::PowerNumber, x::Number) = z*inv(x)
/(x::Number, z::PowerNumber) = x*inv(z)

for OP in (:*, :+, :-, :/)
    @eval begin
        $OP(a::Directed{s}, b::PowerNumber) where {s} = Directed{s}($OP(a.x,b))
        $OP(a::PowerNumber, b::Directed{s}) where {s} = Directed{s}($OP(a,b.x))
    end
end

==(a::PowerNumber, b::PowerNumber) = realpart(a) == realpart(b) && epsilon(a) == epsilon(b) && alpha(a) == alpha(b)
isapprox(a::PowerNumber, b::PowerNumber; opts...) = ≈(realpart(a), realpart(b); opts...) && ≈(epsilon(a), epsilon(b); opts...) && ≈(alpha(a), alpha(b); opts...)

function log(z::PowerNumber)
    x, y, α = realpart(z), epsilon(z), alpha(z)
    tol = 1E-14
    if abs(x) ≤ tol || α < 0.0
        return LogNumber(α,log(y))
    elseif abs(x) >= 1/tol
        return LogNumber(-α,log(y))
    else throw(ArgumentError("Cannot eval at $z"))
    end
end

Complex{Float64}(a::LogNumber) = a(1)

function atanh(z::PowerNumber)
    x, y, α = realpart(z), epsilon(z), alpha(z)
    if x ≈ 1
        LogNumber(-α/2,log(2)/2  - log(-y)/2)
    elseif x ≈ -1
        LogNumber(α/2,-log(2)/2  + log(y)/2)
    else
        atanh(x)
    end
end

log1p(z::PowerNumber) = log(z+1)

for op in (:real, :imag, :conj)
    @eval $op(l::PowerNumber) = PowerNumber($op(realpart(l)), $op(epsilon(l)), alpha(l))
end

abs(l::PowerNumber) = abs(realpart(l))

function sqrt(z::PowerNumber)
    x, y, α = realpart(z), epsilon(z), alpha(z)
    (x!=0) ? PowerNumber(sqrt(x),y/(2*sqrt(x)),α) : PowerNumber(0,sqrt(y),α/2)
end

function ^(z::PowerNumber, p::Number)
    x, y, α = realpart(z), epsilon(z), alpha(z)
    (x!=0) ? error("Not implemented.") : PowerNumber(0,y^p,α*p)
end

# Analytic functions are trivial
for op in (:cos, :sin, :exp) # TODO: use same list as in Dual Numbers
	@eval function $op(z::PowerNumber)
	    x, y, α = realpart(z), epsilon(z), alpha(z)
	    fz = $op(dual(x,y))
	    PowerNumber(realpart(fz), epsilon(fz), α)
	end
end

#speciallog(x::PowerNumber{<:Complex}) = alpha(x) == 1.0 ? (s = sqrt(x); 3(atanh(s)-realpart(s))/realpart(s)^3) : error("Only implemented for alpha = 1")

function SingularIntegralEquations.HypergeometricFunctions.speciallog(x::PowerNumber)
    s = sqrt(x)
    return 3(atanh(s)-realpart(s))/realpart(s)^3
end

Base.show(io::IO, x::PowerNumber) = print(io, "($(realpart(x))) + ($(epsilon(x)))ε^$(alpha(x))")

end # module


text:
using Oceananigans
using Oceananigans.Units
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
using Printf
using GLMakie
using JLD2

grid = RectilinearGrid(CPU(),
                       topology = (Periodic, Bounded, Bounded),
                       size = (128, 128, 16),
                       x = (-500kilometers, 500kilometers),
                       y = (-500kilometers, 500kilometers),
                       z = (-4kilometers, 0),
                       halo = (3, 3, 3))

const Lz = grid.Lz

# Uncomment to put a bump in the grid:
# This will slow down the simulation because we will use a
# matrix-based Poisson solver instead of the FFT-based solver.

const width = 50kilometers
@inline bump(x, y) = - Lz * (1 - 0.5 * exp(-(x^2 + y^2) / 2width^2))
grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bump))
free_surface = ImplicitFreeSurface(solver_method=:HeptadiagonalIterativeSolver)

#free_surface = ImplicitFreeSurface()

# Physics

Δx = grid.Lx / grid.Nx
κ₄h = Δx^4 / 1day
κz = 1e-2

diffusive_closure = VerticalScalarDiffusivity(VerticallyImplicitTimeDiscretization(), ν=κz, κ=κz)
horizontal_closure = HorizontalScalarBiharmonicDiffusivity(ν=κ₄h, κ=κ₄h)

model = HydrostaticFreeSurfaceModel(; grid, free_surface,
                                    coriolis = BetaPlane(latitude = -45),
                                    buoyancy = BuoyancyTracer(),
                                    closure = (diffusive_closure, horizontal_closure),
                                    tracers = (:b, :c),
                                    momentum_advection = WENO5(),
                                    tracer_advection = WENO5())

# Initial condition: a baroclinically unstable situation!
ramp(y, δy) = min(max(0, y/δy + 1/2), 1)

# Parameters
N² = 4e-6 # [s⁻²] buoyancy frequency / stratification
M² = 8e-8 # [s⁻²] horizontal buoyancy gradient

δy = 50kilometers
δz = 400
Lz = grid.Lz

δc = 2δy
δb = δy * M²
ϵb = 1e-2 * δb # noise amplitude

bᵢ(x, y, z) = N² * z + δb * ramp(y, δy) + ϵb * randn()
cᵢ(x, y, z) = exp(-y^2 / 2δc^2) * exp(-(z + Lz/4)^2 / 2δz^2)

set!(model, b=bᵢ, c=cᵢ)

simulation = Simulation(model, Δt=10minutes, stop_time=60days)

wizard = TimeStepWizard(cfl=0.2, max_change=1.1, max_Δt=simulation.Δt)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(1))

print_progress(sim) =
    @printf("Iter: %d, time: %s, wall time: %s, max|u|: %6.3e, m s⁻¹, next Δt: %s\n",
            iteration(sim), prettytime(sim), prettytime(sim.run_wall_time),
            maximum(abs, sim.model.velocities.u), prettytime(sim.Δt))

simulation.callbacks[:print_progress] = Callback(print_progress, IterationInterval(20))

b, c = model.tracers
u, v, w = model.velocities
ζ = Field(∂x(v) - ∂y(u))

simulation.output_writers[:fields] = JLD2OutputWriter(model, (; ζ, b, c),
                                                      schedule = TimeInterval(1hour),
                                                      field_slicer = FieldSlicer(k=grid.Nz),
                                                      prefix = "baroclinic_adjustment",
                                                      force = true)

run!(simulation)

file = jldopen("baroclinic_adjustment.jld2")
iterations = parse.(Int, keys(file["timeseries/t"]))
ζ = [file["timeseries/ζ/$i"][:, :, 1] for i in iterations]
b = [file["timeseries/b/$i"][:, :, 1] for i in iterations]
c = [file["timeseries/c/$i"][:, :, 1] for i in iterations]
t = [file["timeseries/t/$i"] for i in iterations]
close(file)

maxζ = maximum(maximum(abs, ζi) for ζi in ζ)
maxb = maximum(maximum(bi) for bi in b)
minb = minimum(minimum(bi) for bi in b)

Nt = length(t)

fig = Figure(resolution=(1800, 600))

axζ = Axis(fig[1, 1])
axb = Axis(fig[1, 2])
axc = Axis(fig[1, 3])

slider = Slider(fig[3, :], range=1:Nt, startvalue=Nt)
n = slider.value

xζ, yζ, zζ = 1e3 .* nodes((Face, Face, Center), grid)
xc, yc, zc = 1e3 .* nodes((Center, Center, Center), grid)

ζⁿ = @lift ζ[$n]
bⁿ = @lift b[$n]
cⁿ = @lift c[$n]

hmζ = heatmap!(axζ, xζ, yζ, ζⁿ, colormap=:redblue, colorrange=(-maxζ, maxζ))
hmb = heatmap!(axb, xc, yc, bⁿ, colormap=:thermal, colorrange=(minb, maxb))
hmc = heatmap!(axc, xc, yc, cⁿ, colormap=:deep, colorrange=(0, 0.3))

Colorbar(fig[2, 1], hmζ, vertical=false, flipaxis=true, label="Vertical vorticity (s⁻¹)")
Colorbar(fig[2, 2], hmb, vertical=false, flipaxis=true, label="Buoyancy (m s⁻²)")
Colorbar(fig[2, 3], hmc, vertical=false, flipaxis=true, label="Tracer concentration")

title = @lift "Baroclinic adjustment at t = " * prettytime(t[$n])
Label(fig[0, :], title)

display(fig)

# record(fig, "beta_plane_baroclinic_adjustment.mp4", 1:Nt, framerate=12) do nn
#     @info "Rendering frame $nn of $Nt..."
#     n[] = nn
# end


text:
function euler_step(f, t::Number, x, len)
    return len' .* f(t, x)
end

function mod_euler_step(f, t::Number, x, len)
    whole_step = len' .* f(t, x)
    return (whole_step + len' .* f(t .+ len, x .+ whole_step))/2
end

function heum_step(f, t::Number, x, len)
    whole_step = len' .* f(t, x)
    return (whole_step + 3*len' .* f(t .+ 2/3*len, x .+ 2/3*whole_step))/4
end

function RK4_step(f, t::Number, x, len)
    k1 = len' .* f(t, x)
    k2 = len' .* f(t .+ len/2, x .+ k1/2)
    k3 = len' .* f(t .+ len/2, x .+ k2/2)
    k4 = len' .* f(t .+ len, x .+ k3)
    return (k1 + 2*k2 + 2*k3 + k4)/6
end

macro define_order(f, ord::Integer)
    @eval function $f(s::Symbol)
        if s == :order
            return $ord
        end
        throw(ErrorException("Invalid symbol. Use :order"))
    end
end

@define_order(euler_step, 1)
@define_order(mod_euler_step, 2)
@define_order(heum_step, 2)
@define_order(RK4_step, 4)

function fixed_step_ode_solver(
    f,
    t0::Number,
    x0;
    method=RK4_step,
    step::Number = 1e-3,
    iter::Integer = 10000
)
    t = t0
    x = Array{typeof(x0 + step*f(t0, x0))}(undef, iter+1)
    x[1] = x0
    for j = 1:iter
        x[j+1] = x[j] + method(f, t, x[j], step)
        t += step
    end
    return x
end

function fixed_step_ode_solver(
    f,
    t0::Number,
    x0,
    stop_cond;
    method=RK4_step,
    step::Number = 1e-3,
    maxit::Integer = 10000
)
    t = t0
    x = Array{typeof(x0 + method(f, t, x0, step)), 1}()
    push!(x, x0)
    for j = 1:maxit
        push!(x, last(x) + method(f, t, last(x), step))
        t += step
        if stop_cond(t, x)
            break
        end
    end
    return x
end

max_stars_repo_name: useredsa/numeric-differential-equations

max_issues_repo_name: useredsa/numeric-differential-equations

max_forks_repo_name: useredsa/numeric-differential-equations


text:
export NestedEinsum

"""
    parse_nested(s::AbstractString (, iy = []))

return a contraction-tree consisting of `NestedEinsumConstructor`
objects as nodes and `IndexGroup` as leaves.
`s` should only contain the `ixs` part
"""

function parse_nested(s::AbstractString, iy = [])
    count(==('('),s) == count(==(')'),s) || throw(
        ArgumentError("Parentheses don't pair up in $s"))

    _, out = parse_parens(s, firstindex(s), 1)
    append!(out.iy, iy)
    filliys!(out)
    return construct(out)
end

"""
    parse_parens(s::AbstractString, i, narg)

parse one level of parens starting at index `i` where `narg` counts which tensor the
current group of indices, e.g. "ijk", belongs to.
Recursively calls itself for each new opening paren that's opened.
"""
function parse_parens(s::AbstractString, i, narg)
    out = NestedEinsumConstructor{Char}([], [], [])
    g = IndexGroup{Char}([],narg)

    while i <= lastindex(s)
        c = s[i]
        j = nextind(s,i)
        if c === '('
            # opening a parens means that next a new contraction is parsed
            j, out2, narg = parse_parens(s, j, narg)
            push!(out, out2)
            union!(out.inds, out2.inds)
        elseif c === ')' || c === ','
            # either closing a parens or a comma means that a tensors indices
            # are complete and can be added to the contraction-object
            if !isempty(g)
                push!(out, g)
                union!(out.inds, g.inds)
            end
            if  c === ','
                # comma implies that a new tensor is parsed for the current contraction
                narg += 1
                g = IndexGroup{Char}([], narg)
            else
                # parens implies that the current contraction is complete
                return j, out, narg
            end
        elseif isletter(c) # could allow isnumeric() too
            push!(g,c)
        else
            throw(ArgumentError("parsing $s failed, $c is not a valid entry"))
        end
        i = j
    end
    if !isempty(g)
        push!(out, g)
        union!(out.inds, g.inds)
    end
    return i, out
end

"""
    filliys!(neinsum::NestedEinsumConstructor)

goes through all `NestedEinsumConstructor` objects in the tree and saves the correct `iy` in them.
"""
function filliys!(neinsum)
    iy = neinsum.iy
    args = neinsum.args
    for i in 1:length(neinsum.args)
        arg = args[i]
        arg isa IndexGroup && continue
        union!(arg.iy, intersect(arg.inds, iy))
        for j in 1:length(neinsum.args)
            i === j && continue
            union!(arg.iy, intersect(arg.inds, args[j].inds))
        end
        filliys!(arg)
    end
    return neinsum
end

"""
    IndexGroup

Leaf in a contractiontree, contains the indices and the number of the tensor it
describes, e.g. in "ij,jk -> ik", indices "ik" belong to tensor `1`, so
would be described by IndexGroup(['i','k'], 1).
"""
struct IndexGroup{T}
    inds::Vector{T}
    n::Int
end

Base.push!(ig::IndexGroup, c) = (push!(ig.inds,c); ig)
Base.isempty(ig::IndexGroup) = isempty(ig.inds)

"""
    NestedEinsumConstructor

describes a (potentially) nested einsum. Important fields:
- `args`, vector of all inputs, either `IndexGroup` objects corresponding to tensors or `NestedEinsumConstructor`
- `iy`, indices of output
"""
struct NestedEinsumConstructor{T}
    args::Vector{Union{NestedEinsumConstructor{T}, IndexGroup{T}}}
    inds::Vector{T}
    iy::Vector{T}
end

Base.push!(neinsum::NestedEinsumConstructor, x) = (push!(neinsum.args,x); neinsum)

using MacroTools

function _nested_ein_macro(ex; einsum=:einsum)
    @capture(ex, (left_ := right_)) || throw(ArgumentError("expected A[] := B[]... "))
    @capture(left, Z_[leftind__] | [leftind__] ) || throw(
        ArgumentError("can't understand LHS, expected A[i,j] etc."))
    Z === nothing && @gensym Z
    primefix!(leftind)
    allinds = unique(leftind)

    MacroTools.postwalk(right) do x
        @capture(x, A_[inds__]) && union!(allinds, inds)
        x
    end
    primefix!(allinds)

    tensors = Symbol[]
    nein = parse_nested_expr(right, tensors, allinds)
    append!(nein.iy, indexin(leftind,allinds))
    filliys!(nein)
    snein = construct(nein)

    tensornames = map(esc, tensors)
    :($(esc(Z)) = $snein(($(tensornames...),)...))
end

function parse_nested_expr(expr, tensors, allinds)
    if @capture(expr, *(args__))
        einargs = map(x -> parse_nested_expr(x,tensors, allinds), args)
        intinds = union(mapreduce(x -> x.inds, vcat, einargs))
        return NestedEinsumConstructor{Int}(einargs, intinds, Int[])
    elseif @capture(expr, A_[inds__])
        push!(tensors,A)
        return IndexGroup{Int}(indexin(primefix!(inds), allinds), length(tensors))
    end
end

struct NestedEinsum
    args
    eins
end

function construct(nein::NestedEinsumConstructor{T}) where T
    ixs = Tuple(map(extractixs, nein.args))
    iy = Tuple(nein.iy)
    eins = EinCode{ixs,iy}()
    args = Tuple(map(x -> x isa NestedEinsumConstructor ? construct(x) : x.n,nein.args))
    return NestedEinsum(args, eins)
end
extractixs(x::IndexGroup) = Tuple(x.inds)
extractixs(x::NestedEinsumConstructor) = Tuple(x.iy)

function (neinsum::NestedEinsum)(xs...; size_info = nothing)
    mxs = map(x -> extractxs(xs, x), neinsum.args)
    neinsum.eins(mxs...; size_info=size_info)
end

extractxs(xs, x::NestedEinsum) = x(xs...)
extractxs(xs, i::Int) = xs[i]

function match_rule(code::NestedEinsum)
    return (match_rule(code.eins), match_rule.(code.args))
end

match_rule(code::Int64) = code

# Better printing
using AbstractTrees

function AbstractTrees.children(ne::NestedEinsum)
    d = Dict()
    for (k,item) in enumerate(ne.args)
        d[k] = item isa Integer ? _join(OMEinsum.getixs(ne.eins)[k]) : item
    end
    d
end

function AbstractTrees.printnode(io::IO, x::String)
    print(io, x)
end
function AbstractTrees.printnode(io::IO, x::NestedEinsum)
    print(io, x.eins)
end

function Base.show(io::IO, e::EinCode{ixs, iy}) where {ixs, iy}
    s = join([_join(ix) for ix in ixs], ", ") * " -> " * _join(iy)
    print(io, s)
end
function Base.show(io::IO, e::NestedEinsum)
    print_tree(io, e)
end
Base.show(io::IO, ::MIME"text/plain", e::NestedEinsum) = show(io, e)
Base.show(io::IO, ::MIME"text/plain", e::EinCode) = show(io, e)
_join(ix::NTuple{0}) = ""
_join(ix::NTuple{N,Char}) where N = join(ix, "")
_join(ix::NTuple{N,Int}) where N = join(ix, "∘")
_join(ix::NTuple{N,Any}) where N = join(ix, "-")

# flattten nested einsum
function _flatten(code::NestedEinsum, iy=nothing)
    ixs = []
    for i=1:length(code.args)
        append!(ixs, _flatten(code.args[i], OMEinsum.getixs(code.eins)[i]))
    end
    return ixs
end
_flatten(i::Int, iy) = [i=>iy]

flatten(code::EinCode) = code
function flatten(code::NestedEinsum)
    ixd = Dict(_flatten(code))
    EinCode(([ixd[i] for i=1:length(ixd)]...,), OMEinsum.getiy(code.eins))
end

dynamic_extractxs(xs, x::NestedEinsum; size_info) = dynamic_einsum(x, xs; size_info=size_info)
dynamic_extractxs(xs, x::Int; size_info) = xs[x]
function dynamic_einsum(ne::NestedEinsum, xs; size_info=nothing)
    mxs = map(x->dynamic_extractxs(xs, x; size_info=size_info), ne.args)
    dynamic_einsum(getixs(ne.eins), mxs, getiy(ne.eins); size_info=size_info)
end


text:
#Pkg.add("SankeyPlots")
using SankeyPlots

src = [1, 1, 1, 1, 2, 2, 2, 3, 4, 5]
dst = [6, 3, 7, 4, 3, 7, 4, 7, 8, 8]
weights = [0.1, 0.3, 0.5, 0.5, 0.2, 2.8, 1, 0.45, 4.5, 3.3]

SankeyPlots.sankey(src, dst, weights)

# 2

names = [
    "PV",
    "Electricity Buy",
    "Battery",
    "Heat pump",
    "Biomass",
    "Electricity Sell",
    "Electricity Demand",
    "Heat demand",
]
energy_colors = SankeyPlots.palette(:seaborn_colorblind)[[9, 10, 3, 5, 2, 8, 1, 4]]

SankeyPlots.sankey(
    src, dst, weights;
    node_labels=names,
    node_colors=energy_colors,
    edge_color=:gradient,
    label_position=:bottom,
    label_size=7,
    compact=true,
)

max_stars_repo_name: Ellon-M/visualizations

max_issues_repo_name: Ellon-M/visualizations

max_forks_repo_name: Ellon-M/visualizations


text:
function getDrawnNumbers(s)
    numbers = split(s,",")
    return parse.(Int32,numbers)
end

function getBoards(lines)
    nboards = (length(lines)-1) ÷ 6
    boards = zeros(Int32,0)
    for i in 2:length(lines)
        if i % 6 != 2
            nums = parse.(Int32, filter(x->length(x)>0, split(lines[i], " ")))
            append!(boards, nums)
        end
    end
    # note that as Julia has column-major arrays,
    # my boards are transposed
    return (nboards, reshape(boards, (5, 5, nboards)))
end

function processInput(filename)
    lines = readlines(filename)
    return (getDrawnNumbers(lines[1]), getBoards(lines))
end

function isBingo(drawnNumbers, row)
    return issetequal(union(drawnNumbers, row), drawnNumbers)
end

function calcScore(board, drawnNumbers)
    return (sum(board) - sum(intersect(board,drawnNumbers))) * last(drawnNumbers)
end

function bingo(filename)
    (drawnNumbers, (nboards, boards)) = processInput(filename)
    winningBoards = zeros(Int32, 0)
    for n in 1:length(drawnNumbers)
        for board in 1:nboards
            if !in(board, winningBoards)
                for row in 1:5
                    # note our matrices are transcribed,
                    # cols and rows are interpreded just the other way arond
                    if isBingo(drawnNumbers[1:n], boards[:, row, board])
                        @show board, "row: ", row, n, calcScore(boards[:, :, board], drawnNumbers[1:n])
                        append!(winningBoards, board)
                    end
                end
                for col in 1:5
                    # note our matrices are transcribed,
                    # cols and rows are interpreded just the other way arond
                    if isBingo(drawnNumbers[1:n], boards[col, :, board])
                        @show board, "col: ", col, n, calcScore(boards[:, :, board], drawnNumbers[1:n])
                        append!(winningBoards, board)
                    end
                end
            end
        end
    end
end


text:
import SpecialFunctions: airy, airyai, airyaiprime, airyaiprimex, airyaix, airybi, airybiprime, airybiprimex, airybix, airyprime, airyx, besselh, besselhx, besseli, besselix, besselj, besselj0, besselj1, besseljx, besselk, besselkx, bessely, bessely0, bessely1, besselyx, beta, cosint, dawson, digamma, erf, erfc, erfcinv, erfcx, erfi, erfinv, eta, gamma, hankelh1, hankelh1x, hankelh2, hankelh2x, invdigamma, lbeta, lfact, lfactorial, polygamma, sinint, trigamma, zeta # , lgamma, logabsgamma, loggamma

# `airy(k,x)` is deprecated, use `airyai(x)`, `airyaiprime(x)`, `airybi(x)` or `airybiprime(x)` instead.
@primitive airyai(x),dy (dy.*(airyaiprime.(x)))
@primitive airyaiprime(x),dy (dy.*(x .* airyai.(x)))
# airyaiprimex
# airyaix
@primitive airybi(x),dy (dy.*(airybiprime.(x)))
@primitive airybiprime(x),dy (dy.*(x .* airybi.(x)))
# airybiprimex
# airybix
# `airyprime(z::Number)` is deprecated, use `airyaiprime(z)` instead.
# `airyx(k,x)` is deprecated, use `airyaix(x)`, `airyaiprimex(x)`, `airybix(x)` or `airybiprimex(x)` instead.
# besselh
# besselhx
# besseli
# besselix
# besselj
@primitive besselj0(x),dy (dy.*(-(besselj1.(x))))
@primitive besselj1(x),dy (dy.*((besselj0.(x) - besselj.(2,x)) / 2))
# besseljx
# besselk
# besselkx
# bessely
@primitive bessely0(x),dy (dy.*(-(bessely1.(x))))
@primitive bessely1(x),dy (dy.*((bessely0.(x) - bessely.(2,x)) / 2))
# besselyx
# beta
# cosint
@primitive dawson(x),dy,y (dy.*((-2 .* y) .* x .+ 1))
@primitive digamma(x),dy,y (dy.*(trigamma.(x)))
@primitive erf(x),dy,y (dy.*(exp.(-(abs2.(x))) .* convert(eltype(x), 2 / √π)))
@primitive erfc(x),dy,y (dy.*(-(exp.(-(abs2.(x)))) .* convert(eltype(x), 2 / √π)))
@primitive erfcinv(x),dy,y (dy.*(-(exp.(abs2.(y))) .* convert(eltype(x), √π / 2)))
@primitive erfcx(x),dy,y (dy.*((2 .* y) .* x .- convert(eltype(x), 2 / √π)))
@primitive erfi(x),dy,y (dy.*(exp.(abs2.(x)) .* convert(eltype(x), 2 / √π)))
@primitive erfinv(x),dy,y (dy.*(exp.(abs2.(y)) .* convert(eltype(x), √π / 2)))
# eta
@primitive gamma(x),dy,y (dy.*(y .* digamma.(x)))
# hankelh1
# hankelh1x
# hankelh2
# hankelh2x
@primitive invdigamma(x),dy,y (dy.*(1 ./ trigamma.(y)))
# lbeta
# `lfact` is deprecated, use `lfactorial` instead.
# lfactorial
# logabsgamma
# `lgamma` is deprecated, use `(logabsgamma(x))[1]` instead. I use `loggamma` which throws a DomainError if gamma(x) is negative.
# `lgamma` deprecated, using loggamma.
# TODO: remove this once everybody uses loggamma and SpecialFunctions 0.8+
if !isdefined(SpecialFunctions, :loggamma) && isdefined(SpecialFunctions, :lgamma)
    import SpecialFunctions: lgamma
    loggamma(x) = lgamma(x)
end
if isdefined(SpecialFunctions, :loggamma) && !isdefined(SpecialFunctions, :lgamma)
    import SpecialFunctions: loggamma, logabsgamma
    lgamma(x) = loggamma(x)
end
@primitive lgamma(x),dy,y (dy.*(digamma.(x)))
@primitive loggamma(x),dy,y (dy.*(digamma.(x)))
@primitive polygamma(x1,x2),dy,y  nothing  unbroadcast(x2,dy.*polygamma(x1+1,x2))
# sinint
@primitive trigamma(x),dy,y (dy.*(polygamma.(2,x)))
# zeta


text:
#################### Approximate Bayesian Computation ####################

#################### Types and Constructors ####################

mutable struct ABCTune <: SamplerTune
    logf::Function
    params::Vector{Symbol}
    datakeys::Vector{Symbol}
    Tsim::Vector{Vector{Float64}}
    Tobs::Vector{Float64}
    epsilon::Float64
    epsilon_vec::Vector{Float64}
    epsilonprime::Vector{Float64}
    kernel::Function
    dist::Function
    proposal::Function
    maxdraw::Integer
    nsim::Integer
    decay::Real
    randeps::Bool
    obsdata::Function
    simdata::Function
    summarizenodes::Function
    pi_epsilon0::Float64

    function ABCTune()
        new()
    end

    function ABCTune(
        x,
        lf,
        params::Vector{Symbol},
        datakeys::Vector{Symbol},
        epsilon::Float64,
        kernel::Function,
        dist::Function,
        proposal::Function,
        maxdraw::Integer,
        nsim::Integer,
        decay::Real,
        randeps::Bool,
        summarizenodes::Function,
    )
        new(
            lf,
            params,
            datakeys,
            Vector{Vector{Float64}}(),
            Vector{Vector{Float64}}(),
            epsilon,
            Array{Float64}(undef, nsim),
            Array{Float64}(undef, nsim),
            kernel,
            dist,
            proposal,
            maxdraw,
            nsim,
            decay,
            randeps,
            (model, key) -> unlist(model[key]),
            (model, key) -> unlist(model[key], rand(model[key])),
            summarizenodes,
            0.0,
        )
    end

end

const ABCVariate = Sampler{ABCTune, T} where T

#################### Sampler Constructor ####################

"""
    ABC(params::ElementOrVector{Symbol},
    scale::ElementOrVector{T}, summary::Function,
    epsilon::Real; kernel::KernelDensityType=SymUniform,
    dist::Function=(Tsim, Tobs) -> sqrt(sum(abs2, Tsim - Tobs)),
    proposal::SymDistributionType=Normal, maxdraw::Integer=1,
    nsim::Integer=1, decay::Real=1.0, randeps::Bool=false,
    args...)

    Construct a `Sampler` object for ABC sampling. Parameters are assumed to be continuous, but may be constrained or unconstrained.
    Returns a `Sampler{ABCTune}` type object.
    * `params`: stochastic node(s) to be updated with the sampler.  Constrained parameters are mapped to unconstrained space according to transformations defined by the Stochastic `unlist()` function.
    * `scale` : scaling value or vector of the same length as the combined elements of nodes `params` for the `proposal` distribution.  Values are relative to the unconstrained parameter space, where candidate draws are generated.
    * `summary` : function that takes a vector of observed or simulated data and returns a summary statistic or vector of statistics.
    * `epsilon` : target tolerance for determining how similar observed and simulated data summary statistics need to be in order to accept a candidate draw.
    * `kernel` : weighting kernel density of type `Biweight`, `Cosine`, `Epanechnikov`, `Normal`, `SymTriangularDist`, `SymUniform`, or `Triweight` to use in measuring similarity between observed and simulated data summary statistics.  Specified `epsilon` determines the standard deviation of Normal kernels and widths of the others.
    * `dist` : positive function for the kernel density to compute distance between vectors of observed (`Tobs`) and simulated (`Tsim`) data summary statistics (default: Euclidean distance).
    * `proposal` : symmetric distribution of type `Biweight`, `Cosine`, `Epanechnikov`, `Normal`, `SymTriangularDist`, `SymUniform`, or `Triweight` to be centered around current parameter values and used to generate proposal draws.  Specified `scale` determines the standard deviations of Normal proposals and widths of the others.
    * `maxdraw` : maximum number of unaccepted candidates to draw in each call of the sampler.  Draws are generated until one is accepted or the maximum is reached.  Larger values increase acceptance rates at the expense of longer runtimes.
    * `nsim` : number of data sets to simulate in deciding whether to accept a candidate draw.  Larger values lead to closer approximations of the target distribution at the expense of longer runtimes.
    * `decay` : if `0 < decay <= 1`, the rate at which internal tolerances are monotonically decreased from the initial distance between observed and simulated summary statistics toward the maximum of each subsequent distance and `epsilon`; if `decay = 0`, internal tolerances are fixed at `epsilon`.
    * `randeps` : whether to perturb internal tolerances by random exponential variates.
    * `args...` : additional keyword arguments to be passed to the `dist` function.
"""
function ABC(
    params::ElementOrVector{Symbol},
    scale::ElementOrVector{T},
    summary::Function,
    epsilon::Real;
    kernel::KernelDensityType=SymUniform,
    dist::Function=(Tsim, Tobs) -> sqrt(sum(abs2, Tsim - Tobs)),
    proposal::SymDistributionType=Normal,
    maxdraw::Integer=1,
    nsim::Integer=1,
    decay::Real=1.0,
    randeps::Bool=false,
    transform::Bool=true,
    args...,
) where {T <: Real}
    0 <= decay <= 1 || throw(ArgumentError("decay is not in [0, 1]"))

    params = asvec(params)
    kernelpdf = (epsilon, d) -> pdf(kernel(0.0, epsilon), d)
    propose(x) = proposalfun(proposal, scale, x)

    lf(m, args...) = m

    tune = ABCTune(
        Float64[],
        lf,
        params,
        Symbol[],
        epsilon,
        kernelpdf,
        dist,
        propose,
        maxdraw,
        nsim,
        decay,
        randeps,
        summary
    )
    Sampler(Float64[], params, tune, Symbol[], transform)
end

"""
    ABC(params::ElementOrVector{Symbol},
    scale::ElementOrVector{T}, summary::Function,
    epsilon::Real, proposal::Function; kernel::KernelDensityType=SymUniform,
    dist::Function=(Tsim, Tobs) -> sqrt(sum(abs2, Tsim - Tobs)),
    maxdraw::Integer=1,
    nsim::Integer=1, decay::Real=1.0, randeps::Bool=false,
    args...)

    Construct a `Sampler` object for ABC sampling. Parameters are assumed to be continuous, but may be constrained or unconstrained.
    Returns a `Sampler{ABCTune}` type object.
    * `params`: stochastic node(s) to be updated with the sampler.  Constrained parameters are mapped to unconstrained space according to transformations defined by the Stochastic `unlist()` function.
    * `summary` : function that takes a vector of observed or simulated data and returns a summary statistic or vector of statistics.
    * `epsilon` : target tolerance for determining how similar observed and simulated data summary statistics need to be in order to accept a candidate draw.
    * `proposal` : a proposal function which proposes a new value.
    * `kernel` : weighting kernel density of type `Biweight`, `Cosine`, `Epanechnikov`, `Normal`, `SymTriangularDist`, `SymUniform`, or `Triweight` to use in measuring similarity between observed and simulated data summary statistics.  Specified `epsilon` determines the standard deviation of Normal kernels and widths of the others.
    * `dist` : positive function for the kernel density to compute distance between vectors of observed (`Tobs`) and simulated (`Tsim`) data summary statistics (default: Euclidean distance).
    * `maxdraw` : maximum number of unaccepted candidates to draw in each call of the sampler.  Draws are generated until one is accepted or the maximum is reached.  Larger values increase acceptance rates at the expense of longer runtimes.
    * `nsim` : number of data sets to simulate in deciding whether to accept a candidate draw.  Larger values lead to closer approximations of the target distribution at the expense of longer runtimes.
    * `decay` : if `0 < decay <= 1`, the rate at which internal tolerances are monotonically decreased from the initial distance between observed and simulated summary statistics toward the maximum of each subsequent distance and `epsilon`; if `decay = 0`, internal tolerances are fixed at `epsilon`.
    * `randeps` : whether to perturb internal tolerances by random exponential variates.
    * `args...` : additional keyword arguments to be passed to the `dist` function.
"""
function ABC(
    params::ElementOrVector{Symbol},
    summary::Function,
    epsilon::Real,
    proposal::Function;
    kernel::KernelDensityType=SymUniform,
    dist::Function=(Tsim, Tobs) -> sqrt(sum(abs2, Tsim - Tobs)),
    maxdraw::Integer=1,
    nsim::Integer=1,
    decay::Real=1.0,
    randeps::Bool=false,
    transform::Bool=true,
    args...,
)
0 <= decay <= 1 || throw(ArgumentError("decay is not in [0, 1]"))

    params = asvec(params)
    kernelpdf = (epsilon, d) -> pdf(kernel(0.0, epsilon), d)
    propose(x) = proposalfun(proposal, scale, x)

    #lf(m, args...) = m
    lf(args...) = 0
    tune = ABCTune(
        Float64[],
        lf,
        params,
        Symbol[],
        epsilon,
        kernelpdf,
        dist,
        proposal,
        maxdraw,
        nsim,
        decay,
        randeps,
        summary
    )
    Sampler(Float64[], params, tune, Symbol[], transform)

    end

#################### Sampler Functions ####################

function proposalfun(d::SymDistributionType, scale, theta0::T)::T where T<:Real
    theta0 + scale * rand(d(0.0, 1.0), length(theta0))
end

function proposalfun(d::SymDistributionType, scale, theta0::T)::T where T<:AbstractArray
    theta0 .+ scale .* rand(d(0.0, 1.0), length(theta0))
end

function simulate(model::Model, key::Symbol)
    unlist(model[key], rand(model[key]))
end

function sample!(v::ABCVariate, lf::Function; gen::Int, model::Model, kwargs...)

    tune = v.tune

    # TODO: Look for alternative
    # dummy function to get handle on model...
    #model = get_model(0)
    stochastics = keys(model, :stochastic)
    datakeys = intersect(setdiff(v.targets, tune.params), stochastics)
    if gen == 1
        pi_epsilon0 = 0.0
        pi_error0 = 1.0

        v.tune.Tobs = v.tune.summarizenodes(unlist(model, datakeys))
        v.tune.Tsim = Array{Vector{Float64}}(undef, v.tune.nsim)
        v.tune.epsilonprime = Array{Float64}(undef, v.tune.nsim)
        for i = 1:tune.nsim
            ## simulated data summary statistics for current parameter values
            res1 = v.tune.summarizenodes(simulate.(Ref(model), datakeys)...)
            tune.Tsim[i] = res1
            d = tune.dist(tune.Tsim[i], tune.Tobs)

            # starting tolerance
            tune.epsilon_vec[i] = tune.decay > 0 ? max(tune.epsilon, d) : tune.epsilon

            if tune.randeps
                dexp = Exponential(tune.epsilon_vec[i])
                tune.epsilonprime[i] = rand(dexp)
                pi_error0 = pdf(dexp, tune.epsilonprime[i])
            else
                tune.epsilonprime[i] = tune.epsilon_vec[i]
            end

            ## kernel density evaluation
            pi_epsilon0 += tune.kernel(tune.epsilonprime[i], d) * pi_error0
            tune.pi_epsilon0 = pi_epsilon0
        end
    end
    ABC_sample(v, model, datakeys)
end

"""
    Do the ABC sampling. This version uses threads.
"""
function ABC_sample(v::ABCVariate, model::Model, datakeys::Vector{Symbol})

    tune = v.tune

    ## current parameter and density values
    theta0::Array{Float64} = v.value
    logprior0::Float64 = logpdf(model, tune.params, v.transform)

    lo::Base.ReentrantLock = Base.ReentrantLock()
    flag = Threads.Atomic{Bool}(true)

    Threads.@threads for k = 1:tune.maxdraw
        if flag[]
            ## candidate draw and prior density value
            theta1::Array{Float64} = tune.proposal(theta0)
            other_model = deepcopy(model)

            logprior1::Float64 = mapreduce(
                keyval -> logpdf(other_model[keyval[1]], keyval[2], v.transform),
                +,
                relist(other_model, theta1, v.params, v.transform),
            )

            if logprior1 == -Inf
                continue
            end

            relist(other_model, theta1, v.params, v.transform)

            #my_relist(other_model, theta1, v.params, true)

            ## tolerances and kernel density
            pi_epsilon1::Float64 = 0.0
            epsilon1::Vector{Float64} = similar(tune.epsilon_vec)
            epsilonprime1::Vector{Float64} = similar(tune.epsilonprime)
            Tsim1::Vector{Vector{Float64}} = similar(tune.Tsim)
            for i = 1:tune.nsim
                ## simulated data summary statistics for candidate draw
                Tsim1[i] = v.tune.summarizenodes(simulate.(Ref(other_model), datakeys)...)
                d::Float64 = tune.dist(Tsim1[i], tune.Tobs)
                pi_error1::Float64 = 1.0
                ## monotonically decrease tolerance to target
                epsilon1[i] =
                    (1 - tune.decay) * tune.epsilon_vec[i] +
                    tune.decay * max(tune.epsilon, min(d, tune.epsilon_vec[i]))

                if tune.randeps
                    dexp = Exponential(epsilon1[i])
                    epsilonprime1[i] = rand(dexp)

                    pi_error1 = pdf(dexp, tune.epsilonprime1[i])
                else
                    epsilonprime1[i] = epsilon1[i]
                end

                ## kernel density evaluation
                pi_epsilon1 += tune.kernel(epsilonprime1[i], d) * pi_error1
            end

            # accept/reject the candidate draw
            # check if simulation has been accepted in other thread
            if rand() < pi_epsilon1 / tune.pi_epsilon0 * exp(logprior1 - logprior0)
                # lock threads to set values
                if flag[]
                    Threads.atomic_xchg!(flag, false)
                    Base.lock(lo)
                    theta0 = theta1
                    tune.Tsim = Tsim1
                    tune.pi_epsilon0 = pi_epsilon1
                    tune.epsilon_vec = epsilon1
                    tune.epsilonprime = epsilonprime1
                    Base.unlock(lo)
                end

            end
        end
    end
    v[:] = theta0[:]
    #relist(model, theta0, block, true)
end

max_stars_repo_name: erathorn/MCPhylo

max_issues_repo_name: erathorn/MCPhylo

max_forks_repo_name: erathorn/MCPhylo


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

JLLWrappers.@generate_main_file_header("Deno")
JLLWrappers.@generate_main_file("Deno", UUID("04572ae6-984a-583e-9378-9577a1c2574d"))
end  # module Deno_jll