#=
# Stress on Truss
In this example we'll plot an animation of the stress on some truss structure
using `GaphMakie.jl` and [`NetworkDynamics.jl`](https://github.com/PIK-ICoN/NetworkDynamics.jl)

![truss animation](truss.mp4)
=#
using NetworkDynamics
using OrdinaryDiffEq
using Graphs
using GraphMakie
using LinearAlgebra: norm, ⋅
using Printf
using CairoMakie
CairoMakie.activate!()

#=
## Definition of the dynamical system
Define dynamic model for NetworkDynamics
=#
function edge_f!(e, v_s, v_d, (K, L), _)
    Δd = view(v_d, 1:2) - view(v_s, 1:2)
    nd = norm(Δd)
    F = K * ( L - nd )
    e[1:2] = F .* view(Δd, 1:2) ./ nd
end

function vertex_free!(du, u, edges, (M, γ, g), _)
    F = sum(edges) - γ * view(u, 3:4) + [0, -M*g]
    du[3:4] .= F ./ M
    du[1:2] = view(u, 3:4)
end

function vertex_fixed!(du, u, edges, _, _)
    du[1:2] .= 0.0
end

edge = StaticEdge(f = edge_f!, dim=2, coupling=:antisymmetric)
vertex_free = ODEVertex(f = vertex_free!, dim=4, sym=[:x, :y, :v, :w])
vertex_fix  = ODEVertex(f = vertex_fixed!, dim=2, sym=[:x, :y, :v, :w])
nothing #hide

#=
Set up graph topology and initial positions.
=#
N = 5
dx = 1.0
shift = 0.2
g = SimpleGraph(2*N + 1)
for i in 1:N
    add_edge!(g, i, i+N); add_edge!(g, i, i+N)
    if i < N
        add_edge!(g, i+1, i+N); add_edge!(g, i, i+1); add_edge!(g, i+N, i+N+1)
    end
end
add_edge!(g, 2N, 2N+1)
pos0 = zeros(Point2f, 2N + 1)
pos0[1:N] = [Point((i-1)dx,0) for i in 1:N]
pos0[N+1:2*N] = [Point(i*dx + shift, 1) for i in 1:N]
pos0[2N+1] = Point(N*dx + 1, -1)

fixed = [1,4] # set fixed vertices
nothing #hide

# create networkdynamics object
vertices = ODEVertex[vertex_free for i in 1:nv(g)]
for i in fixed
    vertices[i] = vertex_fix # use the fixed vertex for the fixed
end
nd = network_dynamics(vertices, edge, g);
nothing #hide

#=
write some auxiliary functions to map between state vector of DGL `u` and graph data
=#
x_idxs = [findfirst(isequal(Symbol("x_$i")), nd.syms) for i in 1:nv(g)]

"Set positions `pos` inside dgl state `u`"
function set_pos!(u, pos)
    for (i, idx) in enumerate(x_idxs)
        u[idx] = pos[i][1]
        u[idx+1] = pos[i][2]
    end
end

"Extract vector of Points from dgl state `u`"
function to_pos(u)
    pos = Vector{Point2f}(undef, nv(g))
    for (i, idx) in enumerate(x_idxs)
        pos[i] = Point(u[idx], u[idx+1])
    end
    return pos
end

"Calculate load on edge for given state."
function get_load(u, p, t=0.0)
    gd_nd = nd(u, p, t, GetGD) # exposes underlying graph data struct
    force = Vector{Float64}(undef, ne(g))
    pos = to_pos(u)
    for (i,e) in enumerate(edges(g))
        edgeval = get_edge(gd_nd, i)
        fvec = Point(edgeval[1], edgeval[2])
        dir = pos[e.dst] .- pos[e.src]
        force[i] = sign(fvec ⋅ dir) * norm(fvec)
    end
    return force
end
nothing #hide

#=
Set parameters.
=#
M = [10 for i in 1:nv(g)] # mass of the nodes
M[end] = 200 # heavy mass
gc = [9.81 for i in 1:nv(g)] # gravitational constant
γ = [200.0 for i in 1:nv(g)] # damping parameter
γ[end] = 100.0

L = [norm(pos0[e.src] - pos0[e.dst]) for e in edges(g)] # length of edges
K = [0.5e6 for i in 1:ne(g)] # spring constant of edges

## bundle parameters for NetworkDynamics
para = collect.((zip(M, γ, gc),
                 zip(K, L)))
nothing #hide

#=
Set initial state and solve the system
=#
u0 = zeros(length(nd.syms))
set_pos!(u0, pos0)

tspan = (0.0, 12.0)
prob = ODEProblem(nd, u0, tspan, para)
sol  = solve(prob, Tsit5());
nothing #hide

#=
## Plot the solution
=#

fig = Figure(resolution=(1000,550))
fig[1,1] = title = Label(fig, "Stress on truss", textsize=30)
title.tellwidth = false

fig[2,1] = ax = Axis(fig)

## calculate some values for colorscaling
(fmin, fmax) = 0.3 .* extrema([(map(u->get_load(u, para), sol)...)...])
nlabels = [" " for i in 1:nv(g)]
nlabels[fixed] .= "Δ"
elabels = ["edge $i" for i in 1:ne(g)]
p = graphplot!(ax, g;
               edge_width=4.0,
               node_size=sqrt.(M)*3,
               nlabels=nlabels,
               nlabels_align=(:center,:top),
               nlabels_textsize=30,
               elabels=elabels,
               elabels_opposite=[ne(g)],
               edge_color=[0.0 for i in 1:ne(g)],
               edge_attr=(colorrange=(fmin,fmax),
                          colormap=:diverging_bkr_55_10_c35_n256))
hidespines!(ax); hidedecorations!(ax); p[:node_pos][]=to_pos(u0); ax.aspect = DataAspect()
limits!(ax, -0.1, pos0[end][1]+0.3, pos0[end][2]-0.5, 1.15)

## draw colorbar
fig[3,1] = cb = Colorbar(fig, p.plots[1], label = "Axial force", vertical=false)

T = tspan[2]
fps = 30
trange = range(0.0, sol.t[end], length=Int(T * fps))
record(fig, "truss.mp4", trange; framerate=fps) do t
    title.text = @sprintf "Stress on truss (t = %.2f )" t
    p[:node_pos][] = to_pos(sol(t))
    load = get_load(sol(t), para)
    p.elabels = [@sprintf("%.0f", l) for l in load]
    p.edge_color[] = load
end
nothing #hide


using Test
using YaoSym, YaoArrayRegister, YaoBlocks, SymEngine

@test ket"111" + 2ket"111" == 3ket"111"
@test bra"111" * ket"111" == 1
@test (bra"111" + bra"101") * ket"111" == 1

@test ket"111"^2 == ket"111111"
@test ket"101" * ket"111" == ket"101111"
@test bra"110"^2 == bra"110110"
@test bra"110" * bra"111" == bra"110111" == bra"110_111"

# printing does not error
@test print(ket"001") == nothing
@test print(bra"001") == nothing
@test print(bra"111" + bra"101")  == nothing


export Day5, Day5_1, Day5_2
abstract type Day5 <: Day end
abstract type Day5_1 <: Day5 end
abstract type Day5_2 <: Day5 end

struct SeatLoc
    row::Int
    col::Int
    id::Int
    SeatLoc(row,col) = new(row, col, row*8+col)
end

line_parser(::Type{<:Day5}, line) = toseat(line)

solve(::Type{Day5_1}, seats::Vector) = maximum((seat.id for seat in seats))
solve(::Type{Day5_2}, seats::Vector) = begin
    seat_ids = Set((seat.id for seat in seats))
    for i in 1:length(seat_ids)-2
        if (i ∉ seat_ids) & (i-1 ∈ seat_ids) & (i+1 ∈ seat_ids)
            return i
        end
    end
end

row(v) = v |> x-> replace(x,Pair("F","0")) |> x-> replace(x,Pair("B","1")) |> x-> parse(Int,x; base=2)
col(v) = v |> x-> replace(x,Pair("L","0")) |> x-> replace(x,Pair("R","1")) |> x-> parse(Int,x; base=2)

toseat(seat) = SeatLoc(row(seat[1:7]), col(seat[8:10]))


using ModelingToolkit
using Test

struct Lorenz <: AbstractComponent
    x::Variable
    y::Variable
    z::Variable
    σ::Variable
    ρ::Variable
    β::Variable
    eqs::Vector{Equation}
end
function generate_lorenz_eqs(t,x,y,z,σ,ρ,β)
    D = Differential(t)
    [D(x) ~ σ*(y-x)
     D(y) ~ x*(ρ-z)-y
     D(z) ~ x*y - β*z]
end
function Lorenz(t)
    @Unknown x(t) y(t) z(t)
    @Param σ ρ β
    Lorenz(x, y, z, σ, ρ, β, generate_lorenz_eqs(t, x, y, z, σ, ρ, β))
end

@Param t
lz1 = Lorenz(t)
lz2 = Lorenz(t)
lz1.x ~ lz2.x


function kernel(x1, x2, hyp)
    n1, dim = size(x1, 1), size(x1, 2)
    n2 = size(x2, 1)

    b = exp.(hyp[1:end - 2])
    c = exp.(hyp[end - 1]);

    x1 = x1 .* repeat(transpose(sqrt.(b)), n1, 1)
    x2 = x2 .* repeat(transpose(sqrt.(b)), n2, 1)

    K = (-2 * x1 * transpose(x2) +
        repeat(transpose(sum(x2 .* x2, dims=2)), n1, 1) +
        repeat(sum(x1 .* x1, dims=2), 1, n2))
    K = c * exp.(-0.5 * K)

    return K
end


# Adapted from:
# SOSDEMO6 --- MAX CUT
# Section 3.6 of SOSTOOLS User's Manual

@testset "SOSDEMO6 with $(factory.constructor)" for factory in sdp_factories
    isscs(factory) && continue
    @polyvar x[1:5]

    # Number of cuts
    f = 2.5 - 0.5*x[1]*x[2] - 0.5*x[2]*x[3] - 0.5*x[3]*x[4] - 0.5*x[4]*x[5] - 0.5*x[5]*x[1]

    # Boolean constraints
    bc = vec(x).^2 .- 1

    @testset "with γ=$γ it should be $(feasible ? "feasible" : "infeasible")" for (γ, feasible) in [(3.9, false), (4.1, true)]
        model = SOSModel(factory)

        Z = monomials(x, 0:1)
        @variable model p1 SOSPoly(Z)

        Z = monomials(x, 0:2)
        @variable model p[1:5] Poly(Z)

        @constraint model p1*(γ-f) + dot(p, bc) >= (γ-f)^2

        JuMP.optimize!(model)

        if feasible
            @test JuMP.primal_status(model) == MOI.FEASIBLE_POINT
        else
            @test JuMP.dual_status(model) == MOI.INFEASIBILITY_CERTIFICATE
        end
    end
end


# -log(density)
using LinearAlgebra

function normal_(x::Vector{Float64}, μ::Vector{Float64}, Σ::Matrix{Float64})
    # TODO allow Σ a parameter
    d = x - μ
    return 0.5 * (d' * (Σ \ d))
end

function normal_(x::Real, μ::Real, σ::Real)
    # TODO allow σ a parameter
    d = x - μ
    return 0.5 * (d' * d) / (σ * σ)
end

# TODO figure out how to incorporate reparameterizations? (is that what we should call these)
function t_(nu::Float64, s::Float64)
    return function(x)
        n = 0.5 * x[1] * x[1]
        ig = 0.5 * nu * x[2] + 0.5 * nu * s * s * exp(-x[2])
        return n + ig
    end
end

function normal_precision_(x::Vector{Float64}, μ::Vector{Float64}, Ω::Matrix{Float64})
    d = x - μ
    return 0.5 * (d' * Ω * d)
end

function normal_precision_(x::Real, μ::Real, ω::Real)
    d = x - μ
    return 0.5 * d * d * ω
end

## below needs lots of reviewing
# logistic
function logistic(x)
    return 1 / (1 + exp(-x))
end

function genlogistic(N = 1001, K = 3, a = Normal(0, 10), b = Normal(0, 2.5), x = Normal(0, 10))
    α = rand(a)
    β = rand(b, K)
    θ = [α; β]
    X = [ones(N) rand(x, (N, K))]
    μ = X * θ
    p = logistic.(μ)
    y = convert(Vector{Int}, rand.(Bernoulli.(p)))
    return y, X, θ
end

# df = DataFrame(y=Y, x2 = X[:, 2], x3 = X[:, 3], x4 = X[:, 4])
# fit = glm(@formula(y ~ x2 + x3 + x4), df, Bernoulli())

function logisticmodel(y, X, θ)
    μ = X * θ
    β = θ[2:end-1]
    return -y' * μ + sum(log1p.(exp.(μ))) + 0.5 + θ[1] * θ[1] / 100 + 0.5 * β' * β / 10
end
# logisticreg(θ) = logisticmodel(y, X, θ)

function hierlogisticmodel(y, X, θ)
    β = θ[1:(end - 1)]
    K = length(β)
    μ = X * β
    likelihood =
    σ = θ[end]
    κ = exp(σ)
    prior = -0.5 * β' * β / κ ^ 2 - K * σ - 0.1 * κ + σ
    return likelihood + prior
end

heirlogisticreg(θ) = hierlogisticmodel(θ, y, X)

# linear
function genlinear(N = 1001, K = 3, a = Normal(0, 10), b = Normal(0, 2.5), x = Normal(0, 10))
    α = rand(a)
    β = rand(b, K)
    θ = [α; β]
    X = [ones(N) rand(x, (N, K))]
    μ = X * θ
    y = rand.(Normal.(μ, 1))
    return y, X, θ
end

# df = DataFrame(y=Y, x2 = X[:, 2], x3 = X[:, 3], x4 = X[:, 4])
# fit = glm(@formula(y ~ x2 + x3 + x4), df, Normal())

function linearmodel(y, X, θ)
    e = y - X * θ
    return 0.5 * e' * e + 0.5 * θ' * θ / 25
end
# linearreg(θ) = linearmodel(y, X, θ)

function hierlinearmodel(y, X, θ)
    β = θ[1:(end - 1)]
    K = length(β)
    σ = θ[end]
    e = y - X * β
    return 0.5 * e' * e + K*σ + 0.5 * θ' * θ / exp(2 * σ) + 0.5 * exp(2 * σ) / 25 - σ
end


abstract AbstractWeighting

immutable EqualWeighting <: AbstractWeighting end
weight_frequencies(::EqualWeighting, ξs) = ones(ξs)

immutable LoudnessWeighting <: AbstractWeighting
    samplerate::Int
end

function weight_frequencies(weighting::LoudnessWeighting, ξs)
    freqs = weighting.samplerate * ξs
    freqs2 = freqs .* freqs
    numerator = 12200.0 * freqs2
    sqden1 = (freqs2 + 20.6 * 20.6)
    sqden2 = sqrt.((freqs2 + 107.7 * 107.7) .* (freqs2 + 737.9 * 737.9))
    sqden3 = (freqs2 + 12200.0 * 12200.0)
    sqdenominator = sqden1 .* sqden2 .* sqden3
    denominator = sqrt.(sqdenominator)
    return numerator ./ denominator
end


module TestClassifier
	using Base.Test
	using kNN
	using DataArrays
	using DataFrames
    using RDatasets
    using Distances
    using StatsBase

    iris = data("datasets", "iris")
    X = array(iris[:, 1:4])'
    y = array(iris[:, 5])
    model = knn(X, y, metric = Euclidean())

    predict_k1 = predict(model, X, 1)
    predict_k2 = predict(model, X, 2)
    predict_k3 = predict(model, X, 3)
    predict_k4 = predict(model, X, 4)
    predict_k5 = predict(model, X, 5)

    mean(predict_k1 .== y)
    mean(predict_k2 .== y)
    mean(predict_k3 .== y)
    mean(predict_k4 .== y)
    mean(predict_k5 .== y)
end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

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

                     @variable(m, -8.232182974914569 <= q <= 6.070970235703473)

                     add_NL_constraint(m, :(log(1 + exp(-0.17705543456567296 + 0.5634748900178774*log(1 + exp(-0.9628319588335761 + 0.5826659544327617*$(x[1]) + -0.6099922924008303*$(x[2]) + 0.9101904633653426*$(x[3]))) + 0.6106029608527352*log(1 + exp(0.06629760484384972 + -0.7340864295918403*$(x[1]) + 0.2309908240081593*$(x[2]) + 0.7867070707759196*$(x[3]))) + 0.5384388086098779*log(1 + exp(-0.07370832879656497 + -0.32073040453400337*$(x[1]) + -0.07649911788086428*$(x[2]) + 0.03499506906851524*$(x[3]))))) + log(1 + exp(0.08053767290922664 + 0.911477397877237*log(1 + exp(-0.9628319588335761 + 0.5826659544327617*$(x[1]) + -0.6099922924008303*$(x[2]) + 0.9101904633653426*$(x[3]))) + 0.7134526118957916*log(1 + exp(0.06629760484384972 + -0.7340864295918403*$(x[1]) + 0.2309908240081593*$(x[2]) + 0.7867070707759196*$(x[3]))) + -0.34097520379571256*log(1 + exp(-0.07370832879656497 + -0.32073040453400337*$(x[1]) + -0.07649911788086428*$(x[2]) + 0.03499506906851524*$(x[3]))))) + log(1 + exp(0.5992621446321471 + 0.15304126986204158*log(1 + exp(-0.9628319588335761 + 0.5826659544327617*$(x[1]) + -0.6099922924008303*$(x[2]) + 0.9101904633653426*$(x[3]))) + -0.37134199071135*log(1 + exp(0.06629760484384972 + -0.7340864295918403*$(x[1]) + 0.2309908240081593*$(x[2]) + 0.7867070707759196*$(x[3]))) + 0.8747239793986465*log(1 + exp(-0.07370832879656497 + -0.32073040453400337*$(x[1]) + -0.07649911788086428*$(x[2]) + 0.03499506906851524*$(x[3]))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


function GDelta(CallPutFlag, S, X, T, r, b, v)

  d1 = (log(S / X) + (b + v^2 / 2) * T) / (v * sqrt(T))

  if CallPutFlag == "c"
    delta = exp((b - r) * T) * CND(d1)

  elseif CallPutFlag == "p"
    delta = -exp((b - r) * T) * CND(-d1)
  end

end


using GradDescent
using Test

include("constructors.jl")
include("quadatric.jl")
include("normal-mle.jl")
include("linear-regression.jl")
#include("optimizer-list.jl")


@testset "Derivatives" begin
    @testset "Normal finite-differences" begin
        B = FiniteDifferences(5,1.0)
        D = Derivative(axes(B,1))

        ∇ = applied(*, B', D, B)
        ∇² = applied(*, B', D', D, B)

        ∂ = materialize(∇)
        T = materialize(∇²)

        @test ∂ isa Tridiagonal
        @test T isa SymTridiagonal

        @test all(diag(∂) .== 0)
        @test all(diag(∂,1) .== 0.5)
        @test all(diag(∂,-1) .== -0.5)

        @test all(diag(T) .== -2)
        @test all(diag(T,1) .== 1)
        @test all(diag(T,-1) .== 1)
    end

    @testset "Restricted bases" begin
        N = 10
        ρ = 1.0

        R = StaggeredFiniteDifferences(N, ρ)
        r = axes(R, 1)
        D = Derivative(r)

        ∇ = R'D*R
        ∇² = R'D'D*R

        for (sela,selb) in Iterators.product([1:10, 3:6, 8:10, 5:10, 4:5], [1:10, 3:6, 8:10, 5:10, 4:5])
            # In case we have partially overlapping restrictions, the
            # bandwidth will not be the full three-point stencil, and if we
            # have no overlap at all, the bandwidth will be zero.
            overlap = sela ∩ selb
            expected_bandwidth = isempty(overlap) ? 0 : min(3, length(overlap)+1)

            Ra = R[:,sela]
            Rb = R[:,selb]

            ∂ = Ra'D*Rb
            @test ∂ == Matrix(∇)[sela,selb]
            if sela == selb
                @test bandwidths(∂) == (1,1)
            end
            @test ∂ isa BandedMatrix
            @test length(bandrange(∂)) == expected_bandwidth

            ∂² = Ra'D'*D*Rb
            @test ∂² == Matrix(∇²)[sela,selb]
            if sela == selb
                @test bandwidths(∂) == (1,1)
            end
            @test ∂² isa BandedMatrix
            @test length(bandrange(∂²)) == expected_bandwidth
        end
    end

    @testset "Implicit finite differences" begin
        N = 10
        ρ = 0.3
        Z = 1.0
        R = ImplicitFiniteDifferences(N, ρ)
        r = axes(R, 1)
        @testset "Singular origin=$(singular_origin)" for singular_origin = [true, false]
            D̃ = Derivative(r)
            D = singular_origin ? CoulombDerivative(D̃, Z, 0) : D̃

            λ,δβ₁ = if singular_origin
                √3 - 2, -Z*ρ*inv(12 - 10*Z*ρ)
            else
                0,0
            end

            ∂ = R'D*R
            @test ∂ isa CompactBases.ImplicitDerivative

            @test ∂.Δ isa Tridiagonal
            @test ∂.Δ[1,1] ≈ λ/2ρ
            @test all(iszero, ∂.Δ.d[2:end])
            @test all(∂.Δ.dl .≈ -1/2ρ)
            @test all(∂.Δ.du .≈ 1/2ρ)

            @test ∂.M isa SymTridiagonal
            @test ∂.M[1,1] ≈ (4+λ)/6
            @test all(∂.M.dv[2:end] .≈ 4/6)
            @test all(∂.M.ev .≈ 1/6)

            ∂² = R'D'D*R
            @test ∂ isa CompactBases.ImplicitDerivative

            @test ∂².Δ isa SymTridiagonal
            @test ∂².Δ[1,1] ≈ -2*(1+δβ₁)/ρ^2
            @test all(∂².Δ.dv[2:end] .≈ -2/ρ^2)
            @test all(∂².Δ.ev .≈ 1/ρ^2)

            @test ∂².M isa SymTridiagonal
            @test ∂².M[1,1] ≈ (-1/2)*(-(10-2δβ₁)/6)
            @test all(∂².M.dv[2:end] .≈ (-1/2)*(-10/6))
            @test all(∂².M.ev .≈ (-1/2)*(-1/6))
        end

        @testset "Restricted bases" begin
            R̃ = R[:, 1:5]
            D = Derivative(axes(R̃, 1))

            @test_throws DimensionMismatch R̃'D*R
            @test_throws DimensionMismatch R'D*R̃
            @test_throws DimensionMismatch R̃'D'D*R
            @test_throws DimensionMismatch R'D'D*R̃

            ∂ = R̃'D*R̃
            @test ∂ isa CompactBases.ImplicitDerivative
            @test size(∂) == (5,5)

            ∂² = R̃'D'D*R̃
            @test ∂² isa CompactBases.ImplicitDerivative
            @test size(∂²) == (5,5)

            @warn "Need to test correctness of ImplicitDerivative in restricted bases"
        end
    end
end

@testset "Derivative accuracy" begin
    Ns = 2 .^ (5:20)

    d = 1.0
    f,g,h,a,b = derivative_test_functions(d)

    @testset "kind = $B" for (order,B) in [(2,FiniteDifferences),
                                           (4,ImplicitFiniteDifferences)]
        @info "$B derivative accuracy"
        hs,ϵg,ϵh,ϵh′,pg,ph,ph′ = compute_derivative_errors(Ns, f, g, h, 1) do N
            B(range(a, stop=b, length=N+2)[2:end-1])
        end

        @test isapprox(pg, order, atol=0.03) || pg > order
        @test isapprox(ph, order, atol=0.03) || ph > order
    end
    @testset "kind = StaggeredFiniteDifferences" begin
        dd = b-a
        @testset "$(label)" for (order,label,fun) = [(2.0, "Uniform", N -> (N,dd/(N+1))),
                                                     (2.0, "Non-uniform", N -> (dd*(1 .- cos.(π*range(0,stop=1,length=N+2)[2:end-1]))/2,))]
            @info "$(label) StaggeredFiniteDifferences derivative accuracy"
            hs,ϵg,ϵh,ϵh′,pg,ph,ph′ = compute_derivative_errors(Ns, x -> f(x-dd/2), x -> g(x-dd/2), x -> h(x-dd/2), 1) do N
                StaggeredFiniteDifferences(fun(N)...)
            end

            @test isapprox(pg, order, atol=0.03) || pg > order
            @test isapprox(ph, order, atol=0.03) || ph > order
        end
    end
end

@testset "Particle-in-a-box" begin
    @testset "Eigenvalues convergence rate" begin
        Ns = 2 .^ (7:14)
        nev = 3
        L = 1.0
        @testset "kind = $B" for (order,B) in [(2,FiniteDifferences),
                                               (4,ImplicitFiniteDifferences)]
            @info "$B particle-in-a-box eigenvalues convergence rate"
            ϵλ,slopes,elapsed =
                compute_diagonalization_errors(N -> B(N, L/(N+1)), test_particle_in_a_box,
                                               Ns, L, nev,
                                               verbosity=1)
            for (p,o) in zip(slopes, [order,order-1,order-1])
                @test isapprox(p, o, atol=0.1) || p > o
            end
        end
    end
end

# This tests the discretization of the Laplacian, especially near
# the origin where the potential term is singular.
@testset "Hydrogen bound states" begin
    @testset "Eigenvalues and eigenvectors" begin
        rₘₐₓ = 300
        ρ = 0.25
        N = ceil(Int, rₘₐₓ/ρ + 1/2)

        R = StaggeredFiniteDifferences(N, ρ)
        r = axes(R, 1)

        D = CoulombDerivative(Derivative(Base.axes(R,1)), 1, 0)
        Tm = (R'D'D*R) / -2
        @test Tm isa SymTridiagonal

        V = R'*QuasiDiagonal((r -> -inv(r)).(r))*R

        H = Tm + V

        ee = eigen(H)

        n = 1:10
        λₐ = -inv.(2n.^2)

        abs_error = ee.values[n] - λₐ
        rel_error = abs_error ./ abs.(1e-10 .+ abs.(λₐ))

        @test abs_error[1] < 3e-5
        @test all(rel_error .< 1e-3)

        r̃ = locs(R)
        # Table 2.2 Foot (2005)
        Rₐ₀ = [2exp.(-r̃),
               -(1/2)^1.5 * 2 * (1 .- r̃/2).*exp.(-r̃/2), # Why the minus?
               (1/3)^1.5 * 2 * (1 .- 2r̃/3 .+ 2/3*(r̃/3).^2).*exp.(-r̃/3)]
        expected_errors = [1e-3,1e-3,2e-3]

        for i = 1:3
            v = ee.vectors[:,i]
            N = norm(v)*√ρ
            # The sign from the diagonalization is arbitrary; make max lobe positive
            N *= sign(v[argmax(abs.(v))])
            abs_error = v/N .- r̃.*Rₐ₀[i]
            @test norm(abs_error)/abs(1e-10+norm(r̃.*Rₐ₀[i])) < expected_errors[i]
        end
    end

    @testset "Eigenvalues convergence rate" begin
        Ns = 2 .^ (7:14)
        nev = 3
        rₘₐₓ = 100.0
        Z = 1.0
        ℓ = 0

        @testset "kind = $B" for (order,B) in [(2,StaggeredFiniteDifferences),
                                               (4,ImplicitFiniteDifferences)]
            @info "$B hydrogen eigenvalues convergence rate"
            ϵλ,slopes,elapsed = compute_diagonalization_errors(test_singular_scheme, Ns, Z, ℓ, nev, verbosity=1) do N
                ρ = B == StaggeredFiniteDifferences ? rₘₐₓ/(N+0.5) : rₘₐₓ/(N+1)
                B(N, ρ)
            end
            for p in slopes
                @test isapprox(p, order, atol=0.04) || p > order
            end
        end
    end
end


function reduce_subtrees!(v,r,reduce)
    @assert length(r) >= length(v)
    n = length(v)
    if n == 0; return nothing; end

    @inbounds begin
        # Reduce no-child nodes
        for i = n:-1:heapparent(n)+1
            r[i] = reduce(v[i])
        end

        # Reduce one-child node, if present
        for i = heapparent(n+1) : heapparent(n)
            r[i] = reduce(v[i],r[n])
        end

        # Reduce two-children nodes
        for i = heapparent(n)-1:-1:1
            r[i] = reduce(v[i], r[heapleft(i)], r[heapright(i)])
        end
    end
    return nothing
end

function Reducer!(v,r,reduce)
    function ReducerF(i, _, nc)
        if nc == 0; r[i] = reduce(v[i]); end
        if nc == 1; r[i] = reduce(v[i], r[heapleft(i)]); end
        if nc == 2; r[i] = reduce(v[i], r[heapleft(i)], r[heapright(i)]); end
        return :continue
    end
    return ReducerF
end


# ---
# title: 221. Maximal Square
# id: problem221
# author: Indigo
# date: 2021-05-29
# difficulty: Medium
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/maximal-square/description/>
# hidden: true
# ---
#
# Given an `m x n` binary `matrix` filled with `0`'s and `1`'s, _find the
# largest square containing only_ `1`'s _and return its area_.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/11/26/max1grid.jpg)
#
#
#
#     Input: matrix = [["1","0","1","0","0"],["1","0","1","1","1"],["1","1","1","1","1"],["1","0","0","1","0"]]
#     Output: 4
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/11/26/max2grid.jpg)
#
#
#
#     Input: matrix = [["0","1"],["1","0"]]
#     Output: 1
#
#
# **Example 3:**
#
#
#
#     Input: matrix = [["0"]]
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `m == matrix.length`
#   * `n == matrix[i].length`
#   * `1 <= m, n <= 300`
#   * `matrix[i][j]` is `'0'` or `'1'`.
#
#
## @lc code=start
using LeetCode

function maximal_square(matrix::Matrix{Char})
    length(matrix) == 0 && return 0
    maxs = any(==('1'), matrix)
    maxs == 0 && return 0
    m, n = size(matrix)
    dp = fill(0, size(matrix))
    for i in 1:m
        dp[i, 1] = matrix[i, 1] == '1'
    end
    for j in 1:n
        dp[1, j] = matrix[1, j] == '1'
    end
    for i in 2:m, j in 2:n
        if matrix[i, j] == '1'
            dp[i, j] = min(dp[i, j - 1], dp[i - 1, j], dp[i - 1, j - 1]) + 1
            maxs = max(maxs, dp[i, j])
        end
    end
    return maxs * maxs
end
## @lc code=end


using BulletsMaggots, Statistics, ProgressMeter

function run_exp(c::Int, reps::Int = 1000)
    println("Solving for code $(string4(c)) for $reps iterations")
    t = [bm_solver(c, false) for _ = 1:reps]
    println("Average steps to solve:   ", mean(t))
    println("Median steps to solve:    ", median(t))
    println("Standard deviation:       ", std(t))
    println("Range in number of steps: ", minimum(t), "-", maximum(t))
end

"""
`run_all(reps::Int=1000)` runs the solver repeatedly on the five possible
digit patterns for Bullets and Maggots.
"""
function run_all(reps::Int = 1000)
    codes = [0000, 0001, 0011, 0012, 0123]
    for c in codes
        run_exp(c, reps)
        println()
    end
end

function experiment(reps::Int = 1000, first_guess::Function = guess_any)
    println("Playing $reps games")
    t = zeros(Int, reps)
    PM = Progress(reps)
    for i = 1:reps
        t[i] = bm_solver(random_code(), false, first_guess)
        next!(PM)
    end
    println("Average steps to solve:   ", mean(t))
    println("Median steps to solve:    ", median(t))
    println("Standard deviation:       ", std(t))
    println("Range in number of steps: ", minimum(t), "-", maximum(t))
end

function first_move_experiment(reps::Int = 1000)

    sep = "-"^60

    println("Any first guess")
    experiment(reps, guess_any)

    println(sep)

    println("All distinct")
    experiment(reps, guess_distinct)

    println(sep)

    println("One pair")
    experiment(reps, guess_one_pair)

    println(sep)

    println("Two pairs")
    experiment(reps, guess_two_pairs)

    println(sep)
end


using Sinograms: bdd_2d
using Test: @test, @testset, @inferred
# using JLD: load

@testset "bdd_2d" begin
    deg = 1
    geo = (DSD = 100000, DSO = 99700, pSize = 1, dSize = 0.5, nPix = 256, nDet = 1024,
          theta = deg2rad.(0:deg:360-deg), isoX = 0, isoY = 0)
    phantomImg = shepp_logan(geo.nPix, SheppLoganToft())
    sinogramB = projection(phantomImg',geo)
    imageB = backprojection(sinogramB,geo)

    # sinogramB_correct = load("sinomat.jld")["sinomat"]
    # imageB_correct = load("imagemat.jld")["imagemat"]

    sino = @inferred projection(phantomImg',geo)
    @test sino isa Matrix
    @test size(sinogramB) == (360, 1024)
    image = @inferred backprojection(sinogramB,geo)
    @test image isa Matrix
    @test size(imageB) == (256, 256)

    # @test sinogramB == sinogramB_correct
    # @test imageB == imageB_correct

    #=
    using MIRTjim: jim

    p1 = jim(1:geo.nDet, 1:length(geo.theta), sinogramB'; aspect_ratio = :auto, clim=(0,68))
    p2 = jim(1:geo.nPix, 1:geo.nPix, imageB'; aspect_ratio = :auto, clim=(30,95))
    jim(p1,p2)
    =#

end


## Not finished!#

#=
Conway's Game of Life takes place on an infinite two-dimensional board of square cells. Each cell is either dead or alive, and at each tick, the following rules apply:

    Any live cell with less than two live neighbours dies.
    Any live cell with two or three live neighbours remains living.
    Any live cell with more than three live neighbours dies.
    Any dead cell with exactly three live neighbours becomes a live cell.

A cell neighbours another cell if it is horizontally, vertically, or diagonally adjacent.

Implement Conway's Game of Life. It should be able to be initialized with a starting list of live cell coordinates and the number of steps it should run for. Once initialized, it should print out the board state at each step. Since it's an infinite board, print out only the relevant coordinates, i.e. from the top-leftmost live cell to bottom-rightmost live cell.

You can represent a live cell with an asterisk (*) and a dead cell with a dot (.).
=#

function ConwayGameOfLife(livecells, steps::Int)
    board = initialize(livecells)
    for step in 0:steps
        println("Step: $step")
        display(board)
        for cell_i in 1:size(board, 1)
            for cell_j in 1:size(board, 2)
                board[cell_index] = nextstatus(cell_i, cell_j, board)
            end
        end
    end
end

function initialize(livecells)
    board = fill(".", maximum(livecells[:,1]), maximum(livecells[:,2]))
    for living_cell in 1:size(livecells, 1)
        board[livecells[living_cell, 1], livecells[living_cell, 2]] = "*"
    end
    return board
end

function nextstatus(cell_i, cell_j, board)
    previous_status = board[cell_i, cell_j]
    livingneighbour = livingneighbours(cell_i, cell_j, board)
    if previous_status == "*" && livingneighbour < 2
        new_status = "."
    elseif previous_status == "*" && livingneighbour >= 2 && livingneighbour <= 3
        new_status = "*"
    elseif previous_status == "*" && livingneighbour > 3
        new_status = "."
    elseif previous_status == "." && livingneighbour == 3
        new_status = "*"
    else
        new_status = previous_status
    end
    return new_status
end

function livingneighbours(cell_i, cell_j, board)

end

# Create random board
livecells = []
for i in 1:10
    push!(livecells, rand(collect(1:4)),rand(collect(1:4)))
end
livecells = reshape(livecells, 10, 2)


#!/usr/bin/env julia

# Different implementations of gradient computation in Fourier space, with
# performance comparisons.

# Some sample benchmark results (on Julia 1.4):
#
#     Transforms: (RFFT, FFT, FFT)
#     Input type: Float64
#     Global dimensions: (64, 32, 64)  ->  (33, 32, 64)
#     MPI topology: 2D decomposition (2×1 processes)
#
#     gradient_global_view!...            184.853 μs (0 allocations: 0 bytes)
#     gradient_global_view_explicit!...   124.993 μs (0 allocations: 0 bytes)
#     gradient_local!...                  146.369 μs (0 allocations: 0 bytes)
#     gradient_local_parent!...           145.743 μs (0 allocations: 0 bytes)
#     gradient_local_linear!...           145.679 μs (0 allocations: 0 bytes)
#     gradient_local_linear_explicit!...  145.543 μs (0 allocations: 0 bytes)
#
# This was obtained when running julia with the default optimisation level -O2.
#
# If one uses -O1 instead, `gradient_global_view_explicit!` becomes much
# slower than the rest, indicating that there's a lot of compiler optimisations
# going on specifically in that function.

# See also https://jipolanco.github.io/PencilFFTs.jl/dev/examples/gradient/.

using BenchmarkTools
using MPI
using PencilFFTs

using AbstractFFTs: fftfreq, rfftfreq
using Printf: @printf
using Random: randn!

const PA = PencilFFTs.PencilArrays

const INPUT_DIMS = (64, 32, 64)

const DEV_NULL = @static Sys.iswindows() ? "nul" : "/dev/null"

function generate_wavenumbers_r2c(dims::Dims{3})
    box_size = (2π, 2π, 2π)  # Lx, Ly, Lz
    sample_rate = 2π .* dims ./ box_size

    # In our case (Lx = 2π and Nx even), this gives kx = [0, 1, 2, ..., Nx/2].
    kx = rfftfreq(dims[1], sample_rate[1])

    # In our case (Ly = 2π and Ny even), this gives
    # ky = [0, 1, 2, ..., Ny/2-1, -Ny/2, -Ny/2+1, ..., -1] (and similarly for kz).
    ky = fftfreq(dims[2], sample_rate[2])
    kz = fftfreq(dims[3], sample_rate[3])

    (kx, ky, kz)
end

# Compute and return ∇θ in Fourier space, using global views.
function gradient_global_view!(∇θ_hat::NTuple{3,PencilArray},
                               θ_hat::PencilArray, kvec_global)
    # Generate OffsetArrays that take global indices.
    θ_glob = global_view(θ_hat)
    ∇θ_glob = map(global_view, ∇θ_hat)

    @inbounds for (n, I) in enumerate(CartesianIndices(θ_glob))
        i, j, k = Tuple(I)  # global indices

        # Wave number vector associated to current Cartesian index.
        kx = kvec_global[1][i]
        ky = kvec_global[2][j]
        kz = kvec_global[3][k]

        u = im * θ_glob[n]

        ∇θ_glob[1][n] = kx * u
        ∇θ_glob[2][n] = ky * u
        ∇θ_glob[3][n] = kz * u
    end

    ∇θ_hat
end

# This is the fastest implementation, and I'm not sure why!
function gradient_global_view_explicit!(∇θ_hat::NTuple{3,PencilArray},
                                        θ_hat::PencilArray, kvec_global)
    # Generate OffsetArrays that take global indices.
    θ_glob = global_view(θ_hat)
    ∇θ_glob = map(global_view, ∇θ_hat)

    rng = axes(θ_glob)  # (i1:i2, j1:j2, k1:k2)

    # Note: since the dimensions in Fourier space are permuted as (z, y, x), it
    # is faster to loop with `k` as the fastest index.
    @assert permutation(θ_hat) === Permutation(3, 2, 1)

    @inbounds for i in rng[1], j in rng[2], k in rng[3]
        # Wave number vector associated to current Cartesian index.
        kx = kvec_global[1][i]
        ky = kvec_global[2][j]
        kz = kvec_global[3][k]

        u = im * θ_glob[i, j, k]

        ∇θ_glob[1][i, j, k] = kx * u
        ∇θ_glob[2][i, j, k] = ky * u
        ∇θ_glob[3][i, j, k] = kz * u
    end

    ∇θ_hat
end

# Compute and return ∇θ in Fourier space, using local indices.
function gradient_local!(∇θ_hat::NTuple{3,PencilArray}, θ_hat::PencilArray,
                         kvec_local)
    @inbounds for (n, I) in enumerate(CartesianIndices(θ_hat))
        i, j, k = Tuple(I)  # local indices

        # Wave number vector associated to current Cartesian index.
        kx = kvec_local[1][i]
        ky = kvec_local[2][j]
        kz = kvec_local[3][k]

        u = im * θ_hat[n]

        ∇θ_hat[1][n] = kx * u
        ∇θ_hat[2][n] = ky * u
        ∇θ_hat[3][n] = kz * u
    end

    ∇θ_hat
end

# Compute and return ∇θ in Fourier space, using local indices on the raw data
# (which takes permuted indices).
function gradient_local_parent!(∇θ_hat::NTuple{3,PencilArray},
                                θ_hat::PencilArray, kvec_local)
    θ_p = parent(θ_hat) :: Array
    ∇θ_p = parent.(∇θ_hat)

    perm = permutation(θ_hat)

    @inbounds for (n, I) in enumerate(CartesianIndices(θ_p))
        # Unpermute indices to (i, j, k)
        J = perm \ I

        # Wave number vector associated to current Cartesian index.
        i, j, k = Tuple(J)  # local indices
        kx = kvec_local[1][i]
        ky = kvec_local[2][j]
        kz = kvec_local[3][k]

        u = im * θ_p[n]

        ∇θ_p[1][n] = kx * u
        ∇θ_p[2][n] = ky * u
        ∇θ_p[3][n] = kz * u
    end

    ∇θ_hat
end

# Similar to gradient_local!, but avoiding CartesianIndices (slightly faster).
function gradient_local_linear!(∇θ_hat::NTuple{3,PencilArray},
                                θ_hat::PencilArray, kvec_local)
    # We want to iterate over the arrays in memory order to maximise
    # performance. For this we need to take into account the permutation of
    # indices in the Fourier-transformed arrays. By default, the memory order in
    # Fourier space is (z, y, x) instead of (x, y, z), but this is never assumed
    # below. The wave numbers must be permuted accordingly.
    perm = permutation(θ_hat)  # e.g. Permutation(3, 2, 1)
    kvec_perm = perm * kvec_local  # e.g. (kz, ky, kx)

    # Create wave number iterator.
    kvec_iter = Iterators.product(kvec_perm...)

    @inbounds for (n, kvec_n) in enumerate(kvec_iter)
        # Apply inverse permutation to the current wave number vector.
        # Note that this permutation has zero cost, since perm is a
        # compile-time constant!
        # (This can be verified by comparing the performance of this function
        # with the "explicit" variant of `gradient_local_linear`, below.)
        κ = perm \ kvec_n  # = (kx, ky, kz)

        u = im * θ_hat[n]

        # Note that this is very easy to generalise to N dimensions...
        ∇θ_hat[1][n] = κ[1] * u
        ∇θ_hat[2][n] = κ[2] * u
        ∇θ_hat[3][n] = κ[3] * u
    end

    ∇θ_hat
end

# Less generic version of the above, assuming that the permutation is (3, 2, 1).
# It's basically the same but probably easier to understand.
function gradient_local_linear_explicit!(∇θ_hat::NTuple{3,PencilArray},
                                         θ_hat::PencilArray, kvec_local)
    @assert permutation(θ_hat) === Permutation(3, 2, 1)

    # Create wave number iterator in (kz, ky, kx) order, i.e. in the same order
    # as the array data.
    kvec_iter = Iterators.product(kvec_local[3], kvec_local[2], kvec_local[1])

    @inbounds for (n, kvec_n) in enumerate(kvec_iter)
        kz, ky, kx = kvec_n
        u = im * θ_hat[n]
        ∇θ_hat[1][n] = kx * u
        ∇θ_hat[2][n] = ky * u
        ∇θ_hat[3][n] = kz * u
    end

    ∇θ_hat
end

MPI.Init()

# Input data dimensions (Nx × Ny × Nz)
dims = INPUT_DIMS

kvec = generate_wavenumbers_r2c(dims)  # as tuple of Frequencies
kvec_collected = collect.(kvec)        # as tuple of Vector

# Apply a 3D real-to-complex (r2c) FFT.
transform = Transforms.RFFT()

# MPI topology information
comm = MPI.COMM_WORLD
Nproc = MPI.Comm_size(comm)
rank = MPI.Comm_rank(comm)

# Disable output on all but one process.
rank == 0 || redirect_stdout(open(DEV_NULL, "w"))

# Let MPI_Dims_create choose the decomposition.
proc_dims = let pdims = zeros(Int, 2)
    MPI.Dims_create!(Nproc, pdims)
    pdims[1], pdims[2]
end

# Create plan
plan = PencilFFTPlan(dims, transform, proc_dims, comm)
println(plan, "\n")

# Allocate data and initialise field
θ = allocate_input(plan)
randn!(θ)

# Perform distributed FFT
θ_hat = plan * θ

# Compute and compare gradients using different methods.
# Note that these return a tuple of 3 PencilArrays representing a vector
# field.
∇θ_hat_base = allocate_output(plan, Val(3))
∇θ_hat_other = similar.(∇θ_hat_base)

# Local wave numbers: (kx[i1:i2], ky[j1:j2], kz[k1:k2]).
kvec_local = let rng = range_local(θ_hat)  # = (i1:i2, j1:j2, k1:k2)
    getindex.(kvec, rng)
end

gradient_global_view!(∇θ_hat_base, θ_hat, kvec)

@printf "%-40s" "gradient_global_view!..."
@btime gradient_global_view!($∇θ_hat_other, $θ_hat, $kvec_collected)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

@printf "%-40s" "gradient_global_view! (lazy)..."
@btime gradient_global_view!($∇θ_hat_other, $θ_hat, $kvec)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

# Slow version: kvec as a tuple of Vector
@printf "%-40s" "gradient_global_view_explicit!..."
@btime gradient_global_view_explicit!($∇θ_hat_other, $θ_hat, $kvec_collected)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

# Fast version: kvec as a tuple of lazy vectors (Frequencies)
@printf "%-40s" "gradient_global_view_explicit! (lazy)..."
@btime gradient_global_view_explicit!($∇θ_hat_other, $θ_hat, $kvec)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

@printf "%-40s" "gradient_local!..."
@btime gradient_local!($∇θ_hat_other, $θ_hat, $kvec_local)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

@printf "%-40s" "gradient_local_parent!..."
@btime gradient_local_parent!($∇θ_hat_other, $θ_hat, $kvec_local)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

@printf "%-40s" "gradient_local_linear!..."
@btime gradient_local_linear!($∇θ_hat_other, $θ_hat, $kvec_local)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

@printf "%-40s" "gradient_local_linear_explicit!..."
@btime gradient_local_linear_explicit!($∇θ_hat_other, $θ_hat, $kvec_local)
@assert all(∇θ_hat_base .≈ ∇θ_hat_other)

# Get gradient in physical space.
∇θ = plan \ ∇θ_hat_base


module Output

# supports output functionality for the results of the sandpile symulations
# outputs the final grid in CSV format, as well as an image file

using CSV
using DataFrames
using Images

function TrimZeros(A)
  # given an array A trims any zero rows/columns from its borders
  # returns a 4 tuple of integers, i1, i2, j1, j2, where the trimmed array corresponds to A[i1:i2, j1:j2]
  # A can be either numeric or a boolean array

    i1, j1 = 1, 1;
    i2, j2 = size(A);

    zz = typeof(A[1,1])(0);   # comparison of a value takes into account the type as well

    # i1 is the first row which has non zero element
    for i = 1:size(A,1)
      q = false;
      for k = 1:size(A,2)
        if A[i,k] != zz
          q = true;
          #print("found non-zero at i1, k  = ", i1, " ", k , " value = ", A[i1, k], " \n");
          i1 = i;
          break
        end
      end

      if q == true
        break
      end
    end

    # i2 is the first from below row with non zero element
    for i in size(A,1):-1:1
      q = false;
      for k = 1:size(A,2)
        if A[i,k] != zz
          q = true;
          i2 = i;
          break
        end
      end

      if q == true
        break
      end
    end

    # j1 is the first column with non zero element

    for j = 1:size(A,2)
      q = false;
      for k = 1:size(A,1)
        if A[k,j] != zz
          j1 = j;
          q = true;
          break
        end
      end

      if q == true
        break
      end
    end

    # j2 is the last column with non zero element

    for j in size(A,2):-1:1
      q=false;
      for k=1:size(A,1)
        if A[k,j] != zz
          j2 = j;
          q=true;
          break
        end
      end

      if q==true
        break
      end
    end

    return i1, i2, j1, j2;

end

function addLayerofZeros(A, extraLayer)
  # adds layer of zeros from all corners to the given array A

  if extraLayer <= 0
    return A;
  end

  N, M = size(A);

  Z = zeros( typeof(A[1,1]), N + 2*extraLayer, M + 2*extraLayer );
  Z[(extraLayer+1):(N + extraLayer ), (extraLayer+1):(M + extraLayer ) ] = A;

  return Z;

end

function printIntoFile(A, extraLayer, strFileName, TrimSmallValues = false)
  # exports a 2d matrix A into a csv file
  # @extraLayer is an integers adding layer of 0-s sorrounding the output matrix

  # trimming off very small values; tiny values affect the performance of CSV export
  if TrimSmallValues == true
    A = map( x -> if ( abs( x - floor(x)  ) < 0.01   )  floor(x) else x end, A ) ;
  end

  i1, i2, j1, j2  = TrimZeros( A );
  A = A[i1:i2, j1:j2];

  A = addLayerofZeros(A, extraLayer);

  CSV.write(string(strFileName,".csv"), DataFrame(A), writeheader = false);

  return A;

end

function Array_magnifier(A, cell_mag, border_mag)
  # A is the main array; @cell_mag is the magnifying size of the cell,
  # @border_mag is the magnifying size of the border between lattice cells

  # creates a new array where each cell of the original array A appears magnified by size = cell_mag

  total_factor = cell_mag + border_mag;

  A1 = zeros( typeof(A[1,1]), total_factor*size(A,1), total_factor*size(A,2) );

  for i = 1:size(A,1)
       for j = 1:size(A,2)
            for u = ((i-1)*total_factor +1):(i*total_factor)
            for v = ((j-1)*total_factor + 1):(j*total_factor)
               if ( ( u - (i-1)*total_factor <= cell_mag  ) && ( v - (j-1)*total_factor <= cell_mag ) )
                  A1[u, v] = A[i,j] ;
               end
            end
            end
      end
  end

  return A1;

end

function saveAsGrayImage(A, fileName, cell_mag, border_mag, TrimSmallValues = false)
    # given a 2d matrix A, we save it as a gray image after magnifying by the given factors
    A1 = Array_magnifier(A, cell_mag, border_mag);
    A1 = A1/maximum(maximum(A1));

    # trimming very small values from A1 to improve performance
    if TrimSmallValues == true
        A1 = map( x -> if ( x < 0.01 ) 0.0 else round(x, digits = 2) end, A1  ) ;
    end

    save( string(fileName, ".png") , colorview(Gray, A1)) ;
end

function saveAsRGBImage(A, fileName, color_codes, cell_mag, border_mag)
  # color_codes is a dictionary, where key is a value in A and value is an RGB triplet
  # given a 2d array A, and color codes (mapping from values in A to RGB triples), save A
  # into fileName as png image after applying the magnifying factors

  #R = zeros(size(A));
  #G = zeros(size(A));
  #B = zeros(size(A));
  A1 = Array_magnifier(A, cell_mag, border_mag);
  color_mat = zeros(UInt8, (3, size(A1,1), size(A1,2)) );

  for i = 1:size(A1,1)
    for j = 1:size(A1,2)
        color_mat[:, i,j]  = get(color_codes, A1[i,j] , [0,0,0] )
    end
  end

  #A1 = Array_magnifier(A, cell_mag, border_mag)
  save( string(fileName, ".png") , colorview(RGB, color_mat/255)) ;

end

end # end of the module


"""
    nfan(n_out, n_in=1) -> Tuple
    nfan(dims...)
    nfan(dims::Tuple)

For a layer characterized by dimensions `dims`, return a tuple `(fan_in, fan_out)`, where `fan_in`
is the number of input neurons connected to an output one, and `fan_out` is the number of output neurons
connected to an input one.

This function is mainly used by weight initializers, e.g., [`kaiming_normal`](@ref Flux.kaiming_normal).

# Examples

```jldoctest
julia> layer = Dense(10, 20);

julia> Flux.nfan(size(layer.weight))
(10, 20)

julia> layer = Conv((3, 3), 2=>10);

julia> Flux.nfan(size(layer.weight))
(18, 90)
```
"""
nfan() = 1, 1 # fan_in, fan_out
nfan(n) = 1, n # A vector is treated as a n×1 matrix
nfan(n_out, n_in) = n_in, n_out # In case of Dense kernels: arranged as matrices
nfan(dims::Tuple) = nfan(dims...)
nfan(dims...) = prod(dims[1:end-2]) .* (dims[end-1], dims[end]) # In case of convolution kernels

ofeltype(x, y) = convert(float(eltype(x)), y)
epseltype(x) = eps(float(eltype(x)))

"""
    rng_from_array([x])

Create an instance of the RNG most appropriate for `x`.
The current defaults are:
- `x isa CuArray`: `CUDA.default_rng()`, else:
- `x isa AbstractArray`, or no `x` provided:
  - Julia version is < 1.7: `Random.GLOBAL_RNG`
  - Julia version is >= 1.7: `Random.default_rng()`
"""
rng_from_array(::AbstractArray) = rng_from_array()
rng_from_array(::CuArray) = CUDA.default_rng()
if VERSION >= v"1.7"
  rng_from_array() = Random.default_rng()
else
  rng_from_array() = Random.GLOBAL_RNG
end

"""
    glorot_uniform([rng=GLOBAL_RNG], size...; gain = 1) -> Array
    glorot_uniform([rng]; kw...) -> Function

Return an `Array{Float32}` of the given `size` containing random numbers drawn from a uniform
distribution on the interval ``[-x, x]``, where `x = gain * sqrt(6 / (fan_in + fan_out))`.

This method is described in [1] and also known as Xavier initialization.

# Examples
```jldoctest; setup = :(using Random; Random.seed!(0))
julia> Flux.glorot_uniform(3, 4) |> summary
"3×4 Matrix{Float32}"

julia> round.(extrema(Flux.glorot_uniform(10, 100)), digits=3)
(-0.232f0, 0.234f0)

julia> round.(extrema(Flux.glorot_uniform(100, 10)), digits=3)
(-0.233f0, 0.233f0)

julia> round.(extrema(Flux.glorot_uniform(100, 100)), digits=3)
(-0.173f0, 0.173f0)

julia> Dense(3 => 2, tanh; init = Flux.glorot_uniform(MersenneTwister(1)))
Dense(3 => 2, tanh)  # 8 parameters

julia> ans.bias
2-element Vector{Float32}:
 0.0
 0.0
```

# References

[1] Glorot, Xavier, and Yoshua Bengio. "Understanding the difficulty of training deep feedforward neural networks." _Proceedings of the thirteenth international conference on artificial intelligence and statistics_. 2010.
"""
function glorot_uniform(rng::AbstractRNG, dims::Integer...; gain::Real=1)
  scale = Float32(gain) * sqrt(24.0f0 / sum(nfan(dims...)))
  (rand(rng, Float32, dims...) .- 0.5f0) .* scale
end
glorot_uniform(dims::Integer...; kw...) = glorot_uniform(rng_from_array(), dims...; kw...)
glorot_uniform(rng::AbstractRNG=rng_from_array(); init_kwargs...) = (dims...; kwargs...) -> glorot_uniform(rng, dims...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable glorot_uniform(::Any...)

"""
    glorot_normal([rng=GLOBAL_RNG], size...; gain = 1) -> Array
    glorot_normal([rng]; kw...) -> Function

Return an `Array{Float32}` of the given `size` containing random numbers drawn from a normal
distribution with standard deviation `gain * sqrt(2 / (fan_in + fan_out))`,
using [`nfan`](@ref Flux.nfan).

This method is described in [1] and also known as Xavier initialization.

# Examples
```jldoctest; setup = :(using Random; Random.seed!(0))
julia> using Statistics

julia> round(std(Flux.glorot_normal(10, 1000)), digits=3)
0.044f0

julia> round(std(Flux.glorot_normal(1000, 10)), digits=3)
0.044f0

julia> round(std(Flux.glorot_normal(1000, 1000)), digits=3)
0.032f0

julia> Dense(10 => 1000, tanh; init = Flux.glorot_normal(gain=100))
Dense(10 => 1000, tanh)  # 11_000 parameters

julia> round(std(ans.weight), sigdigits=3)
4.45f0
```

# References

[1] Glorot, Xavier, and Yoshua Bengio. "Understanding the difficulty of training deep feedforward neural networks." _Proceedings of the thirteenth international conference on artificial intelligence and statistics_. 2010.
"""
function glorot_normal(rng::AbstractRNG, dims::Integer...; gain::Real=1)
  std = Float32(gain) * sqrt(2.0f0 / sum(nfan(dims...)))
  randn(rng, Float32, dims...) .* std
end
glorot_normal(dims::Integer...; kwargs...) = glorot_normal(rng_from_array(), dims...; kwargs...)
glorot_normal(rng::AbstractRNG=rng_from_array(); init_kwargs...) = (dims...; kwargs...) -> glorot_normal(rng, dims...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable glorot_normal(::Any...)

"""
    kaiming_uniform([rng=GLOBAL_RNG], size...; gain = √2) -> Array
    kaiming_uniform([rng]; kw...) -> Function

Return an `Array{Float32}` of the given `size` containing random numbers drawn from a uniform distribution
on the interval `[-x, x]`, where `x = gain * sqrt(3/fan_in)` using [`nfan`](@ref Flux.nfan).

This method is described in [1] and also known as He initialization.

# Examples
```jldoctest; setup = :(using Random; Random.seed!(0))
julia> round.(extrema(Flux.kaiming_uniform(100, 10)), digits=3)
(-0.774f0, 0.774f0)

julia> round.(extrema(Flux.kaiming_uniform(10, 100)), digits=3)
(-0.245f0, 0.244f0)

julia> round.(extrema(Flux.kaiming_uniform(100, 100)), digits=3)
(-0.245f0, 0.245f0)
```

# References

[1] He, Kaiming, et al. "Delving deep into rectifiers: Surpassing human-level performance on imagenet classification." _Proceedings of the IEEE international conference on computer vision_. 2015.
"""
function kaiming_uniform(rng::AbstractRNG, dims::Integer...; gain::Real = √2)
  bound = Float32(√3 * gain / sqrt(first(nfan(dims...)))) # fan_in
  return (rand(rng, Float32, dims...) .- 0.5f0) .* 2bound
end

kaiming_uniform(dims::Integer...; kwargs...) = kaiming_uniform(rng_from_array(), dims...; kwargs...)
kaiming_uniform(rng::AbstractRNG=rng_from_array(); init_kwargs...) = (dims...; kwargs...) -> kaiming_uniform(rng, dims...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable kaiming_uniform(::Any...)

"""
    kaiming_normal([rng=GLOBAL_RNG], size...; gain = √2) -> Array
    kaiming_normal([rng]; kw...) -> Function

Return an `Array{Float32}` of the given `size` containing random numbers taken from a normal
distribution standard deviation `gain / sqrt(fan_in)`, using [`nfan`](@ref Flux.nfan).

This method is described in [1] and also known as He initialization.

# Examples
```jldoctest; setup = :(using Random; Random.seed!(0))
julia> using Statistics

julia> round(std(Flux.kaiming_normal(10, 1000)), digits=3)
0.045f0

julia> round(std(Flux.kaiming_normal(1000, 10)), digits=3)
0.447f0

julia> round(std(Flux.kaiming_normal(1000, 1000)), digits=3)
0.045f0
```

# References

[1] He, Kaiming, et al. "Delving deep into rectifiers: Surpassing human-level performance on imagenet classification." _Proceedings of the IEEE international conference on computer vision_. 2015.
"""
function kaiming_normal(rng::AbstractRNG, dims::Integer...; gain::Real = √2f0)
  std = Float32(gain / sqrt(first(nfan(dims...)))) # fan_in
  return randn(rng, Float32, dims...) .* std
end

kaiming_normal(dims::Integer...; kwargs...) = kaiming_normal(rng_from_array(), dims...; kwargs...)
kaiming_normal(rng::AbstractRNG; init_kwargs...) = (dims...; kwargs...) -> kaiming_normal(rng, dims...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable kaiming_normal(::Any...)

"""
    truncated_normal([rng=GLOBAL_RNG], size...; mean = 0, std = 1, lo = -2, hi = 2) -> Array
    truncated_normal([rng]; kw...) -> Function

Return an `Array{Float32}` of the given `size` where each element is drawn from a truncated normal distribution.
The numbers are distributed like `filter(x -> lo<=x<=hi, mean .+ std .* randn(100))`.

The values are generated by sampling a Uniform(0, 1) (`rand()`) and then
applying the inverse CDF of the truncated normal distribution.
This method works best when `lo ≤ mean ≤ hi`.

# Examples
```jldoctest
julia> using Statistics

julia> Flux.truncated_normal(3, 4) |> summary
"3×4 Matrix{Float32}"

julia> round.(extrema(Flux.truncated_normal(10^6)); digits=3)
(-2.0f0, 2.0f0)

julia> round(std(Flux.truncated_normal(10^6; lo = -100, hi = 100)))
1.0f0
```
"""
function truncated_normal(rng::AbstractRNG, dims::Integer...; mean = 0, std = 1, lo = -2, hi = 2)
  norm_cdf(x) = 0.5 * (1 + erf(x/√2))
  if (mean < lo - 2 * std) || (mean > hi + 2 * std)
    @warn "Mean is more than 2 std outside the limits in truncated_normal, so the distribution of values may be inaccurate." maxlog=1
  end
  l = norm_cdf((lo - mean) / std)
  u = norm_cdf((hi - mean) / std)
  xs = rand(rng, Float32, dims...)
  broadcast!(xs, xs) do x
    x = x * 2(u - l) + (2l - 1)
    x = erfinv(x)
    x = clamp(x * std * √2 + mean, lo, hi)
  end
  return xs
end

truncated_normal(dims::Integer...; kwargs...) = truncated_normal(rng_from_array(), dims...; kwargs...)
truncated_normal(rng::AbstractRNG=rng_from_array(); init_kwargs...) = (dims...; kwargs...) -> truncated_normal(rng, dims...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable truncated_normal(::Any...)

"""
    orthogonal([rng=GLOBAL_RNG], size...; gain = 1) -> Array
    orthogonal([rng]; kw...) -> Function

Return an `Array{Float32}` of the given `size` which is a (semi) orthogonal matrix, as described in [1].

Cannot construct a vector, i.e. `length(size) == 1` is forbidden.
For `length(size) > 2`, a `prod(size[1:(end - 1)])` by `size[end]` orthogonal matrix
is computed before reshaping it to the original dimensions.

# Examples
```jldoctest; setup = :(using LinearAlgebra)
julia> W = Flux.orthogonal(5, 7);

julia> summary(W)
"5×7 Matrix{Float32}"

julia> W * W' ≈ I(5)
true

julia> W2 = Flux.orthogonal(7, 5);

julia> W2 * W2' ≈ I(7)
false

julia> W2' * W2 ≈ I(5)
true

julia> W3 = Flux.orthogonal(3, 3, 2, 4);

julia> transpose(reshape(W3, :, 4)) * reshape(W3, :, 4) ≈ I(4)
true
```

# References

[1] Saxe, McClelland, Ganguli. "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks", ICLR 2014, https://arxiv.org/abs/1312.6120

"""
function orthogonal(rng::AbstractRNG, rows::Integer, cols::Integer; gain::Real = 1)
  if rows < cols
    return permutedims(orthogonal(rng, cols, rows; gain))
  end
  mat = randn(rng, Float32, rows, cols)
  Q, R = LinearAlgebra.qr(mat)
  mat .= Array(Q) * sign.(LinearAlgebra.Diagonal(R)) .* Float32(gain)
  return mat
end

function orthogonal(rng::AbstractRNG, d1::Integer, ds::Integer...; kwargs...)
  dims = (d1, ds...)
  rows = prod(dims[1:end-1])
  cols = dims[end]
  return reshape(orthogonal(rng, rows, cols; kwargs...), dims)
end

orthogonal(dims::Integer...; kwargs...) = orthogonal(rng_from_array(), dims...; kwargs...)
orthogonal(rng::AbstractRNG=rng_from_array(); init_kwargs...) = (dims::Integer...; kwargs...) -> orthogonal(rng, dims...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable orthogonal(::Any...)

"""
    sparse_init([rng=GLOBAL_RNG], rows, cols; sparsity, std = 0.01) -> Array
    sparse_init([rng]; kw...) -> Function

Return a `Matrix{Float32}` of size `rows, cols` where each column contains a fixed fraction of
zero elements given by `sparsity`. Non-zero elements are normally distributed
with a mean of zero and standard deviation `std`.

This method is described in [1].

# Examples
```jldoctest; setup = :(using Random; Random.seed!(0))
julia> count(iszero, Flux.sparse_init(10, 10, sparsity=1/5))
20

julia> sum(0 .== Flux.sparse_init(10, 11, sparsity=0.9), dims=1)
1×11 Matrix{Int64}:
 9  9  9  9  9  9  9  9  9  9  9

julia> Dense(3 => 10, tanh; init=Flux.sparse_init(sparsity=0.5))
Dense(3 => 10, tanh)  # 40 parameters

julia> count(iszero, ans.weight, dims=1)
1×3 Matrix{Int64}:
 5  5  5
```

# References

[1] Martens, J, "Deep learning via Hessian-free optimization" _Proceedings of the 27th International Conference on International Conference on Machine Learning_. 2010.
"""
function sparse_init(rng::AbstractRNG, dims::Integer...; sparsity, std = 0.01)
  if length(dims) != 2
    throw(ArgumentError("Only 2-dimensional outputs are supported for sparse initialization."))
  end
  rows, cols = dims
  prop_zero = min(1.0, sparsity)
  num_zeros = ceil(Integer, prop_zero * rows)
  sparse_array = randn(rng, Float32, dims...) .* Float32(std)
  sparse_array[1:num_zeros, :] .= 0f0
  return mapslices(shuffle, sparse_array, dims=1)
end

sparse_init(dims::Integer...; kwargs...) = sparse_init(rng_from_array(), dims...; kwargs...)
sparse_init(rng::AbstractRNG=rng_from_array(); init_kwargs...) = (dims...; kwargs...) -> sparse_init(rng, dims...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable sparse_init(::Any...)

"""
    identity_init(size...; gain=1, shift=0) -> Array
    identity_init(; kw...) -> Function

Return an `Array{Float32}` of the given `size` which yields an identity mapping when used as parameters in
most Flux layers. Use `gain` to scale the identity by a constant.

Often useful in the context of transfer learning, i.e when one wants to add more capacity to
a model but start from the same mapping.

Has the following behaviour
*  1D: A `Vector` of `zeros` (useful for an identity bias)
*  2D: An identity matrix (useful for an identity matrix multiplication)
*  More than 2D: A dense block array of center tap spatial filters (useful for an identity convolution)

Some caveats:
* Not all layers will be identity mapping when used with this init. Exceptions
  include recurrent layers and normalization layers.

* Layers must have `input_size == output_size` for identity mapping to be
  possible. When this is not the case, extra dimensions of the array are padded with zeros.

* For convolutional layers, in addition to the above, the kernel sizes must also be odd and
  padding must be applied so that output feature maps have the same size as input feature maps,
  e.g by using [`SamePad`](@ref).

Use keyword `shift` (integer or tuple) to apply circular shift to the output,
equivalent to `Base.circshift(identity_init(size...), shift)`.

For consistency with other initialisers, it accepts `rng::AbstractRNG` as an optional
first argument. But this is ignored, since the result is not random.

# Examples
```jldoctest
julia> Flux.identity_init(3,5)
3×5 Matrix{Float32}:
 1.0  0.0  0.0  0.0  0.0
 0.0  1.0  0.0  0.0  0.0
 0.0  0.0  1.0  0.0  0.0

julia> Dense(5 => 3, relu, init=Flux.identity_init)([1,-2,3,-4,5])
3-element Vector{Float32}:
 1.0
 0.0
 3.0

julia> Flux.identity_init(3,3,2; gain=100)
3×3×2 Array{Float32, 3}:
[:, :, 1] =
   0.0  0.0  0.0
 100.0  0.0  0.0
   0.0  0.0  0.0

[:, :, 2] =
 0.0    0.0  0.0
 0.0  100.0  0.0
 0.0    0.0  0.0

julia> x4 = cat([1 2 3; 4 5 6; 7 8 9]; dims=4);

julia> Conv((2,2), 1 => 1, init=Flux.identity_init(gain=10), pad=SamePad())(x4)
3×3×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
 10.0  20.0  30.0
 40.0  50.0  60.0
 70.0  80.0  90.0
```
"""
identity_init(cols::Integer; gain::Real=1, shift=0) = zeros32(cols) # Assume bias

# Assume matrix multiplication
identity_init(rows::Integer, cols::Integer; gain::Real=1, shift=0) = circshift(Matrix{Float32}(I * gain, rows,cols), shift)

# Assume convolution
function identity_init(dims::Integer...; gain::Real=1, shift=0)
  nin, nout = dims[end-1], dims[end]
  centers = map(d -> cld(d, 2), dims[1:end-2])
  weights = zeros32(dims...)
  for i in 1:min(nin,nout)
    weights[centers..., i, i] = gain
  end
  return circshift(weights, shift)
end

# For consistency, it accepts an RNG, but ignores it:
identity_init(::AbstractRNG, dims::Integer...; kwargs...) = identity_init(dims...; kwargs...)
identity_init(rng::AbstractRNG=rng_from_array(); init_kwargs...) = (args...;kwargs...) -> identity_init(rng, args...; init_kwargs..., kwargs...)

ChainRulesCore.@non_differentiable identity_init(::Any...)

ones32(dims::Integer...) = Base.ones(Float32, dims...)
zeros32(dims::Integer...) = Base.zeros(Float32, dims...)

"""
    ones32(size...) = ones(Float32, size...)
    zeros32(size...) = zeros(Float32, size...)

Return an `Array{Float32}` of the given `size`.
"""
ones32(dims...) = Base.ones(Float32, dims...)

@doc @doc(ones32)
zeros32(dims...) = Base.zeros(Float32, dims...)

"""
    rand32([rng], size...)
    randn32([rng], size...)

Return an `Array{Float32}` of the given `size`, filled like `rand` or `randn`.
When the size is not provided, `rand32(rng::AbstractRNG)` returns a function.
"""
rand32(dims::Integer...) = Base.rand(Float32, dims...)
rand32(rng::AbstractRNG, dims::Integer...) = Base.rand(rng, Float32, dims...)
rand32(rng::AbstractRNG) = (dims...,) -> Base.rand(rng, Float32, dims...)

@doc @doc(rand32)
randn32(dims::Integer...) = Base.randn(Float32, dims...)
randn32(rng::AbstractRNG, dims::Integer...) = Base.randn(rng, Float32, dims...)
randn32(rng::AbstractRNG) = (dims...,) -> Base.randn(rng, Float32, dims...)

"""
    create_bias(weights, bias, size...)

Return a bias parameter for a layer, based on the value given
to the constructor's keyword `bias=bias`.

* `bias == true` creates a trainable array of the given size, of the same type as `weights`, initialised to zero.
* `bias == false` returns `false`, which is understood by AD to be non-differentiable.
* `bias::AbstractArray` uses the array provided, provided it has the correct size.
  It does not at present correct the `eltype` to match that of `weights`.
"""
function create_bias(weights::AbstractArray, bias::Bool, dims::Integer...)
  bias ? fill!(similar(weights, dims...), 0) : false
end
function create_bias(weights::AbstractArray, bias::AbstractArray, dims::Integer...)
  size(bias) == dims || throw(DimensionMismatch("expected bias of size $(dims), got size $(size(bias))"))
  bias
end

# Other

"""
    throttle(f, timeout; leading=true, trailing=false)

Return a function that when invoked, will only be triggered at most once
during `timeout` seconds.

Normally, the throttled function will run as much as it can, without ever
going more than once per `wait` duration; but if you'd like to disable the
execution on the leading edge, pass `leading=false`. To enable execution on
the trailing edge, pass `trailing=true`.
"""
function throttle(f, timeout; leading=true, trailing=false)
  cooldown = true
  later = nothing
  result = nothing

  function throttled(args...; kwargs...)
    yield()

    if cooldown
      if leading
        result = f(args...; kwargs...)
      else
        later = () -> f(args...; kwargs...)
      end

      cooldown = false
      @async try
        while (sleep(timeout); later != nothing)
          later()
          later = nothing
        end
      finally
        cooldown = true
      end
    elseif trailing
      later = () -> (result = f(args...; kwargs...))
    end

    return result
  end
end

"""
    modules(m)

Return an iterator over non-leaf objects
that can be reached by recursing `m` over
the children given by [`functor`](@ref).

Useful for applying a function (e.g. a regularizer)
over specific modules or subsets of the parameters
(e.g. the weights but not the biases).

# Examples

```jldoctest
julia> m1 = Chain(Dense(28^2, 64), BatchNorm(64, relu));

julia> m2 = Chain(m1, Dense(64, 10))
Chain(
  Chain(
    Dense(784 => 64),                   # 50_240 parameters
    BatchNorm(64, relu),                # 128 parameters, plus 128
  ),
  Dense(64 => 10),                      # 650 parameters
)         # Total: 6 trainable arrays, 51_018 parameters,
          # plus 2 non-trainable, 128 parameters, summarysize 200.312 KiB.

julia> Flux.modules(m2)
7-element Vector{Any}:
 Chain(Chain(Dense(784 => 64), BatchNorm(64, relu)), Dense(64 => 10))  # 51_018 parameters, plus 128 non-trainable
 (Chain(Dense(784 => 64), BatchNorm(64, relu)), Dense(64 => 10))
 Chain(Dense(784 => 64), BatchNorm(64, relu))  # 50_368 parameters, plus 128 non-trainable
 (Dense(784 => 64), BatchNorm(64, relu))
 Dense(784 => 64)    # 50_240 parameters
 BatchNorm(64, relu)  # 128 parameters, plus 128 non-trainable
 Dense(64 => 10)     # 650 parameters

julia> L2(m) = sum(sum(abs2, l.weight) for l in Flux.modules(m) if l isa Dense)
L2 (generic function with 1 method)

julia> L2(m2) isa Float32
true
```
"""
modules(m) = [x for x in Functors.fcollect(m) if !isleaflike(x)]

@nograd modules # TODO: is this correct? might fail with explicit parameters.
function ChainRulesCore.rrule(::typeof(modules), m)
  modules(m), dm -> error("Flux.modules is not at present differentiable, sorry")
end

isleaflike(x) = Functors.isleaf(x)
isleaflike(::Tuple{Vararg{<:Number}}) = true
isleaflike(::Tuple{Vararg{<:AbstractArray{<:Number}}}) = true

"""
    patience(predicate, wait)

Return a function that internally counts by one when
`predicate(...) == true`, otherwise the count is reset to zero.
If the count is greater than or equal to `wait`,
the function returns `true`, otherwise it returns `false`.

# Examples
```jldoctest
julia> loss() = rand();

julia> trigger = Flux.patience(() -> loss() < 1, 3);

julia> Flux.@epochs 10 begin
         trigger() && break
       end
[ Info: Epoch 1
[ Info: Epoch 2
[ Info: Epoch 3
```
"""
function patience(predicate, wait)
  let count = 0
    function on_trigger(args...; kwargs...)
      count = predicate(args...; kwargs...) ? count + 1 : 0

      return count >= wait
    end
  end
end

"""
    early_stopping(f, delay; distance = -, init_score = 0, min_dist = 0)

Return a function that internally counts by one when
`distance(best_score, f(...)) <= min_dist`, where
`best_score` is the last seen best value of `f(...)`.
If the count is greater than or equal to `delay`,
the function returns `true`, otherwise it returns `false`.
The count is reset when `distance(best_score, f(...)) > min_dist`.

# Examples
```jldoctest
julia> loss = let l = 0
         () -> l += 1
       end; # pseudo loss function that returns increasing values

julia> es = Flux.early_stopping(loss, 3);

julia> Flux.@epochs 10 begin
         es() && break
       end
[ Info: Epoch 1
[ Info: Epoch 2
[ Info: Epoch 3
```
"""
function early_stopping(f, delay; distance = -, init_score = 0, min_dist = 0)
  trigger = let best_score = init_score
    (args...; kwargs...) -> begin
      score = f(args...; kwargs...)
      Δ = distance(best_score, score)
      best_score = Δ < 0 ? best_score : score

      return Δ < min_dist
    end
  end

  return patience(trigger, delay)
end

"""
    plateau(f, width; distance = -, init_score = 0, min_dist = 1f-6)

Return a function that internally counts by one when
`abs(distance(last_score, f(...))) <= min_dist`, where
`last_score` holds the last value of `f(...)`.
If the count is greater than or equal to `width`,
the function returns `true`, otherwise it returns `false`.
The count is reset when `abs(distance(last_score, f(...))) > min_dist`.

# Examples
```jldoctest
julia> f = let v = 10
         () -> v = v / abs(v) - v
       end; # -9, 8, -7, 6, ...

julia> trigger = Flux.plateau(f, 3; init_score=10, min_dist=18);

julia> Flux.@epochs 10 begin
         trigger() && break
       end
[ Info: Epoch 1
[ Info: Epoch 2
[ Info: Epoch 3
[ Info: Epoch 4
```
"""
function plateau(f, width; distance = -, init_score = 0, min_dist = 1f-6)
  is_plateau = let last_score = init_score
    (args...; kwargs...) -> begin
      score = f(args...; kwargs...)
      Δ = abs(distance(last_score, score))
      last_score = score

      return Δ < min_dist
    end
  end

  return patience(is_plateau, width)
end


# -*- coding: utf-8 -*-
module CategoryTheory

using SemanticModels
import Catlab.Theories: dom, codom
import Base: append!, push!, deleteat!, delete!

export ⊔, AbstractMorph, FinSetMorph, dom, codom, verify, func, Decorated, decorations, undecorate, AbstractSpan, leftob, rightob, apexob, Span, left, right, DoublePushout, AbstractCospan, Cospan, pushout

⊔(a::UnitRange, b::UnitRange) = 1:(length(a)+length(b))
⊔(a::AbstractVector{Int}, b::AbstractVector{Int}) = vcat(a,b)

"""    AbstractMorph

an abstract type for representing morphisms. The essential API for subtypes of AbstractMorph are

1. dom(m)::T
2. codom(m)::T
3. func(m)::Function

where T is the type of the objects in the category. See FinSetMorph for an example.
"""
abstract type AbstractMorph end

"""    FinSetMorph{T,F}

morphisms in the category of Finite Sets. The objects are of type UnitRange{Int}.
func(m::FinSetMorph) is a function that takes `Int -> Int`. FinSetMorphs can be constructed
from a list of numbers. For example, `FinSetMorph([1,3,2,3])` is the morphism that takes
`1->1, 2->3, 3->2, 4->3` on domain `1:4` with codomain `1:3`. When you define a morphism from
a list of integers, the codomain is inferred from the largest element of the list. The domain
must always be the `1:l` where `l` is the length of the input list.

```(f::FinSetMorph)(g::G) where G <: AbstractGraph```

lift a finite set morphism (list of integers) to a graph homomorphism by its action on the vertex
set. The graph `h = f(g)` is defined by taking the edges of `g` and relabeling their src and dst
according to the function of `f`.

This method computes a valid graph homomorphism by definition.
"""
struct FinSetMorph{T,F} <: AbstractMorph
    codom::T
    fun::F
end

FinSetMorph(v::AbstractVector{Int}) = FinSetMorph(1:maximum(v), v)

dom(m::FinSetMorph) = 1:length(m.fun)
codom(m::FinSetMorph) = m.codom
func(m::FinSetMorph) = i->m.fun[i]

"""    ⊔(f::FinSetMorph, g::FinSetMorph)

the union of two morphisms in a finite set.
"""
function ⊔(f::FinSetMorph, g::FinSetMorph)
    Y = codom(f) ⊔ codom(g)
    h = f.fun ⊔ g.fun
    FinSetMorph(Y, h)
end

"""    Decorated{M,T}

a decoration applied to the objects of a morphism, where M is a type of morphism and
type T is the category of the decoration
"""
struct Decorated{M}
    f::M
    d::AbstractArray{AbstractModel}
end

# Handle creating a decorated morphism with an array of a single type
Decorated(f, d::AbstractArray{T}) where T<:AbstractModel = Decorated(f, Vector{AbstractModel}(d))
# Handle creating a decorated morphism from a single decoration
Decorated(f, d::T) where T<:AbstractModel = Decorated(f, Vector{AbstractModel}([d]))

# Get the domain or codomain of a decorated morphism
dom(m::Decorated) = dom(m.f)
codom(m::Decorated) = codom(m.f)

# Get the decorations of a decorated morphism
function decorations(m::Decorated)
  return m.d
end
# Get the decorations of AbstractModel T of a decorated morphism
function decorations(m::Decorated, ::Type{T}) where T<:AbstractModel
  filter(x -> isa(x,T), decorations(m))
end

# Remove the decoration of a decorated morphism, and return the original morphism
function undecorate(m::Decorated)
  return m.f
end

# Add a decoration to a decorated morphism
function push!(m::Decorated, decoration::AbstractModel)
  push!(decorations(m), decoration)
end

# Add a collection of decorations to a decorated morphism
function append!(m::Decorated, decorations)
    append!(m.d, decorations)
end

# remove a decoration from a decorated morphism
function deleteat!(m::Decorated, i)
  deleteat!(decorations(m), i)
end

# Remove the decorations of AbstractModel T from a decorated morphism
function delete!(m::Decorated, ::Type{T}) where T<:AbstractModel
  filter!(x -> !isa(x,T), decorations(m))
end

function left(d::Decorated)
  return left(d.f)
end

function right(d::Decorated)
  return right(d.f)
end

"""    AbstractSpan

an abstract type for representing spans. The essential API for subtypes of AbstractSpan are

1. left(s)::M
2. right(s)::M
3. pushout(s)::C

where M is the type of morphisms in the span, and C is the type of cospan that solves
a pushout of span s. See Span for an example.
"""
abstract type AbstractSpan end

leftob(s::AbstractSpan) = codom(left(s))

rightob(s::AbstractSpan) = codom(right(s))

function apexob(s::AbstractSpan)
    a = dom(left(s))
    b = dom(right(s))
    a == b || error("Inconsistent span")
    return a
end

"""    Span{F,G} <: AbstractSpan

a general span type where types F and G are types of morphisms in the span
"""
struct Span{F,G} <: AbstractSpan
    f::F
    g::G
end

function left(s::Span)
    return s.f
end

function right(s::Span)
    return s.g
end

"""    undecorate(s::Span{T,T}) where T <: Decorated

remove decorations of a span of decorated morphisms
"""
function undecorate(s::Span{T,T}) where T <: Decorated
    return Span(undecorate(left(s)), undecorate(right(s)))
end

struct DoublePushout{S<:AbstractSpan, T<:NTuple{3,AbstractMorph}}
    rule::S
    morphs::T
    application::S
end

# TODO DPO CONSTRUCTOR TO SOLVE UNKNOWN DOUBLEPUSHOUT OF FINSET
# take in Span `top` (l, c, r) and finset `l′ `, f: l -> l′
# Solve for c′ using dropdown
#     c′ = setdiff(l′, f(l)) ⊔ c
# Solve pushout for span of c′ ← c → r to get r′

# function pullback

# +
"""    AbstractCospan

an abstract type for representing cospans. The essential API for subtypes of AbstractCospan are

1. left(c)::M
2. right(c)::M
3. pullback(c)::S

where M is the type of morphisms in the cospan, and S is the type of span that solves the
pullback defined by c. See Cospan for an example.
"""
abstract type AbstractCospan end

leftob(c::AbstractCospan) = codom(left(c))

rightob(c::AbstractCospan) = codom(right(c))

function apexob(c::AbstractCospan)
    a = dom(left(c))
    b = dom(right(c))
    a == b || error("Inconsistent cospan")
    return a
end

"""    Cospan{F,G} <: AbstractCospan

a general cospan type where types F and G are types of morphisms in the cospan
"""
struct Cospan{F,G} <: AbstractCospan
    f::F
    g::G
end

function left(c::Cospan)
    return c.f
end

function right(c::Cospan)
    return c.g
end

"""    undecorate(s::Copan{T,T}) where T <: Decorated

remove decorations of a cospan of decorated morphisms
"""
function undecorate(c::Cospan{T,T}) where T <: Decorated
    return Cospan(undecorate(left(c)), undecorate(right(c)))
end

# TODO: This pushout gets the numbers mixed up. It is right up to permutation
# which is fine for when we want to do one pushout. But when we want to chain
# pushouts in order to do composition, "up to permutation" isn't good enough.
"""    pushout(s::Span{T,T}) where T <: FinSetMorph

treat f,g as a span and compute the pushout that is, the cospan of f=(f⊔g) and g=(a⊔b)
"""
function pushout(s::Span{T,U}) where {T <: FinSetMorph, U <: FinSetMorph}
    f_dict = Dict(a=>i for (i, a) in enumerate(left(s).fun))
    g′ = map(n->n in keys(f_dict) ? func(right(s))(f_dict[n]) : n+length(rightob(s)), leftob(s))

    g_dict = Dict(a=>i for (i, a) in enumerate(right(s).fun))
    f′ = map(n->n in keys(g_dict) ? g′[func(left(s))(g_dict[n])] : n, rightob(s))

    u = union(f′, g′)
    u_dict = Dict(a=>i for (i, a) in enumerate(u))
    f′ = FinSetMorph(1:length(u), map(n->u_dict[n], f′))
    g′ = FinSetMorph(1:length(u), map(n->u_dict[n], g′))

    return Cospan(f′, g′)
end

"""    pushout(s::Span{T, T}) where T <: Decorated

treat f,g as a decorated span and compute the pushout that is, the cospan of f=(f⊔g) and g=(a⊔b),
with the decoration of (f⊔g)(d)
"""
function pushout(s::Span{T, T}) where T <: Decorated
    cs = pushout(undecorate(s))
    decorations = map(x->x[1] ⊔ x[2], zip(left(s).d, right(s).d))
    return Decorated(cs, map((right(cs) ⊔ left(cs)), decorations))
end

end


"""
    2D: unit_normal(p1::T, p2::T) where {T<:AbstractVector}
    3D: unit_normal(p1::T, p2::T, p3::T) where {T<:AbstractVector}
Calculate unit normal vector
"""
function unit_normal(p1::T, p2::T) where {T<:AbstractVector}
    Δ = p2 .- p1
    l = norm(Δ) + 1e-6

    return [-Δ[2], Δ[1]] ./ l
end

function unit_normal(p1::T, p2::T, p3::T) where {T<:AbstractVector}
    v1 = p2 .- p1
    v2 = p3 .- p1

    n = cross(v1, v2)
    l = norm(n) + 1e-6

    return n ./ l
end


# Copyright 2020 Google LLC

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

function part1(filename="day15input.txt", nth=2020)
    starting_numbers = map(s->parse(Int, s), split(String(read(filename)), ","))
    last_said = Dict{Int, Int}()
    prev = nothing
    saying = nothing
    for i ∈ 1:nth
        if i <= length(starting_numbers)
            saying = starting_numbers[i]
        elseif prev ∈ keys(last_said)
            saying = i - last_said[prev] - 1
        else
            saying = 0
        end
        if prev !== nothing
            last_said[prev] = i-1
        end
        prev = saying
    end
    return saying
end

println("Part 1: ", part1())
part2() = part1("day15input.txt", 30000000)
println("Part 2: ", part2())


function farloi01()
#
# Farmer-Loizou
#
    p = conv([1.0;1;2],[1;1;3]);
    p = conv(p,p); p = conv(p,p);
    z = [-0.50000000000000 + 1.65831239517770im
        -0.50000000000000 - 1.65831239517770im
        -0.50000000000000 + 1.32287565553230im
        -0.50000000000000 - 1.32287565553230im];
    z = [z 4*ones(4)];

    p, PolyZeros(z)
end


function db(t,n,k)
    out = binomial(n,k).*(k.*t.^(k.-1).*(1 .-t).^(n-k) .- t.^k.*(n-k).*(1 .-t).^(n-k-1));
    return out
end


# module for representing motion edge

using LinearAlgebra

include(joinpath(split(@__FILE__, "src")[1], "src/common/state_transition/state_transition.jl"))

mutable struct MotionEdge
  t1
  t2
  x1
  x2
  omega_upper_right
  omega_upper_left
  omega_lower_right
  omega_lower_left
  xi_upper
  xi_lower

  # matrix of input noise
  function mat_M(speed, yaw_rate, time, stds)
    return diagm(0 => [stds["nn"]^2*abs(speed)/time + stds["no"]^2*abs(yaw_rate)/time,
                       stds["on"]^2*abs(speed)/time + stds["oo"]^2*abs(yaw_rate)/time])
  end

  # jacobian of input
  function mat_A(speed, yaw_rate, time, theta)
    st, ct = sin(theta), cos(theta)
    stw, ctw = sin(theta + yaw_rate * time), cos(theta + yaw_rate * time)
    return [(stw - st)/yaw_rate  -speed/(yaw_rate^2)*(stw - st) + speed/yaw_rate*time*ctw;
            (-ctw + ct)/yaw_rate -speed/(yaw_rate^2)*(-ctw + ct) + speed/yaw_rate*time*stw;
            0                    time]
  end

  # jacobian of pose
  function mat_F(speed, yaw_rate, time, theta)
    F = diagm(0 => [1.0, 1.0, 1.0])
    F[1, 3] = speed / yaw_rate * (cos(theta + yaw_rate * time) - cos(theta))
    F[2, 3] = speed / yaw_rate * (sin(theta + yaw_rate * time) - sin(theta))
    return F
  end

  # init
  function MotionEdge(t1, t2, pose_list, input_list, delta;
                      motion_noise_stds=Dict("nn"=>0.20, "no"=>0.001, "on"=>0.11, "oo"=>0.20))
    self = new()

    self.t1 = t1 # time
    self.t2 = t2
    self.x1 = pose_list[string(t1)] # pose [x, y, theta]
    self.x2 = pose_list[string(t2)]
    speed = input_list[string(t2)][1] # input [speed, yaw_rate]
    yaw_rate = input_list[string(t2)][2]
    if abs(yaw_rate) < 1e-5
      yaw_rate = 1e-5
    end

    # noise matrix and jacobian
    M = mat_M(speed, yaw_rate, delta, motion_noise_stds)
    A = mat_A(speed, yaw_rate, delta, self.x1[3])
    F = mat_F(speed, yaw_rate, delta, self.x1[3])

    # precision matrix of edge
    omega_edge = inv(A*M*A' + Matrix{Float64}(I,3,3).*0.0001)

    # precision matrix of graph
    self.omega_upper_left = F' * omega_edge * F
    self.omega_upper_right = -F' * omega_edge
    self.omega_lower_left = -omega_edge * F
    self.omega_lower_right = omega_edge

    # coefficient vector
    x2 = state_transition(speed, yaw_rate, delta, self.x1)
    self.xi_upper = F' * omega_edge * (self.x2 - x2)
    self.xi_lower = -omega_edge * (self.x2 - x2)

    return self
  end
end


using Flux
using Flux: mse, throttle
using Base.Iterators: repeated

# input
x = [0 0 1 1
     0 1 0 1
     1 1 1 1]
# output
y = [0 0 1 1]

# model
m = Dense(3,1,σ)    # only one layer (excl. input layer) with 3 inputs, 1 output and a sigmoid function

# loss function
loss(x,y) = mse(m(x), y)           # calculates the mean squared error

dataset = repeated((x, y), 10000)  # Train the model 10000 times with the same dataset
opt = SGD(params(m), 0.25)         # Gradient descent with learning rate 0.25
evalcb = () -> @show(loss(x, y))   # output the progress

Flux.train!(loss, dataset, opt, cb = throttle(evalcb, 0.5))

# output final error
@show(loss(x, y))

# testing the results
@show m([1,0,0])    # should be close to 1
@show m([0,1,0])    # should be close to 0


"""
**FUNÇÃO VOLUME CILINDRO**
Sintaxe:
volcilindro(raio, altura_cilindro)
"""
function volcilindro(raio, altura)
    area_base = areaBC(raio)
    println("Volume do Cilindro")
    return  area_base.*altura
end


const width = height = 1000.0
const trunklength = 400.0
const scalefactor = 0.6
const startingangle = 1.5 * pi
const deltaangle = 0.2 * pi

function tree(fh, x, y, len, theta)
   if len >= 1.0
       x2 = x + len * cos(theta)
       y2 = y + len * sin(theta)
       write(fh, "<line x1='$x' y1='$y' x2='$x2' y2='$y2' style='stroke:rgb(0,0,0);stroke-width:1'/>\n")
       tree(fh, x2, y2, len * scalefactor, theta + deltaangle)
       tree(fh, x2, y2, len * scalefactor, theta - deltaangle)
    end
end

outsvg = open("tree.svg", "w")
write(outsvg,
    """<?xml version='1.0' encoding='utf-8' standalone='no'?>
    <!DOCTYPE svg PUBLIC '-//W3C//DTD SVG 1.1//EN'
    'http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd'>
    <svg width='100%%' height='100%%' version='1.1'
    xmlns='http://www.w3.org/2000/svg'>\n""")

tree(outsvg, 0.5 * width, height, trunklength, startingangle)

write(outsvg, "</svg>\n") # view file tree.svg in browser


#Moore et al 2009 mean vectors:
M=Array{Any,1}(undef,8)
M[1]=vec([0.0234 0.0192 0.0129 0.0075 0.0031 0.0002])
M[2]=vec([0.0162 0.0141 0.0112 0.0073 0.0034 0.0002])
M[3]=vec([0.0107 0.0098 0.0092 0.0070 0.0039 0.0003])
M[4]=vec([0.0065 0.0064 0.0070 0.0064 0.0048 0.0006])
M[5]=vec([0.0033 0.0034 0.0042 0.0042 0.0043 0.0009])
M[6]=vec([0.0064 0.0074 0.0105 0.0116 0.0140 0.0041])
M[7]=vec([0.0121 0.0140 0.0192 0.0204 0.0231 0.0084])
M[8]=vec([0.0184 0.0230 0.0333 0.0359 0.0409 0.0137])

#Moore et al 2009 covariance matrices:
S=Array{Any,1}(undef,8)
S[1]=[0.00000959 0.00000556 0.00000138 -0.00000034 -0.00000024 -0.00000003;
        0.00000556 0.00000493 0.00000193 0.00000060 0.00000023 0.00000001;
        0.00000138 0.00000193 0.00000282 0.00000223 0.00000119 0.00000007;
        -0.00000034 0.00000060 0.00000223 0.00000232 0.00000119 0.00000007;
        -0.00000024 0.00000023 0.00000119 0.00000119 0.00000071 0.00000005;
        -0.00000003 0.00000001 0.00000007 0.00000007 0.00000005 0.00000001]
S[2]=[0.00000346 0.00000186 -0.00000011 -0.00000060 -0.00000062 -0.00000005;
        0.00000186 0.00000228 0.00000086 0.00000033 -0.00000007 0.00000003;
        -0.00000011 0.00000086 0.00000231 0.00000221 0.00000145 0.00000023;
        -0.00000060 0.00000033 0.00000221 0.00000266 0.00000191 0.00000034;
        -0.00000062 -0.00000007 0.00000145 0.00000191 0.00000175 0.00000041;
        -0.00000005 0.00000003 0.00000023 0.00000034 0.00000041 0.00000021]
S[3]=[0.00000241 0.00000144 0.00000035 -0.00000031 -0.00000063 -0.00000006;
        0.00000144 0.00000138 0.00000076 0.00000015 -0.00000021 -0.00000001;
        0.00000035 0.00000076 0.00000161 0.00000156 0.00000121 0.00000016;
        -0.00000031 0.00000015 0.00000156 0.00000227 0.00000209 0.00000031;
        -0.00000063 -0.00000021 0.00000121 0.00000209 0.00000225 0.00000037;
        -0.00000006 -0.00000001 0.00000016 0.00000031 0.00000037 0.00000013]
S[4]=[0.00000166 0.00000091 0.00000034 -0.00000009 -0.00000080 -0.00000025;
        0.00000091 0.00000097 0.00000071 0.00000025 -0.00000041 -0.00000015;
        0.00000034 0.00000071 0.00000118 0.00000103 0.00000072 0.00000003;
        -0.00000009 0.00000025 0.00000103 0.00000137 0.00000162 0.00000025;
        -0.00000080 -0.00000041 0.00000072 0.00000162 0.00000290 0.00000065;
        -0.00000025 -0.00000015 0.00000003 0.00000025 0.00000065 0.00000050]
S[5]=[0.00000178 0.00000132 0.00000104 0.00000081 0.00000018 -0.00000014;
        0.00000132 0.00000127 0.00000121 0.00000099 0.00000034 -0.00000010;
        0.00000104 0.00000121 0.00000150 0.00000142 0.00000110 0.00000013;
        0.00000081 0.00000099 0.00000142 0.00000158 0.00000177 0.00000042;
        0.00000018 0.00000034 0.00000110 0.00000177 0.00000351 0.00000131;
        -0.00000014 -0.00000010 0.00000013 0.00000042 0.00000131 0.00000081]
S[6]=[0.00000715 0.00000586 0.00000409 0.00000292 0.00000005 -0.00000075;
        0.00000586 0.00000589 0.00000520 0.00000398 0.00000027 -0.00000114;
        0.00000409 0.00000520 0.00000634 0.00000541 0.00000188 -0.00000097;
        0.00000292 0.00000398 0.00000541 0.00000528 0.00000392 0.00000070;
        0.00000005 0.00000027 0.00000188 0.00000392 0.00000995 0.00000657;
        -0.00000075 -0.00000114 -0.00000097 0.00000070 0.00000657 0.00000819]
S[7]=[0.00002625 0.00001981 0.00001058 0.00000544 -0.00000654 -0.00001122;
        0.00001981 0.00001745 0.00001314 0.00000822 -0.00000431 -0.00001228;
        0.00001058 0.00001314 0.00001629 0.00001226 0.00000035 -0.00001311;
        0.00000544 0.00000822 0.00001226 0.00001170 0.00000742 -0.00000500;
        -0.00000654 -0.00000431 0.00000035 0.00000742 0.00002241 0.00001782;
        -0.00001122 -0.00001228 -0.00001311 -0.00000500 0.00001782 0.00003987]
S[8]=[0.00001186 0.00001134 0.00001139 0.00000919 0.00000395 -0.00000186;
        0.00001134 0.00001484 0.00002034 0.00001907 0.00001531 0.00000087;
        0.00001139 0.00002034 0.00003467 0.00003546 0.00003555 0.00000708;
        0.00000919 0.00001907 0.00003546 0.00003907 0.00004604 0.00001733;
        0.00000395 0.00001531 0.00003555 0.00004604 0.00007306 0.00004953;
        -0.00000186 0.00000087 0.00000708 0.00001733 0.00004953 0.00006542];


"""
    peakproms(peaks, x;
        strict=true,
        minprom=nothing,
        maxprom=nothing
    ) -> (peaks, proms)

Calculate the prominences of `peaks` in `x`, filtering peaks with prominences less than
`minprom` and greater than `maxprom`, if either are given.

Peak prominence is the absolute height difference between the current peak and the larger of
the two adjacent smallest magnitude points between the current peak and adjacent larger
peaks or signal ends.

The prominence for a peak with a `NaN` or `missing` between the current peak and either
adjacent larger peaks will be `NaN` or `missing` if `strict == true`, or it will be
the larger of the smallest non-`NaN` or `missing` values between the current peak and
adjacent larger peaks for `strict == false`.

See also: [`findminima`](@ref), [`findmaxima`](@ref)

# Examples
```jldoctest
julia> x = [0,5,2,3,3,1,4,0];

julia> xpks = argmaxima(x)
3-element Vector{Int64}:
 2
 4
 7

julia> peakproms(xpks, x)
([2, 4, 7], Union{Missing, Int64}[5, 1, 3])

julia> x = [missing,5,2,3,3,1,4,0];

julia> peakproms(xpks, x)
([2, 4, 7], Union{Missing, Int64}[missing, 1, 3])

julia> peakproms(xpks, x; strict=false)
([2, 4, 7], Union{Missing, Int64}[5, 1, 3])
```
"""
function peakproms(peaks::AbstractVector{Int}, x::AbstractVector{T};
    strict=true, minprom=nothing, maxprom=nothing, strictbounds=nothing
) where T
    if !isnothing(strictbounds)
        Base.depwarn("Keyword `strictbounds` has been renamed to `strict`", :peakproms!)
        strict=strictbounds
    end

    if !isnothing(minprom) || !isnothing(maxprom)
        _peaks = copy(peaks)
    else
        # peaks will not be modified
        _peaks = peaks
    end
    return peakproms!(_peaks, x; strict=strict, minprom=minprom, maxprom=maxprom)
end

"""
    peakproms!(peaks, x;
        strict=true,
        minprom=nothing,
        maxprom=nothing
    ) -> (peaks, proms)

Calculate the prominences of `peaks` in `x`, removing `peaks` with prominences less than
`minprom` or greater than `maxprom`, if either are given. Returns the modified peaks and
their prominences.

See also: [`peakproms`](@ref), [`findminima`](@ref), [`findmaxima`](@ref)
"""
function peakproms!(peaks::AbstractVector{Int}, x::AbstractVector{T};
    strict=true, minprom=nothing, maxprom=nothing
) where T
    if !isnothing(minprom) && !isnothing(maxprom)
        minprom < maxprom || throw(ArgumentError("minprom must be less than maxprom"))
    end
    all(∈(eachindex(x)), peaks) ||
        throw(ArgumentError("peaks contains invalid indices to x"))
    isempty(peaks) && return peaks, T[]

    # if peaks was calculated with strict=false, first(peaks) could be minima at firstindex
    fp = length(peaks) > 1 ? peaks[2] : first(peaks)
    if fp > 1 && ((x[fp] < x[fp-1]) === true)
        pktype = :minima
    else
        pktype = :maxima
    end
    cmp = pktype === :maxima ? (≥) : (≤)
    exm = pktype === :maxima ? minimum : maximum
    exa = pktype === :maxima ? max : min

    _ref = Missing <: T ? missing :
           Float64 <: T ? NaN :
           Float32 <: T ? NaN32 :
           Float16 <: T ? NaN16 :
                          missing

    proms = similar(peaks,promote_type(T,typeof(_ref)))

    if strict
        lbegin, lend = firstindex(x), lastindex(x)

        @inbounds for i in eachindex(peaks, proms)
            # Find left and right bound (self-intersections)
            lb = something(findprev(y -> cmp(y, x[peaks[i]]) === true, x, peaks[i] - 2),
                lbegin)
            rb = something(findnext(y -> cmp(y, x[peaks[i]]) === true, x, peaks[i] + 2),
                lend)

            # Find extremum of left and right bounds
            if isempty(lb:(peaks[i] - 1))
                lref = _ref
            else
                lref = exm(view(x, lb:(peaks[i] - 1)))
            end

            if isempty((peaks[i] + 1):rb)
                rref = _ref
            else
                rref = exm(view(x, (peaks[i] + 1):rb))
            end

            proms[i] = abs(x[peaks[i]] - exa(lref, rref))
        end
    else
        # The extremum search space in the bounding intervals can be reduced by
        # restricting the search space to known peaks/reverse peaks. The cost of
        # finding all peaks/reverse peaks should be mitigated by the fact that
        # the same peaks/reverse peaks will be the pivotal elements for
        # numerous peaks.
        if pktype === :maxima
            peaks′ = argmaxima(x, 1; strict=false)
            notm = argminima(x, 1; strict=false)
        else
            peaks′ = argminima(x, 1; strict=false)
            notm = argmaxima(x, 1; strict=false)
        end

        notmval = x[notm]

        for i in eachindex(peaks, proms)
            j = searchsorted(peaks′, peaks[i])

            # Find left and right bounding peaks
            _lb = findprev(y -> cmp(x[y], x[peaks[i]]) === true, peaks′, first(j)-1)
            peaks′[j] === peaks[i] && (j += 1)
            _rb = findnext(y -> cmp(x[y], x[peaks[i]]) === true, peaks′, last(j)+1)

            # Find left and right reverse peaks just inside the bounding peaks
            lb = isnothing(_lb) ? firstindex(notm) :
                                  searchsortedfirst(notm, peaks′[_lb])
            rb = isnothing(_rb) ? lastindex(notm) :
                                  searchsortedlast(notm, peaks′[_rb])

            k = searchsortedfirst(notm, peaks[i])

            if isempty(lb:(k-1))
                lref = missing
            else
                lref = exm(view(notmval, lb:(k-1)))
            end

            rb > lastindex(notm) && (rb -= 1)
            if isempty(k:rb)
                rref = missing
            else
                rref = exm(view(notmval, k:rb)) # Slice corollary upper side
            end

            proms[i] = abs(x[peaks[i]] - exa(coalesce(lref, rref), coalesce(rref, lref)))
        end
    end

    if !isnothing(minprom) || !isnothing(maxprom)
        lo = something(minprom, zero(eltype(x)))
        up = something(maxprom, typemax(Base.nonmissingtype(eltype(x))))
        matched = findall(x -> !ismissing(x) && !(lo ≤ x ≤ up), proms)
        deleteat!(peaks, matched)
        deleteat!(proms, matched)
    end

    return peaks, proms
end

struct Maxima; end
struct Minima; end

@deprecate peakprom(x::AbstractVector, w::Int=1; strictbounds=true, minprom=nothing) peakproms(argmaxima(x, w; strict=strictbounds), x; strict=strictbounds, minprom=minprom)
@deprecate peakprom(m::Minima, x::AbstractVector, w::Int=1; strictbounds=true, minprom=nothing) peakproms(argminima(x, w; strict=strictbounds), x; strict=strictbounds, minprom=minprom)
@deprecate peakprom(m::Maxima, x::AbstractVector, w::Int=1; strictbounds=true, minprom=nothing) peakproms(argmaxima(x, w; strict=strictbounds), x; strict=strictbounds, minprom=minprom)


"""
    surrplot(ts, surrogate; n_windows_periodogram = 100, gfs = 8, lfs = 6)

Plot a time series along with its surrogate realization, and compares the
periodogram and histogram of the two series.

## Keywords

- **`ts`**: the time series for which to generate a surrogate.

- **`surrogate`** is the surrogate realization of the time series.

- **`n_windows_periodogram`**: sets the number of windows for binning the
    periodogram.

- **`gfs`** and **`lfs`**: fontsizes of the guides and legends, respectively.
"""
function surrplot(ts, surrogate; n_windows_periodogram = 100, gfs = 8, lfs = 6)
    p1 = Plots.plot(ts, label = ""); asurr = surrogate
    Plots.plot!(p1, surrogate, label = "")

    p2 = Plots.plot(autocor(ts), label = "")
    Plots.plot!(p2, autocor(surrogate), label = "")
    p, psurr = DSP.mt_pgram(ts), DSP.mt_pgram(surrogate)

    p3 = Plots.plot(interp([x for x in p.freq], p.power, n_windows_periodogram), label = "")
    Plots.plot!(p3, interp([x for x in psurr.freq], psurr.power, n_windows_periodogram), label = "")

    p4 = Plots.histogram(ts, label = "Original")
    Plots.histogram!(p4, surrogate, label = "Surrogate")

    Plots.xlabel!(p1, "Time step")
    Plots.ylabel!(p1, "Value")

    Plots.xlabel!(p2, "Lag")
    Plots.ylabel!(p2, "Autocorrelation")

    Plots.xlabel!(p3, "Binned frequency")
    Plots.ylabel!(p3, "Power")

    Plots.xlabel!(p4, "Binned value")
    Plots.ylabel!(p4, "Frequency")

    l = Plots.@layout [a{0.3h}; b{0.25h}; c{0.25h}; d{0.2h}]
    Plots.plot(p1, p2, p3, p4,
                layout = l,
                guidefont = (gfs, gfs, gfs, gfs),
                legendfont = (lfs, lfs, lfs, lfs)
                )
end

export surrplot


using DitherPunk
using DitherPunk: gradient_image, test_on_gradient
using Images

# # Gallery
# ## On color spaces
# A simple linear gradient works well to reveal the characteristic patterns of
# different dithering algorithms.
srbg, linear = gradient_image(100, 800);
mosaicview(srbg, linear)
# The pixel intensities in the image `srgb` increase linearly from 0 to 1.
# The second image `linear` has been converted from sRGB to a linear representation,
# which more closely matches our human perception of brightness.
#
# The helper function `test_on_gradient` takes a dithering algorithm and runs it on
# both the `srgb` and the `linear` image.
# It then shows a comparison of both inputs and outputs.
#
# Most dithering algorithms in DitherPunk.jl provide an optional parameter `to_linear`,
# which converts the input image to a linear color space before applying the dithering.
# Select what looks best!
#
# ## Threshold dithering
# #### `ConstantThreshold`
test_on_gradient(ConstantThreshold())

# #### `WhiteNoiseThreshold`
test_on_gradient(WhiteNoiseThreshold())

# ## Ordered dithering
# ### Bayer matrices
# #### `bayer_dithering`
test_on_gradient(Bayer())

# The order of the Bayer-matrix can be specified through the parameter `level`,
# which defaults to `1`.
# **Level 2**
test_on_gradient(Bayer(; level=2))
# **Level 3**
test_on_gradient(Bayer(; level=3))
# **Level 4**
test_on_gradient(Bayer(; level=4))

# ### Clustered / halftone dithering
# The following methods have large characteristic patterns and are therefore
# better suited for large images.
#
# #### `ClusteredDots`
test_on_gradient(ClusteredDots())

# #### `CentralWhitePoint`
test_on_gradient(CentralWhitePoint())

# #### `BalancedCenteredPoint`
test_on_gradient(BalancedCenteredPoint())

# #### `Rhombus`
test_on_gradient(Rhombus())

# #### ImageMagick methods
test_on_gradient(IM_checks())
#
test_on_gradient(IM_h4x4a())
#
test_on_gradient(IM_h6x6a())
#
test_on_gradient(IM_h8x8a())
#
test_on_gradient(IM_h4x4o())
#
test_on_gradient(IM_h6x6o())
#
test_on_gradient(IM_h8x8o())
#
test_on_gradient(IM_c5x5())
#
test_on_gradient(IM_c6x6())
#
test_on_gradient(IM_c7x7())

# ## Error diffusion
# #### `SimpleErrorDiffusion`
test_on_gradient(SimpleErrorDiffusion())

# #### `FloydSteinberg`
test_on_gradient(FloydSteinberg())

# #### `JarvisJudice`
test_on_gradient(JarvisJudice())

# #### `Stucki`
test_on_gradient(Stucki())

# #### `Burkes`
test_on_gradient(Burkes())

# #### `Sierra`
test_on_gradient(Sierra())

# #### `TwoRowSierra`
test_on_gradient(TwoRowSierra())

# #### `SierraLite`
test_on_gradient(SierraLite())

# #### `Fan93()`
test_on_gradient(Fan93())

# #### `ShiauFan()`
test_on_gradient(ShiauFan())

# #### `ShiauFan2()`
test_on_gradient(ShiauFan2())

# #### `Atkinson()`
test_on_gradient(Atkinson())