using Compat
using Distances
using Ripserer
using SparseArrays
using Random
using StaticArrays
using Suppressor
using Test

using Ripserer:
    distances, births, adjacency_matrix, edges, nv, unsafe_simplex, ChainElement, Chain

include("../testdatasets.jl")
include("interfacetest.jl")

@testset "distances" begin
    for points in (
        [(0, 0), (0, 1), (1, 1), (1, 0)],
        [SVector(0, 0), SVector(0, 1), SVector(1, 1), SVector(1, 0)],
    )
        @test distances(points) ≈ [0 1 √2 1; 1 0 1 √2; √2 1 0 1; 1 √2 1 0]
        @test distances(points, Cityblock()) == [0 1 2 1; 1 0 1 2; 2 1 0 1; 1 2 1 0]
    end
end

@testset "Rips points constructor, sparse=false" begin
    filtration = Rips([(sin(x), cos(x)) for x in range(0, 2π; length=101)[1:(end - 1)]])
    adj = adjacency_matrix(filtration)
    @test all(x -> x > 0, adj[i, j] for i in 1:100 for j in (i + 1):100)
    @test eltype(edges(filtration)) === Simplex{1,Float64,Int}

    filtration = Rips{Int32}([
        (sin(x), cos(x)) for x in range(0.0f0, 2.0f0π; length=101)[1:(end - 1)]
    ])
    adj = adjacency_matrix(filtration)
    @test !issparse(adj)
    @test all(x -> x > 0, adj[i, j] for i in 1:100 for j in (i + 1):100)
    @test eltype(edges(filtration)) === Simplex{1,Float32,Int32}
end

@testset "Rips points constructor, sparse=true" begin
    filtration = Rips(
        [(sin(x), cos(x)) for x in range(0, 2π; length=101)[1:(end - 1)]];
        threshold=0.1,
        sparse=true,
    )
    adj = adjacency_matrix(filtration)
    @test issparse(adj)
    @test maximum(adj) ≤ 0.1
    @test maximum(adj) ≤ threshold(filtration)
end

@testset "Construction does not alter input" begin
    dist = [
        0 9 1 2
        9 0 3 4
        1 3 0 4
        2 4 4 0
    ]
    orig_dist = copy(dist)
    Rips(dist; threshold=1)
    @test dist == orig_dist
    Rips(dist; threshold=1, sparse=true)
    @test dist == orig_dist

    dist = sparse(dist)
    orig_dist = copy(dist)
    Rips(dist; threshold=1)
    @test dist == orig_dist
end

@testset "Warns with duplicate points" begin
    pts = [(1, 0), (1, 1), (0, 1), (0, 0), (0, 0)]
    @test @capture_err(Rips(pts)) ≠ ""
    @test @capture_err(Rips(pts[1:4])) == ""
    @test nv(@suppress Rips(pts)) == 4
end

@testset "Errors" begin
    @testset "Non-square matrices throw an error" begin
        @test_throws DimensionMismatch Rips(zeros(3, 2))
        @test_throws DimensionMismatch Rips(zeros(3, 2); sparse=true)
    end
    @testset "Asymmetric matrices throw an error" begin
        @test_throws ArgumentError Rips([1 1 1; 1 1 1; 1 2 1])
        @test_throws ArgumentError Rips([1 1 1; 2 1 1; 1 2 1]; sparse=true)
    end
    @testset "Edge births must be larger than vertex births" begin
        @test_throws ArgumentError Rips([1 1 1; 1 1 1; 1 1 2])
        @test_throws ArgumentError Rips([1 1 1; 1 2 1; 1 1 1]; sparse=true)
    end
end

@testset "ripserer" begin
    @testset "Dense" begin
        @testset "Icosahedron" begin
            d0, d1, d2 = ripserer(icosahedron; dim_max=2)
            @test d0 == [fill((0.0, 1.0), 11); (0.0, Inf)]
            @test d1 == []
            @test d2 == [(1.0, 2.0)]
        end
        @testset "Cycle with various fields" begin
            d0_2, d1_2, d2_2, d3_2 = ripserer(Rips{Int32}, cycle; dim_max=3)
            d0_7, d1_7, d2_7, d3_7 = ripserer(Rips(cycle); dim_max=3, modulus=7)
            d0_r, d1_r, d2_r, d3_r = ripserer(cycle; dim_max=3, field=Rational{Int})

            @test d0_2 == d0_7 == d0_r == [fill((0, 1), size(cycle, 1) - 1); (0, Inf)]
            @test d1_2 == d1_7 == d1_r == [(1, 6)]
            @test d2_2 == d2_7 == d2_r == fill((6, 7), 5)
            @test d3_2 == d3_7 == d3_r == [(7, 8)]
        end
        @testset "RP2 with various fields" begin
            _, d1_2, d2_2 = ripserer(Rips(projective_plane); dim_max=2)
            _, d1_3, d2_3 = ripserer(Rips, projective_plane; dim_max=2, modulus=3)
            _, d1_331, d2_331 = ripserer(projective_plane; dim_max=2, field=Mod{5})
            _, d1_r, d2_r = ripserer(projective_plane; dim_max=2, field=Rational{Int})
            @test d1_2 == [(1, 2)]
            @test d2_2 == [(1, 2)]
            @test d1_3 == d1_331 == d1_r == []
            @test d2_3 == d2_331 == d1_r == []
        end
        @testset "Icosahedron, threshold=1" begin
            d0, d1, d2 = ripserer(Rips(icosahedron; threshold=1); dim_max=2)
            @test d0 == [fill((0.0, 1.0), 11); (0.0, Inf)]
            @test d1 == []
            @test d2 == [(1.0, Inf)]
        end
        @testset "Icosahedron, threshold=0.5" begin
            d0, d1, d2 = ripserer(icosahedron; dim_max=2, threshold=0.5)
            @test d0 == fill((0.0, Inf), 12)
            @test d1 == []
            @test d2 == []
        end
        @testset "RP2 with various fields, threshold=1" begin
            _, d1_2, d2_2 = ripserer(Rips, projective_plane; dim_max=2, threshold=1)
            _, d1_3, d2_3 = ripserer(projective_plane; dim_max=2, modulus=3, threshold=1)
            _, d1_331, d2_331 = ripserer(
                projective_plane; dim_max=2, field=Mod{5}, threshold=1
            )
            _, d1_r, d2_r = ripserer(
                projective_plane; dim_max=2, field=Rational{Int}, threshold=1
            )
            @test d1_2 == [(1, Inf)]
            @test d2_2 == [(1, Inf)]
            @test d1_3 == d1_331 == d1_r == []
            @test d2_3 == d2_331 == d1_r == []
        end
        @testset "Points as input" begin
            for metric in (Euclidean(), Cityblock())
                pts = torus_points(9)
                @test ripserer(pts; metric=metric) == ripserer(Rips(pts; metric=metric))
            end
        end
        @testset "Cutoff" begin
            d0, d1 = ripserer(rand_dist_matrix(20); cutoff=0.5)
            @test all(persistence.(d0) .> 0.5)
            @test all(persistence.(d1) .> 0.5)
        end
    end

    @testset "Sparse" begin
        @testset "Icosahedron" begin
            d0, d1, d2 = ripserer(sparse(icosahedron); dim_max=2)
            @test d0 == [fill((0.0, 1.0), 11); (0.0, Inf)]
            @test d1 == []
            @test d2 == [(1.0, 2.0)]
        end
        @testset "RP2 with various fields, threshold=1" begin
            _, d1_2, d2_2 = ripserer(sparse(projective_plane); dim_max=2, threshold=1)
            _, d1_3, d2_3 = ripserer(
                Rips(projective_plane; threshold=1, sparse=true); dim_max=2, modulus=3
            )
            _, d1_331, d2_331 = ripserer(
                Rips, sparse(projective_plane); dim_max=2, field=Mod{5}, threshold=1
            )
            _, d1_r, d2_r = ripserer(
                Rips,
                projective_plane;
                sparse=true,
                dim_max=2,
                field=Rational{Int},
                threshold=1,
            )
            @test d1_2 == [(1, Inf)]
            @test d2_2 == [(1, Inf)]
            @test d1_3 == d1_331 == d1_r == []
            @test d2_3 == d2_331 == d1_r == []
        end
        @testset "Equal to Rips" begin
            for thresh in (nothing, 1, 0.5, 0.126)
                data = torus_points(100)
                r_res = ripserer(data; threshold=thresh, dim_max=2)
                s_res_1 = ripserer(Rips, data; threshold=thresh, sparse=true, dim_max=2)

                # Add zeros to diagonal. Adding ones first actually changes the structure of
                # the matrix.
                data2 = sparse(Ripserer.distances(data))
                for i in axes(data2, 1)
                    data2[i, i] = 1
                    data2[i, i] = 0
                end
                s_res_2 = ripserer(data2; threshold=thresh, dim_max=2)

                @test r_res == s_res_1 == s_res_2
            end
        end
    end

    @testset "Representatives" begin
        @testset "Known example" begin
            # This example was generated by getting representatives from ripser.
            _, d1, d2 = ripserer(projective_plane; dim_max=2, reps=true)

            @test simplex.(representative(only(d1))) == [
                Simplex{1}((11, 10), 1),
                Simplex{1}((10, 7), 1),
                Simplex{1}((10, 6), 1),
                Simplex{1}((8, 1), 1),
                Simplex{1}((7, 3), 1),
                Simplex{1}((7, 1), 1),
                Simplex{1}((6, 2), 1),
                Simplex{1}((5, 1), 1),
                Simplex{1}((2, 1), 1),
            ]
            @test coefficient.(representative(only(d1))) == fill(Mod{2}(1), 9)
            @test simplex.(representative(only(d2))) == [Simplex{2}((6, 2, 1), 1)]
            @test coefficient.(representative(only(d2))) == [Mod{2}(1)]
        end
        @testset "Types" begin
            d0, d1, d2, d3 = ripserer(cycle; dim_max=3, reps=true)
            @test d1[1].representative isa Chain

            d0, d1, d2, d3 = ripserer(cycle; dim_max=3, reps=true, field=Rational{Int})
            @test d3[1].representative isa Chain
        end
        @testset "Infinite interval very low threshold" begin
            _, d1 = ripserer(cycle; dim_max=1, reps=true, threshold=1, field=Rational{Int})
            rep = representative(only(d1))
            @test simplex(only(rep)) == birth_simplex(only(d1))
            @test rep isa Chain{Rational{Int},Simplex{1,Int,Int}}
        end
        @testset "Infinite interval higher threshold" begin
            _, d1 = ripserer(cycle; dim_max=1, reps=true, threshold=3, field=Rational{Int})
            rep = representative(only(d1))
            @test !isempty(rep)
            @test rep isa Chain{Rational{Int},Simplex{1,Int,Int}}
        end
        @testset "Critical simplices" begin
            result = ripserer(torus(100); reps=true, threshold=0.5)
            for diag in result
                @test birth.(diag) == birth.(birth_simplex.(diag))
                finite = filter(isfinite, diag)
                @test death.(finite) == birth.(death_simplex.(finite))
                infinite = filter(!isfinite, diag)
                @test all(isnothing, death_simplex.(infinite))
            end
        end
    end

    @testset "Diagram metadata" begin
        filtration = Rips(cycle)
        d0, d1, d2, d3 = ripserer(filtration; dim_max=3, reps=true, field=Rational{Int})
        @test d0.dim == 0
        @test d1.dim == 1
        @test d2.dim == 2
        @test d3.dim == 3
        thresh = Float64(threshold(Rips(cycle)))
        @test d0.threshold ≡ thresh
        @test d1.threshold ≡ thresh
        @test d2.threshold ≡ thresh
        @test d3.threshold ≡ thresh
        field = Rational{Int}
        @test d0.field ≡ field
        @test d1.field ≡ field
        @test d2.field ≡ field
        @test d3.field ≡ field
        @test d0.filtration == filtration
        @test d1.filtration == filtration
        @test d2.filtration == filtration
        @test d3.filtration == filtration
    end

    @testset "Zero-dimensional sublevel set persistence" begin
        @testset "with sparse matrix" begin
            data = [1, 0, 1, 2, 3, 4, 3, 2, 3, 2, 1, 2]

            # Create distance matrix from data, where neighboring points are connected by edges
            # and the edge weights are equal to the max of both vertex births.
            n = length(data)
            dists = spzeros(n, n)
            for i in 1:n
                dists[i, i] = data[i]
            end
            for i in 1:(n - 1)
                j = i + 1
                dists[i, j] = dists[j, i] = max(dists[i, i], dists[j, j])
            end
            # 0-dimensional persistence should find values of minima and maxima of our data.
            d0 = ripserer(dists; dim_max=0)[1]
            @test d0 == [(2, 3), (1, 4), (0, Inf)]
        end

        @testset "with Rips" begin
            data = [1, 0, 1, 2, 3, 4, 3, 2, 3, 2, 1, 2]

            n = length(data)
            dists = zeros(Int, (n, n))
            for i in 1:n, j in 1:n
                if abs(i - j) ≤ 1
                    dists[i, j] = max(data[i], data[j])
                else
                    dists[i, j] = 5
                end
            end
            d0 = ripserer(dists; dim_max=0)[1]
            @test d0 == [(2, 3), (1, 4), (0, Inf)]
        end
    end

    @testset "Homology and explicit cohomology" begin
        @testset "Representative cycle" begin
            res_hom = ripserer(cycle; alg=:homology, reps=true, dim_max=3)
            @test vertices.(simplex.(representative(res_hom[2][1]))) ==
                sort!(vcat([(i + 1, i) for i in 1:17], [(18, 1)]))
        end
        @testset "Infinite intervals" begin
            @test_broken ripserer(
                Rips(cycle; threshold=2); alg=:homology, implicit=true
            )[2][1] == (1.0, Inf)
            @test_broken ripserer(
                Rips, cycle; alg=:homology, threshold=2, implicit=false
            )[2][1] == (1.0, Inf)
            @test ripserer(cycle; alg=:involuted, threshold=2)[2][1] == (1.0, Inf)
        end
    end

    @testset "Interface" begin
        test_filtration(Rips, cycle; modulus=3, dim_max=2)
    end

    @testset "Errors" begin
        @testset "Overflow checking" begin
            @test_throws OverflowError ripserer(Rips{Int16}(ones(1000, 1000)))
        end
        @testset "Unsupported algirithms" begin
            @test_throws ArgumentError ripserer(ones(5, 5); alg=:something)
        end
        @testset "Int or Float64 field type" begin
            @test_throws ErrorException ripserer(ones(5, 5); field=Int)
            @test_throws ErrorException ripserer(ones(5, 5); field=Float64)
            @test_throws ErrorException ripserer(ones(5, 5); field=UInt8)
        end
        @testset "Explicit cohomology reperesentatives unsupported" begin
            @test_throws ErrorException ripserer(cycle; implicit=false, reps=true)
        end
    end
