#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5

# Short description of this file: Heat Equation

using NeuralPDE, Flux, ModelingToolkit, Optim, DiffEqFlux
using GalacticOptim
using Base
using SymPy

@parameters x t θ
@variables u(..)
@derivatives Dxx''~x
@derivatives Dt'~t

quasi_dirac = function(x,a)
    return 1/(a*pi^0.5)*exp(-(x/a)*(x/a))
end

quasi_dirac(0,0)

# 3D PDE
eq  = Dt(u(x,t,θ)) ~ Dxx(u(x,t,θ))

# Initial and boundary conditions
bcs = [u(x,0,θ) ~ quasi_dirac(x,1/100000)]
# Space and time domains
domains = [x ∈ IntervalDomain(0.0,2.0),
           t ∈ IntervalDomain(0.0,2.0)]

# Discretization
dx = 0.25; dt = 0.25
# Neural network
chain = FastChain(FastDense(2,16,Flux.σ),FastDense(16,16,Flux.σ),FastDense(16,1))

discretization = NeuralPDE.PhysicsInformedNN([dx,dt],
                                             chain,
                                             strategy = NeuralPDE.GridTraining())
pde_system = PDESystem(eq,bcs,domains,[x,t],[u])
prob = NeuralPDE.discretize(pde_system,discretization)

res = GalacticOptim.solve(prob, ADAM(0.01), progress = false; cb = cb, maxiters=1000)
phi = discretization.phi

ts,xs = [domain.domain.lower:dx:domain.domain.upper for domain in domains]

u_real  = [[analytic_sol_func(t,x)[i] for t in ts for x in xs] for i in 1:3]
u_predict  = [[phi([t,x],res.minimizer)[i] for t in ts for x in xs] for i in 1:3]
diff_u = [abs.(u_real[i] .- u_predict[i] ) for i in 1:3]

for i in 1:3
    p1 = plot(xs, ts, u_real[i], st=:surface,title = "u$i, analytic");
    p2 = plot(xs, ts, u_predict[i], st=:surface,title = "predict");
    p3 = plot(xs, ts, diff_u[i],linetype=:contourf,title = "error");
    plot(p1,p2,p3)
    savefig("sol_u$i")
end


module Quadtrees

export position,
    Boundary,
    Rect,
    Circle,
    Quadtree,
    subdivide!,
    isleaf,
    query,
    query!,
    search,
    search!

import Base: push!, insert!, append!, length, in, +, -, *, /

const Point = Tuple{Real, Real}

"""
        distance²((x₁, y₁), (x₂, y₂)) -> Real

Returns the square distance between two points.

```math
d²(p₁, p₂) = (x₂ - x₁)² + (y₂ - y₁)²
```
"""
distance²((x₁, y₁)::Point, (x₂, y₂)::Point) = (x₂ - x₁) ^ 2 + (y₂ - y₁) ^ 2
"""
        distance((x₁, y₁), (x₂, y₂)) -> Real

Returns the distance between two points.

```math
distance(p₁, p₂) = \\sqrt{(x₂ - x₁)² + (y₂ - y₁)²}
```
"""
distance(P₁::Point, P₂::Point) = √distance²(P₁, P₂)

"""
        position(value::T) -> Tuple{Real, Real}

This function is used for the [in](@ref) method for determine
the intersections between the [Boundary](@ref) objects and points.

This function is not mented to be used directly, but to be extended
for the objects `T` that your [quadtree](@ref) carries.

This function **must** return an tuple of two [Real](@ref) numbers.

For example,

```julia
mutble struct MyType
    pos::Vector{Real}
    # other fields...
end

Quadtrees.position(obj::MyType) = (obj.pos[1], obj.pos[2])
```
"""
function position(::T) where {T}
    throw(ErrorException("You need to implement the Quadtrees.position method for the type $T"))
end

position((x, y)::Point) = (x, y)

function _position(value::T) where {T}
    pos = position(value)
    @assert typeof(pos) <: Point "The return value of the position method must be
    a tuple of two Real numbers, but got type $(typeof(pos)) instead."
    return pos;
end

abstract type Boundary end
"""
Represents a 2D rectangle boundary.
"""
struct Rect <: Boundary
    x::Real
    y::Real
    w::Real
    h::Real
end

+(r::Rect, (x, y)::Point) = Rect(r.x + x, r.y + y, r.w, r.h)
-(r::Rect, (x, y)::Point) = Rect(r.x - x, r.y - y, r.w, r.h)
*(r::Rect, n::Real) = Rect(r.x, r.y, n * r.w, n * r.h)
/(r::Rect, n::Real) = Rect(r.x, r.y, r.w / n, r.h / n)

in(item::T, R::Rect) where {T} = begin
    x, y = _position(item)
    rx, ry = R.w / 2, R.h / 2

    x ≥ R.x - rx &&
    x ≤ R.x + rx &&
    y ≥ R.y - ry &&
    y ≤ R.y + ry
end

function in(A::Rect, B::Rect)
    Arx, Ary = A.w / 2, A.h / 2
    Brx, Bry = B.w / 2, B.h / 2
    return !(
        A.x - Arx > B.x + Brx ||
        A.x + Arx < B.x - Brx ||
        A.y - Ary > B.y + Bry ||
        A.y + Ary < B.y - Bry
    )
end

"""
        vertices(rect) -> Vector

Returns the vertices of a rectangle.
"""
function vertices(R::Rect)
    x, y = R.x, R.y
    rx, ry = R.w / 2, R.h / 2
    return Point[
        # Top-Left corner
        (x - rx, y - ry),
        # Top-Right corner
        (x + rx, y - ry),
        # Bottom-Right corner
        (x + rx, y + ry),
        # Botton-Left corner
        (x - rx, y + ry)
    ]
end

"""
Represents a 2D circular boundary.
"""
struct Circle <: Boundary
    x::Real
    y::Real
    r::Real
end

+(c::Circle, (x, y)::Point) = Circle(c.x + x, c.y + y, c.r)
-(c::Circle, (x, y)::Point) = Circle(c.x - x, c.y - y, c.r)
*(c::Circle, n::Real) = Circle(c.x, c.y, n * c.r)
/(c::Circle, n::Real) = Circle(c.x, c.y, c.r / n)

Base.convert(::Type{Rect}, C::Circle) = Rect(C.x, C.y, C.r, C.r)

in(item::T, C::Circle) where {T} = begin
    x, y = _position(item)
    d² = distance²((x, y), (C.x, C.y))
    return d² ≤ C.r ^ 2
end

in(A::Circle, B::Circle) = begin
    d² = distance²((B.x, B.y), (A.x, A.y))
    return d² < A.r + B.r
end

in(self::Circle, range::Rect) = begin
    dx = abs(range.x - self.x)
    dy = abs(range.y - self.y)

    r = self.r
    w, h = range.w/2, range.h/2
    edges = (dx - w)^2 + (dy - h)^2

    if dx > r + w || dy > r + h return false end
    if dx <= w || dy <= h return true end;

    return edges <= r^2
end

in(R::Rect, C::Circle) = C in R

mutable struct Quadtree{T}
    datatype::Type{T}
    capacity::Int
    bounds::Rect
    items::Vector{T}
    size::Int
    sections::Vector{Quadtree{T}}
    Quadtree{T}(init::Function) where {T} = (self = new{T}(); init(self); self)
end

function Quadtree{T}(bounds::Rect, capacity::Int) where {T}
    return Quadtree{T}() do self
        self.datatype = T
        self.capacity = capacity
        self.bounds = bounds
        self.size = 0
        self.items = Vector{T}()
        self.sections = Vector{Quadtree{T}}()
    end
end

function Quadtree{T}((x, y)::Point, width::Real, height::Real, capacity::Int) where {T}
    return Quadtree{T}(Rect(x, y, width, height), capacity)
end

function Quadtree{T}(init::Function, range::AbstractRange, bounds::Rect, capacity::Int) where {T}
    qt = Quadtree{T}(bounds, capacity)

    foreach(range) do i
        push!(qt, init(i))
    end

    return qt
end

"""
        isleaf(qt::Quadtree) -> Bool

Return true if `qt` is leaf (has no children quadtree).
"""
isleaf(qt::Quadtree) = isempty(qt.sections)

"""
        insert!(qt::Quadtree{T}, item::T) -> Bool

Insert an `item` in the quadtree `qt`.

Returns true if the item was inserted successfully, false otherwise.
"""
function insert!(qt::Quadtree{T}, item::T)::Bool where {T}
    # do not add if its not
    # in the current bounds
    !(item in qt.bounds) && return false;

    if length(qt.items) < qt.capacity
        push!(qt.items, item)
        qt.size += 1
        return true;
    else
        isleaf(qt) && subdivide!(qt)
        for section in qt.sections
            if insert!(section, item)
                qt.size += 1
                return true;
            end
        end
    end
end

"""
        subdivide!(qt::Quadtree)

Subdivides the quadtree `qt` in four sections.
"""
function subdivide!(qt::Quadtree{T}) where {T}
    w, h = qt.bounds.w, qt.bounds.h
    P = vertices(qt.bounds / 2)
    qt.sections = Quadtree{T}.(P, w/2, h/2, qt.capacity)
end

"""
        push!(qt, items...)

Inserts all `items` on `qt`.
"""
push!(qt::Quadtree{T}, items::T...) where {T} = begin
    for item in items
        insert!(qt, item)
    end
end

append!(qt::Quadtree{T}, collection::Vector{T}...) where {T} = begin
    for items in collection
        push!(qt, items...)
    end
end

"""
        length(qt::Quadtree) -> Int

Returns the number of items in `qt`.
"""
length(qt::Quadtree) = qt.size
"""
        size(qt::Quadtree) -> Int

Returns the number of items in `qt`.
"""
size(qt::Quadtree) = qt.size

"""
        query(qt::Quadtree{T}, bounds::Boundary) -> Vector{T}

Returns all items from `qt` that are in the `bounds` boundary or an empty array,
otherwise.
"""
function query(qt::Quadtree{T}, bounds::Boundary)::Vector{T} where {T}
    result = Vector{T}()

    if bounds in qt.bounds
        append!(result, filter(item -> item in bounds, qt.items))
        if !isleaf(qt)
            foreach(qt.sections) do section
                append!(result, query(section, bounds))
            end
        end
    end

    return result;
end
"""
        query!(array, quadtree, bounds)

In place version of the [query](@ref) method.
"""
query!(array::Vector{T}, qt::Quadtree{T}, range::Boundary) where {T} = empty!(array) && append!(array, query(qt, range))

"""
        search(quadtree, x, y, radius) -> Vector

Returns all items from the `quadtree` that falls on a circular boundary,
where `(x, y)` is the position of the search and `r` is the radius.
"""
function search(qt::Quadtree, x::Real, y::Real, radius::Real)
    return query(qt, Circle(x, y, radius))
end

"""
        search!(array, quadtree, x, y, radius)

In place version of the [search](@ref) method.
"""
function search!(array::Vector{T}, qt::Quadtree{T}, x::Real, y::Real, radius::Real) where {T}
    query!(array, qt, Circle(x, y, radius))
end

end # module


m=rand(2:10)
n=rand(1:10)
Z = Frobenius(m, randn(n)) #size m+n

#Special properties
@test Matrix(inv(Z)) ≈ inv(Matrix(Z))

#Matvec product
b = randn(m+n)
@test Z*b ≈ Matrix(Z)*b

#Matmul product
Y1 = Frobenius(m, randn(n)) #Another one of the same column
@test Matrix(Y1*Z) ≈ Matrix(Y1)*Matrix(Z)

n2=rand(1:(m+n-1))
Y2 = Frobenius(m+n-n2, randn(n2)) #Probably not the same column
@test Matrix(Y2*Z) ≈ Matrix(Y2)*Matrix(Z)


""" Functions for the types of random variate draws often used in individual-based epidemiological models.
"""
module Sampling

export bernoulli_sample, choose, delay_geom_sample

using Distributions: Exponential, Geometric, cdf
using Random: randsubseq
using StatsBase: sample

function rate_to_prob(x::AbstractFloat)
    0.0 <= x || throw(ArgumentError("rate $x not in [0,Inf)"))
    cdf(Exponential(), x)
end

"""
    delay_geom_sample(n::Integer, rate::AbstractFloat, dt::AbstractFloat)

Sample time steps until an event fires given a `rate` and `dt` where the waiting time
follows a Geometric distribution. The minimum delay is `1`, otherwise there could be
"instantaneous" events, which are better (more efficiently) simulated as _processes_.
"""
function delay_geom_sample(n::Integer, rate::AbstractFloat, dt::AbstractFloat)
    prob = rate_to_prob(rate * dt)
    rand(Geometric(prob), n) .+ 1
end

