using SparseArrays, DiffEqOperators, LinearAlgebra, Random,
      Test, BandedMatrices, FillArrays

# Analytic solutions to higher order operators.
# Do not modify unless you are completely certain of the changes.

# 4th derivative, 4th order
function fourth_deriv_approx_stencil(N)
    A = zeros(N,N+2)
    A[1,1:8] = [3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3]
    A[2,1:8] = [2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6]

    A[N-1,N-5:end] = reverse([2/3 -11/6 0.0 31/6 -22/3 4.5 -4/3 1/6], dims=2)
    A[N,N-5:end] = reverse([3.5 -56/3 42.5 -54.0 251/6 -20.0 5.5 -2/3], dims=2)

    for i in 3:N-2
        A[i,i-2:i+4] = [-1/6 2.0 -13/2 28/3 -13/2 2.0 -1/6]
    end
    return A
end

function second_deriv_fourth_approx_stencil(N)
    A = zeros(N,N+2)
    A[1,1:6] = [5/6 -15/12 -1/3 7/6 -6/12 5/60]
    A[N,N-3:end] = [1/12 -6/12 14/12 -4/12 -15/12 10/12]
    for i in 2:N-1
        A[i,i-1:i+3] = [-1/12 4/3 -5/2 4/3 -1/12]
    end
    return A
end

function second_derivative_stencil(N)
  A = zeros(N,N+2)
  for i in 1:N, j in 1:N+2
      (j-i==0 || j-i==2) && (A[i,j]=1)
      j-i==1 && (A[i,j]=-2)
  end
  A
end

# analytically derived central difference operators for the irregular grid:
# [0.0, 0.08, 0.1, 0.15, 0.19, 0.26, 0.29], and spacing:
# dx = [0.08, 0.02, 0.05, 0.04, 0.07, 0.03]
# with convention that the first number is the derivative order
# and the second is the approximation order

function analyticCtrOneTwoIrr()
    A = zeros(5,7)
    A[1,1:3] = [-5. / 2., -75. / 2., 40.]
    A[2,2:4] = [-250. / 7., 30., 40. / 7.]
    A[3,3:5] = [-80. / 9., -5., 125. / 9.]
    A[4,4:6] = [-1225. / 77., 825. / 77., 400. / 77.]
    A[5,5:7] = [-30. / 7., -400. / 21., 70. / 3.]
    A
end

function analyticCtrTwoTwoIrr()
    A = zeros(5,7)
    A[1,1:3] = [250., -1250., 1000.]
    A[2,2:4] = [10000. / 7., -2000., 4000. / 7.]
    A[3,3:5] = [4000. / 9., -1000., 5000. / 9.]
    A[4,4:6] = [454. + 42. / 77., -5000. / 7., 20000. / 77.]
    A[5,5:7] = [6000. / 21., -20000. / 21., 2000. / 3.]
    A
end

function analyticCtrTwoFourIrr()
    A = zeros(3,7)
    A[1,1:5] = [14. + 12012. / 13167., 1542. + 2736. / 13167.,
                -2288. - 11704. / 13167., 838. + 1254. / 13167.,
                -106. - 4298. / 13167.]
    A[2,2:6] = [-223. - 461. / 693., 847. + 154. / 693.,
                -1311. - 477. / 693., 699. + 593. / 693.,
                -11. - 502. / 693.]
    A[3,3:7] = [2. + 12166. / 13167., 538. + 12654. / 13167.,
                -912. - 9196. / 13167., 508. + 8664. / 13167.,
                -137. - 11121. / 13167.]
    A
end

function analyticCtrFourTwoIrr()
    A = zeros(3,7)
    A[1,1:5] = [462000000. / 4389.,
                -8550000000. / 4389.,
                11704000000. / 4389.,
                -5016000000. / 4389.,
                1400000000. / 4389.]
    A[2,2:6] = [200000000. / 231.,
                -385000000. / 231.,
                360000000. / 231.,
                -200000000. / 231.,
                25000000. / 231.]
    A[3,3:7] = [770000000. / 4389.,
                -3420000000. / 4389.,
                4180000000. / 4389.,
                -2850000000. / 4389.,
                1320000000 / 4389.]
    A
end

function convert_by_multiplication(::Type{Array}, A::AbstractDerivativeOperator{T}, N::Int=A.dimension) where T
    @assert N >= A.stencil_length # stencil must be able to fit in the matrix
    mat = zeros(T, (N, N+2))
    v = zeros(T, N+2)
    for i=1:N+2
        v[i] = one(T)
        #=
            calculating the effect on a unit vector to get the matrix of transformation
            to get the vector in the new vector space.
        =#
        mul!(view(mat,:,i), A, v)
        v[i] = zero(T)
    end
    return mat
end

# Tests the corrrectness of stencils, along with concretization.
# Do not modify the following test-set unless you are completely certain of your changes.
@testset "Correctness of Uniform Stencils" begin
    N = 20
    L = CenteredDifference(4,4, 1.0, N)
    correct = fourth_deriv_approx_stencil(N)

    # Check that stencils (according to convert_by_multiplication) agree with correct
    @test convert_by_multiplication(Array, L, N) ≈ correct

    # Check that concretization agrees correct
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct

    L = CenteredDifference(2,4, 1.0, N)
    correct = second_deriv_fourth_approx_stencil(N)

    # Check that stencils (according to convert_by_multiplication) agree with correct
    @test convert_by_multiplication(Array, L, N) ≈ correct

    # Check that concretization agrees correct
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct
end

@testset "Correctness of Non-Uniform Stencils" begin
    x = [0., 0.08, 0.1, 0.15, 0.19, 0.26, 0.29]
    nx = length(x)
    dx = diff(x)

    # Second-Order First Derivative
    L = CenteredDifference(1, 2, dx, nx-2)
    correct = analyticCtrOneTwoIrr()

    # Check that stencils agree with correct
    for (i,coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i,correct[i,:] .!= 0.]
    end
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct

    # Second-Order Second Derivative
    L = CenteredDifference(2, 2, dx, nx-2)
    correct = analyticCtrTwoTwoIrr()

    # Check that stencils agree with correct
    for (i,coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i,correct[i,:] .!= 0.]
    end
    @test Array(L) ≈ correct
    @test sparse(L) ≈ correct
    @test BandedMatrix(L) ≈ correct

    # Fourth-Order Second Derivative
    L = CenteredDifference(2, 4, dx, nx-2)
    correct = analyticCtrTwoFourIrr()

    # Check that stencils agree with correct
    for (i,coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i,correct[i,:] .!= 0.]
    end
    @test Array(L)[2:end-1,:] ≈ correct
    @test sparse(L)[2:end-1,:] ≈ correct
    @test BandedMatrix(L)[2:end-1,:] ≈ correct

    # Second-Order Fourth Derivative
    L = CenteredDifference(4, 2, dx, nx-2)
    correct = analyticCtrFourTwoIrr()

    # Check that stencils agree with correct
    for (i,coefs) in enumerate(L.stencil_coefs)
        @test Array(coefs) ≈ correct[i,correct[i,:] .!= 0.]
    end
    @test Array(L)[2:end-1,:] ≈ correct
    @test sparse(L)[2:end-1,:] ≈ correct
    @test BandedMatrix(L)[2:end-1,:] ≈ correct
