text:
module VersusMacroTools

using Benchmarkplotting
using MLStyle
using Statistics
using Gadfly
using MacroTools: @capture

data = [
    :simple1 => quote
        struct Foo
            x :: Int
            y
        end
    end,

    :simple2 => quote
        struct Foo
            x :: Int
            y :: Array{Int}
        end
    end,

    Symbol("N(fields) = 5") => quote
        struct Foo
            x1 :: Int
            x2 :: Int
            x3 :: Int
            x4 :: Float32
            x5 :: Int
        end
    end,
    Symbol("N(fields) = 10") => quote
        struct Foo
            x1 :: Int
            x2 :: Int
            x3 :: Int
            x4 :: Float32
            x5 :: Int
            x6 :: Int
            x7 :: Int
            x8 :: Int
            x9 :: Int
            x10 :: Int
        end
    end
]

implementations = [
    :MLStyle => function (ex)
        @match ex begin
            Do(fields = []) &&
            quote
            $(::LineNumberNode)
            struct $typename
                $(
                Many(
                    ::LineNumberNode                 ||
                    :($name :: $typ) &&
                        Do(push!(fields, (name, typ)))||
                    (a :: Symbol)    &&
                        Do(push!(fields, (a, Any)))
                )...
                )
            end
            end => (typename, fields)
        end
    end,
    :MacroTools => function(ex)
        @capture(ex, struct T_ fields__ end)
        (T, fields)
    end

]

criterion(x) = (meantime = mean(x.times), allocs = 1 + x.allocs)
df = Benchmarkplotting.bcompare(criterion, data, implementations)

theme = Theme(
    guide_title_position = :left,
    colorkey_swatch_shape = :circle,
    minor_label_font = "Consolas",
    major_label_font = "Consolas",
    point_size=6px
)
report_meantime, df_time = report(:meantime, df, Scale.y_log10, theme)
report_allocs, df_allocs = report(:allocs, df, theme)

open("stats/vs-macrotools(ast).txt", "w") do f
    write(f, string(df))
end

draw(SVG("stats/vs-macrotools(ast)-on-time.svg", 10inch, 4inch), report_meantime);
draw(SVG("stats/vs-macrotools(ast)-on-allocs.svg", 10inch, 4inch), report_allocs);

end


text:
@testset "▶ /render                     " begin
    if G.GP_ENV["HAS_BACKEND"]
        #
        # SAVEFIG
        #
        tmpdir = mktempdir()
        # - empty figure
        f = Figure("fig1", reset=true)
        # - base call
        plot(randn(10))
        @test savefig("figsave1", path=tmpdir) == joinpath(tmpdir, "figsave1.png")
        @test isfile(joinpath(tmpdir, "figsave1.png"))
    end
    #
    # PREVIEWFIG (NOTE, in a test environment this will error because not Juno/IJulia)
    #
    # pf = preview()
    # @test pf.fig.id == f.id
    # @test pf.fname == joinpath(G.GP_ENV["TMP_PATH"], "__PREVIEW__.png")
    # rf = render()
    # @test rf.fig.id == f.id
    # @test rf.fname == joinpath(G.GP_ENV["TMP_PATH"], "__PREVIEW__.png")
    #
    # CLEANUP
    #
    clf()
    plot(randn(10))
    G.assemble_figure(gcf())
    @test isfile(joinpath(G.GP_ENV["TMP_PATH"], gcf().id * ".gle"))
    G.cleanup(gcf())
    @test !isfile(joinpath(G.GP_ENV["TMP_PATH"], gcf().id * ".gle"))
end

ext: jl


text:
using AutomotiveSensors
using AutomotiveSimulator
using Test

@testset "occlusion checker" begin
    obs = ConvexPolygon([VecE2(0.0, 0.0), VecE2(10.0, 0.0), VecE2(10.0, 10.0), VecE2(0.0, 10.0)])
    veh1 = VehicleState(VecSE2(-5,-5,0.0), 0.)
    veh2 = VehicleState(VecSE2(15,15,0.0), 0.0)
    veh3 = VehicleState(VecSE2(10.0, -5, 0.0), 0.)

    @test occlusion_checker(veh1, veh2, obs)
    @test !occlusion_checker(veh1, veh3, obs)
end

@testset "perfect sensor" begin
    sensor = PerfectSensor()
    roadway = gen_straight_roadway(1, 100.0)

    ego = Entity(VehicleState(VecSE2(20.0, 0.0, 0.0), 5.0), VehicleDef(), 1)
    veh1 = Entity(VehicleState(VecSE2(30.0, 0.0, 0.0), 5.0), VehicleDef(), 2)
    veh2 = Entity(VehicleState(VecSE2(10.0, 0.0, 0.0), 5.0), VehicleDef(), 3)
    scene = Scene([ego, veh1, veh2])

    obs = measure(sensor, ego, scene, roadway, ConvexPolygon[])

    @test obs[1] == veh1
    @test obs[2] == veh2
end

@testset "gaussian sensor" begin
    sensor = GaussianSensor(false_positive_rate=0.0, false_negative_rate=0.0)
    roadway = gen_straight_roadway(1, 100.0)

    ego = Entity(VehicleState(VecSE2(20.0, 0.0, 0.0), 5.0), VehicleDef(), 1)
    veh1 = Entity(VehicleState(VecSE2(30.0, 0.0, 0.0), 5.0), VehicleDef(), 2)
    veh2 = Entity(VehicleState(VecSE2(10.0, 0.0, 0.0), 5.0), VehicleDef(), 3)
    scene = Scene([ego, veh1, veh2])

    obs = measure(sensor, ego, scene, roadway, ConvexPolygon[])

    @test length(obs) == 2

    @test AutomotiveSensors.obs_weight(sensor, ego.state, obs[1].state, veh1.state, ConvexPolygon[]) > 0.0
end


text:
abstract type AbstractContinuation <: Function end

struct RationalContinuation{T}
    pmax::T
    kmax::Int
    b::T
end
function Base.getproperty(c::RationalContinuation, f::Symbol)
    f === :length && return length(c)
    return getfield(c, f)
end
function RationalContinuation(xmin, steps::Int)
    pmax = (1 - xmin) / xmin
    kmax = steps
    b = 1 / steps * (1 - sqrt(xmin)) / (1 + sqrt(xmin))
    return RationalContinuation(pmax, kmax, b)
end
function Base.iterate(s::RationalContinuation, k=1)
    (k > s.kmax + 1) && return nothing
    b = s.b
    return s(k), k + 1
end
Base.length(s::RationalContinuation) = s.kmax + 1
function (s::RationalContinuation{T})(k) where {T}
    k = clamp(k, 1, s.kmax + 1)
    b = s.b
    return 4 * b * (k - 1) / (1 - b * (k - 1))^2
end

struct FixedContinuation{T} <: AbstractContinuation
    param::T
    length::Int
end
function Base.iterate(s::FixedContinuation, x=1)
    (x > s.length) && return nothing
    return s.param, x + 1
end
Base.length(s::FixedContinuation) = s.length
(s::FixedContinuation{T})(x...) where {T} = s.param

