using Plots, Statistics, Random
using Flux, DiffEqFlux, StochasticDiffEq, CUDA, RegNeuralDE, Tracker
using BSON
using Tracker: data
using ColorBrewer
colors = ColorBrewer.palette("Set2", 8)

dataset = BSON.load("data/sde_demo.bson")
sde_data = dataset[:sde_data]
sde_data_vars = dataset[:sde_data_vars]

u0 = reshape(Float32[2.0; 0.0], :, 1)
datasize = 30
tspan = [0.0f0, 1.0f0]
tsteps = range(tspan[1], tspan[2], length = datasize) |> track

plots1, plots2 = [], []
ptime1, ptime2 = [], []
ttime1, ttime2 = [], []
loss1, loss2 = [], []
nfe1, nfe2 = [], []

neuralsde = nothing
REGULARIZE = false

function loss_function(u0, p, i, c = 0.2f0)
    sol, nfe1, nfe2, sv = neuralsde(u0, p)
    means = mean(sol; dims = 3)
    vars = var(sol; dims = 3, mean = means)
    l2_means = mean(abs2, sde_data .- means)
    l2_vars = mean(abs2, sde_data_vars .- vars)
    reg = REGULARIZE ? c * sum(sv.saveval) : 0.0f0
    loss = l2_means + l2_vars
    if i % 50 == 0
        @show i, data(loss), data(l2_means), data(l2_vars), data(reg), nfe1, nfe2
    end
    i == -1 && return (loss, nfe1)
    return loss + reg
end

for i in [5]#, 35, 1]
    for reg in [false, true]
        Random.seed!(i)

        drift_dudt = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2)) |> track
        diffusion_dudt = Dense(2, 2) |> track

        REGULARIZE = reg

        neuralsde = TrackedNeuralDSDE(
            drift_dudt,
            diffusion_dudt,
            [0.0f0, 1.0f0 + eps(Float32)],
            REGULARIZE,
            SOSRI(),
            saveat = tsteps,
            reltol = 3f-1,
            abstol = 3f-1,
        )

        u0_ = repeat(u0, 1, 100)
        loss_function(u0_ |> track, neuralsde.p, 0)
        Tracker.gradient(p -> loss_function(u0_ |> track, p, 0), neuralsde.p)

        opt = AdaBelief(0.01)
        ps = neuralsde.p
        ps_best = ps
        total_time = 0
        for iter = 1:250
            _t = time()
            l, back = Tracker.forward(p -> loss_function(u0_ |> track, p, iter), ps)
            gs = back(1.0f0)[1]
            total_time += time() - _t
            ps_best = ps
            update_parameters!((ps,), (gs,), opt)
        end

        u0_ = repeat(u0, 1, 100)

        begin
            neuralsde(u0_ |> track, ps_best)
            ptime = @belapsed neuralsde($u0_ |> track, $ps_best)[1]
            preds = neuralsde(u0_ |> track, ps_best)[1]
            preds = preds |> untrack
            means, vars = mean(preds; dims = 3), var(preds, dims = 3)

            plt = plot(
                tsteps |> untrack,
                means[1, :, 1],
                ribbon = vars[1, :, 1],
                color = colors[1],
                fillalpha = 0.2,
                label = "Pred: Dim 1",
                linewidth = 3,
                legend = :bottomleft,
            )
            plot!(
                plt,
                tsteps |> untrack,
                means[2, :, 1],
                ribbon = vars[2, :, 1],
                color = colors[2],
                fillalpha = 0.2,
                label = "Pred: Dim 2",
                linewidth = 3,
            )
            scatter!(
                plt,
                tsteps |> untrack,
                sde_data[1, :],
                yerror = sde_data_vars[1, :],
                label = "Data: Dim 1",
                color = colors[1],
                msc = colors[1],
            )
            scatter!(
                plt,
                tsteps |> untrack,
                sde_data[2, :],
                yerror = sde_data_vars[2, :],
                label = "Data: Dim 2",
                color = colors[2],
                msc = colors[2],
            )

            if REGULARIZE
                push!(plots1, plt)
            else
                push!(plots2, plt)
            end
            display(plt)
        end

        if REGULARIZE
            push!(ttime1, total_time)
            l, n = loss_function(u0_ |> track, ps_best, -1)
            push!(loss1, l)
            push!(ptime1, ptime)
            push!(nfe1, n)

            @show ttime1[end], l, n, ptime[end]
        else
            push!(ttime2, total_time)
            l, n = loss_function(u0_ |> track, ps_best, -1)
            push!(loss2, l)
            push!(ptime2, ptime)
            push!(nfe2, n)

            @show ttime2[end], l, n, ptime2[end]
        end

        GC.gc(true)
        GC.gc(true)
        GC.gc(true)
    end
end

@show mean(loss1), std(loss1), mean(loss2), std(loss2)
@show mean(nfe1), std(nfe1), mean(nfe2), std(nfe2)
@show mean(ptime1), std(ptime1), mean(ptime2), std(ptime2)
@show mean(ttime1), std(ttime1), mean(ttime2), std(ttime2)

title!(plots2[1], "Unregularized Neural SDE")
title!(plots1[1], "Regularized Neural SDE")
xlabel!(plots1[1], "t")
plt = plot(plots2[1], plots1[1], layout = (2, 1))

savefig(plt, "spiral_sde.pdf")


function mp_evaluate(equations, x)
    variables = MP.variables(equations)
    map(equations) do f
        f(variables => x)
    end
end

function mp_jacobian(equations, x)
    variables = MP.variables(equations)
    map(Iterators.product(equations, variables)) do arg
        (f, var) = arg
        MP.differentiate(f, var)(variables => x)
    end
end

function sp_evaluate(equations, x)
    F = SP.PolynomialSystem(equations)
    SP.evaluate(F, x)
end
function sp_jacobian(equations, x)
    F = SP.PolynomialSystem(equations)
    SP.jacobian(F, x)
end

function katsura(n)
  @polyvar x[0:n] # This creates variables x0, x1, ...

  return [
    (sum(x[abs(l)+1]*x[abs(m-l)+1] for l=-n:n if abs(m-l)<=n) -
    x[m+1] for m=0:n-1)...,
    x[1] + 2sum(x[i+1] for i=1:n) - 1
  ]
end

@testset "Katsura Evaluation tests" begin
    for T = [Float64, Complex{Float64}]
        for n=4:12
            nvars = n+1
            eqs = katsura(n)
            x = SVector{nvars}(rand(T, nvars))

            @test norm(sp_evaluate(eqs, x) - mp_evaluate(eqs, x)) < 1e-14
            @test norm(sp_evaluate(eqs, Vector(x)) - mp_evaluate(eqs, x)) < 1e-14

            @test norm(sp_jacobian(eqs, x) - mp_jacobian(eqs, x)) < 1e-14
            @test norm(sp_jacobian(eqs, Vector(x)) - mp_jacobian(eqs, x)) < 1e-14

            F = SP.PolynomialSystem(eqs)
            val, jac = SP.evaluate_and_jacobian(F, x)
            @test norm(val - mp_evaluate(eqs, x)) < 1e-14
            @test norm(F(x) - mp_evaluate(eqs, x)) < 1e-14
            @test norm(jac - mp_jacobian(eqs, x)) < 1e-14

            val, jac = SP.evaluate_and_jacobian(F, Vector(x))
            @test norm(val - mp_evaluate(eqs, x)) < 1e-14
            @test norm(F(x) - mp_evaluate(eqs, x)) < 1e-14
            @test norm(jac - mp_jacobian(eqs, x)) < 1e-14
        end
    end
end


abstract type SamplingOperator <: AbstractOperator
end

dest_space(op::SamplingOperator) = Span(dest(op))

function apply(op::SamplingOperator, f::AbstractVector; options...)
	@warn "This is a funny function to use"
	@assert length(f) == size(op,2)
	f
end

"""
A `GridSampling` is an operator that maps a function to its samples
in a grid.
"""
struct GridSampling <: SamplingOperator
    src     ::  FunctionSpace
    dest    ::  GridBasis
end

GridSampling(grid::AbstractGrid{S}, ::Type{T} = subeltype(S)) where {S,T} =
    GridSampling(GenericFunctionSpace{S,T}(), GridBasis{T}(grid))

GridSampling(gridbasis::GridBasis{S,T}) where {S,T} =
	GridSampling(GenericFunctionSpace{eltype(grid(gridbasis)),T}(), gridbasis)

GridSampling(m::DiscreteWeight, ::Type{T} = numeltype(m)) where T = GridSampling(points(m), T)

grid(op::GridSampling) = grid(dest(op))

(op::GridSampling)(f) = apply(op, f)

dest(op::GridSampling) = op.dest

src_space(op::GridSampling) = op.src

apply!(result, op::GridSampling, f; options...) = sample!(result, grid(op), f; options...)

"Sample the function f on the given grid."
sample(g::AbstractGrid, f, T = numtype(g)) = sample!(zeros(T, size(g)), g, f)

broadcast(f::Function, grid::AbstractGrid) = sample(grid, f)

component(op::GridSampling, i) = GridSampling(component(dest(op),i))
components(op::GridSampling) = map( s -> GridSampling(s), components(dest(op)))

factors(op::GridSampling) = components(op)

function tensorproduct(op1::GridSampling, op2::GridSampling)
	T = promote_type(coefficienttype(dest(op1)),coefficienttype(dest(op2)))
	grid1 = grid(op1)
	grid2 = grid(op2)
	GridSampling(grid1 × grid2, T)
end

function tensorproduct(op1::GridSampling, op2::AbstractOperator)
	@assert dest(op2) isa GridBasis
	@assert ncomponents(op2) == 2
	tensorproduct(IdentityOperator(dest(op1)), component(op2,2)) * tensorproduct(op1, component(op2,1))
end

function tensorproduct(ops::GridSampling...)
	T = promote_type(map(t->coefficienttype(dest(t)), ops)...)
	g = ProductGrid(map(grid, ops)...)
	GridSampling(g, T)
end

function tensorproduct(op1::AbstractOperator, op2::AbstractOperator)
	@assert dest(op1) isa GridBasis
	@assert dest(op2) isa GridBasis
	@assert ncomponents(op1) == 2
	@assert ncomponents(op2) == 2
	T = promote_type(coefficienttype(dest(op1)),coefficienttype(dest(op2)))
	grid1 = grid(dest(op1))
	grid2 = grid(dest(op2))
	g = grid1 × grid2
	tensorproduct(component(op1,2),component(op2,2)) * tensorproduct(component(op1,1),component(op2,1))
end

function tensorproduct(op1::AbstractOperator, op2::AbstractOperator, op3::AbstractOperator)
	@assert dest(op1) isa GridBasis
	@assert dest(op2) isa GridBasis
	@assert dest(op3) isa GridBasis
	# TODO: generalize to longer operators
	@assert ncomponents(op1) == 2
	@assert ncomponents(op2) == 2
	@assert ncomponents(op3) == 2
	tensorproduct(component(op1,2),component(op2,2),component(op3,2)) * tensorproduct(component(op1,1),component(op2,1),component(op3,1))
end

# We don't want to assume that f can be called with a vector argument.
# In order to avoid the overhead of splatting, we capture a number of special cases
call_function_with_vector(f, x::Number) = f(x)
call_function_with_vector(f, x::SVector{1}) = f(x[1])
call_function_with_vector(f, x::SVector{2}) = f(x[1], x[2])
call_function_with_vector(f, x::SVector{3}) = f(x[1], x[2], x[3])
call_function_with_vector(f, x::SVector{4}) = f(x[1], x[2], x[3], x[4])
call_function_with_vector(f, x::SVector{N}) where {N} = f(x...)
call_function_with_vector(f, x::AbstractVector) = f(x...)

function sample!(result, grid, f; options...)
    for i in eachindex(grid)
		result[i] = call_function_with_vector(f, grid[i])
	end
	result
end

show(io::IO, mime::MIME"text/plain", op::GridSampling) = composite_show(io, mime, op)
Display.object_parentheses(op::GridSampling) = false
Display.stencil_parentheses(op::GridSampling) = false
Display.displaystencil(op::GridSampling) = ["GridSampling(", src_space(op), ", ", grid(op), ")"]

"""
A `ProjectionSampling` is an operator that maps a function to its inner products
with a projection basis.
"""
struct ProjectionSampling <: SamplingOperator
    dict		::  Dictionary
	measure		::	Measure
	src_space	::	FunctionSpace
end

const AnalysisOperator = ProjectionSampling

space(dict::Dictionary) = space(measure(dict))

ProjectionSampling(dict::Dictionary) = ProjectionSampling(dict, measure(dict), Span(dict))

ProjectionSampling(dict::Dictionary, measure::Measure) = ProjectionSampling(dict, measure, space(measure))

ProjectionSampling(dict::Dictionary, space::FunctionSpace) = ProjectionSampling(dict, measure(space), space)

dictionary(op::ProjectionSampling) = op.dict

dest(op::ProjectionSampling) = dictionary(op)

src_space(op::ProjectionSampling) = op.src_space

measure(op::ProjectionSampling) = op.measure

apply!(result, op::ProjectionSampling, f; options...) = project!(result, f, dictionary(op), measure(op); options...)

function project!(result, f, dict::Dictionary, measure::Measure; options...)
    for i in eachindex(result)
		result[i] = innerproduct(dict[i], f, measure; options...)
	end
	result
end

show(io::IO, mime::MIME"text/plain", op::ProjectionSampling) = composite_show(io, mime, op)
Display.object_parentheses(op::ProjectionSampling) = false
Display.stencil_parentheses(op::ProjectionSampling) = false
Display.displaystencil(op::ProjectionSampling) = ["ProjectionSampling(", dictionary(op), ", ", measure(op), ")"]