function delay_geom_sample(n::Integer, rate::AbstractVector{T}, dt::AbstractFloat) where {T <: AbstractFloat}
    n == length(rate) || throw(ArgumentError("number of draws must be equal to the length of 'rate'"))
    out = Vector{Int64}(undef, n)
    for i = 1:n
        prob = rate_to_prob(rate[i] * dt)
        out[i] = rand(Geometric(prob)) + 1
    end
    return out
end

"""
    bernoulli_sample(target::AbstractVector, prob::AbstractFloat)

Sample without replacement from `target` with success probability `prob`.
"""
function bernoulli_sample(target::T, prob::AbstractFloat) where {T <: Integer}
    0.0 <= prob <= 1.0 || throw(ArgumentError("probability $prob not in [0,1]"))
    if rand() < prob
        return [target]
    else
        return T[]
    end
end

function bernoulli_sample(target::AbstractVector, prob::AbstractFloat)
    randsubseq(target, prob)
end

"""
    bernoulli_sample(target::AbstractVector, rate::AbstractFloat, dt::AbstractFloat)

Sample without replacement from `target` with success probability calculated
from ``1 - e^{-\\mathrm{rate} * \\mathrm{dt}}``
"""
function bernoulli_sample(target::T, rate::AbstractFloat, dt::AbstractFloat) where {T <: Integer}
    prob = rate_to_prob(rate * dt)
    if rand() < prob
        return [target]
    else
        return T[]
    end
end

function bernoulli_sample(target::AbstractVector, rate::AbstractFloat, dt::AbstractFloat)
    randsubseq(target, rate_to_prob(rate * dt))
end

"""
    bernoulli_sample(target::AbstractVector{T}, prob::Vector{R})

Sample without replacement from `target` where each element has a unique success probability
given in the vector `prob`.
"""
function bernoulli_sample(target::AbstractVector{T}, prob::Vector{R}) where {T <: Integer, R <: AbstractFloat}
    length(target) == length(prob) || throw(ArgumentError("target and prob not of equal length"))
    samp = Vector{T}(undef, length(target))
    runif = rand(length(target))
    j = 0
    for i = 1:length(target)
        0.0 <= prob[i] <= 1.0 || throw(ErrorException("probability $(prob[i]) not in [0,1]"))
        if runif[i] < prob[i]
            j += 1
            samp[j] = target[i]
        end
    end

    # return sampled set
    if j > 0
        return samp[1:j]
    else
        return T[]
    end
end

"""
    bernoulli_sample(target::AbstractVector{T}, rate::Vector{R}, dt::AbstractFloat)

Sample without replacement from `target` where each element's success probability is calculated
from ``1 - e^{-\\mathrm{rate} * \\mathrm{dt}}``.
"""
function bernoulli_sample(target::AbstractVector{T}, rate::Vector{R}, dt::AbstractFloat) where {T <: Integer, R <: AbstractFloat}
    length(target) == length(rate) || throw(ArgumentError("target and rate not of equal length"))
    prob = map((x) -> rate_to_prob(x), rate * dt)
    samp = Vector{T}(undef, length(target))
    runif = rand(length(target))
    j = 0
    for i = 1:length(target)
        if runif[i] < prob[i]
            j += 1
            samp[j] = target[i]
        end
    end

    # return sampled set
    if j > 0
        return samp[1:j]
    else
        return T[]
    end
end

"""
    choose(target::T, K::Integer) where {T <: Integer}
    choose(target::AbstractVector{T}, K::Integer) where {T <: Integer}

Return a vector of size `K` with that number of random elements selected
without replacement from `target`.
"""
function choose(target::T, K::Integer) where {T <: Integer}
    if K > 0 # should K > 1 be an error?
        return target
    else
        return T[]
    end
end

function choose(target::AbstractVector{T}, K::Integer) where {T <: Integer}
    if K > 0
        target[sample(1:length(target), K, replace = false)]
    else
        return T[]
    end
end

end # end module


using LinearAlgebra: lu, I, ldiv!, rmul!, norm, qr, rank, nullspace
using SparseArrays: sparse, spzeros, SparseMatrixCSC, findnz
using Combinatorics: combinations
using MatrixNetworks: bfs, scomponents

export sparse_nullbasis, nullbasis

"""
    closest_permutation(mat::AbstractMatrix)

Given an `m × n` matrix `mat` where `m ≥ n`, find an assignment of rows to columns such
that each column is assigned a distinct row and the absolute values of the entries
corresponding to the assignment is as large as possible under lexicographical ordering.
Return a matrix of Bools where each column has a single true in it corresponding to its
assignment.
"""
function closest_permutation(mat::AbstractMatrix)
    nrows, ncols = size(mat)
    @assert nrows >= ncols
    ans = similar(BitArray, axes(mat)) # initialize answer as all false
    fill!(ans, false)

    ax2 = axes(mat, 2)

    # Do a first pass, ignoring collisions.
    colidxs = similar(ax2)
    for col in ax2
        @inbounds begin
            _, i = findmax(mat[:, col])
            ans[i, col] = true
            colidxs[col] = i
        end
    end

    # Prepare a BitArray `collisions` and mark collided columns in it.
    sp = sortperm(colidxs)
    sortedidxs = colidxs[sp]
    collisions = similar(BitArray, ax2)
    fill!(collisions, false)

    previndex = firstindex(sortedidxs)
    @inbounds for i in eachindex(sortedidxs)
        isequal(i, firstindex(sortedidxs)) && continue
        if isequal(sortedidxs[i], sortedidxs[previndex])
            collisions[sp[i]] = true
            collisions[sp[previndex]] = true
        end
        previndex = i
    end

    @inbounds stack = sp[collisions]
    isempty(stack) && return ans

    # reset columns in `ans` where collisions happened
    for col in stack
        @inbounds ans[:, col] .= false
    end

    # fall-back to single-threaded collision resolution on collided columns
    return _closest_permutation(mat, ans, stack)
end

function _closest_permutation(mat::AbstractMatrix, ans, stack)
    ax1 = axes(mat, 1)
    ax1l = length(ax1)
    possible_assignments = Vector{eltype(ax1)}() # row indices

    @inbounds while length(stack) > 0
        col = pop!(stack)
        resize!(possible_assignments, ax1l)
        possible_assignments .= axes(mat, 1)
        assigned = false
        while !assigned
            value, i = findmax(mat[possible_assignments, col])
            assignment = possible_assignments[i]
            collision = findfirst(ans[assignment, :])
            if collision == nothing
                ans[assignment, col] = true
                assigned = true
            elseif mat[assignment, collision] < value
                ans[assignment, col] = true
                assigned = true
                ans[assignment, collision] = false
                push!(stack, collision) # need to reassign collision
            else
                deleteat!(possible_assignments, i) # try somewhere else
            end
        end
    end
    return ans
end