"""
p(x) = 1/(a + b*ℯ^(-c*x))
"""
struct SigmoidContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
SigmoidContinuation(; kwargs...) = SigmoidContinuation{Float64}(; kwargs...)
function SigmoidContinuation{T}(;
    c::T=T(0.1), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = 1 - T(finish - start) / finish / (ℯ^(-c) - ℯ^(-steps * c)) * ℯ^(-c)
    b = T(finish - start) / finish / (ℯ^(-c) - ℯ^(-steps * c))
    return SigmoidContinuation(a, b, c, steps, min)
end
function Base.iterate(s::SigmoidContinuation, x=1)
    (x > s.length) && return nothing
    return max(1 / (s.a + s.b * ℯ^(-s.c * x)), s.min), x + 1
end
Base.length(s::SigmoidContinuation) = s.length
(s::SigmoidContinuation{T})(x) where {T} = max(1 / (s.a + s.b * ℯ^(-s.c * T(x))), s.min)

"""
p(x) = a(x-b*steps)^3 + c
"""
struct CubicSplineContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
CubicSplineContinuation(; kwargs...) = CubicSplineContinuation{Float64}(; kwargs...)
function CubicSplineContinuation{T}(;
    b::T=T(0.5), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = (finish - start) / (steps^3 * (1 - b)^3 - (1 - b * steps)^3)
    c = start - a * (1 - b * steps)^3
    return CubicSplineContinuation(a, b, c, steps, min)
end

function Base.iterate(s::CubicSplineContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * (x - s.b * s.length)^3 + s.c, s.min), x + 1
end
Base.length(s::CubicSplineContinuation) = s.length
function (s::CubicSplineContinuation{T})(x) where {T}
    return max(s.a * (T(x) - s.b * s.length)^3 + s.c, s.min)
end

"""
p(x) = a*x^b + c
"""
struct PowerContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
PowerContinuation(; kwargs...) = PowerContinuation{Float64}(; kwargs...)
function PowerContinuation{T}(;
    b=T(2), start=T(1), finish=T(5), steps::Int=30, min=-T(Inf)
) where {T}
    a = (finish - start) / max(T(1), steps^b - 1)
    c = start - a
    return PowerContinuation(a, b, c, steps, min)
end
function Base.iterate(s::PowerContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * x^s.b + s.c, s.min), x + 1
end
Base.length(s::PowerContinuation) = s.length
(s::PowerContinuation{T})(x) where {T} = max(s.a * T(x)^s.b + s.c, s.min)

"""
p(x) = a*ℯ^(b*x) + c
"""
struct ExponentialContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
ExponentialContinuation(; kwargs...) = ExponentialContinuation{Float64}(; kwargs...)
function ExponentialContinuation{T}(;
    b::T=T(0.1), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = (finish - start) / max(T(1), ℯ^(b * steps) - ℯ^(b))
    c = start - a * ℯ^(b)
    return ExponentialContinuation(a, b, c, steps, min)
end
function Base.iterate(s::ExponentialContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * ℯ^(s.b * x) + s.c, s.min), x + 1
end
Base.length(s::ExponentialContinuation) = s.length
(s::ExponentialContinuation{T})(x) where {T} = max(s.a * ℯ^(s.b * T(x)) + s.c, s.min)

"""
p(x) = a*log(b*x) + c
"""
struct LogarithmicContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
LogarithmicContinuation(; kwargs...) = LogarithmicContinuation{Float64}(; kwargs...)
function LogarithmicContinuation{T}(;
    b::T=T(1), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = (finish - start) / max(T(1), log(b * steps) - log(b))
    c = start - a * log(b)
    return LogarithmicContinuation(a, b, c, steps, min)
end
function Base.iterate(s::LogarithmicContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * log(s.b * x) + s.c, s.min), x + 1
end
Base.length(s::LogarithmicContinuation) = s.length
(s::LogarithmicContinuation{T})(x) where {T} = max(s.a * log(s.b * T(x)) + s.c, s.min)

Continuation(::PowerPenalty; kwargs...) = Continuation(PowerPenalty; kwargs...)
Continuation(::RationalPenalty; kwargs...) = Continuation(RationalPenalty; kwargs...)
function Continuation(::Type{<:PowerPenalty}; steps, pmax=5.0, kwargs...)
    return PowerContinuation{Float64}(; b=1.0, start=1.0, steps=steps + 1, finish=pmax)
end
function Continuation(::Type{<:RationalPenalty}; steps, xmin=0.001, kwargs...)
    return RationalContinuation(xmin, steps)
end


text:
using LinearAlgebra, SparseArrays
import Base: *, +, -, ^
export
    QUBO,
    +

struct QUBO{T} <: BinaryModel where {T<:Real}
    vars::Vector{AbstractVariable}
    q::SparseMatrixCSC{T,Int}
    offset::T # c_b
    function QUBO{T}(vars::Vector{<:AbstractVariable},
                     q::SparseMatrixCSC{T,Int},
                     offset::T) where {T<:Real} # offset and matrix elements of the same type T
        # TODO: check dimensions and fix q to be proper UpperTriangular
        new{T}(vars, q, offset)
    end
end

###############################################
################# conversions #################
###############################################
# sugar syntax
function QUBO(vars::Vector{<:AbstractVariable},
              q::SparseMatrixCSC{T,Int},
              offset::T) where {T<:Real}
    QUBO{T}(vars, q, offset)
end

# QUBO(1.)
function QUBO(offset::T) where T<:Real
    QUBO(AbstractVariable[], spzeros(0, 0), offset)
end

QUBO() = QUBO(0.)

function QUBO(::Type{T}, x::BitVar) where T<:Real
    q = spzeros(T, 1, 1)
    q[1] = one(T)
    QUBO([x], q, zero(T))
end

QUBO(x::BitVar) = QUBO(Float64, x)

function QUBO(::Type{T}, x::BitVarSubArray) where T<:Real
    @assert all(indexes(x) .> 0) # conversion of non-trivial subarray to QUBO is not specified
    n = prod(size(x))
    ind = mixed_radix_conv(indexes(x), size(x))
    q = spzeros(T, n, n)
    q[ind,ind] = one(T)
    QUBO([x.var], q, zero(T))
end

QUBO(x::BitVarSubArray) = QUBO(Float64, x)

###############################################
################## calls ######################
###############################################

# syntax
# julia> x = BitVar("x")
# julia> a = BitAssignment(x)
# julia> a["x"] = 1
# julia> q = QUBO(x)
# julia> q(a)
function (qubo::QUBO)(a::BitAssignment)
    @assert false "TBD"
end

###############################################
################## utils ######################
###############################################

function enlarge(qubo::QUBO{T}, vars::Vector{<:AbstractVariable}) where T
    vars = sort(collect(unique(vcat(qubo.vars, vars))))
    println(vars)
    n = sum(length.(vars))
    len_qubo = sum(length.(qubo.vars))
    q = spzeros(T, n, n)
    for i = 1:len_qubo
        new_i = findfirst(var -> var.name == qubo.vars[i].name, vars)
        # TODO: fix the loop if we are sure Q is UpperTriangular
        # for j = 1:len_qubo - i + 1
        for j = 1:len_qubo
            new_j = findfirst(var -> var.name == qubo.vars[j].name, vars)
            q[new_i][new_j] = qubo.q[i][j]
        end
    end
    return QUBO(vars, q, qubo.offset)
end

###############################################
################## operations #################
###############################################

function +(qubo::QUBO, x::Real)
    QUBO(qubo.vars, qubo.q, x+qubo.offset)
end

+(x::Real, qubo::QUBO) = qubo + x

function +(qubo1::QUBO, qubo2::QUBO)
    if qubo1.vars == qubo2.vars
        return QUBO(qubo1.vars, qubo1.q + qubo2.q, qubo1.offset + qubo2.offset)
    else
        new_vars = sort(collect(unique(vcat(qubo1.vars, qubo2.vars))))
        qubo1_enl = enlarge(qubo1, qubo2.vars) # change to common dimension
        qubo2_enl = enlarge(qubo2, qubo1.vars)
        return QUBO(new_vars, qubo1_enl.q + qubo2_enl.q, qubo1.offset + qubo2.offset)
    end
end

# TODO implement +, - (binary and unary), *, ^
# remember to export (top of the file)!


text:
using TimeSeries: TimeArray, collapse
using TimeSeries: timestamp, values, colnames
using TimeFrames: TimeFrame, dt_grouper, Begin, End
import Base: sum, getindex
import Statistics: mean
import Statistics: std

abstract type AbstractAction end

struct TimeArrayResampler <: AbstractAction
    ta::TimeArray
    tf::TimeFrame
end

struct GroupBy
    action::AbstractAction
    by::Vector{Symbol}
end

function resample(ta::TimeArray, tf::TimeFrame)
    TimeArrayResampler(ta, tf)
end

function resample(ta::TimeArray, tf)
    resample(ta, TimeFrame(tf))
end

function getindex(action::TimeArrayResampler, by...)
    action = TimeArrayResampler(action.ta[by...], action.tf)
    by = collect(by)
    GroupBy(action, by)
end

function ohlc(grp::GroupBy)
    ohlc(grp.action)
end

function ohlc(resampler::TimeArrayResampler)
    ta = resampler.ta
    f_group = dt_grouper(resampler.tf, eltype(timestamp(ta)))
    ta_o = collapse(ta, f_group, first, first)
    ta_h = collapse(ta, f_group, first, maximum)
    ta_l = collapse(ta, f_group, first, minimum)
    ta_c = collapse(ta, f_group, first, last)
    a_ohlc = hcat(values(ta_o), values(ta_h), values(ta_l), values(ta_c))
    ts = map(f_group, timestamp(ta_o))
    col_ohlc = [:Open, :High, :Low, :Close]
    if length(colnames(ta)) == 1
        _colnames = col_ohlc
    else
        _colnames = Symbol[]
        for col in colnames(ta)
            for col2 in col_ohlc
                new_col = String(col) * "_" * String(col2)
                push!(_colnames, Symbol(new_col))
            end
        end
    end
    ta_ohlc = TimeArray(ts, a_ohlc, _colnames)
end

function mean(grp::GroupBy)
    mean(grp.action)
end

function mean(resampler::TimeArrayResampler)
    f_group = dt_grouper(resampler.tf, eltype(timestamp(resampler.ta)))
    collapse(resampler.ta, f_group, dt -> f_group(first(dt)), mean)
end

function sum(grp::GroupBy)
    sum(grp.action)
end

function sum(resampler::TimeArrayResampler)
    f_group = dt_grouper(resampler.tf, eltype(timestamp(resampler.ta)))
    collapse(resampler.ta, f_group, dt -> f_group(first(dt)), sum)
end

function std(grp::GroupBy)
    std(grp.action)
end

function std(resampler::TimeArrayResampler)
    f_group = dt_grouper(resampler.tf, eltype(timestamp(resampler.ta)))
    collapse(resampler.ta, f_group, dt -> f_group(first(dt)), std)
end


text:
using Test, PolyChaos
import LinearAlgebra: norm, dot

##########################################################
# MEASURE
##########################################################
@testset "Error handling for Measure" begin
    @test_throws AssertionError Measure("jacobi",Dict(:shape_a=>-1.,:shape_b=>2.))
    @test_throws AssertionError Measure("jacobi",Dict(:shape_a=>-1.,:shape_b=>-2.))
    @test_throws AssertionError Measure("jacobi",Dict(:shape_a=>1,:shape_b=>-2))

    @test_throws AssertionError Measure("beta01",Dict(:shape_a=>0.,:shape_b=>2.))
    @test_throws AssertionError Measure("beta01",Dict(:shape_a=>0.,:shape_b=>-2.))
    @test_throws AssertionError Measure("beta01",Dict(:shape_a=>1,:shape_b=>-2))

    @test_throws ErrorException Measure("something")

    @test_throws AssertionError Measure("test",t->t,(3,-1),true)
end

##########################################################
# ORTHOPOLY
##########################################################
@test_throws ErrorException OrthoPoly("idonotexist",3)

quads = [ fejer, fejer2, clenshaw_curtis ]
discs = [ lanczos, stieltjes ]
degs  = 1:5:20
m = Measure("uniform01")
tol = 1e-6
@testset "General constructor for Gaussian OrthoPoly" begin
    for quad in quads, disc in discs, deg in degs
        op = OrthoPoly("uniform01",deg)
        op1 = OrthoPoly("myuniform",deg,m; quadrature=quad, discretization = disc, Nquad = 100*deg)
        @test isapprox(norm(coeffs(op)-coeffs(op1),Inf),0; atol = tol)
        op2 = OrthoPoly("myuniform",deg,t->1,(0,1),true; quadrature=quad, discretization = disc, Nquad = 100*deg)
        @test isapprox(norm(coeffs(op)-coeffs(op2),Inf),0; atol = tol)
    end
end

##########################################################
# QUAD
##########################################################
Ns = 10:5:30
m = Measure("uniform01")
myf(t) = 6t^5
result = 1.
tol = 1e-8
@testset "Last-resort-constructor for Quad using quadgp" begin
    for quad in quads, N in Ns
        op = OrthoPoly("myOP",N,m)
        NW = nw(Quad(N,m;quadrature=quad))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)
        NW = nw(Quad(N,t->1.,(0,1),true;quadrature=quad))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)

        NW = nw(Quad(N,op))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)
        NW = nw(Quad(N,m.w,op.α,op.β,m.dom,m.symmetric,m.pars))
        @test isapprox(dot(myf.(NW[:,1]),NW[:,2]), result; atol=tol)
    end
end

d = 5
ops = [ OrthoPoly("hermite",d),
        OrthoPoly("uniform01",d),
        OrthoPoly("beta01",d,Dict(:shape_a=>3.2,:shape_b=>5.34)) ]
mop = MultiOrthoPoly(ops,d)

is = 0:d
n = length(ops)

x = [1., 2., 3.]

@testset "Evaluation of multivariate basis" begin
    for ind in Iterators.product([collect(0:d) for i = 1:n]...)
        @test isapprox(prod( map(i->evaluate(ind[i],x[i],ops[i]),1:n) ) -  evaluate([ind...],x,mop)[1],0;atol=tol)
    end
end

##########################################################
# TENSOR
##########################################################
deg, M = 5, 1:4
opq = OrthoPolyQ("legendre",deg;Nrec=3*deg)
mop = MultiOrthoPoly([OrthoPolyQ("gaussian",deg;Nrec=3*deg), OrthoPolyQ("logistic",deg;Nrec=3*deg), OrthoPolyQ("uniform01",deg;Nrec=3*deg)], deg)

@test_throws AssertionError Tensor(0,mop)
@testset "Check tensor" begin
    for m in M, op in [opq, mop]
        tensor = Tensor(m,op)
        for ind in Iterators.product([collect(0:deg) for i=1:m]...)
            @test isapprox(tensor.get(collect(ind)), computeSP(collect(ind),op); atol = tol )
        end
    end
end


text:
import Base: display

"""
    BlockingSchedule

Types inheriting from abstract type `BlockingSchedule` define the schedule
according to which blocking is done
"""
abstract type BlockingSchedule end

"""
    struct NoBlocking <: BlockingSchedule end

Regular updates with no blocking
"""
struct NoBlocking <: BlockingSchedule end

"""
    ChequeredBlocking

Struct
```
struct ChequeredBlocking{TP,TWW,TXX} <: BlockingSchedule
    P::TP      # blocking workspace: diffusion law
    WW::TWW    # blocking workspace: accepted Wiener path
    WWᵒ::TWW   # blocking workspace: proposed Wiener path
    XX::TXX    # blocking workspace: accepted diffusion path
    XXᵒ::TXX   # blocking workspace: proposed diffusion path
    Ls         # observation operators for both sets of blocks
    vs         # copied over observations of the process
    Σs         # covariance matrix of the noise for both sets of blocks
    # two sets of knots
    knots::Tuple{Vector{Int64}, Vector{Int64}}
    # two sets of blocks, where each block are indicies of intervals that make up a block
    blocks::Tuple{Vector{Vector{Int64}}, Vector{Vector{Int64}}}
    idx::Int64 # index of set of blocks that are being updated ∈{1,2}
    accpt::Tuple{Vector{Int64}, Vector{Int64}} # tracker for the number of accepted samples
    props::Tuple{Vector{Int64}, Vector{Int64}} # tracker for the number of proposed samples
    # info about the points at which to switch between the systems of ODEs
    changePts::Tuple{Vector{ODEChangePt}, Vector{ODEChangePt}}
end
```
is a blocking schedule in which two sets of blocks are defined in an interlacing
manner. For instance, if all knots consist of {1,2,3,4,5,6,7,8,9}, then set A
will contain {1,3,5,7,9}, whereas set B {2,4,6,8}. These knots will then
uniquely determine the blocks

    ChequeredBlocking(knots::Vector{Int64}, ϵ::Float64, changePt::ODEChangePt,
                      P::TP, WW::TWW, XX::TXX)

Base constructor that takes a set of all `knots` (which it then splits into
two disjoint subsets), the artificial noise parameter `ϵ`, information about
a point at which to switch between ODE solvers for H, Hν, c and L, M⁺, μ, an
object with diffusion laws `P`, container `WW` for the driving Brownian motion,
container `XX` for the sampled path (only the size and type of the latter three
are important, not actual values passed, because they are copied and used
internally as containers).

    ChequeredBlocking(𝔅::ChequeredBlocking{TP̃, TWW, TXX}, P::TP, idx::Int64)

Clone constructor. It creates a new object `ChequeredBlocking` from the old one
`𝔅` by using the same type and size of the containers `WW` and `XX`, but a
changed type of the law `P` (changed from `TP̃` to `TP`). The index of a current
update can also be changed via `idx`.

    ChequeredBlocking()

Empty constructor.
"""
struct ChequeredBlocking{TP,TWW,TXX} <: BlockingSchedule
    P::TP      # blocking workspace: diffusion law
    Pᵒ::TP     # blocking workspace: diffusion law
    WW::TWW    # blocking workspace: accepted Wiener path
    WWᵒ::TWW   # blocking workspace: proposed Wiener path
    XX::TXX    # blocking workspace: accepted diffusion path
    XXᵒ::TXX   # blocking workspace: proposed diffusion path
    Ls         # observation operators for both sets of blocks
    vs         # copied over observations of the process
    Σs         # covariance matrix of the noise for both sets of blocks
    # two sets of knots
    knots::Tuple{Vector{Int64}, Vector{Int64}}
    # two sets of blocks, where each block are indicies of intervals that make up a block
    blocks::Tuple{Vector{Vector{Int64}}, Vector{Vector{Int64}}}
    idx::Int64 # index of set of blocks that are being updated ∈{1,2}
    accpt::Tuple{Vector{Int64}, Vector{Int64}} # tracker for the number of accepted samples
    props::Tuple{Vector{Int64}, Vector{Int64}} # tracker for the number of proposed samples
    # info about the points at which to switch between the systems of ODEs
    changePts::Tuple{Vector{ODEChangePt}, Vector{ODEChangePt}}

    function ChequeredBlocking(knots::Vector{Int64}, ϵ::Float64,
                               changePt::ODEChangePt, P::TP, WW::TWW, XX::TXX
                               ) where {TP,TWW,TXX}
        findKnots(mod, rem) = [k for (i,k) in enumerate(knots) if i % mod == rem]
        knotsA = findKnots(2, 1)
        knotsB = findKnots(2, 0)

        m, d = size(P[end].L)

        findL(knots) = [( k in knots ? SMatrix{d,d}(1.0*I) : p.L) for (k,p) in enumerate(P)]
        LsA = findL(knotsA)
        LsB = findL(knotsB)

        vs = [p.v for p in P]

        findΣ(knots) = [(k in knots ? SMatrix{d,d}(ϵ*I) : p.Σ) for (k,p) in enumerate(P)]
        ΣsA = findΣ(knotsA)
        ΣsB = findΣ(knotsB)

        findChP(knots) = [(k in knots ? deepcopy(changePt) : p.changePt)
                                                    for (k,p) in enumerate(P)]
        chpA = findChP(knotsA)
        chpB = findChP(knotsB)

        """
            knotsToBlocks(knots, idxLast, i)

        Given a list of `knots` fetch the indices of the intervals that together
        make up the `i`-th block. `idxLast` is the index of the last interval.
        """
        function knotsToBlocks(knots, idxLast, i)
            M = length(knots)
            @assert M > 0
            if M >= i > 1
                return (knots[i-1]+1):knots[i]
            elseif i == 1
                return 1:knots[1]
            else
                return (knots[M]+1):idxLast
            end
        end
        blocks = ([collect(knotsToBlocks(knotsA, length(P), i)) for i in 1:length(knotsA)+1],
                  [collect(knotsToBlocks(knotsB, length(P), i)) for i in 1:length(knotsB)+1])

        accpt = (zeros(Int64, length(blocks[1])),
                 zeros(Int64, length(blocks[2])))
        props = (zeros(Int64, length(blocks[1])),
                 zeros(Int64, length(blocks[2])))
        new{TP,TWW,TXX}(deepcopy(P), deepcopy(P), deepcopy(WW), deepcopy(WW),
                        deepcopy(XX), deepcopy(XX), (LsA, LsB), vs, (ΣsA, ΣsB),
                        (knotsA, knotsB), blocks, 1, accpt, props,
                        (chpA, chpB))
    end

    function ChequeredBlocking(𝔅::ChequeredBlocking{TP̃, TWW, TXX}, P::TP, Pᵒ::TP,
                               idx::Int64) where {TP̃,TP,TWW,TXX}
        new{TP,TWW,TXX}(P, Pᵒ, 𝔅.WW, 𝔅.WWᵒ, 𝔅.XX, 𝔅.XXᵒ, 𝔅.Ls, 𝔅.vs, 𝔅.Σs,
                        𝔅.knots, 𝔅.blocks, idx, 𝔅.accpt, 𝔅.props, 𝔅.changePts)
    end

    function ChequeredBlocking()
        new{Nothing, Nothing, Nothing}(nothing, nothing, nothing, nothing,
                                       nothing, nothing, nothing, nothing,
                                       nothing, ([0],[0]),([[0]],[[0]]), 1,
                                       ([0],[0]), ([0],[0]),
                                       ([NoChangePt()],[NoChangePt()])
                                       )
    end
end

"""
    findEndPts(𝔅::ChequeredBlocking, XX, idx)

Determine the observations for the update of the `idx`-th set of blocks. In
particular, on each block with interval indices [a₁,…,aₙ], observations vᵢ with
i∈{1,…,n-1} are made, whereas the last obesrvation is an exactly observed
process Xₜ at the terminal time t of the aₙ-th interval.
"""
function findEndPts(𝔅::ChequeredBlocking, XX, idx)
    [( k in 𝔅.knots[idx] ? X.yy[end] : 𝔅.vs[k]) for (k,X) in enumerate(XX)]
end

"""
    next(𝔅::ChequeredBlocking, XX, θ)

Switch the set of blocks that are being updated. `XX` is the most recently
sampled (accepted) path. `θ` can be used to change parametrisation.
"""
function next(𝔅::ChequeredBlocking, XX, θ)
    newIdx = (𝔅.idx % 2) + 1
    vs = findEndPts(𝔅, XX, newIdx)
    Ls = 𝔅.Ls[newIdx]
    Σs = 𝔅.Σs[newIdx]
    chPts = 𝔅.changePts[newIdx]

    P = [GuidPropBridge(𝔅.P[i], Ls[i], vs[i], Σs[i], chPts[i], θ)
                                            for (i,_) in enumerate(𝔅.P)]
    Pᵒ = [GuidPropBridge(𝔅.Pᵒ[i], Ls[i], vs[i], Σs[i], chPts[i], θ)
                                            for (i,_) in enumerate(𝔅.Pᵒ)]

    ChequeredBlocking(𝔅, P, Pᵒ, newIdx)
end

"""
    display(𝔅::NoBlocking)

Nothing particular to display
"""
function display(𝔅::NoBlocking)
    print("No blocking...\n")
end

"""
    display(𝔅::ChequeredBlocking)

Display the pattern of blocks
"""
function display(𝔅::ChequeredBlocking)
    function printBlocks(knots, idxLast, m)
        M = length(knots)
        getKnot(knots, i) = (M >= i > 0 ? knots[i] : idxLast * (i>0))
        function printRange(from, to)
            for i in from:to
                print("|", getKnot(knots,i-1), "|----",
                      getKnot(knots,i)-getKnot(knots,i-1), "----")
            end
        end

        if m > div(M, 2)
            printRange(1, M+1)
            print("|", idxLast, "|")
        else
            printRange(1,m)
            print("|", getKnot(knots, m) ,"|   ...   ")
            printRange(M+2-m,M+1)
            print("|", idxLast, "|")
        end
        print("  (number of blocks: ", M+1,")")
    end
    print("Chequered Blocking scheme\n",
          "-------------------------\n",
          "Format:\n",
          "block sizes in A: ")
    printBlocks(𝔅.knots[1], length(𝔅.P), 3)
    print("\nblock sizes in B: ")
    printBlocks(𝔅.knots[2], length(𝔅.P), 3)
    print("\n")
end

"""
    registerAccpt(𝔅::BlockingSchedule, i, accepted)

Register whether the block has been accepted
"""
function registerAccpt!(𝔅::BlockingSchedule, i, accepted)
    𝔅.props[𝔅.idx][i] += 1
    𝔅.accpt[𝔅.idx][i] += 1*accepted
end

"""
    displayAcceptance(𝔅::NoBlocking)

Nothing to display
"""
function displayAcceptanceRate(𝔅::NoBlocking) end

"""
    displayAcceptance(𝔅::NoBlocking)

Display acceptance rates
"""
function displayAcceptanceRate(𝔅::BlockingSchedule)
    print("\nAcceptance rates:\n----------------------\n")
    function printAccptRate(accpt, prop)
        m = length(accpt)
        for i in 1:m
            print("b", i, ": ", accpt[i]/prop[i], " | ")
        end
        print("\n- - - - - - - - - - - - - -\n")
    end
    printAccptRate(𝔅.accpt[1], 𝔅.props[1])
    printAccptRate(𝔅.accpt[2], 𝔅.props[2])
end


text:
@testset "Identity map" begin
    m = IdentityMap(2)
    @test outputdim(m) == 2
    @test islinear(m) && isaffine(m)
    @test apply(m, ones(2)) == ones(2)
end

@testset "Constrained identity map" begin
    X = Line([1., -1], 0.) # line x = y
    m = ConstrainedIdentityMap(2, X)
    @test outputdim(m) == 2
    @test islinear(m) && isaffine(m)
    @test stateset(m) == X
    @test apply(m, ones(2)) == ones(2)
end

@testset "Linear map" begin
    A = [1. 1; 1 -1]
    m = LinearMap(A)
    @test outputdim(m) == 2
    @test islinear(m) && isaffine(m)

    # applying the linear map on a vector
    @test apply(m, ones(2)) == [2.0, 0.0]

    # applying the linear map on a set
    X = BallInf(zeros(2), 1.0)
    @test apply(m, X) == LazySets.LinearMap(A, X)
end

@testset "Constrained linear map" begin
    A = [1. 1; 1 -1]
    X = Line([1., -1], 0.) # line x = y
    m = ConstrainedLinearMap(A, X)
    @test outputdim(m) == 2
    @test stateset(m) == X
    @test islinear(m) && isaffine(m)

    # applying the linear map on a vector
    @test apply(m, ones(2)) == [2.0, 0.0]

    # applying the linear map on a set
    X = BallInf(zeros(2), 1.0)
    @test apply(m, X) == LazySets.LinearMap(A, X)
end

@testset "Affine map" begin
    A = [1. 1; 1 -1]
    b = [0.5, 0.5]
    m = AffineMap(A, b)
    @test outputdim(m) == 2

    # applying the affine map on a vector
    @test apply(m, ones(2)) == [2.5, 0.5]

    A = [1. 1.]
    b = [0.5]
    m = AffineMap(A, b)
    @test outputdim(m) == 1
    @test !islinear(m) && isaffine(m)
end

@testset "Constrained affine map" begin
    A = [1. 1; 1 -1]
    b = [0.5, 0.5]
    X = Line([1., -1], 0.) # line x = y
    m = ConstrainedAffineMap(A, b, X)
    @test outputdim(m) == 2
    @test stateset(m) == X

    # applying the affine map on a vector
    @test apply(m, ones(2)) == [2.5, 0.5]

    A = [1. 1.]
    b = [0.5]
    m = ConstrainedAffineMap(A, b, X)
    @test outputdim(m) == 1
    @test !islinear(m) && isaffine(m)
end

@testset "Linear control map" begin
    A = [1. 1; 1 -1]
    B = Matrix([0.5 1.5]')
    m = LinearControlMap(A, B)
    @test outputdim(m) == 2
    @test islinear(m) && isaffine(m)

    # applying the affine map on a vector
    x = ones(2)
    u = [1/2]
    @test apply(m, x, u) == A * x + B * u
end

@testset "Constrained linear control map" begin
    A = [1. 1; 1 -1]
    B = Matrix([0.5 1.5]')
    X = Line([1., -1], 0.) # line x = y
    U = Interval(-1, 1)
    m = ConstrainedLinearControlMap(A, B, X, U)
    @test outputdim(m) == 2
    @test islinear(m) && isaffine(m)

    # applying the affine map on a vector
    x = ones(2) # (it is not checked that x ∈ X in this library)
    u = [1/2]   # same for u ∈ U
    @test stateset(m) == X
    @test apply(m, x, u) == A*x + B*u
end

@testset "Affine control map" begin
    A = [1. 1; 1 -1]
    B = Matrix([0.5 1.5]')
    c = [0.5, 0.5]
    m = AffineControlMap(A, B, c)
    @test outputdim(m) == 2
    @test !islinear(m) && isaffine(m)

    # applying the affine map on a vector
    x = ones(2)
    u = [1/2]
    @test apply(m, x, u) == A*x + B*u + c
end

@testset "Constrained affine control map" begin
    A = [1. 1; 1 -1]
    B = Matrix([0.5 1.5]')
    X = Line([1., -1], 0.) # line x = y
    c = [0.5, 0.5]
    U = Interval(-1, 1)
    m = ConstrainedAffineControlMap(A, B, c, X, U)
    @test outputdim(m) == 2
    @test stateset(m) == X
    @test !islinear(m) && isaffine(m)

    # applying the affine map on a vector
    x = ones(2)
    u = [1/2]
    @test apply(m, x, u) == A * x + B * u  + c
end

@testset "Reset map" begin
    m = ResetMap(10, Dict(9 => 0.))
    @test outputdim(m) == 10

    m = ResetMap(10, 2 => -1., 5 => 1.)
    @test outputdim(m) == 10
    @test !islinear(m) && isaffine(m)

    # applying the affine map on a vector
    x = zeros(10)
    @test apply(m, x) == [0, -1., 0, 0, 1., 0, 0, 0, 0, 0]
end

@testset "Constrained reset map" begin
    X = BallInf(zeros(10), 1.)

    m = ConstrainedResetMap(10, X, Dict(9 => 0.))
    @test outputdim(m) == 10
    @test stateset(m) == X

    m = ConstrainedResetMap(10, X, 2 => -1., 5 => 1.)
    @test outputdim(m) == 10
    @test stateset(m) == X
    @test !islinear(m) && isaffine(m)

    # applying the affine map on a vector
    x = zeros(10)
    @test apply(m, x) == [0, -1., 0, 0, 1., 0, 0, 0, 0, 0]
end


text:
#######################################################################
# Diagonalize methods
#   (All but LinearAlgebraPackage `@require` some package to be loaded)
#######################################################################
abstract type AbstractDiagonalizeMethod end

struct Diagonalizer{S<:AbstractDiagonalizeMethod,C}
    method::S
    codiag::C
    minprojection::Float64
end

diagonalizer(h::Hamiltonian, mesh::Mesh, method, minprojection) =
    Diagonalizer(method, codiagonalizer(h, mesh), minprojection)

## Diagonalize methods ##

function defaultmethod(h::Hamiltonian)
    if eltype(h) <: Number
        # method = issparse(h) ? ArpackPackage() : LinearAlgebraPackage()
        method = LinearAlgebraPackage()
    else
        # method = KrylovKitPackage()
        throw(ArgumentError("Methods for generic Hamiltonian eltypes not yet implemented. Consider using `flatten` on your Hamiltonian."))
    end
    return method
end

checkloaded(package::Symbol) = isdefined(Main, package) ||
    throw(ArgumentError("Package $package not loaded, need to be `using $package`."))

## LinearAlgebra ##
struct LinearAlgebraPackage{K<:NamedTuple} <: AbstractDiagonalizeMethod
    kw::K
end

LinearAlgebraPackage(; kw...) = LinearAlgebraPackage(values(kw))

function diagonalize(matrix, method::LinearAlgebraPackage)
    ϵ, ψ = eigen!(matrix; (method.kw)...)
    return ϵ, ψ
end

similarmatrix(h, ::LinearAlgebraPackage) = Matrix(similarmatrix(h))

## Arpack ##
struct ArpackPackage{K<:NamedTuple} <: AbstractDiagonalizeMethod
    kw::K
end

ArpackPackage(; kw...) = (checkloaded(:Arpack); ArpackPackage(values(kw)))

function diagonalize(matrix, method::ArpackPackage)
    ϵ, ψ = Main.Arpack.eigs(matrix; (method.kw)...)
    return ϵ, ψ
end

similarmatrix(h, ::ArpackPackage) = similarmatrix(h)

## IterativeSolvers ##

struct KrylovKitPackage{K<:NamedTuple} <: AbstractDiagonalizeMethod
    kw::K
end

KrylovKitPackage(; kw...) = (checkloaded(:KrylovKit); KrylovKitPackage(values(kw)))

function diagonalize(matrix::AbstractMatrix{M}, method::KrylovKitPackage) where {M}
    ishermitian(matrix) || throw(ArgumentError("Only Hermitian matrices supported with KrylovKitPackage for the moment"))
    origin = get(method.kw, :origin, 0.0)
    howmany = get(method.kw, :howmany, 1)
    kw´ = Base.structdiff(method.kw, NamedTuple{(:origin,:howmany)}) # Remove origin option
    lmap = shiftandinvert(matrix, origin)
    T = eltypevec(matrix)
    n = size(matrix, 2)
    x0 = rand(T, n)
    ϵ, ψ, _ = Main.KrylovKit.eigsolve(x -> lmap * x, x0, howmany, :LM; kw´...)

    ϵ´ = invertandshift(ϵ, origin)
    resize!(ϵ´, howmany)

    dimh = size(matrix, 2)
    ψ´ = Matrix{T}(undef, dimh, howmany)
    for i in 1:howmany
        copyslice!(ψ´, CartesianIndices((1:dimh, i:i)), ψ[i], CartesianIndices((1:dimh,)))
    end

    return ϵ´, ψ´
end

similarmatrix(h, ::KrylovKitPackage) = similarmatrix(h)

#######################################################################
# shift and invert methods
#######################################################################

function shiftandinvert(matrix::AbstractMatrix{Tv}, origin) where {Tv}
    cols, rows = size(matrix)
    # Shift away from real axis to avoid pivot point error in factorize
    matrix´ = diagshift!(parent(matrix), origin + im)
    F = factorize(matrix´)
    lmap = LinearMap{Tv}((x, y) -> ldiv!(x, F, y), cols, rows,
                         ismutating = true, ishermitian = false)
    return lmap
end

function diagshift!(matrix::AbstractMatrix, origin)
    matrix´ = parent(matrix)
    vals = nonzeros(matrix´)
    rowval = rowvals(matrix´)
    for col in 1:size(matrix, 2)
        found_diagonal = false
        for ptr in nzrange(matrix´, col)
            if col == rowval[ptr]
                found_diagonal = true
                vals[ptr] -= origin * I  # To respect non-scalar eltypes
                break
            end
        end
        found_diagonal || throw(error("Sparse work matrix must include the diagonal. Possible bug in `similarmatrix`."))
    end
    return matrix
end

function invertandshift(ϵ::Vector{T}, origin) where {T}
    ϵ´ = similar(ϵ, real(T))
    ϵ´ .= real(inv.(ϵ) .+ (origin + im))  # Caution: we assume a real spectrum
    return ϵ´
end

#######################################################################
# resolve_degeneracies
#######################################################################
# Tries to make states continuous at crossings. Here, ϵ needs to be sorted
function resolve_degeneracies!(ϵ, ψ, ϕs, matrix, codiag)
    issorted(ϵ, by = real) || sorteigs!(codiag.perm, ϵ, ψ)
    hasapproxruns(ϵ, codiag.degtol) || return ϵ, ψ
    ranges, ranges´ = codiag.rangesA, codiag.rangesB
    resize!(ranges, 0)
    pushapproxruns!(ranges, ϵ, 0, codiag.degtol) # 0 is an offset
    for n in 1:num_codiag_matrices(codiag)
        v = codiag_matrix!(matrix, n, codiag, ϕs)
        resize!(ranges´, 0)
        for (i, r) in enumerate(ranges)
            subspace = view(ψ, :, r)
            vsubspace = subspace' * v * subspace
            veigen = eigen!(Hermitian(vsubspace))
            if hasapproxruns(veigen.values, codiag.degtol)
                roffset = minimum(r) - 1 # Range offset within the ϵ vector
                pushapproxruns!(ranges´, veigen.values, roffset, codiag.degtol)
            end
            subspace .= subspace * veigen.vectors
        end
        ranges, ranges´ = ranges´, ranges
        isempty(ranges) && break
    end
    return ψ
end

# Could perhaps be better/faster using a generalized CoSort
function sorteigs!(perm, ϵ::Vector, ψ::Matrix)
    resize!(perm, length(ϵ))
    p = sortperm!(perm, ϵ, by = real)
    # permute!(ϵ, p)
    sort!(ϵ, by = real)
    Base.permutecols!!(ψ, p)
    return ϵ, ψ
end

#######################################################################
# Codiagonalizer
#######################################################################
codiagonalizer(h, mesh) = VelocityCodiagonalizer(h, meshdelta(mesh))

meshdelta(mesh::Mesh{<:Any,T}) where {T} = T(0.1) * first(minmax_edge_length(mesh))

## VelocityCodiagonalizer
## Uses velocity operators along different directions. If not enough, use finite differences
struct VelocityCodiagonalizer{S,T,H<:Hamiltonian}
    h::H
    directions::Vector{S}
    degtol::T
    delta::T                        # Finite differences
    rangesA::Vector{UnitRange{Int}} # Prealloc buffer for degeneray ranges
    rangesB::Vector{UnitRange{Int}} # Prealloc buffer for degeneray ranges
    perm::Vector{Int}               # Prealloc for sortperm!
end

function VelocityCodiagonalizer(h::Hamiltonian{<:Any,L}, delta;
                                direlements = -0:1, onlypositive = true, kw...) where {L}
    directions = vec(SVector{L,Int}.(Iterators.product(ntuple(_ -> direlements, Val(L))...)))
    onlypositive && filter!(ispositive, directions)
    unique!(normalize, directions)
    sort!(directions, by = norm, rev = false)
    degtol = sqrt(eps(realtype(h)))
    delta´ = delta === missing ? degtol : delta
    VelocityCodiagonalizer(h, directions, degtol, delta´, UnitRange{Int}[], UnitRange{Int}[], Int[])
end

num_codiag_matrices(codiag) = 2 * length(codiag.directions) + 1
function codiag_matrix!(matrix, n, codiag, ϕs)
    ndirs = length(codiag.directions)
    if n <= ndirs
        bloch!(matrix, codiag.h, ϕs, dn -> im * codiag.directions[n]' * dn)
    elseif n <= 2ndirs # resort to finite differences
        bloch!(matrix, codiag.h, ϕs + codiag.delta * codiag.directions[n - ndirs])
    else # use a fixed random matrix
        randomfill!(matrix)
    end
    return matrix
end

function randomfill!(matrix::SparseMatrixCSC, seed = 1234)
    Random.seed!(seed)
    rand!(nonzeros(matrix))  ## CAREFUL: this will be non-hermitian
    return matrix
end

function randomfill!(matrix::AbstractArray{T}, seed = 1234) where {T}
    Random.seed!(seed)
    fill!(matrix, zero(T))
    for i in 1:minimum(size(matrix))
        r = rand(T)
        @inbounds matrix[i, i] = r + r'
    end
    return matrix
end

max_issues_repo_name: pablosanjose/QBox

max_forks_repo_name: pablosanjose/QBox


text:
function LRect(parent::Scene; bbox = nothing, kwargs...)

    default_attrs = default_attributes(LRect, parent).attributes
    theme_attrs = subtheme(parent, :LRect)
    attrs = merge!(merge!(Attributes(kwargs), theme_attrs), default_attrs)

    @extract attrs (color, visible, valign, halign, padding, strokewidth,
        strokevisible, strokecolor)

    layoutobservables = LayoutObservables{LRect}(attrs.width, attrs.height, attrs.tellwidth, attrs.tellheight,
        halign, valign, attrs.alignmode; suggestedbbox = bbox)

    strokecolor_with_visibility = lift(strokecolor, strokevisible) do col, vis
        vis ? col : RGBAf0(0, 0, 0, 0)
    end

    ibbox = @lift(round_to_IRect2D($(layoutobservables.computedbbox)))

    r = poly!(parent, ibbox, color = color, visible = visible, raw = true,
        strokecolor = strokecolor_with_visibility, strokewidth = strokewidth)[end]

    # trigger bbox
    layoutobservables.suggestedbbox[] = layoutobservables.suggestedbbox[]

    LRect(parent, layoutobservables, r, attrs)
end

function Base.delete!(lr::LRect)
    GridLayoutBase.disconnect_layoutobservables!(lr.layoutobservables.gridcontent)
    GridLayoutBase.remove_from_gridlayout!(lr.layoutobservables.gridcontent)
    empty!(lr.layoutobservables.suggestedbbox.listeners)
    empty!(lr.layoutobservables.computedbbox.listeners)
    empty!(lr.layoutobservables.reportedsize.listeners)
    empty!(lr.layoutobservables.autosize.listeners)
    empty!(lr.layoutobservables.protrusions.listeners)

    # remove the plot object from the scene
    delete!(lr.parent, lr.rect)
end


text:
using Pkg
cd(@__DIR__)
Pkg.activate(@__DIR__)
using StaticArrays
using LinearAlgebra
using SparseArrays
using Random
Random.seed!(1)
const d = 2
const Point = SArray{Tuple{d},Float32,1,d}
using CairoMakie, AbstractPlotting
# using Makie
T = (0:0.1f0:200)[2:end]
n = 50
σ = 0.1f0
μ = 0.2f0
a = 0.8f0
b = 0.15f0
B = sparse((-μ - a + b)*I, n, n)
B1 = sparse(-1f0*I, n, n)
B2 = sparse(1f0*I, n, n)

for i in 1:n
    k1 = rand(1:n-1)
    k1 += (i <= k1)
    k2 = rand(1:n-2)
    k2 += (min(i, k1) <= k2)
    k2 += (max(i, k1) <= k2)
    @assert i ≠ k1
    @assert i ≠ k2
    @assert k1 ≠ k2

    B[i, k1] = a
    B[i, k2] = -b

    B1[i,k1] = 1
    B2[i, k2] = -1

end

x0 = randn(Point, n)
x = copy(x0)
xs = [x]
ts = [0.0]
t = 0.0
r = 1.
nlz(x) = x/norm(x)
for i in eachindex(T)
    global t, x, xs
    local dt = T[i] - t
    @assert dt > 0
    x = x + (B*x)*dt + σ*sqrt(dt)*randn(Point, n)
    t = T[i]
    push!(ts, t)
    push!(xs, copy(x))
end

flatten = Base.Iterators.flatten
function posterior(ts, xs, n, μ, σ)
    v = [c for c in CartesianIndices((n,n)) if c[1] != c[2]]
    a1 = 2
    a2 = 1
    i = 1
    x = xs[i]
    s = (ts[i+1] - ts[i])*σ^(-2.0) # dt/σ^2
    dx = xs[i+1] - (xs[i] + xs[i]*(ts[i+1] - ts[i])*(-μ)) # dxt + xt*dt*-μ

    z = [(x[c[a2]] - x[c[a1]])'*(σ^(-2.0)*dx[c[a1]]) for c in v]
    Γ = [(c1[a1]==c2[a1])*s*(x[c1[a2]] - x[c1[a1]])'*(x[c2[a2]] - x[c2[a1]]) for c1 in v, c2 in v]

    for i in 2:length(ts)-1

        x = xs[i]
        s = (ts[i+1] - ts[i])*σ^(-2.0)

        for i1 in eachindex(v) # same as z += [(x[c[a2]] - x[c[a1]])'*(σ^(-2.0)*dx[c[a1]]) for c in v]
            c = v[i1]
            dx = xs[i+1][c[a1]] - (xs[i][c[a1]] + xs[i][c[a1]]*(ts[i+1] - ts[i])*(-μ))
            z[i1] += (x[c[a2]] - x[c[a1]])'*(σ^(-2.0)*dx)
        end
        for i1 in eachindex(v) # same as Γ += [(c1[a1]==c2[a1])*s*(x[c1[a2]] - x[c1[a1]])'*(x[c2[a2]] - x[c2[a1]]) for c1 in v, c2 in v]
            c1 = v[i1]
            i2 = i1
            while i2 <= length(v)
                c2 = v[i2]
                c1[a1]!=c2[a1] && break
                uu = s*(x[c1[a2]] - x[c1[a1]])'*(x[c2[a2]] - x[c2[a1]])
                Γ[i1, i2] += uu
                if i1 != i2
                    Γ[i2, i1] += uu
                end
                i2 += 1
            end
        end
    end
    Γ, z
end

Γ, z = posterior(ts, xs, n, μ, σ)
size(Γ)
50*49
γ0 = 0.02
post = cholesky(Hermitian(Γ + γ0*I))\z
(Γ + γ0*I)*post - z

function img(x, n)
    B2 = zeros(n, n)
    v = [c for c in CartesianIndices((n,n)) if c[1] != c[2]]
    for (p, c) in zip(x, v)
        B2[c] = p
    end
    B2
end
B2a = img(post, n)
p1 = heatmap([B2a' Matrix(B - Diagonal(B))], colormap=:berlin, colorrange = (-1, 1))

[B2a' Matrix(B - Diagonal(B))]

using StatsBase, Statistics

Γ0 = sparse(Γ)
μ = cholesky(Hermitian(Γ + γ0*I))\z

using ZigZagBoomerang
# ϕ(x', Γ) = 0.5*x'*Γ*x - z'*x # not needed

# Define ∇ϕ(x, i, Γ) giving the partial derivative of ϕ(x) with respect to x[i]
∇ϕ(x, i, Γ, z) = ZigZagBoomerang.idot(Γ, i, x) - z[i] # more efficient that dot(Γ[:, i], x)
# Random initial values
t0 = 0.0
m = length(z)
x0 = μ
θ0 = rand([-1.0,1.0], m)

# Rejection bounds
c = [1e-4 for i in 1:m]

# Define ZigZag
Z = ZigZag(Γ0, μ)

# Run sparse ZigZag for T time units and collect trajectory
T0 = 20.
@time trace, (tT, xT, θT), (acc, num) = spdmp(∇ϕ, t0, x0, θ0, T0, c, Z, Γ0, z; structured = true, adapt = true)
@time traj0 = collect(discretize(trace, 0.1))

κ = 0.05*ones(length(x0))
@time trace, (tT, xT, θT), (acc, num) = ZigZagBoomerang.sspdmp(∇ϕ, t0, x0, θ0, T0, c, Z, κ, Γ0, z; structured = true, strong_upperbounds = false, adapt = true)
@time traj = collect(discretize(trace, 0.1))

using ColorSchemes

A = Node(img(traj[1].second, n))

p2 = heatmap(A, colormap=:berlin, colorrange = (-1, 1))
display(p2)
for (t, x) in traj
    A[] = img(x, n)
    yield()
    sleep(0.01)
end

using Statistics

xhat = [median(getindex.(last.(traj), i)) for i in eachindex(μ)]
xhat0 = [median(getindex.(last.(traj0), i)) for i in eachindex(μ)]

p3 = heatmap([img(xhat, n)'  Matrix(B - Diagonal(B))], colormap=:berlin, colorrange = (-1, 1))

norm(img(xhat, n)' - Matrix(B - Diagonal(B)))
norm(img(μ, n)' - Matrix(B - Diagonal(B)))

using GeometryBasics

fi0 = linesegments(repeat(1:n*n, inner = 2), vec([(vec(Matrix(B)-Diagonal(B))) vec(img(post,n)')]'), linewidth = 0.3, resolution = (1200,900))
scatter!(1:n*n, vec(img(post,n)'), color=(sqrt∘abs).(vec(Matrix(B)-Diagonal(B))), colormap=:berlin, markersize=7,
            strokewidth = 0.5,marker = map(x -> x != 0 ? GeometryBasics.HyperSphere{2} : :x, vec(Matrix(B)-Diagonal(B))))
fi0

# fi0 = scatter(1:n*n, vec(img(xhat0,n)'), color=(sqrt∘abs).(vec(Matrix(B)-Diagonal(B))), colormap=:berlin, markersize=5, alpha=0.3)
fi = linesegments(repeat(1:n*n, inner = 2), vec([(vec(Matrix(B)-Diagonal(B))) vec(img(xhat,n)')]'), linewidth = 0.3, resolution = (1200,900))
scatter!(1:n*n, vec(img(xhat,n)'), color=(sqrt∘abs).(vec(Matrix(B)-Diagonal(B))), colormap=:berlin, markersize=7,
            strokewidth = 0.5,marker = map(x -> x != 0 ? GeometryBasics.HyperSphere{2} : :x, vec(Matrix(B)-Diagonal(B))))
fi
save("figures/sparseinteraction.png", fi0)
save("figures/sparseinteractionsticky.png", fi)

# show sum of errors
sum(abs.(vec(img(xhat,n)') .- vec(Matrix(B)-Diagonal(B))))
sum(abs.(vec(img(xhat0,n)') .- vec(Matrix(B)-Diagonal(B))))

function confusion(xhat)
    Dict(
    "false negative" => sum((vec(img(xhat,n)') .== 0) .& (vec(Matrix(B)-Diagonal(B)) .!= 0)),
    "true positive" => sum((vec(img(xhat,n)') .!= 0) .& (vec(Matrix(B)-Diagonal(B)) .!= 0)),
    "true negative" => sum((vec(img(xhat,n)') .== 0) .& (vec(Matrix(B)-Diagonal(B)) .== 0)),
    "false positive" => sum((vec(img(xhat,n)') .!= 0) .& (vec(Matrix(B)-Diagonal(B)) .== 0)),
    )
end

confusion(xhat)
confusion(abs.(xhat0) .> 0.1265)


text:
"""
    non_backtracking_matrix(g)

Return a non-backtracking matrix `B` and an edgemap storing the oriented
edges' positions in `B`.

Given two arcs ``A_{i j}` and `A_{k l}` in `g`, the
non-backtraking matrix ``B`` is defined as

``B_{A_{i j}, A_{k l}} = δ_{j k} * (1 - δ_{i l})``
"""
function non_backtracking_matrix(g::AbstractGraph)
    # idedgemap = Dict{Int,Edge}()
    edgeidmap = Dict{Edge,Int}()
    m = 0
    for e in edges(g)
        m += 1
        edgeidmap[e] = m
    end

    if !is_directed(g)
        for e in edges(g)
            m += 1
            edgeidmap[reverse(e)] = m
        end
    end

    B = zeros(Float64, m, m)

    for (e, u) in edgeidmap
        i, j = src(e), dst(e)
        for k in inneighbors(g, i)
            k == j && continue
            v = edgeidmap[Edge(k, i)]
            B[v, u] = 1
        end
    end

    return B, edgeidmap
end

"""
    Nonbacktracking{G}

A compact representation of the nonbacktracking operator.

The Nonbacktracking operator can be used for community detection.
This representation is compact in that it uses only ne(g) additional storage
and provides an implicit representation of the matrix B_g defined below.

Given two arcs ``A_{i j}` and `A_{k l}` in `g`, the
non-backtraking matrix ``B`` is defined as

``B_{A_{i j}, A_{k l}} = δ_{j k} * (1 - δ_{i l})``

This type is in the style of GraphMatrices.jl and supports the necessary operations
for computed eigenvectors and conducting linear solves.

Additionally the `contract!(vertexspace, nbt, edgespace)` method takes vectors
represented in the domain of ``B`` and represents them in the domain of the
adjacency matrix of `g`.
"""
struct Nonbacktracking{G <: AbstractGraph}
    g::G
    edgeidmap::Dict{Edge,Int}
    m::Int
end

function Nonbacktracking(g::AbstractGraph)
    edgeidmap = Dict{Edge,Int}()
    m = 0
    for e in edges(g)
        m += 1
        edgeidmap[e] = m
    end
    if !is_directed(g)
        for e in edges(g)
            m += 1
            edgeidmap[reverse(e)] = m
        end
    end
    return Nonbacktracking(g, edgeidmap, m)
end

size(nbt::Nonbacktracking) = (nbt.m, nbt.m)
eltype(nbt::Nonbacktracking) = Float64
issymmetric(nbt::Nonbacktracking) = false

function *(nbt::Nonbacktracking, x::Vector{T}) where T <: Number
    length(x) == nbt.m || error("dimension mismatch")
    y = zeros(T, length(x))
    for (e, u) in nbt.edgeidmap
        i, j = src(e), dst(e)
        for k in inneighbors(nbt.g, i)
            k == j && continue
            v = nbt.edgeidmap[Edge(k, i)]
            y[v] += x[u]
        end
    end
    return y
end
function mul!(C, nbt::Nonbacktracking, B)
    # computs C = A * B
    for i in 1:size(B, 2)
        C[:, i] = nbt * B[:, i]
    end
    return C
end

function coo_sparse(nbt::Nonbacktracking)
    m = nbt.m
    #= I,J = zeros(Int, m), zeros(Int, m) =#
    I, J = zeros(Int, 0), zeros(Int, 0)
    for (e, u) in nbt.edgeidmap
        i, j = src(e), dst(e)
        for k in inneighbors(nbt.g, i)
            k == j && continue
            v = nbt.edgeidmap[Edge(k, i)]
            #= J[u] = v =#
            #= I[u] = u =#
            push!(I, v)
            push!(J, u)
        end
    end
    return I, J, 1.0
end

sparse(nbt::Nonbacktracking) = sparse(coo_sparse(nbt)..., nbt.m, nbt.m)

function *(nbt::Nonbacktracking, x::AbstractMatrix)
    y = zero(x)
    for i in 1:nbt.m
        y[:, i] = nbt * x[:, i]
    end
    return y
end

"""
    contract!(vertexspace, nbt, edgespace)

The mutating version of `contract(nbt, edgespace)`. Modifies `vertexspace`.
"""
function contract!(vertexspace::Vector, nbt::Nonbacktracking, edgespace::Vector)
    T = eltype(nbt.g)
    for i = one(T):nv(nbt.g), j in neighbors(nbt.g, i)
        u = nbt.edgeidmap[i > j ? Edge(j, i) : Edge(i, j)]
        vertexspace[i] += edgespace[u]
    end
end

#TODO: documentation needs work. sbromberger 20170326
"""
    contract(nbt, edgespace)

Integrate out the edges by summing over the edges incident to each vertex.
"""
function contract(nbt::Nonbacktracking, edgespace::Vector)
    y = zeros(eltype(edgespace), nv(nbt.g))
    contract!(y, nbt, edgespace)
    return y
end


text:
"""
	FluPop(f::Union{AbstractString,IO}, sequencetype::Symbol, headerfields;
	flulineage=missing, segment=missing, strainfilters = [!is_flu_outlier(flulineage), BioTools.hasdate], separator = '|')

Call `readfastastrains` to read `f`. Store the result in a `FluPop` object.
"""
function FluPop(f::Union{AbstractString,IO},
	sequencetype,
	headerfields;
	flulineage=missing,
	segment=missing,
	strainfilters = flu_usual_filters(flulineage),
	separator = '|',
	ignore_read_errors=false)

	strains = readfastastrains(f, sequencetype, headerfields, separator = separator, strainfilters = strainfilters, ignore_read_errors=ignore_read_errors)
	return FluPop(strains = Dict{String, eltype(strains)}(String(x[:strain])=>x for x in strains))
end
"""
	AAFluPop(f::Union{AbstractString,IO}, headerfields; [kwargs...])

Call `FluPop(f, :aa, headerfields, [kwargs...])
"""
function AAFluPop(f::Union{AbstractString,IO}, headerfields;
	flulineage=missing,
	segment=missing,
	strainfilters = flu_usual_filters(flulineage),
	separator = '|')
	return FluPop(f, :aa, headerfields, flulineage=flulineage, segment=segment, strainfilters=strainfilters, separator=separator)
end

"""
	read_mutations!(tree::Tree, mutfile::String)

Read mutations from a JSON file outputted by augur. Store them into `tree`.

## Note
This is in the `Flu` module because it is explicitely designed for use of augur on flu. In particular, it uses the `gene_positions` global.
"""
function read_mutations!(tree::Tree, mutfile::String; type=:aa, lineage="h3n2", segment="ha")
	if type == :aa
		read_mutations_aa!(tree, mutfile, lineage, segment)
	else
		@warn "Only implemented for `type=:aa`"
	end
end

"""
"""
function read_mutations_aa!(tree, mutfile::String, lineage, segment)
	muts = JSON.Parser.parsefile(mutfile)["nodes"]
	for (label, mut) in muts
		if haskey(tree.lnodes, label)
			tmp = Array{TreeTools.Mutation}(undef, 0)
			for (gene, pos) in gene_offsets[lineage, segment]
				if haskey(mut["aa_muts"], gene)
					for m in mut["aa_muts"][gene]
						push!(tmp, _parse_aa_mut(m))
						tmp[end].i = tmp[end].i + pos[1] # Offset for different genes
					end
				end
			end
			tree.lnodes[label].data.mutations = tmp
		end
	end
end
"""
Parse a string of format `XiY` into a `TreeTools.Mutation` object with fields `i`, `X` and `Y`.
"""
function _parse_aa_mut(m::String)
	if length(m) < 3
		error("Can't parse mutation string of length smaller than 3")
	end
	i = parse(Int64, m[2:end-1])
	old = AminoAcid(m[1])
	new = AminoAcid(m[end])
	out =  TreeTools.Mutation(i, old, new)
	return out
end

max_stars_repo_name: PierreBarrat/FluPredictibility

max_issues_repo_name: PierreBarrat/FluPredictibility

max_forks_repo_name: PierreBarrat/FluPredictibility


text:
include("types.jl")
include("define_env.jl")
include("simulation.jl")


text:
abstract type NeuralDELayer <: Function end

"""
Constructs a neural ODE with the gradients computed using the adjoint
method[1]. At a high level this corresponds to solving the forward
differential equation, using a second differential equation that propagates
the derivatives of the loss  backwards in time.
This first solves the continuous time problem, and then discretizes following
the rules specified by the numerical ODE solver.
On the other hand, the 'neural_ode_rd' first disretizes the solution and then
computes the adjoint using automatic differentiation.

Ref
[1]L. S. Pontryagin, Mathematical Theory of Optimal Processes. CRC Press, 1987.

Arguments
≡≡≡≡≡≡≡≡
model::Chain defines the ̇x
x<:AbstractArray initial value x(t₀)
args arguments passed to ODESolve
kwargs key word arguments passed to ODESolve; accepts an additional key
    :callback_adj in addition to :callback. The Callback :callback_adj
    passes a separate callback to the adjoint solver.

"""
struct NeuralODE{M,P,RE,T,S,A,K} <: NeuralDELayer
    model::M
    p::P
    re::RE
    tspan::T
    solver::S
    args::A
    kwargs::K

    function NeuralODE(model,tspan,solver=nothing,args...;kwargs...)
        p,re = Flux.destructure(model)
        new{typeof(model),typeof(p),typeof(re),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(
            model,p,re,tspan,solver,args,kwargs)
    end

    function NeuralODE(model::FastChain,tspan,solver=nothing,args...;kwargs...)
        p = initial_params(model)
        re = nothing
        new{typeof(model),typeof(p),typeof(re),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(
            model,p,re,tspan,solver,args,kwargs)
    end
end

Flux.@functor NeuralODE

function (n::NeuralODE)(x,p=n.p)
    dudt_(u,p,t) = n.re(p)(u)
    prob = ODEProblem{false}(dudt_,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;n.kwargs...)
end

function (n::NeuralODE{M})(x,p=n.p) where {M<:FastChain}
    dudt_(u,p,t) = n.model(u,p)
    prob = ODEProblem{false}(dudt_,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;
                                sensealg=InterpolatingAdjoint(
                                autojacvec=DiffEqSensitivity.ReverseDiffVJP(true)),
                                n.kwargs...)
end

struct NeuralDSDE{M,P,RE,M2,RE2,T,S,A,K} <: NeuralDELayer
    p::P
    len::Int
    model1::M
    re1::RE
    model2::M2
    re2::RE2
    tspan::T
    solver::S
    args::A
    kwargs::K
    function NeuralDSDE(model1,model2,tspan,solver=nothing,args...;kwargs...)
        p1,re1 = Flux.destructure(model1)
        p2,re2 = Flux.destructure(model2)
        p = [p1;p2]
        new{typeof(model1),typeof(p),typeof(re1),typeof(model2),typeof(re2),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(p,
            length(p1),model1,re1,model2,re2,tspan,solver,args,kwargs)
    end

    function NeuralDSDE(model1::FastChain,model2::FastChain,tspan,solver=nothing,args...;kwargs...)
        p1 = initial_params(model1)
        p2 = initial_params(model2)
        re1 = nothing
        re2 = nothing
        p = [p1;p2]
        new{typeof(model1),typeof(p),typeof(re1),typeof(model2),typeof(re2),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(p,
            length(p1),model1,re1,model2,re2,tspan,solver,args,kwargs)
    end
end

Flux.@functor NeuralDSDE

function (n::NeuralDSDE)(x,p=n.p)
    dudt_(u,p,t) = n.re1(p[1:n.len])(u)
    g(u,p,t) = n.re2(p[(n.len+1):end])(u)
    prob = SDEProblem{false}(dudt_,g,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

function (n::NeuralDSDE{M})(x,p=n.p) where {M<:FastChain}
    dudt_(u,p,t) = n.model1(u,p[1:n.len])
    g(u,p,t) = n.model2(u,p[(n.len+1):end])
    prob = SDEProblem{false}(dudt_,g,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

struct NeuralSDE{P,M,RE,M2,RE2,T,S,A,K} <: NeuralDELayer
    p::P
    len::Int
    model1::M
    re1::RE
    model2::M2
    re2::RE2
    tspan::T
    nbrown::Int
    solver::S
    args::A
    kwargs::K
    function NeuralSDE(model1,model2,tspan,nbrown,solver=nothing,args...;kwargs...)
        p1,re1 = Flux.destructure(model1)
        p2,re2 = Flux.destructure(model2)
        p = [p1;p2]
        new{typeof(p),typeof(model1),typeof(re1),typeof(model2),typeof(re2),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(
            p,length(p1),model1,re1,model2,re2,tspan,nbrown,solver,args,kwargs)
    end

    function NeuralSDE(model1::FastChain,model2::FastChain,tspan,nbrown,solver=nothing,args...;kwargs...)
        p1 = initial_params(model1)
        p2 = initial_params(model2)
        re1 = nothing
        re2 = nothing
        p = [p1;p2]
        new{typeof(p),typeof(model1),typeof(re1),typeof(model2),typeof(re2),
            typeof(tspan),typeof(solver),typeof(args),typeof(kwargs)}(
            p,length(p1),model1,re1,model2,re2,tspan,nbrown,solver,args,kwargs)
    end
end

Flux.@functor NeuralSDE

function (n::NeuralSDE)(x,p=n.p)
    dudt_(u,p,t) = n.re1(p[1:n.len])(u)
    g(u,p,t) = n.re2(p[(n.len+1):end])(u)
    prob = SDEProblem{false}(dudt_,g,x,n.tspan,p,noise_rate_prototype=zeros(Float32,length(x),n.nbrown))
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

function (n::NeuralSDE{P,M})(x,p=n.p) where {P,M<:FastChain}
    dudt_(u,p,t) = n.model1(u,p[1:n.len])
    g(u,p,t) = n.model2(u,p[(n.len+1):end])
    prob = SDEProblem{false}(dudt_,g,x,n.tspan,p,noise_rate_prototype=zeros(Float32,length(x),n.nbrown))
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

struct NeuralCDDE{P,M,RE,H,L,T,S,A,K} <: NeuralDELayer
    p::P
    model::M
    re::RE
    hist::H
    lags::L
    tspan::T
    solver::S
    args::A
    kwargs::K
end

function NeuralCDDE(model,tspan,hist,lags,solver=nothing,args...;kwargs...)
    p,re = Flux.destructure(model)
    NeuralCDDE(p,model,re,hist,lags,tspan,solver,args,kwargs)
end

function NeuralCDDE(model::FastChain,tspan,hist,lags,solver=nothing,args...;kwargs...)
    p = initial_params(model)
    re = nothing
    NeuralCDDE(p,model,re,hist,lags,tspan,solver,args,kwargs)
end

Flux.@functor NeuralCDDE

function (n::NeuralCDDE)(x,p=n.p)
    function dudt_(u,h,p,t)
        _u = vcat(u,(h(p,t-lag) for lag in n.lags)...)
        n.re(p)(_u)
    end
    prob = DDEProblem{false}(dudt_,x,n.hist,n.tspan,p,constant_lags = n.lags)
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

function (n::NeuralCDDE{P,M})(x,p=n.p) where {P,M<:FastChain}
    function dudt_(u,h,p,t)
        _u = vcat(u,(h(p,t-lag) for lag in n.lags)...)
        n.model(_u,p)
    end
    prob = DDEProblem{false}(dudt_,x,n.hist,n.tspan,p,constant_lags = n.lags)
    concrete_solve(prob,n.solver,x,p,n.args...;sensealg=TrackerAdjoint(),n.kwargs...)
end

struct NeuralDAE{P,M,M2,D,RE,T,S,DV,A,K} <: NeuralDELayer
    model::M
    constraints_model::M2
    p::P
    du0::D
    re::RE
    tspan::T
    solver::S
    differential_vars::DV
    args::A
    kwargs::K

    function NeuralDAE(model,constraints_model,tspan,solver=nothing,du0=nothing,args...;differential_vars=nothing,kwargs...)
        p,re = Flux.destructure(model)
        new{typeof(p),typeof(model),typeof(constraints_model),typeof(du0),typeof(re),
            typeof(tspan),typeof(solver),typeof(differential_vars),typeof(args),typeof(kwargs)}(
            model,constraints_model,p,du0,re,tspan,solver,differential_vars,args,kwargs)
    end
end

Flux.@functor NeuralDAE

function (n::NeuralDAE)(x,du0=n.du0,p=n.p)
    function f(du,u,p,t)
        nn_out = n.re(p)(u)
        alg_out = n.constraints_model(u,p,t)
        v_out = []
        for (j,i) in enumerate(n.differential_vars)
            if i
                push!(v_out,nn_out[j])
            else
                push!(v_out,alg_out[j])
            end
        end
        return v_out
    end
    dudt_(du,u,p,t) = f
    prob = DAEProblem(dudt_,du0,x,n.tspan,p,differential_vars=n.differential_vars)
    concrete_solve(prob,n.solver,x,p,n.args...;sensalg=TrackerAdjoint(),n.kwargs...)
end

struct NeuralODEMM{M,M2,P,RE,T,S,MM,A,K} <: NeuralDELayer
    model::M
    constraints_model::M2
    p::P
    re::RE
    tspan::T
    solver::S
    mass_matrix::MM
    args::A
    kwargs::K

    function NeuralODEMM(model,constraints_model,tspan,mass_matrix,solver=nothing,args...;kwargs...)
        p,re = Flux.destructure(model)
        new{typeof(model),typeof(constraints_model),typeof(p),typeof(re),
            typeof(tspan),typeof(solver),typeof(mass_matrix),typeof(args),typeof(kwargs)}(
            model,constraints_model,p,re,tspan,solver,mass_matrix,args,kwargs)
    end

    function NeuralODEMM(model::FastChain,constraints_model,tspan,mass_matrix,solver=nothing,args...;kwargs...)
        p = initial_params(model)
        re = nothing
        new{typeof(model),typeof(constraints_model),typeof(p),typeof(re),
            typeof(tspan),typeof(solver),typeof(mass_matrix),typeof(args),typeof(kwargs)}(
            model,constraints_model,p,re,tspan,solver,mass_matrix,args,kwargs)
    end
end

Flux.@functor NeuralODEMM

function (n::NeuralODEMM)(x,p=n.p)
    function f(u,p,t)
        nn_out = n.re(p)(u)
        alg_out = n.constraints_model(u,p,t)
        vcat(nn_out,alg_out)
    end
    dudt_= ODEFunction(f,mass_matrix=n.mass_matrix)
    prob = ODEProblem(dudt_,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;n.kwargs...)
end

function (n::NeuralODEMM{M})(x,p=n.p) where {M<:FastChain}
    function f(u,p,t)
        nn_out = n.model(u,p)
        alg_out = n.constraints_model(u,p,t)
        vcat(nn_out,alg_out)
    end
    dudt_= ODEFunction(f,mass_matrix=n.mass_matrix)
    prob = ODEProblem(dudt_,x,n.tspan,p)
    concrete_solve(prob,n.solver,x,p,n.args...;
                   sensealg=InterpolatingAdjoint(
                            autojacvec=DiffEqSensitivity.ReverseDiffVJP(true)),
                            n.kwargs...)
end


text:
# TODO currently does not work

abstract type KeyView{V, E_VAL, G <: AbstractValueGraph{V}} <: AbstractValueGraph{V, E_VAL} end

# TODO maybe also add adjlist to struct
struct ValueGraphKeyView{V, E_VAL, G <: ValueGraph{V}} <: KeyView{V, E_VAL, G}
    edgevals::Adjlist{E_VAL}
    parent::G # mainly used as reference for ne(g)
end

struct ValueOutDiGraphKeyView{V, E_VAL, G <: ValueOutDiGraph{V}} <: KeyView{V, E_VAL, G}
    edgevals::Adjlist{E_VAL}
    parent::G # mainly used as reference for ne(g)
end

struct ValueDiGraphKeyView{V, E_VAL, G <: ValueDiGraph{V}} <: KeyView{V, E_VAL, G}
    edgevals::Adjlist{E_VAL}
    redgevals::Adjlist{E_VAL}
    parent::G # mainly used as reference for ne(g)
end

function keyview(g::ValueGraph{V}, key) where {V}
    edgevals = g.edgevals[key]
    E_VAL = eltype(eltype(edgevals))
    return ValueGraphKeyView{V, E_VAL, typeof(g)}(edgevals, g)
end

function keyview(g::ValueOutDiGraph{V}, key) where {V}
    edgevals = g.edgevals[key]
    E_VAL = eltype(eltype(edgevals))
    return ValueOutDiGraphKeyView{V, E_VAL, typeof(g)}(edgevals, g)
end

function keyview(g::ValueDiGraph{V}, key) where {V}
    edgevals = g.edgevals[key]
    redgevals = g.redgevals[key]
    E_VAL = eltype(eltype(edgevals))
    return ValueOutDiGraphKeyView{V, E_VAL, typeof(g)}(edgevals, redgevals, g)
end

#=
    LightGraphs.edges
    LightGraphs.has_edge
=#

for f in (:ne, :nv, :is_directed, :degree, :vertices)
    @eval $f(kview::KeyView) = $f(kview.parent)
end

for f in (:degree, :inneighbors, :outneighbors, :neighbors)
    @eval $f(kview::KeyView, arg1::Integer) = $f(kview.parent, arg1)
end

is_directed(::Type{<:ValueGraphKeyView}) = false
is_directed(::Type{<:ValueOutDiGraphKeyView}) = true
is_directed(::Type{<:ValueDiGraphKeyView}) = true

eltype(kview::KeyView{V}) where {V} = V

#  ------------------------------------------------------
#  get_edgeval
#  ------------------------------------------------------

function get_edgeval(kview::ValueGraphKeyView, s::Integer, d::Integer, default=nothing)
    verts = vertices(kview)
    (s in verts && d in verts) || return default
    fadjlist = kview.parent.fadjlist
    @inbounds list_s = fadjlist[s]
    @inbounds list_d = fadjlist[d]
    if length(list_s) > length(list_d)
        s, d = d, s
        list_s = list_d
    end
    index = searchsortedfirst(list_s, d)
    @inbounds if index <= length(list_s) && list_s[index] == d
        return kview.edgevals[s][index]
    end
    return default
end

function get_edgeval(kview::ValueOutDiGraphKeyView, s::Integer, d::Integer, default=nothing)
    verts = vertices(kview)
    (s in verts && d in verts) || return default
    @inbounds list_s = kview.parent.fadjlist[s]
    index = searchsortedfirst(list_s, d)
    @inbounds if index <= length(list_s) && list_s[index] == d
        return kview.edgevals[s][index]
    end
    return default
end

# TODO could probably be made faster by checking the shorter list
function get_edgeval(kview::ValueDiGraphKeyView, s::Integer, d::Integer, default=nothing)
    verts = vertices(kview)
    (s in verts && d in verts) || return default
    @inbounds list = kview.parent.fadjlist[s]
    index = searchsortedfirst(list, d)
    @inbounds if index <= length(list) && list[index] == d
        return kview.edgevals[s][index]
    end
    return default
end

get_edgeval(kview::KeyView, e::SimpleEdge, default=nothing) = get_edgeval(kview, src(e), dst(e), default)

#  ------------------------------------------------------
#  set_edgeval!
#  ------------------------------------------------------

function set_edgeval!(kview::ValueGraphKeyView, s::Integer, d::Integer, value)
    verts = vertices(kview)
    (s in verts && d in verts) || return false
    fadjlist = kview.parent.fadjlist
    @inbounds list = fadjlist[s]
    index = searchsortedfirst(list, d)
    @inbounds index <= length(list) && list[index] == d || return false
    edgevals = kview.edgevals
    @inbounds edgevals[s][index] = value

    @inbounds list = fadjlist[d]
    index = searchsortedfirst(list, s)
    @inbounds edgevals[d][index] = value

    return true
end

function set_edgeval!(kview::ValueOutDiGraphKeyView, s::Integer, d::Integer, value)
    verts = vertices(kview)
    (s in verts && d in verts) || return false
    @inbounds list = kview.parent.fadjlist[s]
    index = searchsortedfirst(list, d)
    @inbounds index <= length(list) && list[index] == d || return false
    @inbounds g.edgevals[s][index] = value

    return true
end

function set_edgeval!(kview::ValueDiGraphKeyView, s::Integer, d::Integer, value)
    verts = vertices(kview)
    (s in verts && d in verts) || return false
    @inbounds list = kview.parent.fadjlist[s]
    index = searchsortedfirst(list, d)
    @inbounds index <= length(list) && list[index] == d || return false
    @inbounds kview.edgevals[s][index] = value

    @inbounds list = kview.parent.badjlist[d]
    index = searchsortedfirst(list, s)
    @inbounds kview.redgevals[s][index] = value

    return true
end

set_edgeval!(kview::KeyView, e::SimpleEdge, u) = set_edgeval!(kview, src(e), dst(e), u)

#  ------------------------------------------------------
#  iterators
#  ------------------------------------------------------

@inline function iterate(eit::ValueEdgeIter{<:ValueGraphKeyView{V}}, state=(one(V), 1) ) where {V}
    kview = eit.g
    fadjlist = kview.parent.fadjlist
    edgevals = kview.edgevals
    n = V(nv(kview))
    u, i = state

    @inbounds while u < n
        list_u = fadjlist[u]
        if i > length(list_u)
            u += one(u)
            i = searchsortedfirst(fadjlist[u], u)
            continue
        end
        e = ValueEdge(u, list_u[i], edgevals[u][i])
        state = (u, i + 1)
        return e, state
    end

    # i > length(fadjlist[end]) || fadlist[end][i] == n

    @inbounds (n == 0 || i > length(fadjlist[n])) && return nothing

    @inbounds e = SimpleEdge(n, n, edgevals[n][end])
    state = (u, i + 1)
    return e, state
end

function iterate(
            eit::ValueEdgeIter{<:Union{ValueOutDiGraphKeyView{V}, ValueDiGraphKeyView{V}} },
            state=(one(V), 1)
    ) where {V}

    kview = eit.g
    fadjlist = kview.parent.fadjlist
    edgevals = kview.edgevals
    n = V(nv(g))
    u, i = state

    @inbounds while u <= n
        if i > length(fadjlist[u])
            u == n && return nothing

            u += one(u)
            i = 1
            continue
        end
        e = ValueEdge(u, fadjlist[u][i], edgevals[u][i])
        return e, (u, i + 1)
    end

    return nothing
end

#TODO
# - has_edge


text:
module IonBase

using Comonicon
using MatchCore
# for precompile
using Pkg
using GitHub
using Downloads
using Comonicon.Tools: cmd_error

Comonicon.set_brief_length!(120)

# GITHUB_TOKEN is used in github actions
# GITHUB_AUTH is suggested by GitHub.jl
const ENV_GITHUB_TOKEN_NAMES = ["GITHUB_TOKEN", "GITHUB_AUTH"]

# TODO: we might want to change PackageCompiler to set our own
# depot path
function ion_dir()
    if haskey(ENV, "DOT_ION_PATH")
        return ENV["DOT_ION_PATH"]
    elseif Sys.isapple() || Sys.islinux()
        return joinpath(homedir(), ".ion")
    elseif Sys.iswindows()
        return expanduser(raw"~\AppData\Local\Ion")
    end
end

dot_ion(xs...) = joinpath(ion_dir(), xs...)
ca_roots() = dot_ion("cert.pem")

function templates(xs...)
    if haskey(ENV, "ION_TEMPLATE_PATH")
        return ENV["ION_TEMPLATE_PATH"]
    else
        return dot_ion("templates", xs...)
    end
end

function copy_templates(dst::String=templates())
    ispath(dst) || mkpath(dst)
    ion_template_dir = joinpath(pkgdir(IonBase), "templates")
    cp(ion_template_dir, dst; force=true, follow_symlinks=true)
end

function copy_cacert(dst::String=ca_roots())
    # Linux is fine, we just ship julia's ca
    # with Mac build
    if Sys.isapple()
        cp(Downloads.ca_roots_path(), dst; force=true, follow_symlinks=true)
    end
end

function copy_assets(root::String=dot_ion())
    ispath(root) || mkpath(root)
    template_dir = joinpath(root, "templates")
    @info "copying assets templates to $template_dir"
    copy_templates(template_dir)
    cacert_dir = joinpath(root, "cert.pem")
    @info "copying assets cacert to $cacert_dir"
    copy_cacert(cacert_dir)
end

ion_toml() = dot_ion("ion.toml")

function init_dot_ion()
    dot_ion_dir = dot_ion()
    if !ispath(dot_ion_dir)
        mkpath(dot_ion_dir)
        return true
    end
    return false
end

function read_github_auth()
    for key in ENV_GITHUB_TOKEN_NAMES
        if haskey(ENV, key)
            return ENV[key]
        end
    end

    buf = Base.getpass("GitHub Access Token (https://github.com/settings/tokens)")
    auth = read(buf, String)
    Base.shred!(buf)
    return auth
end

function gitcmd(path::AbstractString; kw...)
    cmd = ["git", "-C", path]
    for (n,v) in kw
        push!(cmd, "-c")
        push!(cmd, "$n=$v")
    end
    return Cmd(cmd)
end

include("options.jl")
# level 1 commands
include("command/install.jl")
include("command/activate.jl")
include("command/pkg.jl")

# level 2 commands
include("templates/templates.jl")
include("command/search.jl")
include("command/release.jl")
include("command/clone.jl")

# extensions
include("command/doc.jl")
include("command/plugin.jl")

# include("precompile.jl")
# _precompile_()

# "The Ion manager."
# @main

end


text:
function cell_1d(xs::ANY...)
    n = length(xs)
    a = Array(Any,n)
    for i=1:n
        arrayset(a,xs[i],i)
    end
    a
end

function cell_2d(nr, nc, xs::ANY...)
    a = Array(Any,nr,nc)
    for i=1:(nr*nc)
        arrayset(a,xs[i],i)
    end
    a
end

# map cell array
map(f::Callable, a::Array{Any,1}) = { f(a[i]) for i=1:length(a) }
map(f::Callable, a::Array{Any,1}, b::Array{Any,1}) =
    { f(a[i],b[i]) for i=1:length_checked_equal(a, b) }
function map(f::Callable, as::Array{Any,1}...)
    n = length_checked_equal(as...)
    { f(map(a->a[i],as)...) for i=1:n }
end

cell(dims::(Integer...))   = Array(Any, map(int, dims))
cell(dims::Integer...)     = Array(Any, map(int, dims))
cell(d1::Integer)          = Array(Any, int(d1))
cell(d1::Integer, d2::Integer) = Array(Any, int(d1), int(d2))

setindex!(A::Array{Any}, x::ANY, i::Real) = arrayset(A, x, to_index(i))

max_stars_repo_name: tobig/julia

max_issues_repo_name: tobig/julia

max_forks_repo_name: tobig/julia


text:
__precompile__()
module DrinkGuiden
using StringDistances

const SIMILARITY_THRESHOLD = 0.8 # Similar names if above

include("types.jl")
include("datalayer.jl")
include("logiclayer.jl")
include("presentationlayer.jl")
include("show.jl")

include("main.jl") # Run the program

end


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

"""
    abstract type AbstractVariateTransform <: Function

Abstract type for change-of-variables transformations.

Subtypes (e.g. `SomeTrafo <: AbstractVariateTransform`) must support (with
`trafo::SomeTrafo`):

```julia
    (trafo)(v_prev::SomeVariate) == v_new
    (trafo)(v_prev::SomeVariate, ladj_prev::Real)) == (v = v_new, ladj = ladj_new)
    (trafo)(s_prev::DensitySample)::DensitySample
    ((trafo2 ∘ trafo1)(v)::AbstractVariateTransform)(v) == trafo2(trafo1(v))
    inv(trafo)(trafo(v)) == v
    inv(inv(trafo)) == trafo

    ValueShapes.varshape(trafo)::ValueShapes.AbstractValueShape
```

with `varshape(v_prev) == varshape(trafo)`.

`ladj` must be `logabsdet(jacobian(trafo, v))`.
"""
abstract type AbstractVariateTransform end
export AbstractVariateTransform

"""
    ladjof(r::NamedTuple{(...,:ladj,...)})::Real

Extract the `log(abs(det(jacobian)))` value that is part of a result `r`.

Examples:

```julia
ladjof((..., ladj = some_ladj, ...)) == some_ladj
ladjof(trafo)(v) = trafo(v, )
```
"""
function ladjof end
export ladjof

ladjof(x::NamedTuple) = x.ladj

struct LADJOfVarTrafo{T<:AbstractVariateTransform} <: Function
    trafo::T
end

(ladjof_trafo::LADJOfVarTrafo)(v::Any) = ladjof(trafo(v, 1))
(ladjof_trafo::LADJOfVarTrafo)(v::Any, prev_ladj::Real) = ladjof(trafo(v, prev_ladj))

"""
    ladjof(trafo::AbstractVariateTransform)::Function

Returns a function that computes the `log(abs(det(jacobian)))` of `trafo` for
a given variate `v`:

```julia
    ladjof(trafo)(v) == ladjof(trafo(v, 1))
    ladjof(trafo)(v, prev_ladj) == ladjof(trafo(v, prev_ladj))
```
"""
ladjof(trafo::AbstractVariateTransform) = LADJOfVarTrafo(trafo)

"""
    abstract type VariateSpace <: Function

*BAT-internal, not part of stable public API.*

Abstract type for variate spaces.
"""
abstract type VariateSpace end

struct UnitSpace <: VariateSpace end
struct InfiniteSpace <: VariateSpace end
struct MixedSpace <: VariateSpace end

product_varspace(s::VariateSpace) = s

product_varspace(::UnitSpace, ::UnitSpace) = MixedSpace
product_varspace(::InfiniteSpace, ::InfiniteSpace) = MixedSpace
product_varspace(::VariateSpace, ::VariateSpace) = MixedSpace

function product_varspace(s1::VariateSpace, s2::VariateSpace, spcs::VariateSpace...)
    product_varspace(product_varspace(s1, s2), spcs...)
end

"""
    abstract type VariateTransform{VF:<VariateForm,ST<:VariateSpace,SF<:VariateSpace}

*BAT-internal, not part of stable public API.*

Abstract parameterized type for change-of-variables transformations.

Subtypes (e.g. `SomeTrafo <: VariateTransform`) must implement:

* `BAT.target_space(trafo::SomeTrafo, v)`
* `BAT.source_space(trafo::SomeTrafo, v)`
* `BAT.apply_vartrafo_impl(trafo::SomeTrafo, v)`
* `BAT.apply_vartrafo_impl(inv_trafo::InverseVT{SomeTrafo}, v)`
* `ValueShapes.varshape(trafo::SomeTrafo)`

for real values and/or real-valued vectors `v`.
"""
abstract type VariateTransform{
    VF<:VariateForm,ST<:VariateSpace,SF<:VariateSpace
} <: AbstractVariateTransform end

function target_space end

function source_space end

function apply_vartrafo end

function apply_vartrafo_impl end

apply_vartrafo(trafo::VariateTransform{Univariate}, v::Real, prev_ladj::Real) =
    apply_vartrafo_impl(trafo, v, prev_ladj)

function apply_vartrafo(trafo::VariateTransform{Univariate}, v::AbstractArray{<:Real,0}, prev_ladj::Real)
    r = apply_vartrafo_impl(trafo, v[], prev_ladj)
    (v = fill(r.v), ladj = r.ladj)
end

apply_vartrafo(trafo::VariateTransform{Multivariate}, v::AbstractVector{<:Real}, prev_ladj::Real) =
    apply_vartrafo_impl(trafo, v, prev_ladj)

apply_vartrafo(trafo::VariateTransform{Matrixvariate}, v::AbstractMatrix{<:Real}, prev_ladj::Real) =
    apply_vartrafo_impl(trafo, v, prev_ladj)

apply_vartrafo(trafo::VariateTransform{ValueShapes.NamedTupleVariate{names}}, v::NamedTuple{names}, prev_ladj::Real) where names =
    apply_vartrafo_impl(trafo, v, prev_ladj)

apply_vartrafo(trafo::VariateTransform{ValueShapes.NamedTupleVariate{names}}, v::ShapedAsNT{<:NamedTuple{names}}, prev_ladj::Real) where names =
    apply_vartrafo_impl(trafo, v, prev_ladj)

(trafo::VariateTransform)(v::Any) = apply_vartrafo(trafo, v, Float32(NaN)).v
(trafo::VariateTransform)(v::Any, prev_ladj::Real) = apply_vartrafo(trafo, v, prev_ladj)

# ToDo: Move to ValueShapes.jl:
_variate_form(shape::ScalarShape{<:Real}) = Univariate
_variate_form(shape::ArrayShape{<:Real,1}) = Multivariate
_variate_form(shape::ArrayShape{<:Real,2}) = Matrixvariate
_variate_form(shape::NamedTupleShape{names}) where names = ValueShapes.NamedTupleVariate{names}

function var_trafo_result(trg_v::Real, src_v::Real, trafo_ladj::Real, prev_ladj::Real)
    R = float(typeof(src_v))
    ladj_sum = convert(R, trafo_ladj + prev_ladj)
    trg_ladj = if !isnan(ladj_sum)
        ladj_sum
    else
        # Should be safe to assume that target dist goes to zero at infinity, should win out over infinite prev_ladj:
        ladjs_should_cancel = (trafo_ladj == R(-Inf) && prev_ladj == R(+Inf) && isinf(trg_v))
        ladjs_should_cancel ? zero(R) : ladj_sum
    end
    (v = convert(R, trg_v), ladj = trg_ladj)
end

function var_trafo_result(trg_v::Real, src_v::Real)
    R = float(typeof(src_v))
    (v = convert(R, trg_v), ladj = convert(R, NaN))
end

function var_trafo_result(trg_v::AbstractVector{<:Real}, src_v::AbstractVector{<:Real}, trafo_ladj::Real, prev_ladj::Real)
    R = float(eltype(src_v))
    ladj_sum = convert(R, trafo_ladj + prev_ladj)
    trg_ladj = if !isnan(ladj_sum)
        ladj_sum
    else
        # Should be safe to assume that target dist goes to zero at infinity, should win out over infinite prev_ladj:
        ladjs_should_cancel = (trafo_ladj == R(-Inf) && prev_ladj == R(+Inf) && any(isinf, trg_v))
        ladjs_should_cancel ? zero(R) : ladj_sum
    end
    (v = convert_eltype(R, trg_v), ladj = trg_ladj)
end

function var_trafo_result(trg_v::AbstractVector{<:Real}, src_v::AbstractVector{<:Real})
    R = float(eltype(src_v))
    (v = convert_eltype(R, trg_v), ladj = convert(R, NaN))
end

function (trafo::VariateTransform)(s::DensitySample)
    r = trafo(s.v, zero(Float32))
    v = stripscalar(r.v)  # ToDo: Do we want to use stripscalar here?
    logd = s.logd - r.ladj
    DensitySample(v, logd, s.weight, s.info, s.aux)
end

# Custom broadcast(::VariateTransform, DensitySampleVector), multithreaded:
function Base.copy(
    instance::Base.Broadcast.Broadcasted{
        <:Base.Broadcast.AbstractArrayStyle{1},
        <:Any,
        <:VariateTransform,
        <:Tuple{DensitySampleVector}
    }
)
    trafo = instance.f
    s_src = instance.args[1]
    vs_trg = valshape(trafo(first(s_src.v)))
    s_src_us = unshaped.(s_src)

    n = length(eachindex(s_src_us))
    s_trg_unshaped = DensitySampleVector((
        nestedview(similar(flatview(s_src_us.v), totalndof(vs_trg), n)),
        zero(s_src_us.logd),
        deepcopy(s_src_us.weight),
        deepcopy(s_src_us.info),
        deepcopy(s_src_us.aux),
    ))
    @assert axes(s_trg_unshaped) == axes(s_src)
    @assert s_trg_unshaped.v isa ArrayOfSimilarArrays
    @threads for i in eachindex(s_trg_unshaped, s_src)
        r = trafo(s_src.v[i], zero(Float32))
        s_trg_unshaped.v[i] .= unshaped(r.v)
        s_trg_unshaped.logd[i] -= r.ladj
    end
    vs_trg.(s_trg_unshaped)
end

struct InvVT{
    VF <: VariateForm,
    ST <: VariateSpace,
    SF <: VariateSpace,
    FT <: VariateTransform{VF,ST,SF}
} <: VariateTransform{VF,SF,ST}
    orig::FT
end

Base.inv(trafo::VariateTransform) = InvVT(trafo)
Base.inv(trafo::InvVT) = trafo.orig

ValueShapes.varshape(trafo::InvVT) = varshape(trafo.orig)

target_space(trafo::InvVT) = source_space(trafo.orig)

source_space(trafo::InvVT) = target_space(trafo.orig)

const InverseVT{FT} = InvVT{VF,ST,SF,FT} where {VF,ST,SF}

struct IdentityVT{
    VF <: VariateForm,
    S <: VariateSpace,
    VS <: AbstractValueShape
} <: VariateTransform{VF,S,S}
    space::S
    varshape::VS
end

function IdentityVT(space::VariateSpace, shape::AbstractValueShape)
    VF = _variate_form(shape)
    S = typeof(space)
    VS = typeof(shape)
    IdentityVT{VF,S,VS}(space, shape)
end

Base.inv(trafo::IdentityVT) = trafo

ValueShapes.varshape(trafo::IdentityVT) = trafo.varshape

function apply_vartrafo_impl(trafo::Union{IdentityVT,InvVT{<:IdentityVT}}, v::Any, prev_ladj::Real)
    (v = v, ladj = prev_ladj)
end

(trafo::Union{IdentityVT,InvVT{<:IdentityVT}})(s::DensitySample) = s

# Custom broadcast(::IdentityVT, DensitySampleVector), multithreaded:
function Base.copy(
    instance::Base.Broadcast.Broadcasted{
        <:Base.Broadcast.AbstractArrayStyle{1},
        <:Any,
        <:Union{IdentityVT,InvVT{<:IdentityVT}},
        <:Tuple{DensitySampleVector}
    }
)
    instance.args[1]
end


text:
module TSDecomposition
export filterTSFFT
using ..Cubes
using ..DAT
using ..Proc
using FFTW
using Statistics
import ..DAT: AnyMissing

#Looks like linreg is broken in 0.7, here is a custom version, this should be replaced soon:
function linreg(x,y)
  b = cov(x,y)/var(x)
  a = mean(y) - b*mean(x)
  a,b
end

function detrendTS!(outar::AbstractMatrix,xin::AbstractVector{T}) where T
    x=T[i for i in 1:length(xin)]
    a,b=linreg(x,xin)
    for i in eachindex(xin)
        outar[i,1]=a+b*x[i]
    end
end

function tscale2ind(b::Float64,l::Int)
    i=round(Int,l/b+1)
    return i
end
mirror(i,l)=l-i+2

"""
    filterTSFFT(c::AbstractCubeData)

Filter each time series using a Fourier filter and return the decomposed series
in 4 time windows (Trend, Long-Term Variability, Annual Cycle, Fast Oscillations)

**Input Axis** `Time`axis

**Output Axes** `Time`axis, `Scale`axis

"""
function filterTSFFT(c::AbstractCubeData;kwargs...)
  indims = InDims(TimeAxis,filter=AnyMissing())
  outdims = OutDims(TimeAxis,(c,p)->ScaleAxis(["Trend", "Long-Term Variability", "Annual Cycle", "Fast Oscillations"]))
  ntime = length(getAxis("Time",c))
  testar = zeros(Complex{Base.nonmissingtype(eltype(c))},ntime)
  fftplan = plan_fft!(testar)
  ifftplan = inv(fftplan)
  mapCube(filterTSFFT,c,getNpY(c),fftplan,ifftplan;indims=indims,outdims=outdims,kwargs...)
end

function filterTSFFT(outar::AbstractMatrix,y::AbstractVector, annfreq::Number,
  fftplan, ifftplan; nharm::Int=3)

  any(ismissing,y) && return outar[:].=missing

    size(outar) == (length(y),4) || error("Wrong size of output array")

    detrendTS!(outar,y)
    l        = length(y)

    fy       = Complex{Base.nonmissingtype(eltype(y))}[y[i]-outar[i,1] for i=1:l]
    fftplan * fy
    fyout    = similar(fy)
    czero    = zero(eltype(fy))

    #Remove annual cycle
    fill!(fyout,zero(eltype(fyout)))

    for jharm = 1:nharm
        iup=tscale2ind(annfreq*1.1/jharm,l)

        idown=tscale2ind(annfreq*0.9/jharm,l)

        for i=iup:idown
            fyout[i]  = fy[i]
            i2        = mirror(i,l)
            fyout[i2] = fy[i2]
            fy[i]     = czero
            fy[i2]    = czero
        end
    end

    ifftplan * fyout
    for i=1:l
        outar[i,3]=real(fyout[i])
    end

    iup   = tscale2ind(annfreq*1.1,l)
    idown = tscale2ind(annfreq*0.9,l)
    #Now split the remaining parts
    fill!(fyout,czero)
    for i=2:iup-1
        i2        = mirror(i,l)
        fyout[i]  = fy[i]
        fyout[i2] = fy[i2]
        fy[i]     = czero
        fy[i2]    = czero
    end
    ifftplan * fyout
    ifftplan * fy

    for i=1:l
        outar[i,2]=real(fyout[i])
        outar[i,4]=real(fy[i])
    end
end
end


text:
using PyPlot
using PyCall
mpl = pyimport("tikzplotlib")
using DelimitedFiles
using MAT

force_scales = [5.0, 6.0, 8.0, 10.0, 20.0]
tids = [106, 206, 300]

Multiscale = zeros(5,3)
NN = zeros(5,3)

for (k1, force_scale) in enumerate(force_scales)
    for (k2, tid) in enumerate(tids)
        Multiscale[k1, k2] = readdlm("Multiscale$(force_scale)$tid.txt")[2]
        NN[k1, k2] = readdlm("NN$(force_scale)$tid.txt")[2]
    end
end

matwrite("cpu.mat", Dict("Multiscale"=>Multiscale, "NN"=>NN))

close("all")

cols = "rgb"
for i = 1:3
    semilogy(8 ./ force_scales, Multiscale[:,i], "o$(cols[i])-", linewidth=3)
    semilogy(8 ./ force_scales, NN[:,i], "o$(cols[i])", linestyle="dotted", linewidth=3)
end

a = ones(5); fill!(a, NaN)
semilogy(8 ./ force_scales, a, "k-", label="Fiber resolved DNS", linewidth=3)
semilogy(8 ./ force_scales, a, "k", linestyle="dotted",  label="CholNN", linewidth=3)

legend()
xlabel("Load (GPa)")
ylabel("CPU Time (sec)")
grid("on")
savefig("cpu.pdf")
savefig("cpu.png")
mpl.save("cpu.tex")