end


module ITAE

export itae

"""
    ```julia
    itae(error, time)
    ```
    ITAE (Integral of Time-Weighted Absolute Error) from error vector.
    Inputs:
    error -> error vector.
    time  -> time vector.
"""
function itae(error::Vector{T}, time::Vector{T}) where T
    sum(time .* abs.(error))
end

"""
    ```julia
    itae(error, time)
    ```
    ITAE (Integral of Time-Weighted Absolute Error) from vector of error vector.
    Inputs:
    error -> vector of error vector.
    time  -> time vector.
"""
function itae(error::Vector{Vector{T}}, time::Vector{T}) where T
    map(x->sum(time .* abs.(x)), error)
end

"""
    ```julia
    itae(x, t, desired)
    ```
    ITAE (Integral of Time-Weighted Absolute Error) from error betwwen desired value and controler output.
    Inputs:
    x -> controler output vector.
    t -> time vector.
    desired -> desired value function.
"""
function itae(x::Vector{T}, t::Vector{T}, desired::Function) where T
    x = x .- map(desired, t)
    sum(t .* abs.(x))
end

end


import NNlib
using AutoGrad
using Statistics: mean
using LinearAlgebra: lmul!
export pool, mean

#= NOTES
- maxpool/meanpool vs op=maximum, op=mean (would require Statistics)
- adaptivexxx vs specifying output=1 or window=Inf.
- torch.nn: MaxPoolNd, MaxUnpoolNd, AvgPoolNd, FractionalMaxPoolNd, LPPoolNd, AdaptiveMaxPoolNd, AdaptiveAvgPoolNd
- torch.nn.functional: avg_poolNd, max_poolNd, max_unpoolNd, lp_poolNd, adaptive_max_poolNd, adaptive_avg_poolNd
- adaptive versions take output_size as argument, regular versions take kernel_size.
- torch also supports dilation? cudnn does not.
- torch supports maxunpool but not avgunpool.
- existence of lp may make the op kwarg more preferrable (other ops in the future?)
- keras has MaxPoolingNd, AveragePoolingNd, GlobalMaxPoolingNd, GlobalAveragePoolingNd: only output=1 supported
- keras GlobalPooling includes reshape.
=#

"""
    pool(x; kwargs...)

Compute pooling of input values (i.e., the maximum or average of several adjacent values) to
produce an output with smaller height and/or width.

If `x` has dimensions `(X1,X2,...,Cx,N)`, the result `y` will have dimensions
`(Y1,Y2,...,Cx,N)` where

    Yi=1+floor((Xi+2*padding[i]-window[i])/stride[i])

Here `Cx` is the number of input channels, `N` is the number of instances, and `Xi,Yi` are
spatial dimensions. `window`, `padding`, `stride` are keyword arguments that can be
specified as a single number (in which case they apply to all dimensions), or an array/tuple
with entries for each spatial dimension. If `window[i] > Xi+2*padding[i]` it is truncated to
fit the padded input, e.g. one can use `window=typemax(Int)` to implement global pooling.

# Keywords:

* `op=maximum`: pooling operation, can be `maximum` or `mean`
* `window=2`: the pooling window size for each dimension.
* `padding=0`: the number of extra zeros implicitly concatenated at the start and at the end of each dimension.
* `stride=window`: the number of elements to slide to reach the next pooling window.
* `propagateNaN=false`: if false do not propagate NaN values in max pooling.
* `includePadding=false`: if false do not include padded values in the count for average.
* `channelmajor = false`: assume channel-major format tensors if true
* `alpha=1`: can be used to scale the result.

"""
function pool(x; op=maximum, window=2, padding=0, stride=window, propagateNaN=false, includePadding=false, channelmajor=false, alpha=1)
    checkpoolopts(x, op, window, padding, stride, propagateNaN, includePadding, channelmajor, alpha)
    if channelmajor; x = permutedims(x,(2,3,1,4)); end
    N = ndims(x)
    window = NNlib.expand(Val(N-2), window)
    stride = NNlib.expand(Val(N-2), stride)
    padding = NNlib.expand(Val(N-2), padding)
    window = min.(window, size(x)[1:N-2] .+ 2 .* padding)
    pdims = NNlib.PoolDims(x, window; padding = padding, stride = stride)
    y = (op == maximum ? NNlib.maxpool(x, pdims) : NNlib.meanpool(x, pdims))
    if channelmajor; y = permutedims(y, (3,1,2,4)); end
    alpha == 1 ? y : lmul!(eltype(y)(alpha), y)
end

function ∇pool(x,y,dy; op=maximum, window=2, padding=0, stride=window, propagateNaN=false, includePadding=false, channelmajor=false, alpha=1)
    checkpoolopts(x, op, window, padding, stride, propagateNaN, includePadding, channelmajor, alpha)
    if alpha != 1; y = y ./ eltype(y)(alpha); end
    if channelmajor; x,y,dy = (a->permutedims(a,(2,3,1,4))).((x,y,dy)); end
    N = ndims(x)
    window = NNlib.expand(Val(N-2), window)
    stride = NNlib.expand(Val(N-2), stride)
    padding = NNlib.expand(Val(N-2), padding)
    window = min.(window, size(x)[1:N-2] .+ 2 .* padding)
    pdims = NNlib.PoolDims(x, window; padding = padding, stride = stride)
    dx = (op == maximum ? NNlib.∇maxpool(dy, y, x, pdims) : NNlib.∇meanpool(dy, y, x, pdims))
    if channelmajor; dx = permutedims(dx, (3,1,2,4)); end
    alpha == 1 ? dx : lmul!(eltype(dx)(alpha), dx)
end

@primitive1 pool(x;o...),dy,y  ∇pool(x,y,dy;o...)
@primitive1 ∇pool(x,y,dy;o...),ddx,dx  nothing  nothing  pool(ddx;o...)

function checkpoolopts(x, op, window, padding, stride, propagateNaN, includePadding, channelmajor, alpha)
    @assert op ∈ (maximum, mean) "Pooling op=$op should be one of (maximum,mean)"
    if op == maximum && !propagateNaN
        @warn "propagateNaN=false not yet implemented in NNlib, NaNs will be propagated. See https://github.com/FluxML/NNlib.jl/issues/218" maxlog=1
    end
    if op == mean && any(padding .> 0) && !includePadding
        @warn "includePadding=false not yet implemented in NNlib, pads will be included. See https://github.com/FluxML/NNlib.jl/issues/218" maxlog=1
    end
    if channelmajor
        @warn "channelmajor is not yet implemented in NNlib, using slow permutedims, see NNlib#267" maxlog=1
    end
end


# # Integration with SimpleSDMLayers

# The `SpatialBoundaries` package works really well with `SimpleSDMLayers`, so
# that you can (i) apply wombling and boundaries finding to a `SimpleSDMLayer`
# object, and (ii) convert the output of a `Womble` object to a *pair* of
# `SimpleSDMLayer` corresponding to the rate and direction of change.

using SimpleSDMLayers
using SpatialBoundaries
using StatsPlots

# Note that the warning about dependencies is a side-effect of loading some
# functionalities for `SimpleSDMLayers` as part of `SpatialBoundaries`, and can
# safely be ignored.

# In this example, we will look at temperature data over North America, and
# measure the rate of change, as well as the direction of it. These data are
# extracted from the WorldClim database.

temperature = SimpleSDMPredictor(
    WorldClim, BioClim, 1; left=-130.0, right=-65.0, bottom=5.0, top=65.0
)

# We can have a look at this layer, after setting a few defaults:

default(; dpi=500, size=(600, 600), aspectratio=1, c=:davos, frame=:box)
plot(temperature)

# There is an overload of the `wombling` method for SDM layers, so we can call
# it directly -- this method might result in a bit more memory usage than
# expected, as it requires to transform the `nothing` values into `NaN`s, which
# in turn might require to convert the inner elements of the layer grid.

W = wombling(temperature)

# By default, this returns a `LatticeWomble`. Let's look at the direction of
# change -- mapping this information is difficult, so we will focus on the

stephist(
    deg2rad.(sort(vec(W.θ)));
    proj=:polar,
    lab="",
    c=:teal,
    fill=(0, 0.2, :teal),
    nbins=100,
)

# The values for the direction of change are concentrated around 180° - note
# that the direction of change is reported as a wind direction, meaning that
# values tend to increase on a south-north axis.

# We can also map the rate of change. This is far easier to do with a proper SDM
# layer, so we will convert the wombling output:

Lr, Ld = SimpleSDMPredictor(W)

# Note that we do not use `convert` here, because this call returns *two* layers
# in a tuple -- this is a slight deviation from what we expect with
# `SimpleSDMLayers`, but it makes the code a little easier to write, and so is
# considered an acceptable trade-off.

plot(Lr)

# With the `Ld` layer storing the direction of change, we *can* do a map - the
# main issue with mapping directions is colors, which need to have the same
# begin and endpoint, while remaining accessible to people with color vision
# defficiencies. [Cyclic color schemes](https://www.fabiocrameri.ch/colourmaps/)
# work for this purpose:

plot(Ld, c=:brocO, clim=(0, 360))


using Test

ProjDir = joinpath(dirname(@__FILE__), "..", "examples", "BernoulliInitTheta")
cd(ProjDir) do

  isdir("tmp") &&
    rm("tmp", recursive=true);

    bernoullimodel = "
    data {
      int<lower=1> N;
      int<lower=0,upper=1> y[N];
    }
    parameters {
      real<lower=0,upper=1> theta;
    }
    model {
      theta ~ beta(1,1);
      y ~ bernoulli(theta);
    }
    "

    bernoullidata = [
      Dict("N" => 10, "y" => [0, 1, 0, 1, 0, 0, 0, 0, 0, 1]),
      Dict("N" => 10, "y" => [0, 1, 0, 0, 0, 0, 1, 0, 0, 1]),
      Dict("N" => 10, "y" => [0, 0, 0, 0, 0, 0, 1, 0, 1, 1]),
      Dict("N" => 10, "y" => [0, 0, 0, 1, 0, 0, 0, 1, 0, 1])
    ]

    theta_init = [0.1, 0.4, 0.5, 0.9]
    bernoulliinit = Dict[
      Dict("theta" => theta_init[1]),
      Dict("theta" => theta_init[2]),
      Dict("theta" => theta_init[3]),
      Dict("theta" => theta_init[4]),
    ]

    monitor = ["theta", "lp__", "accept_stat__"]

    stanmodel = Stanmodel(name="bernoulli",
      model=bernoullimodel,
      init=Stan.Init(init=bernoulliinit),
      adapt=1);

      sim1 = stan(stanmodel, bernoullidata, CmdStanDir=CMDSTAN_HOME)
      global sim1_values = round.(sim1.value[1, end, :], 1)

  isdir("tmp") &&
    rm("tmp", recursive=true);

  sim1_values |> display

  @assert sim1_values == theta_init "Not close to theta input values."
end # cd


"""
Scale Transform
```
    l = 2.0
    tr = ScaleTransform(l)
```
Multiply every element of the input by `l`
"""
struct ScaleTransform{T<:Real} <: Transform
    s::Vector{T}
end

function ScaleTransform(s::T=1.0) where {T<:Real}
    @check_args(ScaleTransform, s, s > zero(T), "s > 0")
    ScaleTransform{T}([s])
end

set!(t::ScaleTransform,ρ::Real) = t.s .= [ρ]
params(t::ScaleTransform) = first(t.s)
dim(str::ScaleTransform) = 1

transform(t::ScaleTransform,x::AbstractVecOrMat,obsdim::Int=defaultobs) = first(t.s) * x

Base.isequal(t::ScaleTransform,t2::ScaleTransform) = isequal(first(t.s),first(t2.s))


using PSLQ, LinearAlgebra
using Test

@testset "PSLQ" begin include("pslq.jl") end


"""
    cudnnNormalizationForward(x, xmean, xvar, bias, scale; o...)
    cudnnNormalizationForward!(y, x, xmean, xvar, bias, scale; o...)

Return batch normalization applied to `x`:

    y .= ((x .- mean(x; dims)) ./ sqrt.(epsilon .+ var(x; dims))) .* scale .+ bias  # training
    y .= ((x .- xmean) ./ sqrt.(epsilon .+ xvar)) .* scale .+ bias                  # inference

Bias and scale are trainable parameters, xmean and xvar are modified to collect statistics
during training and treated as constants during inference. Note that during inference the
values given by xmean and xvar arguments are used in the formula whereas during training the
actual mean and variance of the minibatch are used in the formula: the xmean/xvar arguments
are only used to collect statistics. In the original paper bias is referred to as beta and
scale as gamma (Batch Normalization: Accelerating Deep Network Training by Reducing Internal
Covariate Shift, S. Ioffe, C. Szegedy, 2015).

Keyword arguments:
* `epsilon = 1e-5`: epsilon value used in the normalization formula
* `exponentialAverageFactor = 0.1`: factor used in running mean/variance calculation: `runningMean = runningMean*(1-factor) + newMean*factor`
* `training = false`: boolean indicating training vs inference mode
* `mode::cudnnNormMode_t = CUDNN_NORM_PER_CHANNEL`: Per-channel layer is based on the paper. In this mode `scale` etc. have dimensions (1,1,C,1). The other alternative is `CUDNN_NORM_PER_ACTIVATION` where `scale` etc. have dimensions `(W,H,C,1)`.
* `algo::cudnnNormAlgo_t = CUDNN_NORM_ALGO_STANDARD`: The other alternative, `CUDNN_NORM_ALGO_PERSIST`, triggers the new semi-persistent NHWC kernel when certain conditions are met (see cudnn docs).
* `normOps::cudnnNormOps_t = CUDNN_NORM_OPS_NORM`: Currently the other alternatives, `CUDNN_NORM_OPS_NORM_ACTIVATION` and `CUDNN_NORM_OPS_NORM_ADD_ACTIVATION` are not supported.
* `z = nothing`: for residual addition to the result of the normalization operation, prior to the activation (will be supported when CUDNN_NORM_OPS_NORM_ADD_ACTIVATION is supported)
* `groupCnt = 1`: Place holder for future work, should be set to 1 now
* `alpha = 1; beta = 0`: scaling parameters: return `alpha * new_y + beta * old_y`

"""
cudnnNormalizationForward, cudnnNormalizationForward!