using BlockArrays, BandedMatrices, BlockBandedMatrices, FillArrays, SparseArrays, Test, ArrayLayouts , LinearAlgebra
import BlockBandedMatrices: _BandedBlockBandedMatrix, blockcolsupport, blockrowsupport, colsupport, rowsupport, isbandedblockbanded, bandeddata

@testset "BandedBlockBandedMatrix" begin
    @testset "constructors" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 4
        cols = rows = 1:N

        @test Matrix(BandedBlockBandedMatrix(Zeros(sum(rows),sum(cols)), rows,cols, (l,u), (λ,μ))) ==
            Array(BandedBlockBandedMatrix(Zeros(sum(rows),sum(cols)), rows,cols, (l,u), (λ,μ))) ==
            zeros(Float64, 10, 10)

        @test Matrix(BandedBlockBandedMatrix{Int}(Zeros(sum(rows),sum(cols)), rows,cols, (l,u), (λ,μ))) ==
            zeros(Int, 10, 10)

        @test Matrix(BandedBlockBandedMatrix(Eye(sum(rows)), rows,cols, (l,u), (λ,μ))) ==
            Matrix{Float64}(I, 10, 10)

        @test Matrix(BandedBlockBandedMatrix{Int}(Eye(sum(rows)), rows,cols, (l,u), (λ,μ))) ==
            Matrix{Int}(I, 10, 10)

        @test Matrix(BandedBlockBandedMatrix(I, rows,cols, (l,u), (λ,μ))) ==
            Matrix{Float64}(I, 10, 10)

        @test Matrix(BandedBlockBandedMatrix{Int}(I, rows,cols, (l,u), (λ,μ))) ==
            Matrix{Int}(I, 10, 10)

        A = [1 2 3 4 5; 6 7 8 9 10; 11 12 13 14 15; 16 17 18 19 20; 21 22 23 24 25]
        B = BandedBlockBandedMatrix(A, [2,3], [2,3], (0,1), (1,1))
        @test Matrix(B) ==
            [1 2 3 4 0; 6 7 8 9 10; 0 0 13 14 0; 0 0 18 19 20; 0 0 0 24 25]

        l,u,λ,μ = 0,0,4,0
        cols = 1:2:3
        rows = 2*cols
        A = BandedBlockBandedMatrix(fill(1.0,sum(rows),sum(cols)), rows,cols, (l,u), (λ,μ))
        B = [1.0 0.0 0.0 0.0;
             1.0 0.0 0.0 0.0;
             0.0 1.0 0.0 0.0;
             0.0 1.0 1.0 0.0;
             0.0 1.0 1.0 1.0;
             0.0 1.0 1.0 1.0;
             0.0 1.0 1.0 1.0;
             0.0 0.0 1.0 1.0]
        @test Matrix(A) == B

        @test Matrix(sparse(A)) == B
    end

    @testset "col and rowranges" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 4
        cols = rows = 1:N
        data = reshape(collect(1:(λ+μ+1)*(l+u+1)*sum(cols)), ((λ+μ+1)*(l+u+1), sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))
        @test Matrix(BlockBandedMatrix(A)) == Matrix(A)

        @test blockbandwidths(A) == (l,u)
        @test BlockBandedMatrices.subblockbandwidths(A) == (l,u)

        @test blockrowsupport(A, 1) == Block.(1:2)
        @test blockrowsupport(A, 2) == Block.(1:3)
        @test blockrowsupport(A, 3) == Block.(2:4)

        @test blockcolsupport(A, 1) == Block.(1:2)
        @test blockcolsupport(A, 2) == Block.(1:3)
        @test blockcolsupport(A, 3) == Block.(2:4)

        @test rowsupport(A,1) == 1:3
        @test rowsupport(A,2) == 1:6
        @test rowsupport(A,3) == 1:6
        @test rowsupport(A,4) == 2:10

        @test colsupport(A,1) == 1:3
        @test colsupport(A,2) == 1:6
        @test colsupport(A,3) == 1:6
        @test colsupport(A,4) == 2:10

        l , u = 2,1
        λ , μ = 1,2
        N = M = 4
        cols = rows = 1:N
        data = reshape(Vector(1:(λ+μ+1)*(l+u+1)*sum(cols)), (λ+μ+1, (l+u+1)*sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))

        @test A.l == l == blockbandwidth(A,1)
        @test A.u == u == blockbandwidth(A,2)
        @test blockbandwidths(A) == (l, u)

        @test A.λ == λ == subblockbandwidth(A,1)
        @test A.μ == μ == subblockbandwidth(A,2)
        @test subblockbandwidths(A) == (λ, μ)

        @test blockrowsupport(A, 1) == Block.(1:2)
        @test blockrowsupport(A, 2) == Block.(1:3)
        @test blockrowsupport(A, 3) == Block.(1:4)
        @test blockrowsupport(A, 4) == Block.(2:4)

        @test blockcolsupport(A, 1) == Block.(1:3)
        @test blockcolsupport(A, 2) == Block.(1:4)
        @test blockcolsupport(A, 3) == Block.(2:4)
        @test blockcolsupport(A, 4) == Block.(3:4)

        @test rowsupport(A,1) == 1:3
        @test rowsupport(A,2) == 1:6
        @test rowsupport(A,3) == 1:6
        @test rowsupport(A,4) == 1:10

        @test colsupport(A,1) == 1:6
        @test colsupport(A,2) == 1:10
        @test colsupport(A,3) == 1:10
        @test colsupport(A,4) == 2:10

        l , u = -1,1
        λ , μ = 0,1
        rows = 1:5
        cols = 1:6

        data = reshape(Vector{Float64}(1:(λ+μ+1)*(l+u+1)*sum(cols)), (λ+μ+1, (l+u+1)*sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))
        @test_throws BandError A[1,1] = 5

        @test A[1,2] == 4
        @test A[Block(2,2)] == [0 0; 0 0]
        @test A[Block(2,3)]  == [8.0 9.0 0.0; 0.0 10.0 11.0]
        @test bandwidths(view(A,Block(2,3))) == bandwidths(A[Block(2,3)]) == (0,1)

        @test blockrowsupport(A, 1) == Block.(2:2)
        @test blockrowsupport(A, 2) == Block.(3:3)
        @test blockrowsupport(A, 3) == Block.(4:4)
        @test blockrowsupport(A, 4) == Block.(5:5)

        @test blockcolsupport(A, 1) == Block.(1:0)
        @test blockcolsupport(A, 2) == Block.(1:1)
        @test blockcolsupport(A, 3) == Block.(2:2)

        @test rowsupport(A,1) == 2:3
        @test rowsupport(A,2) == 4:6
        @test rowsupport(A,3) == 4:6
        @test rowsupport(A,4) == 7:10

        @test colsupport(A,1) == 1:0
        @test colsupport(A,2) == 1:1
        @test colsupport(A,3) == 1:1
        @test colsupport(A,4) == 2:3

        l , u = -1,1
        λ , μ = -1,1
        rows = 1:5
        cols = 1:6

        data = reshape(Vector{Float64}(1:(λ+μ+1)*(l+u+1)*sum(cols)), (λ+μ+1, (l+u+1)*sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))
        @test_throws BandError A[1,1] = 5

        @test blockcolsupport(A, 1) == Block.(1:0)
        @test blockcolsupport(A, 2) == Block.(1:1)
        @test blockcolsupport(A, 3) == Block.(2:2)

        @test colsupport(A,1) == 1:0
        @test colsupport(A,2) == 1:1
        @test colsupport(A,3) == 1:1
        @test colsupport(A,4) == 2:3
    end

    @testset "block indexing" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 4
        cols = rows = 1:N
        data = reshape(collect(1:(λ+μ+1)*(l+u+1)*sum(cols)), ((λ+μ+1)*(l+u+1), sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))

        @test A[Block(1), Block(1)] isa BandedMatrix
        @test A[Block(1), Block(1)] == A[Block(1,1)] == BlockArrays.getblock(A, 1, 1) == BandedMatrix(view(A, Block(1,1)))
        @test A[1,1] == view(A,Block(1),Block(1))[1,1] == view(A,Block(1,1))[1,1] == A[Block(1,1)][1,1]  == A[Block(1),Block(1)][1,1] == 5
        @test A[2,1] == view(A,Block(2),Block(1))[1,1] == view(A,Block(2,1))[1,1] == 8
        @test A[3,1] == view(A,Block(2),Block(1))[2,1] == 9
        @test A[4,1] == 0
        @test A[1,2] == view(A,Block(1,2))[1,1] == 11
        @test A[1,3] == view(A,Block(1,2))[1,2] == view(A,Block(1,2))[2] == 19

        @test bandwidths(view(A, Block(3),Block(1))) == (-720,-720)
        @test isempty(bandeddata(view(A, Block(3),Block(1))))

        @test A[Block(3,1)] == view(A, Block(3),Block(1)) == zeros(3,1)
        @test A[Block(3,1)] ≈ view(A, Block(3),Block(1)) ≈ zeros(3,1)
        @test_throws BandError view(A, Block(3),Block(1))[1,1] = 4
        @test_throws BlockBoundsError view(A, Block(5,1))
    end

    @testset "indexing" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 10
        cols = rows = 1:N
        data = reshape(Vector{Float64}(1:(λ+μ+1)*(l+u+1)*sum(cols)), ((λ+μ+1)*(l+u+1),sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))

        A[1,1] = 5
        @test A[1,1] == 5

        @test_throws BandError A[1,4] = 5
        A[1,4] = 0
        @test A[1,4] == 0

        # TODO: return a BandedMatrix
        @test A[1:10,1:10] ≈ Matrix(A)[1:10,1:10]
    end

    @testset "banded matrix interface for blocks" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 10
        cols = rows = 1:N
        data = reshape(collect(1:(λ+μ+1)*(l+u+1)*sum(cols)), ((λ+μ+1)*(l+u+1), sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))

        # test blocks
        V = view(A, Block(1,1))
        @test_throws BoundsError V[2,1]

        # BandedMatrix interface
        @test isbanded(V)
        @test bandwidths(V) == BlockBandedMatrices.subblockbandwidths(A)
        @test V[band(0)] ≈ view(V, band(0)) ≈ A[1:1,1:1]

        # test views of blocks fulfill BnadedMatrix interface
        @test BandedMatrices.inbands_getindex(V, 1, 1) == V[1,1] == 5
        BandedMatrices.inbands_setindex!(V, -1, 1, 1)
        @test A[1,1] == -1
        # these should throw errors but inbands turns it off

        V = view(A, Block(3,4))
        @test V[3,1] == 0
        @test_throws BandError V[3,1] = 5

        view(V, band(0)) .= -3
        @test all(A[Block(3,4)][band(0)] .== -3)

        @test BandedMatrix(V) isa BandedMatrix{Int,Matrix{Int}}
        @test BandedMatrix{Float64}(V) isa BandedMatrix{Float64,Matrix{Float64}}
        @test BandedMatrix{Float64}(BandedMatrix(V)) == BandedMatrix{Float64}(V) ==
                convert(BandedMatrix{Float64}, V)
        @test A[4:6,7:10] ≈ BandedMatrix(V)

        @test A[Block(3,4)].l == A.λ
        @test A[Block(3,4)].u == A.μ

        A[Block(3,4)] = BandedMatrix(Ones{Int}(3,4),(1,1))
        @test A[Block(3,4)] == BandedMatrix(Ones{Int}(3,4),(1,1))

        l , u = 2,1
        λ , μ = 1,2
        N = M = 4
        cols = rows = 1:N
        data = reshape(Vector(1:(λ+μ+1)*(l+u+1)*sum(cols)), (λ+μ+1, (l+u+1)*sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))

        @test_throws BandError A[1,4] = 5
        @test_throws BandError view(A, Block(1,3))[2] = 5

        lu = (l , u) = -1,1
        λμ = (λ , μ) = -1,1
        rows = 1:5
        cols = 1:6

        data = reshape(Vector{Float64}(1:(λ+μ+1)*(l+u+1)*sum(cols)), (λ+μ+1, (l+u+1)*sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))
        @test_throws BandError A[1,1] = 5

        @test A[1,3] == 3
        @test A[Block(2,2)] == [0 0; 0 0]
        @test A[Block(2,3)]  == [0 5 0; 0 0 6]
        @test bandwidths(A[Block(2,3)]) == (-1,1)
    end

    @testset "banded-block-banded matrix interface for blockranges" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 10
        cols = rows = 1:N
        data = reshape(collect(1.0:(λ+μ+1)*(l+u+1)*sum(cols)), ((λ+μ+1)*(l+u+1), sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))
        V = view(A, Block.(2:3), Block.(3:4))
        @test isbandedblockbanded(V)

        B = BandedBlockBandedMatrix(V)
        @test B isa BandedBlockBandedMatrix
        @test blockbandwidths(V) == blockbandwidths(B) == (2,0)
        @test subblockbandwidths(V) == subblockbandwidths(B) == subblockbandwidths(A) == (λ,μ)
        @test B == V == A[Block.(2:3), Block.(3:4)]

        @test A[Block.(2:3), Block.(3:4)] isa BandedBlockBandedMatrix

        x = randn(size(B,2))
        y = similar(x, size(B,1))
        @test all((similar(y) .= MulAdd(B, x)) .=== (similar(y) .= MulAdd(V,x)))
    end

    @testset "BLAS arithmetic" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 10
        cols = rows = fill(1000,N)
        data = reshape(Vector{Float64}(1:(λ+μ+1)*(l+u+1)*sum(cols)), ((λ+μ+1)*(l+u+1), sum(cols)))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))

        V = view(A, Block(N,N))

        AN = A[Block(N,N)]
        BLAS.axpy!(2.0, V, V)
        @test A[Block(N,N)] ≈ 3AN

        Y = zeros(cols[N], cols[N])
        BLAS.axpy!(2.0, V, Y)
        @test Y ≈ 2A[Block(N,N)]

        Y = BandedMatrix(Zeros(cols[N], cols[N]), (λ, μ))
        BLAS.axpy!(2.0, V, Y)
        @test Y ≈ 2A[Block(N,N)]

        Y = BandedMatrix(Zeros(cols[N], cols[N]), (λ+1, μ+1))
        BLAS.axpy!(2.0, V, Y)
        @test Y ≈ 2A[Block(N,N)]

        Y = BandedMatrix(Zeros(cols[N], cols[N]), (0, 0))
        @test_throws BandError BLAS.axpy!(2.0, V, Y)
    end

    @testset "Float32"  begin
        A = BandedBlockBandedMatrix{Float32}(Zeros{Float32}(10,10),
                                    fill(1,10), fill(1,10), (1,1), (1,1))

        @test eltype(A) == Float32

        A = BandedBlockBandedMatrix(Zeros{Float32}(10,10),
                                    fill(1,10), fill(1,10), (1,1), (1,1))

        @test eltype(A) == Float32
    end

    @testset "fill and copy" begin
        l , u = 1,1
        λ , μ = 1,1
        N = M = 10
        cols = rows = 1:N
        data = randn((λ+μ+1)*(l+u+1), sum(cols))
        A = _BandedBlockBandedMatrix(data, rows,cols, (l,u), (λ,μ))

        V = view(A,Block(2,3))
        @test_throws BandError fill!(V, 2.0)
        fill!(V, 0)
        @test A[2:3,4:6] == zeros(2,3)

        dataA = randn((λ+μ+1)*(l+u+1), sum(cols))
        A = _BandedBlockBandedMatrix(copy(dataA), rows,cols, (l,u), (λ,μ))

        K,J = 2,1
        V = view(A,Block(K),Block(J))
        fill!(V, 2.0)
        @test Matrix(V) == fill(2.0, 2,1)
        fill!(view(A,Block(K),Block(J)), 0.0)
        @test Matrix(view(A,Block(K),Block(J))) == zeros(2,1)

        K,J = 3,1
        @test_throws BandError fill!(view(A,Block(K),Block(J)), 2.0)
        fill!(view(A,Block(K),Block(J)), 0.0)
        @test Matrix(view(A,Block(K),Block(J))) == zeros(3,1)

        @test_throws BandError fill!(A, 2.0)
        fill!(A, 0.0)
        @test Matrix(A) == zeros(size(A))

        dataA = randn((λ+μ+1)*(l+u+1), sum(cols))
        A = _BandedBlockBandedMatrix(copy(dataA), rows,cols, (l,u), (λ,μ))

        dataB = randn((λ+μ+3)*(l+u+3), sum(cols))
        B = _BandedBlockBandedMatrix(copy(dataB), rows,cols, (l+1,u+1), (λ+1,μ+1))

        B = _BandedBlockBandedMatrix(copy(dataB), rows,cols, (l+1,u+1), (λ+1,μ+1))
        copyto!(view(B, Block(N,N)), view(A, Block(N,N)))
        @test B[Block(N,N)] == A[Block(N,N)]

        B = _BandedBlockBandedMatrix(copy(dataB), rows,cols, (l+1,u+1), (λ+1,μ+1))
        copyto!(B, A)
        @test norm(B[Block(3,1)]) == 0
        @test  B ≈ A
        @test Matrix(B) ≈ Matrix(A)
    end

    @testset "Zero bands" begin
        B = BandedBlockBandedMatrix{Float64}(undef, 1:5,1:5, (-1,-1), (-1,-1))
        @test Matrix(B) == zeros(size(B))

        B = BandedBlockBandedMatrix{Float64}(undef, 1:5,1:5, (-1,-1), (1,-1))
        @test Matrix(B) == zeros(size(B))

        B = BandedBlockBandedMatrix{Float64}(undef, 1:5,1:5, (1,-1), (-1,-1))
        @test Matrix(B) == zeros(size(B))
    end

    @testset "BandedBlockBanded with BlockMatrix" begin
        WithBlockMatrix{T} = BandedBlockBandedMatrix{T, BlockMatrix{T, Matrix{Matrix{T}}}}
        args = [1, 2, 3], [2, 2, 1], (1, 1), (1, 1)
        A = WithBlockMatrix{Int64}(undef, args...)
        B = BandedBlockBandedMatrix{Int64}(undef, axes(A), blockbandwidths(A), subblockbandwidths(A))

        @test eltype(A) === eltype(B) === Int64
        @test typeof(A.data) <: BlockArray
        @test typeof(B.data) <: PseudoBlockArray
        @test size(A) == size(B)
        @test bandrange(A) == bandrange(B)
        @test blockbandwidths(A) == blockbandwidths(B)
        @test blocksize(A) == blocksize(B)

        A = WithBlockMatrix{Int64}(Zeros{Int64}(sum(args[1]),sum(args[2])), args...)
        B = BandedBlockBandedMatrix{Int64}(Zeros{Int64}(sum(args[1]),sum(args[2])), args...)
        @test typeof(A.data) <: BlockArray
        @test typeof(B.data) <: PseudoBlockArray
        @test A == B

        A = WithBlockMatrix{Int64}(Ones{Int64}(sum(args[1]),sum(args[2])), args...)
        B = BandedBlockBandedMatrix{Int64}(Ones{Int64}(sum(args[1]),sum(args[2])), args...)
        @test typeof(A.data) <: BlockArray
        @test typeof(B.data) <: PseudoBlockArray
        @test A == B
        @test (A .+ 1) .* 2 == B .* 2 .+ 2
    end

    @testset "Sparse dimensions" begin
        A = BandedBlockBandedMatrix{Float64}(undef, 1:5,1:5, (-1,1), (-1,1))
        @test size(sparse(A)) == size(A) == (15,15)
    end

    @testset "non-standard blocks" begin
        A = BandedBlockBandedMatrix{Float64}(undef, Int[], 1:5,(-1,1), (-1,1))
        @test BlockBandedMatrices.colstart(A,1) == 1
        A = BandedBlockBandedMatrix{Float64}(undef, 1:2, 1:5,(-1,1), (-1,1))
        A.data .= randn.()
        V = view(A, Block(2,3))
        @test MemoryLayout(typeof(V)) == BandedMatrices.BandedColumnMajor()
        @test isbanded(V)
        @test bandwidths(V) == (-1,1)
        @test BandedMatrix(V) == A[2:3,4:6]
        V2 = view(V, :, 2:3)
        @test MemoryLayout(typeof(V2)) == BandedMatrices.BandedColumnMajor()
        @test bandwidths(V2) == (0,0)
        @test BandedMatrix(V2) == A[2:3,5:6]
        V2 = view(V, 2:2, :)
        @test MemoryLayout(typeof(V2)) == BandedMatrices.BandedColumnMajor()
        @test bandwidths(V2) == (-2,2)
        @test BandedMatrix(V2) == A[3:3,4:6]
        V2 = view(V, 2:2, 2:3)
        @test MemoryLayout(typeof(V2)) == BandedMatrices.BandedColumnMajor()
        @test bandwidths(V2) == (-1,1)
        @test BandedMatrix(V2) == A[3:3,5:6]
    end

    @testset "OneTo axes" begin
        ret = BandedBlockBandedMatrix{Int}(undef, (Base.OneTo(5),Base.OneTo(5)), (0,0), (1,1))
        V = view(ret, Block(1,1))
        @test bandwidths(V) == (1,1)
        A = brand(5,5,1,1)
        B = BandedBlockBandedMatrix(A)
        @test A == B
        @test subblockbandwidths(B) == (1,1)
        @test blockbandwidths(B) == (0,0)
    end