end

# tests for full and sparse function
@testset "Full and Sparse functions:" begin
    N = 10
    d_order = 2
    approx_order = 2
    correct = second_derivative_stencil(N)
    A = CenteredDifference(d_order,approx_order,1.0,N)

    @test convert_by_multiplication(Array,A,N) == correct
    @test Array(A) == second_derivative_stencil(N)
    @test sparse(A) == second_derivative_stencil(N)
    @test BandedMatrix(A) == second_derivative_stencil(N)
    @test opnorm(A, Inf) == opnorm(correct, Inf)

    # testing higher derivative and approximation concretization
    N = 20
    d_order = 4
    approx_order = 4
    A = CenteredDifference(d_order,approx_order,1.0,N)
    correct = convert_by_multiplication(Array,A,N)

    @test Array(A) ≈ correct
    @test sparse(A) ≈ correct
    @test BandedMatrix(A) ≈ correct

    N = 26
    d_order = 8
    approx_order = 8
    A = CenteredDifference(d_order,approx_order,1.0,N)
    correct = convert_by_multiplication(Array,A,N)

    @test Array(A) ≈ correct
    @test sparse(A) ≈ correct
    @test BandedMatrix(A) ≈ correct

    # testing correctness of multiplication
    N = 1000
    d_order = 4
    approx_order = 10
    y = collect(1:1.0:N+2).^4 - 2*collect(1:1.0:N+2).^3 + collect(1:1.0:N+2).^2;
    y = convert(Array{BigFloat, 1}, y)

    A = CenteredDifference(d_order,approx_order,one(BigFloat),N)
    correct = convert_by_multiplication(Array,A,N)
    @test Array(A) ≈ correct
    @test sparse(A) ≈ correct
    @test BandedMatrix(A) ≈ correct
    @test A*y ≈ Array(A)*y
end

@testset "Indexing tests" begin
    N = 1000
    d_order = 4
    approx_order = 10

    A = CenteredDifference(d_order,approx_order,1.0,N)
    @test A[1,1] == Array(A)[1,1]
    @test A[10,20] == 0

    correct = Array(A)
    for i in 1:N
        @test A[i,i] == correct[i,i]
    end

    # Indexing Tests
    N = 1000
    d_order = 2
    approx_order = 2

    A = CenteredDifference(d_order,approx_order,1.0,N)
    M = Array(A,1000)
    @test A[1,1] == M[1,1]
    @test A[1:4,1] == M[1:4,1]
    @test A[5,2:10] == M[5,2:10]
    @test A[60:100,500:600] == M[60:100,500:600]
end