# Public methods
cudnnNormalizationForward(x, xmean, xvar, bias, scale; o...) = cudnnNormalizationForwardWithDefaults(x, xmean, xvar, bias, scale; o...)
cudnnNormalizationForward!(y, x, xmean, xvar, bias, scale; o...) = cudnnNormalizationForwardWithDefaults(x, xmean, xvar, bias, scale; y, o...)

# Private method
function cudnnNormalizationForwardWithDefaults(
    x, mean, variance, bias, scale;

    # Inference parameters:
    y = similar(x),
    z = nothing, # for residual addition to the result of the normalization operation, prior to the activation
    mode::cudnnNormMode_t = CUDNN_NORM_PER_CHANNEL, # Per-channel layer is based on the paper Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift, S. Ioffe, C. Szegedy, 2015.
    normOps::cudnnNormOps_t = CUDNN_NORM_OPS_NORM,  # Currently CUDNN_NORM_OPS_NORM_ACTIVATION and CUDNN_NORM_OPS_NORM_ADD_ACTIVATION are only supported in the NHWC layout (training,backward), not supported (inference)
    algo::cudnnNormAlgo_t = CUDNN_NORM_ALGO_STANDARD, # trigger the new semi-persistent NHWC kernel when CUDNN_NORM_ALGO_PERSIST
    alpha::Real = 1,
    beta::Real = 0,
    epsilon::Real = Cdouble(1e-5), # Has to be >= 0. Should be the same in forward and backward functions.
    groupCnt::Integer = Cint(1),   # Place hold for future work, should be set to 1 now

    # Training-only parameters:
    training = false,
    exponentialAverageFactor::Real = Cdouble(0.1),
    savedMean = nothing, # Optionally save intermediate results from the forward pass here - can be reused to speed up backward pass. NULL if unused.
    savedInvVariance = nothing,

    # Activation parameters:
    activationMode::cudnnActivationMode_t = CUDNN_ACTIVATION_IDENTITY,
    activationReluNanOpt::cudnnNanPropagation_t = CUDNN_NOT_PROPAGATE_NAN,
    activationCoef::Real = 1,
    activationDesc::Union{Nothing,cudnnActivationDescriptor} = (normOps == CUDNN_NORM_OPS_NORM ? nothing : cudnnActivationDescriptor(activationMode, activationReluNanOpt, Cdouble(activationCoef))),

    # Tensor descriptors:
    format::cudnnTensorFormat_t = CUDNN_TENSOR_NCHW,
    xDesc::cudnnTensorDescriptor = cudnnTensorDescriptor(x; format),
    yDesc::cudnnTensorDescriptor = cudnnTensorDescriptor(y; format),
    zDesc::Union{Nothing,cudnnTensorDescriptor} = (z === nothing ? nothing : cudnnTensorDescriptor(z; format)),
    normScaleBiasDesc::cudnnTensorDescriptor = cudnnTensorDescriptor(scale; format),
    normMeanVarDesc::Union{Nothing,cudnnTensorDescriptor} = (mean === nothing ? nothing : cudnnTensorDescriptor(mean; format)),

    # Temporary space used in training:
    workspace = nothing,
    reserveSpace = nothing,
    dx = Ref{Any}(),
    dscale = Ref{Any}(),
    dbias = Ref{Any}(),
    dz = Ref{Any}(),
)
    @assert epsilon >= 0 && exponentialAverageFactor >= 0  "epsilon and exponentialAverageFactor should be non-negative."
    @assert groupCnt == 1  "Currently only groupCnt=1 is supported."
    @assert normOps === CUDNN_NORM_OPS_NORM "Currently only normOps=CUDNN_NORM_OPS_NORM is supported."
    alpha, beta = (a->scalingParameter(eltype(x),a)).((alpha, beta))
    # Backward called separately on each variable. We will calculate all gradients on first call. Use `dready` to avoid subsequent calls.
    dready = Ref{Bool}(false)   # this will be turned to `true` by the first backward call.
    cudnnNormalizationForwardAD(x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready)
end

# AD method:
function cudnnNormalizationForwardAD(x, scale, bias, z; training, mean, variance, y, mode, normOps, algo, alpha, beta, epsilon, groupCnt, exponentialAverageFactor, savedMean, savedInvVariance, activationDesc, xDesc, yDesc, zDesc, normScaleBiasDesc, normMeanVarDesc, workspace, reserveSpace, dx, dscale, dbias, dz, dready)
    issimilar(x,y) = (typeof(x) === typeof(y) && (x === nothing || size(x) === size(y)))
    if training
        mean === nothing ? savedMean = nothing : savedMean === nothing ? savedMean = similar(mean) : @assert issimilar(mean, savedMean)
        variance === nothing ? savedInvVariance = nothing : savedInvVariance === nothing ? savedInvVariance = similar(variance) : @assert issimilar(variance, savedInvVariance)
        workspaceSize, reserveSpaceSize = cudnnNormalizationTempSpaceSizes(mode, normOps, algo, xDesc, zDesc, yDesc, normScaleBiasDesc, activationDesc, normMeanVarDesc, groupCnt)
        if reserveSpaceSize > 0 && reserveSpace === nothing; reserveSpace = cudnnTempSpace(reserveSpaceSize); end
        @assert sizeof(reserveSpace) >= reserveSpaceSize  "reserveSpace should be at least $(reserveSpaceSize) bytes"
        if workspaceSize > 0 && workspace === nothing; workspace = cudnnTempSpace(workspaceSize); end
        @assert sizeof(workspace) >= workspaceSize  "workspace should be at least $(workspaceSize) bytes"
        cudnnNormalizationForwardTraining(handle(), mode, normOps, algo, alpha, beta, xDesc, x, normScaleBiasDesc, scale, bias, exponentialAverageFactor, something(normMeanVarDesc,C_NULL), something(mean,CU_NULL), something(variance,CU_NULL), epsilon, something(savedMean,CU_NULL), something(savedInvVariance,CU_NULL), something(activationDesc,C_NULL), something(zDesc,C_NULL), something(z,CU_NULL), yDesc, y, something(workspace,CU_NULL), sizeof(workspace), something(reserveSpace,CU_NULL), sizeof(reserveSpace), groupCnt)
    else
        @assert mean !== nothing && variance !== nothing && normMeanVarDesc !== nothing "normalization mean and variance are required in inference mode."
        cudnnNormalizationForwardInference(handle(), mode, normOps, algo, alpha, beta, xDesc, x, normScaleBiasDesc, scale, bias, normMeanVarDesc, mean, variance, something(zDesc,C_NULL), something(z,CU_NULL), something(activationDesc,C_NULL), yDesc, y, epsilon, groupCnt)
    end
    return y
end

# Helper functions
function cudnnNormalizationTempSpaceSizes(mode, normOps, algo, xDesc, zDesc, yDesc, normScaleBiasDesc, activationDesc, normMeanVarDesc, groupCnt)
    workspaceSize, reserveSpaceSize = Ref{Csize_t}(0), Ref{Csize_t}(0)
    cudnnGetNormalizationForwardTrainingWorkspaceSize(handle(), mode, normOps, algo, xDesc, something(zDesc,C_NULL), yDesc, normScaleBiasDesc, something(activationDesc,C_NULL), something(normMeanVarDesc,C_NULL), workspaceSize, groupCnt)
    cudnnGetNormalizationTrainingReserveSpaceSize(handle(), mode, normOps, algo, something(activationDesc,C_NULL), xDesc, reserveSpaceSize, groupCnt)
    workspaceSize[], reserveSpaceSize[]
end


function solve_005()
    D = Dict{Int64, Int64}()
    for i = 1:20
        D[i] = 0
    end
    for n = 4:20
        V = get_prime_factors(n)
        for p = 1:20
            multiplicity = sum(V.==p)
            if multiplicity > D[p] D[p] = multiplicity end
        end
    end
    prod = 1
    for p = 1:20
        prod *= p^D[p]
    end
    return prod
end


module Born

import GeoPhyInv: AGeom, SrcWav
import GeoPhyInv: Medium
import GeoPhyInv.Data
using SpecialFunctions
using FFTW
using DSP

mutable struct Param
end

"""
k = wavenumber = 2pif/v0
"""
function G0_homo_acou(x::T, z::T, k::T, rho0::T) where {T<:AbstractFloat}
	sqrt(x*x + z*z) == 0.0 ? error("distance cannot be zero") : nothing
	k == 0.0 ? error("wavenumber cannot be zero") : nothing
	G0 = -0.25 * rho0 * im * complex.(besselj0(k*sqrt(x*x + z*z)), -bessely0(k*sqrt(x*x + z*z)))
	return G0
end

function mod(;
	     vp0::Float64=2000.0,
	     rho0::Float64=2000.0,
	     medium_pert::Medium=nothing,
             born_flag::Bool=false,
	     tgridmod::StepRangeLen=nothing,
	     tgrid::StepRangeLen = tgridmod,
	     ageom::AGeom=nothing,
	     srcwav::SrcWav=nothing,
	     src_flag::Int64=2,
		  )
	(length(ageom) != length(srcwav))  ? error("different supersources") : nothing
	@info "fix this"
#	(length(ageom) != getfield(srcwav,:ns))  ? error("different sources") : nothing

	nt = (length(tgridmod) == length(srcwav[1].grid)) ? length(tgrid) : error("srcwav tgrid")
	np2 = nextpow(2, 2*nt);

	if(born_flag)
		mesh_x=medium_pert.mgrid[2]
		mesh_z=medium_pert.mgrid[1]
		nz=length(mesh_z)
		nx=length(mesh_x)
		δx = step(mesh_x)
		δz = step(mesh_z)

		δmodtt = medium_pert[:KI] - (vp0 * vp0 * rho0)^(-1)
		δmodrr = medium_pert[:rhoI] - (rho0)^(-1)
	end

	fnpow2grid = FFTW.fftfreq(np2,inv(step(tgridmod)));

	nss = length(ageom)
	data=Data(tgridmod, ageom, [:P])

	for iss in 1:nss, ifield in 1:length(data[iss].d)
		sx = ageom[iss].sx
		rx = ageom[iss].rx
		sz = ageom[iss].sz
		rz = ageom[iss].rz

		for ir = 1:ageom[iss].nr

		dtemp=zeros(nt)
		dpow2all=complex.(zeros(np2), zeros(np2));
		wpow2=complex.(zeros(np2), zeros(np2));

		for is=1:ageom[iss].ns
			# zero pad wavelet
			for it in 1:nt
				wpow2[it]=complex(srcwav[iss].d[ifield][it,is])
			end
			FFTW.fft!(wpow2) # source wavelet in the frequency domain

			x = sx[is] - rx[ir]
			z = sz[is] - rz[ir]
			# analytical expression for every frequency, except zero
			dpow2 = complex.(zeros(np2), zeros(np2));
			dpow2[1] = complex.(0.0, 0.0)
			for iω in 2:np2
				ω = 2. * pi * abs(fnpow2grid[iω])
				k = ω / vp0

				if(born_flag)
					term = complex(0., 0.)
					for ix=1:nx
						@simd for iz=1:nz
							if(δmodtt[iz,ix] ≠ 0.0)
								term += (G0_homo_acou(sx[is]-mesh_x[ix], sz[is]-mesh_z[iz], k, rho0)[1]
									 * G0_homo_acou(rx[ir]-mesh_x[ix], rz[ir]-mesh_z[iz], k, rho0)[1] .* ω .* ω .* δmodtt[iz,ix]) * δx * δz
								# factor due to integration
							end
						end
					end

				else
					if(src_flag == 2)
						term = G0_homo_acou(x, z, k, rho0);
					elseif(src_flag == 1)
						dpow2[iω] = G0_homo_acou(x, z, k, rho0) * im * abs(fnpow2grid[iω])
					else
						error("invalid src_flag")
					end
				end
				if(fnpow2grid[iω] > 0)
					dpow2[iω] = term;
				elseif(fnpow2grid[iω] < 0)
					dpow2[iω] = conj(term);
				end

			end
			# convolution and stack over simultaneous sources
			dpow2all += dpow2 .* wpow2;
		end

		# back to time domain
		FFTW.ifft!(dpow2all)

		# truncate
		for it in 1:nt
			data[iss].d[ifield][it,ir] = real(dpow2all[it])
		end
		end
	end
	return data
end

end # module


using Test
using ForwardDiff
using Rotations
include("./BEMUtils.jl")

function primitive(x, h, r)
  η = sqrt(h^2 + r^2 * x^2 + x^2)
  # return x * log(η / sqrt(h^2 + x^2) + r * x / sqrt(h^2 + x^2)) -
  # 1/2 * h * atan((h * (η^2 - r^3 * x * η)) / ((r^2 * (x^2 + h^2) - η^2) * η^2 - h^2 * r^4 * x^2))

  return x * log(η / sqrt(h^2 + x^2) + r * x / sqrt(h^2 + x^2)) -
  1/2 * h * angle((r^2 - 1) * h^2 - (r^2 + 1) * x^2 + im * 2 * h *r * η)

end

function deriv(x, h, r)
  return log(sqrt(1 + r^2 * x^2 / (x^2 + h^2)) + r * x / sqrt(x^2 + h^2))
end

function grad_primitive(x, h, r)
  return ForwardDiff.derivative(x -> primitive(x, h, r), x)
end

function ∫Φ_rectangle_numeric(a, b, h)
  Φ(xy) = 1/sqrt(xy[1]^2 + xy[2]^2 + h^2)
  cb_result = hcubature(Φ, (0.0,0.0), (a,b))
  return cb_result[1]
end