"""
    match_vectors(port_vectors::AbstractMatrix{<:Number},
        eigenvectors::AbstractMatrix{<:Number})

For each column of `port_vectors` find a corresponding column of `eigenvectors` that has a
similar "shape". Match each `port_vector` to a different `eigenvector`.
"""
function match_vectors(port_vectors::AbstractMatrix{<:Number},
    eigenvectors::AbstractMatrix{<:Number})
    @assert size(port_vectors, 1) == size(eigenvectors, 1)
    @assert size(port_vectors, 2) <= size(eigenvectors, 2)
    mat = abs2.(eigenvectors' * port_vectors) # tall skinny matrix
    _, indices = findmax(closest_permutation(mat), dims=1)
    return [index[1] for index in vcat(indices...)]
end

"""
    inv_power_eigen(M::AbstractMatrix{<:Number};
        v0::AbstractVector{<:Number}=rand(Complex{real(eltype(M))}, size(M, 1)),
        λ0::Union{Number, Nothing}=nothing,
        pred::Function=(x,y) -> norm(x-y) < eps())

Use the inverse power method to find an eigenvalue and eigenvector.

# Arguments
- `v0::AbstractVector{<:Number}`: an initial value for the eigenvector. A random vector is
    used if `v0` is not given.
- `λ0::Union{Number, Nothing}`: an initial value for the eigenvalue. `v' * M * v` is used
    if `λ0` is not given.
- `pred::Function`: a function taking two successive eigenvalue estimates. The algorithm
    halts when `pred` returns true.
"""
function inv_power_eigen(M::AbstractMatrix{<:Number};
    v0::AbstractVector{<:Number}=rand(Complex{real(eltype(M))}, size(M, 1)),
    λ0::Union{Number, Nothing}=nothing,
    pred::Function=(x,y) -> norm(x-y) < eps())
    v = v0/norm(v0)
    l = length(v)
    @assert size(M) == (l, l)
    λ = (λ0 == nothing) ? v' * M * v : λ0
    done = false
    μ = λ
    A = lu(M - λ * I)
    while !done
        ldiv!(A, v)
        rmul!(v, 1/norm(v))
        μ = v' * M * v
        done = pred(λ, μ)::Bool
        λ = μ
    end
    return μ, v
end

#######################################################
# sparse nullbases
#######################################################

"""
    row_column_graph(mat::SparseMatrixCSC{<:Number, Int})

Let `mat` be a matrix with the following properties
    - for any two rows of `mat` there is at most one column on which both rows are nonzero
    - any column of `mat` has at most 2 nonzero values
The rows of `mat` form a simple graph where two rows have an edge if and only if they share
a nonzero column in `mat`. We add one additional vertex to this graph and for each column
with exactly one nonzero value, we add an edge between its nonzero row and the additional
vertex. In this graph every column with at least one nonzero value is represented by an
edge. Produce a matrix `A` where `A[i,j]` is the column shared by rows `i` and `j` if such a
column exists, and `0` otherwise. If `i == size(mat, 1) + 1`, then `A[i,j]` is chosen
arbitrarily from the columns whose unique nonzero value is in row `j`, or 0 if no such
columns exist. The case where `j == size(mat, 1) + 1` is analogous. `A` is an adjacency
matrix for the graph with extra information relating the graph to `mat`.
"""
function row_column_graph(mat::SparseMatrixCSC{<:Number, Int})
    num_vertices = size(mat, 1) + 1
    row_inds, col_inds, vals = Int[], Int[], Int[]
    matT = sparse(transpose(mat)) # SparseMatrixCSC row access is slow
    for (r1, r2) in combinations(1:(num_vertices-1), 2)
        inds = findall((matT[:, r1] .!= 0) .& (matT[:, r2] .!= 0))
        @assert length(inds) <= 1 # any two rows share at most 1 common nonzero column
        if length(inds) == 1
            c = inds[1]
            push!(row_inds, r1); push!(col_inds, r2); push!(vals, c)
            push!(row_inds, r2); push!(col_inds, r1); push!(vals, c)
        end
    end
    root = num_vertices
    connected_to_root = Set{Int}()
    for c in 1:size(mat, 2)
        inds = findall(mat[:, c] .!= 0)
        @assert length(inds) <= 2 # any column has at most 2 nonzero values
        if length(inds) == 1
            r = inds[1]
            if !(r in connected_to_root)
                push!(connected_to_root, r)
                push!(row_inds, root); push!(col_inds, r); push!(vals, c)
                push!(row_inds, r); push!(col_inds, root); push!(vals, c)
            end
        end
    end
    adj = sparse(row_inds, col_inds, vals, num_vertices, num_vertices)
    return adj
end

"""
    connected_row_column_graph(mat::SparseMatrixCSC{<:Number, Int})

Let `mat` be a matrix with the following properties
    - for any two rows of `mat` there is at most one column on which both rows are nonzero
    - any column of `mat` has at most 2 nonzero values
Produce a matrix `mat_connected` with the same nullspace as `mat` such that
`row_column_graph(mat_connected)` encodes a connected graph. Return `mat_connected` and
`row_column_graph(mat_connected)`.
"""
function connected_row_column_graph(mat::SparseMatrixCSC{<:Number, Int})
    adj = row_column_graph(mat)
    components = scomponents(adj)
    if components.number > 1 # adj is not connected
        root_component = components.map[end]
        remove_rows = Int[]
        zero_columns = Set{Int}()
        matT = sparse(transpose(mat)) # SparseMatrixCSC row access is slow
        for component in 1:components.number
            if component != root_component
                inds = findall(components.map .== component)
                num_dependent_rows = length(inds) - find_rank(matT[:, inds])
                if num_dependent_rows > 0 # some rows are redundant, remove them
                    append!(remove_rows, inds[1:num_dependent_rows])
                else # submat has full rank and trivial nullspace
                    append!(remove_rows, inds) # remove these rows
                    for i in inds # set columns in these rows to 0
                        push!(zero_columns, findall(matT[:, i] .!= 0)...)
                    end
                end
            end
        end
        num_zeros = length(zero_columns)
        zero_rows = sparse(1:num_zeros, collect(zero_columns), ones(num_zeros),
            num_zeros, size(mat, 2))
        remaining_inds = setdiff(1:size(mat, 1), remove_rows)
        mat = vcat(transpose(matT[:, remaining_inds]), zero_rows)
        adj = row_column_graph(mat) # adj is now connected
    end
    return mat, adj
end

find_rank(m::Matrix) = rank(m)
find_rank(m::SparseMatrixCSC) = rank(qr(1.0 * m))

"""
    bfs_tree(adj::AbstractMatrix{Int}, root::Int)

Produce the vertices and edges of the tree found by performing breadth first on the graph
with adjacency matrix `adj` starting from the vertex `root`. The vertices includes all
vertices of the tree besides the root.
"""
function bfs_tree(adj::AbstractMatrix{Int}, root::Int)
    dists, _, predecessors = bfs(adj, root)
    tree_vertices, tree_edges = Int[], Tuple{Int, Int}[]
    for i1 in reverse(sortperm(dists))
        i2 = predecessors[i1]
        if dists[i1] > 0 # exclude unreachables and the root
            push!(tree_vertices, i1)
            push!(tree_edges, tuple(sort([i1, i2])...))
        end
    end
    return tree_vertices, tree_edges
end

# Find the nullspace of mat by solving the equations corresponding to each of its rows. We
# solve row `tree_rows[i]`` by eliminating variable `tree_columns[i]`.
function sparse_nullbasis(mat::SparseMatrixCSC{<:Number, Int},
    tree_rows::AbstractVector{Int}, tree_columns::AbstractVector{Int})
    @assert length(tree_rows) == length(tree_columns) == size(mat, 1)
    non_tree_columns = setdiff(1:size(mat, 2), tree_columns)
    # first create mapping that preserves the non_tree_columns
    row_inds, col_inds, vals = Int[], Int[], Float64[]
    for (col_ind, row_ind) in enumerate(non_tree_columns)
        push!(row_inds, row_ind); push!(col_inds, col_ind); push!(vals, 1)
    end
    # work with transposes because SparseMatrixCSC row access is slow
    nullbasisT = sparse(col_inds, row_inds, vals, length(non_tree_columns), size(mat, 2))
    matT = sparse(transpose(mat))
    # now solve row constraints by walking the tree from the leaves to the root
    for (mat_row, mat_col) in zip(tree_rows, tree_columns)
        null_row = spzeros(size(nullbasisT, 1))
        for other_mat_col in findall(matT[:, mat_row] .!= 0)
            if other_mat_col != mat_col
                null_row += (-mat[mat_row, other_mat_col]/mat[mat_row, mat_col]) *
                    nullbasisT[:, other_mat_col]
            end
        end
        nullbasisT[:, mat_col] = null_row
    end
    return sparse(transpose(nullbasisT))
end

"""
    sparse_nullbasis(mat::SparseMatrixCSC{<:Number, Int})

Let `mat` be a matrix with the following properties
    - for any two rows of `mat` there is at most one column on which both rows are nonzero
    - any column of `mat` has at most 2 nonzero values
Produce a sparse matrix whose columns form a basis for the nullspace of `mat`.
"""
function sparse_nullbasis(mat::SparseMatrixCSC{<:Number, Int})
    mat, adj = connected_row_column_graph(mat)
    tree_rows, tree_edges = bfs_tree(adj, size(adj, 1))
    tree_columns = [adj[r1, r2] for (r1, r2) in tree_edges]
    return sparse_nullbasis(mat, tree_rows, tree_columns)
end

"""
    nullbasis(mat::AbstractMatrix{<:Number}; warn::Bool=true)

Return a `SparseMatrixCSC` whose columns form a basis for the null space of mat. If `mat`
satisfies the conditions of `sparse_nullbasis`, a sparse matrix is returned. Otherwise a
warning is issued and the `LinearAlgebra.nullspace` function is used.
"""
function nullbasis(mat::AbstractMatrix{<:Number}; warn::Bool=true)
    try
        return sparse_nullbasis(sparse(mat))
    catch AssertionError
        if warn
            @warn "sparse nullbasis not found"
        end
        return sparse(nullspace(collect(mat)))
    end
end


# Input the number of samples we wish to model (must match below)
nSamples = 4
# Make an instance of a ChronSection object for nSamples
smpl = NewChronAgeData(nSamples)
smpl.Name           = ("Sample 1", "Sample 2", "Sample 3", "Sample 4") # Et cetera
smpl.Age_14C       .= [ 6991,  7088,  7230,  7540,] # Measured ages
smpl.Age_14C_sigma .= [   30,    70,    50,    50,] # Measured 1-σ uncertainties
smpl.Height        .= [ -355,  -380,-397.0,-411.5,] # Depths below surface should be negative
smpl.Height_sigma  .= fill(0.01, nSamples) # Usually assume little or no sample height uncertainty
smpl.Age_Sidedness .= zeros(nSamples) # Sidedness (zeros by default: geochron constraints are two-sided). Use -1 for a maximum age and +1 for a minimum age, 0 for two-sided
smpl.Age_Unit = "Years BP" # Unit of measurement for ages
smpl.Height_Unit = "m" # Unit of measurement for Height and Height_sigma

# Calculate calendar age PDFs for each sample
smpl.Params = fill(NaN, length(intcal13["Age_Calendar"]), nSamples)
for i = 1:nSamples
    # The likelihood that a measured 14C age could result from a sample of
    # a given calendar age is proportional to the intergral of the product
    # of the two respective distributions
    likelihood = normproduct.(smpl.Age_14C[i], smpl.Age_14C_sigma[i], intcal13["Age_14C"], intcal13["Age_sigma"])
    likelihood ./= sum(likelihood) * intcal13["dt"] # Normalize

    samples = draw_from_distribution(likelihood, 10^4) .* maximum(intcal13["Age_Calendar"])
    smpl.Age[i] = mean(samples)
    smpl.Age_sigma[i] = std(samples)

    # Populate smpl.Params with log likelihood for each sample
    smpl.Params[:,i] = normproduct_ll.(smpl.Age_14C[i], smpl.Age_14C_sigma[i], intcal13["Age_14C"], intcal13["Age_sigma"])
end

# Run stratigraphic model
# Configure the stratigraphic Monte Carlo model
config = NewStratAgeModelConfiguration()
# If you in doubt, you can probably leave these parameters as-is
config.resolution = 5 # Same units as sample height. Smaller is slower!
config.bounding = 0.5 # how far away do we place runaway bounds, as a fraction of total section height
(bottom, top) = extrema(smpl.Height)
npoints_approx = round(Int,length(bottom:config.resolution:top) * (1 + 2*config.bounding))
config.nsteps = 100000 # Number of steps to run in distribution MCMC
config.burnin = 100000*npoints_approx # Number to discard
config.sieve = round(Int,npoints_approx) # Record one out of every nsieve steps
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# Run the stratigraphic MCMC model
@time (mdl, agedist, lldist) = StratMetropolis14C(smpl, config)

# Test that results match expectation, within some tolerance
@test isa(mdl.Age, Array{Float64,1})
@test isapprox(mdl.Age, [8319.5, 8235.7, 8154.8, 8072.3, 8025.8, 7982.2, 7938.8, 7914.7, 7891.3, 7868.5, 7845.1], atol=10)
@test isapprox(mdl.Age_025CI, [8205.0, 8045.4, 7998.9, 7975.5, 7893.6, 7859.9, 7838.6, 7802.2, 7779.3, 7761.3, 7747.5], atol=15)
@test isapprox(mdl.Age_975CI, [8403.9, 8385.1, 8345.7, 8163.4, 8150.0, 8123.4, 8028.4, 8014.9, 8002.3, 7984.7, 7957.1], atol=15)

# Data about hiatuses
nHiatuses = 2 # The number of hiatuses you have data for
hiatus = NewHiatusData(nHiatuses) # Struct to hold data
hiatus.Height         = [-371.5, -405.0 ]
hiatus.Height_sigma   = [   0.0,    0.0 ]
hiatus.Duration       = [ 100.0,  123.0 ]
hiatus.Duration_sigma = [  30.5,   20.0 ]

# Run the model. Note the additional `hiatus` arguments
@time (mdl, agedist, hiatusdist, lldist) = StratMetropolis14C(smpl, hiatus, config); sleep(0.5)

# Test that results match expectation, within some tolerance
@test isa(mdl.Age, Array{Float64,1})
@test isapprox(mdl.Age, [8336.6, 8173.9, 8120.0, 8068.4, 8031.4, 7995.5, 7957.1, 7937.2, 7860.6, 7841.4, 7825.1], atol=10)
@test isapprox(mdl.Age_025CI, [8214.9, 8027.9, 7993.8, 7975.3, 7915.4, 7884.2, 7865.3, 7832.7, 7757.4, 7744.4, 7734.2], atol=15)
@test isapprox(mdl.Age_975CI, [8407.2, 8288.1, 8253.8, 8161.2, 8150.2, 8128.7, 8062.0, 8042.4, 7968.9, 7945.2, 7922.8], atol=15)


## Banded LDLᵀ decomposition

@inline bandwidth(A::LDLt{T,Symmetric{T,M}}, k) where {T,M<:BandedMatrix{T}} = bandwidth(A.data, k)

factorize(S::Symmetric{<:Any,<:BandedMatrix}) = ldlt(S)
if VERSION ≥ v"1.4-"
    LinearAlgebra._factorize(S::Symmetric{<:Any,<:BandedMatrix}; check=false) = ldlt(S)
end

function ldlt!(F::Symmetric{T,M}) where {T,M<:BandedMatrix{T}}
    A = F.data
    n = size(A, 1)
    b = bandwidth(F)
    if F.uplo == 'L'
        @inbounds for j = 1:n
            @simd for k = max(1,j-b-1):j-1
                A[j,j] -= abs2(A[j,k])*A[k,k]
            end
            for i = j+1:min(n,j+b)
                @simd for k = max(1,j-b-1):j-1
                    A[i,j] -= A[i,k]*A[j,k]*A[k,k]
                end
                A[i,j] /= A[j,j]
            end
        end
    elseif F.uplo == 'U'
        @inbounds for j = 1:n
            @simd for k = max(1,j-b-1):j-1
                A[j,j] -= abs2(A[k,j])*A[k,k]
            end
            for i = j+1:min(n,j+b)
                @simd for k = max(1,j-b-1):j-1
                    A[j,i] -= A[k,i]*A[k,j]*A[k,k]
                end
                A[j,i] /= A[j,j]
            end
        end
    end
    return LDLt{T,Symmetric{T,M}}(F)
end

function ldlt(A::Symmetric{T,M}) where {T,M<:BandedMatrix{T}}
    S = typeof(zero(T)/one(T))
    uplo = A.uplo == 'U' ? :U : :L
    return S == T ? ldlt!(copy(A)) : ldlt!(Symmetric(BandedMatrix{S}(A), uplo))
end

function ldiv!(S::LDLt{T,Symmetric{T,M}}, B::AbstractVecOrMat{T}) where {T,M<:BandedMatrix{T}}
    require_one_based_indexing(B)
    n, nrhs = size(B, 1), size(B, 2)
    if size(S, 1) != n
        throw(DimensionMismatch("Matrix has dimensions $(size(S)) but right hand side has first dimension $n"))
    end
    A = S.data
    b = bandwidth(A, 1)
    @inbounds for l = 1:nrhs
        for j = 1:n
            @simd for k = max(1,j-b-1):j-1
                B[j,l] -= A[j,k]*B[k,l]
            end
        end
        @simd for j = 1:n
            B[j,l] /= A[j,j]
        end
        for j = n:-1:1
            @simd for k = j+1:min(n,j+b)
                B[j,l] -= A[k,j]*B[k,l]
            end
        end
    end
    B
end

function logabsdet(F::LDLt{T,Symmetric{T,M}}) where {T,M<:BandedMatrix{T}}
    it = (F.data[i,i] for i in 1:size(F, 1))
    return sum(log∘abs, it), prod(sign, it)
end

function getproperty(F::LDLt{<:Any,<:Symmetric{<:Any,<:AbstractBandedMatrix}}, d::Symbol)
    Fdata = getfield(F, :data)
    if d === :d
        return diag(Fdata)
    elseif d === :D
        return Diagonal(diag(Fdata))
    elseif d === :L
        return UnitLowerTriangular(Fdata)
    elseif d === :Lt
        return UnitUpperTriangular(Fdata)
    else
        return getfield(F, d)
    end
end


immutable QuasiHyperbolicIntercept{T} <: Model
    β₀::T
    a::T
    β::T
    δ::T
end

function Base.show(io::IO, ::Type{QuasiHyperbolicIntercept})
    return print(io, "Intercept Quasi-Hyperbolic")
end

function Base.show(io::IO, m::QuasiHyperbolicIntercept)
    @printf(
        io,
        "Intercept Quasi-Hyperbolic (β₀ = %.3f, a = %.3f, β = %.3f, δ = %.3f)",
        m.β₀,
        m.a,
        m.β,
        m.δ,
    )
end

function initial_parameters(::Type{QuasiHyperbolicIntercept})
    return Float64[0.0, log(1.0), logit(0.5), logit(0.5)]
end

canonical_parameters(m::QuasiHyperbolicIntercept) = (m.β₀, m.a, m.β, m.δ)

function QuasiHyperbolicIntercept(parameters::Vector)
    β₀ = parameters[1]
    a = exp(parameters[2])
    β = invlogit(parameters[3])
    δ = invlogit(parameters[4])
    return QuasiHyperbolicIntercept(β₀, a, β, δ)
end

@inline function predict(
    m::QuasiHyperbolicIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, β, δ = m.β₀, m.a, m.β, m.δ
    u1 = x1 * ifelse(t1 == 0.0, 1.0, β) * δ^t1
    u2 = x2 * ifelse(t2 == 0.0, 1.0, β) * δ^t2
    z = β₀ + a * (u2 - u1)
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::QuasiHyperbolicIntercept,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    β₀, a, β, δ = m.β₀, m.a, m.β, m.δ
    βt2 = ifelse(t2 == 0.0, 1.0, β)
    βt1 = ifelse(t1 == 0.0, 1.0, β)
    Δβt2 = ifelse(t2 == 0.0, 0.0, 1.0)
    Δβt1 = ifelse(t1 == 0.0, 0.0, 1.0)
    δt2 = δ^t2
    δt1 = δ^t1
    gr = (
        1.0,
        (x2 * βt2 * δt2 - x1 * βt1 * δt1) * a,
        a * (x2 * Δβt2 * δt2 - x1 * Δβt1 * δt1) * β * (1 - β),
        a * ((x2 * βt2 * t2 * δt2 - x1 * βt1 * t1 * δt1) / δ) * δ * (1 - δ),
    )
    return gr
end


using DelimitedFiles

function readmat()
	t = vec(readdlm("17.txt", String))
	n = length(t)
	m = length(t[1])
	mat = Array{Char,2}(undef, n,m)
	for i=1:n, j=1:m
		mat[i,j] = t[i][j]
	end
	mat
end

function countNeighbours(vol, i,j,k)
	s = vol[i,j,k] == '#' ? -1 : 0
	for a=max(i-1,1):min(i+1,size(vol,1)), b=max(j-1,1):min(j+1,size(vol,2)), c=max(k-1,1):min(k+1,size(vol,3))
		s += vol[a,b,c] == '#' ? 1 : 0
	end
	s
end

function countNeighbours2(vol, i,j,k,l)
	s = vol[i,j,k,l] == '#' ? -1 : 0
	for a=max(i-1,1):min(i+1,size(vol,1)), b=max(j-1,1):min(j+1,size(vol,2)), c=max(k-1,1):min(k+1,size(vol,3)), d=max(l-1,1):min(l+1,size(vol,3))
		s += vol[a,b,c,d] == '#' ? 1 : 0
	end
	s
end

# prints solution of both parts
function x()
	mat = readmat()
	n = size(mat,1)
	m = size(mat,2)

	vol = Array{Char,3}(undef, n+12,m+12,13)
	fill!(vol, '.')
	for i=1:n, j=1:m
		vol[i+6,j+6,7] = mat[i,j]
	end
	vol
	for it=1:6
		vol2 = copy(vol)
		for i=1:n+12,j=1:m+12,k=(7-it):(7+it)
			N = countNeighbours(vol,i,j,k)
			if vol[i,j,k] == '.' && N == 3
				vol2[i,j,k] = '#'
			elseif vol[i,j,k] == '#' && !(2 <= N <= 3)
				vol2[i,j,k] = '.'
			end
		end
		vol = copy(vol2)
	end

	println(count(c->c=='#', vol))

	Vol = Array{Char,4}(undef, n+12,m+12,13,13)
	fill!(Vol, '.')
	for i=1:n, j=1:m
		Vol[i+6,j+6,7,7] = mat[i,j]
	end
	for it=1:6
		Vol2 = copy(Vol)
		for i=1:n+12,j=1:m+12,k=(7-it):(7+it),l=(7-it):(7+it)
			N = countNeighbours2(Vol,i,j,k,l)
			if Vol[i,j,k,l] == '.' && N == 3
				Vol2[i,j,k,l] = '#'
			elseif Vol[i,j,k,l] == '#' && !(2 <= N <= 3)
				Vol2[i,j,k,l] = '.'
			end
		end
		Vol = copy(Vol2)
	end

	println(count(c->c=='#', Vol))

end


function testnaiveimpulsersp(N_tests::Int)

    discrepancy = 0.0
    for n = 1:N_tests
        ω = rand()*900
        h = randn(5)
        Hn = h[1] + h[2]*exp(-im*ω) + h[3]*exp(-im*2*ω) + h[4]*exp(-im*3*ω) + h[5]*exp(-im*4*ω)
        Ht = naiveimpulsersp(h,ω)

        discrepancy += abs(Hn-Ht)
    end

    println("testnaiveimpulsersp: total discrepancy = ", discrepancy)
    return discrepancy
end

function testscaletoChebyinterval(N_tests::Int)

    for n = 1:N_tests
        a = abs(randn())
        b = abs(randn())
        α = rand()

        x = (1-α)*a+α*b
        y = scaletoChebyinterval(x,a,b)
        println(y)

        if !( -1 <= y <= 1)
            println("testscaletoChebyinterval failed")
            return false
        end
    end

    println("testscaletoChebyinterval passed")
    return true
end

function testscaletoChebyinterval2(N_tests::Int, N_test_locations::Int)

    function test(N_test_locations)
        discrepancy = 0.0
        for α in collect(LinRange(0.0,1.0,N_test_locations))
            x = (1-α)*a+α*b
            y = scaletoChebyinterval(x,a,b)
            xx = scalefromChebyinterval(y,a,b)

            discrepancy += abs(xx-x)
        end

        return discrepancy
    end

    discrepancy = 0.0
    for n = 1:N_tests
        a = abs(randn())
        b = abs(randn())

        discrepancy += test(a,b, N_test_locations)
    end

    return discrepancy
end

function testget𝑤scaled(𝓧::Vector{T})::Vector{T} where T

    w = get𝑤scaled(𝓧)
    w2 = get𝑤(𝓧)

    C = 0.5
    w1 = w/C^L

    return LinearAlgebra.norm(w1-w2)
end


#*******************************************************/
#* Copyright(c) 2018 by Artelys                        */
#* This source code is subject to the terms of the     */
#* MIT Expat License (see LICENSE.md)                  */
#*******************************************************/

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#  This example demonstrates how to use Knitro to solve the following
#  simple quadratic programming problem (QP).
#
#     minimize     0.5*(x0^2+x1^2+x2^2) + 11*x0 + x2
#     subject to   -6*x2  <= 5
#                 0 <= x0
#                 0 <= x1
#                -3 <= x2 <= 2
#
#  The optimal solution is:
#     obj=-0.4861   x=[0,0,-5/6]
#
#  The purpose is to illustrate how to invoke Knitro using the C
#  language API.
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

using KNITRO
using Test

function example_qp1(; verbose=true)
    # Used to specify whether linear and quadratic objective
    # terms are loaded separately or together in this example.
    bSeparate = false

    # Create a new Knitro solver instance.
    kc = KNITRO.KN_new()

    # Illustrate how to override default options by reading from
    # the knitro.opt file.
    options = joinpath(dirname(@__FILE__), "..", "examples", "knitro.opt")
    KNITRO.KN_load_param_file(kc, options)
    kn_outlev = verbose ? KNITRO.KN_OUTLEV_ITER : KNITRO.KN_OUTLEV_NONE
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_OUTLEV, kn_outlev)

    # Initialize Knitro with the problem definition.

    # Add the variables and set their bounds.
    # Note: unset bounds assumed to be infinite.
    KNITRO.KN_add_vars(kc, 3)
    KNITRO.KN_set_var_lobnds(kc, [0.0, 0.0, -3.0])
    KNITRO.KN_set_var_upbnds(kc, 2, 2.0)

    # Add the constraint and set the bound and coefficient.
    KNITRO.KN_add_cons(kc, 1)
    KNITRO.KN_set_con_upbnd(kc, 0, 5.0)
    KNITRO.KN_add_con_linear_struct(kc, 0, 2, -6.0)

    # Set the coefficients for the objective -
    # can either set linear and quadratic objective structure
    # separately or all at once.  We show both cases below.
    # Change the value of "bSeparate" to try both cases.

    if (bSeparate)
        # Set linear and quadratic objective structure separately.
        # First set linear objective structure.
        lobjIndexVars = Int32[0, 2]
        lobjCoefs = [11.0, 1.0]
        KNITRO.KN_add_obj_linear_struct(kc, lobjIndexVars, lobjCoefs)
        # Now set quadratic objective structure.
        qobjIndexVars1 = Int32[0, 1, 2]
        qobjIndexVars2 = Int32[0, 1, 2]
        qobjCoefs = [0.5, 0.5, 0.5]
        KNITRO.KN_add_obj_quadratic_struct(kc, qobjIndexVars1, qobjIndexVars2, qobjCoefs)
    else
        # Example of how to set linear and quadratic objective
        # structure at once. Setting the 2nd variable index in a
        # quadratic term to be negative, treats it as a linear term.
        indexVars1 = Int32[0, 1, 2, 0, 2]
        indexVars2 = Int32[0, 1, 2, -1, -1]  # -1 for linear coefficients
        objCoefs = [0.5, 0.5, 0.5, 11.0, 1.0]
        KNITRO.KN_add_obj_quadratic_struct(kc, indexVars1, indexVars2, objCoefs)
    end

    # Set minimize or maximize (if not set, assumed minimize)
    KNITRO.KN_set_obj_goal(kc, KNITRO.KN_OBJGOAL_MINIMIZE)

    # Enable iteration output and crossover procedure to try to
    # get more solution precision
    KNITRO.KN_set_param(kc, KNITRO.KN_PARAM_BAR_MAXCROSSIT, 5)

    # Solve the problem.
    #
    # Return status codes are defined in "kn_defines.jl" and described
    # in the Knitro manual.
    nStatus = KNITRO.KN_solve(kc)
    nStatus, objSol, x, lambda_ =  KNITRO.KN_get_solution(kc)

    # An example of obtaining solution information.
    if verbose
        println("Knitro converged with final status = ", nStatus)
        println("  optimal objective value  = ", objSol)
        println("  optimal primal values x  = ",  x)
        println("  feasibility violation    = ", KNITRO.KN_get_abs_feas_error(kc))
        println("  KKT optimality violation = ", KNITRO.KN_get_abs_opt_error(kc))
    end
    # Delete the Knitro solver instance.
    KNITRO.KN_free(kc)

    @testset "Example QP1" begin
        @test nStatus == 0
        @test objSol ≈ -0.4861 atol=1e-4
        @test x ≈ [0., 0., -5/6] atol=1e-5
    end
end

example_qp1(; verbose=isdefined(Main, :KN_VERBOSE) ? KN_VERBOSE : true)


#=

  From srl-pp-tutorial-wasp-stockholm.pdf
  "Statistical Relational learning and Probabilistic Programming"
  by Luc De Raedt, Anton Dries, and Angelika Kimmig
  https://dtai.cs.kuleuven.be/problog/wasp17-tutorial.html

  slide 394f:
  """
  - Flipping a coin with unknown weight
  - Prior: uniform distribution on [0,1]
  - Observation: 5x heads in a row
  """

  The ProbLog model return the following which corresponds to the
  density of the probabilist
  """
   weight(c1,0.1): 3.3994697e-13
   weight(c1,0.3): 2.1679411e-06
   weight(c1,0.5): 0.0041497433
   weight(c1,0.7): 0.1317485
   weight(c1,0.9): 0.86409959         <----
   weight(c2,0.1): 3.2276726e-06
   weight(c2,0.3): 0.024109997
   weight(c2,0.5): 0.66724754         <----
   weight(c2,0.7): 0.30628626
   weight(c2,0.9): 0.0023529733
  """

  Here we observe 13 tosses and detecting the probability of throwing a head.

  Cf ~/problog/coins_learning.pl
     ~/blog/coins_learning.blog
     ~/webppl/coins_learning.wppl

  Note that this model is much simpler than the ProbLog model since ProbLog
  don't support continous distributions.
  Which makes it clear that I might have missed the point of the ProbLog model. :-)

  For a model which is much more close to the ProbLog model, see coins_learning2.jl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function coins_learning(obs)
    n = length(obs)
    p ~ Uniform(0,1)
    obs ~ filldist(flip(p),n)
    # post ~ filldist(flip(p),n)

