# Euler's Totient function, φ(n) [sometimes called the phi function], is used
# to determine the number of numbers less than n which are relatively prime to
# n. For example, as 1, 2, 4, 5, 7, and 8, are all less than nine and relatively
# prime to nine, φ(9)=6.
#
# n   Relatively Prime  φ(n)   n/φ(n)
# --------------------------------------
# 2   1                 1      2
# 3   1,2               2      1.5
# 4   1,3               2      2
# 5   1,2,3,4           4      1.25
# 6   1,5               2      3
# 7   1,2,3,4,5,6       6      1.1666...
# 8   1,3,5,7           4      2
# 9   1,2,4,5,7,8       6      1.5
# 10  1,3,7,9           4      2.5
#
# It can be seen that n=6 produces a maximum n / φ(n) for n ≤ 10.
#
# Find the value of n ≤ 1,000,000 for which n / φ(n) is a maximum.

using ProjectEulerSolutions

# Brute force solution, considering every value less than n.
function p069solution_bruteforce(n::Integer=1_000)::Integer

    max_val = 0.0
    max_index = 0
    for i in 2:n
        test_val = i / totient(i)
        if max_val < test_val
            max_val = test_val
            max_index = i
        end
    end

    return max_index
end

# Only consider values that are likely candidates, which are primodial numbers.
# Primodial numbers are the product of consecutive primes starting at 2, so
# 2, 6, 30, 210, etc.  We only consider primordial numbers up to that who's
# product more than n since they have maximal n / φ(n) values.
function p069solution_factors(n::Integer=1_000)::Integer

    primes = sieve_eratosthenes(30)
    index = 2
    p = primes[1]
    while p * primes[index] < n
        p *= primes[index]
        index += 1
    end
    return p
end

p069 = Problems.Problem(Dict("Brute force" => p069solution_bruteforce,
                             "Factors" => p069solution_factors))

Problems.benchmark(p069, 1_000_000)


"""
	mask_elements(img, threshold, connectivity)
Remove every element that is below the `threshold` and
adjacent to more than the `connectivity` amount of other
elements similarly below the threshold. Returns a mask where
all elements are 1 for elements that should be kept based on
`threshold` and `connectivity` criteria.

# Arguments
- img: array which will be used to determine mask
- threshold: cutoff value for elements that should be
    removed
- connectivity: minimum number of adjacent elements that are
    required to remove an element
"""
function mask_elements(img, threshold, connectivity)
    mask = img .< threshold
    label = Images.label_components(mask)
    n = length(unique(label))
    indices = []
    for lbl in 0:(n - 1)
        lbl_mask = map(x -> x == lbl, label)
        num = count(lbl_mask)
        if num < connectivity
            rmv_indices = findall(x -> x == lbl, label)
            push!(indices, rmv_indices)
        end
    end
    for i in 1:length(indices)
        idx = indices[i]
        mask[idx] .= 0
    end
    return iszero.(mask)
end


@doc raw"""
    Otsu <: AbstractImageBinarizationAlgorithm
    Otsu()

    binarize([T,] img, f::Otsu)
    binarize!([out,] img, f::Otsu)

Under the assumption that the image histogram is bimodal the binarization
threshold is set so that the resultant between-class variance is maximal.

# Output

Return the binarized image as an `Array{Gray{T}}` of size `size(img)`. If
`T` is not specified, it is inferred from `out` and `img`.

# Details

Let ``f_i`` ``(i=1 \ldots I)`` denote the number of observations in the
``i``th bin of the image histogram. Then the probability that an observation
belongs to the ``i``th bin is given by  ``p_i = \frac{f_i}{N}`` (``i = 1,
\ldots, I``), where ``N = \sum_{i=1}^{I}f_i``.

The choice of a threshold ``T`` partitions the data into two categories, ``C_0``
and ``C_1``. Let
```math
P_0(T) = \sum_{i = 1}^T p_i \quad \text{and} \quad P_1(T) = \sum_{i = T+1}^I p_i
```
denote the cumulative probabilities,
```math
\mu_0(T) = \sum_{i = 1}^T i \frac{p_i}{P_0(T)} \quad \text{and} \quad \mu_1(T) = \sum_{i = T+1}^I i \frac{p_i}{P_1(T)}
```
denote the means, and
```math
\sigma_0^2(T) = \sum_{i = 1}^T (i-\mu_0(T))^2 \frac{p_i}{P_0(T)} \quad \text{and} \quad \sigma_1^2(T) = \sum_{i = T+1}^I (i-\mu_1(T))^2 \frac{p_i}{P_1(T)}
```
denote the variances of categories ``C_0`` and ``C_1``, respectively.
Furthermore, let
```math
\mu = P_0(T)\mu_0(T) + P_1(T)\mu_1(T),
```
represent the overall mean,
```math
\sigma_b^2(T) = P_0(T)(\mu_0(T) - \mu)^2 + P_1(T)(\mu_1(T) - \mu)^2,
```
the between-category variance, and
```math
\sigma_w^2(T) = P_0(T) \sigma_0^2(T) +  P_1(T)\sigma_1^2(T)
```
the within-category variance, respectively.

Finding the discrete value ``T`` which maximises the function ``\sigma_b^2(T)``
produces the sought-after threshold value (i.e. the bin which determines the
threshold). As it turns out, that threshold value is equal to the threshold
decided by minimizing the within-category variances criterion ``\sigma_w^2(T)``.
Furthermore, that threshold is also the same as the threshold calculated by
maximizing the ratio of between-category variance to within-category variance.

# Arguments

The function argument is described in more detail below.

##  `img::AbstractArray`

The image that needs to be binarized.  The image is automatically converted
to `Gray` in order to construct the requisite graylevel histogram.

# Example

Binarize the "cameraman" image in the `TestImages` package.

```julia
using TestImages, ImageBinarization

img = testimage("cameraman")
img_binary = binarize(img, Otsu())
```

# Reference

1. Nobuyuki Otsu (1979). “A threshold selection method from gray-level histograms”. *IEEE Trans. Sys., Man., Cyber.* 9 (1): 62–66. [doi:10.1109/TSMC.1979.4310076](http://dx.doi.org/doi:10.1109/TSMC.1979.4310076)
"""
struct Otsu <: AbstractImageBinarizationAlgorithm end

function (f::Otsu)(out::GenericGrayImage, img::GenericGrayImage)
    edges, counts = build_histogram(img,  256)
    t = find_threshold(HistogramThresholding.Otsu(), counts[1:end], edges)
    @simd for i in CartesianIndices(img)
        out[i] = img[i] < t ? 0 : 1
    end
    out
end

(f::Otsu)(out::GenericGrayImage, img::AbstractArray{<:Color3}) =
    f(out, of_eltype(Gray, img))


using Plots
plotlyjs()

# Rossler space equation
function rossler(t::Float64, x::Vector{Float64}, dx::Vector{Float64},
                 a::Float64, b::Float64, c::Float64)
    dx[1] = -x[2] - x[3]
    dx[2] = x[1] + a*x[2]
    dx[3] = b + x[3]*(x[1]-c)
    return dx
end

"""
    rk4(f, x0, time_span[, h, args])

4-th order Runge-Kutta. ``f`` is a handler to the
space state funtion, ``x0`` is a vector of initial
conditions, ``time_span`` is a tuple containing
the start and end simulation times and ``h`` is the
integration interval.

The function ``f`` is called as:

    dx = f(t, x, args...)
"""
function rk4(f::Function, x0::Vector{Float64}, time_span; h=(time_span[2]-time_span[1])/1e4, args=())
    t = time_span[1]:h:time_span[2]
    N = length(t)
    n = length(x0)
    # Allocate buffers
    x = Vector{Vector{Float64}}(N)
    for i = 1:N
        x[i] = Vector{Float64}(n)
    end
    xd = Vector{Float64}(n)
    x_aux = Vector{Float64}(n)
    # Simulate
    copy!(x[1], x0)
    @inbounds for i = 1:N-1
        copy!(x[i+1], x[i])
        # 1-st evaluation
        f(t[i], x[i], xd, args...)
        Base.LinAlg.axpy!(h/6, xd, x[i+1])
        # 2-nd evaluation
        copy!(x_aux, x[i])
        Base.LinAlg.axpy!(h/2, xd, x_aux)
        f(t[i]+h/2, x_aux, xd, args...)
        Base.LinAlg.axpy!(h/6, xd, x[i+1])
        # 3-rd evaluation
        copy!(x_aux, x[i])
        Base.LinAlg.axpy!(h/2, xd, x_aux)
        f(t[i]+h/2, x_aux, xd, args...)
        Base.LinAlg.axpy!(h/6, xd, x[i+1])
        # 4-th evaluation
        copy!(x_aux, x[i])
        Base.LinAlg.axpy!(h, xd, x_aux)
        f(t[i]+h, x_aux, xd, args...)
        Base.LinAlg.axpy!(h/6, xd, x[i+1])
    end
    return x
end

# Simulate Rossler
time_span = (0, 1000)
x0 = [1.0, 1.0, 2.0]
args = (0.2, 0.2, 5.7)

@time x = rk4(rossler, x0, time_span, h=0.01, args=args)

x = hcat(x...)'

plot(x)

plot(x[:, 1], x[:,2], x[:, 3])


module SimpleDrawing
using Plots, LinearAlgebra

export newdraw, draw_circle, draw_arc, draw_segment, draw_point, draw_vector
export draw_rectangle, find_center, non_colinear_check
export finish, draw

"""
`newdraw()` clears the screen and sets up a blank drawing canvas with
no axis, grid, or legend, and sets the aspect ratio to 1.
"""
function newdraw()
    plot(aspectratio = 1, legend = false, axis = false, grid = false, ticks = false)
end

function draw() end

"""
    draw_circle(x,y,r)

Draw a circle centered
at `(x,y)` with radius `r`.

Also: `draw_circle(z,r)` where `z` is complex.
"""
function draw_circle(x::Real, y::Real, r::Real; opts...)
    f(t) = r * cos(t) + x
    g(t) = r * sin(t) + y
    plot!(f, g, 0, 2pi; opts...)
end

draw_circle(z::Complex, r::Real; opts...) = draw_circle(real(z), imag(z), r; opts...)

"""
    draw_circle(a,b,c)

Draw a circle through the points given by
the three complex arguments.
"""
function draw_circle(a::Complex, b::Complex, c::Complex; opts...)
    z = find_center(a, b, c)
    r = abs(a - z)
    x, y = reim(z)
    draw_circle(x, y, r; opts...)
end

export draw_disc

"""
    draw_disc(x,y,r)

Draw a disc centered at `(x,y)` and with radius `r`. May also be invoked as:
* `draw_disc(z,r)` where `z` is complex.
* `draw_disc(a,b,c)` where `a,b,c` are complex which gives a disc passing through `a`, `b`, and `c`.

For example: `draw_disc(0,1,3; color=:yellow, linecolor=:red)`
"""
function draw_disc(x::Real, y::Real, r::Real; opts...)
    draw_circle(x, y, r, seriestype = [:shape], fillalpha = 1; opts...)
end

function draw_disc(z::Complex, r::Real; opts...)
    draw_disc(real(z), imag(z), r; opts...)
end

function draw_disc(a::Complex, b::Complex, c::Complex; opts...)
    z = find_center(a, b, c)
    r = abs(a - z)
    x, y = reim(z)
    draw_disc(x, y, r; opts...)
end

"""
`draw_arc(x,y,r,t1,t2)` draws the shorter arc centered at `(x,y)`
with radius `r` between angles `t1` and `t2`.
"""
function draw_arc(x::Real, y::Real, r::Real, t1::Real, t2::Real; opts...)
    f(t) = r * cos(t) + x
    g(t) = r * sin(t) + y

    plot!(f, g, t1, t2; opts...)
end

"""
`draw_arc(a,b,c)` where the arguments are `Complex`: Draw the
arc from `a` through `b` to `c`.
"""
function draw_arc(a::Complex, b::Complex, c::Complex; opts...)
    m2pi(x::Real) = mod(x, 2π)

    z = find_center(a, b, c)
    x, y = reim(z)
    r = abs(b - z)
    ta = m2pi(angle(a - z))
    tb = m2pi(angle(b - z))
    tc = m2pi(angle(c - z))

    if ta > tc
        ta, tc = tc, ta
    end

    if ta < tb < tc
        return draw_arc(x, y, r, ta, tc; opts...)
    end

    return draw_arc(x, y, r, tc - 2pi, ta; opts...)
end

"""
`draw_segment(x,y,xx,yy)` draws a line segment from `(x,y)` to
`(xx,yy)`. Also `draw_segment(z,zz)` for `Complex` arguments.
"""
function draw_segment(a::Real, b::Real, c::Real, d::Real; opts...)
    plot!([a, c], [b, d]; opts...)
end

function draw_segment(a::Complex, b::Complex; opts...)
    x, y = reim(a)
    xx, yy = reim(b)
    draw_segment(x, y, xx, yy; opts...)
end

"""
`draw_point(x,y)` [or `draw_point(z)`] draws a point at coordinates `(x,y)`
[or complex location `z`].

`draw_point(list)` plots all the points in the one-dimensional `list` of
complex values.
"""
function draw_point(x::Real, y::Real; opts...)
    plot!([x], [y]; marker = 1, opts...)
end

draw_point(z::Complex; opts...) = draw_point(real(z), imag(z); opts...)

function draw_point(pts::Array{Complex{T},1}; opts...) where {T}
    for p in pts
        draw_point(p; opts...)
    end
    plot!()
end

"""
`draw_rectangle(x,y,xx,yy)` draws a rectangle with opposite corners
`(x,y)` and `(xx,yy)`.  May also be called with complex arguments
`draw_rectangle(w,z)`.
"""
function draw_rectangle(x::Real, y::Real, xx::Real, yy::Real; opts...)
    xlist = [x, xx, xx, x, x]
    ylist = [y, y, yy, yy, y]
    plot!(xlist, ylist; opts...)
end

function draw_rectangle(w::Complex, z::Complex; opts...)
    draw_rectangle(real(w), imag(w), real(z), imag(z); opts...)
end

"""
`draw_vector` is used to draw vectors (line segments with an arrow at one end).
The variations are:
+ `draw_vector(x,y)` draws a vector from the origin to `(x,y)`.
+ `draw_vector(x,y,basex,basey)` draws a vector from `(basex,basey)` to
`(basex+x,basey+y)`.
+ `draw_vector(z)` draws a vector from the origin to the complex value `z`.
+ `draw_vector(z,basez)` draws a vector from the complex location `basez` to
`z+basez`.
"""
draw_vector(x::Real, y::Real; opts...) = draw_segment(0, 0, x, y; arrow = :arrow, opts...)