function ∫Φ_rectangle(a, b, h)
  return ∫Φ_right_angle(a, b, h) + ∫Φ_right_angle(b, a, h)
end

function ∫Φ_right_angle_test()
  a, b, h = rand(3)
  # @show ∫Φ_rectangle_numeric(a, b, h) - ∫Φ_rectangle(a, b, h)
  return ∫Φ_rectangle_numeric(a, b, h) ≈ ∫Φ_rectangle(a, b, h)
end

function ∫Φ_general_angle_test()
  a, b, d, h = rand(4) .- 0.5
  a⃗ = SVector(a, d)
  α⃗ = SVector(a, 0.)
  b⃗ = SVector(b, d)
  β⃗ = SVector(b, 0.)

  return abs(∫Φ_rectangle_numeric(abs(a), abs(d), h) -
             sign(a * b) * ∫Φ_rectangle_numeric(abs(b), abs(d), h)) ≈
         abs(sign(a * b) * ∫Φ_general_angle(α⃗, a⃗, h) -
             ∫Φ_general_angle(b⃗, β⃗, h)) +
         ∫Φ_general_angle(a⃗, b⃗, h)
end

function ∫Φ_square_test()
  c1, c2, h = rand(3) .-0.5; d = rand()
  c1 *= 4; c2 *= 4
  c⃗ = SVector(c1, c2)

  Φ(xy) = 1/sqrt(xy[1]^2 + xy[2]^2 + h^2)
  cb_result = hcubature(Φ, (c1 - d/2, c2 - d/2), (c1 + d/2, c2 + d/2))
  cb_result[1]
  ∫Φ_square(c⃗, d, h)
  return ∫Φ_square(c⃗, d, h) ≈ cb_result[1]
end

function matrix_element_point_scaling_test()
  d = rand()
  v_point = rand(3)
  v_scaling = rand(3)
  O = rand(RotMatrix{3})
  point = Point(v_point)
  scaling = rotate(HaarScaling(d, v_scaling), O)

  @show constructMatrixElement(point, scaling)
  @show constructMatrixElement_numeric(point, scaling)

  return constructMatrixElement(point, scaling) ≈ constructMatrixElement_numeric(point, scaling)
end

function matrix_element_scaling_scaling_test()
  d1, d2 = rand(), rand()
  v1, v2 = rand(3), rand(3)
  O1, O2 = rand(RotMatrix{3}), rand(RotMatrix{3})
  e1 = rotate(HaarScaling(d1, v1), O1)
  e2 = rotate(HaarScaling(d2, v2), O2)

  @show constructMatrixElement(e1, e2)
  @show constructMatrixElement_numeric(e1, e2)
  return constructMatrixElement(e1, e2) ≈ constructMatrixElement_numeric(e1, e2)
end

@test ∫Φ_right_angle_test()
@test ∫Φ_general_angle_test()
@test ∫Φ_square_test()
@test matrix_element_point_scaling_test()
@test matrix_element_scaling_scaling_test()


# This file is a part of Julia. License is MIT: https://julialang.org/license

##### mean #####

"""
    mean(f::Function, v)

Apply the function `f` to each element of `v` and take the mean.

```jldoctest
julia> mean(√, [1, 2, 3])
1.3820881233139908

julia> mean([√1, √2, √3])
1.3820881233139908
```
"""
function mean(f::Callable, iterable)
    y = iterate(iterable)
    if y == nothing
        throw(ArgumentError("mean of empty collection undefined: $(repr(iterable))"))
    end
    count = 1
    value, state = y
    f_value = f(value)
    total = reduce_first(add_sum, f_value)
    y = iterate(iterable, state)
    while y !== nothing
        value, state = y
        total += f(value)
        count += 1
        y = iterate(iterable, state)
    end
    return total/count
end
mean(iterable) = mean(identity, iterable)
mean(f::Callable, A::AbstractArray) = sum(f, A) / _length(A)

"""
    mean!(r, v)

Compute the mean of `v` over the singleton dimensions of `r`, and write results to `r`.

# Examples
```jldoctest
julia> v = [1 2; 3 4]
2×2 Array{Int64,2}:
 1  2
 3  4

julia> mean!([1., 1.], v)
2-element Array{Float64,1}:
 1.5
 3.5

julia> mean!([1. 1.], v)
1×2 Array{Float64,2}:
 2.0  3.0
```
"""
function mean!(R::AbstractArray, A::AbstractArray)
    sum!(R, A; init=true)
    x = max(1, _length(R)) // _length(A)
    R .= R .* x
    return R
end

"""
    mean(v; dims)

Compute the mean of whole array `v`, or optionally along the given dimensions.

!!! note
    Julia does not ignore `NaN` values in the computation. Use the [`missing`](@ref) type
    to represent missing values, and the [`skipmissing`](@ref) function to omit them.
"""
mean(A::AbstractArray; dims=:) = _mean(A, dims)

_mean(A::AbstractArray{T}, region) where {T} = mean!(reducedim_init(t -> t/2, +, A, region), A)
_mean(A::AbstractArray, ::Colon) = sum(A) / _length(A)

##### median & quantiles #####

"""
    middle(x)

Compute the middle of a scalar value, which is equivalent to `x` itself, but of the type of `middle(x, x)` for consistency.
"""
middle(x::Union{Bool,Int8,Int16,Int32,Int64,Int128,UInt8,UInt16,UInt32,UInt64,UInt128}) = Float64(x)
# Specialized functions for real types allow for improved performance
middle(x::AbstractFloat) = x
middle(x::Real) = (x + zero(x)) / 1

"""
    middle(x, y)

Compute the middle of two reals `x` and `y`, which is
equivalent in both value and type to computing their mean (`(x + y) / 2`).
"""
middle(x::Real, y::Real) = x/2 + y/2

"""
    middle(range)

Compute the middle of a range, which consists of computing the mean of its extrema.
Since a range is sorted, the mean is performed with the first and last element.

```jldoctest
julia> middle(1:10)
5.5
```
"""
middle(a::AbstractRange) = middle(a[1], a[end])

"""
    middle(a)

Compute the middle of an array `a`, which consists of finding its
extrema and then computing their mean.

```jldoctest
julia> a = [1,2,3.6,10.9]
4-element Array{Float64,1}:
  1.0
  2.0
  3.6
 10.9

julia> middle(a)
5.95
```
"""
middle(a::AbstractArray) = ((v1, v2) = extrema(a); middle(v1, v2))

"""
    median!(v)

Like [`median`](@ref), but may overwrite the input vector.
"""
function median!(v::AbstractVector)
    isempty(v) && throw(ArgumentError("median of an empty array is undefined, $(repr(v))"))
    if eltype(v)<:AbstractFloat
        @inbounds for x in v
            isnan(x) && return x
        end
    end
    inds = axes(v, 1)
    n = _length(inds)
    mid = div(first(inds)+last(inds),2)
    if isodd(n)
        return middle(partialsort!(v,mid))
    else
        m = partialsort!(v, mid:mid+1)
        return middle(m[1], m[2])
    end
end
median!(v::AbstractArray) = median!(vec(v))

"""
    median(v; dims)

Compute the median of an entire array `v`, or, optionally,
along the given dimensions. For an even number of
elements no exact median element exists, so the result is
equivalent to calculating mean of two median elements.

!!! note
    Julia does not ignore `NaN` values in the computation. Use the [`missing`](@ref) type
    to represent missing values, and the [`skipmissing`](@ref) function to omit them.
"""
median(v::AbstractArray; dims=:) = _median(v, dims)

_median(v::AbstractArray, dims) = mapslices(median!, v, dims)

_median(v::AbstractArray{T}, ::Colon) where {T} = median!(copyto!(Array{T,1}(undef, _length(v)), v))

# for now, use the R/S definition of quantile; may want variants later
# see ?quantile in R -- this is type 7
"""
    quantile!([q, ] v, p; sorted=false)

Compute the quantile(s) of a vector `v` at the probability or probabilities `p`, which
can be given as a single value, a vector, or a tuple. If `p` is a vector, an optional
output array `q` may also be specified. (If not provided, a new output array is created.)
The keyword argument `sorted` indicates whether `v` can be assumed to be sorted; if
`false` (the default), then the elements of `v` may be partially sorted.

The elements of `p` should be on the interval [0,1], and `v` should not have any `NaN`
values.

Quantiles are computed via linear interpolation between the points `((k-1)/(n-1), v[k])`,
for `k = 1:n` where `n = length(v)`. This corresponds to Definition 7 of Hyndman and Fan
(1996), and is the same as the R default.

!!! note
    Julia does not ignore `NaN` values in the computation: `quantile!` will
    throw an `ArgumentError` in the presence of `NaN` values in the data array.
    Use the [`missing`](@ref) type to represent missing values, and the
    [`skipmissing`](@ref) function to omit them.

* Hyndman, R.J and Fan, Y. (1996) "Sample Quantiles in Statistical Packages",
  *The American Statistician*, Vol. 50, No. 4, pp. 361-365
"""
function quantile!(q::AbstractArray, v::AbstractVector, p::AbstractArray;
                   sorted::Bool=false)
    if size(p) != size(q)
        throw(DimensionMismatch("size of p, $(size(p)), must equal size of q, $(size(q))"))
    end
    isempty(q) && return q

    minp, maxp = extrema(p)
    _quantilesort!(v, sorted, minp, maxp)

    for (i, j) in zip(eachindex(p), eachindex(q))
        @inbounds q[j] = _quantile(v,p[i])
    end
    return q
end

quantile!(v::AbstractVector, p::AbstractArray; sorted::Bool=false) =
    quantile!(similar(p,float(eltype(v))), v, p; sorted=sorted)

quantile!(v::AbstractVector, p::Real; sorted::Bool=false) =
    _quantile(_quantilesort!(v, sorted, p, p), p)

function quantile!(v::AbstractVector, p::Tuple{Vararg{Real}}; sorted::Bool=false)
    isempty(p) && return ()
    minp, maxp = extrema(p)
    _quantilesort!(v, sorted, minp, maxp)
    return map(x->_quantile(v, x), p)
end

# Function to perform partial sort of v for quantiles in given range
function _quantilesort!(v::AbstractArray, sorted::Bool, minp::Real, maxp::Real)
    isempty(v) && throw(ArgumentError("empty data vector"))

    if !sorted
        lv = length(v)
        lo = floor(Int,1+minp*(lv-1))
        hi = ceil(Int,1+maxp*(lv-1))

        # only need to perform partial sort
        sort!(v, 1, lv, Sort.PartialQuickSort(lo:hi), Base.Sort.Forward)
    end
    isnan(v[end]) && throw(ArgumentError("quantiles are undefined in presence of NaNs"))
    return v
end

# Core quantile lookup function: assumes `v` sorted
@inline function _quantile(v::AbstractVector, p::Real)
    0 <= p <= 1 || throw(ArgumentError("input probability out of [0,1] range"))

    lv = length(v)
    f0 = (lv - 1)*p # 0-based interpolated index
    t0 = trunc(f0)
    h  = f0 - t0
    i  = trunc(Int,t0) + 1

    T  = promote_type(eltype(v), typeof(v[1]*h))

    if h == 0
        return convert(T, v[i])
    else
        a = v[i]
        b = v[i+1]
        if isfinite(a) && isfinite(b)
            return convert(T, a + h*(b-a))
        else
            return convert(T, (1-h)*a + h*b)
        end
    end
end

"""
    quantile(v, p; sorted=false)

Compute the quantile(s) of a vector `v` at a specified probability or vector or tuple of
probabilities `p`. The keyword argument `sorted` indicates whether `v` can be assumed to
be sorted.

The `p` should be on the interval [0,1], and `v` should not have any `NaN` values.

Quantiles are computed via linear interpolation between the points `((k-1)/(n-1), v[k])`,
for `k = 1:n` where `n = length(v)`. This corresponds to Definition 7 of Hyndman and Fan
(1996), and is the same as the R default.

!!! note
    Julia does not ignore `NaN` values in the computation: `quantile` will
    throw an `ArgumentError` in the presence of `NaN` values in the data array.
    Use the [`missing`](@ref) type to represent missing values, and the
    [`skipmissing`](@ref) function to omit them.

- Hyndman, R.J and Fan, Y. (1996) "Sample Quantiles in Statistical Packages",
  *The American Statistician*, Vol. 50, No. 4, pp. 361-365
"""
quantile(v::AbstractVector, p; sorted::Bool=false) =
    quantile!(sorted ? v : copymutable(v), p; sorted=sorted)


using Test, Lattices

lattice = Square(4, 5)

test_sites = []
for i = 1:5
    for j = 1:4
        push!(test_sites, (j, i))
    end
end
test_sites
@test collect(sites(lattice)) == test_sites
edges(lattice)

function check_index(lattice::Square{Tuple{h, w}, Periodic}) where {h, w}
    count = 1
    for i = 1:w
        for j = 1:h
            @test lattice[j, i] == count
            count += 1
        end
    end

    @test lattice[h+1, 2] == lattice[1, 2]
    @test lattice[2, w+1] == lattice[2, 1]

    true
end

function check_index(lattice::Square{Tuple{h, w}, Fixed}) where {h, w}
    count = 1
    for i = 1:w
        for j = 1:h
            @test lattice[j, i] == count
            count += 1
        end
    end
    true
end

@test check_index(Square(4, 5, boundary=Periodic))
@test check_index(Square(4, 5, boundary=Fixed))

lattice = Square(4, 5, boundary=Fixed)
@test_throws BoundsError lattice[1, 6]


"""
$(TYPEDEF)

The abstract representation of the cost for a player in a game.
"""
abstract type PlayerCost end

"""
    $(FUNCTIONNAME)(x::SVector{nx}, u::SVector{nu}, t::AbstractFloat)

A functor to evaluate the cost representation.
"""
function (pc::PlayerCost) end

"""
    $FUNCTIONNAME(pc::PlayerCost, x::SVector{nx}, u::SVector{nu}, t::AbstractFloat)

Computes the quadratic approximation of the cost.

We provide a convencience default below.
"""
function quadraticize end

"""
    $(FUNCTIONNAME)(pc::PlayerCost)

Returns the index of the player this cost is associated with.
"""
function player_id end

"""
$(TYPEDEF)

Represents a parameter free player cost that is defined by a funciton f.
"""
struct FunctionPlayerCost{F} <: PlayerCost
    f::F