end

# p: 0.9344
# obs = [true for _ in 1:13]
# p: 0.5984
obs = [true,false,true,true,true,true,true,false,false,true,false,false,true]
model = coins_learning(obs)
num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5), 10_000)
chns = sample(model, SMC(1000), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))


using CSPSolver
using Random
using Test

@testset "eliminate tests" begin
    constrs = [
        CageConstraint(3, [(1, 1), (2, 1)], +),
        CageConstraint(9, [(1, 2), (2, 2), (2, 3)], *),
        CageConstraint(2, [(3, 2), (2, 2)], ÷),
        CellConstraint(2, (1, 3)),
        CellConstraint(3, (3, 1)),
    ]
    kk = KenKen(constrs , 3)

    eliminate!(kk)
    fill_in!(kk)
    eliminate!(kk)
    fill_in!(kk)

    @test consistent(kk)
end

@testset "fill_in tests" begin
    constrs = [
        CageConstraint(3, [(1, 1), (2, 1)], +),
        CageConstraint(9, [(1, 2), (2, 2), (2, 3)], *),
        CageConstraint(2, [(3, 2), (2, 2)], ÷),
        CellConstraint(2, (1, 3)),
        CellConstraint(3, (3, 1)),
    ]
    kk = KenKen(constrs , 3)

    kk.domains[1, 1] = [1, 0, 0]
    kk.domains[2, 2] = [0, 1, 0]
    kk.domains[3, 3] = [0, 0, 1]

    fill_in!(kk)
    @test kk.variables[1, 1] == 1
    @test kk.variables[2, 2] == 2
    @test kk.variables[3, 3] == 3