draw_vector(z::Complex; opts...) = draw_segment(0 + 0im, z; arrow = :arrow, opts...)

draw_vector(x::Real, y::Real, basex::Real, basey::Real; opts...) =
    draw_segment(basex, basey, basex + x, basey + y; arrow = :arrow, opts...)

draw_vector(z::Complex, basez::Complex; opts...) =
    draw_segment(basez, basez + z; arrow = :arrow, opts...)

"""
`finish()` is used to clean up a drawing after various calls to `draw`.
It removes the axes and the grid lines, and sets the aspect ratio to one.
"""
function finish()
    plot!(aspectratio = 1, legend = false, axis = false, grid = false, ticks = false)
end

export find_center, non_colinear_check

"""
`find_center(a,b,c)`: Given three points in the complex plane,
find the point `z` that is equidistant from all three. If the three
points are collinear then return `Inf + Inf*im`.
"""
function find_center(a::Complex, b::Complex, c::Complex)::Complex
    if !non_colinear_check(a, b, c)
        return Inf + im * Inf
    end
    A = collect(reim(a))
    B = collect(reim(b))
    C = collect(reim(c))
    AB = 0.5 * (A + B)
    BC = 0.5 * (B + C)
    M = [(A - B)'; (B - C)']
    rhs = [dot(AB, A - B), dot(BC, B - C)]

    Z = M \ rhs
    return Z[1] + im * Z[2]
end

"""
`non_colinear_check(a,b,c)`: test if the complex numbers are distinct
and noncollinear.
"""
function non_colinear_check(a::Complex, b::Complex, c::Complex)::Bool
    if a == b || b == c || a == c
        return false
    end
    z = (b - a) / (c - a)
    return imag(z) != 0
end

"""
    resize_gr_window(wide=800, tall=600)

Change the size of the `gr` drawing window. This is particularly
useful in VS code.
"""
function resize_gr_window(wide::Int = 800, tall::Int = 600)
    gr(size = (wide, tall))
end
export resize_gr_window

include("cubic.jl")
include("curve.jl")
include("my_spy.jl")
include("axes.jl")
end


"""
    smooth(z,sdim)

Smoothen field `z` with a circular kernel of radius `sdim` using convolution
"""
	function smooth(zi,sdim)

        (nsx,nsy)=size(zi);
        #sdim=div(div(ns,nas),2);

        mask=zeros(nsx,nsy);
        nsx2=div(nsx,2)
        nsy2=div(nsy,2)

        for i=1:nsx
           for j=1:nsy
               kx=i-1;
               ky=j-1;
               if(i>nsx2+1)
                  kx=i-nsx-1 ;
               end
               if(j>nsy2+1)
                  ky=j-nsy-1 ;
               end
               r2=kx*kx+ky*ky;
               if(r2<=sdim*sdim)
                  mask[i,j]=1.0;
                end
           end
        end
        zf=real(ifft(fft(mask).*fft(zi)))/sum(mask)
        return zf
   end


mutable struct EmpiricalBayesianRegressor
    _alpha::Float64
    _beta::Float64
    _w_mean::AbstractArray{Float64,1}
    _w_precision::AbstractArray{Float64,2}
    _n_features::Int64
    function EmpiricalBayesianRegressor(alpha::Float64, beta::Float64, n_features)
        new(
            alpha,
            beta,
            zeros(n_features) * 1.0,
            zeros(n_features, n_features) * 1.0,
            n_features,
        )
    end
end

function fitting(
    regressor::EmpiricalBayesianRegressor,
    Phi::AbstractArray{Float64,2},
    t::AbstractArray{Float64,1},
    max_iter::Int64 = 100,
)
    # Phi is the design matrix  [phi(x_1), phi(x_2), ,,, phi(x_N)]^T of size (n_samples, n_features)
    n_features = size(Phi)[2]
    n_samples = size(Phi)[1]
    @assert n_samples == size(t)[1]

    alpha = regressor._alpha
    beta = regressor._beta
    M = transpose(Phi) * Phi
    eigenvalues = eigvals(M)
    @assert n_features == size(eigenvalues)[1]

    E = Matrix(I, n_features, n_features)
    w_mean = zeros(n_features)
    w_precision = zeros(n_features, n_features)
    for i = 1:max_iter
        params = [alpha beta]
        w_precision = alpha * E + beta * M
        w_mean = beta * inv(w_precision) * transpose(Phi) * t
        gamma = sum(eigenvalues ./ (eigenvalues .+ alpha))

        alpha = gamma / norm(w_mean)^2
        beta = (n_samples - gamma) / sum((Phi * w_mean - t) .^ 2)
        if isapprox(params, [alpha beta])
            break
        end
    end

    # update
    regressor._alpha = alpha
    regressor._beta = beta
    regressor._w_mean = w_mean
    regressor._w_precision = w_precision
end

function log_evidence(
    regressor::EmpiricalBayesianRegressor,
    Phi::AbstractArray{Float64,2},
    t::AbstractArray{Float64,1},
)
    # Phi is the design matrix [phi(x_1), phi(x_2), ,,, phi(x_N)]^T of size (n_samples, n_features)
    n_features = size(Phi)[2]
    n_samples = size(Phi)[1]
    M = n_features
    N = n_samples
    @assert n_samples == size(t)[1]

    alpha = regressor._alpha
    beta = regressor._beta
    E = Matrix(I, n_features, n_features)
    A = alpha * E + beta * transpose(Phi) * Phi
    w = regressor._w_mean
    Ew = beta / 2.0 * sum((t - Phi * w) .^ 2) + alpha / 2.0 * (transpose(w) * w)

    return M / 2.0 * log(alpha) + N / 2.0 * log(beta) - Ew - log(det(A)) / 2.0 -
           N / 2.0 * log(2 * pi)
end

function predict(
    regressor::EmpiricalBayesianRegressor,
    Phi::AbstractArray{Float64,2},
    return_std::Bool = true,
)
    # Phi is the design matrix [phi(x_1), phi(x_2), ,,, phi(x_N)]^T of size (n_samples, n_features)
    N = size(Phi)[1]
    w_mean = regressor._w_mean
    w_cov = inv(regressor._w_precision)

    y = Phi * w_mean

    if return_std == true
        y_vars =
            1.0 / regressor._beta .+ [transpose(Phi[i, :]) * w_cov * Phi[i, :] for i = 1:N]
        y_std = sqrt.(y_vars)

        return y, y_std
    end

    return y
end


@testset "Fitting" begin
  wl = geostatsimage("WalkerLake")
  TI = asarray(wl, :Z)[1:20,1:20]
  d = georef((z=TI,))
  g = EmpiricalVariogram(d, :z, maxlag=15.)

  # all fits lead to similar sill
  γ₁ = fit(GaussianVariogram, g)
  γ₂ = fit(SphericalVariogram, g)
  γ₃ = fit(ExponentialVariogram, g)
  γ₄ = fit(MaternVariogram, g)
  @test isapprox(sill(γ₁), 0.054, atol=1e-3)
  @test isapprox(sill(γ₂), 0.054, atol=1e-3)
  @test isapprox(sill(γ₃), 0.058, atol=1e-3)
  @test isapprox(sill(γ₄), 0.057, atol=1e-3)

  # best fit is a Gaussian variogram
  γ = fit(Variogram, g)
  @test γ isa GaussianVariogram
  @test isapprox(sill(γ), 0.054, atol=1e-3)

  if visualtests
    plts = map([γ₁,γ₂,γ₃,γ₄]) do γ
      plot(g, legend=false)
      plot!(γ, 0., 15.)
    end
    @test_reference "data/fitting.png" plot(plts...)
  end

  # make sure convenient methods work
  γ₁ = fit(GaussianVariogram, g, h -> 1/h)
  γ₂ = fit(Variogram, g, h -> 1/h)
  @test sill(γ₁) > 0
  @test sill(γ₂) > 0

  # unitful types
  wl = geostatsimage("WalkerLake")
  TI = asarray(wl, :Z)[1:20,1:20]
  d = georef((z=TI*u"K",))
  g = EmpiricalVariogram(d, :z, maxlag=15.)
  γ = fit(Variogram, g)
  @test unit(sill(γ)) == u"K^2"
  @test unit(nugget(γ)) == u"K^2"
end


using HACTRep
using Test

results=HJBequation(0.03,0.5)

@test size(results.A)[1]==length(results.a)*length(results.z)
@test results.c[1]≈0.0955
#@test results.c[500]≈0.35751765597047436
@test (results.c[500]>=0.2) & (results.c[500]<=0.4)
#@test results.c[1000]≈0.3612279938702292
@test (results.c[1000]>=0.2) & (results.c[1000]<=0.4)
@test results.c[200]≈0.24780538173577077
@test results.v[1]≈-143.49029922114445
@test results.v[500]≈-69.22947350451497
@test results.v[1000]≈-68.54369002729676
@test results.v[200]≈-97.12267792948414

resultsKF=KF(results.a, results.z, results.A)

@test resultsKF[1]<=0.01
@test resultsKF[10]≈1.4366548720872103
@test resultsKF[200]≈0.0
@test resultsKF[500]≈0.0
@test resultsKF[510]≈1.0320935781711122
@test resultsKF[550]≈0.6793836305961967
@test resultsKF[700]≈-0.0

@test equilibriumR(0.5)≈0.021091194152832025
@test equilibriumR(0.05)≈0.03619252443313599
@test equilibriumR(0.8)≈0.026094157079133797
@test equilibriumR(1.0)≈0.03117997884750366
@test equilibriumR(1.2)≈0.034192388057708745
@test equilibriumR(1.5)≈0.037621619701385514

results3=TimeHJBequation(0.03*ones(100),results,100,0.5)

@test results3.v[1,1]≈-143.4902993264915
@test results3.v[200,10]≈-97.12267792838718
@test results3.s[300,50]≈-0.11604803029950189

results4=TimeKF(results3.a,results3.z,results3.A,resultsKF,0.03*ones(100),20)

@test results4[20,20]≈1.337443770972458
@test results4[20,50]≈1.3374437709723683
@test results4[10,20]≈1.436654872087002
@test results4[40,30]≈0.5845834701237425

resultsS=shockSim()

@test resultsS[1]<=-0.1
@test resultsS[end]>=0.029

resultsh=HJBequationHousing(0.03,0.5)

@test resultsh.h[200]≈2.3927855711422845
@test resultsh.h[1]≈0.0
@test resultsh.h[500]≈5.6126932999996875
@test resultsh.v[250]≈-76.42190579246638

resultsKFT=KFHousing(resultsh.a,resultsh.z,resultsh.A)

@test resultsKFT[1]>=4.0
@test resultsKFT[2]<=2.0
@test resultsKFT[end]≈0.0

#@testset "Heterogeneous-Agents-Continuous-Time.jl" begin
    # Write your own tests here.

#end


##include( "../src/lossFunctions.jl" )
using Statistics
using Test

println( "running loss functions' unit tests" );

#== testing for 0,1 labels error functions ==#
a = zeros( Float64, 10000 )
b = ones( Float64, 10000 )

@test mean( lossFunctions.MSE( a, b ) ) == 1.0
@test mean( lossFunctions.MSE( a, a ) ) == 0.0
@test mean( lossFunctions.MSE( b, b ) ) == 0.0
@test mean( lossFunctions.zeroOneLoss( a, b ) ) == 1.0
@test mean( lossFunctions.zeroOneLoss( a, a ) ) == 0.0
@test mean( lossFunctions.zeroOneLoss( b, b ) ) == 0.0

#== testing for 1,-1 labels error functions ==#
a       = Bool.( rand( 0:1, 10000 ) );
b       = .!a;
index_a = ( a .== false );
a       = Float64.( a );
a[index_a] .= -1.0;
index_b     = ( b .== false );
b           = Float64.( b );
b[index_b] .= -1.0;
@test mean( lossFunctions.hingeLoss( a, b ) ) == 2.0
@test mean( lossFunctions.hingeLoss( a, a ) ) == 0.0
@test mean( lossFunctions.hingeLoss( b, b ) ) == 0.0

@test lossFunctions.MSE( 1.0, 1.0 ) == 0.0
@test lossFunctions.MSE( 0.0, 0.0 ) == 0.0
@test lossFunctions.MSE( 1.0, 0.0 ) == 1.0
@test lossFunctions.MSE( 0.0, 1.0 ) == 1.0
@test lossFunctions.MSE( 0.0, 10.0 ) == 100.0
@test lossFunctions.MSE( 10.0, 0.0 ) == 100.0


# generic solve_hjb( ) Implicit method.
# TODO: ns states/nc controls
# TODO: add uncertainty, nx shocks...

using LinearAlgebra, SparseArrays, Plots

# NGM ###############################
if 1==1
    σ= 2.0; ρ = 0.05; δ = 0.05; A = 1.0; α= 0.3;

    # Closed form. Doesn't work.
    # σ= 4.0;  δ = 0.05; A = 1.0; α= 0.3;
    # ρ = δ*(α*σ - 1) # restriction for closed form
    # FOC(v;σ=σ) = abs(v)^(-1.0/σ)                  # FOC, u'^-1()

    r(c; σ=σ)  = (c^(1-σ))/(1-σ)             # return fcn.      make it r(s,c)
    μ(s,c; α=α,δ=δ,A=A) = A*(s^α) - δ*s -c;  # Transition fcn.  ṡ=μ(s,c)

    dr(c;σ=σ) = c^(-σ)                       # derivative of return fcn.
    FOC(v;σ=σ) = v^(-1.0/σ)                  # FOC, u'^-1()

    μ_inv(s,ṡ; α=α,δ=δ,A=A) = A*(s^α) - δ*s -ṡ; # Inv Trans c=μ(s,ṡ), solve ṡ=μ(s,c) for c
    # con0(s;α=α,δ=δ,A=A)     = A*(s^α) - δ*s      # solve: ṡ=μ(s,c) for c(s) if ṡ=0.
    # μ_inv(s,0) = con0(s)

    s_ss = (α*A/(ρ+δ))^(1/(1-α))
    s_min = 0.001*s_ss
    s_max = 2.000*s_ss
    H = 10_000;
    s = LinRange(s_min, s_max, H)
    s = convert(Array, s)
    ds = (s_max-s_min)/(H-1)
end