end

if false # turned off since tests have check-bounds=yes
    # test that @inbounds is working properly
    exceed_band(V, k, j) = @inbounds return V[k,j]
    @test exceed_band(V, 2,1) == 8

    @test BandedMatrices.inbands_getindex(V, 2, 1) == 8
    BandedMatrices.inbands_setindex!(V, -2, 5, 1)
    @test A[2,1] == -2
end


immutable DataFunction{T<:Number,D<:Number,A<:AbstractArray} <: AbstractData
    index::Vector
    values::A
    modifyvalues::Bool
    generator::Function
    args::Tuple
    DataFunction(i::AbstractVector{T},d::AbstractArray{D},w::Bool,g::Function,args...) = new(collect(i),d,w,g,args)
end

function _data(::Type{Val{:function!}},i::AbstractVector,d::AbstractArray,g::Function,args...)
    @assert length(i) == size(d,1)
    DataFunction{eltype(i),eltype(d),typeof(d)}(i,d,true,g,args...)
end

function _data(::Type{Val{:function}},i::AbstractVector,g::Function,args...)
    d = g(args...)::AbstractArray
    @assert length(i) == size(d,1)
    DataFunction{eltype(i),eltype(d),typeof(d)}(i,d,false,g,args...)
end

@inline _generate(::Type{Val{true}},d::DataFunction) = d.generator(d.values,d.args...)
@inline _generate(::Type{Val{false}},d::DataFunction) = copy!(d.values,d.generator(d.args...))

numvalues(d::DataFunction) = size(d.values,1)
numvars(d::DataFunction) = size(d.values,2)
eltype(d::DataFunction) = eltype(d.values)
generate!(d::DataFunction) = (_generate(Val{d.modifyvalues},d) ; d.values)
dataindex(d::DataFunction) = d.index
datavalues(d::DataFunction) = d.values
datatypename(d::DataFunction) = "DataFunction"


# The original license on the t-tests from HypothesisTests.jl
# is as follows. We're including it here because we're mostly copying the
# documentation of the functions.
#
# kolmogorov_smirnov.jl
# Kolmogorov–Smirnov
#
# Copyright (C) 2014   Christoph Sawade
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import HypothesisTests.ExactOneSampleKSTest
import HypothesisTests.ApproximateTwoSampleKSTest

"""
    ExactOneSampleKSTest(uv::AbstractUncertainValue,
        d::UnivariateDistribution, n::Int = 1000) -> ExactOneSampleKSTest

Perform a one-sample exact Kolmogorov–Smirnov test of the null hypothesis that
a draw of `n` realisations of the uncertain value `uv` comes from the
distribution `d` against the alternative hypothesis that the sample is not
drawn from `d`.
"""
function ExactOneSampleKSTest(uv::AbstractUncertainValue,
        d::UnivariateDistribution, n::Int = 1000)
    x = resample(uv, n)
    ExactOneSampleKSTest(x, d)
end

"""
    ExactOneSampleKSTestPooled(ud::UncertainDataset,
        d::UnivariateDistribution, n::Int = 1000) -> ExactOneSampleKSTest

First, draw `n` realisations of each uncertain value in `ud` and pool them
together. Then perform a one-sample exact Kolmogorov–Smirnov test of the null
hypothesis that the pooled values comes from the distribution `d` against the
alternative hypothesis that the sample is not drawn from `d`.
"""
function ExactOneSampleKSTestPooled(ud::UncertainDataset,
        d::UnivariateDistribution, n::Int = 1000)
    x = vcat(resample(ud, n)...,)
    ExactOneSampleKSTest(x, d)
end

"""
    ExactOneSampleKSTestElementWise(ud::UncertainDataset,
        d::UnivariateDistribution, n::Int = 1000) -> Vector{ExactOneSampleKSTest}

First, draw `n` realisations of each uncertain value in `ud`, keeping one
pool of values for each uncertain value.

Then, perform an element-wise (pool-wise) one-sample exact Kolmogorov–Smirnov
test of the null hypothesis that each value pool comes from the distribution `d`
against the alternative hypothesis that the sample is not drawn from `d`.
"""
function ExactOneSampleKSTestElementWise(ud::UncertainDataset,
        d::UnivariateDistribution, n::Int = 1000)
    [ExactOneSampleKSTest(resample(ud[i], n), d) for i = 1:length(ud)]
end

"""
    ApproximateTwoSampleKSTest(d1::AbstractUncertainValue,
        d2::AbstractUncertainValue, n::Int = 1000) -> ApproximateTwoSampleKSTest

Perform an asymptotic two-sample Kolmogorov–Smirnov-test of the null hypothesis
that the distribution furnishing the uncertain value `d1` represent the same
distribution as the distribution furnishing the uncertain value `d2`
against the alternative hypothesis that the furnishing distributions are
different.
"""
function ApproximateTwoSampleKSTest(d1::AbstractUncertainValue,
                                    d2::AbstractUncertainValue,
                                    n::Int = 1000)
    x = resample(d1, n)
    y = resample(d2, n)
    ApproximateTwoSampleKSTest(x, y)
end

"""
    ApproximateTwoSampleKSTestPooled(d1::UncertainDataset,
        d2::UncertainDataset, n::Int = 1000) -> ApproximateTwoSampleKSTest

First, draw `n` realisations of each uncertain value in `d1`, then separately
draw `n` realisations of each uncertain value in `d2`. Then, pool all
realisations for `d1` together and all realisations of `d2` together.

On the pooled realisations, perform an asymptotic two-sample
Kolmogorov–Smirnov-test of the null hypothesis that the distribution furnishing
the `d1` value pool represents the same distribution as the distribution
furnishing the `d2` value pool, against the alternative hypothesis that the
furnishing distributions are different.
"""
function ApproximateTwoSampleKSTestPooled(d1::UncertainDataset,
                                        d2::UncertainDataset,
                                        n::Int = 1000)
    x = vcat(resample(d1, n)...,)
    y = vcat(resample(d2, n)...,)
    ApproximateTwoSampleKSTest(x, y)
end