end

@testset "backtrack 4x4 easy tests" begin

    constrs = [
        CageConstraint(5, [(1, 1), (1, 2)], +),
        CageConstraint(3, [(1, 3), (1, 4)], -),
        CageConstraint(3, [(2, 1), (3, 1)], -),
        CageConstraint(6, [(2, 2), (3, 2)], +),
        CageConstraint(2, [(2, 3), (2, 4)], -),
        CageConstraint(3, [(4, 2), (4, 3)], -),
        CageConstraint(5, [(3, 4), (4, 4)], +),
        CellConstraint(2, (3, 3)),
        CellConstraint(3, (4, 1)),
    ]
    kk = KenKen(constrs , 4)

    res = backtrack(kk, 1)
    @test !isnothing(res)
    @test res.variables == [2 3 1 4; 4 2 3 1; 1 4 2 3; 3 1 4 2]
end

@testset "backtrack 4x4 hard tests" begin

    constrs = [
        CageConstraint(7, [(1, 1), (2, 1)], +),
        CageConstraint(8, [(1, 2), (1, 3), (1, 4)], *),
        CageConstraint(6, [(2, 2), (3, 2), (3, 1)], +),
        CageConstraint(1, [(2, 3), (2, 4)], -),
        CageConstraint(1, [(4, 1), (4, 2)], -),
        CageConstraint(2, [(3, 3), (4, 3)], ÷),
        CageConstraint(4, [(3, 4), (4, 4)], +),
    ]
    kk = KenKen(constrs , 4)

    res = backtrack(kk, 1)
    @test !isnothing(res)
    @test consistent(res) && complete(res)
end

@testset "backtrack 6x6 hard test 1" begin
    constrs = [
        CageConstraint(1, [(1, 1), (2, 1)], -),
        CageConstraint(11, [(1, 2), (2, 2), (1, 3)], +),
        CageConstraint(10, [(1, 4), (1, 5), (1, 6)], *),
        CageConstraint(6, [(2, 3), (2, 4)], *),
        CageConstraint(1, [(2, 5), (2, 6)], -),
        CageConstraint(2, [(3, 1), (3, 2)], ÷),
        CageConstraint(3, [(3, 3), (3, 4)], -),
        CageConstraint(24, [(3, 5), (3, 6), (4, 6)], *),
        CageConstraint(4, [(4, 1), (4, 2)], -),
        CageConstraint(12, [(4, 3), (5, 3)], *),
        CageConstraint(20, [(4, 4), (5, 4)], *),
        CageConstraint(1, [(4, 5), (5, 5)], -),
        CageConstraint(15, [(5, 1), (5, 2), (6, 1), (6, 2)], +),
        CageConstraint(1, [(4, 5), (5, 5)], -),
        CageConstraint(11, [(6, 3), (6, 4), (6, 5)], +),
        CageConstraint(3, [(5, 6), (6, 6)], ÷),
    ]
    kk = KenKen(constrs, 6)

    res = backtrack(kk, 1)
    @test !isnothing(res)
    @test consistent(res) && complete(res)
end

@testset "backtrack 6x6 hard test 2" begin
    constrs = [
        CageConstraint(2, [(1, 1), (1, 2)], ÷),
        CageConstraint(18, [(1, 3), (1, 4), (1, 5)], *),
        CageConstraint(15, [(3, 6), (2, 6), (1, 6)], +),
        CageConstraint(10, [(2, 1), (2, 2), (3, 1)], +),
        CageConstraint(8, [(2, 3), (2, 4), (2, 5)], +),
        CageConstraint(72, [(3, 3), (3, 2), (3, 4), (4, 3)], *),
        CageConstraint(3, [(3, 5), (4, 5)], -),
        CageConstraint(2, [(4, 1), (5, 1)], -),
        CageConstraint(3, [(4, 2), (5, 2)], ÷),
        CageConstraint(15, [(4, 4), (5, 4)], *),
        CellConstraint(1, (4, 6)),
        CageConstraint(3, [(5, 3), (6, 3)], -),
        CageConstraint(96, [(5, 5), (6, 5), (6, 4)], *),
        CageConstraint(5, [(5, 6), (6, 6)], +),
    ]
    kk = KenKen(constrs, 6)

    res = backtrack(kk, 1)
    @test !isnothing(res)
    @test consistent(res) && complete(res)
end

@testset "generate kenken" begin
    Random.seed!(1)
    for k = 3:7
        kk, sol = generate_random_kenken(k)
        solution = backtrack(kk, 2)
        @test length(solution) == 1
        @test solution[1].variables == sol
    end
end

# @testset "backtrack 9x9 hard tests" begin
#     constrs = [
#         CageConstraint(16, [(1, 1), (2, 1), (3, 1), (2, 2)], +),
#         CageConstraint(2, [(1, 2), (1, 3)], -),
#         CageConstraint(72, [(1, 4), (1, 5)], *),
#         CageConstraint(11, [(1, 6), (2, 6)], +),
#         CageConstraint(2, [(1, 7), (1, 8)], ÷),
#         CageConstraint(1080, [(1, 9), (2, 7), (2, 8), (2, 9)], *),
#         CageConstraint(105, [(2, 3), (3, 3), (3, 4)], *),
#         CageConstraint(3, [(2, 4), (2, 5)], -),
#         CageConstraint(16, [(3, 2), (4, 1), (4, 2)], +),
#         CageConstraint(5, [(3, 5), (3, 6)], -),
#         CageConstraint(5, [(3, 7), (4, 7)], -),
#         CageConstraint(8, [(3, 8), (3, 9)], -),
#         CageConstraint(120, [(4, 3), (5, 3), (5, 2)], *),
#         CageConstraint(45, [(4, 4), (5, 4)], -),
#         CageConstraint(36, [(4, 5), (5, 5), (4, 6)], +),
#         CageConstraint(3, [(4, 8), (5, 8)], -),
#         CageConstraint(6, [(4, 9), (5, 9)], -),
#         CageConstraint(5, [(5, 1), (6, 1)], -),
#         CageConstraint(6, [(5, 6), (6, 6), (5, 7)], *),
#         CageConstraint(20, [(6, 2), (6, 3), (7, 3)], +),
#         CageConstraint(8, [(6, 4), (7, 4), (6, 5)], *),
#         CageConstraint(1, [(6, 7), (7, 7)], -),
#         CageConstraint(2, [(6, 8), (7, 8)], ÷),
#         CageConstraint(11, [(6, 9), (7, 9)], +),
#         CageConstraint(3, [(7, 1), (8, 1)], ÷),
#         CellConstraint(2, (7, 2)),
#         CageConstraint(21, [(7, 5), (7, 6), (8, 5)], +),
#         CageConstraint(7, [(8, 2), (8, 3)], -),
#         CageConstraint(10, [(8, 4), (9, 4), (9, 3)], +),
#         CellConstraint(5, (8, 6)),
#         CageConstraint(2, [(8, 7), (9, 7)], ÷),
#         CageConstraint(4, [(8, 8), (9, 8)], -),
#         CageConstraint(3, [(8, 9), (9, 9)], -),
#         CageConstraint(14, [(9, 1), (9, 2)], +),
#         CageConstraint(36, [(9, 6), (9, 5)], *),
#     ]
#
#     kk = GridCSP(constrs, 9)
#
#     res = backtrack(kk)
#     @test !isnothing(res)
#     display(res.variables)
# end


using Plots

function square(c,side)
  x = [ c[1]-side/2, c[1]+side/2, c[1]+side/2, c[1]-side/2, c[1]-side/2]
  y = [ c[2]-side/2, c[2]-side/2, c[2]+side/2, c[2]+side/2, c[2]-side/2]
  return x, y
end

# Generate points
p = [ 10*rand(2) for i in 1:100 ]

# Plot a square around point 1
plot(square(p[1],10),seriestype=[:shape],
            linewdith=0.5,fillalpha=0.2,
            color="green",label="")

# Images
images = [ -2, -1, 0, 1, 2 ]
for j in images, k in images
  if j != 0 || k != 0 # Do not plot (0,0) again
    xim = [ p[i][1] + j*10 for i in 1:100 ]
    yim = [ p[i][2] + k*10 for i in 1:100 ]
    scatter!(xim,yim,color="orange",label="")
  end
end

# Central plot
x = [ p[i][1] for i in 1:100 ]
y = [ p[i][2] for i in 1:100 ]
scatter!(x,y,xlabel="x",ylabel="y",label="",color="lightblue")

# Mark point one as red
scatter!([p[1][1]],[p[1][2]],color="red",label="",size=4)

plot!(xlim=[-15,25],ylim=[-15,25],
      size=(400,400),framestyle=:box)

hline!(10*images,color="black",label="") # horizontal lines
vline!(10*images,color="black",label="") # vertical lines

savefig("pbc2.pdf")


#=

  From BLOG example/csi:
  """
  Model that illustrates context-specific independence.  This is a
  parameterization of the model shown in Fig. 3 of [Milch et al.,
  AISTATS 2005]: X depends on W when U is true, but on V when U is false.

  The probability of X being true given no evidence is P(X=true) = 0.596.
  """

  Summary Statistics
  parameters      mean       std   naive_se      mcse          ess      rhat   ess_per_sec
      Symbol   Float64   Float64    Float64   Float64      Float64   Float64       Float64

           u    0.2964    0.4567     0.0046    0.0045    9896.4994    0.9999     2267.2393
           v    0.9005    0.2993     0.0030    0.0030    9747.9184    0.9999     2233.2001
           w    0.0971    0.2961     0.0030    0.0034   10124.1395    0.9999     2319.3905
           x    0.5931    0.4913     0.0049    0.0049    9952.7108    0.9999     2280.1170

  Cf ~/webppl/csi.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function csi()
    u ~ flip(0.3)
    v ~ flip(0.9)
    w ~ flip(0.1)

    x ~
        (u) ? (w == true ? flip(0.8) : flip(0.2)) :
        (v) ? flip(0.8) : flip(0.2)

end

model = csi()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(5), 10_000)
chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 100_000)

display(chns)
# display(plot(chns))


using FullyThreadedTree
using PrettyTables

max_steps = 12
error_tolerance = 1E-2
R = pi / 12.

# Define a smooth function which has a sharp gradient at some interface
interface(x) = sum(x.^2) - R^2 * (1 + 0.25 * cos(5 * atan(x[end], x[1]) + 1))^2
shape(x) = tanh(32 * pi * interface(x))