end
(pc::FunctionPlayerCost)(args...) = pc.f(args...)

"""
$(TYPEDEF)

Represents the quadratic players costs for a singel player in a game.

# Fields

$(TYPEDFIELDS)
"""
struct QuadraticPlayerCost{nx, nu, TL<:SVector{nx}, TQ<:SMatrix{nx, nx},
                           TRL<:SVector{nu}, TRQ<:SMatrix{nu, nu}} <: PlayerCost
    "The linear state cost"
    l::TL
    "The qudratic state cost matrix"
    Q::TQ
    "The linear term of the input cost."
    r::TRL
    "A square matrix to represent the quadratic control cost for this player"
    R::TRQ
end

(pc::QuadraticPlayerCost)(::Any, x::SVector, u::SVector) = pc(x, u)

function (pc::QuadraticPlayerCost)(x::SVector, u::SVector)
    return 1//2 * x'*pc.Q*x + pc.l'*x + 1//2 * u'*pc.R*u + pc.r'*u
end


module Flight

using Reexport
@reexport using BenchmarkTools

#general
include("utils.jl")
include("misc.jl")
include("modeling.jl")
include("plotting.jl")

#math
include("quaternions.jl")
include("attitude.jl")

#environment & dynamics
include("geodesy.jl")
include("terrain.jl")
include("atmosphere.jl")

include("kinematics.jl")
include("dynamics.jl")

#aircraft components
include("components.jl")
include("airdata.jl")
include("propulsion.jl")
include("landinggear.jl")
include("aerodynamics.jl")

#IO interfaces
include("input.jl")
include("output.jl")

# #aircraft
include("aircraft.jl")
include("c172.jl")

@reexport using .Utils
@reexport using .Misc
@reexport using .Modeling
@reexport using .Plotting

@reexport using .Quaternions
@reexport using .Attitude

@reexport using .Geodesy
@reexport using .Terrain
@reexport using .Atmosphere

@reexport using .Kinematics
@reexport using .Dynamics

@reexport using .Components
@reexport using .Airdata
@reexport using .Propulsion
@reexport using .LandingGear
@reexport using .Aerodynamics

@reexport using .Input
@reexport using .Output

@reexport using .Aircraft
@reexport using .C172

export ftest

ftest() = println("Welcome to Flight")

end


infile = "./src/Day17/input.txt"
input = readlines(infile)

using DataStructures

cube = Dict{CartesianIndex, Bool}()
for x = 1:length(input[1]), y in 1:length(input)
    z = 0
    if input[y][x] == '#'
        cube[CartesianIndex(x,y,z)] = true
    end
end

neighbourgrid = CartesianIndices((-1:1, -1:1, -1:1))

function countactiveneighbours(cube, loc, neighbourgrid)
    grid = loc .+ neighbourgrid
    neighbours = sum(cube[i] for i in grid if haskey(cube, i))
    if haskey(cube, loc)
        count = cube[loc] ? neighbours - 1 : neighbours
    else
        count = neighbours
    end
    return count
end

test = countactiveneighbours(cube, CartesianIndex(8,6,0), neighbourgrid)

function getnext(cube, neighbourgrid)
    nextcube = Dict{CartesianIndex, Bool}()
    for centre in keys(cube)
        # println("Centre is $centre")
        # ps = centre .+ neighbourgrid
        for p in (centre .+ neighbourgrid)
            # println("Current point: $p")
            if haskey(nextcube, p)
                # already pre-calculated
                # println("Pre-done")
                continue
            end
            mystate = haskey(cube, p) ? cube[p] : false
            neighbourcount = countactiveneighbours( cube, p, neighbourgrid)
            if mystate
                if (neighbourcount == 2) || ( neighbourcount == 3)
                    active = true
                else
                    active = false
                end
            else
                if neighbourcount == 3
                    active = true
                else
                    active = false
                end
            end
            nextcube[p] = active
        end
    end
    return nextcube
end

t = getnext(cube, neighbourgrid)

cycles = 6
function part1(cube, neighbourgrid, cycles)
    for i in 1:cycles
        cube = getnext(cube, neighbourgrid)
    end
    return cube
end

cube
sixthcube = part1(cube, neighbourgrid, cycles)
active6 = sum(v for v in values(sixthcube))

neighbourgrid2 = CartesianIndices((-1:1, -1:1, -1:1, -1:1))

cube2 = Dict{CartesianIndex, Bool}()
for x = 1:length(input[1]), y in 1:length(input)
    z = 0
    w = 0
    if input[y][x] == '#'
        cube2[CartesianIndex(x,y,z,w)] = true
    end
end

using BenchmarkTools
@profview part2 = part1(cube2, neighbourgrid2, cycles)
activepart2 = sum(v for v in values(part2))
@benchmark part1(cube2, neighbourgrid2, 3)


#---------------------------------------------------------
# # [Shepp-Logan Phantoms](@id 07-shepp)
#---------------------------------------------------------