"""
    ApproximateTwoSampleKSTestElementWise(d1::UncertainDataset,
        d2::UncertainDataset, n::Int = 1000) -> Vector{ApproximateTwoSampleKSTest}

Assuming `d1` and `d2` contain the same number of uncertain observations,
draw `n` realisations of each uncertain value in `d1`, then separately and
separately draw `n` realisations of each uncertain value in `d2`.

Then, perform an asymptotic two-sample Kolmogorov–Smirnov-test of
the null hypothesis that the uncertain values in `d1` and `d2` come from
the same distribution against the alternative hypothesis that the (element-wise)
values in  `d1` and `d2` come from different distributions.

The test is performed pairwise, i.e. ApproximateTwoSampleKSTest(d1[i], d2[i])
with `n` draws for the ``i``-ith pair of uncertain values.
"""
function ApproximateTwoSampleKSTestElementWise(d1::UncertainDataset,
                                                d2::UncertainDataset,
                                                n::Int = 1000)
    N = length(d1)

    [ApproximateTwoSampleKSTest(resample(d1[i], n), resample(d2[i], n)) for i = 1:N]
end

export
ExactOneSampleKSTest,
ExactOneSampleKSTestPooled,
ExactOneSampleKSTestElementWise,

ApproximateTwoSampleKSTest,
ApproximateTwoSampleKSTestPooled,
ApproximateTwoSampleKSTestElementWise


using LazIO

function countsgrid(bbox, cellsizes)
    min_x, min_y, min_z, max_x, max_y, max_z = bbox
    rows = Int(cld(max_x - min_x, cellsizes[1]))
    cols = Int(cld(max_y - min_y, cellsizes[2]))
    # heights = Int(cld(max_z - min_z, cellsizes[3]))
    zeros(Int64, rows, cols) #, heights)
end


include("../utils.jl")
include("group_utils.jl")

@testset "Heisenberg group" begin
    G = HeisenbergGroup(1)
    @test repr(G) == "HeisenbergGroup(1)"

    pts = [
        [1.0 2.0 3.0; 0.0 1.0 -1.0; 0.0 0.0 1.0],
        [1.0 4.0 -3.0; 0.0 1.0 3.0; 0.0 0.0 1.0],
        [1.0 -2.0 1.0; 0.0 1.0 1.1; 0.0 0.0 1.0],
    ]
    Xpts = [
        [0.0 2.0 3.0; 0.0 0.0 -1.0; 0.0 0.0 0.0],
        [0.0 4.0 -3.0; 0.0 0.0 3.0; 0.0 0.0 0.0],
        [0.0 -2.0 1.0; 0.0 0.0 1.1; 0.0 0.0 0.0],
    ]

    ge = allocate(pts[1])
    identity_element!(G, ge)
    @test isapprox(ge, I; atol=1e-10)

    @test check_point(G, [0.0 2.0 3.0; 0.0 1.0 -1.0; 0.0 0.0 1.0]) isa DomainError
    @test check_point(G, [1.0 2.0 3.0; 0.0 1.0 -1.0; 1.0 0.0 1.0]) isa DomainError
    @test check_point(G, [1.0 2.0 3.0; 0.0 2.0 -1.0; 0.0 0.0 1.0]) isa DomainError
    @test check_point(G, [1.0 2.0 3.0; 0.0 1.0 -1.0; 0.0 0.0 2.0]) isa DomainError
    @test check_point(G, [1.0 2.0 3.0; 0.0 1.0 -1.0; 0.0 1.0 1.0]) isa DomainError
    @test check_point(G, [1.0 2.0 3.0; 1.0 1.0 -1.0; 0.0 0.0 1.0]) isa DomainError

    @test check_vector(G, pts[1], [1.0 2.0 3.0; 0.0 0.0 -1.0; 0.0 0.0 0.0]) isa DomainError
    @test check_vector(G, pts[1], [0.0 2.0 3.0; 1.0 0.0 -1.0; 0.0 0.0 0.0]) isa DomainError
    @test check_vector(G, pts[1], [0.0 2.0 3.0; 0.0 0.0 -1.0; 0.0 0.0 2.0]) isa DomainError

    @test compose(G, pts[1], pts[2]) ≈ pts[1] * pts[2]
    test_group(
        G,
        pts,
        Xpts,
        Xpts;
        test_diff=true,
        test_invariance=true,
        test_lie_bracket=true,
        test_adjoint_action=true,
    )

    test_manifold(
        G,
        pts;
        basis_types_to_from=(DefaultOrthonormalBasis(),),
        parallel_transport=true,
        test_injectivity_radius=true,
        test_project_point=true,
        test_project_tangent=true,
        test_musical_isomorphisms=false,
    )
end


immutable ShiftedTensorOperator{S,P,T}<:AbstractTensorMap{S,P,T}
    map::AbstractTensorMap{S,P}
    shift::T
    function ShiftedTensorOperator(A::AbstractTensorMap{S,P},shift::T)
        domain(A)==codomain(A) || throw(SpaceError("Not an operator, i.e. domain != codomain"))
        promote_type(eltype(A),T)==T || throw(InexactError())
        new(A,shift)
    end
end
ShiftedTensorOperator{S,P,T1,T2}(A::AbstractTensorMap{S,P,T1},shift::T2)=ShiftedTensorOperator{S,P,promote_type(T1,T2)}(A,convert(promote_type(T1,T2),shift))

# create from adding scalars to tensormaps
+(A::ShiftedTensorOperator,shift::Number)=shift==0 ? A : ShiftedTensorOperator(A.map,A.shift+shift)
+(shift::Number,A::ShiftedTensorOperator)=shift==0 ? A : ShiftedTensorOperator(A.map,A.shift+shift)
+(A::AbstractTensorMap,shift::Number)=shift==0 ? A : ShiftedTensorOperator(A,shift)
+(shift::Number,A::AbstractTensorMap)=shift==0 ? A : ShiftedTensorOperator(A,shift)
-(A::ShiftedTensorOperator,shift::Number)=shift==0 ? A : ShiftedTensorOperator(A.map,A.shift-shift)
-(shift::Number,A::ShiftedTensorOperator)=shift==0 ? -A : ShiftedTensorOperator(-A.map,-A.shift+shift)
-(A::AbstractTensorMap,shift::Number)=shift==0 ? A : ShiftedTensorOperator(A,-shift)
-(shift::Number,A::AbstractTensorMap)=shift==0 ? -A : ShiftedTensorOperator(-A,shift)

+(A::AbstractTensorMap,shift::UniformScaling)=+(A,shift[1,1])
+(shift::UniformScaling,A::AbstractTensorMap)=+(shift[1,1],A)
-(A::AbstractTensorMap,shift::UniformScaling)=-(A,shift[1,1])
-(shift::UniformScaling,A::AbstractTensorMap)=-(shift[1,1],A)

# properties
domain(A::ShiftedTensorOperator)=domain(A.map)
codomain(A::ShiftedTensorOperator)=codomain(A.map)

Base.isreal(A::ShiftedTensorOperator)=isreal(A.map) && isreal(A.shift)
Base.issym(A::ShiftedTensorOperator)=issym(A.map)
Base.ishermitian(A::ShiftedTensorOperator)=ishermitian(A.map)
Base.isposdef(A::ShiftedTensorOperator)=isposdef(A.map) && isposdef(A.shift)

# comparison
==(A::ShiftedTensorOperator,B::ShiftedTensorOperator)=(A.map==B.map && A.shift==B.shift)

# multiplication with vector
Base.A_mul_B!{S,P}(y::AbstractTensor{S,P},A::ShiftedTensorOperator{S,P},x::AbstractTensor{S,P})=(Base.A_mul_B!(y,A.map,x);Base.LinAlg.axpy!(A.shift,x,y))
Base.At_mul_B!{S,P}(y::AbstractTensor{S,P},A::ShiftedTensorOperator{S,P},x::AbstractTensor{S,P})=(Base.At_mul_B!(y,A.map,x);Base.LinAlg.axpy!(A.shift,x,y))
Base.Ac_mul_B!{S,P}(y::AbstractTensor{S,P},A::ShiftedTensorOperator{S,P},x::AbstractTensor{S,P})=(Base.Ac_mul_B!(y,A.map,x);Base.LinAlg.axpy!(conj(A.shift),x,y))


"""
    subspace_construction(model, cost, data, opt; T = 10, c = 1, M = 3, print_freq = 1)

Izmailov, P., Maddox, W. J., Kirichenko, P., Garipov, T., Vetrov, D., & Wilson, A. G. (2020, August). Subspace inference for Bayesian deep learning.
In Uncertainty in Artificial Intelligence (pp. 1169-1179). PMLR.

To construct subspace from pretrained weights.

# Input Arguments
- `model` 	 : Machine learning model. Eg: Chain(Dense(10,2)). Model should be created with Chain in Flux
- `cost`  	 : Cost function. Eg: L(x, y) = Flux.Losses.mse(m(x), y)
- `data` 	 : Inputs and outputs. Eg:	X = rand(10,100); Y = rand(2,100); data = DataLoader(X,Y);
- `opt`		 : Optimzer. Eg: opt = ADAM(0.1)

# Keyword Arguments
- `T` 		  : Number of steps for subspace calculation. Eg: T= 1
- `c` 		  : Moment update frequency. Eg: c = 1
- `M` 		  : Maximum number of columns in deviation matrix. Eg: M= 2
- `print_freq`: Loss printing frequency

# Outputs
- `W_swa`    : Mean weights
- `P` 		 : Projection Matrix
"""
function subspace_construction(model, cost, data, opt; T = 10, c = 1, M = 3, print_freq = 1
)
	training_loss = 0.0

	ps = Flux.params(model)
	W_swa = zeros(length(extract_params(ps)))
	all_len = length(W_swa)
	A = Array{eltype(W_swa)}(undef,0) #initialize deviation matrix

	# #initaize weights with mean
	all_weights = []
	for i in 1:T
		for d in data
			gs = gradient(ps) do
				training_loss = cost(model, d...)
				return training_loss
			end
			Flux.update!(opt, ps, gs)
			if mod(i,c) == 0
				W = extract_params(ps)
				n = i/c
				W_swa = (n.*W_swa + W)./(n+1)
				# if(length(A) >= M*all_len)
				# 	A = A[1:(end - all_len)]
				# end
				W_dev =  W - W_swa
				append!(A, W_dev)
			end
		end
		#print loss based o print frequency
		if (mod(i,print_freq) == 0 )|| (i == T)
			println("Traing loss: ", training_loss," Epoch: ", i)
		end
	end

	A = reshape(A, all_len, :)
	#take pseudo SVD and take only first M columns
	U,s,V = psvd(A)
	#generate projection matrix
	P = U[:,1:M]*LinearAlgebra.Diagonal(s[1:M])
	return W_swa, P
end

"""
	auto_encoder_subspace(model, cost, data, opt, encoder, decoder; T = 10, c = 1, M = 3, print_freq = 1
)

To construct subspace from pretrained weights using autoencoders.

# Input Arguments
- `model` 	 : Machine learning model. Eg: Chain(Dense(10,2)). Model should be created with Chain in Flux
- `cost`  	 : Cost function. Eg: L(x, y) = Flux.Losses.mse(m(x), y)
- `data` 	 : Inputs and outputs. Eg:	X = rand(10,100); Y = rand(2,100); data = DataLoader(X,Y);
- `opt`		 : Optimzer. Eg: opt = ADAM(0.1)
- `encoder`	 : Encoder to generate subspace from NN or Neural ODE parameters
- `decoder`	 : Decoder to generate NN or Neural ODE parameters from subspace

# Keyword Arguments
- `T` 		  : Number of steps for subspace calculation. Eg: T= 1
- `c` 		  : Moment update frequency. Eg: c = 1
- `M` 		  : Maximum number of columns in deviation matrix. Eg: M= 2
- `print_freq`: Loss printing frequency

# Outputs
- `W_swa`    : Mean weights
- `decoder`  : Trained decoder to generate NN or Neural ODE parameters from subspace
"""
function auto_encoder_subspace(model, cost, data, opt, encoder, decoder; T = 10, c = 1, M = 3, print_freq = 1
)
	training_loss = 0.0

	ps = Flux.params(model)
	W_swa = zeros(length(extract_params(ps)))
	all_len = length(W_swa)
	A = Array{eltype(W_swa)}(undef,0) #initialize deviation matrix

	# #initaize weights with mean
	all_weights = []
	for i in 1:T
		for d in data
			gs = gradient(ps) do
				training_loss = cost(model, d...)
				return training_loss
			end
			Flux.update!(opt, ps, gs)
			if mod(i,c) == 0
				W = extract_params(ps)
				n = i/c
				W_swa = (n.*W_swa + W)./(n+1)
				W_dev =  W - W_swa
				append!(A, W_dev)
			end
		end
		#print loss based o print frequency
		if (mod(i,print_freq) == 0 )|| (i == T)
			println("Traing loss: ", training_loss," Epoch: ", i)
		end
	end

	re_weight = reshape(A, all_len, :)
	# model
	sae = Chain(encoder,decoder)
	autops = Flux.params(sae);
	# loss
	autoloss(X) = Flux.mse(sae(X), X)
	+ sum(sqnorm, autops)

	# optimization
	opt = ADAM()
	# parameters

	cb() = @show autoloss(re_weight)
	#data
	wdata = DataLoader(re_weight, batchsize = 5, shuffle = true)
	Flux.train!(autoloss, autops, wdata, opt; cb=cb)
	return W_swa, decoder

end