function adaptive_refinement(fun::Function, max_steps::Int, error_tolerance; plot_error::Bool = false, print_table::Bool = false, dim::Int = 2)
    tree = Tree(zeros(dim), cell_state = fun)

    collect_data = print_table || plot_error
    if collect_data
        max_error = Vector()
        integral = Vector()
        nr_of_cells = Vector()
        nr_of_active_cells = Vector()
        nr_marked_cells = Vector()
    end

    marked = [tree]
    for steps = 1:max_steps

        refine!(marked, cell_state = fun, recurse = true)
        cells = collect_cells(tree)
        if collect_data
            push!(integral, integrate(tree))
            push!(nr_of_cells, length(cells))
            push!(nr_of_active_cells, length(filter(active, cells)))
            push!(max_error, 0.)
            push!(nr_marked_cells, 0)
        end

        marked = Vector{Tree{dim}}()
        for cell ∈ filter(parent_of_active, cells)
            error_estimate = abs(cell.state - sum([child.state for child ∈ cell.children]) / (1<<dim))
            if collect_data max_error[end] = max(error_estimate, max_error[end]) end
            if error_estimate > error_tolerance
                if collect_data  nr_marked_cells[end] += 1<<dim end
                push!(marked, cell)
            end
        end

        if length(marked) == 0
            max_steps = steps
            break
        end
    end

    if print_table
        formatter = Dict(0 => (v, i) -> typeof(v) == Int ? Int(v) : round(v; digits=5))
        pretty_table(hcat(1 : max_steps, nr_of_active_cells, nr_of_cells, nr_marked_cells, max_error, integral), ["step", "# active_cells", "# cells", "# marked", "error", "integral"], tf = markdown, formatter = formatter)
    end
    if plot_error
        display(plot(nr_of_active_cells, max_error, xaxis = :log, yaxis = :log))
    end

    return tree
end

tree = adaptive_refinement(shape, max_steps, error_tolerance, print_table=true)
display(plot(tree))


import ThorinDistributions
using DynamicPolynomials, Random
Random.seed!(123)