# con-sav ###############################
if 1==1
    a_bar = -0.02;
    #a_bar = 0.1;
    σ= 2.0; ρ = 0.05; A=0.045; w = 0.1;
    r(c;σ=σ)  = (c^(1-σ))/(1-σ)            # return fcn.
    μ(s,c; w=w,A=A) = w + A*s - c;         # Transition fcn.  ṡ=μ(s,c)

    dr(c;σ=σ) = c^(-σ)                     # deriv return fcn.
    FOC(v;σ=σ) = v^(-1.0/σ)                # FOC, u'^-1()

    μ_inv(s,ṡ; w=w,A=A) = w + A*s -ṡ;  # Inv Trans c=μ(s,ṡ), solve ṡ=μ(s,c) for c
    # μ_inv(s,0) = con0(s)
    #con0(s;α=α,δ=δ,A=A) = w + A*s          # solve: ṡ=μ(s,c) for c(s) if ṡ=0.

    s_min = a_bar
    μ_inv(s_min,0) > 0

    s_max = 1.0
    H = 10_000;
    s = LinRange(s_min, s_max, H)
    s = convert(Array, s)
    ds = (s_max-s_min)/(H-1)
end

##################################
Δ = 1_000
maxit = 10_000
ε = 10e-6
dVf, dVb     = [zeros(H,1) for i in 1:2]
dV_Upwind, c = [zeros(H,1) for i in 1:2]
v0 = @. r(μ_inv(s,0))/ρ #initial guess for V
v = v0
dist=[]

for n=1:maxit
    #println(n)
	V=v
    dV = (V[2:H]-V[1:H-1])/ds

    # forward difference
	dVf[1:H-1] = dV
	dVf[H]     = dr(μ_inv(s_max,0))      # u'(c(s_max)) state constraint, for stability
    #dVf[H]= 0;

	# backward difference
	dVb[2:H] = dV
	dVb[1]   = dr(μ_inv(s_min,0))        # state constraint, for stability
    #dVb[1] = (A*s_min)^(-σ)

	I_concave = dVb .> dVf
    # scatter(I_concave) #1 everywhere EXCEPT @ last point H.

    # consumption and savings with forward difference
    cf  = FOC.(dVf)
    μ_f = μ.(s, cf)

    # consumption and savings with backward difference
    cb  = FOC.(dVb)
    μ_b = μ.(s, cb)

    # c if k̇=0 & V = u'(c)
	c0  = μ_inv.(s,0)   # con0.(s)
    dV0 = dr.(c0)

    # Now to make a choice between forward and backward difference
    If = μ_f .> 0
    Ib = μ_b .< 0
    I0 = (1.0 .- If - Ib)

    dV_Upwind= dVf.*If + dVb.*Ib + dV0.*I0
    c = FOC.(dV_Upwind)
    u = r.(c)

    # plot(dVf, lab="dVf")
    # plot!(dVb, lab="dVb")
    # plot!(dV0, lab="dV0")
    # plot!(dV_Upwind, lab="dV_Upwind")

    # plot(cf, lab="dVf")
    # plot!(cb, lab="dVb")
    # plot!(c0, lab="dV0")
    # plot!(c, lab="dV_Upwind")

    # plot(μ_f, lab="dVf")
    # plot!(μ_b, lab="dVb")
    # plot!(zero(μ_f), lab="dV0")

    # create the transition matrix
    X = -min.(μ_b,0)/ds
    Y = -max.(μ_f,0)/ds + min.(μ_b,0)/ds
    Z = max.(μ_f,0)/ds

    a1 = sparse(Diagonal((Y[:])))
    a2 = [zeros(1,H); sparse(Diagonal((X[2:H]))) zeros(H-1,1)]
    a3 = [zeros(H-1,1) sparse(Diagonal((Z[1:H-1]))); zeros(1,H)]
    AA = a1 + a2 + a3

    B = (ρ + 1/Δ)*sparse(I,H,H) - AA
    b = u + V./Δ

    # Solve: B V = b
    # => (ρ + 1/Δ)*sparse(I,H,H)*V - AA*V = u + V./Δ
    # => (ρ + 1/Δ)*V = u + V./Δ +  AA*V
    # => (ρ)*V + V./Δ = u + V./Δ +  AA*V
    # => (ρ)*V = u  +  AA*V

    V = B \ b
    V_change = V-v
    v = V

	push!(dist,findmax(abs.(V_change))[1])
    println(n, " ", dist[n])
	if dist[n] .< ε
		println("Value Function Converged Iteration=")
		println(n)
		break
	end
end

s_dot = μ.(s,c)
v_err = r.(c) + dV_Upwind.*s_dot - ρ.*v # approx @ borrowing constraint
# TODO: SIMULATE!

plot(dist,
		xlabel="Iteration", ylabel="||V^{n+1} - V^n||",
		#ylims=(-0.001,0.030),
		legend=false, title="")
#png("Convergence")

plot(s, v_err,
		xlabel="s", ylabel="Error in the HJB equation",
		xlims=(s_min,s_max),
		legend=false, title="")
#png("HJB_error")

plot(s, v,
		xlabel="k", ylabel="V(k)",
		xlims=(s_min,s_max),
		legend=false, title="")
#png("Value_function_vs_k")

plot(s, c,
		xlabel="k", ylabel="c(k)",
		xlims=(s_min,s_max),
		legend=false, title="")
#png("c(k)_vs_k")

plot(s, s_dot,
		xlabel="k", ylabel="s(k)",
		xlims=(s_min,s_max), title="", label="s(k)", legend=false)
plot!(s, zeros(H,1), label="", line=:dash)
#png("stateconstraint")

# Simulate.
# s_0 = 0.10;
# c_0 # = interpolate(s_0, (s,c(s)))
# sp = μ(s_0, c)

# Special case of NGM has closed-form sol. Compare!
σ= 4.0;  δ = 0.05; A = 1.0; α= 0.3;
ρ = δ*(α*σ - 1) # restriction for closed form

# Policy: c = (1-s)(k^α)
cp(s) = (1-((α*δ)/(ρ+δ)))*s^α
s_d(s) = (s^α) - δ*s - cp(s)
plot(s, cp)
plot!(s, s_d)

#
s_ss = (α*A/(ρ+δ))^(1/(1-α))
z_ss = (s_ss)^(1-α)
c_ss = (s_ss)^(α)  - δ*s_ss

s_0 = 5;
z_0 = (s_0)^(1-α)

λ=(1-α)*δ
z(t) = z_ss + exp(-λ*t)*(z_0 - z_ss)
kk(t) = z(t)^(1/(1-α))
cc(t) = (1-((α*δ)/(ρ+δ)))*kk(t)^α

t_sim = 0.01:0.01:100
plot(legend=:topleft)
plot!(t_sim, z, lab="z")
plot!(t_sim, kk, lab="k")
plot!(t_sim, cc, lab="c")
plot!([z_ss s_ss c_ss],  seriestype = :hline, lab="", color="grey")


export lorenz63, lorenz96

"""

    lorenz63(du, u, p, t)

Lorenz-63 dynamical model ``u = [x, y, z]`` and ``p = [\\sigma, \\rho, \\mu]``:
```math
\\frac{dx}{dt} = σ(y-x) \\\\
\\frac{dy}{dt} = x(ρ-z) - y \\\\
\\frac{dz}{dt} = xy - βz \\\\
```

- [Example Catalog](@ref)
- [Lorenz system on wikipedia](https://en.wikipedia.org/wiki/Lorenz_system)
"""
function lorenz63(du, u, p, t)

    du[1] = p[1] * (u[2] - u[1])
    du[2] = u[1] * (p[2] - u[3]) - u[2]
    du[3] = u[1] * u[2] - p[3] * u[3]

end

"""
    lorenz96(S, t, F, J)

Lorenz-96 dynamical model. For ``i=1,...,N``:

```math
\\frac{dx_i}{dt} = (x_{i+1}-x_{i-2})x_{i-1} - x_i + F
```

where it is assumed that ``x_{-1}=x_{N-1},x_0=x_N`` and ``x_{N+1}=x_1``.
Here ``x_i`` is the state of the system and ``F`` is a forcing constant.

- [Lorenz 96 model on wikipedia](https://en.wikipedia.org/wiki/Lorenz_96_model)
"""
function lorenz96(dx, x, p, t)
    F = p[1]
    N = Int64(p[2])
    # 3 edge cases
    dx[1] = (x[2] - x[N-1]) * x[N] - x[1] + F
    dx[2] = (x[3] - x[N]) * x[1] - x[2] + F
    dx[N] = (x[1] - x[N-2]) * x[N-1] - x[N] + F
    # then the general case
    for n = 3:(N-1)
        dx[n] = (x[n+1] - x[n-2]) * x[n-1] - x[n] + F
    end
end


using Plots
using BenchmarkTools

function vfsolve(vnew, kgrid, tolerance, imax)
    α = 0.37
    δ = 0.1
    β = 0.95
    σ = 1.5

    v = vnew .+ 2*tolerance
    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1

    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = (c .^ (1-σ) .- 1) ./ (1-σ)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while maximum(abs.(v - vnew)) > tolerance && i <= imax
        v = vnew;

        (vnew, cartesianindex) = findmax(u .+ β*v', dims = 2);
        i += 1;
    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex)
end

kupper = 5
klower = 0.05
n = 1000
kgrid = collect(range(klower, stop = kupper, length = n))
(v, kprime, kprimeindex) = vfsolve(zeros(n), kgrid, 0.001, 1000);

plot(kgrid, kprime, label = "k'")

function policy(x)
    e=getindex(findall(kgrid .== x),1)
    return kprime[e]
end

T = 100
kpath = zeros(T)
k = kgrid[92]
for i in 1:T
    k = policy(k)
    kpath[i] = k
end
time = collect(range(1, T, length = T))
plot(time, kpath, label = "kpath")

# Check k bounds, stepsize, tolerance, imax
findall(kprime .== maximum(kgrid))
any(kprime .== minimum(kgrid))

function vfsolve1(vnew, kgrid, imax)
    α = 1.5
    δ = 0.1
    β = 0.95
    σ = 1.5

    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1

    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = (c .^ (1-σ) .- 1) ./ (1-σ)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while i <= imax

        (vnew, cartesianindex) = findmax(u .+ β*vnew', dims = 2);
        i += 1;

    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex)
end

(v1, kprime1, kprimeindex1) = vfsolve1(zeros(n), kgrid, 1);
(v5, kprime5, kprimeindex5) = vfsolve1(zeros(n), kgrid, 5);
(v10, kprime10, kprimeindex10) = vfsolve1(zeros(n), kgrid, 10);
scatter(kgrid, v1, label = "v1")
scatter(kgrid, v5, label = "v5")
scatter(kgrid, v10, label = "v10")