"""
    diffusion_subspace(model, cost, data, opt; T = 10, c = 1, M = 3, print_freq = 1)

To construct subspace from pretrained weights using diffusion maps.

# Input Arguments
- `model` 	 : Machine learning model. Eg: Chain(Dense(10,2)). Model should be created with Chain in Flux
- `cost`  	 : Cost function. Eg: L(x, y) = Flux.Losses.mse(m(x), y)
- `data` 	 : Inputs and outputs. Eg:	X = rand(10,100); Y = rand(2,100); data = DataLoader(X,Y);
- `opt`		 : Optimzer. Eg: opt = ADAM(0.1)

# Keyword Arguments
- `T` 		  : Number of steps for subspace calculation. Eg: T= 1
- `c` 		  : Moment update frequency. Eg: c = 1
- `M` 		  : Maximum number of columns in deviation matrix. Eg: M= 2
- `print_freq`: Loss printing frequency

# Outputs
- `W_swa`    : Mean weights
- `P` 		 : Projection Matrix
"""
function diffusion_subspace(model, cost, data, opt; T = 10, c = 1, M = 3, print_freq = 1
)
	training_loss = 0.0

	ps = Flux.params(model)
	W_swa = zeros(length(extract_params(ps)))
	all_len = length(W_swa)
	A = Array{eltype(W_swa)}(undef,0) #initialize deviation matrix

	# #initaize weights with mean
	all_weights = []
	for i in 1:T
		for d in data
			gs = gradient(ps) do
				training_loss = cost(model, d...)
				return training_loss
			end
			Flux.update!(opt, ps, gs)
			if mod(i,c) == 0
				W = extract_params(ps)
				n = i/c
				W_swa = (n.*W_swa + W)./(n+1)
				# if(length(A) >= M*all_len)
				# 	A = A[1:(end - all_len)]
				# end
				W_dev =  W - W_swa
				append!(A, W_dev)
			end
		end
		#print loss based o print frequency
		if (mod(i,print_freq) == 0 )|| (i == T)
			println("Traing loss: ", training_loss," Epoch: ", i)
		end
	end

	A = reshape(A, all_len, :)
	#diffusion map
	diff_M = fit(DiffMap, A', maxoutdim=M) # construct diffusion map model
	P = transform(diff_M)
	return W_swa, P'
end


module utilitiesOCIM

# package code goes here
"""
    Iabove, Ibelow = buildIndexMatrices()

Build the matrices of indices above and below.

These are shifted-diagonal matrices that allow for
easy construction of the particle flux divergence `PFD` operator.
"""
function buildIaboveIbelow()
  n = nlon * nlat * (ndepth+1)
  In = speye(n)
  idx = zeros(nlat,nlon,ndepth+1)
  idx[:] = 1:n
  idxabove = idx[:,:,[ndepth+1; 1:ndepth]] # use a periodic upward shift
  idxbelow = idx[:,:,[2:ndepth+1; 1]]      # use a periodic downward shift
  Iabove = In[idxabove[:],:]
  Ibelow = In[idxbelow[:],:]
  # keep only indices of wet boxes
  Iabove = Iabove[iwet,iwet]
  Ibelow = Ibelow[iwet,iwet]
  return Iabove, Ibelow
end
const Iabove, Ibelow = buildIaboveIbelow()

"""
    Iabove = buildIabove(wet3d)

Build the shifted-diagonal sparse matrix of the indices of above neighbours.

Iabove[i,j] = 1 if the box represented by the linear index i
lies directly above the box represented by the linear index j.
"""
function buildIabove(wet3d)
  nlat, nlon, ndepth = size(wet3d)
  iwet = find(wet3d)
  n = nlon * nlat * (ndepth + 1)
  In = speye(n)
  idx = zeros(nlat, nlon, ndepth + 1)
  idx[:] = 1:n
  idx .= idx[:, :, [ndepth + 1 ; 1:ndepth]] # shift
  return In[idx[:], :][iwet, iwet]
end

function buildIabove(wet3d)

end

end # module


abstract RedundantWaveletClass

immutable Gammatone <: RedundantWaveletClass
    order::Int
end
Gammatone() = Gammatone(4)
isdyadic(::Gammatone) = false

immutable MexicanHat <: RedundantWaveletClass end
isdyadic(::MexicanHat) = true
issteerable(::MexicanHat) = false

uncertainty(::MexicanHat) = 2.0

immutable Morlet <: RedundantWaveletClass end

"""In the special case `n_filters_per_octave=1`, we manually set `ξ=0.39`. That
is more accurate with the Littlewood-Paley energy conservation criterion than
the generic fallback `ξ=0.4`, which is only valid when the wavelet has a
symmetric profile in the Fourier domain. This is no longer the case for
n_filters_per_octave==max_qualityfactor==1 as the Morlet low-frequency corrective
term is no longer negligible."""
default_motherfrequency(class::Morlet, n_filters_per_octave) =
    n_filters_per_octave==1 ? 0.39 : inv(3.0 - exp2(-1.0/n_filters_per_octave))

isdyadic(::Morlet) = false
issteerable(::Morlet) = true

"""
By neglecting the low-frequency corrective term, we write the Morlet wavelet as
a Gaussian of variance σ in the Fourier domain. Its 3 dB bandwidth, defined as
the full width at half maximum (FWHM) of the squared magnitude in the Fourier
domain, is then equal to b = 2σ*sqrt(log(2)).

Therefore, for a given center frequency ω and a quality factor Q, the variance
σ of the Gaussian is equal to σ = b / (2*sqrt(log(2))). In the spatial domain,
this amounts to a Gabor wavelet (a Gaussian modulated by a sine wave, without
any low-frequency corrective term) of variance 1/σ. Its spatial scale (FWTM)
is equal to s = 2*sqrt(log(10))/σ. We conclude that the uncertainty
constant of the Morlet wavelet, defined as the (scale*bandwidth) constant,
is equal to
    h = b*s = sqrt(log(10)/log(2)) = 1.8226...
"""
uncertainty(class::Morlet) = sqrt(log(10.0) / log(2.0))


# A.*B

FAhadamardprodB(A,B)=(A.*B,nothing)

function FAhadamardprodB_inplace(value,aux,A,B)
    copy!(value,A.*B)
end

function DAhadamardprodB(derivativeIDX,f_c,faux_c,grad_c,grad_n,A,B)
    if derivativeIDX==1
        axpy!(1.0,grad_c.*B,grad_n)
    elseif derivativeIDX==2
        axpy!(1.0,grad_c.*A,grad_n)
    end
end

if (PROC=="GPU") || (PROC=="GPU32")

    function FAhadamardprodB(A::CudaArray,B::CudaArray)
        tmp=CudaArray(Float64,size(A))
        vmult!(1.0,A,B,tmp)
        return (tmp,nothing)
    end

    function FAhadamardprodB_inplace(value::CudaArray,aux,A::CudaArray,B::CudaArray)
        vmult!(1.0,A,B,value)
    end

    function DAhadamardprodB(derivativeIDX,f_c,faux_c,grad_c,grad_n,A::CudaArray,B::CudaArray)
    if derivativeIDX==1
        vmultupdate!(1.0,grad_c,B,grad_n)
    elseif derivativeIDX==2
        vmultupdate!(1.0,grad_c,A,grad_n)
    end
end

end

Derivative[FAhadamardprodB]=DAhadamardprodB
Inplace[FAhadamardprodB]=FAhadamardprodB_inplace

import Base. .*
.*(A::ADnode,B::ADnode)=ADnode(FAhadamardprodB,[A B])

export .*

export FAhadamardprodB


function trapezoidal_quadrature(w::Matrix{RT}, v::Union{Array{RT,2},SharedArray{RT,2}}, grid::Grid2d{M,N,RT}) where {M,N,RT}
    @assert size(v) == size(w)

    local result::RT = 0

    # @inbounds for j in 1:size(v,2)
    #     for i in 1:size(v,1)
    #         result += w[i,j] * v[i,j]
    #     end
    # end

    @inbounds for i in eachindex(v,w)
          result += w[i] * v[i]
    end

    result * (grid.x2 - grid.x1) / M * (grid.y2 - grid.y1) / N
end

function trapezoidal_quadrature(w::Matrix{Matrix{RT}}, v::Union{Array{RT,2},SharedArray{RT,2}}, grid::Grid2d{M,N,RT}) where {M,N,RT}
    local result::Matrix{CT} = zeros(RT, size(w,1), size(w,2))

    for l in 1:size(w,2)
        for k in 1:size(w,1)
            result[k,l] = trapezoidal_quadrature(w[k,l], v, grid)
        end
    end

    result
end

function trapezoidal_quadrature(w::Matrix{Matrix{RT}}, v::Union{Vector{Array{RT,2}},Vector{SharedArray{RT,2}}}, grid::Grid2d{M,N,RT}) where {M,N,RT}
    @assert size(w,2) == length(v)

    local result::Vector{CT} = zeros(RT, size(w,1))

    for k in 1:size(w,1)
        for l in 1:size(w,2)
            result[k] += trapezoidal_quadrature(w[k,l], v[l], grid)
        end
    end

    result
end


using QSM: lap, wrapped_laplacian

@testset "Wrapped Laplacian" for T in (Float32, Float64)
    Random.seed!(501)
    for sz in [(11,12,13), (12,13,14,5)]
        x, dx = rand(T, sz), Tuple(rand(T, 3))
        @test wrapped_laplacian(x, dx) ≈ lap(x, dx)
    end
end


function ilqrallocs(solver)
    allocs = @allocated J_prev = TO.cost(solver, solver.Z̄)

    # Calculate expansions
    # TODO: do this in parallel
    allocs += @allocated Altro.errstate_jacobians!(solver.model, solver.G, solver.Z̄)
    allocs += @allocated Altro.dynamics_expansion!(solver, solver.Z̄)
    allocs += @allocated Altro.cost_expansion!(solver.obj, solver.Efull, solver.Z̄)
    allocs += @allocated Altro.error_expansion!(solver.model, solver.Eerr, solver.Efull, solver.G, solver.Z̄)

    # Get next iterate
    allocs += @allocated Altro.backwardpass!(solver)
    allocs += @allocated Jnew = Altro.forwardpass!(solver, J_prev)

    # Accept the step and update the current trajectory
    # This is kept out of the forward pass function to make it easier to
    # benchmark the forward pass
    allocs += @allocated copyto!(solver.Z, solver.Z̄)

    # Calculate the gradient of the new trajectory
    dJ = J_prev - Jnew
    grad = Altro.gradient!(solver)

    # Record the iteration
    Altro.record_iteration!(solver, Jnew, dJ, grad)

    # Check convergence
    exit = Altro.evaluate_convergence(solver)
    return allocs
end

@testset "iLQR Solver Allocations" begin
solver = Altro.iLQRSolver(Problems.Cartpole()..., use_static=Val(true))
ilqrallocs(solver)
@test ilqrallocs(solver) == 0
solver = Altro.iLQRSolver(Problems.Cartpole()..., use_static=Val(false))
ilqrallocs(solver)
@test ilqrallocs(solver) == 0
end


"""
    resample2(v, rate)

Performs time-domain resampling of `v` by the ratio `rate` using cubic interpolation.
When downsampling the signal is lowpassed before resampling. When up-sampling it is
lowpassed afterwards to remove interpolation noise.

Similarities to `DSP.resample`

- The signals are aligned such that the first sample of `v` corresponds to the first
  sample of `resample2(v, rate)`.

Differences from `DSP.resample`:

- this guarantees that `length(resample2(v, rate)) == round(Int, length(v)*rate)`
- `resample2` uses :reflect_at boundary conditions to avoid ringing when the signal
  starts far from 0.
"""
function resample2(v, rate, lpf=resample_filter2(rate))
    L = length(v)
    Lf = length(lpf)
    @assert isodd(Lf)
    fcomp = Lf÷2
    # reflecting adds Lf-1 samples
    # convolution adds Lf-1 samples
    if rate <= 1
        vf = conv(lpf, pad(v, :reflect_at, fcomp))
        map(0:round(Int, L*rate)-1) do i
            cubic(vf, i/rate+Lf, :reflect_at)
        end
    else
        out = map(0:round(Int, L*rate)-1) do i
            cubic(v, i/rate+1, :reflect_at)
        end
        conv(lpf, pad(out, :reflect_at, fcomp))[Lf:end-Lf+1]
    end
end

function pad(v, bc, pre, post=pre)
    # TODO: support padding longer than the signal, which is useful for large resampling
    # factors
    prepad, postpad = if bc == :reflect_at
        reverse(v[2:pre+1]), reverse(v[end-post:end-1])
    else
        throw(ArgumentError("Unsupported boundary condition type"))
    end

    [prepad; v; postpad]
end

function resample_filter2(rate)
    # DSP.jl doesn't force FIRWindow to be odd-length, so we do that here
    # see https://github.com/JuliaDSP/DSP.jl/issues/345
    cutoff = rate <= 1 ? rate : 1/rate
    ord, α = kaiserord(0.1*cutoff)
    iseven(ord) && (ord += 1)
    digitalfilter(Lowpass(cutoff), FIRWindow(kaiser(ord, α)))
end

# generate data without information too close to nyquist
# x = resample2(rand(100).+10, 2)
# r = 0.5
# plot(0:length(x)-1, x)
# x2 = resample2(x, r)
# plot!((0:length(x2)-1)/r, x2)
#
# x = randn(100000)
# @benchmark resample(x, 0.05)
# @benchmark resample2(x, 0.05)
# currently `resample2` is about 4-5x faster than `resample` for this case


export SymReal

struct SymReal <: Real
    ex::Expression
end

SymReal(name::Symbol) = SymReal(Variable(name))
SymReal(x::SymReal) = x
SymReal(x::Real) = SymReal(Numeric(x))
SymReal(::Irrational{sym}) where sym = SymReal(Constant(sym))
DOMAIN_TYPES[:Real] = SymReal

term(x::SymReal) = x.ex
Base.show(io::IO, t::SymReal) = print(io, t.ex)

for (M, f, arity) in DiffRules.diffrules()
    op = :($M.$f)

    if arity == 1
        @eval $op(x::SymReal) = track($op, x)
    elseif arity == 2
        @eval $op(x::SymReal, y::SymReal) = track($op, x, y)
    end
end

Base.promote_rule(::Type{<:SymReal}, ::Type{<:Real}) = SymReal
Base.convert(::Type{<:SymReal}, x::SymReal) = x
Base.convert(::Type{<:SymReal}, x::Real) = SymReal(x)
Base.convert(::Type{<:SymReal}, x::Irrational) = SymReal(x)


 #=
Author:= Shreeda Bhat
License:= MIT
UTF-8 =#

#=
In Julia the for loop is simple, Unlike python, matlab it won't use vectorized code.

    `for.....end`

=#

println("***********************1st Example*************************")
for num = 1:10 # here the colon is used represent the range like in python
    println("number is $num")
end

println("***********************2nd Example*************************")

x = ["shreeda",1,2,0.3,4]

for i in x
    println("what i'm passing now is $i")
end

#=
Note: in the condition it should be `variable = iterable`
=#

println("***********************3rd Example*************************")

for i in 1:10
    println(i)
end

println("***********************4th Example*************************")

my_bio = ["shreeda","20","Astaqc Consulting"]

for bio in my_bio
    println("let's see this array $bio")
end

println("*********************** 5th Example*************************")
m,n = 3,3
array = zeros(m,n)
println(array)

for i in 1:m
    for j in 1:n
        array[i,j]=i+j
    end
end
println(array)

println("*********************** 6th Example*************************")
# Let's do list comprehension
B = zeros(3,3)
for i in 1:m,j in 1:n
    B[i,j]=i+j
end
println(B)

println("*********************** 7th Example*************************")
# Sorry its array comprehension
c= [i+j for i in 1:n, j in 1:m]
print(c)


module Rank

using DataStructures
using Base: tail

export extrema_filter

# Max-min filter

# This is an implementation of the Lemire max-min filter
# http://arxiv.org/abs/cs.DS/0610046

# Monotonic wedge
immutable Wedge{T}
    L::CircularDeque{T}
    U::CircularDeque{T}
end
(::Type{Wedge{T}}){T}(n::Integer) = Wedge(CircularDeque{T}(n), CircularDeque{T}(n))

function Base.push!(W::Wedge, i::Integer)
    push!(W.L, i)
    push!(W.U, i)
    W
end

function addtoback!(W::Wedge, A, i, J)
    mn, mx = A[i, J]
    @inbounds while !isempty(W.L) && mn < A[back(W.L), J][1]
        pop!(W.L)
    end
    @inbounds while !isempty(W.U) && mx > A[back(W.U), J][2]
        pop!(W.U)
    end
    push!(W.L, i)
    push!(W.U, i)
    W
end

function Base.empty!(W::Wedge)
    empty!(W.L)
    empty!(W.U)
    W
end

@inline function getextrema(A, W::Wedge, J)
    (A[front(W.L), J][1], A[front(W.U), J][2])
end

"""
    extrema_filter(A, window) --> Array{(min,max)}

Calculate the running min/max over a window of width `window[d]` along
dimension `d`, centered on the current point. The returned array has
the same indices as the input `A`.
"""
function extrema_filter{T,N}(A::AbstractArray{T,N}, window::NTuple{N,Integer})
    _extrema_filter!([(a,a) for a in A], window...)
end
extrema_filter(A::AbstractArray, window::AbstractArray) = extrema_filter(A, (window...,))
extrema_filter(A::AbstractArray, window) = error("`window` must have the same number of entries as dimensions of `A`")

extrema_filter{T,N}(A::AbstractArray{T,N}, window::Integer) = extrema_filter(A, ntuple(d->window, Val{N}))

function _extrema_filter!(A::Array, w1, w...)
    if w1 > 1
        a = first(A)
        cache = ntuple(i->a, w1>>1)
        _extrema_filter1!(A, w1, cache)
    end
    _extrema_filter!(permutedims(A, [2:ndims(A);1]), w...)
end
_extrema_filter!(A::Array) = A

# Extrema-filtering along "columns" (dimension 1). This implements Lemire
# Algorithm 1, with the following modifications:
#   - multidimensional array support by looping over trailing dimensions
#   - working with min/max pairs rather than plain values, to
#     facilitate multidimensional processing
#   - output for all points of the array, handling the edges as max-min
#     over halfwindow on either side
function _extrema_filter1!{T}(A::AbstractArray{Tuple{T,T}}, window::Int, cache)
    # Initialise the internal wedges
    # U[1], L[1] are the location of the global (within the window) maximum and minimum
    # U[2], L[2] are the maximum and minimum over (U1, end] and (L1, end], respectively
    W = Wedge{Int}(window+1)
    tmp = Array{Tuple{T,T}}(window)
    c = z = first(cache)

    inds = indices(A)
    inds1 = inds[1]
    halfwindow = window>>1
    iw = min(last(inds1), first(inds1)+window-1)
    for J in CartesianRange(tail(inds))
        empty!(W)
        # Leading edge. We can't overwrite any values yet in A because
        # we'll need them again in later computations.
        for i = first(inds1):iw
            addtoback!(W, A, i, J)
            c, cache = cyclecache(cache, getextrema(A, W, J))
        end
        # Process the rest of the "column"
        for i = iw+1:last(inds1)
            A[i-window, J] = c
            if i == window+front(W.U)
                shift!(W.U)
            end
            if i == window+front(W.L)
                shift!(W.L)
            end
            addtoback!(W, A, i, J)
            c, cache = cyclecache(cache, getextrema(A, W, J))
        end
        for i = last(inds1)-window+1:last(inds1)-1
            if i >= first(inds1)
                A[i, J] = c
            end
            if i == front(W.U)
                shift!(W.U)
            end
            if i == front(W.L)
                shift!(W.L)
            end
            c, cache = cyclecache(cache, getextrema(A, W, J))
        end
        A[last(inds1), J] = c
    end
    A
end

# This is slightly faster than a circular buffer
@inline cyclecache(b, x) = b[1], (Base.tail(b)..., x)

end


include("./params.jl")

function pprocess_inhomopoisson(lambda)
    # Example usage:
    # L = 100000;
    # ts = linspace(0, 10, L);
    # f = 1.0;
    # lambda = 20*(1+sin(2.0 * pi * f * ts));
    # x, rs = pprocess_inhomopoisson(lambda)
    #
    # Input argument
    # lambda: A vector that specifies a time-dependent spike-rate.
    #
    # Output argument
    # x:    Spike time.
    # rs:   Intensity function.

    ################# WRITE SOMETHING #################

    return x, rs
end


### A Pluto.jl notebook ###
# v0.18.4

using Markdown
using InteractiveUtils

# ╔═╡ 3e14817e-85f8-4d5b-8a8f-075d1799ed27
using AbstractTrees, PlutoUI, Random, Wordlegames

# ╔═╡ ed649770-a53a-11ec-0434-1bbbc561f8cb
# hideall
title = "Wordle games as a tree";

# ╔═╡ 021364a8-f40b-49fe-9a02-406c0044c9ed
"""
+++
title = "$title"
+++
""" |> Base.Text

# ╔═╡ c49e1040-67c8-4386-a6ce-bbbeace8bc90
md"""
# $title

As described in the previous tutorial, strategies such as maximizing the entropy or minimizing the expected pool size for the next stage can be used to select guesses automatically in Wordle or related games.

When doing so the possible games can be represented in a data structure called a [tree](https://en.wikipedia.org/wiki/Tree_(data_structure)).

Some of the terminology used with these structures is based on concepts of a family tree.

First attach some packages that will be used
"""

# ╔═╡ cec480d3-7348-4d14-9628-007893713a07
md"and create an instance of `wordle` where the guesses are chosen to maximize the entropy, which is the default criterion."

# ╔═╡ 9586bcbd-16aa-4598-8faa-8220755566e6
begin
	datadir = joinpath(pkgdir(Wordlegames), "data")
	wordle = GamePool(collect(readlines(joinpath(datadir, "Wordletargets.txt"))))
end;

# ╔═╡ ddc457c3-f207-4ed1-8a8a-7f1c4cfdaed7
md"Finally, we create a tree from the games for a random selection of 25 targets."

# ╔═╡ 52ed4682-a0d7-44cf-b878-27f1cf9cc91d
gametree25 = tree(wordle, Random.seed!(1234321), 25);

# ╔═╡ 86d0bb82-ed8f-4b03-b90c-31c55e556bae
md"""
## The AbstractTrees package

The [AbstractTrees](https://github.com/JuliaCollections/AbstractTrees.jl) package provides many methods for working with tree data structures.
One of the most useful is `print_tree` which, as the name suggests, prints the tree in a special format.
(Because the content for these tutorials is generated as [Pluto](https://github.com/fonsp/Pluto.jl.git) notebooks, we need to wrap the call to `print_tree` in `with_terminal() do ... end` to have the output displayed.
Outside of Pluto this is not necessary.
"""

# ╔═╡ 30ba17b1-9c0c-4fd6-9f40-3c3c1c986e93
with_terminal() do
	print_tree(gametree25; maxdepth=8)
end

# ╔═╡ 384abc84-3fda-4ba3-b52e-2d46052b8ce6
md"""
Each guess in a game constitutes a `"node"` in the tree.
The initial guess in any of the games is `"raise"`, which is the `"root"` node for the tree.
A node can have zero or more `"children"` which are its immediate descendents.

The nodes in this tree are each a `GameNode` struct with a `"children"` field.
"""

# ╔═╡ 0c24b1ca-b36a-46f6-9da7-e5c3e3670d29
typeof(gametree25)

# ╔═╡ 68f9375f-31af-4916-933f-fb41905819f8
fieldnames(GameNode)

# ╔═╡ e1735528-469e-4a17-8b8a-d95dca03f46c
length(gametree25.children)  # number of children of the root node

# ╔═╡ 112b0452-7ae4-49b4-bbdd-54fb00ff8bdb
md"""
The `score` field of a `GameNode` is similar to the elements of the `guesses` field of a `GamePool` object but with one important difference.
In a `GameNode` the `score` and `sc` fields are the score that will produce the guess, as opposed to the score for the guess,

That is, the first child of `"raise"` is `"pilot"` which is the next guess in a game in which `"raise"` returns a score of `"🟫🟫🟨🟫🟫"` as tiles or 9 as a decimal number.
"""

# ╔═╡ f8f1521c-fbbb-4648-af9c-0a15c7e787ff
first(gametree25.children).score

# ╔═╡ 1dd5c9fb-820f-41c6-b796-494d68479f1c
md"""
The 16 children of the root node from these 25 games are
"""

# ╔═╡ 8302e26a-53ae-4f9f-9579-26a1c839c865
[child.score.guess for child in gametree25.children]

# ╔═╡ 7f4a7f14-fe77-43fa-b2dd-37ba70cf9da3
md"""
Some of these children have many descendents.
When generating the tree the children of a node are ordered according to the size of the tree rooted at that node.

Because a subtree is exactly the same type of structure as a tree, we can print a subtree with `print_tree`.
"""

# ╔═╡ 088ca1f6-44c3-4d0d-a68b-c71881813410
with_terminal() do
	print_tree(first(gametree25.children))
end

# ╔═╡ 785fe89b-4008-437a-ac73-2bdaa26e7cf6
md"""
We see that the size of the tree rooted at `"pilot"` is 10.
"""

# ╔═╡ 8d83a193-aed2-4c1d-8bc5-614a81d16548
md"""
The "leaves" of a tree are the terminal nodes, i.e. the nodes that do not have children.
"""

# ╔═╡ 576d0904-3f3e-49e8-9ce7-ccdd8767bf1b
[leaf.score.guess for leaf in Leaves(gametree25)]

# ╔═╡ 144a9d04-f971-48d9-ace4-dd0156f591c2
md"""
It happens in this case that all of the targets that generated the tree are leaves in this tree, but that is not necessarily the case.
"""

# ╔═╡ 105a0522-729a-49ea-b0a5-62b868ba5e45
length(collect(Leaves(gametree25)))

# ╔═╡ efcb3473-6ea1-459c-9995-32b0a407a2f4
md"""
## Creating the tree structure

As a language Julia gets high marks for "composability" - the ability to adapt one package to use concepts from another package.
The use of generic functions and multiple dispatch is central to this enhanced compositibility.

All that is necessary to use many of the functions in `AbstractTrees.jl` on the trees created from a collection of games using a particular `GamePool` is to define the `GameNode` struct, the method of generating the tree, and methods for `AbstractTrees.children`, `AbstractTrees.nodetype` and `AbstractTrees.printnode`
"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
Wordlegames = "1cb69566-e1cf-455f-a587-fd79a2e00f5a"

[compat]
AbstractTrees = "~0.3.4"
PlutoUI = "~0.7.38"
Wordlegames = "~0.3.0"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.8.0-beta3"
manifest_format = "2.0"
project_hash = "2c8800ff06b3096b67485eb059a224666198ba30"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "8eaf9f1b4921132a4cff3f36a1d9ba923b14a481"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.4"

[[deps.AbstractTrees]]
git-tree-sha1 = "03e0550477d86222521d254b741d470ba17ea0b5"
uuid = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
version = "0.3.4"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
version = "1.1.1"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "96b0bc6c52df76506efc8a441c6cf1adcb1babc4"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.42.0"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"
version = "0.5.2+0"

[[deps.Crayons]]
git-tree-sha1 = "249fe38abf76d48563e2f4556bebd215aa317e15"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.1.1"

[[deps.DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[deps.DataFrames]]
deps = ["Compat", "DataAPI", "Future", "InvertedIndices", "IteratorInterfaceExtensions", "LinearAlgebra", "Markdown", "Missings", "PooledArrays", "PrettyTables", "Printf", "REPL", "Reexport", "SortingAlgorithms", "Statistics", "TableTraits", "Tables", "Unicode"]
git-tree-sha1 = "ae02104e835f219b8930c7664b8012c93475c340"
uuid = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
version = "1.3.2"

[[deps.DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "3daef5523dd2e769dad2365274f760ff5f282c7d"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.11"

[[deps.DataValueInterfaces]]
git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
uuid = "e2d170a0-9d28-54be-80f0-106bbe20a464"
version = "1.0.0"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[deps.Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[deps.Downloads]]
deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
version = "1.6.0"

[[deps.FileWatching]]
uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Formatting]]
deps = ["Printf"]
git-tree-sha1 = "8339d61043228fdd3eb658d86c926cb282ae72a8"
uuid = "59287772-0a20-5a39-b81b-1366585eb4c0"
version = "0.4.2"

[[deps.Future]]
deps = ["Random"]
uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.InvertedIndices]]
git-tree-sha1 = "bee5f1ef5bf65df56bdd2e40447590b272a5471f"
uuid = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
version = "1.1.0"

[[deps.IteratorInterfaceExtensions]]
git-tree-sha1 = "a3f24677c21f5bbe9d2a714f95dcd58337fb2856"
uuid = "82899510-4779-5014-852e-03e436cf321d"
version = "1.0.0"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "3c837543ddb02250ef42f4738347454f95079d4e"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.3"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"
version = "0.6.3"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"
version = "7.81.0+0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"
version = "1.10.2+0"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"
version = "2.28.0+0"

[[deps.Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "bf210ce90b6c9eed32d25dbcae1ebc565df2687f"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.0.2"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
version = "2022.2.1"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
version = "1.2.0"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
version = "0.3.20+0"

[[deps.OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "621f4f3b4977325b9128d5fae7a8b4829a0c2222"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.2.4"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
version = "1.8.0"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "ColorTypes", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "670e559e5c8e191ded66fa9ea89c97f10376bb4c"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.38"

[[deps.PooledArrays]]
deps = ["DataAPI", "Future"]
git-tree-sha1 = "28ef6c7ce353f0b35d0df0d5930e0d072c1f5b9b"
uuid = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
version = "1.4.1"

[[deps.PrettyTables]]
deps = ["Crayons", "Formatting", "Markdown", "Reexport", "Tables"]
git-tree-sha1 = "dfb54c4e414caa595a1f2ed759b160f5a3ddcba5"
uuid = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
version = "1.3.1"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"
version = "0.7.0"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SortingAlgorithms]]
deps = ["DataStructures"]
git-tree-sha1 = "b3363d7460f7d098ca0912c69b082f75625d7508"
uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
version = "1.0.1"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
version = "1.0.0"

[[deps.TableTraits]]
deps = ["IteratorInterfaceExtensions"]
git-tree-sha1 = "c06b2f539df1c6efa794486abfb6ed2022561a39"
uuid = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
version = "1.0.1"

[[deps.Tables]]
deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "OrderedCollections", "TableTraits", "Test"]
git-tree-sha1 = "5ce79ce186cc678bbb5c5681ca3379d1ddae11a1"
uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
version = "1.7.0"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
version = "1.10.0"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Wordlegames]]
deps = ["AbstractTrees", "DataFrames", "Random", "Tables"]
git-tree-sha1 = "4c463de78d2f3f9447b695e241eba43cc945f866"
uuid = "1cb69566-e1cf-455f-a587-fd79a2e00f5a"
version = "0.3.0"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"
version = "1.2.12+1"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"
version = "5.1.0+0"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"
version = "1.41.0+1"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
version = "16.2.1+1"
"""

# ╔═╡ Cell order:
# ╟─ed649770-a53a-11ec-0434-1bbbc561f8cb
# ╟─021364a8-f40b-49fe-9a02-406c0044c9ed
# ╟─c49e1040-67c8-4386-a6ce-bbbeace8bc90
# ╠═3e14817e-85f8-4d5b-8a8f-075d1799ed27
# ╟─cec480d3-7348-4d14-9628-007893713a07
# ╠═9586bcbd-16aa-4598-8faa-8220755566e6
# ╟─ddc457c3-f207-4ed1-8a8a-7f1c4cfdaed7
# ╠═52ed4682-a0d7-44cf-b878-27f1cf9cc91d
# ╟─86d0bb82-ed8f-4b03-b90c-31c55e556bae
# ╠═30ba17b1-9c0c-4fd6-9f40-3c3c1c986e93
# ╟─384abc84-3fda-4ba3-b52e-2d46052b8ce6
# ╠═0c24b1ca-b36a-46f6-9da7-e5c3e3670d29
# ╠═68f9375f-31af-4916-933f-fb41905819f8
# ╠═e1735528-469e-4a17-8b8a-d95dca03f46c
# ╟─112b0452-7ae4-49b4-bbdd-54fb00ff8bdb
# ╠═f8f1521c-fbbb-4648-af9c-0a15c7e787ff
# ╟─1dd5c9fb-820f-41c6-b796-494d68479f1c
# ╠═8302e26a-53ae-4f9f-9579-26a1c839c865
# ╟─7f4a7f14-fe77-43fa-b2dd-37ba70cf9da3
# ╠═088ca1f6-44c3-4d0d-a68b-c71881813410
# ╟─785fe89b-4008-437a-ac73-2bdaa26e7cf6
# ╟─8d83a193-aed2-4c1d-8bc5-614a81d16548
# ╠═576d0904-3f3e-49e8-9ce7-ccdd8767bf1b
# ╟─144a9d04-f971-48d9-ace4-dd0156f591c2
# ╠═105a0522-729a-49ea-b0a5-62b868ba5e45
# ╟─efcb3473-6ea1-459c-9995-32b0a407a2f4
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


#=
MIT License

Copyright (c) 2019 Jagadeesh Kotra

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
=#

using ArgParse

function parse_commandline_args()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--switch"
           help = "to SWITCH or NOT"
           action = :store_true

        "--simulations"
           help = "N no of simulations"
           required = true
           arg_type = Int64

    end
    return parse_args(s)
end

parsed_args = parse_commandline_args()

n_doors = 3
n_siml = parsed_args["simulations"]
switch = parsed_args["switch"]

w = 0
l = 0

for i = 1:n_siml

    doors = rand(0:0,n_doors)

    #Keep the money in random door
    money_door = rand(1:length(doors))
    doors[money_door] = 1

    chosen_door = rand(1:length(doors))

    #DO YOU WANT TO SWITCH?!?!?!
    if switch
        doors_cpy = copy(doors)
        doors_cpy[chosen_door] = -1

        #open and mark money-less door!
        host_choice = findall(x->x!=1 && x!=-1, doors_cpy)[1]
        doors_cpy[host_choice] = -1

        #set chosen door index
        chosen_door = findall(x->x!=-1, doors_cpy)[1]
    end

    if doors[chosen_door] == 1
        global w += 1
    else
        global l += 1
    end

    println("Chosen:",chosen_door," Money In:", money_door, " outcome:", doors[chosen_door] == 1)

end

println("---END---")
println(w,"-",l," Win %:", 100*(w/(w+l)))


using DiffEqJump, DiffEqBase, OrdinaryDiffEq, StochasticDiffEq, Statistics
using Test
using StableRNGs
rng = StableRNG(12345)

rate = (u,p,t) -> 1.0*u[1]
affect! = function (integrator)
  integrator.u[1] = 1.0
end
jump1 = ConstantRateJump(rate,affect!)

prob = DiscreteProblem([10],(0.0,50.0))
prob_control = DiscreteProblem([10],(0.0,50.0))

jump_prob = JumpProblem(prob,Direct(),jump1; rng=rng)
jump_prob_control = JumpProblem(prob_control,Direct(),jump1; rng=rng)

coupling_map = [(1, 1)]
coupled_prob = SplitCoupledJumpProblem(jump_prob,jump_prob_control,Direct(),coupling_map; rng=rng)

@time sol =  solve(coupled_prob,FunctionMap())
@time solve(jump_prob,FunctionMap())
@test [s[1]-s[2] for s in sol.u] == zeros(length(sol.t)) # coupling two copies of the same process should give zero

rate = (u,p,t) -> 1.0
affect! = function (integrator)
  integrator.u[1] = 1.0
end
jump1 = ConstantRateJump(rate,affect!)
rate = (u,p,t) -> 2.0
jump2 = ConstantRateJump(rate,affect!)

f = function (du,u,p,t)
  du[1] = u[1]
end
g = function (du,u,p,t)
  du[1] = 0.1
end

# Jump ODE to jump ODE
prob = ODEProblem(f,[1.],(0.0,1.0))
prob_control = ODEProblem(f,[1.],(0.0,1.0))
jump_prob = JumpProblem(prob,Direct(),jump1; rng=rng)
jump_prob_control = JumpProblem(prob_control,Direct(),jump2; rng=rng)
coupled_prob = SplitCoupledJumpProblem(jump_prob,jump_prob_control,Direct(),coupling_map; rng=rng)
sol =  solve(coupled_prob,Tsit5())
@test mean([abs(s[1]-s[2]) for s in sol.u])<=5.

# Jump SDE to Jump SDE
prob = SDEProblem(f,g,[1.],(0.0,1.0))
prob_control = SDEProblem(f,g,[1.],(0.0,1.0))
jump_prob = JumpProblem(prob,Direct(),jump1; rng=rng)
jump_prob_control = JumpProblem(prob_control,Direct(),jump1; rng=rng)
coupled_prob = SplitCoupledJumpProblem(jump_prob,jump_prob_control,Direct(),coupling_map; rng=rng)
sol =  solve(coupled_prob,SRIW1())
@test mean([abs(s[1]-s[2]) for s in sol.u])<=5.

# Jump SDE to Jump ODE
prob = ODEProblem(f,[1.],(0.0,1.0))
prob_control = SDEProblem(f,g,[1.],(0.0,1.0))
jump_prob = JumpProblem(prob,Direct(),jump1; rng=rng)
jump_prob_control = JumpProblem(prob_control,Direct(),jump1; rng=rng)
coupled_prob = SplitCoupledJumpProblem(jump_prob,jump_prob_control,Direct(),coupling_map; rng=rng)
sol =  solve(coupled_prob,SRIW1())
@test mean([abs(s[1]-s[2]) for s in sol.u])<=5.

# Jump SDE to Discrete
rate = (u,p,t) -> 1.
affect! = function (integrator)
  integrator.u[1] += 1.
end
prob = DiscreteProblem([1.],(0.0,1.0))
prob_control = SDEProblem(f,g,[1.],(0.0,1.0))
jump_prob = JumpProblem(prob,Direct(),jump1; rng=rng)
jump_prob_control = JumpProblem(prob_control,Direct(),jump1; rng=rng)
coupled_prob = SplitCoupledJumpProblem(jump_prob,jump_prob_control,Direct(),coupling_map; rng=rng)
sol =  solve(coupled_prob,SRIW1())

# test mass action jumps coupled to ODE
# 0 -> A (stochasic) and A -> 0 (ODE)
rate        = [100.]
react_stoch = [Vector{Pair{Int,Int}}()]
net_stoch   = [[1 => 1]]
majumps     = MassActionJump(rate, react_stoch, net_stoch)
f = function (du,u,p,t)
  du[1] = -1.0*u[1]
end
odeprob     = ODEProblem(f,[10.0],(0.0,10.0))
jump_prob   = JumpProblem(odeprob, Direct(), majumps, save_positions=(false,false); rng=rng)
Nsims = 8000
Amean = 0.
for i in 1:Nsims
  global Amean
  local sol = solve(jump_prob,Tsit5(),saveat=10.)
  Amean += sol[1,end]
end
Amean /= Nsims
actmean = 100. + (10.0-100.0)*exp(-1.0*10.0)
#println(abs(Amean-actmean)/actmean)
@test abs(actmean - Amean) < .02 * actmean


using SudokuSolver
using Test

@testset "SudokuSolver.jl" begin
    s6ans = [3 7 9 8 1 6 2 5 4; 8 6 4 2 7 5 1 9 3; 5 2 1 9 4 3 6 8 7; 6 5 7 4 3 1 9 2 8;
    4 9 8 5 6 2 7 3 1; 1 3 2 7 8 9 5 4 6; 2 1 3 6 9 4 8 7 5; 7 4 5 1 2 8 3 6 9; 9 8 6 3 5 7 4 1 2]
    @test solvesudoku(SudokuSolver.PuzzleExamples[6]) == s6ans
    @test solvesudoku(SudokuSolver.PuzzleExamples[6], solver=SudokuSolver.naivesolver!) == s6ans
    @test solvesudoku(SudokuSolver.PuzzleExamples[6], solver=SudokuSolver.reordersolver!) == s6ans
    @test solvesudoku(SudokuSolver.PuzzleExamples[6], solver=SudokuSolver.prioritysolver!) == s6ans
    invalidpuzzle = "000011054800000000000000000650400000000002730000000000210000800700000300000350000"
    @test solvesudoku(invalidpuzzle) === nothing
    @test solvesudoku(invalidpuzzle, solver=SudokuSolver.naivesolver!, check=false) == SudokuSolver.loadpuzzle(invalidpuzzle)
    @test solvesudoku(invalidpuzzle, solver=SudokuSolver.reordersolver!, check=false) == SudokuSolver.loadpuzzle(invalidpuzzle)
    @test solvesudoku(invalidpuzzle, solver=SudokuSolver.prioritysolver!, check=false) |> size == (9, 9)
end


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.1
#     language: julia
#     name: julia-1.6
# ---

# %% [markdown] tags=[]
# Optimization by specialization to argument types comes with the tradeoff of increased compilation time. An extreme case of this is shown in the following example.

# %%
f(x) = x^2 + 3x + 2
g(::Val{x}) where x = f(x)
f(10^6), g(Val(10^6))

# %%
using BenchmarkTools
@btime f(10^6)
@btime g($(Val(10^6)));

# %% [markdown]
# `g(Val(10^6))` is ultra fast because it is specialized to the argument type `Val{10^6}` and compiled to `return 1000003000002`.

# %%
@code_typed debuginfo=:none g(Val(10^6))

# %% [markdown]
# However, `g(Val(k))` is compiled separately for each different `k`.  So if you run `g(Val(k))` for a large number of different `k`'s,  it will perform a large number of compilations and will be very slow. (After compilation, though, it will be explosively fast.)

# %%
F(n) = [f(k) for k in 1:n]
G(n) = [g(Val(k)) for k in 1:n]

@time F(10^4)
@time G(10^4);

# %% [markdown]
# The first execution of `G(10^4)` is very slow.

# %%
F(10^4) == G(10^4)

# %% [markdown]
# Thus, it is not reasonable to try to optimize by specialization to the argument types by different large numbers of `Val{k}` types.
#
# On the other hand, the native code specialized to argument types is very fast, so if compilation time is not an issue, optimization by specialization to argument types should be done aggressively.
#
# In short, it is a matter of trade-off.

# %%


# should be constriants not projection, projection is a function on the contrsrinats
abstract type Constraints end
struct BoxConstraint{T} <: Constraints
    lower::T
    upper::T
end
abstract type Scaling end
struct NoScaling <: Scaling end
scale(S::NoScaling, gx) = gx
struct ConstantPDScaling{T} <: Scaling
    S::T
end
scale(S::ConstantPDScaling, gx) = S.S \ gx

function update_xλ(prob::OptimizationProblem, S::Scaling, λ, x, gx)
    bounds = bounds(prob)
    min.(max.(bounds.lower, x .- λ .* scale(S, gx)), bounds.upper)
end
#update_xλ(p::BoxConstraint, λ, x, gx) = min.(max.(p.lower, x .- λ.*gx), p.upper)

# 5.13 in Kelley IMO
# what's it in Bertsekas?
function sufficient_decrease(fλ, fx, λ, α, x, xλ, gx)
    lhs = fλ
    # We use (22) from Bersekas (1976) to avoid dividing by λ
    rhs = fx - dot(gx, x .- xλ) * α
    lhs <= rhs
end
function sufficient_decrease(fλ, fx, λ, α, gx)
    lhs = fλ
    rhs = fx - norm(gx, 2) * α * λ
    lhs <= rhs
end

function projected_gradient(
    prob::OptimizationProblem,
    f,
    g,
    x0,
    sc::Union{Scaling,NoScaling} = NoScaling();
    itermax,
    ls_itermax,
    β = 0.5,
    α = 1e-6,
    print_trace = true,
)

    f0, g0 = f(x0), g(x0)
    if norm(x0 .- update_xλ(prob, sc, 1, x0, g0), Inf) <= 1e-6
        return x0, f0, g0
    end

    x = x0

    # 5.12 in Kelley IMO
    # what's it in Bertsekas?
    local fx, gx
    for i = 1:itermax
        fx, gx = f(x), g(x)
        if print_trace
            println("i = $i    fx: $(fx)    |gx|: $(norm(gx,Inf))")
            println("x1 = $(x[1])")
            println("x2 = $(x[2])")
        end
        if norm(x .- update_xλ(prob, sc, 1, x, gx)) <= 1e-6
            return x, fx, gx
        end
        for j = 1:ls_itermax
            λ = β^j
            xλ = update_xλ(P, sc, λ, x, gx)
            fλ = f(xλ)
            if sufficient_decrease(fλ, fx, λ, α, x, xλ, gx)
                x = xλ
                break
            end
        end
    end
    x, fx, gx
end


cd(@__DIR__)
using Pkg
# you may need to use Pkg.instantiate() the first time you run the file
Pkg.activate("")
using Plots
include("functions.jl")

# current chain
x = [.8,.5]
# random chain
z = [1.5,2.5]
# difference of current and random z
# forms dotted line
d = x - z

# random chains r1 and r2
r1 = [0,-1.5]
r2 = [-.3,-2.1]

# project r1 onto line spanned by d = x - z
proj1 = project(r1, d)
# project r2 onto line spanned by d = x - z
proj2 = project(r2, d)

# difference scalar set to 1 for simplicity
γ = 1.0

# difference in projections
dproj1_2 = proj1 - proj2

# proposal x + γ * (proj1 - proj2)
proposal = x + γ * dproj1_2

scatter(
    [x[1]],
    [x[2]],
    xlims = (-3.0, 3.0),
    ylims = (-3.0, 3.0),
    color = :black,
    linewidth = 2,
    label = "",
    framestyle = :origin,
)
annotate!(x[1]-.1,x[2]+.2, text("X", :black, 12))

scatter!(
    [z[1]],
    [z[2]],
    xlims = (-3.0, 3.0),
    ylims = (-3.0, 3.0),
    color = :black,
    linewidth = 2,
    label = "",
    framestyle = :origin,
)
annotate!(z[1]-.1,z[2] +.1, text("Z", :black, 12))

xvals = -30:.1:30
β = get_coefs(x, z)
yvals = make_line(β..., xvals)
plot!(
    xvals,
    yvals,
    color = :black,
    linestyle = :dash,
    linewidth = 2,
    label = "",
    framestyle = :origin,
)

scatter!(
    [r1[1]],
    [r1[2]],
    color = :black,
    linewidth = 2,
    label = "",
    framestyle = :origin,
)
annotate!(r1[1]-.1,r1[2] +.2, text("R1", :black, 12))

scatter!(
    [r2[1]],
    [r2[2]],
    color = :black,
    linewidth = 2,
    label = "",
    framestyle = :origin,
)
annotate!(r2[1]-.1,r2[2] +.2, text("R2", :black, 12))

scatter!(
    [proposal[1]],
    [proposal[2]],
    color = :red,
    linewidth = 2,
    label = "",
    framestyle = :origin,
)
annotate!(proposal[1] - .1,proposal[2] +.1, text("X*", :black, 12))


@testset "kinetic energy zero at zero velocity" begin
    θ_t = deg2rad(30.0)
    θ′_t = 0.0
    θ_k = deg2rad(22.0)
    θ′_k = 0.0
    γ_OA = deg2rad(60.0)
    γ′_OA = 0.0
    T = @test_nowarn kinetic_energy(0.0, 0.0, 0.0, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA)
    @test T == 0.0
end

@testset "kinetic energy non-negative on various inputs" begin
    θ_t = deg2rad(30.0)
    θ′_t = 0.0
    θ_k = deg2rad(22.0)
    θ′_k = 0.1
    γ_OA = deg2rad(60.0)
    γ′_OA = -0.2
    @test kinetic_energy(-0.1, 0.0, 0.0, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    @test kinetic_energy(0.1, 0.0, 0.0, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    @test kinetic_energy(0.0, -0.2, 0.0, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    @test kinetic_energy(0.1, 0.12, 0.0, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    @test kinetic_energy(0.0, 0.0, -0.22, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    @test kinetic_energy(0.0, 0.4, 0.22, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    @test kinetic_energy(-1000.0, 100.4, 10.22, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    @test kinetic_energy(-1000.0, -100.4, -1000.22, θ_t, θ′_t, θ_k, θ′_k, γ_OA, γ′_OA) ≥ 0
    x = randn(9, 12) .* [ones(3) .* 1e3; ones(6) .* 0.1]
    for p in eachcol(x)
        @test kinetic_energy(p...) ≥ 0
    end
end


using IntegralArrays, LinearAlgebra, IntervalSets
using ColorTypes, ColorVectorSpace
using ColorTypes.FixedPointNumbers
using OffsetArrays
using Test
using Documenter

@testset "meta-quality" begin
    if VERSION >= v"1.6"
        doctest(IntegralArrays; manual = false)
    end
end

@testset "IntegralArrays" begin
    a = zeros(10, 10)
    int_array = IntegralArray(a)
    @test all(int_array == a)

    a = ones(10,10)
    int_array = IntegralArray(a)
    chk = Array(1:10)
    @test all([vec(int_array[i, :]) == chk * i for i in 1:10])

    int_sum = int_array[1..5, 1..2]
    @test int_sum == 10.0
    int_sum = int_array[3 ± 2, 1..2]
    @test int_sum == 10.0
    int_sum = int_array[3 ± 2, 3 ± 2]
    @test int_sum == 25.0
    int_sum = int_array[1..2, 1..5]
    @test int_sum == 10.0
    int_sum = int_array[4..8, 4..8]
    @test int_sum == 25.0
    int_sum = int_array[6 ± 2, 6 ± 2]
    @test int_sum == 25.0
    int_sum = int_array[6 ± 2, 6 ± 2]
    @test int_sum == 25.0

    a = Array(reshape(1:100, 10, 10))
    int_array = IntegralArray(a)
    @test diag(int_array) == Array([1, 26,  108,  280,  575, 1026, 1666, 2528, 3645, 5050])

    int_sum = int_array[1..3, 1..3]
    @test int_sum == 108
    int_sum = int_array[2 ± 1, 2 ± 1]
    @test int_sum == 108
    int_sum = int_array[2 ± 1, 2 ± 1]
    @test int_sum == 108
    int_sum = int_array[1..5, 1..2]
    @test int_sum == 80
    int_sum = int_array[3 ± 2, 1..2]
    @test int_sum == 80
    int_sum = int_array[4..8, 4..8]
    @test int_sum == 1400
    int_sum = int_array[6 ± 2, 6 ± 2]
    @test int_sum == 1400
    int_sum = int_array[6 ± 2, 6 ± 2]
    @test int_sum == 1400

    I, Δ = CartesianIndex(6, 6), CartesianIndex(2, 2)
    @test int_array[I-Δ..I+Δ] == 1400

    @testset "OffsetArray" begin
        X = rand(1:5, 5, 5)
        iX = IntegralArray(X)
        @test axes(iX) == (1:5, 1:5)
        Xo = OffsetArray(X, -1, -1)
        iXo = IntegralArray(Xo)
        @test axes(iXo) == (0:4, 0:4)

        @test_throws DimensionMismatch IntegralArray(X, Xo)
    end

    @testset "in-place with buffer" begin
        X = collect(reshape(1:25, 5, 5))
        iX = similar(X)
        IntegralArray(iX, X)
        IntegralArray(X, X)
        @test iX == X == IntegralArray(reshape(1:25, 5, 5))
    end
end

@testset "Color integral arrays" begin
    A = [RGB{N0f8}(0.2, 0.8, 0), RGB{N0f8}(0.5, 0.3, 0)]
    intA = IntegralArray(A)
    a1, a2 = A
    @test intA[2] == RGB(Float32(red(a1))+Float32(red(a2)),
                         Float32(green(a1))+Float32(green(a2)),
                         Float32(blue(a1))+Float32(blue(a2)))
    @test intA[1..2] == intA[2]
    @test intA[2..2] ≈ A[2]
end


## This makes implementing operators simpler
# but overrided BandedMatrix directly is likely to be faster

abstract TridiagonalOperator{T} <: Operator{T}
abstract SymTridiagonalOperator{T} <: TridiagonalOperator{T}
abstract BidiagonalOperator{T} <: TridiagonalOperator{T}
abstract DiagonalOperator{T} <: BidiagonalOperator{T}
# override getindex

bandinds(::TridiagonalOperator)=-1,1
bandinds(::BidiagonalOperator)=0,1
bandinds(::DiagonalOperator)=0,0

Base.Tridiagonal{T}(J::TridiagonalOperator{T},n::Integer)=Tridiagonal(T[J.γ(k) for k=2:n],
T[J.α(k) for k=1:n],
T[J.β(k) for k=1:n-1])

function symmetrize{T}(J::TridiagonalOperator{T},n::Integer)
    d=Array(T,n)
    d[1]=1
    for k=2:n
        d[k]=sqrt(J[k,k-1]/J[k-1,k])*d[k-1]
    end

   SymTridiagonal(
    T[J[k,k] for k=1:n],
    T[J[k,k+1]*d[k+1]/d[k] for k=1:n-1])
end

immutable DiagIteratorOperator{IT,T} <: DiagonalOperator{T}
    iterator::IT
end

DiagIteratorOperator(it) = DiagIteratorOperator{typeof(it),eltype(it)}(it)

getindex(D::DiagIteratorOperator,k::Integer,j::Integer) =
    k==j? D.iterator[k] : zero(eltype(D))

domainspace(D::DiagIteratorOperator) = ℓ⁰
rangespace(D::DiagIteratorOperator) = ℓ⁰

Base.diagm(c::AbstractCount) = DiagIteratorOperator(c)


export quick_sort!, hoare_qsort!

function quick_sort!(arr::Vector, h::Int64, t::Int64)
    len = t-h+1
    if len < 2
        return arr
    end

    n, i = h, h
    pivot = arr[t]
    for n = h:t
        if arr[n] <= pivot
            arr[i], arr[n] = arr[n], arr[i]
            i += 1
        end
    end

    if i > t
        i = t
    end

    quick_sort!(arr, h, i-1)
    quick_sort!(arr, i, t)

    return arr
end

function hoare_qsort!(arr::Vector, p::Int64, r::Int64)
    if p >= r
        return arr
    end
    x = arr[p]
    i = p - 1
    j = r + 1
    while true
        while true
            j -= 1
            arr[j] > x || break
        end
        while true
            i += 1
            arr[i] < x || break
        end

        if i < j
            arr[i], arr[j] = arr[j], arr[i]
        else
            break
        end
    end

    hoare_qsort!(arr, p, j);
    hoare_qsort!(arr, j+1, r);

    return arr
end