function get_hankels_truth(s,a,b)
    m = length(s)-1
    if iseven(m)
        # m = 2n
        n = Int(m//2)
        H_low = [s[i+j-1] for i in 1:n+1,j in 1:n+1]
        H_up = [(a+b)*s[i+j] - s[i+j+1] - a*b*s[i+j-1] for i in 1:n,j in 1:n]
    else
        # m = 2n+1
        n = Int((m-1)//2)
        H_low = [s[i+j] - a*s[i+j-1] for i in 1:n+1,j in 1:n+1]
        H_up = [b*s[i+j-1] - s[i+j] for i in 1:n+1,j in 1:n+1]
    end
    return H_low,H_up
end

@testset "HankelMatrices.jl tests" begin

    m=20
    @polyvar x[1:m+1]
    H_low, H_up = vec.(ThorinDistributions.get_hankels(x,2,3))
    H_low_op, H_up_op = ThorinDistributions.vec_H.(ThorinDistributions.hankel_operators(m,2,3))

    @testset "coherence of hankel_operator" begin
        @test H_low_op*x == H_low
        @test H_up_op*x == H_up
    end

    a = rand()
    b = rand()
    a,b = min(a,b),max(a,b)

    @testset "coherence of get_hankels for random bounds" begin
        @test get_hankels_truth(x,3,4) == ThorinDistributions.get_hankels(x,3,4)
        @test get_hankels_truth(x,a,b) == ThorinDistributions.get_hankels(x,a,b)
    end

    x = randn(100)

    @testset "coherence for random values" begin
        @test get_hankels_truth(x,0,1) == ThorinDistributions.get_hankels(x,0,1)
        @test get_hankels_truth(x,a,b) == ThorinDistributions.get_hankels(x,a,b)
    end
end


"""
`randomSDR()` creates a dataset of N SDRs with sparsity between 2-20%
"""
function randomSDR(n; sparsity=nothing)
  len= 32*32
  data= falses(len,n)
  for i= 1:n
    if isnothing(sparsity)
      k= floor(Int, (rand()*0.18 + 0.02) * len)  # number of 1s
    else
      k= floor(Int, sparsity * len)  # number of 1s
    end
    @views data[ randperm(len)[1:k] ,i].= true
  end
  data
end


using FastTransforms, ToeplitzMatrices

@testset "BigFloat Toeplitz" begin
    T = Toeplitz(BigFloat[1,2,3,4,5], BigFloat[1,6,7,8,0])
    @test T*ones(BigFloat,5) ≈ [22,24,19,16,15]

    let n = 512
        r = map(BigFloat,rand(n))
        T = Toeplitz(r,[r[1];map(BigFloat,rand(n-1))])
        @test T*ones(BigFloat,n) ≈ Matrix(T)*ones(BigFloat,n)

        T = TriangularToeplitz(BigFloat[1,2,3,4,5],:L)
        @test T*ones(BigFloat,5) ≈ Matrix(T)*ones(BigFloat,5)

        r = map(BigFloat,rand(n))
        T = TriangularToeplitz(r,:L)
        @test T*ones(BigFloat,n) ≈ Matrix(T)*ones(BigFloat,n)

        T = TriangularToeplitz(BigFloat[1,2,3,4,5],:U)
        @test T*ones(BigFloat,5) ≈ Matrix(T)*ones(BigFloat,5)

        r = map(BigFloat,rand(n))
        T = TriangularToeplitz(r,:U)
        @test T*ones(BigFloat,n) ≈ Matrix(T)*ones(BigFloat,n)
    end
end


#=
Created on 13 August 2019
Last update: 14 July 2020

@author: Michiel Stock
michielfmstock@gmail.com

Functions to simulate species interaction networks with optimal transport.
=#

"""
    zerodiv(a, b)

Division which results 0 if both arguments are 0.
"""
zerodiv(a, b) = a == b == zero(a) ? zero(a) : a / b

function optimaltransport_(M::AbstractMatrix, a::AbstractVector; λ::Real=1.0)
    @assert length(a) == size(M, 1) throw(DimensionMismatch(
                "a ($(size(a))) does not match M ($(size(M)))"))
    K = exp.(λ * M)
    return Diagonal(a ./ sum(K, dims=2)[:]) * K
end

function optimaltransport_inplace(M::AbstractMatrix, a::AbstractVector, b::AbstractVector;
                        λ::Real=1.0, maxitter::Int=100, ϵ::Real=1e-5)
    Q = exp.(λ * M)
    iter = 0
    while iter < maxitter && maximum(abs.(sum(Q, dims=1)[:] - b)) > ϵ
        iter += 1
        lmul!(Diagonal(a ./ sum(Q, dims=2)[:]), Q)
        rmul!(Q, Diagonal(b ./ sum(Q, dims=1)[:]))
    end
    return Q
end

# version that is differentiable
function optimaltransport_(M::AbstractMatrix, a::AbstractVector, b::AbstractVector;
                        λ::Real=1.0, maxitter::Int=100, ϵ::Real=1e-5)
    Q = exp.(λ * M)
    iter = 0
    while iter < maxitter && maximum(abs.(sum(Q, dims=2)[:] - a)) > ϵ
        iter += 1
        Q = Diagonal(zerodiv.(a, sum(Q, dims=2)[:])) * Q
        Q = Q * Diagonal(zerodiv.(b, sum(Q, dims=1)))
    end
    return Q
end

function h(A, x, n=10)
    res = exp.(x)
    for i in 1:n
        res = A * x
    end
    return sum(res.^2)
end

Ps = [rand(10, 20) |> x -> x/sum(x) for i in 1:11]
M = randn(10, 20)
A = hcat([sum(P, dims=2) for P in Ps]...)
B = hcat([sum(P, dims=1)' for P in Ps]...)

"""Parallel version of optimal transport"""
function optimaltransport(M::AbstractMatrix, A::AbstractMatrix, B::AbstractMatrix;
                        λ::Real=1.0, maxitter::Int=10, ϵ::Real=1e-5)
    K = exp.(λ * M)
    U, V = A, B
    for itter in 1:maxitter
        U = A ./ (K * V)
        V = B ./ (K' * U)
    end
    return U, V
end

function par_kl(Ps, M, U, V)
    #Q = similar(first(Ps))
    K = exp.(M)
    l = 0.0
    for (i, P) in enumerate(Ps)
        Q = Diagonal(U[:,i]) * K * Diagonal(V[:,i])
        l += relative_entropy(P, Q)
    end
    return l
end

function optimaltransport(M::AbstractMatrix,
                a::Union{AbstractVector,Nothing}=nothing,
                b::Union{AbstractVector,Nothing}=nothing;
                λ::Real=1.0, maxitter::Int=1000, ϵ::Real=1e-5)
    if a == nothing && b == nothing
        return exp.(λ * M) |> x -> x / sum(x)
    elseif b == nothing
        return optimaltransport_(M, a, λ=λ)
    elseif a == nothing
        return optimaltransport_(M', b, λ=λ)'
    else
        return optimaltransport_(M, a, b, λ=λ, maxitter=maxitter, ϵ=ϵ)
    end
end

optimaltransport(a::AbstractVector, b::AbstractVector; λ::Real=0.0) = a * b'


using QuasiCopula, LinearAlgebra, GLM
using Random, Distributions, DataFrames, ToeplitzMatrices
using Test, BenchmarkTools

BLAS.set_num_threads(1)
Threads.nthreads()

p = 3    # number of fixed effects, including intercept

# true parameter values
Random.seed!(12345)
βtrue = rand(Uniform(-2, 2), p)
# βtrue = 0.1 * ones(p)
σ2true = [0.5]
ρtrue = [0.5]

function get_V(ρ, n)
    vec = zeros(n)
    vec[1] = 1.0
    for i in 2:n
        vec[i] = ρ
    end
    V = ToeplitzMatrices.SymmetricToeplitz(vec)
    V
end

#simulation parameters
samplesize = 10000

st = time()
currentind = 1
d = Poisson()
link = LogLink()
D = typeof(d)
Link = typeof(link)
T = Float64

gcs = Vector{GLMCopulaCSObs{T, D, Link}}(undef, samplesize)

ni = 25 #  number of observations per individual
V = get_V(ρtrue[1], ni)

# true Gamma
Γ = σ2true[1] * V

# for reproducibility I will simulate all the design matrices here
Random.seed!(12345)
X_samplesize = [randn(ni, p - 1) for i in 1:samplesize]

for i in 1:samplesize
    X = [ones(ni) X_samplesize[i]]
    η = X * βtrue
    μ = exp.(η)
    vecd = Vector{DiscreteUnivariateDistribution}(undef, ni)
    for i in 1:ni
        vecd[i] = Poisson(μ[i])
    end
    nonmixed_multivariate_dist = NonMixedMultivariateDistribution(vecd, Γ)
    # simuate single vector y
    y = Vector{Float64}(undef, ni)
    res = Vector{Float64}(undef, ni)
    rand(nonmixed_multivariate_dist, y, res)
    # push!(Ystack, y)
    V = [Float64.(Matrix(I, ni, ni))]
    # V = [ones(ni, ni)]
    gcs[i] = GLMCopulaCSObs(y, X, d, link)
end

# form model
gcm = GLMCopulaCSModel(gcs)
# precompile
println("precompiling Poisson CS fit")
gcm2 = deepcopy(gcm);
QuasiCopula.fit!(gcm2, IpoptSolver(print_level = 0, max_iter = 20));

fittime = @elapsed QuasiCopula.fit!(gcm, IpoptSolver(print_level = 5, max_iter = 100, tol = 10^-8, accept_after_max_steps = 2, limited_memory_max_history = 50, hessian_approximation = "limited-memory"))

loglikelihood!(gcm, true, true)
@show fittime
@show gcm.β
@show gcm.σ2
@show gcm.ρ
@show gcm.∇β
@show gcm.∇σ2
@show gcm.∇ρ

vcov!(gcm)
@show QuasiCopula.confint(gcm)
mseβ, mseρ, mseσ2 = MSE(gcm, βtrue, ρtrue, σ2true)
@show mseβ
@show mseσ2
@show mseρ

@test mseβ < 0.01
@test mseσ2 < 1
@test mseρ < 1

println("checking memory allocation for Poisson CS")
# logl_gradient_memory = @benchmark loglikelihood!($gcm, true, false) # this will allocate for threads
logl_gradient_memory = @benchmark loglikelihood!($gcm.data[1], $gcm.β, $gcm.ρ[1], $gcm.σ2[1], true, false)
@test logl_gradient_memory.memory == 0.0


# https://en.wikipedia.org/wiki/Day_count_convention

export yearlength, daycount, yearfraction

export Thirty360Bond, Thirty360US, ThirtyE360, ThirtyE360ISDA
export Actual365Fixed, Actual364, Actual360
export ActualActualISDA

abstract type DayCountConvention end

abstract type FixedYear{N} <: DayCountConvention end
abstract type ActualYear <: DayCountConvention end

abstract type Thirty360 <: FixedYear{360} end

struct Thirty360Bond <: Thirty360 end
struct Thirty360US <: Thirty360
    eom::Bool
end
struct ThirtyE360 <: Thirty360 end
struct ThirtyE360ISDA <: Thirty360
    maturity::Date
end

abstract type ActualFixedYear{N} <: FixedYear{N} end

struct Actual365Fixed <: ActualFixedYear{365} end
struct Actual364 <: ActualFixedYear{364} end
struct Actual360 <: ActualFixedYear{360} end

abstract type ActualActual <: ActualYear end

struct ActualActualISDA <: ActualActual end

yearlength(::FixedYear{N}, year = 0) where {N} = N

yearlength(::ActualYear, year) = isleapyear(year) ? 366 : 365

daycount(::Union{ActualFixedYear, ActualActual}, a, b) = Dates.value(b - a)

daycount(c::Thirty360, a, b) = begin
    y1, m1 = yearmonth(a)
    y2, m2 = yearmonth(b)

    d1, d2 = adjustdays(c, a, b)

    360 * (y2 - y1) + 30 * (m2 - m1) + (d2 - d1)
end

adjustdays(::Thirty360Bond, a, b) = begin
    d1, d2 = day(a), day(b)

    d1 = min(d1, 30)
    d1 == 30 && (d2 = min(d2, 30))

    d1, d2
end

adjustdays(c::Thirty360US, a, b) = begin
    m1, d1 = monthday(a)
    m2, d2 = monthday(b)

    if c.eom && m1 == 2 && d1 == day(lastdayofmonth(a))
        m2 == 2 && d2 == day(lastdayofmonth(b)) && (d2 = 30)
        d1 = 30
    end

    d1 = min(d1, 30)
    d1 == 30 && (d2 = min(d2, 30))

    d1, d2
end

adjustdays(::ThirtyE360, a, b) = begin
    d1, d2 = day(a), day(b)

    d1 = min(d1, 30)
    d2 = min(d2, 30)

    d1, d2
end

adjustdays(c::ThirtyE360ISDA, a, b) = begin
    y1, m1, d1 = yearmonthday(a)
    y2, m2, d2 = yearmonthday(b)

    d1 == day(lastdayofmonth(a)) && (d1 = 30)
    d2 == day(lastdayofmonth(b)) && (m2 != 2 || b != c.maturity) && (d2 = 30)

    d1, d2
end

yearfraction(c::Union{Thirty360, ActualFixedYear}, a, b, s::Type{S} = Float64) where {S<:Real} = convert(S, daycount(c, a, b)) / yearlength(c)

yearfraction(c::ActualActualISDA, a, b, ::Type{S} = Float64) where {S<:Real} = begin
    firstyear = 1 - yearfraction(c, a)
    betweenyears = year(b) - year(a) - 1
    lastyear = yearfraction(c, b)

    convert(S, firstyear + betweenyears + lastyear)
end

yearfraction(c::ActualActualISDA, date, ::Type{S} = Float64) where {S<:Real} = (dayofyear(date) - 1) // yearlength(c, date)


export EpsilonGreedySelector, get_prob

using StatsBase: sample
using .Utils: findallmax

"""
    EpsilonGreedySelector{T}(;kwargs...)

> Epsilon-greedy strategy: The best lever is selected for a proportion `1 - epsilon` of the trials, and a lever is selected at random (with uniform probability) for a proportion epsilon . [Multi-armed_bandit](https://en.wikipedia.org/wiki/Multi-armed_bandit)

Two kinds of epsilon-decreasing strategy are implmented here (`linear` and `exp`).

> Epsilon-decreasing strategy: Similar to the epsilon-greedy strategy, except that the value of epsilon decreases as the experiment progresses, resulting in highly explorative behaviour at the start and highly exploitative behaviour at the finish.  - [Multi-armed_bandit](https://en.wikipedia.org/wiki/Multi-armed_bandit)

# Keywords

- `T::Symbol`: defines how to calculate the epsilon in the warmup steps. Supported values are `linear` and `exp`.
- `step::Int = 1`: record the current step.
- `ϵ_init::Float64 = 1.0`: initial epsilon.
- `warmup_steps::Int=0`: the number of steps to use `ϵ_init`.
- `decay_steps::Int=0`: the number of steps for epsilon to decay from `ϵ_init` to `ϵ_stable`.
- `ϵ_stable::Float64`: the epsilon after `warmup_steps + decay_steps`.

# Example

```julia
s = EpsilonGreedySelector{:linear}(ϵ_init=0.9, ϵ_stable=0.1, warmup_steps=100, decay_steps=100)
plot([RL.get_ϵ(s, i) for i in 1:500], label="linear epsilon")
```
![](../assets/img/linear_epsilon_greedy_selector.png)

```julia
s = EpsilonGreedySelector{:exp}(ϵ_init=0.9, ϵ_stable=0.1, warmup_steps=100, decay_steps=100)
plot([RL.get_ϵ(s, i) for i in 1:500], label="exp epsilon")
```
![](../assets/img/exp_epsilon_greedy_selector.png)
"""
Base.@kwdef mutable struct EpsilonGreedySelector{T} <: AbstractDiscreteActionSelector
    ϵ_stable::Float64
    ϵ_init::Float64 = 1.0
    warmup_steps::Int = 0
    decay_steps::Int = 0
    step::Int = 1
end

"""
    EpsilonGreedySelector(ϵ) -> EpsilonGreedySelector{:linear}(; ϵ_stable = ϵ)
"""
EpsilonGreedySelector(ϵ) = EpsilonGreedySelector{:linear}(; ϵ_stable = ϵ)

function get_ϵ(s::EpsilonGreedySelector{:linear}, step)
    if step <= s.warmup_steps
        s.ϵ_init
    elseif step >= (s.warmup_steps + s.decay_steps)
        s.ϵ_stable
    else
        steps_left = s.warmup_steps + s.decay_steps - step
        s.ϵ_stable + steps_left / s.decay_steps * (s.ϵ_init - s.ϵ_stable)
    end
end

function get_ϵ(s::EpsilonGreedySelector{:exp}, step)
    if step <= s.warmup_steps
        s.ϵ_init
    else
        n = step - s.warmup_steps
        s.ϵ_stable + (s.ϵ_init - s.ϵ_stable) * exp(-1.0 * n / s.decay_steps)
    end
end

get_ϵ(s::EpsilonGreedySelector) = get_ϵ(s, s.step)

"""
    (s::EpsilonGreedySelector)(values; step) where T

!!! note
    If multiple values with the same maximum value are found.
    Then a random one will be returned!

    `NaN` will be filtered unless all the values are `NaN`.
    In that case, a random one will be returned.
"""
function (s::EpsilonGreedySelector)(values)
    ϵ = get_ϵ(s)
    s.step += 1
    rand() > ϵ ? sample(findallmax(values)[2]) : rand(1:length(values))
end

"""
    get_prob(s::EpsilonGreedySelector, values)

Return the probability of selecting each action given the estimated `values` of each action.
"""
function get_prob(s::EpsilonGreedySelector, values)
    ϵ, n = get_ϵ(s), length(values)
    probs = fill(ϵ / n, n)
    max_val_inds = findallmax(values)[2]
    for ind in max_val_inds
        probs[ind] += (1 - ϵ) / length(max_val_inds)
    end
    probs
end


using Statistics, StatsBase, Random, LinearAlgebra

function ARMAmodel(θ, rndseed=1234)
    Random.seed!(rndseed)
    N = 300
    burnin = 50
    α = θ[1]
    β = θ[2]
    σ = sqrt(θ[3])
    data = zeros(N)
    xlag = 0.0
    flag = 0.0
    @inbounds for i = 1:(N+burnin)
        f = σ*randn()
        x = α*xlag + f - β*flag
        xlag = x
        flag = f
        if i > burnin
            data[i-burnin] = x
        end
    end
    data
end

function auxstat(θ, reps)
    data = [ARMAmodel(θ, rand(1:Int64(1e12))) for i = 1:reps]  # reps draws of data
    auxstat.(data)
end

@views function auxstat(x)
    s = std(x)
    c = cor(x[2:end,:],x[1:end-1])
    n = size(x,1)
    b, varb, u, junk1, rsq  = ols(x[2:end,:], x[1:end-1,:])
    ϕ, varϕ, u, junk2, rsq2 = ols(u[2:end,:], u[1:end-1,:])
    σ = std(u)
    sqrt(300.0).*vcat(s, c, b, varb, rsq, ϕ, varϕ, rsq2, σ, pacf(x,collect(1:4)))
end

function TrueParameters()
    [0.95, 0.5, 1.0]
end

function PriorSupport()
    lb = [-0.99, -0.99, 0.0001]
    ub = [0.99, 0.99, 4.0]
    lb,ub
end

function InSupport(θ)
    lb,ub = PriorSupport()
    all(θ .>= lb) & all(θ .<= ub)
end

function PriorMean()
    lb, ub = PriorSupport()
    (ub - lb) ./ 2.0
end

function PriorDraw()
    lb, ub = PriorSupport()
    (ub - lb) .* rand(3) + lb
end

# prior just checks that we're in the bounds
function Prior(theta)
    lb, ub = PriorSupport()
    a = 0.0
    if(all((theta .>= lb) .& (theta .<= ub)))
        a = 1.0
    end
    return a
end

# from Econometrics
# compute ols coefficients, fitted values, and errors
function lsfit(Y, X)
    beta = X\Y
    fit = X*beta
    errors = Y - fit
    return beta, fit, errors
end

function ols(y, x)
    n,k = size(x)
    b, fit, e = lsfit(y,x)
    df = n-k
    sigsq = (e'*e/df)[1,1]
    xx_inv = inv(x'*x)
    ess = (e' * e)[1,1]
    xe = x.*e
    varb = xx_inv*xe'xe*xx_inv
    seb = sqrt.(diag(varb))
    seb = seb.*(seb.>1e-16) # round off to zero when there are restrictions
    t = b ./ seb
    tss = y .- mean(y)
    tss = (tss'*tss)[1,1]
    rsq = (1.0 - ess / tss)
    return b, varb, e, ess, rsq
end


include("plotter_setup.jl")

using LambertW
using Statistics

function theoretical_p0(p0)
    x = (1-p0)/exp(1)/p0
    lambertw(x, 0)/lambertw(x, -1)
end

## p0 theoretical only

data_x = (1.:.1:10000)
data_p = theoretical_p0.(data_x)

cla()
figure(figsize=[4, 2.])
plot(data_x, data_p, "-k")
ylim(0, 1)
xlim(.9, 11000)
xscale("log")
xlabel(L"p_0")
ylabel("succ. prob. lower bound")
yticks([0, 0.5, 1.])
savefig("plots/hiding_search/p0_theoretical.pdf", bbox_inches="tight", metadata = Dict("CreationDate" => nothing))

## p0 real scenarios
p0s = .5:.5:2.
p(p0, n) = p0*log(n)/n
repeating = 40

data_n = 100:100:2000

fig, axes = subplots(1, length(p0s), figsize=[6, 2.])
println(p0s)
cla()
for (p0, ax) = zip(p0s, axes)
    data = npzread("hiding_search/sparse_ER_probs_and_bounds_$p0.npz")
    data_real_p = zeros(length(data_n))
    data_real_p_up = zeros(length(data_n))
    data_real_p_down = zeros(length(data_n))
    data_real_bound = zeros(length(data_n))
    data_real_bound_up = zeros(length(data_n))
    data_real_bound_down = zeros(length(data_n))
    for m =1:length(data_n)
        data_b = data[m,:,1]
        data_p = data[m,:,2]

        data_real_bound[m] = mean(data_b)
        data_real_bound_up[m] = maximum(data_b)
        data_real_bound_down[m] = minimum(data_b)

        data_real_p[m] = mean(data_p)
        data_real_p_up[m] = maximum(data_p)
        data_real_p_down[m] = minimum(data_p)
    end

    if p0 > 1
        ax.plot(data_n, data_real_bound, "-b", markersize=.8, linewidth=.8)
        ax.fill_between(data_n, data_real_bound_up, data_real_bound_down, color="b", alpha=.2)
    end
    ax.plot(data_n, data_real_p, ":r", markersize=.8, linewidth=.8)
    ax.fill_between(data_n, data_real_p_up, data_real_p_down, color="r", alpha=.2)

    if p0 > 1
        ax.plot([minimum(data_n), maximum(data_n)], fill(theoretical_p0(p0), 2), "--k", linewidth=.8)
    end
    ax.set_title(latexstring("p_0=$p0"))
    ax.set(xlabel=L"n", ylim=[0,1], yticks=0:0.5:1.)
    ax.label_outer()
end
axes[1].set(ylabel="succ. prob.")
savefig("plots/hiding_search/real_statistics_p0.pdf", bbox_inches="tight", metadata = Dict("CreationDate" => nothing))


# This file includes portions that were formerly part of Julia in modified form.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module PrimesIterator

export ModulePrimesIterator
export Primes, PrimePi, PrimeSieve

"""

* Primes, PrimePi, PrimeSieve
"""
const ModulePrimesIterator = ""

# Primes generating functions

const wheel = [4, 2, 4, 2, 4, 6, 2, 6]
const wheel_primes = [7, 11, 13, 17, 19, 23, 29, 31]
const wheel_indices = [
    0,    0,    0,    0,    0,    0,    0,    1,    1,
    1,    1,    2,    2,    3,    3,    3,    3,    4,
    4,    5,    5,    5,    5,    6,    6,    6,    6,
    6,    6,    7,    7
]

@inline function wheel_index(n)
    d, r = Base.divrem(n - 1, 30)
    return 8d + wheel_indices[r+2]
end

@inline function wheel_prime(n)
    d, r = (n - 1) >>> 3, (n - 1) & 7
    return 30d + wheel_primes[r+1]
end

#Internal function
function mask(limit::Int)
    limit < 7 && throw(ArgumentError("The condition limit ≥ 7 must be met."))
    n = wheel_index(limit)
    m = wheel_prime(n)
    sieve = ones(Bool, n)
    @inbounds for i = 1:wheel_index(isqrt(limit))
        if sieve[i]
            p = wheel_prime(i)
            q = p^2
            j = (i - 1) & 7 + 1
            while q ≤ m
                sieve[wheel_index(q)] = false
                q += wheel[j] * p
                j = j & 7 + 1
            end
        end
    end
    return sieve
end

#Internal function
function mask(lo::Int, hi::Int)
    7 ≤ lo ≤
    hi || throw(ArgumentError("The condition 7 ≤ lo ≤ hi must be met."))
    lo == 7 && return mask(hi)
    wlo, whi = wheel_index(lo - 1), wheel_index(hi)
    m = wheel_prime(whi)
    sieve = ones(Bool, whi - wlo)
    hi < 49 && return sieve
    small_sieve = mask(isqrt(hi))
    @inbounds for i = 1:length(small_sieve)  # don't use eachindex here
        if small_sieve[i]
            p = wheel_prime(i)
            j = wheel_index(2 * div(lo - p - 1, 2p) + 1)
            q = p * wheel_prime(j + 1)
            j = j & 7 + 1
            while q ≤ m
                sieve[wheel_index(q)-wlo] = false
                q += wheel[j] * p
                j = j & 7 + 1
            end
        end
    end
    return sieve
end

"""

Return the prime sieve, as a `BitArray`, of the positive integers (from `lo`, if specified) up to `hi`. Useful when working with either primes or composite numbers.
"""
function PrimeSieve(lo::Int, hi::Int)
    0 < lo ≤
    hi || throw(ArgumentError("The condition 0 < lo ≤ hi must be met."))
    sieve = falses(hi - lo + 1)
    lo ≤ 2 ≤ hi && (sieve[3-lo] = true)
    lo ≤ 3 ≤ hi && (sieve[4-lo] = true)
    lo ≤ 5 ≤ hi && (sieve[6-lo] = true)
    hi < 7 && return sieve
    wheel_sieve = mask(max(7, lo), hi)
    lsi = lo - 1
    lwi = wheel_index(lsi)
    @inbounds for i = 1:length(wheel_sieve)   # don't use eachindex here
        sieve[wheel_prime(i + lwi)-lsi] = wheel_sieve[i]
    end
    return sieve
end

function PrimeSieve(lo::T, hi::T) where {T<:Integer}
    lo ≤ hi ≤ typemax(Int) && return PrimeSieve(Int(lo), Int(hi))
    throw(ArgumentError("Both endpoints of the interval to sieve must be ≤ $(typemax(Int)), got $lo and $hi."))
end

PrimeSieve(limit::Int) = PrimeSieve(1, limit)

function PrimeSieve(n::Integer)
    n ≤ typemax(Int) && return PrimeSieve(Int(n))
    throw(ArgumentError("Requested number of primes must be ≤ $(typemax(Int)), got $n."))
end

"""

Return the collection of the prime numbers (from `lo`, if specified) up to `hi`.
"""
function Primes(lo::Int, hi::Int)
    lo ≤ hi || throw(ArgumentError("The condition lo ≤ hi must be met."))
    list = Int[]
    lo ≤ 2 ≤ hi && push!(list, 2)
    lo ≤ 3 ≤ hi && push!(list, 3)
    lo ≤ 5 ≤ hi && push!(list, 5)
    hi < 7 && return list
    lo = max(2, lo)
    sizehint!(
        list,
        5 +
        floor(Int, hi / (log(hi) - 1.12) - lo / (log(lo) - 1.12 * (lo > 7)))
    )
    sieve = mask(max(7, lo), hi)
    lwi = wheel_index(lo - 1)
    # don't use eachindex here
    @inbounds for i = 1:length(sieve)
        sieve[i] && push!(list, wheel_prime(i + lwi))
    end
    return list
end

Primes(n::Int) = Primes(1, n)

"""

Return the number of primes ``≤ n``.
"""
PrimePi(n::Int) = length(Primes(1, n))

#START-TEST-########################################################

using Test

function test()
    @testset "Primes" begin
        @test PrimePi(2^14) == 1900
        @test PrimePi(10^6) == 78498
    end
end

function demo()
    println(Primes(0, 100))
    println(Primes(90, 100))
    println(PrimePi(10000))
end

"""

primepi = PrimePi(1000000)
0.005080 seconds (5 allocations: 876.141 KiB)
"""
function perf()
    @time primepi = PrimePi(1000000)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module

#=

[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]
[97]
1229

  0.005080 seconds (5 allocations: 876.141 KiB)

=#


using OVERTVerify
using LazySets
using Dates
using JLD2
#ENV["JULIA_DEBUG"] = Main

controller_name = ARGS[1]
controller = "nnet_files/jmlr/car_"*controller_name*"_controller.nnet"
println("Controller is: ", controller)

query = OvertQuery(
	SimpleCar,  # problem
	controller,    # network file
	Id(),      	# last layer activation layer Id()=linear
	"MIP",     	# query solver, "MIP" or "ReluPlex"
	10,        	# ntime
	0.2,       	# dt
	-1,        	# N_overt
	)

input_set = Hyperrectangle(low=[9.5, -4.5, 2.1, 1.5], high=[9.55, -4.45, 2.11, 1.51])
concretization_intervals = [5, 5]
t1 = Dates.time()
concrete_state_sets, symbolic_state_sets, concrete_meas_sets, symbolic_meas_sets = symbolic_reachability_with_concretization(query, input_set, concretization_intervals)
t2 = Dates.time()
dt = (t2-t1)
print("elapsed time= $(dt) seconds")

# clean up sets
init_set0, reachable_state_sets = clean_up_sets(concrete_state_sets, symbolic_state_sets, concretization_intervals, dims=[1,2])

# In this example, our property is the following:
# We want the car to reach the box [-.6, .6] [-.2,.2]
# at SOME point in the time history
# constraint 1  x1 >= -0.6  --> -x1 <= 0.6
c1 = HalfSpace([-1.0, 0.0, 0.0, 0.0], 0.6)
# constraint 2 x1 <= 0.6
c2 = HalfSpace([1.0, 0.0, 0.0, 0.0], 0.6)
# constraint 3  x2 >= -0.2   -> -x2 <= 0.2
c3 = HalfSpace([-1.0, 0.0, 0.0, 0.0], 0.2)
# constraint 4  x2 <= 0.2
c4 = HalfSpace([1.0, 0.0, 0.0, 0.0], 0.2)
goal_set = HPolyhedron([c1, c2, c3, c4])
t1 = time()
goal_reached_steps = reachable_state_sets .⊆ Ref(goal_set)
goal_reached = any(goal_reached_steps)
dt_check = time() - t1
println("Goal reached: $goal_reached")
println("Goal reached at step: $(findfirst(goal_reached_steps))")

JLD2.@save "src/examples/jmlr/data/car_reachability_"*string(controller_name)*"_controller_data.jld2" query input_set concretization_intervals goal_set concrete_state_sets symbolic_state_sets concrete_meas_sets symbolic_meas_sets reachable_state_sets dt goal_reached goal_reached_steps dt_check


using Test, MLToolkit
using Distributions: Normal
using MLToolkit.DistributionsX: test_stat
using Statistics: mean, var
using Flux: gpu

@testset "Noise" begin
    dim1, dim2 = 10, 5
    n_samples, atol = 5_000, 0.01

    @testset "UniformNoise" begin
        noise = UniformNoise(dim1) |> gpu
        @test size(noise) == (dim1,)
        @test size(rand(noise)) == size(noise)
        @test size(rand(noise, dim2)) == (size(noise)..., dim2)
        @test size(rand(UniformNoise(dim1, dim2))) == (dim1, dim2)

        x = test_stat(mean, noise, n_samples, atol * prod(size(noise)))
        test_stat(var, noise, x, atol * prod(size(noise)))

        @test logpdf(noise, x) ≈ log.(ones(dim1, n_samples) / 2)
    end

    @testset "GaussianNoise" begin
        noise = GaussianNoise(dim1) |> gpu
        @test size(noise) == (dim1,)
        @test size(rand(noise)) == size(noise)
        @test size(rand(noise, dim2)) == (size(noise)..., dim2)
        @test size(rand(GaussianNoise(dim1, dim2))) == (dim1, dim2)

        x = test_stat(mean, noise, n_samples, atol * prod(size(noise)))
        test_stat(var, noise, x, atol * prod(size(noise)))

        @test logpdf(noise, x) ≈ map(xi -> logpdf(Normal(0, 1), xi), x)
    end
end


# ---
# title: 34. Find First and Last Position of Element in Sorted Array
# id: problem34
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Binary Search
# link: <https://leetcode.com/problems/find-first-and-last-position-of-element-in-sorted-array/description/>
# hidden: true
# ---
#
# Given an array of integers `nums` sorted in ascending order, find the starting
# and ending position of a given `target` value.
#
# If `target` is not found in the array, return `[-1, -1]`.
#
# **Follow up:**  Could you write an algorithm with `O(log n)` runtime
# complexity?
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [5,7,7,8,8,10], target = 8
#     Output: [3,4]
#
#
# **Example 2:**
#
#
#
#     Input: nums = [5,7,7,8,8,10], target = 6
#     Output: [-1,-1]
#
#
# **Example 3:**
#
#
#
#     Input: nums = [], target = 0
#     Output: [-1,-1]
#
#
#
#
# **Constraints:**
#
#   * `0 <= nums.length <= 105`
#   * `-109 <= nums[i] <= 109`
#   * `nums` is a non-decreasing array.
#   * `-109 <= target <= 109`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


# Adapted from https://ccrma.stanford.edu/~jos/filters/Partial_Fraction_Expansion_residuez_m.html

function [r, p, f, m] = residuez(B, A, tol=0.001)
NUM = vec(B); DEN = vec(A);
# Matlab's residue does not return m (implied by p):

[r,p,f,m]=residue(conj(flipdim(NUM,1)),conj(flipdim(DEN,1)),tol);
p = 1 ./ p;
r = r .* ((-p) .^m);
if f, f = conj(flipdim(f,1)); end


module SavitskyGolayFilters

import DSP: conv

vandermonde(x, k) = x .^ transpose(0:k)
fit_matrix(n, k) = vandermonde(-floor(n/2):floor(n/2), k)

"""
    savitskygolayfilter(n,k,d=0)

Compute the Savitsky-Golay filter of length `n`, degree `k` and derivative
order `d`.

See also: [`savitskygolay`](@ref)
"""
savitskygolayfilter(n,k,d=0) = begin
    M = fit_matrix(n, k)
    B = inv(transpose(M)*M)*transpose(M)
    B[d+1,end:-1:1]
end

"""
    savitskygolay(x, n, k, d=0)
    savitskygolay(x, filter)

Apply a Savitsky-Golay filter on `x`. The applied filter is either the already
calculated `filter` or a newly calculated filter of length `n`, degree `k`
and derivative order `d`.

# Examples
```jldoctest
julia> using Random

julia> x = -3π:0.1:3π;

julia> y = sin.(x) .+ randn(length(x));

julia> filtered = savitskygolay(y, 10, 2);

julia> size(filtered)
(158,)

julia> size(y)
(158,)

julia> filtered_derivative = savitskygolay(y, 10, 2, 1);

julia> size(filtered_derivative)
(158,)

```

See also: [`savitskygolayfilter`](@ref)
"""
function savitskygolay end

savitskygolay(x, n, k, d=0) = savitskygolay(x, savitskygolayfilter(n,k,d))

savitskygolay(x, filter::AbstractArray) = begin
    len_filter = length(filter)
    border_repeat = floor(Int, len_filter/2)
    conv(
         reverse(filter),
        [repeat([x[1]],border_repeat);x;repeat([x[end]],border_repeat)],
       )[len_filter:end-len_filter+1]
end

export savitskygolayfilter, savitskygolay

end