#=
This page illustrates the Shepp-Logan phantoms in the Julia package
[`ImagePhantoms`](https://github.com/JuliaImageRecon/ImagePhantoms.jl).

This page was generated from a single Julia file:
[07-shepp.jl](@__REPO_ROOT_URL__/07-shepp.jl).
=#

#md # In any such Julia documentation,
#md # you can access the source code
#md # using the "Edit on GitHub" link in the top right.

#md # The corresponding notebook can be viewed in
#md # [nbviewer](http://nbviewer.jupyter.org/) here:
#md # [`07-shepp.ipynb`](@__NBVIEWER_ROOT_URL__/07-shepp.ipynb),
#md # and opened in [binder](https://mybinder.org/) here:
#md # [`07-shepp.ipynb`](@__BINDER_ROOT_URL__/07-shepp.ipynb).

# ### Setup

# Packages needed here.

using ImagePhantoms
using ImageGeoms: ImageGeom, axesf
using MIRTjim: jim, prompt

# The following line is helpful when running this example.jl file as a script;
# this way it will prompt user to hit a key after each figure is displayed.

isinteractive() ? jim(:prompt, true) : prompt(:draw);

# ### Overview

# There are several variations of the Shepp-Logan phantom available.

# ### CT version

#=
Original version from:
Larry A Shepp, Benjamin F Logan,
"The Fourier reconstruction of a head section,"
IEEE Transactions on Nuclear Science, 21(3):21-42, June 1974.
[doi](http://doi.org/10.1109/TNS.1974.6499235)

This (default) version has low soft tissue contrast,
so it usually should be displayed
with a narrow window the using `clim` option:
=#

image1 = shepp_logan(256) # CT version by default
jim(image1, "SheppLogan (original CT version)", clim=(0.95, 1.05))

# ### Over-sampling

#=
When generating ellipse phantoms,
it is generally preferable
to "over-sample" the ellipse values within each pixel
to account for partial volume effect.
A factor of 3 over-sampling (along both axes) typically suffices,
so this factor is the default for the `shepp_logan` method.
Here is how the phantom image looks without over-sampling:
=#

image1o = shepp_logan(256; oversample=1)
jim(image1o, "No over sampling", clim=(0.95,1.05))

#=
Note that boundaries of the interior ellipses look smoother
in the original version with the default 3× over-sampling.
Most of the remaining examples use the recommended default over-sampling.
=#

# ### Toft version

#=
This version is from:
Toft, Peter Aundal & Sørensen, John Aasted
"The Radon transform-theory and implementation,"
Technical University of Denmark (DTU), 1996. Page 201.
[pdf](https://files.openpdfs.org/1ra51GP6gJO.pdf)
=#

image2 = shepp_logan(256, SheppLoganToft())
jim(image2, "SheppLoganToft")

# ### Emission tomography version

#=
This version has low intensity for the skull
because typical PET/SPECT radiotracers
do not accumulate in bone regions.
It is probably also useful for MRI,
because typical MRI scans
have low signal from bone.
=#

image3 = shepp_logan(256, SheppLoganEmis())
jim(image3, "SheppLoganEmis")

# ### BrainWeb version

#=
This version was inspired by the
[BrainWeb](https://brainweb.bic.mni.mcgill.ca)
phantoms that have integer indices
for each of the different regions.
It should not be used directly,
but rather one should assign meaningful intensity values
to each of the integer indices.
=#

image4 = shepp_logan(256, SheppLoganBrainWeb())
jim(image4, "SheppLoganBrainWeb")

# For the BrainWeb version, there is no over-sampling by default,
# to preserve the integer indices.

# ### Disjoint middle ellipses

# Sometimes it can be more convenient
# to have the middle ellipses be non-overlapping:

ob = ellipse_parameters(SheppLoganBrainWeb(), disjoint=true)
ob[:,end] = 1:10
ob = Ellipse(ob)
x = LinRange(-0.4, 0.4, 206)
y = LinRange(-0.5, 0.5, 256)
oversample = 3
image5 = phantom(x, y, ob, oversample)
jim(x, y, image5, "Disjoint")

# ### Comedy version

image6 = shepp_logan(256, SouthPark(); fovs=(1,1))
jim(image6, "SouthPark")


import Pkg
Pkg.activate(".")
using Revise
using SimpleSDMLayers # dev version
using Statistics

temperature = worldclim(1)
precipitation = worldclim(12)
wc_vars = worldclim(1:19)

# Extrema overload
minimum(temperature)
maximum(temperature)
extrema(temperature)

# Broadcast type
broadcast(abs, temperature)
maximum(abs(temperature))

# Max-min
@time min(temperature, precipitation)
@time max(temperature, precipitation)
@time test = reduce(max, wc_vars)
wc_vars_resp = convert.(SimpleSDMResponse, wc_vars)
@time test = reduce(max, wc_vars_resp)

minifloat = SimpleSDMPredictor([-1 0 1 nothing])
miniint = SimpleSDMResponse([-1.0 0.0 1.0 nothing])
mini32 = SimpleSDMPredictor([Float32.([-1.0 0.0 1.0])... nothing])
min(miniint, minifloat)
@time reduce(min, [miniint, minifloat, mini32])

# Mean-std
mean(wc_vars)
std(wc_vars)
# Attempt 1
mininothing = SimpleSDMPredictor([nothing 1.0 nothing nothing])
coordsnothing = map(x -> findall(isnothing, x.grid), [minifloat, miniint, mini32, mininothing])
uniquenothing = unique(coordsnothing)
unique(reduce(vcat, uniquenothing)) |> sort
# Need InvertedIndex...
# Attempt 2
coordsnotnothing = map(x -> findall(!isnothing, x.grid), [minifloat, miniint, mini32, mininothing])
uniquenot = unique(coordsnotnothing)
@time reduce(intersect, uniquenot)
@time intersect(uniquenot...)

# Math operations
temperature + precipitation
temperature - precipitation
temperature * precipitation
temperature + precipitation
reduce(+, wc_vars)


module TestSingularLRRL

using Modia
@usingModiaPlot

SingularLRRL = Model(
    v0   = 10,
    L1_L = 0.1,
    L2_L = 0.2,
    R1_R = 10,
    R2_R = 20,

    L2_p_i = Var(init = 0.1),
    R2_p_i = Var(start = 0.0),
    V_n_v = Var(start = 0.0),

    equations =:[
        # Inductor 1
        L1_v = L1_p_v - L1_n_v
        0 = L1_p_i + L1_n_i
        L1_L*der(L1_p_i) = L1_v

        # Inductor 2
        L2_v = L2_p_v - L2_n_v
        0 = L2_p_i + L2_n_i
        L2_L*der(L2_p_i) = L2_v

        # Resistor 1
        R1_v = R1_p_v - R1_n_v
        0 = R1_p_i + R1_n_i
        R1_v = R1_R*R1_p_i

        # Resistor 2
        R2_v = R2_p_v - R2_n_v
        0 = R2_p_i + R2_n_i
        R2_v = R2_R*R2_p_i

        # Voltage source
        V_v = V_p_v - V_n_v
        0   = V_p_i + V_n_i

        # Connect equations
        #   connect(V.p, L1.p)
        #   connect(L1.n, R1.p)
        #   connect(L1.n, R2.p)
        #   connect(R1.n, L2.p)
        #   connect(R2.n, L2.p)
        #   connect(L2.n, V.n)
        V_p_v = L1_p_v
        0 = V_p_i + L1_p_i
        L1_n_v = R1_p_v
        L1_n_v = R2_p_v
        0 = L1_n_i + R1_p_i + R2_p_i
        R1_n_v = L2_p_v
        R2_n_v = L2_p_v
        0 = R1_n_i + R2_n_i + L2_p_i
        L2_n_v = V_n_v
        0 = L2_n_i + V_n_i #+ ground_v_i
        #V_n_v = 0
        V_v = v0 ]
)

singularLRRL = @instantiateModel(SingularLRRL, unitless=true)

simulate!(singularLRRL, Tsit5(), stopTime = 1.0, requiredFinalStates = [1.4999999937318995])

plot(singularLRRL, [("L1_p_i", "L2_p_i", "R1_p_i", "R2_p_i"), ("L1_v", "L2_v", "R1_v", "R2_v")])

include("../models/Electric.jl")

setLogMerge(false)

SingularLRRL2 = Model(
    R1 = Resistor | Map(R=10u"Ω"),
    R2 = Resistor | Map(R=20u"Ω", start = Map(v=0.0u"V")),
    L1 = Inductor | Map(L=0.1u"H", i = Var(init=nothing)),
    L2 = Inductor | Map(L=0.2u"H", start = Map(v=0.0u"V")),
    V = ConstantVoltage | Map(V=10u"V"),
    connections = :[
        (V.p, L1.p)
        (L1.n, R1.p, R2.p)
        (L2.p, R1.n, R2.n)
        (L2.n, V.n) ]
)

#=
singularLRRL2 = @instantiateModel(SingularLRRL2, unitless=true, log=true, logCode=true, logExecution=true)

simulate!(singularLRRL2, Tsit5(), stopTime = 1.0)

plot(singularLRRL2, [("L1.p.i", "L2.p.i", "R1.p.i", "R2.p.i"), ("L1.v", "L2.v", "R1.v", "R2.v")])
=#

SingularLRRL3 = Model(
    R1 = Resistor | Map(R=1),
    R2 = Resistor | Map(R=2),
    R3 = Resistor | Map(R=3),
    R4 = Resistor | Map(R=4),
    L1 = Inductor | Map(L=0.1),
    L2 = Inductor | Map(L=0.2, i = Var(init=nothing)),
    V = ConstantVoltage | Map(V=10),
#    ground = Ground,
    connections = :[
        (V.p, L1.p)
        (L1.n, R1.p, R2.p)
        (R1.n, R2.n, R3.p, R4.p)
        (L2.p, R3.n, R4.n)
        (L2.n, V.n) ] #, ground.p) ]
)

singularLRRL3 = @instantiateModel(SingularLRRL3, unitless=true, log=false, logDetails=false)

simulate!(singularLRRL3, Tsit5(), stopTime = 1.0, requiredFinalStates = [4.198498632779839])

plot(singularLRRL3, [("L1.i", "L2.i", "R1.i", "R2.i", "R3.i", "R4.i"), ("L1.v", "L2.v", "R1.v", "R2.v", "R3.v", "R4.v")])

end


export srcubature

struct SphericalRadialCubature <: AbstractApproximationMethod end

approximation_name(::SphericalRadialCubature)       = "SphericalRadial"
approximation_short_name(::SphericalRadialCubature) = "SR"

function srcubature()
    return SphericalRadialCubature()
end

function getweights(::SphericalRadialCubature, mean::AbstractVector{T}, covariance::AbstractMatrix{T}) where { T <: Real }
    d = length(mean)
    return Base.Generator(1:2d + 1) do i
        return i === (2d + 1) ? 1.0 / (d + 1) : 1.0 / (2.0(d + 1))
    end
end

function getpoints(::SphericalRadialCubature, mean::AbstractVector{T}, covariance::AbstractMatrix{T}) where { T <: Real }
    d = length(mean)
    L = cholsqrt(covariance)

    tmpbuffer = zeros(d)
    sigma_points = Base.Generator(1:2d + 1) do i
        if i === (2d + 1)
            fill!(tmpbuffer, 0.0)
        else
            tmpbuffer[rem((i - 1), d) + 1] = sqrt(d + 1) * (-1)^(div(i - 1, d))
            if i !== 1
                tmpbuffer[rem((i - 2), d) + 1] = 0.0
            end
        end
        return tmpbuffer
    end

    tbuffer = similar(mean)
    return Base.Generator(sigma_points) do point
        copyto!(tbuffer, mean)
        return mul!(tbuffer, L, point, 1.0, 1.0) # point = m + 1.0 * L * point
    end
end


using Distributions
using Base.Test

v = 7.0
S = eye(2)
S[1, 2] = S[2, 1] = 0.5

W = Wishart(v,S)
IW = InverseWishart(v,S)

for d in [W,IW]
    @test size(d) == size(rand(d))
    @test length(d) == length(rand(d))
    @test typeof(d)(params(d)...) == d
    @test partype(d) == Float64

    ### conversions
    D = typeof(d).name.primary
    @test typeof(convert(D{Float32}, d)) == typeof(D(Float32(v), Array{Float32}(S)))

    @test typeof(convert(D{Float32}, params(d)...)) == typeof(D(Float32(v), Array{Float32}(S)))
end

@test partype(Wishart(7, eye(Float32, 2))) == Float32
@test partype(InverseWishart(7, eye(Float32, 2))) == Float32

@test_approx_eq_eps mean(rand(W,100000)) mean(W) 0.1
@test_approx_eq_eps mean(rand(IW,100000)) mean(IW) 0.1

v = 3.0

@test_approx_eq_eps pdf(Wishart(v,S), S) 0.04507168 1e-8
@test_approx_eq_eps pdf(Wishart(v,S), inv(S)) 0.01327698 1e-8
@test_approx_eq_eps pdf(Wishart(v,inv(S)),S) 0.0148086 1e-8
@test_approx_eq_eps pdf(Wishart(v,inv(S)),inv(S)) 0.01901462 1e-8

@test_approx_eq logpdf(Wishart(v,S), S) log(pdf(Wishart(v,S), S))
@test_approx_eq logpdf(Wishart(v,S), inv(S)) log(pdf(Wishart(v,S), inv(S)))
@test_approx_eq logpdf(Wishart(v,inv(S)), S) log(pdf(Wishart(v,inv(S)), S))
@test_approx_eq logpdf(Wishart(v,inv(S)), inv(S)) log(pdf(Wishart(v,inv(S)), inv(S)))

@test_approx_eq_eps pdf(InverseWishart(v,S), S) 0.04507168 1e-8
@test_approx_eq_eps pdf(InverseWishart(v,S), inv(S)) 0.006247377 1e-8
@test_approx_eq_eps pdf(InverseWishart(v,inv(S)),S)  0.03147137 1e-8
@test_approx_eq_eps pdf(InverseWishart(v,inv(S)),inv(S)) 0.01901462 1e-8

@test_approx_eq logpdf(InverseWishart(v,S), S) log(pdf(InverseWishart(v,S), S))
@test_approx_eq logpdf(InverseWishart(v,S), inv(S)) log(pdf(InverseWishart(v,S), inv(S)))
@test_approx_eq logpdf(InverseWishart(v,inv(S)), S) log(pdf(InverseWishart(v,inv(S)), S))
@test_approx_eq logpdf(InverseWishart(v,inv(S)), inv(S)) log(pdf(InverseWishart(v,inv(S)), inv(S)))


mutable struct VTNode{T,Id}
    nentries::Int
    values::Matrix{T}
    ids::Vector{Id}
    children::Vector{VTNode{T,Id}}
    distances::Vector{Tuple{Float64,Int}}
end

function VTNode(::Type{T}, ::Type{Id}, d::Int; capacity::Int) where {T,Id}
    values = Matrix{T}(undef, d, capacity)
    ids = Vector{Id}(undef, capacity)
    children = Vector{VTNode{T,Id}}(undef, capacity)
    distances = Vector{Tuple{Float64,Int}}(undef, capacity)
    VTNode(0, values, ids, children, distances)
end

function VTNode(v::AbstractVector{T}, id::Id; kwargs...) where {T,Id}
    node = VTNode(T, Id, length(v); kwargs...)
    node.nentries = 1
    node.values[:, 1] .= v
    node.ids[1] = id
    node
end

Base.length(node::VTNode) = node.nentries
Base.isempty(node::VTNode) = length(node) == 0
capacity(node::VTNode) = length(node.children)

function compute_distances!(node, x, metric::M) where {M}
    for j = 1:length(node)
        node.distances[j] = (metric(x, view(node.values, :, j)), j)
    end
    node.distances
end

function search_in_radius(
    node::VTNode{T,Id},
    x::AbstractVector,
    tol::Real,
    metric::M,
) where {T,Id,M}
    !isempty(node) || return nothing

    n = length(node)

    # We compute now the distance to every other element in the node
    # If the distance to one element is smaller than tol, we are done.
    # Otherwise, we need to look into those children whose distance
    # is closest (let's say with distance `d`) and all those
    # children for whose distance `dᵢ` we have |d-dᵢ| < 2*tol
    # This fact can be shown by using the triangle inequality.
    # We therefore simply compute the distances to every point and sort
    # afterwards

    # We go through the elements and compute the distances,
    # keeping track of the three smallest elements
    m₁ = m₂ = m₃ = (Inf, 1)
    # we have a distances cache per thread
    distances = compute_distances!(node, x, metric)
    for i = 1:n
        dᵢ = first(distances[i])
        # early exit
        if dᵢ < tol
            # we rely on the distances for insertion, so place at the first place the smallest element
            distances[1] = (dᵢ, i)
            return node.ids[i]
        end

        # check three smallest elements and update if necessary
        if dᵢ < first(m₁)
            m₃ = m₂
            m₂ = m₁
            m₁ = (dᵢ, i)
        elseif dᵢ < first(m₂)
            m₃ = m₂
            m₂ = (dᵢ, i)
        elseif dᵢ < first(m₃)
            m₃ = (dᵢ, i)
        end
    end

    # Now we computed all distances and also kept track of the two smallest elements
    # Now there are three cases
    # 1) m₁[1] + 2tol < m₂[1] -- > The element can only be in one subtree
    # 2) m₁[1] + 2tol < m₃[1] -- > The element can only be in the first or second subtree
    # 3) else -> The element can also be in more subtrees,
    #            we need to sort the distances vector and look through everything

    # Check smallest element first
    if isassigned(node.children, last(m₁))
        retid = search_in_radius(node.children[last(m₁)], x, tol, metric)
        if !isnothing(retid)
            # we rely on the distances for insertion, so place the smallest element first
            distances[1] = m₁
            return retid::Id
        end
    end
    # Case 1)
    if m₂[1] - m₁[1] > 2tol
        distances[1] = m₁
        return nothing # already checked first tree
    end
    # Case 2) We know m₂[1] - m₁[1] ≤ 2tol
    if isassigned(node.children, last(m₂))
        retid = search_in_radius(node.children[last(m₂)], x, tol, metric)
        if !isnothing(retid)
            # we rely on the distances for insertion, so place the smallest element first
            distances[1] = m₁
            return retid::Id
        end
    end

    if m₃[1] - m₁[1] > 2tol
        distances[1] = m₁
        return nothing # Checked first and second case
    end

    # Since we know also the third element, let's check it
    if isassigned(node.children, last(m₃))
        retid = search_in_radius(node.children[last(m₃)], x, tol, metric)
        if !isnothing(retid)
            # we rely on the distances for insertion, so place at the first place the smallest element
            distances[1] = m₁
            return retid::Id
        end
    end

    # Case 3)
    # We need to sort distances
    sort!(view(distances, 1:n), Base.Sort.InsertionSort, Base.Sort.By(first))

    # We can start at 4 since we already checked the smallest 3
    d = m₃[1]
    for k ∈ 4:n
        dᵢ, i = distances[k]
        if dᵢ - m₁[1] < 2tol
            if isassigned(node.children, i)
                retid = search_in_radius(node.children[i], x, tol, metric)
                if !isnothing(retid)
                    return retid::Id
                end
            end
        else
            break
        end
    end

    return nothing
end

function _insert!(
    node::VTNode{T,Id},
    v::AbstractVector,
    id::Id,
    metric::M;
    use_distances::Bool = false,
) where {T,Id,M}
    # if not filled so far, just add it to the current node
    if length(node) < capacity(node)
        k = (node.nentries += 1)
        node.values[:, k] .= v
        node.ids[k] = id
        return nothing
    end

    if use_distances
        dᵢ, minᵢ = first(node.distances)
    else
        compute_distances!(node, v, metric)
        dᵢ, minᵢ = findmin(node.distances)
    end

    if !isassigned(node.children, minᵢ)
        node.children[minᵢ] = VTNode(v, id; capacity = capacity(node))
    else # a node already exists, so recurse
        _insert!(node.children[minᵢ], v, id, metric)
    end

    nothing
end

function identifiers!(ids, node::VTNode)
    for i = 1:length(node)
        push!(ids, node.ids[i])
        if isassigned(node.children, i)
            identifiers!(ids, node.children[i])
        end
    end
    ids
end

"""
     VoronoiTree(
    v::AbstractVector{T},
    id::Id;
    metric = EuclideanNorm(),
    capacity = 8,
)

Construct a Voronoi tree data structure for vector `v` of element type `T` and with identifiers
`Id`. Each node has the given `capacity` and distances are measured by the given `metric`.
"""
mutable struct VoronoiTree{T,Id,M}
    root::VTNode{T,Id}
    nentries::Int
    metric::M
end

function VoronoiTree{T,Id}(d::Int; metric = EuclideanNorm(), capacity::Int = 8) where {T,Id}
    root = VTNode(T, Id, d; capacity = capacity)
    VoronoiTree(root, 0, metric)
end

function VoronoiTree(v::AbstractVector{T}, id::Id; kwargs...) where {T,Id}
    VoronoiTree{T,Id}(length(v); kwargs...)
end

Base.length(T::VoronoiTree) = T.nentries
Base.broadcastable(T::VoronoiTree) = Ref(T)

"""
    insert!(tree::VoronoiTree, v::AbstractVector, id)

Insert in the tree the point `v` with identifier `id`.
"""
function Base.insert!(
    tree::VoronoiTree{T,Id},
    v::AbstractVector,
    id::Id;
    use_distances::Bool = false,
) where {T,Id}
    _insert!(tree.root, v, id, tree.metric; use_distances = use_distances)
    tree.nentries += 1
    tree
end

"""
    search_in_radius(tree::VoronoiTree, v::AbstractVector, tol::Real)

Search whether the given `tree` contains a point `p` with distances at most `tol` from `v`.
Returns `nothing` if no point exists, otherwise the identifier of `p` is returned.
"""
function search_in_radius(tree::VoronoiTree, v::AbstractVector, tol::Real)
    search_in_radius(tree.root, v, tol, tree.metric)
end

"""
    add!(tree::VoronoiTree, v, id, tol)

Insert in the tree the point `v` with identifier `id` if `search_in_radius(tree v, tol)` is
`nothing`. If this is the case the tuple `(id, true)` is returned, otherwise the obtained
identifier and `false` is returned.
"""
function add!(tree::VoronoiTree{T,Id}, v, id::Id, tol::Real) where {T,Id}
    found_id = search_in_radius(tree, v, tol)
    if isnothing(found_id)
        insert!(tree, v, id; use_distances = true)
        return (id, true)
    else
        return (found_id::Id, false)
    end
end

function Base.collect(tree::VoronoiTree{T,Id}) where {T,Id}
    ids = Id[]
    identifiers!(ids, tree.root)
    ids
end


using Neighborhood
using Distances

export uzal_cost
export uzal_cost_local

"""
    uzal_cost(Y::Dataset; kwargs...) → L
Compute the L-statistic `L` for input dataset `Y` according to Uzal et al.[^Uzal2011], based on
theoretical arguments on noise amplification, the complexity of the
reconstructed attractor and a direct measure of local stretch which constitutes
an irrelevance measure. It serves as a cost function of a state space
trajectory/embedding and therefore allows to estimate a "goodness of a
embedding" and also to choose proper embedding parameters, while minimizing
`L` over the parameter space. For receiving the local cost function `L_local`
(for each point in state space - not averaged), use `uzal_cost_local(...)`.

## Keyword arguments

* `samplesize = 0.5`: Number of considered fiducial points v as a fraction of
  input state space trajectory `Y`'s length, in order to average the conditional
  variances and neighborhood sizes (read algorithm description) to produce `L`.
* `K = 3`: the amount of nearest neighbors considered, in order to compute σ_k^2
  (read algorithm description).
  If given a vector, minimum result over all `k ∈ K` is returned.
* `metric = Euclidean()`: metric used for finding nearest neigbhors in the input
  state space trajectory `Y.
* `w = 1`: Theiler window (neighbors in time with index `w` close to the point,
  that are excluded from being true neighbors). `w=0` means to exclude only the
  point itself, and no temporal neighbors.
* `Tw = 40`: The time horizon (in sampling units) up to which E_k^2 gets computed
  and averaged over (read algorithm description).

## Description
The `L`-statistic is based on theoretical arguments on noise amplification, the
complexity of the reconstructed attractor and a direct measure of local stretch
which constitutes an irrelevance measure. Technically, it is the logarithm of
the product of `σ`-statistic and a normalization statistic `α`:

L = log10(σ*α)

The `σ`-statistic is computed as follows. `σ = √σ² = √(E²/ϵ²)`.
`E²` approximates the conditional variance at each point in state space and
for a time horizon `T ∈ Tw`, using `K` nearest neighbors. For each reference
point of the state space trajectory, the neighborhood consists of the reference
point itself and its `K+1` nearest neighbors. `E²` measures how strong
a neighborhood expands during `T` time steps. `E²` is averaged over many time
horizons `T = 1:Tw`. Consequently, `ϵ²` is the size of the neighborhood at the
reference point itself and is defined as the mean pairwise distance of the
neighborhood. Finally, `σ²` gets averaged over a range of reference points on
the attractor, which is controlled by `samplesize`. This is just for performance
reasons and the most accurate result will obviously be gained when setting
`samplesize=1.0`

The `α`-statistic is a normalization factor, such that `σ`'s from different
embeddings can be compared. `α²` is defined as the inverse of the sum of
the inverse of all `ϵ²`'s for all considered reference points.

[^Uzal2011]: Uzal, L. C., Grinblat, G. L., Verdes, P. F. (2011). [Optimal reconstruction of dynamical systems: A noise amplification approach. Physical Review E 84, 016223](https://doi.org/10.1103/PhysRevE.84.016223).
"""
function uzal_cost(Y::AbstractDataset{D, ET};
        Tw::Int = 40, K::Int = 3, w::Int = 1, samplesize::Real = 0.5,
        metric = Euclidean()
    ) where {D, ET}

    # select a random state space vector sample according to input samplesize
    NN = length(Y)-Tw;
    NNN = floor(Int, samplesize*NN)
    ns = sample(1:NN, NNN; replace=false) # the fiducial point indices

    vs = Y[ns] # the fiducial points in the data set

    vtree = KDTree(Y[1:end-Tw], metric)
    allNNidxs, allNNdist = all_neighbors(vtree, vs, ns, K, w)
    ϵ² = zeros(NNN)             # neighborhood size
    E²_avrg = zeros(NNN)        # averaged conditional variance
    E² = zeros(Tw)
    ϵ_ball = zeros(ET, K+1, D) # preallocation
    u_k = zeros(ET, D)

    # loop over each fiducial point
    for (i,v) in enumerate(vs)
        NNidxs = allNNidxs[i] # indices of k nearest neighbors to v
        # pairwise distance of fiducial points and `v`
        pdsqrd = fiducial_pairwise_dist_sqrd(view(Y.data, NNidxs), v, metric)
        ϵ²[i] = (2/(K*(K+1))) * pdsqrd  # Eq. 16
        # loop over the different time horizons
        for T = 1:Tw
            E²[T] = comp_Ek2!(ϵ_ball, u_k, Y, ns[i], NNidxs, T, K, metric) # Eqs. 13 & 14
        end
        # Average E²[T] over all prediction horizons
        E²_avrg[i] = mean(E²)                   # Eq. 15
    end
    σ² = E²_avrg ./ ϵ² # noise amplification σ², Eq. 17
    σ²_avrg = mean(σ²) # averaged value of the noise amplification, Eq. 18
    α² = 1 / mean(ϵ².^(-1)) # for normalization, Eq. 21
    L = log10(sqrt(σ²_avrg)*sqrt(α²))
end

function fiducial_pairwise_dist_sqrd(fiducials, v, metric)
    pd = zero(eltype(fiducials[1]))
    pd += evaluate(metric, v, v)^2
    for (i, v1) in enumerate(fiducials)
        pd += evaluate(metric, v1, v)^2
        for j in i+1:length(fiducials)
            @inbounds pd += evaluate(metric, v1, fiducials[j])^2
        end
    end
    return sum(pd)
end

"""
    comp_Ek2!(ϵ_ball,u_k,Y,v,NNidxs,T,K,metric) → E²(T)
Returns the approximated conditional variance for a specific point in state space
`ns` (index value) with its `K`-nearest neighbors, which indices are stored in
`NNidxs`, for a time horizon `T`. This corresponds to Eqs. 13 & 14 in [^Uzal2011].
The norm specified in input parameter `metric` is used for distance computations.
"""
function comp_Ek2!(ϵ_ball, u_k, Y, ns::Int, NNidxs, T::Int, K::Int, metric)
    # determine neighborhood `T` time steps ahead
    ϵ_ball[1, :] .= Y[ns+T]
    @inbounds for (i, j) in enumerate(NNidxs)
        ϵ_ball[i+1, :] .= Y[j + T]
    end

    # compute center of mass
    @inbounds for i in 1:size(Y)[2]; u_k[i] = sum(view(ϵ_ball, :, i))/(K+1); end # Eq. 14

    E²_sum = 0
    @inbounds for j = 1:K+1
        E²_sum += (evaluate(metric, view(ϵ_ball, j, :), u_k))^2
    end
    E² = E²_sum / (K+1)         # Eq. 13
end

"""
    uzal_cost_local(Y::Dataset; kwargs...) → L_local
Compute the local L-statistic `L_local` for input dataset `Y` according to
Uzal et al.[^Uzal2011]. The length of `L_local` is `length(Y)-Tw` and
denotes a value of the local cost-function to each of the points of the
state space trajectory.

In contrast to [`uzal_cost`](@ref) `σ²` here does not get averaged over all the
state space reference points on the attractor. Therefore, the mean of 'L_local'
is different to `L`, when calling `uzal_cost`, since the averaging is performed
before logarithmizing.

Keywords as in [`uzal_cost`](@ref).

[^Uzal2011]: Uzal, L. C., Grinblat, G. L., Verdes, P. F. (2011). [Optimal reconstruction of dynamical systems: A noise amplification approach. Physical Review E 84, 016223](https://doi.org/10.1103/PhysRevE.84.016223).
"""
function uzal_cost_local(Y::AbstractDataset{D, ET};
        Tw::Int = 40, K::Int = 3, w::Int = 1, samplesize::Real = 0.5,
        metric = Euclidean()
    ) where {D, ET}

    # select a random state space vector sample according to input samplesize
    NN = length(Y)-Tw;
    NNN = floor(Int, samplesize*NN)
    ns = sample(1:NN, NNN; replace=false) # the fiducial point indices

    vs = Y[ns] # the fiducial points in the data set

    vtree = KDTree(Y[1:end-Tw], metric)
    allNNidxs, allNNdist = all_neighbors(vtree, vs, ns, K, w)
    ϵ² = zeros(NNN)             # neighborhood size
    E²_avrg = zeros(NNN)        # averaged conditional variance
    E² = zeros(Tw)
    ϵ_ball = zeros(ET, K+1, D) # preallocation
    u_k = zeros(ET, D)

    # loop over each fiducial point
    for (i,v) in enumerate(vs)
        NNidxs = allNNidxs[i] # indices of k nearest neighbors to v
        # pairwise distance of fiducial points and `v`
        pdsqrd = fiducial_pairwise_dist_sqrd(view(Y.data, NNidxs), v, metric)
        ϵ²[i] = (2/(K*(K+1))) * pdsqrd  # Eq. 16
        # loop over the different time horizons
        for T = 1:Tw
            E²[T] = comp_Ek2!(ϵ_ball, u_k, Y, ns[i], NNidxs, T, K, metric) # Eqs. 13 & 14
        end
        # Average E²[T] over all prediction horizons
        E²_avrg[i] = mean(E²)                   # Eq. 15
    end
    σ² = E²_avrg ./ ϵ² # noise amplification σ², Eq. 17
    α² = 1 / mean(ϵ².^(-1)) # for normalization, Eq. 21
    L_local = log10.(sqrt.(σ²).*sqrt(α²))
    return L_local
end


# PolyDerivEvaluator: used to compute the derivatives once instead of every call
struct PolyDerivEvaluator{F, D<:Tuple{Vararg{Polynomial}}}
    f::F
    derivs::D
end

@inline function PolyDerivEvaluator(p::Polynomial, ::Val{num_derivs}) where num_derivs
    PolyDerivEvaluator(p, make_poly_derivs(p, Val(num_derivs)))
end

make_poly_derivs(p::Polynomial, ::Val{0}) = ()
@inline function make_poly_derivs(p::Polynomial, ::Val{num_derivs}) where num_derivs
    p′ = SUP.derivative(p)
    (p′, make_poly_derivs(p′, Val(num_derivs - 1))...)
end

@generated function (pde::PolyDerivEvaluator)(θ, ::Val{num_derivs}) where num_derivs
    num_derivs >= 0 || return :(throw(ArgumentError("num_derivs must be nonnegative")))
    exprs = [:(pde.f(θ))]
    for i = 1 : num_derivs
        push!(exprs, :(pde.derivs[$i](θ)))
    end
    return :(tuple($(exprs...)))
end


using Distributed

function calc_napier_parallel()
    n_trial =1e9
    total_count = @distributed (+) for i in 1:n_trial
        counter = 0
        accumulated = 0.0
        while true
            accumulated += rand()
            counter += 1
            if accumulated >= 1.0
                break
            end
        end
        counter
    end
    println(total_count/n_trial)
end

addprocs()


# https://adventofcode.com/2020/day/10

# Each of your joltage adapters is rated for a specific output joltage (your puzzle input).
# Any given adapter can take an input 1, 2, or 3 jolts lower than its rating and still produce
# its rated output joltage.

# In addition, your device has a built-in joltage adapter rated for 3 jolts higher than the
# highest-rated adapter in your bag. (If your adapter list were 3, 9, and 6, your device's
# built-in adapter would be rated for 12 jolts.)

# Treat the charging outlet near your seat as having an effective joltage rating of 0.

function read_array(string::AbstractString)::Vector{Int}
    return parse.(Int, split(string, '\n', keepempty=false))
end

example1 = read_array("
16
10
15
5
1
11
7
19
6
12
4
")

example2 = read_array("
28
33
18
42
31
14
46
20
48
47
24
23
49
45
19
38
39
11
1
32
25
35
8
17
7
9
4
2
34
10
3
")

# What is the number of 1-jolt differences multiplied by the number of 3-jolt differences?

function part1(numbers)
    sorted = sort(numbers)
    sorted = [0; sorted; sorted[end] + 3]
    differences = diff(sorted)
    @assert minimum(differences) == 1
    @assert maximum(differences) == 3

    return sum(differences .== 1) * sum(differences .== 3)
end

@assert part1(example2) == 220

# What is the total number of distinct ways you can arrange the adapters to connect
# the charging outlet to your device?

"""
Cache for holding the partial solution, it maps: node => solutions_count.
"""
Cache = Dict{Int,Int}

"""
Use memoized shortcuts when traversing the tree. On the way, update the cache
in place. Returns the number of paths from the beggining of the sorted list
of nodes.
"""
function count_solutions!(adapters, cache::Cache=Cache())
    socket = adapters[1]
    if socket in keys(cache)
        return cache[socket]
    end

    solutions_found = 0
    for i = 2:4
        if i > length(adapters)
            break
        end
        if adapters[i] <= (socket + 3)
            if i == length(adapters)
                solutions_found += 1
            else
                # updates cache in place, will be available for next iteration
                solutions_found += count_solutions!(adapters[i:end], cache)
            end
        end
    end
    cache[socket] = solutions_found
    return solutions_found
end

@assert count_solutions!([0; sort(example1); 22], Cache()) == 8
@assert count_solutions!([0; sort(example2); 52], Cache()) == 19208

"""
Traverse the tree starting from the back, to build the memoization
cache.
"""
function init_cache(adapters, step_size=0.05)
    n = length(adapters)
    step = Int(max(1, round(n * step_size)))
    cache = Cache()

    pos = n
    while pos >= n / 3
        pos -= step
        count_solutions!(adapters[pos:end], cache)
    end
    return cache
end

function part2(numbers)
    sorted = sort(numbers)
    full_sequence = [0; sorted; sorted[end] + 3]
    cache = init_cache(full_sequence)
    return count_solutions!(full_sequence, cache)
end

@assert part2(example1) == 8
@assert part2(example2) == 19208

test = read_array(read("data/day-10.txt", String))
println("Part 1: $(result1 = part1(test))")
println("Part 2: $(result2 = part2(test))")

part1(test) == 2738
part2(test) == 74049191673856

println()
print("Example 1: ")
@time @assert part2(example1) == 8
print("Example 2: ")
@time @assert part2(example2) == 19208
print("Part 2:    ")
@time @assert result2 == 74049191673856


using DiffEqProblemLibrary.SDEProblemLibrary: importsdeproblems; importsdeproblems()
using DiffEqProblemLibrary.SDEProblemLibrary: prob_sde_lorenz, oval2ModelExample

oval2ModelExample()

oval2ModelExample(useBigs=true,largeFluctuations=true)


using Faulhaber
using Test

@testset "Faulhaber.jl" begin
    @test_throws ArgumentError faulhaber(10, Val(-7))

    for m in 1:200, p in 0:7, T in [Int, Int128, Float64]
        if T == Float64
            @test faulhaber(float(m), Val(p)) ≈ sum(i->i^p, 1:big(m))
        else
            @test faulhaber(T(m), Val(p)) === T(sum(i->i^p, 1:big(m)))
        end
    end
    for m in 1:13, p in 8:15
        @test faulhaber(m, Val(p)) == sum(i->i^p, 1:big(m))
    end
end


# Convolutional Neural Network
# Andrew Bannerman 5/2/2019

# RelU activation functions variants https://medium.com/tinymind/a-practical-guide-to-relu-b83ca804f1f7 + Julia: https://int8.io/neural-networks-in-julia-hyperbolic-tangent-and-relu/
# Convolutional Neural Network: https://www.youtube.com/watch?v=YRhxdVk_sIs
# CNN - ground up https://towardsdatascience.com/convolutional-neural-networks-from-the-ground-up-c67bb41454e1
# CNN - http://cs231n.github.io/convolutional-networks/
# CNN - https://adeshpande3.github.io/A-Beginner%27s-Guide-To-Understanding-Convolutional-Neural-Networks/
# CNN - https://github.com/wiseodd/hipsternet/blob/master/hipsternet/neuralnet.py
# CNN layer julia = https://github.com/FluxML/Flux.jl/blob/master/src/layers/conv.jl

#----------------------
# max pooling expalined: https://www.youtube.com/watch?v=ZjM_XQa5s6s
# - Zero padding = https://www.youtube.com/watch?v=qSTv_m-KFk0
#----------------------

# Needed packages
using DataFrames
using CSV
using Distributions
using Plots
using Random
using Images, CoordinateTransformations, TestImages, OffsetArrays, CoordinateTransformations # For data augmentation
using PaddedViews # for cropping images

# Train and test set data preparation

# Load MNIST Data set from .csv format
train_set = CSV.read("C:/Users/Andrew.Bannerman/Desktop/Julia/My_First_Neural_Network/Data/mnist_train.csv",header=false,types=fill(Float64,785))
test_set = CSV.read("C:/Users/Andrew.Bannerman/Desktop/Julia/My_First_Neural_Network/Data/mnist_test.csv",header=false,types=fill(Float64,785))

# Plot Example image
example = Float64.(train_set[2,2:size(train_set,2)])
example = rotl90(reshape(example,28,28))
heatmap(example)

# Extract labels
train_label = train_set[:,1]
test_label = test_set[:,1]

# Extract Data and exclude the label in the first column
train_set = train_set[:,2:size(train_set,2)]
test_set = test_set[:,2:size(test_set,2)]

# Pre process - Scale the inputs to a range of 0.01 to 1.00
m_train = Array{Float64}(zeros(nrow(train_set),size(train_set,2))) # empty matrix
@inbounds for i in 1:nrow(train_set)
    @inbounds for j in 1:size(train_set,2)
    m_train[i,j] = (train_set[i,j] / 255.0 * 0.99) + 0.01
    end
end

m_test = Array{Float64}(zeros(nrow(test_set),size(test_set,2))) # empty matrix
@inbounds for i in 1:nrow(test_set)
    @inbounds for j in 1:size(test_set,2)
    m_test[i,j] = (test_set[i,j] / 255.0 * 0.99) + 0.01
    end
end

# Convert matrix to DataFrame and append the labels back to the train and test set
scaled_train_inputs = DataFrame(hcat(train_label,m_train))
scaled_test_inputs = DataFrame(hcat(test_label,m_test))

# Plot the new scale example
example = reshape(Float64.(scaled_train_inputs[5,2:size(scaled_train_inputs,2)]),784,1)
example = rotl90(reshape(example,28,28))
print(example)
heatmap(example,c = :greys)

function train(wih::Array{Float64,2}, who::Array{Float64,2},inputs::Array{Float64,2}, targets::Array{Float64,2}; lr::Float64=0.1, drop_out::Bool=true, drop_out_p::Float64=.8) #wih, who, lr, inputs, targets)

  A = example

# Set weights and filter
# Initialize filter using a normal distribution with mean 0
# standard deviation inversely proportional the square root of the number of units
# Resource: https://stats.stackexchange.com/questions/200513/how-to-initialize-the-elements-of-the-filter-matrix
scale = 1.0
stdev = scale/sqrt(size(A,2)*(size(A,1)))
filt = reshape(rand(Normal(0, stdev), 9),3,3)

# Build the convolution layer
"""
    conv(A::Array{Float64,2}; m::Int64=3,n::Int64=3, s::Int64=1, p::Int64=1)
m x n = image convolution block size
s = stride
p = padding
padding = "valid" = no padding, "same" = keep conv dim output same as input dim, "custom" = custom p value
Input should be divisible by 2
"""
m=3
n=3
s=1
p=3
function conv_layer(A::Array{Float64,2}; m::Int64=3,n::Int64=3, s::Int64=1, p::Int64=0, padding::String="valid") # mxn: block_size / p::String="valid"
# Note need to add a check so that the full conv box with stride fits into the image
# https://github.com/JuliaArrays/PaddedViews.jl/issues/9
  M,N = size(A)
  mc = M+(p*2)-m+1          # no. vertical blocks (rows) # adjusted for the padding dim
  nc = N+(p*2)-n+1          # no. horizontal blocks (cols)
  if padding == "same" # want output dim of conv to be same as input dims
  p=0 # reset p
  #Calculate the first conv layer output dimensions
  w = size(A,2)
  f = size(filt,2)
  #M_out,N_out = (M-f+1) , (N-f+1)
  conv_out_dim = Int64.((w-f+2*p) / s+1) # final convoltion output dim
  p = Int64.((f-1) / 2) # Calculate correct padding value to keep the output the same dims as the input
  pad = p+1 # this is first element position within pad in A for PaddedView() PaddedView.jl
    A_pad = Float64.(PaddedView(0.0, A,(M+(p*2),N+(p*2)), (pad,pad)))
elseif padding == "valid"
  p=0 # reset p
  pad = p+1 # this is first element position within pad in A
    A_pad = Float64.(PaddedView(0.0, A,(M+(p*2),N+(p*2)), (pad,pad)))
elseif padding == "custom"
    p = p # keep as function input
    pad = p+1 # this is first element position within pad in A
      A_pad = Float64.(PaddedView(0.0, A,(M+(p*2),N+(p*2)), (pad,pad)))
  end
  #A_pad = Float64.(PaddedView(0.0, A,(M+(p*2),N+(p*2)), (pad,pad))) # p*2 to increase dim in both directions pad,pad = position of the first element to pad around
  #B = Array{eltype(A)}(undef, m*n, mc*nc)
  M,N = size(A_pad)
  mc = M-m+1          # no. vertical blocks (rows)
  nc = N-n+1          # no. horizontal blocks (cols)
  #B = zeros(m*n, mc*nc)
  B = zeros(m*n, 0)
  # Create custom index to allow for stride
  stride_col_index = collect(1:s:nc)
  stride_row_index = collect(1:s:mc)
  #B = Array{eltype(A)}(undef)
  # split the image into small convolution blocks
  # This function is similar to im2col
  i = 1
  j = 1
  @inbounds for j = 1:size(stride_col_index,1) # Loop over columns horizontal
    @inbounds for i = 1:size(stride_row_index,1) # loop over rows vertical
      @views block = A_pad[stride_row_index[i]:stride_row_index[i]+m-1, stride_col_index[j]:stride_col_index[j]+n-1]
      # Consider adding extra for loop to populate output memory wise...
        shapes = reshape(block,m*n,1)
       B = hcat(B, shapes) # Append to output matrix
    #B[:, j] = shapes
      end
  end
    # multiply the filter by the convolution block
    # Calculate spatial size of the volume
    # w = Input volume size
    # f = size of the Conv Layer neurons
    # s = Stride
    # p = Amount of Zero padding
    # Neuron fit = (W−F+2P)/S+1
    w = size(A_pad,2)
    f = size(filt,2)
    conv_out_dim = Int64.((w-f) / s+1) # This is conv output dimension (excluding padding)
    conv_out = zeros(conv_out_dim,conv_out_dim) # initialize output
    @inbounds for j in 1:size(B,2) # loop over column dim (horz.)
      conv_out[j] = sum(filt .* reshape(B[:,j],m,n)) #+ bias # element wise multiplication
    end
         return conv_out
  end

      @time out = conv_layer(A,s=1,p=3,padding="same")

  A = Float64.(reshape(1.0:81.0,9,9))
  # initialize output
  B = zeros(Int64, 2, 0)
  j = 1
  for j in 1:size(A,2) # loop cols
    out = [sum(A[:,j]),j]
    out = reshape(out,2,1)
    # append out to the B
    global B = hcat(B,out) # this grows...
  end

  function f2()
      A = reshape(1:81,9,9)
      B = zeros(Int64, 2, 9)
      for j in 1:size(A,2) # loop cols
          B[:, j] .= (sum(view(A, :, j)), j)
      end
      B
  end

    @time out = conv_layer(A,s=1,p=3,padding="same")

  heatmap(conv_out,c = :greys)
m=2
n=2
s=2
pool_type = "max"
"""
    pooling_layer(A::Array{Float64,2}; m::Int64=2,n::Int64=2, s::Int64=1, pool_type::String="max")
m x n = image pooling block size
s = stride
pool_type = "max" | "avg"
"""
  function pooling_layer(A::Array{Float64,2}; m::Int64=2,n::Int64=2, s::Int64=1, pool_type::String="max")
    M,N = size(A)
    mc = M-m+1          # no. vertical blocks (rows)
    nc = N-n+1          # no. horizontal blocks (cols)
    B = zeros(Int8, m*n, 0) # output
    # Create custom index to allow for stride
    stride_col_index = collect(1:s:nc)
    stride_row_index = collect(1:s:mc)
    @inbounds for j = 1:size(stride_col_index,1) # Loop over columns horizontal
      @inbounds for i = 1:size(stride_row_index,1) # loop over rows vertical
          @views block = A[stride_row_index[i]:stride_row_index[i]+m-1, stride_col_index[j]:stride_col_index[j]+n-1]
          shapes = rotr90(reshape(block,1,m*n))
            B = hcat(B, shapes) # Append to output matrix
        end
    end
    h = size(A,1) # rows
    w = size(A,2) # cols
    f = m
    pooling_out_h = Int64.((h-f)/s+1) # rows
    pooling_out_w = Int64.((w-f)/s+1) # cols
    pooling_out = zeros(pooling_out_h,pooling_out_w) # initialize output
    if pool_type == "max"
    for j in 1:size(B,2) # loop over column dim (horz.)
      @inbounds pooling_out[j] = maximum(B[:,j]) #+ bias # element wise multiplication
    end
  elseif pool_type == "avg"
    for j in 1:size(B,2) # loop over column dim (horz.)
      @inbounds pooling_out[j] = mean(B[:,j]) #+ bias # element wise multiplication
    end
  end
    return pooling_out
  end

  @time pool_out = pooling_layer(A,m=2,n=2,s=2)
    heatmap(testing, c = :greys)

    conv2 = conv_layer(pool_out,s=1,p=3,padding="same")
heatmap(conv2, c = :greys)
  pool_out2 = pooling_layer(conv2,m=2,n=2,s=1)
heatmap(pool_out2, c = :greys)

"""
input ---> conv_layer ---> pooling ---> conv_layer ---> pooling ---> fully connected layer
"""
# Create Fully-Connected Layer
pooled_shape = size(pool_out2)[1] * size(pool_out2)[2]
flatten = reshape(pool_out2,pooled_shape,1) # flatten pool layer

# inpout for a normal NN below

       @time A = reshape(collect(1.0:1.0:81.0),9,:)
       @time B = DataFrame(conv(A,s=1,p=0))

       plots = conv(A,s=1,p=0)
       heatmap(plots)

       testing = DataFrame(A_pad)

       CSV.write("C:/Users/Andrew.Bannerman/Desktop/Julia/test_outB1.csv", testing;delim=',')


module TestRuns

using Test
using RLEVectors

@testset begin

# numruns single
@test numruns([1,2,2,3,3]) == 3
@test numruns([1,1,2,2,3]) == 3

# numruns pair
@test numruns(["a","b","c","d","d","e","f","g"],[3,6,9,12,15,18,21,24]) == 7 # Run of 2 runvalues
@test numruns(["a","b","c","d","d","d","e","f"],[3,6,9,12,15,18,21,24]) == 6 # Run of 3 runvalues
@test numruns(["a","b","c","d","e","f","g"],[3,6,9,12,12,15,17]) == 6 # run of 2 runends
@test numruns(["a","b","c","d","e","f"],[3,6,12,12,12,15]) == 4 # run of 3 runends
@test numruns([1,2,3],[1,2,2]) == 2
@test_throws ArgumentError numruns([1,2,3],[3,2,4])

# one vector
@test ree( [1,1,2,2,3,3,4,4,4] ) == ([1,2,3,4],[2,4,6,9])
@test ree( [1,1,2,2,3,3,4,4,7] ) == ([1,2,3,4,7],[2,4,6,8,9])
@test ree( [0,1,2,2,3,3,4,4,7] ) ==([0,1,2,3,4,7],[1,2,4,6,8,9])

# Clean up an RLEVector
@test ree(["a","b","c","d","d","e","f","g"],[3,6,9,12,15,18,21,24]) == (["a","b","c","d","e","f","g"],[3,6,9,15,18,21,24]) # Run of 2 runvalues
@test ree(["a","b","c","d","d","d","e","f"],[3,6,9,12,15,18,21,24]) == (["a","b","c","d","e","f"],[3,6,9,18,21,24]) # Run of 3 runvalues
@test ree(["a","b","c","d","e","f","g"],[3,6,9,12,12,15,17]) == (["a","b","c","d","f","g"],[3,6,9,12,15,17]) # run of 2 runends
@test ree(["a","b","c","d","e","f"],[3,6,12,12,12,15]) == (["a","b","c","f"],[3,6,12,15]) # run of 3 runends
@test ree([1,2,3],[1,2,2]) == ([1,2],[1,2])

# two vectors
@test ree([9],[1]) == ([9],[1])
@test ree([4,3],[2,4]) == ([4,3],[2,4])
@test ree([4,3],[0,4]) == ([3],[4])

# scalar
@test ree(5) == ([5],[1])
@test ree([5]) == ([5],[1])

# inverse_ree
@test inverse_ree([4,5,6],[2,5,7]) == [4,4,5,5,5,6,6]
@test inverse_ree(Real[],Int64[]) == Real[]
@test_throws ArgumentError inverse_ree([],[1])

end # testset

end # module