function vfsolve2(vnew, kgrid, tolerance, imax)
    α = 1.5
    δ = 0.1
    β = 0.95
    σ = 0.5

    v = vnew .+ 2*tolerance
    cartesianindex = Array{CartesianIndex{2}, length(v)}
    i = 1

    c = kgrid.^α + (1-δ)*kgrid .- kgrid'
    c[c .< 0] .= 0
    u = (c .^ (1-σ) .- 1) ./ (1-σ)
    u[kgrid.^α + (1-δ)*kgrid .- kgrid' .< 0] .= -Inf

    while maximum(abs.(v - vnew)) > tolerance && i <= imax
        v = vnew;

        (vnew, cartesianindex) = findmax(u .+ β*v', dims = 2);
        i += 1;
    end
    kprimeindex = getindex.(cartesianindex, 2)
    return (v = vnew, kprime = kgrid[kprimeindex], kprimeindex = kprimeindex)
end

(v, kprime, kprimeindex) = vfsolve2(zeros(n), kgrid, 0.001, 1000);
scatter(kgrid, v, label = "v")
scatter(kgrid, kprime, label = "k'")


"""tests if point p1 and p2 are on the same side of line ab"""
function same_side(p1::Vector{U},p2::Vector{U}, a::Vector{U}, b::Vector{U}) where U <: AbstractFloat
    cp1 = cross(b-a, p1-a)
    cp2 = cross(b-a, p2-a)
    dot(cp1, cp2) >= 0
end

""" tests if a point p is Left|On|Right of an infinite line ab
- >0 for p left of the line ab
- =0 for p on the line
- <0 for p right of the line
"""
function is_left(p::Vector{U}, a::Vector{U}, b::Vector{U}) where U <: Real
    (b[1]-a[1])*(p[2]-a[2]) - (p[1]-a[1])*(b[2]-a[2])
end

"""calculates 2D convec hull for set of points
returns points which determine the hull and its indices in the input list

based on Andrew's Monotone Chain Algorithm (Andrew, 1979), explained here:
http://geomalgorithms.com/a10-_hull-1.html"""
function get_hull(P::Vector{Vector{T}}) where T <: Real
    np = length(P)
    # sort indices, first by x then y
    pm = zeros(T,(np, 3))
    [pm[i,:] = [P[i][1], P[i][2], i] for i in 1:np] # to 2d array
    pm = sortrows(pm, by=x->(x[1],x[2]))
    P = [pm[i,1:2] for i in 1:np]
    idx = Int[pm[i,3] for i in 1:np]
    # get indices of points with 1st x min or max and 2nd y min or max
    min_min = 1 # min x, min y
    i = 0 ## initiate so it is known outside loop, arghh...
    for i in 2:np
        P[i][1] != P[1][1] && break
    end
    min_max = i-1
    for i in np-1:-1:1
        P[i][1] != P[np][1] && break
    end
    max_min = i+1
    max_max = np
    # initiate convex hull stack with point indices
    stack = Int[]

    ## LOWER HULL
    push!(stack, min_min)
    top = 1 # no. of points in stack
    for i in 2:max_min-1 # loop through points with increasing x sequence
        is_left(P[i], P[min_min], P[max_min]) >= 0 && continue # ignore P[i] above or on the lower line
        while top >= 2 # at least two points in stack
            is_left(P[i], P[stack[top-1]], P[stack[top]]) > 0 && break # P[i] in hull
            pop!(stack)
            top -= 1
        end
        push!(stack, i)
        top += 1
    end
    push!(stack, max_min)
    top += 1

    ## UPPER HULL
    # if  distinct xmax points push max_min point onto stack
    if max_max != max_min
        push!(stack, max_max)
        top += 1
    end
    bot = top # bottom point of upper hull
    for i in np-1:-1:2 # loop through points with decreasing x sequence
        is_left(P[i], P[max_max], P[min_max]) >= 0 && continue # ignore P[i] below or on the upper line
        while bot > top # at least two points in upper stack
            is_left(P[i], P[stack[top-1]], P[stack[top]]) > 0 && break # P[i] in hull
            pop!(stack)
            top -= 1
        end
        push!(stack, i)
        top += 1
    end
    push!(stack, min_max)
    top += 1

    # return points of hull and indices
    P[stack], idx[stack]
end

"""calculates bbox coordinates for set of points
returns the list with coordinates the indices of the points which determine the bbox"""
function get_bbox(P::Vector{Vector{T}}) where T <: Real
    @assert length(P) >= 2
    xmin, xmax, ymin, ymax = Inf, -Inf, Inf, -Inf
    xmini, xmaxi, ymini, ymaxi = 0, 0, 0, 0  # declare indices of points that make up bbox
    i = 0
    @inbounds for p in P
        addp = false
        i += i  # index of point in vector P
        if p[1] < xmin
            xmin = p[1]
            xmini = i
        end
        if p[1] > xmax
            xmax = p[1]
            xmaxi = i
        end
        if p[2] < ymin
            ymin = p[2]
            ymini = i
        end
        if p[2] > ymax
            ymax = p[2]
            ymaxi = p
        end
    end

    # return bbox and indices that make up bbox
    [xmin,xmax,ymin,ymax], [xmini, xmaxi, ymini, ymaxi]
end

function triangle_area(a::Vector{U}, b::Vector{U}, c::Vector{U}) where U <: AbstractFloat
    abs(a[1]*(b[2]-c[2]) + b[1]*(c[2]-a[2]) + c[1]*(a[2]-b[2])) / 2.
end

"""check if point p [x,y] in triangle made out of points a,b,c"""
function in_triangle(p::Vector{U}, a::Vector{U},b::Vector{U},c::Vector{U}) where U <: AbstractFloat
    p,a,b,c = [[x[1],x[2],0.] for x in [p,a,b,c]]
    same_side(p,a, b,c) && same_side(p,b, a,c) && same_side(p,c, a,b)
end

"""check if point (x,y) in traingle made out of points vector"""
function in_triangle(x::AbstractFloat, y::AbstractFloat, points::Vector{Vector{U}}) where U <: AbstractFloat
    length(points)
    npoints != 3 && warn("more than three points in list, check based on first three points only")
    a,b,c = points[1], points[2], points[3]
    xypoint_in_triangle([x, y], a, b,c)
end

"""check if point (x,y) in bbox"""
function in_bbox(x::AbstractFloat, y::AbstractFloat, points::Vector{Vector{U}}) where U <: AbstractFloat
    xmin,xmax,ymin,ymax = get_bbox(points)[1]
    xmin <= x <= xmax && ymin <= y <= ymax
end

"""check if point in hull of all points based on hull computation"""
function in_hull(x::Real, y::Real, points::Vector{Vector{U}}) where U <: Real
    npoints = length(points)
    @assert npoints >= 3 ["point vector should contain at least three points"]
    push!(points, [x,y])
    p_hull, idx_hull = get_hull(points)
    !in(npoints+1, idx_hull)
end


export Companion

immutable Companion{T} <: AbstractArray{T, 2}
    c :: Vector{T}
end
#From polynomial

using Polynomials
#Generate companion matrix from a polynomial

function Companion(P::Poly)
   n = length(P)
   c = Array(Number,n-1)
   d=P.a[n]
   for i=1:n-1
       c[i]=P.a[i]/d
   end
   Companion(c)
end

#Basic property computations
size(C::Companion, r::Int) = (r==1 || r==2) ? length(C.c) :
    throw(ArgumentError("Companion is of rank 2"))

function size(C::Companion)
    n = length(C.c)
    n, n
end

#XXX Inefficient but works
getindex(C::Companion, i, j) = getindex(full(C), i, j)
isassigned(C::Companion, i, j) = isassigned(full(C), i, j)

function full{T}(C::Companion{T})
    M = zeros(T, size(C)...)
    M[:,end]=-C.c
    for i=1:size(C,1)-1
    	M[i+1, i] = one(T)
	end
    M
end

#Linear algebra stuff
function A_mul_B!{T}(C::Companion{T}, b::Vector{T})
	x = b[end]
	y = -C.c[1]*x
	b[2:end] = b[1:end-1]-C.c[2:end]*x
	b[1] = y
    b
end
*{T}(C::Companion{T}, b::Vector{T}) = A_mul_B!(C, copy(b))

function A_mul_B!{T}(A::Matrix{T}, C::Companion{T})
	v = Array(T, size(A,1))
	for i=1:size(A,1)
		v[i] =(A[i,:]*-C.c)[1]
	end
	for i=1:size(A,1), j=1:size(A,2)-1
		A[i,j] = A[i,j+1]
	end
	A[:,end] = v
	A
end
*{T}(A::Matrix{T}, C::Companion{T}) = A_mul_B!(copy(A), C)

function inv{T}(C::Companion{T})
	M = zeros(T, size(C)...)
    for i=1:size(C,1)-1
    	M[i, i+1] = one(T)
	end
	d = M[end, 1] = -one(T)/C.c[1]
	M[1:end-1, 1] = d*C.c[2:end]
    M
end


"""
    mat_C1(q,N)

Generate the matrix corresponding to even, π-periodic solutions to the
angular Mathieu equation.

* `q`: real parameter of Mathieu’s equation
* `N`: truncated matrix dimension (matrix is N×N)
"""
mat_C1(q::Number,N::Integer) = mat_C1(blasfloat(q),N)
function mat_C1(q::T,N::Integer) where T<:BlasReal
    dd = [T(2k)^2 for k = 0:N-1]
    du = fill(q,N-1)
    du[1] *= √2
    return SymTridiagonal(dd,du)
end

"""
    mat_C2(q,N)

Generate the matrix corresponding to even, 2π-periodic (π-antiperiodic)
solutions to the angular Mathieu equation.

* `q`: real parameter of Mathieu’s equation
* `N`: truncated matrix dimension (matrix is N×N)
"""
mat_C2(q::Number,N::Integer) = mat_C2(blasfloat(q),N)
function mat_C2(q::T,N::Integer) where T<:BlasReal
    dd = [T(2k+1)^2 for k = 0:N-1]
    dd[1] += q
    du = fill(q,N-1)
    return SymTridiagonal(dd,du)
end

"""
    mat_C3(q,N)

Generate the matrix corresponding to odd, 2π-periodic (π-antiperiodic)
solutions to the angular Mathieu equation.

* `q`: real parameter of Mathieu’s equation
* `N`: truncated matrix dimension (matrix is N×N)
"""
mat_C3(q::Number,N::Integer) = mat_C3(blasfloat(q),N)
function mat_C3(q::T,N::Integer) where T<:BlasReal
    dd = [T(2k+1)^2 for k = 0:N-1]
    dd[1] -= q
    du = fill(q,N-1)
    return SymTridiagonal(dd,du)
end

"""
    mat_C4(q,N)

Generate the matrix corresponding to odd, π-periodic solutions to the
angular Mathieu equation.

* `q`: real parameter of Mathieu’s equation
* `N`: truncated matrix dimension (matrix is N×N)
"""
mat_C4(q::Number,N::Integer) = mat_C4(blasfloat(q),N)
function mat_C4(q::T,N::Integer) where T<:BlasReal
    dd = [T(2k+2)^2 for k = 0:N-1]
    du = fill(q,N-1)
    return SymTridiagonal(dd,du)
end


@parallel function diffusion3D_step!(T2, T, Ci, lam, dt, dx, dy, dz)
    @inn(T2) = @inn(T) + dt*(lam*@inn(Ci)*(@d2_xi(T)/dx^2 + @d2_yi(T)/dy^2 + @d2_zi(T)/dz^2));
    return
end

function diffusion3D()
# Physics
lam        = 1.0;                                        # Thermal conductivity
cp_min     = 1.0;                                        # Minimal heat capacity
lx, ly, lz = 10.0, 10.0, 10.0;                           # Length of computational domain in dimension x, y and z

# Numerics
nx, ny, nz = 8, 8, 8;                              # Number of gridpoints in dimensions x, y and z
nt         = 3;                                          # Number of time steps
dx         = lx/(nx-1);                                  # Space step in x-dimension
dy         = ly/(ny-1);                                  # Space step in y-dimension
dz         = lz/(nz-1);                                  # Space step in z-dimension

# Array initializations
T   = @zeros(nx, ny, nz);
T2  = @zeros(nx, ny, nz);
Ci  = @zeros(nx, ny, nz);

# Initial conditions (heat capacity and temperature with two Gaussian anomalies each)
Ci .= 1.0./( cp_min .+ Data.Array([5*exp(-(((ix-1)*dx-lx/1.5))^2-(((iy-1)*dy-ly/2))^2-(((iz-1)*dz-lz/1.5))^2) +
                                   5*exp(-(((ix-1)*dx-lx/3.0))^2-(((iy-1)*dy-ly/2))^2-(((iz-1)*dz-lz/1.5))^2) for ix=1:size(T,1), iy=1:size(T,2), iz=1:size(T,3)]) )
T  .= Data.Array([100*exp(-(((ix-1)*dx-lx/2)/2)^2-(((iy-1)*dy-ly/2)/2)^2-(((iz-1)*dz-lz/3.0)/2)^2) +
                   50*exp(-(((ix-1)*dx-lx/2)/2)^2-(((iy-1)*dy-ly/2)/2)^2-(((iz-1)*dz-lz/1.5)/2)^2) for ix=1:size(T,1), iy=1:size(T,2), iz=1:size(T,3)])
T2 .= T;                                                 # Assign also T2 to get correct boundary conditions.

# Time loop
dt = min(dx^2,dy^2,dz^2)*cp_min/lam/8.1;                 # Time step for the 3D Heat diffusion
for it = 1:nt
    @parallel diffusion3D_step!(T2, T, Ci, lam, dt, dx, dy, dz);
    T, T2 = T2, T;
end

return true
end


import Random: GLOBAL_RNG

"""Random number generator for Float32 in [0,1) that samples from
42*2^23 float32s in [0,1) compared to 2^23 for rand(Float32)."""
function randfloat(rng::Random.AbstractRNG,::Type{Float32})
    # create exponent bits in 0000_0000 to 0111_1110
    # at following chances
    # e=01111110 at 50.0% for [0.5,1.0)
    # e=01111101 at 25.0% for [0.25,0.5)
    # e=01111100 at 12.5% for [0.125,0.25)
    # ...
    ui = rand(rng,UInt64)

    # count leading zeros of random UInt64
    # 0 leading zeros at 50% chance
    # 1 leading zero at 25% chance
    # 2 leading zeros at 12.5% chance etc.
    # then convert leading zeros to exponent bits of Float32
    lz = leading_zeros(ui)
    e = ((126 - lz) % UInt32) << 23

    # for 64 leading zeros the smallest float32 that can be created is 2.7105054f-20
    # use last 23 bits for signficand, only when they are not part of the leading zeros
    # to sample from all floats in 2.7105054f-20 to prevfloat(1f0)
    ui = lz > 40 ? rand(rng,UInt64) : ui

    # combine exponent and signficand
    return reinterpret(Float32,e | ((ui % UInt32) & 0x007f_ffff))
end

"""Random number generator for Float64 in [0,1) that samples from
64*2^52 floats compared to 2^52 for rand(Float64)."""
function randfloat(rng::Random.AbstractRNG,::Type{Float64})
    # create exponent bits in 000_0000_0000 to 011_1111_1110
    # at following chances
    # e=01111111110 at 50.0% for [0.5,1.0)
    # e=01111111101 at 25.0% for [0.25,0.5)
    # e=01111111100 at 12.5% for [0.125,0.25)
    # ...
    ui = rand(rng,UInt64)

    # count leading zeros of random UInt64 in several steps
    # 0 leading zeros at 50% chance
    # 1 leading zero at 25% chance
    # 2 leading zeros at 12.5% chance etc.
    # then convert leading zeros to exponent bits of Float64
    lz = leading_zeros(ui)
    e = ((1022 - lz) % UInt64) << 52

    # for 64 leading zeros the smallest float64 that
    # can be created is 2.710505431213761e-20

    # draw another UInt64 for significant bits in case the leading
    # zeros reach into the bits that would be used for the significand
    # (in which case the first signifcant bits would always be zero)
    ui = lz > 11 ? rand(rng,UInt64) : ui

    # combine exponent and significand (sign always 0)
    return reinterpret(Float64,e | (ui & 0x000f_ffff_ffff_ffff))
end

"""Random number generator for Float16 in [0,1) that samples from
all 15360 float16s in that range."""
function randfloat(rng::Random.AbstractRNG,::Type{Float16})
    # create exponent bits in 00000 to 01110
    # at following chances
    # e=01110 at 50.0% for [0.5,1.0)
    # e=01101 at 25.0% for [0.25,0.5)
    # e=01100 at 12.5% for [0.125,0.25)
    # ...
    ui = rand(rng,UInt32) | 0x0002_0000
    # set 15th bit to 1 to have at most 14 leading zeros.

    # count leading zeros of random UInt64 in several steps
    # 0 leading zeros at 50% chance
    # 1 leading zero at 25% chance
    # 2 leading zeros at 12.5% chance etc.
    # then convert leading zeros to exponent bits of Float16
    lz = leading_zeros(ui)
    e = ((14 - lz) % UInt32) << 10

    # combine exponent and significand (sign always 0)
    return reinterpret(Float16,(e | (ui & 0x0000_03ff)) % UInt16)
end

# use stdlib default RNG as a default here too
randfloat(::Type{T}=Float64) where {T<:Base.IEEEFloat} = randfloat(GLOBAL_RNG,T)
randfloat(rng::Random.AbstractRNG) = randfloat(rng,Float64)

# randfloat for arrays - in-place
function randfloat!(rng::Random.AbstractRNG, A::AbstractArray{T}) where T
    for i in eachindex(A)
        @inbounds A[i] = randfloat(rng, T)
    end
    A
end

# randfloat for arrays with memory allocation
randfloat(rng::Random.AbstractRNG, ::Type{T}, dims::Integer...) where {T<:Base.IEEEFloat} = randfloat!(rng, Array{T}(undef,dims))
randfloat(rng::Random.AbstractRNG,            dims::Integer...)                           = randfloat!(rng, Array{Float64}(undef,dims))
randfloat(::Type{T}, dims::Integer...) where {T<:Base.IEEEFloat} = randfloat!(GLOBAL_RNG, Array{T}(undef,dims))
randfloat(           dims::Integer...)                           = randfloat!(GLOBAL_RNG, Array{Float64}(undef,dims))


using ApproxFun, BandedMatrices, BlockBandedMatrices, LinearAlgebra, BlockArrays, MultivariateOrthogonalPolynomials, Test
using SparseVolterraExamples

#############################
## Tests for Volterra integral operator
## with Clenshaw Jacobi expansion of the kernel.
####
## Includes tests for 1 to x and 1 to 1-x limits of integration.
#############################

@testset "Set #1.1 - Volterra, Clenshaw sin kern, flip = true" begin
    K(x,y)=sin(x+y);
    f=Fun(y->y, Jacobi(0,1,0..1));
    @test triVolterraFullKernelFunP01(f,K,25,true) ≈ Fun(x->-x*cos(2*x)-sin(x)+sin(2*x), Jacobi(1,0,0..1));
end

@testset "Set #1.2 - Volterra, Clenshaw sin kern, flip = false, val test" begin
    K(x,y)=sin(x+y);
    val = rand(0..1);
    f=Fun(y->y, Jacobi(0,1,0..1));
    @test triVolterraFullKernelFunP01(f,K,30,false)(val) ≈ Fun(x->((-1)*cos(1)+x*cos(1)+sin(1)-sin(x)), Jacobi(1,0,0..1))(val);
end

@testset "Set #2 - Volterra, Clenshaw sinh kern, flip = true" begin
    K(x,y)=(y-x)*sinh(x);
    f=Fun(y->y, Jacobi(0,1,0..1));
    @test triVolterraFullKernelFunP01(f,K,25,true) ≈ Fun(x->-1/6*x^3*sinh(x), Jacobi(1,0,0..1));
end

@testset "Set #3 - Volterra, Clenshaw sinh kern, flip = false" begin
    K(x,y)=(y-x)*sinh(x^2);
    f=Fun(y->y, Jacobi(0,1,0..1));
    @test triVolterraFullKernelFunP01(f,K,25,false) ≈ Fun(x->(sinh(x^2))/3-3/2*x*sinh(x^2)+2*x^2*sinh(x^2)-5/6*x^3*sinh(x^2), Jacobi(0,1,0..1));
end


# This file is part of GenericSchur.jl, released under the MIT "Expat" license

# The method in this file is derived from LAPACK's zlacon.
# LAPACK is released under a BSD license, and is
# Copyright:
# Univ. of Tennessee
# Univ. of California Berkeley
# Univ. of Colorado Denver
# NAG Ltd.

# Hager's one-norm estimator, with modifications by N.J. Higham
"""
    norm1est!(applyA!,applyAH!,y::Vector) => γ

Estimate the 1-norm of a linear operator `A` expressed as functions which
apply `A` and `adjoint(A)` to a vector such as `y`.

cf. N.J. Higham, SIAM J. Sci. Stat. Comp. 11, 804 (1990)
"""
function norm1est!(applyA!, applyAH!, y::AbstractVector{Ty}; maxiter=5) where Ty
    n = length(y)
    x = fill(one(Ty)/n,n)
    y .= zero(Ty)
    applyA!(x)
    (n == 1) && return abs(x[1])
    γ = norm(x,1)
    tiny = safemin(real(Ty))
    for i=1:n
        absxi = abs(x[i])
        if absxi > tiny
            x[i] /= absxi
        else
            x[i] = one(Ty)
        end
    end
    applyAH!(x)
    ax0,j0 = _findamax(x)
    for iter = 1:maxiter
        x .= zero(Ty)
        x[j0] = one(Ty)
        applyA!(x)
        y .= x
        oldγ = γ
        γ = norm(y,1)
        if γ <= oldγ
            break
        end
        for i=1:n
            absxi = abs(x[i])
            if absxi > tiny
                x[i] /= absxi
            else
                x[i] = one(Ty)
            end
        end
        applyAH!(x)
        jlast = j0
        ax0, j0 = _findamax(x)
        if abs(x[jlast]) == ax0
            break
        end
    end
    # alternative estimate for tricky cases (see Higham 1990)
    v = x # reuse workspace
    isign = 1
    for i in 1:n
        v[i] = isign * (1+(i-1)/(n-1))
        isign = -isign
    end
    applyA!(v)
    t = 2*norm(v,1) / (3*n)
    return max(t,γ)
end

function _findamax(x::AbstractVector{T}) where T
    ax0 = abs(x[1])
    i0 = 1
    for i=2:length(x)
        ax = abs(x[i])
        if ax > ax0
            ax0 = ax
            i0 = i
        end
    end
    return ax0,i0
end

function _findamax(x::AbstractVector{T}) where T <: Complex
    ax0 = abs2(x[1])
    i0 = 1
    for i=2:length(x)
        ax = abs2(x[i])
        if ax > ax0
            ax0 = ax
            i0 = i
        end
    end
    return sqrt(ax0),i0
end


export CyclicRange

#= make CyclicRange For ScaLAPACK =#

struct CyclicRange{T <: Integer} <: AbstractRange{T}
    start::T # full array's first
    stop::T # full array's last
    blocksize::T # blocksize
    n_split::T # number of process
    split_index::T # what process

    CyclicRange(;start::T, stop::T, blocksize::T, n_split::T, split_index::T) where T <: Integer = new{T}(start, stop, blocksize, n_split, split_index)
end

number_element(r::CyclicRange) = r.stop - r.start + 1

function number_element_local(r::CyclicRange)
    n_element = number_element(r)
    n_block_local_min = number_block_local_min(r)
    n_extra_block = number_block_global_extra(r)

    n_element_local = r.blocksize*n_block_local_min
    n_extra_element = n_element % r.blocksize

    # add extra block
    if r.split_index <= n_extra_block
        n_element_local += r.blocksize
    end
    # add extra element,
    if r.split_index == n_extra_block + 1 && n_extra_element != 0
        n_element_local += n_extra_element
    end
    return n_element_local
end

function number_block(r::CyclicRange)
    n_element = number_element(r)
    n_block = n_element ÷ r.blocksize
    return n_block
end

function number_block_local_min(r::CyclicRange{T}) where T <: Integer
    n_block = number_block(r)
    n_block_local = n_block ÷ r.n_split
    return n_block_local
end

function number_block_global_extra(r::CyclicRange{T}) where T <: Integer
    n_block = number_block(r)
    n_extra_block = n_block % r.n_split
    return n_extra_block
end

function local_index(r::CyclicRange, i::Integer)
    i_temp = i - 1
    block_position = i_temp ÷ r.blocksize
    local_index = (block_position ÷ r.n_split)*r.blocksize + (i_temp % r.blocksize) + 1
    split_index = block_position % r.n_split + 1
    # mpi rank = split_index - 1, in 0-based local index = extra_element - 1
    return split_index, local_index
end

Base.firstindex(::CyclicRange) = 1

Base.first(r::CyclicRange) = getindex(r, 1)

Base.last(r::CyclicRange) = getindex(r, length(r))

Base.length(r::CyclicRange{T}) where T <: Integer = number_element_local(r)

Base.getindex(r::CyclicRange, i::Integer) = r.blocksize*(r.split_index-1) + r.n_split*r.blocksize*div(i-1, r.blocksize) + mod(i-1, r.blocksize) + 1

Base.show(io::IO, r::CyclicRange) = print(io, "$(r.start):[$(r.split_index) $(r.blocksize) $(r.n_split)]:$(r.stop)")


using Revise
using Basins
using Plots
using DynamicalSystems
using DifferentialEquations

#
#
# ω=1.
# F = 0.2
# ds =Systems.duffing([0.1, 0.25]; ω = ω, f = F, d = 0.15, β = -1)
# integ_df  = integrator(ds; alg=Tsit5(),  reltol=1e-8, save_everystep=false)
# xg = range(-2.2,2.2,length=150)
# yg = range(-2.2,2.2,length=150)
#
# @time bsn = Basins.basins_map2D(xg, yg, integ_df; T=2*pi/ω)
#
# sa,sb = compute_saddle(integ_df, bsn, [1], [2]; N=1000)

# plot(xg,yg,bsn.basin', seriestype=:heatmap)
# s = Dataset(sa)
# plot!(s[:,1],s[:,2],seriestype=:scatter)

# io = open("myfile.txt", "w");
# for v in s
#     write(io, string(v[1], " ", v[2], "; \n"));
# end
# close(io);

# Forced pendulum

# Equations of motion:
function forced_pendulum!(du, u, p, t)
    d = p[1]; F = p[2]; omega = p[3]
    du[1] = u[2]
    du[2] = -d*u[2] - sin(u[1])+ F*cos(omega*t)
end

# We have to define a callback to wrap the phase in [-π,π]
function affect!(integrator)
    if integrator.u[1] < 0
        integrator.u[1] += 2*π
    else
        integrator.u[1] -= 2*π
    end
end

condition(u,t,integrator) = (integrator.u[1] < -π  || integrator.u[1] > π)

cb = DiscreteCallback(condition,affect!)

#d, F ,w
F = 1.66
ω = 1.
d=0.2
p=[d, F, ω]
#p=[0.15, 0.2, 0.1]
df = ODEProblem(forced_pendulum!,rand(2),(0.0,20.0), p)
integ_df  = init(df, alg=AutoTsit5(Rosenbrock23()); reltol=1e-9, abstol=1e-9, save_everystep=false, callback=cb)

xres=200
yres=200

# range for forced pend
xg = range(-pi,pi,length=xres)
yg = range(-2.,4.,length=yres)

# compute basin
@time bsn = Basins.basins_map2D(xg, yg, integ_df; T=2*pi/ω)

Na = length(unique(bsn.basin))

sa,sb = compute_saddle(integ_df, bsn, [1], [2,3]; N=1000)

plot(xg,yg,bsn.basin', seriestype=:heatmap)
s = Dataset(sa)
plot!(s[:,1],s[:,2],seriestype=:scatter, markercolor=:blue)


#=
jinc.jl
=#

export jinc

using SpecialFunctions: besselj1

"""
    jinc(x)

Return `jinc(x) = J1(π*x)/(2x)`,
where `J1` is a Bessel function of the first kind.

The argument `x` must be unitless.

Return type is `promote_type(typeof(x), Float32)`.
"""
function jinc(x::X) where {X <: Real}
    T = promote_type(X, Float32)
    if (x == 0) return convert(T, π/4) end
    y = abs(x)
    convert(T, besselj1(π*y) / (2*y))
end


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

example = "
939
7,13,x,x,59,x,31,19
"

function read_input(string)
    earliest_timestamp, bus_numbers = split(string, '\n', keepempty=false)
    earliest_timestamp = parse(Int, earliest_timestamp)
    bus_numbers = map(x -> tryparse(Int, x), split(bus_numbers, ','))
    return earliest_timestamp, bus_numbers
end

@assert read_input(example) == (939, [7, 13, nothing, nothing, 59, nothing, 31, 19])

function closest_arrival(bus_number, timestamp)
    divisor, reminder = divrem(timestamp, bus_number)
    if reminder == 0
        return 0
    else
        return bus_number * (divisor + 1) - timestamp
    end
end

function part1(input)
    earliest_timestamp, bus_numbers = read_input(input)
    bus_numbers = convert(Vector{Int}, filter(x -> !isnothing(x), bus_numbers))
    closest_arrivals = closest_arrival.(bus_numbers, earliest_timestamp)
    i = argmin(closest_arrivals)
    return closest_arrivals[i] * bus_numbers[i]
end

@assert part1(example) == 295

"""
Solve Chinese reminder theorem problem using the inverse modulo algorithm

See:
* https://en.wikipedia.org/wiki/Chinese_remainder_theorem
* https://www.geeksforgeeks.org/chinese-remainder-theorem-set-2-implementation/
* https://rosettacode.org/wiki/Chinese_remainder_theorem
"""
function chinese_reminder(modulus, reminder)
    N = prod(modulus)
    Ni = div.(N, modulus)
    i = invmod.(Ni, modulus)
    return mod(sum(reminder .* Ni .* i), N)
end

@assert chinese_reminder([3, 4, 5], [2, 3, 1]) == 11

function part2(input)
    _, bus_numbers = read_input(input)

    nonmissing = map(x -> !isnothing(x), bus_numbers)
    indexes = collect(0:(length(bus_numbers) - 1))
    indexes = indexes[nonmissing]
    bus_numbers = bus_numbers[nonmissing]
    reminders = maximum(indexes) .- indexes

    return chinese_reminder(bus_numbers, reminders) - maximum(reminders)
end

@assert part2(example) == 1068781
@assert part2("0\n17,x,13,19") == 3417
@assert part2("0\n67,7,59,61") == 754018
@assert part2("0\n67,x,7,59,61") == 779210
@assert part2("0\n67,7,x,59,61") == 1261476
@assert part2("0\n1789,37,47,1889") == 1202161486

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

@assert result1 == 4207
@assert result2 == 725850285300475


# Symmetric Toeplitz solver
function durbin!{T<:BlasReal}(r::AbstractVector{T}, y::AbstractVector{T})
    n = length(r)
    n <= length(y) || throw(DimensionMismatch("Auxiliary vector cannot be shorter than data vector"))
    y[1] = -r[1]
    β = one(T)
    α = -r[1]
    for k = 1:n-1
        β *= one(T) - α*α
        α = -r[k+1]
        for j = 1:k
            α -= r[k-j+1]*y[j]
        end
        α /= β
        for j = 1:div(k,2)
            tmp = y[j]
            y[j] += α*y[k-j+1]
            y[k-j+1] += α*tmp
        end
        if isodd(k) y[div(k,2)+1] *= one(T) + α end
        y[k+1] = α
    end
    return y
end
durbin{T<:BlasReal}(r::AbstractVector{T}) = durbin!(r, zeros(T, length(r)))

function levinson!{T<:BlasReal}(r::AbstractVector{T}, b::AbstractVector{T}, x::AbstractVector{T})
    n = length(b)
    n == length(r) || throw(DimensionMismatch("Vectors must have same length"))
    n <= length(x) || throw(DimensionMismatch("Auxiliary vector cannot be shorter than data vector"))
    x[1] = b[1]
    b[1] = -r[2]/r[1]
    β = one(T)
    α = -r[2]/r[1]
    for k = 1:n-1
        β *= one(T) - α*α
        μ = b[k+1]
        for j = 2:k+1
            μ -= r[j]/r[1]*x[k-j+2]
        end
        μ /= β
        for j = 1:k
            x[j] += μ*b[k-j+1]
        end
        x[k+1] = μ
        if k < n - 1
            α = -r[k+2]
            for j = 2:k+1
                α -= r[j]*b[k-j+2]
            end
            α /= β*r[1]
            for j = 1:div(k,2)
                tmp = b[j]
                b[j] += α*b[k-j+1]
                b[k-j+1] += α*tmp
            end
            if isodd(k) b[div(k,2)+1] *= one(T) + α end
            b[k+1] = α
        end
    end
    for i = 1:n
        x[i] /= r[1]
    end
    return x
end
levinson{T<:BlasReal}(r::AbstractVector{T}, b::AbstractVector{T}) = levinson!(r, copy(b), zeros(T, length(b)))


# Functions to create standard graph matrices
#import Base.eigvals
export adjacency, char_poly, laplace, incidence, dist_matrix, eigvals
using SparseArrays
# Adjaceny Matrix

"""
`adjacency(G)` returns the adjacency matrix of `G`.

Note: If the vertices can be sorted by `sort`, then the first row of
the adjacency matrix corresponds to the first vertex (in order) in `G`
and so forth. However, if the vertices are not sortable in this way,
the mapping between vertices and rows/columns of the matrix is
unpredictable.
"""
function adjacency(G::SimpleGraph)
    n = NV(G)
    A = zeros(Int,(n,n))

    # create a table from V to 1:n
    d = vertex2idx(G)

    for e in G.E
        i = d[e[1]]
        j = d[e[2]]
        A[i,j]=1
        A[j,i]=1
    end

    return A
end

# Laplace matrix

"""
`laplace(G)` returns the Laplacian matrix of `G`. This is the
adjacency matrix minus the (diagonal) degree matrix. See `adjacency`
to understand how vertices correspond to rows/columns of the resulting
matrix.
"""
function laplace(G::SimpleGraph)
    A = adjacency(G)
    d = collect(sum(A,dims=1))[:]
    # D = diagm(d) # deprecated version
    D = Matrix(Diagonal(d))
    L = D-A
    return L
end

# incidence matrix
"""

`incidence(G)` returns the vertex-edge incidence matrix of `G`.

Notes:

* The result is a sparse matrix. Wrap in `full` to convert to nonsparse.

* Each column of the matrix has exactly one `+1` and one `-1`. If `G`
is undirected and an unsigned incidence matrix is desired, use
`incidence(G,false)`.
"""
function incidence(G::SimpleGraph, signed::Bool = true)
    n = NV(G)
    m = NE(G)
    M = spzeros(Int,n,m)
    d = vertex2idx(G)
    E = elist(G)
    a = 1
    b = signed ? -1 : 1

    idx = 0
    for e in E
        i = d[e[1]]
        j = d[e[2]]
        idx += 1
        M[i,idx] = a
        M[j,idx] = b
    end

    return M
end

# Create the n-by-n distance matrix
"""
`dist_matrix(G)` returns a matrix whose `i,j`-entry is the distance
from the `i`th vertex to the `j`th vertex. If there is no `i,j`-path,
that entry is `-1`.
"""
function dist_matrix(G::AbstractSimpleGraph)
    if cache_check(G,:dist_matrix)
      return cache_recall(G,:dist_matrix)
    end
    vtcs = vlist(G)
    n = length(vtcs)
    dd = dist(G)

    A = zeros(Int,n,n)

    for i = 1:n
        u = vtcs[i]
        for j = 1:n
            v = vtcs[j]
            A[i,j] = dd[(u,v)]
        end
    end
    cache_save(G,:dist_matrix,A)
    return A
end

"""
`char_poly(G)` returns the characteristic polynomial of
`adjacency(G)`. Use `char_poly(G,function)` for other
possible integer matrix functions such as `laplace`.
"""
function char_poly(G::AbstractSimpleGraph, func::Function=adjacency)
    if cache_check(G,:char_poly) && func==adjacency
      return cache_recall(G,:char_poly)
    end
    evs = eigvals(G,func)
    P = poly(evs)
    cs = round.(Int,real(coeffs(P)))
    P =  Poly(cs)
    if func==adjacency
        cache_save(G,:char_poly,P)
    end
    return P
end

"""
`eigvals(G)` for a `SimpleGraph` returns the eigenvalues of `G`'s
adjacency matrix. More generally, `eigvals(G,mat)` returns the eigenvalues
of `mat(G)` where `mat` is a matrix-valued function of `G`. In particular,
one can use `mat(G,laplace)` to find the eigenvalues of `G`'s Laplacian.
"""
function LinearAlgebra.eigvals(G::SimpleGraph, mat::Function = adjacency)
    return eigvals(mat(G))
end


using LinearAlgebra
include("approx_inequality.jl")
# export ProximableFunction
export prox, prox!

# abstract type ProximableFunction end
export IndNonnegativeBallL2

struct IndNonnegativeBallL2{R <: Real}
    r::R
    function IndNonnegativeBallL2{R}(r::R) where {R <: Real}
        if r <= 0
            error("parameter r must be positive")
        else
            new(r)
        end
    end
end

is_convex(f::IndNonnegativeBallL2) = true
is_set(f::IndNonnegativeBallL2) = true

IndNonnegativeBallL2(r::R=1.0) where {R <: Real} = IndNonnegativeBallL2{R}(r)

function (f::IndNonnegativeBallL2)(x::AbstractArray{T}) where {R <: Real, T <: RealOrComplex{R}}
    if isapprox_le(norm(x), f.r, atol=eps(R), rtol=sqrt(eps(R)))
        for k in eachindex(x)
            if x[k] < 0
                return R(Inf)
            end
        end
        return R(0)
    end
    return R(Inf)
end

function prox!(y::AbstractArray{T}, f::IndNonnegativeBallL2, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    for k in eachindex(x)
        if x[k] < 0
            y[k] = R(0)
        else
            y[k] = x[k]
        end
    end

    scal = f.r/norm(y)
    if scal > 1
        return R(0)
    end
    for k in eachindex(y)
        y[k] = scal*y[k]
    end
    return R(0)
end

function prox(f::IndNonnegativeBallL2, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    y = copy(x)
    for k in eachindex(x)
        if x[k] < 0
            y[k] = R(0)
        else
            y[k] = x[k]
        end
    end

    scal = f.r/norm(y)
    if scal > 1
        return y, R(0)
    end
    for k in eachindex(y)
        y[k] = scal*y[k]
    end
    return y, R(0)
end

fun_name(f::IndNonnegativeBallL2) = "indicator of an L2 norm ball and nonnegative orthant"
fun_dom(f::IndNonnegativeBallL2) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndNonnegativeBallL2) = "x ↦ 0 if ||x|| ⩽ r & all(0 ⩽ x), +∞ otherwise"
fun_params(f::IndNonnegativeBallL2) = "r = $(f.r)"

function prox_naive(f::IndNonnegativeBallL2, x::AbstractArray{T}, gamma::R=R(1)) where {R <: Real, T <: RealOrComplex{R}}
    y = copy(x)
    for k in eachindex(x)
        if x[k] < 0
            y[k] = R(0)
        end
    end

    normx = norm(y)
    if normx > f.r
        y = (f.r/normx)*y
    end
    return y, R(0)
end


n = 1000
targets = rand(0:1, n)
scores = rand(n)
thres = 0.7
predicts = scores .>= thres

p  = sum(targets .== 1)
n  = sum(targets .== 0)
tp = sum(targets .* predicts)
tn = sum((1 .- targets) .* (1 .- predicts))
fp = sum((1 .- targets) .* predicts)
fn = sum(targets .* (1 .- predicts))

cm = ConfusionMatrix(p, n, tp, tn, fp, fn)
cm2 = ConfusionMatrix(2*p, 2*n, 2*tp, 2*tn, 2*fp, 2*fn)

set_encoding(OneZero())

encs = [
    OneZero(),
    OneMinusOne(),
    OneTwo(),
    OneVsOne(3,4),
    OneVsOne(:three,:four),
    OneVsOne("three","four"),
    OneVsRest(1, [2,3,4]),
    OneVsRest(:one, [:two, :three]),
    OneVsRest("one", ["two", "three"]),
    RestVsOne([1,2,3], 4),
    RestVsOne([:one, :two], :three),
    RestVsOne(["one", "two"], "three")
]

@testset "ConfusionMatrix +" begin
    @test cm + cm == cm2
end

@testset "ConfusionMatrix constructors for $enc encoding" for enc in encs
    global targets, predicts

    targets = recode.(current_encoding(), enc, targets)
    predicts = recode.(current_encoding(), enc, predicts)
    set_encoding(enc)

    @test ConfusionMatrix(targets, predicts) == cm
    @test ConfusionMatrix(enc, targets, predicts) == cm
    @test ConfusionMatrix(targets, scores, thres) == cm
    @test ConfusionMatrix(enc, targets, scores, thres) == cm
    @test ConfusionMatrix(targets, scores, [thres]) == [cm]
    @test ConfusionMatrix(enc, targets, scores, [thres]) == [cm]
end


#=
The cube, 41063625 (345^3), can be permuted to produce two other cubes: 56623104 (384^3) and 66430125 (405^3).
In fact, 41063625 is the smallest cube which has exactly three permutations of its digits which are also cube.

Find the smallest cube for which exactly five permutations of its digits are cube.
=#
include("projecteulerutils.jl")

function largest_permutation(i)
  nd(reverse(sort(reverse(digits(i)))))
end

function calc()
  i, max_permutations = 1, 5
  hash = Dict{Int64, Array{Int64}}()  #largest_permutation => [smallest_cube, count]
  while true
    current_cube = i ^ 3
    largest = largest_permutation(current_cube)
    v = get(hash, largest, [current_cube, 0])
    count = v[2] + 1
    hash[largest] = [v[1], count]
    if count == max_permutations
      return v[1]
    end
    i = i + 1
  end
end
@time println(calc())


#
# Denoise an SPD Example with Cyclic Proximal Point to minimize the
#
# L2-TV functional with anisotropic TV
#
# where the example is the same data and cost as for SPD_Image_CP_vs_CPPA
#
# This example and its data is used for comparison in Section 6.2 of
#
# > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal Núñez:
# > _Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds_,
# > arXiv: [1908.02022](https://arxiv.org/abs/1908.02022)
#
using Manopt, Manifolds
using Images, CSV, DataFrames, LinearAlgebra, JLD2
#
# Settings
experiment_name = "SPD_Image_CPPA"
export_orig = true # export and render input image
export_result = true # export resulting image
export_table = true # export a table of data recorded through the iterations
export_function_value = true # export the final cost function value for ChambollePock
asy_render_detail = 4 # resolution in asymptote, a high resolution is e.g. 4
results_folder = joinpath(@__DIR__, "Image_TV")
!isdir(results_folder) && mkdir(results_folder)

#
# Algorithm Parameters
α = 6.0
maxIterations = 4000

#
# Manifold and Data
f = artificial_SPD_image2(32)
pixelM = SymmetricPositiveDefinite(3)
if export_orig
    orig_filename = joinpath(results_folder, experiment_name * "-orig.asy")
    asymptote_export_SPD(orig_filename; data=f, scaleAxes=(7.5, 7.5, 7.5))
    render_asymptote(orig_filename; render=asy_render_detail)
end

#
# Build Problem for L2-TV & proximal maps
M = PowerManifold(pixelM, NestedPowerRepresentation(), size(f)...)
d = length(size(f))
rep(d) = (d > 1) ? [ones(Int, d)..., d] : d
fidelity(x) = 1 / 2 * distance(M, x, f)^2
Λ(x) = forward_logs(M, x) # on T_xN
prior(x) = norm(norm.(Ref(pixelM), repeat(x, rep(d)...), Λ(x)), 1)
cost(x) = (1 / α) * fidelity(x) + prior(x)
proxes = ((λ, x) -> prox_distance(M, λ, f, x, 2), (λ, x) -> prox_TV(M, α * λ, x, 1))

x0 = f
@time o = cyclic_proximal_point(
    M,
    cost,
    proxes,
    x0;
    λ=i -> 4 / i,
    stopping_criterion=StopAfterIteration(maxIterations),
    debug=[
        :Iteration,
        " | ",
        DebugProximalParameter(),
        " | ",
        :Change,
        " | ",
        :Cost,
        "\n",
        100,
        :Stop,
    ],
    record=[:Iteration, :Iterate, :Cost],
    return_options=true,
)
y = get_solver_result(o)
yRec = get_record(o)
#
# Results
if export_result
    result_filename = joinpath(results_folder, experiment_name * "-result.asy")
    asymptote_export_SPD(result_filename; data=y, scaleAxes=(7.5, 7.5, 7.5))
    render_asymptote(result_filename; render=asy_render_detail)
end
if export_table
    A = cat([y[1] for y in yRec], [y[3] for y in yRec]; dims=2)
    table_filename = joinpath(results_folder, experiment_name * "-recorded-cost.csv")
    CSV.write(table_filename, DataFrame(A); writeheader=false)
end
if export_function_value
    fctval_filename = joinpath(results_folder, experiment_name * "-cost.jld2")
    values = Dict("cost_function_value" => last(yRec)[3], "iterations" => length(yRec) - 1)
    save(fctval_filename, values)
end


"""
Problem 1:

Given a path given by an input array of commands (direction + integer),
compute the final position relative to the start.

Problem 2:
Same, but changing slightly the rules.

"""

using DelimitedFiles

# Read input
function readInput() :: Array{Tuple{String,Int}}
  data = readdlm("./inputs/day2.txt", ' ', String)
  data = map(x -> (x[1],parse(Int,x[2])), data[i,:] for i in 1:size(data)[1] )
  return data
end

# Converts direction from string to cartesian index.
function toIndex(s::String)
  if s == "up"
    return CartesianIndex(0,-1)
  elseif s == "down"
    return CartesianIndex(0,1)
  elseif s == "forward"
    return CartesianIndex(1,0)
  else
    return CartesianIndex(-1,0)
  end
end

# Follows a path given by an array of steps (distance + direction). Return final position.
function followPath(path:: Array{Tuple{String,Int}}) :: CartesianIndex{2}
  pos = CartesianIndex(0,0)
  for step in path
    pos += step[2] * toIndex(step[1])
  end
  return pos
end

# Modified version for the second part of the problem.
function followPath2(path:: Array{Tuple{String,Int}}) :: CartesianIndex{2}
  aim = 0
  pos = CartesianIndex(0,0)
  for step in path
    dir = toIndex(step[1])
    aim += step[2] * dir[2]
    pos += CartesianIndex(step[2] * dir[1] ,aim * step[2] * dir[1])
  end
  return pos
end

path = readInput()

pos = followPath(path)
println("Problem 1:")
println(pos[1]*pos[2])

pos = followPath2(path)
println("Problem 2:")
println(pos[1]*pos[2])


include("CG_Convergence_Utilities.jl")

# ==============================================================================
# Exact Arithematic error bound.
# ==============================================================================

function PerformCGFor(A, b, immitate_exact=true)
    cgm = ConjGradModified(A, b)
    if !immitate_exact
        TurnOffReorthgonalize(cgm)
    end
    ResidualNorm = GetResidualNorm(cgm)
    ResidualNorms = Vector{Float64}()
    Guesses = Vector{Vector{Float64}}()
    push!(ResidualNorms, ResidualNorm)
    while ResidualNorm > 1e-10
        ResidualNorm = cgm()
        push!(ResidualNorms, ResidualNorm)
        push!(Guesses, cgm.x)
    end
    return ResidualNorms, Guesses
end

function Run1()

return end

N = 512
A = GetUniformPSDMatrix(N)
b = rand(N)
ResidualNorms, Guesses = PerformCGFor(A, b)
ResidualEnergy = ResRelEnergyNorm(A, b,Guesses)
k = length(ResidualNorms)
ErrorsBound = [TheoreticalErrorBound(A, idx) for idx in 1: k]

fig1 = plot(log10.(ResidualEnergy), label="Relative Energy (exact)", legend=:bottomleft)
plot!(fig1, log10.(ErrorsBound), label="Theoretical Error Bound")
display(fig1)
SaveFigToCurrentScriptDir(fig1, "fig1.png")

# ==============================================================================
# Floating Points Error Bound.
# ==============================================================================

A = GetUniformPSDMatrix(N)
b = rand(N)
ResidualNorms, Guesses = PerformCGFor(A, b, false)
ResidualEnergy = ResRelEnergyNorm(A, b,Guesses)
k = length(ResidualNorms)
ErrorsBound = [TheoreticalErrorBound(A, idx) for idx in 1: k]

fig2 = plot(log10.(ResidualEnergy), label="Relative Energy (floats)", legend=:bottomleft)
plot!(fig2, log10.(ErrorsBound), label="Theoretical Error Bound(exact)")
display(fig2)
SaveFigToCurrentScriptDir(fig2, "fig2.png")


# ints_trap = map(hs) do h
#     x = y1:h:y2
#     -(g(x[1]) + g(x[end]))*h/2.0 + sum(g.(x)*h)
# end

function integration_scheme(x::AbstractVector{T}; scheme::Symbol = :trapezoidal, kws...) where T<:AbstractFloat
    if length(collect(x)) == 1
        return [0.]
    end

    if scheme == :trapezoidal
        trap_scheme(x; kws...)
    elseif scheme == :simpson
        simpson_scheme(x; kws...)
    else
        @warn("Integration scheme $scheme unknown, will use trapezoidal")
        trap_scheme(x; kws...)
    end
end

function trap_scheme(x::AbstractVector{T}; x0::T = first(x), xn::T = last(x)) where T<:AbstractFloat

    inds = axes(x,1)

    h = (x[inds[2]]-x[inds[1]])
    σs = similar(x)
    σs .= h

    σs[inds[1]] -= h/T(2)

    σs[end] -= h/T(2)
    # accounts for the ends
    σs[inds[1]] +=   (x[inds[1]] - x0)*(x[inds[2]] - x0 + h)/(2.0h)
    σs[inds[2]] +=  -(x0 - x[inds[1]])^T(2)/(T(2)*h)
    σs[end] +=  (xn - x[end])*(xn - x[end-1] + h)/(2h)
    σs[end-1] += - (xn - x[end])^T(2)/(T(2)*h)

    return σs
end

function simpson_scheme(x::AbstractVector{T}; x0::T = first(x), xn::T = last(x)) where T<:AbstractFloat
    inds = axes(x,1)

    if iseven(length(inds))
        @warn("Simpson's integration scheme is designed for an odd number of mesh points")
    end

    h = (x[inds[2]]-x[inds[1]])
    σs = similar(x).*zero(T)
    for j in 1:length(inds)
        σs[inds[j]] = iseven(j) ? 4.0 : 2.0
    end
    σs = σs.*(h/3.0)

    # accounts for the ends
    σs[inds[1]] = h/3.0
    σs[inds[1]] +=  (x[inds[1]] - x0)*(x[inds[2]] - x0 + h)/(2.0h)
    σs[inds[2]] +=  - (x0-x[inds[1]])^2.0/(2.0h)
    σs[inds[end]] = h/3.0
    σs[inds[end]] +=  (xn-x[inds[end]])*(xn - x[inds[end-1]] + h)/(2h)
    σs[inds[end-1]] += - (xn-x[inds[end]])^2.0/(2h)

    return σs
end


###########
# This gives an example of Gauss–Seidel
#
###########

using BlockBandedMatrices, BandedMatrices, ArrayLayouts, FillArrays, LazyBandedMatrices, LazyArrays

function finitedifference_2d(n)
    h = 1/n
    D² = BandedMatrix(0 => Fill(-2,n), 1 => Fill(1,n-1), -1 => Fill(1,n-1))/h^2
    D_xx = BandedBlockBandedMatrix(Kron(D², Eye(n)))
    D_yy = BandedBlockBandedMatrix(Kron(Eye(n), D²))
    D_xx + D_yy
end

function _gaussseidel(L, U, b, x=copy(b), y=copy(b), M=5)
    for _=1:M
        @view(y[1:end-1]) .= Mul(U , @view(x[2:end]))
        y[end] = 0
        x .= b .- y
        x .= Ldiv(L, x)
    end
    x
end

function _gaussseidel2(L, U, b, x=copy(b), y=copy(b), M=5)
    for _=1:M
        mul!(@view(y[1:end-1]) , U , @view(x[2:end]))
        y[end] = 0
        y .= b .- y
        x .= L\ y
    end
    x
end

function gaussseidel(A, b, M)
    n = Int(sqrt(length(b)))
    L = LowerTriangular(A)
    U = BandedBlockBandedMatrix(UpperTriangular(@view(A[1:end-1,2:end])), ([fill(n,n-1); n-1], [n-1; fill(n,n-1)]),
                                            (0,1), (0,1))
    x = copy(b)
    _gaussseidel(L,U, b, x, copy(x), M) # 1.6s
    x
end

n = 1000;
    Δt = (1/n^2)/4; @time Δ = finitedifference_2d(n);
    @time A = I - Δt*Δ  # 1m x 1m discretization
    b = randn(n^2)
    L = LowerTriangular(A)
    @time U = BandedBlockBandedMatrix(UpperTriangular(@view(A[1:end-1,2:end])), ([fill(n,n-1); n-1], [n-1; fill(n,n-1)]),
                                            (0,1), (0,1))
    x = copy(b)
    y = copy(x)
    @time _gaussseidel(L,U, b, x, y, 20) # 0.4s

    @time L̃ = sparse(BandedBlockBandedMatrix(L))
    @time Ũ = sparse(U)
    x = copy(b)
    y = copy(x)
    @time _gaussseidel2(L̃,Ũ, b, x, y, 20)

Ã = sparse(A)
@time qr(Ã);
@time Ã \b ;

A = randn(n,n)

n = 1000;
    As = [randn(2k+1,k) for k=1:n];

    @time for A in As
        qr(A)
    end
n = 1000;
    # As = [randn(2k+1,k) for k=1:n];

    @time for A in As
        lu(A)
    end


using DelayEmbeddings
using Test, Random, DelimitedFiles

println("\nTesting pecuzal_method.jl...")
@testset "PECUZAL" begin

## Test Lorenz example
# For comparison reasons using Travis CI we carry out the integration on a UNIX
# OS and save the resulting time series
# See https://github.com/JuliaDynamics/JuliaDynamics for the storage of
# the time series used for testing
#
# u0 = [0, 10.0, 0.0]
# lo = Systems.lorenz(u0; σ=10, ρ=28, β=8/3)
# tr = trajectory(lo, 100; Δt = 0.01, Ttr = 100)
tr = readdlm(joinpath(tsfolder, "test_time_series_lorenz_standard_N_10000_multivariate.csv"))
tr = Dataset(tr)

@testset "Univariate example" begin

    s = vec(tr[:, 1]) # input timeseries = x component of lorenz
    w = estimate_delay(s, "mi_min")
    Tmax = 100

    @time Y, τ_vals, ts_vals, Ls , εs = pecuzal_embedding(s[1:5000];
                                        τs = 0:Tmax , w = w, L_threshold = 0.05, econ=true)

    @test -1.08 < Ls[1] < -1.04
    @test -0.51 < Ls[2] < -0.47
    @test -0.18 < Ls[3] < -0.14

    @test τ_vals[2] == 19
    @test τ_vals[3] == 96
    @test τ_vals[4] == 80
    @test length(ts_vals) == 4

    @time Y, τ_vals, ts_vals, Ls , εs = pecuzal_embedding(s;
                                        τs = 0:Tmax , w = w, L_threshold = 0.05)
    @test -0.95 < Ls[1] < -0.91
    @test -0.48 < Ls[2] < -0.44

    @test τ_vals[2] == 18
    @test τ_vals[3] == 9
    @test length(ts_vals) == 3

    YY1 = DelayEmbeddings.hcat_lagged_values(Dataset(s), s, 21)
    YY2 = DelayEmbeddings.hcat_lagged_values(YY1, s, 13)
    @test length(YY1) == length(YY2)
    @test YY1 == YY2[:,1:2]

end

@testset "Multivariate example" begin
    ## Test of the proposed PECUZAL-embedding-method (multivariate case)

    w1 = estimate_delay(tr[:,1], "mi_min")
    w2 = estimate_delay(tr[:,2], "mi_min")
    w3 = estimate_delay(tr[:,3], "mi_min")
    w = w1
    Tmax = 100

    @time Y, τ_vals, ts_vals, Ls , ε★ = pecuzal_embedding(tr[1:5000,:];
                                        τs = 0:Tmax , w = w, econ = true)

    @test length(ts_vals) == 4
    @test ts_vals[2] == ts_vals[3] == ts_vals[4] == 1
    @test ts_vals[1] == 3
    @test τ_vals[1] == 0
    @test τ_vals[2] == 9
    @test τ_vals[3] == 64
    @test τ_vals[4] == 53
    @test -1.40 < Ls[1] < -1.36
    @test -0.76 < Ls[2] < -0.72
    @test -0.1 < Ls[3] < -0.06

    @time Y, τ_vals, ts_vals, Ls , ε★ = pecuzal_embedding(tr[1:5000,:];
                                        τs = 0:Tmax , w = w, econ = true, L_threshold = 0.2)
    @test -1.40 < Ls[1] < -1.36
    @test -0.76 < Ls[2] < -0.72
    @test size(Y,2) == 3

end

@testset "Dummy example" begin
    # Dummy input
    Random.seed!(1234)
    d1 = randn(1000)
    d2 = rand(1000)
    Tmax = 50
    dummy_set = Dataset(hcat(d1,d2))

    w1 = estimate_delay(d1, "mi_min")
    w2 = estimate_delay(d2, "mi_min")
    w = w1

    @time Y, τ_vals, ts_vals, Ls , ε★ = pecuzal_embedding(dummy_set;
                                    τs = 0:Tmax , w = w, econ = true)

    @test size(Y,2) == 1
end

end


# A simple example to demonstrate the use of cross validation
#
# Here, we consider a simple model: using a mean vector to represent
# a set of samples. The goodness of the model is assessed in terms
# of the RMSE (root-mean-square-error) evaluated on the testing set
#

using MLBase

# functions

function compute_center(X::Matrix{Float64})
	c = vec(mean(X, 2))
	@printf("training on %d samples => (%.4f, %.4f)\n",
		size(X,2), c[1], c[2])
	return c
end

function compute_rmse(c::Vector{Float64}, X::Matrix{Float64})
	v = sqrt(mean(sum(abs2(X .- c),1)))
	@printf("RMSE on test set: %.6f\n\n", v)
	return v
end

# data

const n = 200
const data = [2., 3.] .+ randn(2, n)

# cross validation

(c, v, inds) = cross_validate(
	inds -> compute_center(data[:, inds]),     		# training function
	(c, inds) -> compute_rmse(c, data[:, inds]),    # evaluation function
	n,    			# total number of samples
	Kfold(n, 5), 	# cross validation plan: 5-fold
	Reverse)		# smaller score indicates better model

# display results

@printf("best model = (%.4f, %.4f), score = %.6f\n", c[1], c[2], v)


"""
    explicit midpoint integrator
"""
function midpoint(model, x, u, w, h)
    x + h * f(model, x + 0.5 * h * f(model, x, u, w), u, w)
end

"""
    implicit midpoint integrator
"""
function midpoint_implicit(model, x⁺, x, u, w, h)
    x⁺ - (x + h * f(model, 0.5 * (x + x⁺), u, w))
end

"""
    explicit RK3 integrator
"""
function rk3(model, z, u, w, h)
    k1 = k2 = k3 = zero(z)
    k1 = h * f(model, z, u, w)
    k2 = h * f(model, z + 0.5 * k1, u, w)
    k3 = h * f(model, z - k1 + 2.0 * k2, u, w)
    z + (k1 + 4.0 * k2 + k3) / 6.0
end

"""
    discrete explicit dynamics
"""
function fd(model, x, u, w, h, t)
    midpoint(model, x, u, w, h)
end

"""
    discrete implicit dynamics
"""
function fd(model, x⁺, x, u, w, h, t)
    midpoint_implicit(model, x⁺, x, u, w, h)
end


module SimplePCHIP
#
# Simple PCHIP implementation following Fritsch and Carlson, SIAM J. NUMER. ANAL. 17 (1980) 238-246.
# Derivative calculated in a fashion similar to SciPy's PchipInterpolate
#

export interpolate

struct _pchip
    N :: Int
    xs :: Vector{Float64}
    ys :: Vector{Float64}
    ds :: Vector{Float64}
end

ϕ(t) = 3t^2 - 2t^3
ψ(t) = t^3 - t^2

function _interp(pchip :: _pchip, x :: Number)
    i = _pchip_index(pchip, x)
    x1, x2 = pchip.xs[i:i+1]
    y1, y2 = pchip.ys[i:i+1]
    d1, d2 = pchip.ds[i:i+1]
    h = x2 - x1

    (y1 * ϕ((x2-x)/h)
     + y2 * ϕ((x-x1)/h)
     - d1*h * ψ((x2-x)/h)
     + d2*h * ψ((x-x1)/h))
end

function _pchip_index(pchip :: _pchip, x)
    N = pchip.N
    if N < 200  # Approximate performance cross-over on my old intel i7-3517U
        i = _pchip_index_linear_search(pchip, x)
    else
        i = _pchip_index_bisectional_search(pchip, x)
    end
    if i == N
        # Treat right endpoint as part of rightmost interval
        @assert (x ≈ pchip.xs[N])
        i = N-1
    end
    i
end

function _pchip_index_linear_search(pchip :: _pchip, x)
    xmin = pchip.xs[1]
    @assert (x >= xmin)

    i = 1
    N = pchip.N
    while i < N  &&  x >= pchip.xs[i+1]
        i = i + 1
    end
    i
end

function _pchip_index_bisectional_search(pchip :: _pchip, x)
    N = pchip.N
    imin, imax = 1, N
    xmin = pchip.xs[imin]
    xmax = pchip.xs[imax]
    @assert (x >= xmin && x <= xmax)

    i = imin + div(imax - imin + 1, 2)
    while imin < imax
        if x < pchip.xs[i]
            imax = i - 1
        elseif x >= pchip.xs[i+1]
            imin = i + 1
        else
            break
        end
        i = imin + div(imax - imin + 1, 2)
    end
    i
end

"Similar to how SciPy's PCHIP does it"
function _initial_ds_scipy(xs, ys)
    h(i) = xs[i+1]-xs[i]
    Δ(i) = (ys[i+1]-ys[i]) / h(i)

    N = length(xs)
    ds = similar(xs)
    if N == 2
        ds[:] .= Δ(1)
    else
        Δl = Δ(1)
        hl = h(1)
        for i ∈ 2:N-1
            Δr = Δ(i)
            hr = h(i)
            if sign(Δl) != sign(Δr) || Δl ≈ 0.0 || Δr ≈ 0.0
                ds[i] = 0.0
            else
                wl = 2hl + hr
                wr = hl + 2hr
                ds[i] = (wl + wr) / (wl/Δl + wr/Δr)
            end
            Δl = Δr
            hl = hr
        end
        ds[1] = _edge_derivative(h(1), h(2), Δ(1), Δ(2))
        ds[N] = _edge_derivative(h(N-1), h(N-2), Δ(N-1), Δ(N-2))
    end
    ds
end

function _edge_derivative(h1, h2, Δ1, Δ2)
    d = ((2h1 + h2)*Δ1 - h2*Δ2) / (h1 + h2)
    if sign(d) != sign(Δ1)
        d = 0.0
    elseif sign(Δ1) != sign(Δ2)  &&  abs(d) > abs(3Δ1)
        d = 3Δ1
    end
    d
end

function interpolate(xs, ys)
    xs_ = [x for x ∈ xs]
    ys_ = [y for y ∈ ys]
    _assert_xs_ys(xs_, ys_)
    ds = _initial_ds_scipy(xs_, ys_)
    pchip = _pchip(length(xs_), xs_, ys_, ds)

    x -> _interp(pchip, x)
end

function _assert_xs_ys(xs, ys)
    N = length(xs)
    @assert (N > 1)
    @assert (N == length(ys))
    assert_monotonic_increase(xs)
end

function assert_monotonic_increase(xs)
    foldl((a,b) -> (@assert (a < b); b), xs)
end

end  # module


using AdFem

for degree in [1, 2]
    m = 100; n = 100; h = 1/n
    mesh = Mesh(m, n, h, degree = degree)

    function f(x, y)
        2 * x^2 + y
    end

    f_fem = eval_f_on_fem_pts(f, mesh)
    f_gauss = fem_to_gauss_points(constant(f_fem), mesh)
    sess = Session(); init(sess)
    f_gauss = run(sess, f_gauss)

    f_gauss2 = eval_f_on_gauss_pts(f, mesh)

    @show maximum(abs.(f_gauss .- f_gauss2))
end


using QuadGK
using FastClosures
"""
	WharmonicR(t, tf, coeffs)
Harmonic correction Hamiltonian in the rotating frame.
"""
function WharmonicR(t, tf, coeffs)
	ε1, ε2, Δ = coeffs[1], coeffs[2], coeffs[3]
	cos2t, sin2t = cos(2*t), sin(2*t)
	g = (ε1*cos2t + ε2*sin2t)*(1 - cos(2*pi*t/tf))
	W = zeros(Float64, 3)
	W[1] = g*cos2t
	W[2] = g*sin2t
	W[3] = Δ/2
	return W
end
"""
	WharmonicI(t, tf, coeffs)
Harmonic correction Hamiltonian in the interaction picture.
"""
function WharmonicI(t, tf, coeffs)
	ε1, ε2, Δ = coeffs[1], coeffs[2], coeffs[3]
	cos2t, sin2t = cos(2*t), sin(2*t)
	g = (ε1*cos2t + ε2*sin2t)*(1 - cos(2*pi*t/tf))
	Λt = Λ(t,tf)
	cosh2Λ, sinh2Λ = cosh(2*Λt), sinh(2*Λt)
	W = zeros(Float64, 3)
	W[1] = g*cos2t*cosh2Λ + (Δ/2 + g)*sinh2Λ
	W[2] = g*sin2t
	W[3] = g*cos2t*sinh2Λ + (Δ/2 + g)*cosh2Λ
	return W
end
"""
	correction_matrix(tf)
Correction matrix for the harmonic correction.
"""
function correction_matrix(tf)
	M = zeros(Float64, 3, 3)
	h(t) = 1 - cos(2*pi*t/tf)

	func11 = @closure t-> h(t)*((1 + cos(4*t))*cosh(2*Λ(t,tf))
								+ 2*cos(2*t)*sinh(2*Λ(t,tf)))/2
	M[1, 1], err = quadgk(func11, 0.0, tf, rtol=1e-10, atol=1e-14)
	func12 = @closure t-> h(t)*(sin(4*t)*cosh(2*Λ(t,tf))
								+ 2*sin(2*t)*sinh(2*Λ(t,tf)))/2
	M[1, 2], err = quadgk(func12, 0.0, tf, rtol=1e-10, atol=1e-14)
	func13 = @closure t-> sinh(2*Λ(t,tf))/2
	M[1, 3], err = quadgk(func13, 0.0, tf, rtol=1e-10, atol=1e-14)
	func21 = t-> h(t)*sin(4*t)/2
	M[2, 1], err = quadgk(func21, 0.0, tf, rtol=1e-10, atol=1e-14)
	func22 = t-> h(t)*(1 - cos(4*t))/2
	M[2, 2], err = quadgk(func22, 0.0, tf, rtol=1e-10, atol=1e-14)
	func31 = @closure t-> h(t)*((1 + cos(4*t))*sinh(2.0*Λ(t,tf))
								+ 2.0*cos(2*t)*cosh(2.0*Λ(t,tf)))/2
	M[3, 1], err = quadgk(func31, 0.0, tf, rtol=1e-10, atol=1e-14)
	func32 = @closure t-> h(t)*(sin(4*t)*sinh(2.0*Λ(t,tf))
								+ 2.0*sin(2*t)*cosh(2.0*Λ(t,tf)))/2
	M[3, 2], err = quadgk(func32, 0.0, tf, rtol=1e-10, atol=1e-14)
	func33 = @closure t-> cosh(2*Λ(t,tf))/2
	M[3, 3], err = quadgk(func33, 0.0, tf, rtol=1e-10, atol=1e-14)

	return M
end
"""
	coefficients_W(Omega, M)
Calculate the coefficients of the harmonic correction.
"""
function coefficients_W(Omega, M)
	ih_vec = Omega
	coeffs = inv(M)*ih_vec
	return coeffs
end
"""
	get_coeffs(M, coeffs, order, tf)
Calculate the coefficients of the Harmonic correction.
"""
function get_coeffs(M, coeffs, order, tf)
	A = @closure (t-> -V_I(t, tf) .- WharmonicI(t, tf, coeffs))

    Ndims = 3
    Omegaf = solveOmega(tf, order, A, comm!, Ndims)

    vec = sum(Omegaf, dims=2)
    coeffs = coefficients_W(vec, M)
    return coeffs
end
"""
	get_coeffs(M, order, tf)
Calculate the coefficients of the Harmonic correction.
"""
function get_coeffs(M, order, tf)
    coeffs = zeros(3)
    for k = 1: order
		coeffs2 = get_coeffs(M, coeffs, order, tf)
        coeffs = coeffs .+ coeffs2
    end
    return coeffs
end


######### Stan program example  ###########

using Mamba, Stan

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  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" => 1, "y" => [0]),
    Dict("N" => 1, "y" => [0]),
    Dict("N" => 1, "y" => [0]),
    Dict("N" => 1, "y" => [0]),
  ]

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

  global stanmodel, rc, sim1, sim
  stanmodel = Stanmodel(num_samples=1200, thin=2, name="bernoulli", model=bernoullimodel);

  println("\nStanmodel that will be used:")
  stanmodel |> display
  println()
  println("Input observed data dictionary:")
  bernoullidata |> display
  println()

  rc, sim1 = stan(stanmodel, bernoullidata, ProjDir, diagnostics=false, CmdStanDir=CMDSTAN_HOME);

  if rc == 0
    ## Subset Sampler Output to variables suitable for describe().
    sim = sim1[1:size(sim1, 1), monitor, 1:size(sim1, 3)]
    describe(sim)
    println()

    ## Brooks, Gelman and Rubin Convergence Diagnostic
    try
      gelmandiag(sim, mpsrf=true, transform=true) |> display
    catch e
      #println(e)
      gelmandiag(sim, mpsrf=false, transform=true) |> display
    end
    println()

    ## Geweke Convergence Diagnostic
    gewekediag(sim) |> display

    ## Highest Posterior Density Intervals
    hpd(sim) |> display

    ## Cross-Correlations
    cor(sim) |> display

    ## Lag-Autocorrelations
    autocor(sim) |> display

    ## Deviance Information Criterion
    #dic(sim) |> display

    ## Plotting
    p = plot(sim, [:trace, :mean, :density, :autocor], legend=true);
    draw(p, ncol=4, filename="$(stanmodel.name)-summaryplot", fmt=:svg)
    draw(p, ncol=4, filename="$(stanmodel.name)-summaryplot", fmt=:pdf)

  	# Pairwise contour plots
  	simmod = ModelChains(sim,Mamba.Model())
  	p = plot(simmod, :contour)
  	draw(p, nrow=2, ncol=2, filename="$(stanmodel.name)-contourplot.svg")
  	draw(p, nrow=2, ncol=2, filename="$(stanmodel.name)-contourplot.pdf", fmt=:pdf)

  end # rc == 0
end # cd


#! /usr/bin/julia

# Rosetta Code, Permutation test

delta_mean{T<:Real}(a::Array{T,1}, b::Array{T,1}) = mean(a) - mean(b)

function bifurcate{T<:Integer}(a::AbstractVector, sel::Array{T,1})
    x = a[sel]
    asel = trues(length(a))
    asel[sel] = false
    y = a[asel]
    return (x, y)
end

function perm_sig_test{T<:Real}(treat::Array{T,1}, control::Array{T,1})
    base_effect = delta_mean(treat, control)
    pool = [treat, control]
    tlen = length(treat)
    plen = length(pool)
    better = 0
    worse = 0
    for s in combinations(1:plen, tlen)
        (t, c) = bifurcate(pool, s)
        if base_effect < delta_mean(t, c)
            better += 1
        else
            worse += 1
        end
    end
    return (better, worse)
end

treat = [85, 88, 75, 66, 25, 29, 83, 39, 97]
control = [68, 41, 10, 49, 16, 65, 32, 92, 28, 98]

(better, worse) = perm_sig_test(treat, control)

tot = better + worse

println("Permutation test using the following data:")
println("Treated:  ", treat)
println("Control:  ", control)

println()
println("There are ", tot, " different permuted groups of these data.")
print(@sprintf("%8d, %5.2f%% ", better, 100better/tot))
println("showed better than actual results.")
print(@sprintf("%8d, %5.2f%% ", worse, 100worse/tot))
println("showed equalivalent or worse results.")


@inline permute(i, N) = isodd(i) ? floor(Int, i/2) + 1 : N - floor(Int, (i-1)/2)

@inline unpermute(i, N) = i <= ceil(N/2) ? 2i-1 : 2(N-i+1)

function dct_makhoul_1d(A::CuArray)
    B = similar(A)
    N = length(A)

    for k in 1:N
        B[permute(k, N)] = A[k]
    end

    B = CUDA.CUFFT.fft(B)

    for k in 1:N
        B[k] = 2 * ω(4N, k-1) * B[k]
    end

    return real(B)
end

function idct_makhoul_1d(A::CuArray)
    B = similar(A, complex(eltype(A)))
    N = length(A)

    B[1] = 1/2 * ω(4N, 0) * A[1]
    for k in 2:N
        B[k] = ω(4N, 1-k) * A[k]
    end

    B = CUDA.CUFFT.ifft(B)

    C = similar(A)
    for k in 1:N
        C[unpermute(k, N)] = real(B[k])
    end

    return C
end

function dct_makhoul_2d(A::CuArray)
    Nx, Ny = size(A)

    # DCT along dimension 1

    B = similar(A)

    for j in 1:Ny, i in 1:Nx
        B[permute(i, Nx), j] = A[i, j]
    end

    B = CUDA.CUFFT.fft(B, 1)

    for j in 1:Ny, i in 1:Nx
        B[i, j] = 2 * ω(4Nx, i-1) * B[i, j]
    end

    B = real(B)

    # DCT along dimension 2

    C = similar(A)

    for j in 1:Ny, i in 1:Nx
        C[i, permute(j, Ny)] = B[i, j]
    end

    C = CUDA.CUFFT.fft(C, 2)

    for j in 1:Ny, i in 1:Nx
        C[i, j] = 2 * ω(4Ny, j-1) * C[i, j]
    end

    return real(C)
end

function idct_makhoul_2d(A::CuArray)
    Nx, Ny = size(A)

    # IDCT along dimension 1

    B = similar(A, complex(eltype(A)))

    for j in 1:Ny
        B[1, j] = 1/2 * ω(4Nx, 0) * A[1, j]
    end

    for j in 1:Ny, i in 2:Nx
        B[i, j] = ω(4Nx, 1-i) * A[i, j]
    end

    B = CUDA.CUFFT.ifft(B, 1)

    C = similar(A)
    for j in 1:Ny, i in 1:Nx
        C[unpermute(i, Nx), j] = real(B[i, j])
    end

    # IDCT along dimension 2

    D = similar(A, complex(eltype(A)))

    for i in 1:Nx
        D[i, 1] = 1/2 * ω(4Ny, 0) * C[i, 1]
    end

    for j in 2:Ny, i in 1:Nx
        D[i, j] = ω(4Ny, 1-j) * C[i, j]
    end

    D = CUDA.CUFFT.ifft(D, 2)

    E = similar(A)
    for j in 1:Ny, i in 1:Nx
        E[i, unpermute(j, Ny)] = real(D[i, j])
    end

    return E
end