@testset "Operations on matrices" begin
    N = 51
    M = 101
    d_order = 2
    approx_order = 2

    xarrs = Vector{Union{Array,StepRangeLen}}(undef,0)
    yarrs = Vector{Union{Array,StepRangeLen}}(undef,0)
    push!(xarrs, range(0,stop=1,length=N))
    push!(yarrs, range(0,stop=1,length=M))
    push!(xarrs, vcat(xarrs[1][1:floor(Int,N/2)].^0.2020, xarrs[1][ceil(Int,N/2):end].^2.015))
    push!(yarrs, vcat(yarrs[1][1:floor(Int,M/2)].^1.793, yarrs[1][ceil(Int,M/2):end].^2.019))

    for (i,xarr,yarr) in zip(1:2,xarrs,yarrs)
        if i == 1
            dx = xarr[2] - xarr[1]
            dy = yarr[2] - yarr[1]
        elseif i == 2
            dx = diff(xarr)
            dy = diff(yarr)
        end
        F = [x^2+y for x = xarr, y = yarr]

        A = CenteredDifference(d_order,approx_order,dx,length(xarr)-2)
        B = i == 1 ? CenteredDifference(d_order,approx_order,dy,length(yarr)) :
            CenteredDifference(d_order,approx_order,dy,length(yarr)-2)

        @test A*F ≈ 2*ones(N-2,M) atol=1e-9

        # Operators are defined such that
        # B must be applied to F from the left
        @test all(abs.(B * F') .<= 1e-8)
        @test all(abs.(A * (B * F')') .<= 1e-4)

        G = [x^2+y^2 for x = xarr, y = yarr]

        @test A*G ≈ 2*ones(N-2,M) atol=1e-9
        @test B * G' ≈ 2 * ones(M-2,N) atol=1e-8
        @test all(abs.(A * (B * G')') .<= 1e-4)
    end
end

@testset "Scalar multiplication with operators" begin
    A = CenteredDifference(2, 2, 10., 3)

    scalar_ans = [0.033 -0.066 0.033  0.     0.;
                  0.    0.033  -0.066 0.033  0.;
                  0.    0.     0.033  -0.066 0.033]
    @test Array(3.3 * A) ≈ scalar_ans
end

@testset "Left-multiplying operators with a vector of coefficients" begin
    A = CenteredDifference(2, 2, 1., 3)
    B = UpwindDifference(1, 1, 1., 3, 1.)
    C = UpwindDifference(1, 1, 1., 3, [2., 3., 4.])
    c = [1., 1., 3.]
    cA =  c * A
    cB =  c * B
    cC =  c * C

    @test cA.coefficients == c
    @test cB.coefficients == c
    @test cC.coefficients == [2., 3., 12.]

    x = [1., 2., 3., 4., 5.]
    @test (A * x) .* c ≈ cA * x
    @test (B * x) .* c ≈ cB * x
    @test (C * x) .* c ≈ cC * x

    @test_throws DimensionMismatch ones(10) * A
    @test_throws DimensionMismatch ones(10) * cA
    @test_throws DimensionMismatch ones(10) * B
    @test_throws DimensionMismatch ones(10) * C
end

# These tests are broken due to the implementation 2.2*LD creating a DerivativeOperator
# rather than an Array
#=
@testset "Linear combinations of operators" begin
    N = 10
    Random.seed!(0); LA = DiffEqArrayOperator(rand(N,N+2))
    LD = CenteredDifference(2,2,1.0,N)
    L = 1.1*LA - 2.2*LD + 3.3*Eye(N,N+2)
    # Builds convert(L) the brute-force way
    fullL = zeros(N,N+2)
    v = zeros(N+2)
    for i = 1:N+2
        v[i] = 1.0
        fullL[:,i] = L*v
        v[i] = 0.0
    end
    @test_broken convert(AbstractMatrix,L) ≈ fullL
    for p in [1,2,Inf]
        @test_broken opnorm(L,p) ≈ opnorm(fullL,p)
    end
end
=#


abstract type AbstractStochasticProcess end
abstract type AbstractMarkovChain <: AbstractStochasticProcess end

"""
    MarkovChain

A discrete-time, discrete-space, finite-space Markov chain.

##### Fields
- `transitionmatrix::Matrix{Float64}` : stochastic matrix in which element `[i,j]` denotes
the probability that the Markov chain transitions from state `i` to state `j`.
- `initialdist::Vector{Float64}` : stochastic vector in which element `i` denotes the
probability that the Markov chain's initial state is `i`.

##### Complexity (for the inner constructor)
- Time complexity: `O(n^2)`, where `n` denotes the number of rows (or columns) in a square
matrix.
"""
struct MarkovChain <: AbstractMarkovChain
    transitionmatrix::Matrix{Float64}
    initialdist::Vector{Float64}
    # ensure that the transition matrix and the initial distribution are stochastic
    function MarkovChain(transitionmatrix::AbstractMatrix{<:Real},
                         initialdist::AbstractVector{<:Real})
        size(transitionmatrix)[1] != length(initialdist) &&
            ArgumentError("Uh-oh! The transition matrix must have the same number of rows as
                          the initial distribution has elements.") |> throw
        !isstochastic(initialdist) || !isstochastic(transitionmatrix) &&
            ArgumentError("Oops! The transition matrix and the initial distribution must be
                          stochastic.") |> throw
        return new(transitionmatrix, initialdist)
    end
end

"""
    Base.show(io::IO, markovchain::MarkovChain) -> nothing

Show a `MarkovChain` in a human-friendly manner.
"""
function Base.show(io::IO, markovchain::MarkovChain)
    println(io, typeof(markovchain))
    println(io, " Initial distribution:")
    println(io, "  ", markovchain.initialdist)
    println(io, " Transition matrix:")
    println(io, "  ", markovchain.transitionmatrix)
    return nothing
end

"""
    MarkovChain(transitionmatrix::AbstractMatrix{<:Real},
                initialdist::AbstractMatrix{<:Real}) -> MarkovChain

Construct a `MarkovChain` when the initial distribution `initialdist` is specified as either
a one-row or a one-column matrix.

##### Complexity
- Time complexity: `O(n^2)`, where `n` denotes the number of rows (or columns) in a square
matrix.
"""
function MarkovChain(transitionmatrix::AbstractMatrix{<:Real},
                     initialdist::AbstractMatrix{<:Real})
    nrows, ncols = size(initialdist)[1], size(initialdist)[2]
    if nrows == 1
        # make initialdist a vector
        initialdist = initialdist[1,:]
    elseif ncols == 1
        # make initialdist a vector
        initialdist = initialdist[:,1]
    else
        throw(ArgumentError("The initial distribution must be a vector."))
    end
    # multiple dispatch
    return MarkovChain(transitionmatrix, initialdist)
end

"""
    isstochastic(matrix::AbstractMatrix, tol::Float64 = 1e-5) -> Bool

Decide if `matrix` is stochastic.

`tol` is the tolerance associated with the row sums. We consider the row sums to be valid if
and only if they are in the interval `[1-tol, 1+tol]`. This is important when the element
type is `Float64`, for example.

##### Complexity
- Time complexity: `O(n^2)`, where `n` denotes the number of rows (or columns) in a square
matrix.
"""
function isstochastic(matrix::AbstractMatrix, tol::Float64 = 1e-5)
    nrows, ncols = size(matrix)[1], size(matrix)[2]
    # if matrix has only one row
    if nrows == 1
        vec = matrix[1,:]::Vector
        # multiple dispatch
        return isstochastic(vec, tol)
    # if matrix has only one column
    elseif ncols == 1
        vec = matrix[:,1]::Vector
        # multiple dispatch
        return isstochastic(vec, tol)
    # if matrix is non-square
    elseif nrows != ncols
        return false
    end
    # check if elements are in [0,1]
    for element in matrix
        if element < 0 || element > 1
            return false
        end
    end
    # check row sums
    for row in eachrow(matrix)
        rowsum = sum(row)
        if rowsum < 1 - tol || rowsum > 1 + tol
            return false
        end
    end
    return true  # this line is reached only if all above checks pass
end

"""
    isstochastic(vector::AbstractVector, tol::Float64 = 1e-5) -> Bool

Decide if `vector` is stochastic.

`tol` is the tolerance associated with the sum over `vector`'s elements. We consider the sum
to be valid if and only if it is in the interval `[1-tol, 1+tol]`. This is important when
the element type is `Float64`, for example.

##### Complexity
- Time complexity: `O(n)`.
"""
function isstochastic(vector::AbstractVector, tol::Float64 = 1e-5)
    # check if elements are in [0,1]
    for element in vector
        if element < 0 || element > 1
            return false
        end
    end
    vecsum = sum(vector)
    if vecsum < 1 - tol || vecsum > 1 + tol
        return false
    end
    return true  # this line is reached only if all above checks pass
end


using Oceananigans.Grids

using PyPlot

include("ConvergenceTests/ConvergenceTests.jl")

using .ConvergenceTests
using .ConvergenceTests.TwoDimensionalDiffusion: run_and_analyze

defaultcolors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
removespine(side) = gca().spines[side].set_visible(false)
removespines(sides...) = [removespine(side) for side in sides]

function convergence_test(Nx, Δt, stop_iteration, topo)
    results = [run_and_analyze(Nx=N, Δt=Δt, stop_iteration=stop_iteration, topo=topo,
                               output=false) for N in Nx]

    L₁ = map(r -> r.L₁, results)
    L∞ = map(r -> r.L∞, results)

    return (L₁=L₁, L∞=L∞)
end

# Setup and run simulations
Nx = [8, 16, 32, 64, 128, 256]
stop_time = 1e-4

# Calculate time step based on diffusive time-step constraint for finest mesh
        min_Δx = 2π / maximum(Nx)
   proposal_Δt = 1e-3 * min_Δx^2 # proposal time-step
stop_iteration = round(Int, stop_time / proposal_Δt)
            Δt = stop_time / stop_iteration # ensure time-stepping to exact finish time.

topologies = (
              (Periodic, Periodic, Bounded),
              (Periodic, Bounded, Bounded),
              (Bounded, Bounded, Bounded)
             )

errors = [convergence_test(Nx, Δt, stop_iteration, topo) for topo in topologies]

fig, axs = subplots()

for (itopo, topo) in enumerate(topologies)
    L₁ = errors[itopo].L₁
    L∞ = errors[itopo].L∞
    name ="$(topo[1]), $(topo[2])"

    loglog(Nx, L₁, "o", basex=2, color=defaultcolors[itopo], mfc="None", label="\$L_1\$-norm, $name")
    loglog(Nx, L∞, "^", basex=2, color=defaultcolors[itopo], mfc="None", label="\$L_\\infty\$-norm, $name")
end

L₁ = errors[1].L₁
L∞ = errors[1].L∞
loglog(Nx, L₁[1] * (Nx[1]./Nx).^2, "k-", basex=2, alpha=0.6, linewidth=1, label=L"\sim N_x^{-2}")

xlabel(L"N_x")
ylabel("\$L\$-norms of \$ | c_\\mathrm{sim} - c_\\mathrm{analytical} |\$")
title("Two dimensional diffusion convergence test")
removespines("top", "right")
legend(loc="upper right")

filepath = joinpath(@__DIR__, "figs", "two_dimensional_diffusion_convergence.png")
savefig(filepath, dpi=480)

for (itopo, topo) in enumerate(topologies)
    L₁ = errors[itopo].L₁
    L∞ = errors[itopo].L∞
    test_rate_of_convergence(L₁, Nx, expected=-2.0, atol=0.01, name="2D diffusion $topo L₁")
    test_rate_of_convergence(L∞, Nx, expected=-2.0, atol=0.06, name="2D diffusion $topo L∞")
end


"""
    v = normalized_legendre(x,n)

Calculates a single value of the `n`-th order Legendre polynomial of `x`.
It calls `Jacobi.legendre`.
"""
function normalized_legendre(x,n)
   a = sqrt((2*n+1)/2)
   return a*legendre(x,n)
end

"""
    M = normalized_legendre_matrix(t,n)

Calculates a matrix with the `n`-th order Legendre coefficients on time points `t`, a vector or range.

```juliadoctests
# up to 3rd order coefficients for 5 to 12
julia> M = normalized_legendre_matrix(5:12,3)

# same expression
julia> M = normalized_legendre_matrix([5,6,7,8,9,10,11,12],3)
```
"""
function normalized_legendre_matrix(t::Union{UnitRange,Vector},n)
   if !issorted(t)
      throw(ArgumentError("time points not sorted"))
   end
   tmin = minimum(t)
   tmax = maximum(t)
   x = -1 .+ 2*( (t .- tmin)/(tmax-tmin) )
   M = [normalized_legendre(y,k) for y=x,k=0:n]
   return M
end

"""
    vare = hrv_class(ve,rangeset)

Generate a vector containing residual variances from a vector of heterogeneous residual variances `ve` and a set of ranges `rangeset`.

```juliadoctests
julia> vare = hrv_class([10.0,20.0],[3:5,6:9])
7-element Vector{Float64}:
10.0
10.0
10.0
20.0
20.0
20.0
20.0
```
"""
function hrv_class(ve,rangeset)
   if length(ve)!=length(rangeset)
      throw(DimensionMismatch("ve and rangeset"))
   end
   tmin = minimum(rangeset[1])
   tmax = maximum(rangeset[end])
   vare = zeros(tmax-tmin+1)
   for i in 1:length(rangeset)
      for t in rangeset[i]
         vare[t-tmin+1] = ve[i]
      end
   end
   return vare
end


# ## クロスバリデーション
# まず線形回帰の場合のクロスバリデーションの関数を定義します。
# ```julia
# function cv_linear(X::AbstractArray,y::Vector,K::Int)
#     n = length(y); m = round(Int,n/K)
#     S = 0
#     for j in 1:K
#         test = j*m-m+1:j*m
#         train = setdiff(1:n,test)
#         β̂ = multiple_regression(X[train,:],y[train])
#         ŷ = insert_ones(X[test,:])*β̂
#         S += dot(y[test]-ŷ,y[test]-ŷ)
#     end
#     return S/n
# end
# ```
# ### 例39
using Joe, Random, Plots
using Plots.PlotMeasures # hide
Plots.reset_defaults() # hide
default( # hide
    titlefont  = font("JuliaMono", default(:titlefontsize),  ), # hide
    guidefont  = font("JuliaMono",  default(:guidefontsize),  ), # hide
    tickfont   = font("JuliaMono", default(:tickfontsize),   ), # hide
    legendfont = font("JuliaMono",  default(:legendfontsize), ), # hide
    left_margin = 30px, # hide
    bottom_margin = 30px # hide
) # hide
using Joe: cv_linear
n = 100; p = 5
Random.seed!(1)
X_39 = randn(n, p)
β = randn(p+1); β[2:3] .= 0
y_39 = insert_ones(X_39)*β + randn(n)
@show cv_linear(X_39[:,[3,4,5]],y_39,10);
@show cv_linear(X_39,y_39,10);
# あれ？変数選択した方が誤差でかいぞ？試行回数を増やして全変数を用いた場合のCV値と
# 変数選択を行った場合のCV値を可視化します。
using Joe # hide
U=Float64[]; V = Float64[]
for _ in 1:100
    global U, V, X_39, β, n # hide
    local y # hide
    y = insert_ones(X_39)*β + randn(n)
    push!(U,cv_linear(X_39[:,[3,4,5]],y,10))
    push!(V,cv_linear(X_39,y,10))
end

p39 = scatter(U,V,xlabel="変数4,5,6を選んだ時の二乗誤差",
                ylabel="全変数を選んだ時の二乗誤差",
                title="変数を多く選びすぎて過学習")
plot!(p39, x->x,xlims=(0.7,1.5),legend=false)
savefig(p39,joinpath(@OUTPUT,"fig3-1.svg")) # hide
# \fig{fig3-1}
# \lineskip
# $y=x$ の直線の上側に大多数の点がプロットされているので、
# 変数選択をした方が誤差が少ないことが分かりました。
# ただ、あまり二乗誤差の違いは大きくないです。
# ### 例40
# kの値を変えたときのCVの予測誤差の変化
n = 100; p = 5;
p40 = plot(ylims=(0.3,1.5),xlabel="k",ylabel="CVの値",
            title="k-foldのkとCVの値の関係",legend=false)
Random.seed!(1)
for _ in 1:10
    global p40 # hide
    local X, β, y, U, V # hide
    X = randn(n,p)
    β = randn(p+1)
    y = insert_ones(X) * β + randn(n)
    U = Int[]
    V = Float64[]
    for k in 2:n
        if n%k==0
            push!(U,k)
            push!(V, cv_linear(X,y,k))
        end
    end
    plot!(p40,U,V)
end
savefig(p40,joinpath(@OUTPUT,"fig3-2.svg")) # hide
# \fig{fig3-2}
# ### 例41 (FInsherのあやめ)
# あやめのデータセットでK近傍法のｋ毎に誤り率を評価します。

using ScikitLearn, StatsBase, Random, Plots
using Joe:knn
@sk_import datasets: load_iris
iris = load_iris()
X_41 = iris["data"]
y_41 = iris["target"]
n = length(y_41);
# データセットをランダムに並べ替えておきます。
Random.seed!(1)
order = StatsBase.sample(1:n,n,replace=false);
X_41 = X_41[order,:];
y_41 = y_41[order];
# 10-foldクロスバリデーションなので、テストデータを15毎に交代させることになります。
errorRate = Vector{Float64}(undef,10)
for k in 1:10
    global errorRate, X_41, y_41 # hide
    S = 0
    for top in 1:15:150
        test = top:top+14
        train = setdiff(1:150,test)
        knn_ans = knn(X_41[train,:],y_41[train],X_41[test,:],k)
        S += sum(y_41[test] .!= knn_ans)
    end
    S /= n
    errorRate[k] = S
end
p41 = plot(errorRate,xlabel="K",ylabel="誤り率",
            legend=false, title="CVによる誤り率の評価")
savefig(p41,joinpath(@OUTPUT,"fig3-3.svg")) # hide
# \fig{fig3-3}
# \lineskip
# $K \ge 3$で大きく誤り率が低下することが分かりました。
# pythonでは実行に10分程度かかったとありますが、いくらなんでも遅すぎでは？(juliaは瞬殺)
# ## 線形回帰の場合の公式(CVの公式)
# クロスバリデーションでの誤差の分析は、訓練データ毎に重回帰分析を行う必要があるので、
# 処理に時間がかかります。$\hat{\beta}_{-S}$をあらわに求めることなく、CVを
# 楽して評価できるのがこの公式です。($S$はテストデータの集合、$-S$は訓練データの集合)
# \lineskip
# $$\sum_{S} || (I - H_S)^{-1}e_S||^2$$
# ただし、$H_S = X_S(X^\mathsf{T}X)^{-1}X_S^\mathsf{T}$、
# $e_S = y_S - X_S\hat{\beta}$ です。プログラムを実装して、ベンチマークをしてみます。
# cv_fast関数は次のように実装しました。
# ```julia
# function cv_fast(X::AbstractArray,y::Vector,K::Int)
#     n = length(y); m = round(Int,n/K)
#     X = insert_ones(X)
#     H = X*((X'X)\X')
#     e = (I(n) - H)*y
#     S = 0
#     @views for j in 1:K
#         test = j*m-m+1:j*m;
#         err = (I(m)-H[test,test]) \ e[test]
#         S += dot(err,err)
#     end
#     return S/n
# end
# ```
using Joe, Random, Plots
using Joe:cv_linear,cv_fast

n = 1000; p = 5
Random.seed!(1)
X = randn(n, p)
β = randn(p+1); β[2:3] .= 0
y = insert_ones(X)*β + randn(n);
# ループのネストを浅くしたいので、K-foldのkのリストを作って、ループすることにします。
K = [i for i in 2:1000 if 1000%i ==0]
t_linear = Vector{Float64}(undef,length(K));
t_fast = similar(t_linear);
# まずは関数の慣らし運転を
@show cv_linear(X,y,10)
@show cv_fast(X,y,10)
# 値が一致することは確認できました。それでは、いざ本番!
for i in eachindex(K)
    global X, y, t_fast, t_linear,K # hide
    start = time_ns();cv_fast(X,y,K[i]);stop = time_ns()
    t_fast[i] = (stop-start)/1e9
    start = time_ns();cv_linear(X,y,K[i]);stop = time_ns()
    t_linear[i] = (stop-start)/1e9
end
p42 = plot(ylims=(0.0,0.5), xlabel = "k" ,ylabel="実行時間",
           label="cv_linear",title="cv_fastとcv_linearの比較",
           legend=:topleft)
plot!(p42,K, t_linear, label="cv_linear")
plot!(p42,K, t_fast, label="cv_fast")
savefig(p42,joinpath(@OUTPUT,"fig3-4.svg")) # hide
# \fig{fig3-4}
# \lineskip
# $n\ge100$ではcv_fastの方が実行時間が短いことが分かります。
# 原著と比べると、マシンの違いがあるとはいえ、juliaは速いということもよくわかります。
# ## ブートストラップ
# 以下のような複合型とbootstrap関数を定義しました。
# 原著のコードよりは可読性が向上したかなぁ。
# ```julia
# using Parameters
# @with_kw mutable struct bt{T<:Number}
#     original::T = 0.0
#     bias::T = 0.0
#     stderr::T = 0.0
# end
# function bootstrap(df::AbstractArray{T}, F::Function, r::Int) where T <:Number
#     m,_ = size(df)
#     original = F(df,1:m)
#     u = Vector{T}(undef, r)
#     for i in 1:r
#         index = sample(1:m,m,replace=true)
#         u[i] = F(df,index)
#     end
#     bias = mean(u) - original; stderr = std(u)
#     result = bt()
#     @pack! result = original, bias, stderr
# end
# ```
# ### 例43
# データは[ここ](https://bitbucket.org/prof-joe/statistical_learning_with_python/src/master/)からダウンロードできます。
using DelimitedFiles, Statistics
using Joe:bootstrap
function func_1(data,index)
    X=data[index,1]; Y= data[index,2]
    (var(Y) - var(X))/(var(X) + var(Y) -2cov(X,Y))
end
Portfolio = readdlm(joinpath("_assets","data","Portfolio.csv"),',', skipstart=1)
bootstrap(Portfolio,func_1,1000)

# ### 例44
# データの読み込み
using Joe,DelimitedFiles
df = readdlm(joinpath("_assets","data","crime.txt"))
X = df[:,3:4]; y =df[:,1]
β = multiple_regression(X,y)

# 回帰係数のそれぞれについて、ブートストラップ法を適用します。ここでは標準偏差だけを出力します。
std_bt = ones(3)
for j in 1:3
    function func_2(data,index)
        X = data[index,3:4];y = data[index,1]
        β = multiple_regression(X,y)
        return β[j]
    end
    std_bt[j] = bootstrap(df,func_2,1000).stderr
end
std_bt
# GLM.jlで回帰係数を推定したときの標準偏差と比較してみます。
using GLM, DataFrames
data = DataFrame(df,:auto)
ols = lm(@formula(x1 ~ x3 + x4), data)
# 標準偏差がほぼ一致することが分かりました。


### A Pluto.jl notebook ###
# v0.12.7

using Markdown
using InteractiveUtils

# ╔═╡ acd7da78-1d69-11eb-2517-e1bafdc14f4a
using Pkg, DrWatson

# ╔═╡ acd82a8c-1d69-11eb-10c9-0902b833d79e
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Plots, LaTeXStrings, Measures
end

# ╔═╡ 2e6c02cc-1d69-11eb-283a-1bcf9e0044bf
md"## Listing 1.10"

# ╔═╡ acd8c1cc-1d69-11eb-3164-5b9e66fc3664
begin
	f(x,y) = x^2 + y^2
	f0(x) = f(x,0)
	f2(x) = f(x,2)
end

# ╔═╡ ace3c478-1d69-11eb-0a4d-f191cf032e78
begin
	x, y = -5:0.1:5 , -5:0.1:5
	plot(x, [f0.(x), f2.(x)],
		c=[:blue :red], xlims=(-5, 5), legend=:top,
		ylims=(-5, 25), ylabel=L"f(x,⋅)", label=["f(x, 0)" "f(x, 2)"])
	p1 = annotate!(-1, -0.2, text("Minimum = (0,0)", :left, :top, 10))
end

# ╔═╡ ace462fa-1d69-11eb-2b31-c1c9efe2376f
begin
	z = [ f(x,y) for y in y, x in x ]
	p2 = surface(x, y, z, c=cgrad([:blue, :red]), legend=:none, ylabel="y", zlabel="f(x,y)")
end

# ╔═╡ acf332f0-1d69-11eb-09df-43d188f0b58d
begin
	M = z[1:10,1:10]
	p3 = heatmap(M, c=cgrad([:blue, :red]), yflip=true, ylabel="y",
		xticks=([1:10;], x), yticks=([1:10;], y))
end

# ╔═╡ ad012874-1d69-11eb-24a0-e936c7c266d5
md"## End of listing 1.10"

# ╔═╡ Cell order:
# ╟─2e6c02cc-1d69-11eb-283a-1bcf9e0044bf
# ╠═acd7da78-1d69-11eb-2517-e1bafdc14f4a
# ╠═acd82a8c-1d69-11eb-10c9-0902b833d79e
# ╠═acd8c1cc-1d69-11eb-3164-5b9e66fc3664
# ╠═ace3c478-1d69-11eb-0a4d-f191cf032e78
# ╠═ace462fa-1d69-11eb-2b31-c1c9efe2376f
# ╠═acf332f0-1d69-11eb-09df-43d188f0b58d
# ╟─ad012874-1d69-11eb-24a0-e936c7c266d5


# ---
# title: 60. Permutation Sequence
# id: problem60
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Math, Backtracking
# link: <https://leetcode.com/problems/permutation-sequence/description/>
# hidden: true
# ---
#
# The set `[1, 2, 3, ..., n]` contains a total of `n!` unique permutations.
#
# By listing and labeling all of the permutations in order, we get the following
# sequence for `n = 3`:
#
#   1. `"123"`
#   2. `"132"`
#   3. `"213"`
#   4. `"231"`
#   5. `"312"`
#   6. `"321"`
#
# Given `n` and `k`, return the `kth` permutation sequence.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 3, k = 3
#     Output: "213"
#
#
# **Example 2:**
#
#
#
#     Input: n = 4, k = 9
#     Output: "2314"
#
#
# **Example 3:**
#
#
#
#     Input: n = 3, k = 1
#     Output: "123"
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 9`
#   * `1 <= k <= n!`
#
#
## @lc code=start
using LeetCode

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


# # Nonlocal diffusion on $\mathbb{S}^2$
# This example calculates the spectrum of the nonlocal diffusion operator:
# ```math
# \mathcal{L}_\delta u = \int_{\mathbb{S}^2} \rho_\delta(|\mathbf{x}-\mathbf{y}|)\left[u(\mathbf{x}) - u(\mathbf{y})\right] \,\mathrm{d}\Omega(\mathbf{y}),
# ```
# defined in Eq. (2) of
#
# R. M. Slevinsky, H. Montanelli, and Q. Du, [A spectral method for nonlocal diffusion operators on the sphere](https://doi.org/10.1016/j.jcp.2018.06.024), *J. Comp. Phys.*, **372**:893--911, 2018.
#
# In the above, $0<\delta<2$, $-1<\alpha<1$, and the kernel:
# ```math
# \rho_\delta(|\mathbf{x}-\mathbf{y}|) = \frac{4(1+\alpha)}{\pi \delta^{2+2\alpha}} \frac{\chi_{[0,\delta]}(|\mathbf{x}-\mathbf{y}|)}{|\mathbf{x}-\mathbf{y}|^{2-2\alpha}},
# ```
# where $\chi_I(\cdot)$ is the indicator function on the set $I$.
#
# This nonlocal operator is diagonalized by spherical harmonics:
# ```math
# \mathcal{L}_\delta Y_\ell^m(\mathbf{x}) = \lambda_\ell(\alpha, \delta) Y_\ell^m(\mathbf{x}),
# ```
# and its eigenfunctions are given by the generalized Funk--Hecke formula:
# ```math
# \lambda_\ell(\alpha, \delta) = \frac{(1+\alpha) 2^{2+\alpha}}{\delta^{2+2\alpha}}\int_{1-\delta^2/2}^1 \left[P_\ell(t)-1\right] (1-t)^{\alpha-1} \,\mathrm{d} t.
# ```
# In the paper, the authors use Clenshaw--Curtis quadrature and asymptotic evaluation of Legendre polynomials to achieve $\mathcal{O}(n^2\log n)$ complexity for the evaluation of the first $n$ eigenvalues. With a change of basis, this complexity can be reduced to $\mathcal{O}(n\log n)$.
#
# First, we represent:
# ```math
# P_n(t) - 1 = \sum_{j=0}^{n-1} \left[P_{j+1}(t) - P_j(t)\right] = -\sum_{j=0}^{n-1} (1-t) P_j^{(1,0)}(t).
# ```
# Then, we represent $P_j^{(1,0)}(t)$ with Jacobi polynomials $P_i^{(\alpha,0)}(t)$ and we integrate using [DLMF 18.9.16](https://dlmf.nist.gov/18.9.16):
# ```math
# \int_x^1 P_i^{(\alpha,0)}(t)(1-t)^\alpha\,\mathrm{d}t = \left\{ \begin{array}{cc} \frac{(1-x)^{\alpha+1}}{\alpha+1} & \mathrm{for~}i=0,\\ \frac{1}{2i}(1-x)^{\alpha+1}(1+x)P_{i-1}^{(\alpha+1,1)}(x), & \mathrm{for~}i>0.\end{array}\right.
# ```
# The code below implements this algorithm, making use of the Jacobi--Jacobi transform `plan_jac2jac`.
# For numerical stability, the conversion from Jacobi polynomials $P_j^{(1,0)}(t)$ to $P_i^{(\alpha,0)}(t)$ is divided into conversion from $P_j^{(1,0)}(t)$ to $P_k^{(0,0)}(t)$, before conversion from $P_k^{(0,0)}(t)$ to $P_i^{(\alpha,0)}(t)$.

using FastTransforms, LinearAlgebra

function oprec!(n::Integer, v::AbstractVector, alpha::Real, delta2::Real)
    if n > 0
        v[1] = 1
    end
    if n > 1
        v[2] = (4*alpha+8-(alpha+4)*delta2)/4
    end
    for i = 1:n-2
        v[i+2] = (((2*i+alpha+2)*(2*i+alpha+4)+alpha*(alpha+2))/(2*(i+1)*(2*i+alpha+2))*(2*i+alpha+3)/(i+alpha+3) - delta2/4*(2*i+alpha+3)/(i+1)*(2*i+alpha+4)/(i+alpha+3))*v[i+1] - (i+alpha+1)/(i+alpha+3)*(2*i+alpha+4)/(2*i+alpha+2)*v[i]
    end
    return v
end

function evaluate_lambda(n::Integer, alpha::T, delta::T) where T
    delta2 = delta*delta
    scl = (1+alpha)*(2-delta2/2)

    lambda = Vector{T}(undef, n)

    if n > 0
        lambda[1] = 0
    end
    if n > 1
        lambda[2] = -2
    end

    oprec!(n-2, view(lambda, 3:n), alpha, delta2)

    for i = 2:n-1
        lambda[i+1] *= -scl/(i-1)
    end

    p = plan_jac2jac(T, n-1, zero(T), zero(T), alpha, zero(T))

    lambda[2:end] .= p'lambda[2:end]

    for i = 2:n-1
        lambda[i+1] = ((2i-1)*lambda[i+1] + (i-1)*lambda[i])/i
    end

    for i = 2:n-1
        lambda[i+1] += lambda[i]
    end

    return lambda
end

# The spectrum in `Float64`:
lambda = evaluate_lambda(10, -0.5, 1.0)

# The spectrum in `BigFloat`:
lambdabf = evaluate_lambda(10, parse(BigFloat, "-0.5"), parse(BigFloat, "1.0"))

# The $\infty$-norm relative error:
norm(lambda-lambdabf, Inf)/norm(lambda, Inf)


#=
  Regression fallacy.

  From Gelman et.al "Regression and other stories", page 88f
  """
  Regression to the mean can be confusing and it has led people to mistakenly attribute causality.
  To see how this can happen, we move from heights of parents and children to the mathematically
  equivalent scenario of students who take two tests.
  [...]
  Rather than using real data, we have simulated exam scores using the following simple process
  representing signal and noise:6
  1. Each student is assumed to have a true ability drawn from a distribution with
     mean 50 and standard deviation 10.
  2. Each student’s score on the midterm exam is the sum of two components: the student’s true ability,
     and a random component with mean 0 and standard deviation 10, reflecting that performance on
     any given test will be unpredictable: a midterm exam is far from a perfect measuring instrument
  3. Likewise, each student’s score on the final exam is his or her true ability, plus another, independent,
     random component.
  """

  1) We simulate with rand and Distributions.jl and then fit with lm(),
     were we see that it's a "regression to the mean", i.e. slope < 1.

    score_final ~ 1 + score_midterm

    Coefficients:
    ───────────────────────────────────────────────────────────────────────────
                    Coef.  Std. Error      t  Pr(>|t|)  Lower 95%  Upper 95%
    ───────────────────────────────────────────────────────────────────────────
    (Intercept)    22.1621     1.44193    15.37    <1e-47  19.3326    24.9917
    score_midterm   0.565276   0.0278395  20.30    <1e-76   0.510646   0.619907
    ───────────────────────────────────────────────────────────────────────────

  2) Then we try to the recover the parameters with Turing.
      mu_ability = 50
      sigma_ability = 10

      mu_score = 0
      sigma_score = 10

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

        mu_ability   49.7729    0.3167     0.0100    0.0090   948.7999    0.9997      129.3877
     sigma_ability    9.7822    0.2086     0.0066    0.0063   670.0839    1.0013       91.3792
          mu_score    0.1490    0.1154     0.0036    0.0039   625.4673    0.9999       85.2949
       sigma_score   10.1247    0.1610     0.0051    0.0045   683.0471    0.9991       93.1470

    Quantiles
        parameters      2.5%     25.0%     50.0%     75.0%     97.5%
            Symbol   Float64   Float64   Float64   Float64   Float64

        mu_ability   49.1503   49.5634   49.7748   49.9773   50.3799
     sigma_ability    9.3879    9.6409    9.7798    9.9150   10.2067
          mu_score    0.0040    0.0605    0.1272    0.2114    0.4748
       sigma_score    9.8106   10.0148   10.1218   10.2331   10.4427

=#
using Turing, StatsPlots, DataFrames, Distributions, Random
using GLM
include("jl_utils.jl")

#
# It's much faster to use random and Distributions to generate
# the data (instead of using Turing for this).
#
# It's almost the same model, except for the ~'s.
#
function regression_fallacy_sim(n=1000)

    mu_ability = 50
    sigma_ability = 10

    mu_score = 0
    sigma_score = 10

    ability = rand(Normal(mu_ability,sigma_ability), n)
    score_midterm = Vector{Real}(undef, n)
    score_final = Vector{Real}(undef, n)
    # for i in 1:n
    #     score_midterm[i] = rand(Normal(ability[i]+mu_score,sigma_score))
    #     score_final[i] = rand(Normal(ability[i]+mu_score,sigma_score))
    # end
    # Neater:
    @. score_midterm = rand(Normal(ability+mu_score,sigma_score))
    @. score_final = rand(Normal(ability+mu_score,sigma_score))

    return ability, score_midterm, score_final
end

n = 1000

ability, score_midterm, score_final = regression_fallacy_sim(n)

#
# Fit the data
#

df = DataFrame(ability=ability,score_midterm=score_midterm, score_final=score_final)
fit = lm(@formula(score_final ~ score_midterm), df)
display(fit)
println()

#=
pred = predict(fit)
# display(plot(xlabel="midterm", ylabel="final", legend=:bottomright))
display(plot(df.score_midterm, df.score_final, label="data", seriestype=:scatter,ms=1,xlabel="midterm", ylabel="final", legend=:bottomright))
display(plot!(score_midterm, pred, label="model", linewidth=3))
=#

#=
# What are the mean for those students that got
# final < / > midterm scores
mean_midterm = mean(score_midterm)
mean_final = mean(score_final)

ix_gt = score_final .> score_midterm
ix_lt = score_final .< score_midterm
@show mean(score_final[ix_gt])
@show mean(score_final[ix_lt])
@show mean(score_final[ix_gt] .- score_midterm[ix_gt])
@show mean(score_final[ix_lt] .- score_midterm[ix_lt])
=#

#
# Now we try to recover the values.
#
# Note: Generating the data using this Turing model takes too long time.
# Hence the dedicated simulation model above.
#
@model function regression_fallacy(ability, score_midterm, score_final)
    n = length(ability)
    # Use uninformed priors so we are not accused of cheating. :-)
    mu_ability ~ Uniform(1,1000)
    sigma_ability ~ Uniform(1,50)

    mu_score ~ Uniform(0,100)
    sigma_score ~ Uniform(1,50)

    ability .~ Normal(mu_ability,sigma_ability)
    @. score_midterm ~ Normal(ability+mu_score,sigma_score)
    @. score_final   ~ Normal(ability+mu_score,sigma_score)
end

model = regression_fallacy(ability,score_midterm,score_final)

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

# chns = sample(model, NUTS(1000,0.65), 1_000)
chns = sample(model, NUTS(), 1_000)
# chns = sample(model, HMC(0.1,5), 1_000)

display(chns)


"""
    SqExponentialARD(lnℒ::AbstractVector)
    <: RadialKernel{WeightedSqEuclidean}

The squared exponential kernel (aka the Gaussian kernel) with
vector valued length scale
``\\exp(2 \\ln\\mathscr{L}_i) = \\mathscr{L}_i^2 > 0``.

Finding an optimal ``\\mathscr{L}`` is refered to as automatic
relevance detetermination.

For the matrix ``\\Sigma_{ii} = \\mathscr{L}_i^{-2}``

```math
k_{\\mathscr{L}}(x, y) = k_{\\mathscr{L}}(\\lVert x - y\\rVert_\\Sigma^2) =
\\exp\\Big\\lbrace -\\frac{1}{2}\\lVert x - y\\rVert_\\Sigma^2\\Big\\rbrace
```

"""
mutable struct SqExponentialARD{
    F<:AbstractFloat,
    V<:AbstractVector{F},
} <: RadialKernel{WeightedSqEuclidean}
    dist::WeightedSqEuclidean
    lnℒ::V
    diagΣ::V
end

function SqExponentialARD(lnℒ::AbstractVector)
    diagΣ = exp.(-2 .* lnℒ)
    SqExponentialARD(WeightedSqEuclidean(diagΣ), lnℒ, diagΣ)
end

@inline (k::SqExponentialARD)(τ::AbstractFloat) = exp(-τ / 2)

GaussianARD = SqExponentialARD

numparams(k::SqExponentialARD) = (length(k.lnℒ),)
paramtypes(k::SqExponentialARD{F}) where F = (F,)
params(k::SqExponentialARD) = (lnℒ = k.lnℒ,)
function setparams!(k::SqExponentialARD{F}, lnℒ::AbstractVector{F}) where F
    diagΣ = exp.(-2 .* lnℒ)
    copyto!(k.lnℒ, lnℒ)
    copyto!(k.diagΣ, diagΣ)
    k.dist = WeightedSqEuclidean(diagΣ)
end


using TuringModels

Turing.setadbackend(:reversediff)

# Turing model
@model m8_3(y) = begin
    α ~ Normal(1, 10)
    σ ~ truncated(Cauchy(0, 1), 0, Inf)

    y .~ Normal(α, σ)
end

y = [-1,1]

# Sample

chns = sample(m8_3(y), NUTS(0.95), 1000);

# Results rethinking

m83rethinking = "
      mean   sd  5.5% 94.5% n_eff Rhat
alpha 0.09 1.63 -2.13  2.39   959    1
sigma 2.04 2.05  0.68  4.83  1090    1
";

# Describe the posterior samples

describe(chns)

# End of `08/m8.3t.jl`


###############################################
#FEniCS tutorial demo program: Deflection of a membrane.
#  -Laplace(w) = p  in the unit circle
#            w = 0  on the boundary
#The load p is a Gaussian function centered at (0, 0.6).
###############################################

module ft02

using FenicsPy

# Create mesh and define function space
Ω = Circle(Point(0, 0), 1)
mesh = generate_mesh(Ω, 64)
V = FunctionSpace(mesh, "P", 2)

# Define boundary condition
w_D = Constant(0)

bc = DirichletBC(V, w_D, "on_boundary")

# Define load
β = 8
R0 = 0.6
p = Expression("4*exp(-pow(beta, 2)*(pow(x[0], 2) + pow(x[1] - R0, 2)))",
               degree=1, beta=β, R0=R0)

# Define variational problem
w = TrialFunction(V)
v = TestFunction(V)
a = dot(grad(w), grad(v))*dx
L = p*v*dx

# Compute solution
w = FeFunction(V)
solve(a == L, w, bc)

p = interpolate(p, V)

# Save solution to file in VTK format
vtkfile_w = File("poisson_membrane/deflection.pvd")
vtkfile_w << w
vtkfile_p = File("poisson_membrane/load.pvd")
vtkfile_p << p

# Plot solution
#plot(w, title="Deflection")
#plot(p, title="Load")

# Curve plot along x = 0 comparing p and w
#import numpy as np
#import matplotlib.pyplot as plt
#tol = 0.001  # avoid hitting points outside the domain
#y = np.linspace(-1 + tol, 1 - tol, 101)
#points = [(0, y_) for y_ in y]  # 2D points
#w_line = np.array([w(point) for point in points])
#p_line = np.array([p(point) for point in points])
#plt.plot(y, 50*w_line, 'k', linewidth=2)  # magnify w
#plt.plot(y, p_line, 'b--', linewidth=2)
#plt.grid(True)
#plt.xlabel('$y$')
#plt.legend(['Deflection ($\\times 50$)', 'Load'], loc='upper left')
#plt.savefig('poisson_membrane/curves.pdf')
#plt.savefig('poisson_membrane/curves.png')

end # module ft02


using FFTW

function FourierCoeffsOdd(f,ta,tb,M)
  N = 2*M+1
  h = (tb-ta)/N
  tt = range(ta,step=h,length=N)
  ff = f.(tt)
  hf = 1/N*rfft(ff)
  return (tt,ff,hf)
end

function FourierEvalOdd(ta,tb,hf,N)
    M = length(hf)-1
    if (N<2M+1)
        error("$M=length(hf)-1 should be less than (N+1)/2=$div(N+1,2)")
    end
    h = (tb-ta)/N
    TT = range(ta,step=h,length=N)
    MM = div(N,2)
    hF = zeros(Complex{Float64},MM+1)
    hF[1:(M+1)] = hf[1:(M+1)]
    PP = N*real.(irfft(hF,N))
    return (TT, PP)
end

function FourierCoeffsEven(f,ta,tb,M)
  N = 2*M
  h = (tb-ta)/N
  tt = range(ta,step=h,length=N)
  ff = f.(tt)
  hf = 1/N*rfft(ff)
  return (tt,ff,hf)
end

function FourierEvalEven(ta,tb,hf,N)
    M = length(hf)-1
    if (N<2M)
        error("$M=length(hf)-1 should be less than or equal to N/2=$div(N,2)")
    end
    h = (tb-ta)/N
    TT = range(ta,step=h,length=N)
    MM = div(N,2)
    hF = zeros(Complex{Float64},MM+1)
    hF[1:M] = hf[1:M]
    hF[M+1] = hf[M+1]/2
    PP = N*real.(irfft(hF,N))
    return (TT, PP)
end

function FourierCoeffs(tt,ff)
    N = length(ff)
    hf = 1/N*rfft(ff)
    return (tt,ff,hf)
end

function FourierEval(ta,tb,hf,N)
    if mod(N, 2) == 0
        return FourierEvalEven(ta,tb,hf,N)
    else
        return FourierEvalOdd(ta,tb,hf,N)
    end
end

function DiffMatrix(N)
    D = zeros(N,N)
    isodd(N) ? auxfcn = csc : auxfcn = cot
    for j in 1:N
        for k in 1:j-1
            D[j,k] = 0.5*(-1)^(j-k)*auxfcn((j-k)*pi/N)
            D[k,j] = -D[j,k]
        end
    end
    return D
end


@inline function rand_point()
    return 2.0 .* rand(3) .- 1.0
end

function sample_in_sphere(R::Number, r::Number, n::Integer, max_retry_times::Integer = 100)
    r > R && throw(ArgumentError("particle radius must be smaller than region radius!"))

    points = Vector{Float64}[]
    r′ = r / R
    for _ = 1:n
        retry_times = 0
        new_point = rand_point()
        while norm(new_point) + r′ > 1.0
            new_point = rand_point()
        end
        min_dist = minimum(norm(point - new_point) for point in points; init = Inf)
        while min_dist < 2r′ && retry_times < max_retry_times
            retry_times += 1
            new_point = rand_point()
            while norm(new_point) + r′ > 1.0
                new_point = rand_point()
            end
            min_dist = minimum(norm(point - new_point) for point in points; init = Inf)
        end
        if min_dist >= 2r′
            push!(points, new_point)
        end
    end

    length(points) < n &&
        @warn "Failed to generate enough points. $(length(points))/$n generated."

    return [point ./ r′ for point in points]
end

export sample_in_sphere


# Load Turing.
using Turing
# Load other dependencies
using Distributions, RDatasets

# Example data set - Fisher's Iris data set
data = dataset("datasets", "iris")
species = data[!, "Species"]
dat = data[!, 1:4]

@model pPCA(x, ::Type{T} = Float64) where {T} = begin

  # Dimensionality of the problem.
  N, D = size(x)

  # latent variable z
  z = Matrix{T}(undef, D, N)
  for d in 1:D
    z[d, :] ~ MvNormal(N, 1.)
  end

  # weights/loadings w
  w = Matrix{T}(undef, D, D)
  for d in 1:D
    w[d, :] ~ MvNormal(D, 1.)
  end

  mu = w * z

  for d in 1:D
    x[:,d] ~ MvNormal(mu[d,:], 1.)
  end

end

ppca = pPCA(dat)

# Hamiltonian Monte Carlo (HMC) sampler parameters
n_iterations = 10000
ϵ = 0.05
τ = 10

chain = sample(ppca, HMC(ϵ, τ), n_iterations)

# Extract paramter estimates for plotting.
w = reshape(mean(group(chain, :w))[:,2], (4,4))
z = reshape(mean(group(chain, :z))[:,2], size(dat))

#  plot(z[1, :], z[2,:])

## Extend to ARD
## Original paper
#
#  #  for (d in 1:D){
    #  w[d] ~ normal(0, sigma * alpha);
  #  }
  #  sigma ~ lognormal(0, 1);
  #  alpha ~ inv_gamma(1, 1);