using Plots: plot3d, plot, Animation, frame, gif, GridLayout, EmptyLayout, create_grid
using StaticArrays

"""
    Body(m, x, y, z, v)

Create a celestial body object with mass `m` kilograms, position `x` `y` `z`,
and 3-component velocity vector `v` (m/s).
"""
mutable struct Body{T <: Float64}
	# The mass of the body
	m::T

	# Its coords
	x::T
	y::T
	z::T

	# Its velocity as a vector of x y z values
	v::MVector{3, T}
end

"""
    getforce(b1, b2) -> float

Return the magnitude of the force between Body objects `b1` and `b2`.

Uses the standard gravitational equation.

See also: [`getforcevector`](@ref)
"""
function getforce(b1::Body, b2::Body)::Float64
	a = b1.x - b2.x
	b = b1.y - b2.y
	c = b1.z - b2.z

	r2 = a^2 + b^2 + c^2
	if r2 == 0; error("Cannot divide by 0. The positions of two bodies cannot be the same."); end

	# G \frac{m_1 m_2}{r^2}
	(6.674e-11 * b1.m * b2.m) / r2
end

"""
    getforcevector(b1, b2) -> [float, float, float]

Return the 3D vector of the force from Body object `b1` towards `b2`.

Expressed as a 3-component vector. Magnitude can be computed
by `getforce(b1, b2)`. Angles must be computed separately.

See also: [`getforce`](@ref)
"""
function getforcevector(b1::Body, b2::Body)::MVector{3, Float64}
	# azimuth is the angle on the x-y plane
	Δx = b2.x - b1.x
	Δy = b2.y - b1.y

	azimuth = atan(Δy, Δx)

	# altitude is the angle from the x-y plane in the z axis
	hyp = sqrt(Δx^2 + Δy^2)
	Δz = b2.z - b1.z
	altitude = atan(Δz, hyp)

	F = getforce(b1, b2)

	# We use F to compute the z component and the new hypotenuse to use
	# for computing the x and y components
	vec_z = F * sin(altitude)
	xy_plane_hyp = F * cos(altitude)

	# We then do some more trig to compute the x and y components
	vec_y = xy_plane_hyp * sin(azimuth)
	vec_x = xy_plane_hyp * cos(azimuth)

	MVector{3, Float64}(vec_x, vec_y, vec_z)
end

"""
    step!(bodies, Δt)

Apply forces between all Body objects in `bodies` over time step `Δt` seconds.
"""
function step!(bodies::Vector{Body}, Δt::Float64)
	# Calculate the force on each body and update its velocity accordingly
	for i in 1:length(bodies)
		b = bodies[i]

		forceonbody = copy(b.v)
		# We loop over every body and if it's not the same one, we add the add the force
		for j in 1:length(bodies)
			if j != i
				# We add the body's velocity to the force applied by gravity
				# to get the resultant force
				forceonbody .+= getforcevector(b, bodies[j])
			end

		end

		# Update the velocity of the body

		# a = (v - u)/t and F = ma ⟹  F = m * (v - u)/t
		# ⟹  v = Ft/m + u

		# Since u is the initial value, we can just use +=
		# We actually need .+= here to broadcast over the MVector
		b.v .+= ((forceonbody * Δt) / b.m)
	end

	# We need to wait until we've calculated all the forces before we start to change the positions,
	# otherwise we would change a position and then calculate a force based on the new position,
	# rather than the old one
	for i in 1:length(bodies)
		b = bodies[i]

		# Then we update the position of the body using its velocity
		b.x += Δt * b.v[1]
		b.y += Δt * b.v[2]
		b.z += Δt * b.v[3]
	end
end

"""
    drawframes(bodies, framecount, Δt)

Return frame data for the list of Body objects `bodies` with time step `Δt`.

This frame data is a Vector{Vector{SVector{3, Float64}}}. It can be thought of as a list of frames,
where each frame has a list of bodies, and each body has a list of x y z coordinates.

See also: [`drawgif`](@ref), [`creategif`](@ref)
"""
function drawframes(bodies::Vector{Body}, framecount::Int, Δt::Float64)::Vector{Vector{SVector{3, Float64}}}
	frames::Vector{Vector{SVector{3, Float64}}} = [[SVector{3, Float64}(bodies[i].x, bodies[i].y, bodies[i].z) for i in 1:length(bodies)]]

	for _ in 1:framecount
		step!(bodies, Δt)

		# We build the position data for this frame
		thisframe::Vector{SVector{3, Float64}} = [SVector{3, Float64}(bodies[i].x, bodies[i].y, bodies[i].z) for i in 1:length(bodies)]

		push!(frames, thisframe)
	end

	return frames
end

"""
    drawgif(positions, cube=false, bounds=nothing, filename=nothing)

Draw a GIF using the position data.

`positions` is a Vector{Vector{SVector{3, Float64}}}. It can be thought of as a list of frames,
where each frame has a list of bodies, and each body has a list of x y z coordinates.

`cube` is a bool for whether to draw the plot bounds as a cube.

`bounds` is a Vector{Tuple{Float64, Float64}} - a list of xlimits, ylimits, and zlimits.
If it's nothing (by default), then the bounds will be auto-generated.

`filename` is an optional filename. If none is provided, one will be programmatically generated,
following the pattern of `out.gif`, `out_1.gif`, `out_2.gif`, etc.

See also: [`drawframes`](@ref), [`creategif`](@ref)
"""
function drawgif(positions::Vector{Vector{SVector{3, Float64}}}, cube::Bool, bounds::Union{Nothing, Vector{Tuple{Float64, Float64}}}=nothing, filename::Union{String, Nothing}=nothing)
	# This is a multiplier to make the bounding box just a bit bigger than strictly necessary,
	# just to make it look a bit nicer
	multiplier = 1.05

	if cube && !isnothing(bounds)
		# We get the maximum absolute value in the bounds and use that to build the cube
		maxbound = multiplier * max(abs.(Iterators.flatten(bounds))...)

		xlimits = (-1 * maxbound, maxbound)
		ylimits = (-1 * maxbound, maxbound)
		zlimits = (-1 * maxbound, maxbound)

	elseif !cube && !isnothing(bounds)
		# We just use the bounds given, multiplied by the multiplier
		xlimits = multiplier .* bounds[1]
		ylimits = multiplier .* bounds[2]
		zlimits = multiplier .* bounds[3]

	elseif cube && isnothing(bounds)
		# This line gets absolute value of every float in positions and then gets the max
		# Because we're dealing with a nested vector, we have to flatten it, and then
		# flatten the result, and then splat it for max
		maxbound = multiplier * max(abs.(Iterators.flatten(Iterators.flatten(positions)))...)

		xlimits = (-1 * maxbound, maxbound)
		ylimits = (-1 * maxbound, maxbound)
		zlimits = (-1 * maxbound, maxbound)

	# Not cube or explicit bounds
	else
		# These list comprehensions loop through the list to get the coords
		# i is every frame in the positions list, and j is every body in each frame
		# We only need to check length(positions[1]) to get j because the length of
		# every sub-list is the same because every frame has the same number of bodies
		xs = [positions[i][j][1] for i in 1:length(positions) for j in 1:length(positions[1])]
		ys = [positions[i][j][2] for i in 1:length(positions) for j in 1:length(positions[1])]
		zs = [positions[i][j][3] for i in 1:length(positions) for j in 1:length(positions[1])]

		xlimits = multiplier .* extrema(xs)
		ylimits = multiplier .* extrema(ys)
		zlimits = multiplier .* extrema(zs)
	end

	# This is just the number of bodies
	n = length(positions[1])

	# We have 4 different subplots. xyz is the 3D perspective camera angle. The others are the orthogonal views.
	# We then put these subplots into a full plot to show everything comprehensively
	xyz = plot3d(n; legend=false, lw=1, title="$n Body Gravity Simulation", xlim=xlimits, ylim=ylimits, zlim=zlimits)
	# We're using standard 2D line plots here to avoid issues with cameras in 3D space
	xy = plot(n; legend=false, lw=0.5, title="Plan", titlefontsize=8, xlim=xlimits, ylim=ylimits, tickfontsize=4)
	xz = plot(n; legend=false, lw=0.5, title="Front", titlefontsize=8, xlim=xlimits, ylim=zlimits, tickfontsize=4)
	yz = plot(n; legend=false, lw=0.5, title="Side", titlefontsize=8, xlim=ylimits, ylim=zlimits, tickfontsize=4)

	# This is the expanded form of the @gif macro over a for loop
	# We're using the expanded form rather than the macro itself because
	# that lets us control the filename
	anim = Animation()
	counter = 1
	# For each frame
	for i = 1:length(positions)
		# For each body in this frame
		for j in 1:length(positions[i])
			# We add the position data of every body to its respective series in the plots
			bodydata = positions[i][j]
			push!(xyz, j, bodydata[1], bodydata[2], bodydata[3])
			push!(xy, j, bodydata[1], bodydata[2])
			push!(xz, j, bodydata[1], bodydata[3])
			push!(yz, j, bodydata[2], bodydata[3])
		end

		# Only if this counter is a multiple of 10, do we add this frame to the animation
		if counter % 10 == 0
			# I can't directly import the @layout macro, so we do it this way
			layout = eval(create_grid(:([a{0.7h}; b c d])))
			frame(anim, plot(xyz, xy, xz, yz; layout = layout, size = (1000, 1000)))
		end

		counter += 1
	end

	if isnothing(filename)
		if !in("out.gif", readdir())
			filename = "out.gif"
		else
			i = 1
			while in("out_$i.gif", readdir())
				i += 1
			end
			filename = "out_$i.gif"
		end
	end

	gif(anim, filename)
end

"""
    creategif(bodies, framecount, Δt, cube=false, bounds=nothing, filename=nothing)

Generate a GIF from the list of Body objects `bodies`, with `framecount` frames and a time step of `Δt` seconds.

`cube` is a bool for whether to draw the plot bounds as a cube.

`bounds` is a Vector{Tuple{Float64, Float64}} - a list of xlimits, ylimits, and zlimits.
If it's nothing (by default), then the plot bounds will be auto-generated.

`filename` is an optional filename. If none is provided, one will be programmatically generated,
following the pattern of `out.gif`, `out_1.gif`, `out_2.gif`, etc.

See also: [`drawgif`](@ref), [`drawframes`](@ref)
"""
creategif(bodies::Vector{Body}, framecount::Int, Δt::Float64, cube::Bool=false, bounds::Union{Nothing, Vector{Tuple{Float64, Float64}}}=nothing, filename::Union{String, Nothing}=nothing) = drawgif(drawframes(bodies, framecount, Δt), cube, bounds, filename)


### A Pluto.jl notebook ###
# v0.15.1

using Markdown
using InteractiveUtils

# ╔═╡ 28042615-8f2b-41fa-9b14-28d69180debb
using Pkg; Pkg.develop(path=".."); Pkg.activate("..")

# ╔═╡ 3ec6aae0-c233-11eb-355d-3dfe780e33bc
begin
	using DataFrames
	using JuMP
	using Ipopt
	using StatsPlots
	plotly()
end

# ╔═╡ a6997248-aee3-4ebb-8a1e-f791ad503aa8
begin
	n = 100
	μ1 = 10
	μ2 = 15

	df = append!(
		DataFrame(x1=randn(n).+μ1, x2=randn(n).+μ1, y=ones(n)),
		DataFrame(x1=randn(n).+μ2, x2=randn(n).+μ2, y=-ones(n))
	)

	p = @df df scatter(:x1, :x2, group=:y, legend=:topleft)
end

# ╔═╡ bf7d0459-4b47-4839-800d-01bdcd500633
md"
$arg\ min_{𝐰, b}\ \frac{1}{2}𝐰^T𝐰$

$\text{subject to }\forall i, y_i(𝐰^T𝐱_i+b) \geq 1$
"

# ╔═╡ 27dfce4e-3ccd-4221-8eca-144ae87e2998
md"
respect to

$arg\ min_{𝐰, b}\ \frac{1}{2}𝐮^T 𝐐 𝐮 + 𝐏^T 𝐮$

$\text{subject to }\forall i, 𝐚_i^T 𝐮) \geq c_i$

Therefore,

$𝐮 = \begin{bmatrix}b \\ 𝐰\end{bmatrix}$
$𝐐 = \begin{bmatrix}0 & 0^T_d \\ 0^T_d & 𝐈_d\end{bmatrix}$
$𝐏 = 0_{d+1}$
$𝐚^T_i = y_i\begin{bmatrix}1 & 𝐱^T_i\end{bmatrix}$
$c_i = 1$
"

# ╔═╡ 5056cc0a-f832-47b3-a6a5-f7ee2d7bc5ba
begin
	x = Matrix(df[:, 1:2])
	y = Vector(df[:, 3])

	function solve_svm(x, y)
		svm = Model(Ipopt.Optimizer)
		@variable(svm, w[1:2])
		@variable(svm, b)
		@objective(svm, Min, 0.5 * w' * w)
		@constraint(svm, y.*(x * w .+ b) .>= 1)
		optimize!(svm)

		@show objective_value(svm)

		return value.(w), value(b)
	end

	w, b = solve_svm(x, y)
end

# ╔═╡ e5334dab-56b2-4dd0-a12c-e17f1c8e584b
ŷ(x1, x2) = sign([x1, x2]' * w + b)

# ╔═╡ 16448a5a-090f-4a6a-a9a0-55192a77dfb3
heatmap!(p, 5:0.1:20, 5:0.1:20, ŷ, color=:coolwarm)

# ╔═╡ Cell order:
# ╟─28042615-8f2b-41fa-9b14-28d69180debb
# ╠═3ec6aae0-c233-11eb-355d-3dfe780e33bc
# ╠═a6997248-aee3-4ebb-8a1e-f791ad503aa8
# ╟─bf7d0459-4b47-4839-800d-01bdcd500633
# ╟─27dfce4e-3ccd-4221-8eca-144ae87e2998
# ╠═5056cc0a-f832-47b3-a6a5-f7ee2d7bc5ba
# ╠═e5334dab-56b2-4dd0-a12c-e17f1c8e584b
# ╠═16448a5a-090f-4a6a-a9a0-55192a77dfb3


function solve_002()
    fib_vec = [1, 2]
    i = 3
    # Generate fibonacci numbers
    while fib_vec[end] < 4000000
        push!(fib_vec, fib_vec[i-1] + fib_vec[i-2])
        i += 1
    end
    # Sum even terms
    sum(filter(x -> x % 2 == 0, fib_vec))
end


module SampleProject

export sin_sum, cos_sum, geometric_series, pi_gregory_leibniz, pi_nilakantha

sin_sum(n) = sum(sin, range(0, 2pi, length=n))

cos_sum(n) = sum(cos, range(0, 2pi, length=n))

geometric_series(a, r, n) = sum(k -> a * r^k, 0:n)

pi_gregory_leibniz(n) = sum(k -> 4/(4k+1) - 4/(4k+3), 0:n)

pi_nilakantha(n) = 3 + sum(k -> 4/((4k-2)*(4k-1)*4k) - 4/(4k*(4k+1)*(4k+2)), 1:n)

end


module Shamir

using Polynomials

function l(x, j, k)
"""
    Create a Lagrange basis polynomial
    Reference: https://wikimedia.org/api/rest_v1/media/math/render/svg/6e2c3a2ab16a8723c0446de6a30da839198fb04b
"""
    polys = []
    for m in 1:k
        if m != j
            d = x[j] - x[m]
            r = Poly([-1 * x[m], 1]) / d
            push!(polys, r)
        end
    end

    #println(polys)

    return prod(polys)
end

function L(x, y, k)
"""
    Create a linear combination of Lagrange basis polynomials
    Reference: https://wikimedia.org/api/rest_v1/media/math/render/svg/d07f3378ff7718c345e5d3d4a57d3053190226a0
"""

    s = []

    for j in 1:k
        r = y[j] * l(x, j, k)
        push!(s, r)
    end

    #println(sum(s))

    return sum(s)
end

function construct_shares(n, production_poly)
"""
    Create shares of the secret

    Parameters:
    n: Total number of shares
    production_poly : The Polynomial created with the secret
                    to produce secret shares
"""

    share = []

    for i in 1:n
        push!(share, [i, production_poly(i)])
    end

    return share
end

function recover_secret(shares, n, k, p)
"""
    Recover the secret by finding the coefficient of x^0 i.e.
    when the x=0 in the polynomial.

    Parameters:
    shares : An array of shares of the individuals.
    n : total number of shares
    k: minimum number of shares required to unravel the secret
    p: Field number (to restrict the computation space)
"""
    if length(shares) < k
        throw("Need more parties!")
    end

    x = []
    y = []
    for i in 1:n
        push!(x, shares[i][1])
        push!(y, shares[i][2])
    end

    f = L(x, y, k)
    f = mod(f(0), p)

    return f
end

export Shamir

end


export mean_squared_error

"""
    MeanSquaredError <: Func
"""
@func mutable struct MeanSquaredError end

function forward(f::MeanSquaredError, x1, x2)
    diff = x1 .- x2
    y = sum(diff.^2) / length(diff)
    return y
end

function backward(f::MeanSquaredError, gy)
    x1, x2 = f.args
    diff = x1 - x2
    gx1 = gy .* diff .* (2.0 / length(diff))
    gx2 = -gx1
    return gx1, gx2
end

mean_squared_error(x1, x2) = MeanSquaredError()(x1, x2)


function e006(N)
    return difference(N)
end

function difference(N)
    sum_squares = begin
        s = 0
        for i = 1:N
            s += i ^ 2
        end
        s
    end

    square_sum = begin
        s = 0
        for i = 1:N
            s += i
        end
        s ^ 2
    end

    return square_sum - sum_squares
end

function solution006()
    N = 100
    answer = e006(N)
    println("The answer to problem 006 is $answer")
end


N, faces = 10^6, 1:6

numSol = sum([iseven(i+j) for i in faces, j in faces]) / length(faces)^2
mcEst  = sum([iseven(rand(faces) + rand(faces)) for i in 1:N]) / N

println("Numerical solution = $numSol \nMonte Carlo estimate = $mcEst")


"""
    𝑛(v::Vector{T}) where T<:Number

Scales a length-``n`` vector ``v``  such that the last component
of the vector is one, provided that the last component is not zero. If the last
component is zero then the vector is left unchanged.

# Details
Suppose the length-``n`` vector ``v`` represents the homogeneous coordinates  of
a point in a projective space. The corresponding Cartesian coordinates  usually
just the first ``n-1`` numbers of homogeneous coordinates divided by  the last
component. So if the last component is one, then the first  ``n-1`` homogeneous
coordinates can be interpreted as Cartesian.  The exceptional case is when the
last component of the homogenenous coordinates is zero. These homogeneous
coordinates are associated with so-called *points at infinity* and have no
Cartesian counterparts.

# Example
```julia
h = [4, 4 , 2]
c = 𝑛(h)

3-element Array{Float64,1}:
 2.0
 2.0
 1.0
```

"""
function 𝑛(v::AbstractArray)
    if v[end] != 0 && v[end] != 1
        v .= v ./ v[end]
    else
        v
    end
end

function 𝑛(v::SVector)
    if v[end] != 0 && v[end] != 1
        v / v[end]
    else
        v
    end
end

function hom⁻¹(v::SVector)
    if isapprox(v[end], 0.0; atol = 1e-14)
        pop(v)
    else
        pop(v / v[end])
    end
end

function hom(v::SVector)
    push(v,1)
end

function ∂hom⁻¹(𝐧::SVector)
    k = length(𝐧)
    𝐞ₖ = push(zeros(SVector{k-1}),1.0)
    𝐈 = SMatrix{3,3}(1.0I)
    1/𝐧[k]*𝐈 - 1/𝐧[k]^2 * 𝐧 * 𝐞ₖ'
end

function ∂𝑛(𝐧::AbstractArray)
    k = length(𝐧)
    𝐞ₖ = fill(0.0,(k,1))
    𝐞ₖ[k] = 1
    1/𝐧[k]*Matrix{Float64}(I, k, k) - 1/𝐧[k]^2 * 𝐧 * 𝐞ₖ'
end

function smallest_eigenpair(A::AbstractArray)
    F = eigen(A)
    index = argmin(F.values)
    (F.values[index], F.vectors[:,index])
end

function smallest_eigenpair(A::AbstractArray,B::AbstractArray)
    F = eigfact(A,B)
    index = indmin(F[:values])
    (F[:values][index], F[:vectors][:,index])
end

function vec2antisym(v::AbstractArray)
    if length(v) != 3
         throw(ArgumentError("The operation is only defined for a length-3 vector."))
    end
    𝐒  = @SMatrix [   0  -v[3]    v[2] ;
                    v[3]    0    -v[1] ;
                   -v[2]  v[1]      0]

end


"""
    fdweights(z, x, m)

Calculates FD (finite difference) weights for arbitrarily spaced nodes in 1D.

Return an array of size m+1 by length(x) containing, in successive rows, the
weights for derivatives 0, 1, ..., m. The zero (m = 0) derivative corresponds
to interpolation.

# Arguments
- `z`: location at which to approximate the derivative (may but need to be a grid point)
- `x`: vector of x-coordinates (grid points; distinct, otherwise arbitrary)
- `m`: highest derivative for which weights are sought (default value is 0)

# Examples
```julia-repl
julia> fdweights(0, -2:2, 6)
5×5 adjoint(::Matrix{Float64}) with eltype Float64:
 -0.0         0.0        1.0   0.0       -0.0
  0.0833333  -0.666667   0.0   0.666667  -0.0833333
 -0.0833333   1.33333   -2.5   1.33333   -0.0833333
 -0.5         1.0        0.0  -1.0        0.5
  1.0        -4.0        6.0  -4.0        1.0
```
"""
function fdweights(z, x, m = 0)
    n = length(x)
    if m >= n
        m = n - 1 # set to highest possible derivative
    end

    c1 = 1
    c4 = x[1] - z
    C = zeros(n, m+1)
    C[1, 1] = 1

    for i = 1:n-1
        i1 = i + 1
        mn = min(i, m)
        c2 = 1
        c5 = c4
        c4 = x[i1] - z
        for j = 0:i-1
            j1 = j + 1
            c3 = x[i1] - x[j1]
            c2 = c2 * c3
            if j == i-1
                for s = mn:-1:1
                    s1 = s + 1
                    C[i1, s1] = c1 * (s * C[i1-1, s1-1] - c5 * C[i1-1, s1]) / c2
                end
                C[i1, 1] = -c1 * c5 * C[i1-1, 1] / c2
            end
        for s = mn:-1:1
            s1 = s + 1
            C[j1, s1] = (c4 * C[j1, s1] - s*C[j1, s1-1]) / c3
        end
        C[j1, 1] = c4 * C[j1,1] / c3
        end
        c1 = c2
    end
    return C'
end


module kuramoto_sivashinsky

using SparseArrays
using LinearAlgebra

"""
    B = calcBernoulli(n)

Return the `n`th Bernoulli number.  Follows the convention that `n=1`
Bernoulli number is equal to -1/2.
"""
function calcBernoulli(n)
    A = Vector{Rational{BigInt}}(undef, n + 1)
    for m = 0:n
        A[m + 1] = 1 // (m + 1)
        for j = m:-1:1
            A[j] = j * (A[j] - A[j + 1])
        end
    end
    if n == 1
        return -A[1]
    else
        return A[1]
    end
end

"""
    quad = buildQuadrature(order [, T=Float64])

Return a Gregory-type quadrture rule for uniformly spaced data.
"""
function buildQuadrature(order::Int)
    # Form the linear system to solve
    numbnd = order - 1
    A = zeros(Rational{BigInt}, order-1, numbnd)
    b = zeros(Rational{BigInt}, order-1)
    for j = 1:order-1
        for i = 1:numbnd
            A[j,i] = j*(numbnd - (i-1))^(j-1)
        end
        b[j] = numbnd^j - ((-1)^j)*calcBernoulli(j)
    end
    w = A\b
    return w
    # if order <= 2
    #     quad = [1//2]
    # elseif order <= 4
    #     quad = [17//48; 59//48; 43//48; 49//48]
    # elseif order <= 6
    #     quad = [13649//43200; 12013//8640; 2711//4320; 5359//4320; 7877//8640;
    #             43801//43200]
    # else
    #     quad = [1498139//5080320; 1107307//725760; 20761//80640;
    #             1304999//725760; 299527//725760; 103097//80640;
    #             670091//725760; 5127739//5080320]
    # end
    # return convert(Array{T,1}, quad)
end

"""
Data structure for finite-difference operators on uniformly spaced points
"""
struct FiniteDiff{T}
    """Defines the nodes in the stencil; node 0 is rel. location of operator"""
    stencil::Array{Int,1}
    """Actual finite-difference coefficients"""
    coeffs::Array{T,1}
end

"""
    fd = buildFiniteDifference(deriv, order [, T=Float64, h=1.0])

Returns a finite difference operator struct for a derivative of order `deriv`
and order of accuracy of `order`. The optional argument `T` controls the
precision used, and `h` sets the (uniform) mesh spacing.
"""
function buildFiniteDifference(deriv::Int, order::Int, T::Type=Float64;
                               h=convert(T,1) )
    num_pts = deriv + order - 1
    if iseven(num_pts) num_pts += 1 end
    stencil = Array{Int}(-div(num_pts-1,2):div(num_pts-1,2))
    V = zeros(T, num_pts, num_pts)
    for i = 1:num_pts
        V[i,:] = stencil.^(i-1)
    end
    b = zeros(T, num_pts)
    b[deriv+1] = factorial(deriv)
    c = V\b
    c /= (h^deriv)
    return FiniteDiff{T}(stencil, c)
end

"""
    interp = buildInterp(order [, T=Float64])

Returns a centered interpolation operator of order `order` that does **not**
use data at the central node.
"""
function buildInterp(order::Int, T::Type=Float64)
    num_pts = order
    if isodd(num_pts) num_pts += 1 end
    stencil = Array{Int}([-div(num_pts,2):-1;1:div(num_pts,2)])
    V = zeros(T, num_pts, num_pts)
    for i = 1:num_pts
        V[i,:] = stencil.^(i-1)
    end
    b = zeros(T, num_pts)
    b[1] = 1.0
    c = V\b
    return FiniteDiff{T}(stencil, c)
end

"""
    num_ghost, num_bndry, ghost_op = buildGhostOp(order [, T=Float64])

Constructs a matrix operator `ghost_op` that, when applied to `num_bndry` nodes
at the boundary, produces `order` accurate `num_ghost` values.  **Note**, this
operator is designed to enforce a zero function value and zero derivative value
at the boundary, as per the Blonigan and Wang paper.

The ghost operator is constructed by solving a linearly constrained minimum
norm problem for the ghost node values.
"""
function buildGhostOp(order::Int, T::Type=Float64)
    # first, build the constraint Jacobian that enforces u=0 and du/dx = 0
    # A*u_ghost - B*u_true = c
    fd = buildFiniteDifference(1, order, T)
    #interp = buildInterp(order, T)
    num_nodes = length(fd.stencil)
    num_ghost = div(num_nodes-1,2) + 1 # this includes node at boundary
    num_bndry = num_nodes - num_ghost
    A = zeros(T, 2, num_ghost)
    A[1,end] = 1.0 # last node is the node on the boundary
    #A[1,1:end-1] = interp.coeffs[1:num_ghost-1]
    A[2,:] = fd.coeffs[1:num_ghost]
    B = zeros(T, 2, num_bndry)
    #B[1,:] = interp.coeffs[num_ghost:end]
    B[2,:] = fd.coeffs[num_ghost+1:end]
    return num_ghost, num_bndry, -A'*( (A*A') \ B )

    if false
        # This version uses a TVD-type regularization
        # first, build the constraint Jacobian that enforces u=0 and du/dx = 0
        # A*u_ghost = B*u_true + c
        fd = buildFiniteDifference(1, order, T)
        #interp = buildInterp(order, T)
        num_nodes = length(fd.stencil)
        num_ghost = div(num_nodes-1,2) + 1 # this includes node at boundary
        num_bndry = num_nodes - num_ghost
        A = zeros(T, 2, num_ghost)
        A[1,end] = 1.0 # last node
        #A[1,1:end-1] = interp.coeffs[1:num_ghost-1]
        A[2,:] = fd.coeffs[1:num_ghost]
        B = zeros(T, 2, num_bndry)
        #B[1,:] = interp.coeffs[num_ghost:end]
        B[2,:] = fd.coeffs[num_ghost+1:end]

        D = zeros(T, num_ghost, num_ghost)
        D += diagm(2*ones(T,num_ghost))
        D -= diagm(ones(T, num_ghost-1), 1)
        D -= diagm(ones(T, num_ghost-1), -1)
        K = [D A'; A zeros(T, 2, 2)]
        G = -inv(K)[1:num_ghost,end-1:end]*B
        return num_ghost, num_bndry, G
    end
end

"""
Data for KS function evaluation to avoid reallocation
"""
struct KSData{T}
    """Number of internal nodes (nodes we are solving for)"""
    num_nodes::Int
    """Number of ghost nodes at each end; include x=0 and x=L"""
    num_ghost::Int
    """Number of internal nodes used to determine ghost values"""
    num_bndry::Int
    """Stores the boundary first-derivative operator"""
    bndry_op::Array{T,2}
    """Stores the linear operatrors"""
    linear_op::SparseMatrixCSC{T,Int}
    """finite difference operator for the first-derivative"""
    fd::FiniteDiff{T}
end

"""
    ks = buildKSDataPeriodic(order, num_nodes [, T=Float64])

Returns a KSData structure that can be used to solve the Kuramoto-Sivashinsky
equations on _periodic domains_, while avoiding reallocation of arrays.
"""
function buildKSDataPeriodic(order::Int, num_nodes::Int, T::Type=Float64)
    # domain length used in Blonigan and Wang paper, but here periodic
    Lx = convert(T, 128)
    h = Lx/num_nodes
    # generate the matrix used for the linear operators
    A = zeros(T, num_nodes, num_nodes)
    # second, add the Laplacian contributions
    fd = buildFiniteDifference(2, order, T)
    fac = 1/(h*h)
    for i = 1:num_nodes
        A[i, mod.(fd.stencil .+ i .- 1, num_nodes) + 1] -= fd.coeffs*fac
    end
    # finally, add the BiLaplacian terms
    fd = buildFiniteDifference(4, order, T)
    fac *= fac
    for i = 1:num_nodes
        A[i, mod.(fd.stencil .+ i .- 1, num_nodes) + 1] -= fd.coeffs*fac
    end
    Asparse = sparse(A)
    # finally, get the first-derivative operator for nonlinear terms
    fd = kuramoto_sivashinsky.buildFiniteDifference(1, order, T, h=h)
    return KSData{T}(num_nodes, 0, 0, zeros(0,0), Asparse, fd)
end

"""
    getKSFuncPeriodic!(ks, t, u, du)

A function that returns the perodic spatial discretization of the
Kuramoto-Sivashinsky PDE for use in a time marching method.
"""
function getKSFuncPeriodic!(ks::KSData{Tks}, t, u::AbstractArray{T,1},
                            du::AbstractArray{T,1}) where {Tks, T}
    # apply the linear operators
    du[:] = ks.linear_op*u
    # add the nonlinear terms; this uses a skew-symmetric split-form
    third = convert(T, 1//3)
    for i = 1:ks.num_nodes
        for k = 1:size(ks.fd.stencil,1)
            j = mod(ks.fd.stencil[k]+i-1, ks.num_nodes) + 1
            du[i] -= third*u[i]*ks.fd.coeffs[k]*u[j]
            du[i] -= third*ks.fd.coeffs[k]*u[j]*u[j]
        end
        # This version uses central differencing for the nonlinear term
        #du[i] -= u[i]*dot(ks.fd.coeffs,
        #                  u[mod.(ks.fd.stencil+i-1,ks.num_nodes)+1])
    end
    return nothing
end

"""
    getKSJacPeriodic!(ks, t, u, Jac)

A function that returns the Jacobian of the perodic spatial discretization of
the Kuramoto-Sivashinsky PDE for use in a time marching method.
"""
function getKSJacPeriodic!(ks::KSData{T}, t, u::AbstractArray{T,1},
                           Jac::SparseMatrixCSC{T,Int}) where {T}
    Jac[:,:] = sparse(ks.linear_op)
    third = convert(T, 1//3)
    for i = 1:ks.num_nodes
        for k = 1:size(ks.fd.stencil,1)
            j = mod(ks.fd.stencil[k]+i-1, ks.num_nodes) + 1
            #du[i] -= fac*u[i]*ks.fd.coeffs[k]*u[j]
            Jac[i,i] -= third*ks.fd.coeffs[k]*u[j]
            Jac[i,j] -= third*u[i]*ks.fd.coeffs[k]
            #du[i] -= fac*ks.fd.coeffs[k]*u[j]*u[j]
            Jac[i,j] -= 2*third*ks.fd.coeffs[k]*u[j]
        end
        # This version uses central differencing for the nonlinear term
        #Jac[i,i] -= dot(ks.fd.coeffs,
        #                u[mod.(ks.fd.stencil+i-1,ks.num_nodes)+1])
        #Jac[i,mod.(ks.fd.stencil+i-1,ks.num_nodes)+1] -= u[i]*ks.fd.coeffs
    end
end

"""
    ks = buildKSData(order, num_nodes [, T=Float64])

Returns a KSData structure that can be used to solve the Kuramoto-Sivashinsky
equations while avoiding reallocation of arrays.  This version imposes
homogeneous Dirichlet and Neumann conditions at the ends of the spatial domain.
"""
function buildKSData(order::Int, num_nodes::Int, T::Type=Float64)
    # domain length used in Blonigan and Wang paper
    Lx = convert(T, 128)
    h = Lx/(num_nodes+1)
    # get the ghost node operator, call it Bg
    num_ghost, num_bndry, Bg = buildGhostOp(order, T)
    # println("Bg = ", Bg)

    # generate the matrix used for the linear operators
    A = zeros(T, num_nodes, num_nodes)
    # first, add the Laplacian contributions
    fd = buildFiniteDifference(2, order, T)
    fac = 1/(h*h)
    for i = 1:num_nodes
        idx = [k for k = 1:size(fd.stencil,1)
               if fd.stencil[k] + i >= 1 && fd.stencil[k] + i <= num_nodes]
        A[i, fd.stencil[idx] .+ i] -= fd.coeffs[idx]*fac
    end

    # eliminate ghost node contributions for Laplacian
    Ag = zeros(T, num_bndry, num_ghost)
    for i = 1:num_bndry
        idx = [k for k = 1:size(fd.stencil,1) if fd.stencil[k] + i < 1]
        Ag[i, fd.stencil[idx] .+ num_ghost .+ i] -= fd.coeffs[idx]*fac
    end
    # println("Ag (Laplacian) = ", Ag)
    G = Ag*Bg
    A[1:num_bndry,1:num_bndry] += G
    A[end:-1:end-num_bndry+1,end:-1:end-num_bndry+1] += G

    # next, add the BiLaplacian terms
    fd = buildFiniteDifference(4, order, T)
    fac = 1/(h*h*h*h)
    for i = 1:num_nodes
        idx = [k for k = 1:size(fd.stencil,1)
               if fd.stencil[k] + i >= 1 && fd.stencil[k] + i <= num_nodes]
        A[i, fd.stencil[idx] .+ i] -= fd.coeffs[idx]*fac
    end

    # eliminate ghost node contributions for BiLaplacian
    fill!(Ag, zero(T))
    for i = 1:num_bndry
        idx = [k for k = 1:size(fd.stencil,1) if fd.stencil[k] + i < 1]
        Ag[i, fd.stencil[idx] .+ num_ghost .+ i] -= fd.coeffs[idx]*fac
    end
    # println("Ag (BiLaplacian) = ", Ag)
    G = Ag*Bg
    A[1:num_bndry,1:num_bndry] += G
    A[end:-1:end-num_bndry+1,end:-1:end-num_bndry+1] += G

    # println("A[1,1:5] = ",A[1,1:5])
    # println("A[end,end-4:end] = ",A[end,end-4:end])

    Asparse = sparse(A)

    # finally, get the first-derivative operator for nonlinear terms
    fd = kuramoto_sivashinsky.buildFiniteDifference(1, order, T, h=h)
    # Get the first-derivative operator at the boundary
    fill!(Ag, zero(T))
    G = zeros(T, num_bndry, div(size(fd.stencil,1)+1,2) + num_bndry - 1 )
    for i = 1:num_bndry
        idx = [k for k = 1:size(fd.stencil,1) if fd.stencil[k] + i < 1]
        Ag[i, fd.stencil[idx] .+ num_ghost .+ i] += fd.coeffs[idx]
        idx = [k for k = 1:size(fd.stencil,1) if fd.stencil[k] + i >= 1]
        G[i, fd.stencil[idx] .+ i] += fd.coeffs[idx]
    end
    # println("Ag (first-derivative) = ", Ag)
    # println("G before = ",G)
    G[:,1:num_bndry] += Ag*Bg
    # println("G after = ",G)

    return KSData{T}(num_nodes, num_ghost, num_bndry, G, Asparse, fd)
end

"""
    getKSFunc!(ks, t, u, du)

A function that returns the spatial discretization of the
Kuramoto-Sivashinsky PDE for use in a time marching method.  Homogeneous BCs
for u and dudx are used at both ends of the domain.
"""
function getKSFunc!(ks::KSData{Tks}, t, u::AbstractArray{T,1},
                    du::AbstractArray{T,1}) where {Tks, T}
    # apply the linear operators
    du[:] = ks.linear_op*u
    # add the nonlinear terms to interior nodes
    fac = convert(T, 1//3)
    for i = ks.num_bndry + 1:ks.num_nodes - ks.num_bndry
        for k = 1:size(ks.fd.stencil,1)
            j = mod(ks.fd.stencil[k]+i-1, ks.num_nodes) + 1
            du[i] -= fac*u[i]*ks.fd.coeffs[k]*u[j]
            du[i] -= fac*ks.fd.coeffs[k]*u[j]*u[j]
        end
        #du[i] -= u[i]*dot(ks.fd.coeffs,
        #                  u[mod.(ks.fd.stencil+i-1,ks.num_nodes)+1])
    end
    # add the nonlinear terms to the boundary nodes
    for i = 1:size(ks.bndry_op,1)
        for k = 1:size(ks.bndry_op,2)
            du[i] -= fac*u[i]*ks.bndry_op[i,k]*u[k]
            du[i] -= fac*ks.bndry_op[i,k]*u[k]*u[k]
            du[end-i+1] += fac*u[end-i+1]*ks.bndry_op[i,k]*u[end-k+1]
            du[end-i+1] += fac*ks.bndry_op[i,k]*u[end-k+1]*u[end-k+1]
        end
    end
    return nothing
end

"""
    getKSJac!(ks, t, u, Jac)

A function that returns the Jacobian of the spatial discretization of
the Kuramoto-Sivashinsky PDE for use in a time marching method. Homogeneous BCs
for u and dudx are used at both ends of the domain.
"""
function getKSJac!(ks::KSData{T}, t, u::AbstractArray{T,1},
                   Jac::SparseMatrixCSC{T,Int}) where {T}
    # Jacobian contribution due to linear terms
    Jac[:,:] = sparse(ks.linear_op)
    #Jac[:,:] = deepcopy(ks.linear_op)
    # Jacobian due to nonlinear terms on interior
    fac = convert(T, 1//3)
    for i = ks.num_bndry+1:ks.num_nodes - ks.num_bndry
        for k = 1:size(ks.fd.stencil,1)
            j = mod(ks.fd.stencil[k]+i-1, ks.num_nodes) + 1
            # du[i] -= fac*u[i]*ks.fd.coeffs[k]*u[j]
            Jac[i,i] -= fac*ks.fd.coeffs[k]*u[j]
            Jac[i,j] -= fac*u[i]*ks.fd.coeffs[k]
            # du[i] -= fac*ks.fd.coeffs[k]*u[j]*u[j]
            Jac[i,j] -= 2*fac*ks.fd.coeffs[k]*u[j]
        end
        # This version uses central differencing for the nonlinear term
        #Jac[i,i] -= dot(ks.fd.coeffs,
        #                u[mod.(ks.fd.stencil+i-1,ks.num_nodes)+1])
        #Jac[i,mod.(ks.fd.stencil+i-1,ks.num_nodes)+1] -= u[i]*ks.fd.coeffs
    end
    # Jacobian due to nonlinear terms on boundary
    for i = 1:size(ks.bndry_op,1)
        for k = 1:size(ks.bndry_op,2)
            # du[i] -= fac*u[i]*ks.bndry_op[i,k]*u[k]
            Jac[i,i] -= fac*ks.bndry_op[i,k]*u[k]
            Jac[i,k] -= fac*u[i]*ks.bndry_op[i,k]
            # du[i] -= fac*ks.bndry_op[i,k]*u[k]*u[k]
            Jac[i,k] -= 2*fac*ks.bndry_op[i,k]*u[k]
            # du[end-i+1] += fac*u[end-i+1]*ks.bndry_op[i,k]*u[end-k+1]
            Jac[end-i+1,end-i+1] += fac*ks.bndry_op[i,k]*u[end-k+1]
            Jac[end-i+1,end-k+1] += fac*u[end-i+1]*ks.bndry_op[i,k]
            # du[end-i+1] += fac*ks.bndry_op[i,k]*u[end-k+1]*u[end-k+1]
            Jac[end-i+1,end-k+1] += 2*fac*ks.bndry_op[i,k]*u[end-k+1]

            # This version uses central differencing for the nonlinear term
            #du[i] -= u[i]*ks.bndry_op[i,k]*u[k]
            #  Jac[i,i] -= ks.bndry_op[i,k]*u[k]
            #  Jac[i,k] -= u[i]*ks.bndry_op[i,k]
            #du[end-i+1] -= u[end-i+1]*ks.bndry_op[i,k]*u[end-k+1]
            #  Jac[end-i+1,end-i+1] += ks.bndry_op[i,k]*u[end-k+1]
            #  Jac[end-i+1,end-k+1] += u[end-i+1]*ks.bndry_op[i,k]
        end
    end
end

"""
    Jac = getJacobian(u, ks)

This function uses the complex-step method to get the Jacobian; used for
verification of getKSJac!.  **Do not use** in simulations, due to cost.
"""
function getJacobian(u::AbstractArray{T,1}, ks::KSData{T}) where{T}
    fd_c = FiniteDiff{Complex128}(ks.fd.stencil, ks.fd.coeffs)
    ks_c = KSData{Complex128}(ks.num_nodes, ks.num_ghost, ks.num_bndry,
                              ks.bndry_op, ks.linear_op, fd_c)
    Jac = zeros(size(u,1),size(u,1))
    u_c = complex.(u, 0.0)
    du_c = zeros(u_c)
    ceps = 1e-60
    for i = 1:ks.num_nodes
        u_c[i] += complex(0.0, ceps)
        getKSFunc!(ks_c, 0.0, u_c, du_c)
        Jac[:,i] = imag.(du_c)/ceps
        u_c[i] -= complex(0.0, ceps)
    end
    return Jac
end

"""
    sol = solveUsingMidpoint(ks, Time, num_stpes, u)

Solve the KS problem on the nonperiodic domain using the midpoint rule.  The
problem is run from t=0 to t=`Time` using `num_steps` steps.  The initial
condition is `u`.  Returns `sol`, the solution over all space and temporal
nodes; `sol[i,n]` stores the `i`th space node at time step `n`.  Note that
`sol` stores the boundary nodes, whereas `u` does not.
"""
function solveUsingMidpoint(ks::KSData{T}, Time::T, num_steps::Int,
                            u::Array{T,1}) where {T}
    num_nodes = size(u,1)
    # storage for solution over all time steps
    dt = Time/num_steps
    sol = zeros(num_nodes+2, num_steps+1)
    sol[2:end-1,1] = u

    t = convert(T, 0.0)
    Jac = similar(ks.linear_op)
    F = lu(ks.linear_op) # initialize memory for the lu factorization
    u_mid = zeros(size(u))
    u_old = zeros(size(u))
    r = zeros(size(u))
    newt_tol = 1e-10 # We will need to ensure this is sufficiently small
    max_newt_iter = 50
    for n = 1:num_steps
        #println("Iteration ",n,": time ",t)
        # Newton iterations
        u_old = u
        norm0 = one(T)
        for k = 1:max_newt_iter
            u_mid = 0.5*u + 0.5*u_old
            kuramoto_sivashinsky.getKSFunc!(ks, t, u_mid, r)
            r = u - u_old - dt*r
            #println("\tNewton iter ",k,": res norm = ",norm(r))
            if k == 1
                norm0 = norm(r)
            else
                if norm(r) < newt_tol*norm0
                    break
                end
            end
            kuramoto_sivashinsky.getKSJac!(ks, t, u_mid, Jac)
            Jac *= 0.5*dt
            Jac -= I #speye(T, num_nodes)
            lu!(F, Jac)
            du = F\r
            u += du
        end
        t += dt
        sol[2:end-1,n+1] = u
    end
    return sol
end

"""
    avg = calcSolutionAverage(order, sol, func)

Returns the space and time average of `func(sol[:,:])` using an `order` order
accurate diagonal norm scheme.
"""
function calcSolutionAverage(order::Int, sol::AbstractArray{T,2},
                             func::Function) where {T}
    avg = zero(T)
    quad = convert(Vector{T}, buildQuadrature(order))
    num_bndry = size(quad,1)
    for n = 1:size(sol,2)
        wt = one(T)
        if n <= num_bndry
            wt *= quad[n]
        elseif n >= size(sol,2) - num_bndry + 1
            wt *= quad[size(sol,2)-n+1]
        end
        for i = 1:size(sol,1)
            if i <= num_bndry
                avg += func(sol[i,n])*wt*quad[i]
            elseif i >= size(sol,1) - num_bndry + 1
                avg += func(sol[i,n])*wt*quad[size(sol,1)-i+1]
            else
                avg += func(sol[i,n])*wt
            end
        end
    end
    avg /= (size(sol,1)-1)*(size(sol,2)-1)
    return avg
end

# function factorBandedSparse(A::SparseMatrixCSC{T,Int64}) where {T}

#     fac = one(T)
#     # loop over the elementary matrix operations
#     for k = 1:A.m-1
#         # update column k of the lower triangular matrix
#         # This assumes the rows are sorted, which Julia matrices should be
#         for ptr = A.colptr[k]:A.colptr[k+1]-1
#             if rowval[ptr] > k
#                 continue
#             elseif rowval[ptr] == k
#                 fac = one(T)/A.nzval[ptr]
#             else
#                 A.nzval[ptr] *= fac
#             end
#         end
#         # apply the elementary matrix to the lower block matrix
#         for j = k+1:A.n
#             if A.rowval[A.colptr[j]] > k
#                 break # there are no further nonzeros in row k beyond col j
#             end
#             for ptr = A.colptr[j]:A.colptr[j+1]-1
#                 i = A.rowval[ptr]
#                 if i <= k continue end
#                 A.nzval[ptr] -= A[i,k]*A[k,j] # this is not efficient
#             end
#         end
#     end
# end

end # module


module IRSpectrum

using FFTW

export autocorrelation,
       read_cp2k_dipoles,
       spectrum

"""
read_cp2k_dipoles(fname) -> Matrix{Float64}

Read dipole trajectory file created by CP2K and return dipolemoments as
a two dimensional array.
"""
function read_cp2k_dipoles(fname::AbstractString)
    out = Float64[]
    open(fname,"r") do file
        re = r"X=\ +(?<X>-?\d+.?\d+)\ +Y=\ +(?<Y>-?\d+.?\d+)\ +Z=\ +(?<Z>-?\d+.?\d+)"
        for line in eachline(file)
            m = match(re, line)
            if m !== nothing
                push!(out,parse(Float64,m[:X]))
                push!(out,parse(Float64,m[:Y]))
                push!(out,parse(Float64,m[:Z]))
            end
        end
    end
    if length(out)%3 == 0
        return reshape(out, (3, Int(length(out)/3)))
    else
        error("Error in reading file")
    end
end

"""
autocorrelation(μ::AbstractMatrix) -> Vector

Calculates autocorrelation function using [Wiener-Khinchin Theorem](http://mathworld.wolfram.com/Wiener-KhinchinTheorem.html).
"""
function autocorrelation(μ::AbstractMatrix)
    tmp = [ fft(abs2.(ifft(r))) for r in eachrow(μ)]
    return real.(reduce(+,tmp))
end

"""
spectrum(μ::AbstractMatrix; tstep=0.5, maxfreq=4000) -> Dict

Calculates spectrum from dipolemoment trajectory. Calculates autocorrelation function and
its Fourier transform.

# Arguments
- `μ::AbstractMatrix` :  data where autocorrelation funtion is calculated for each row

# Keywords
- `tstep` : timestep for trajectory in fempto seconds
- `maxfreq` : maximum frequency in wavenumbers for resulting spectrum

# Returns
- `Dict` : with fields `"absorption"` and `"wavenumber"`.
"""
function spectrum(μ::AbstractMatrix; tstep=0.5, maxfreq=4000)
    l = size(μ)[2]
    lh = Int(floor(l/2))
    acor = autocorrelation(μ)
    s = abs.(fft(acor))[1:lh]
    # freq = n/NΔ and wavenumber = freq/100c
    w = [ n/(l*tstep*1E-15*299792458E2) for n in 1:lh ]
    i = w .<= maxfreq
    return Dict("absorption"=>s[i], "wavenumber"=>w[i])
end

"""
spectrum(fname::AbstractString; tstep=0.5, maxfreq=4000) -> Dict

Read dipolemoments from CP2K dipole file `fname` and calculate spectrum.
"""
function spectrum(fname::AbstractString; tstep=0.5, maxfreq=4000)
    μ = read_cp2k_dipoles(fname)
    return spectrum(μ; tstep=tstep, maxfreq=maxfreq)
end

end # module


# single qubit Cliffords

export localclifford

const pX = [0 1; 1 0]
const pY = [0 -im; im 0]
const pZ = [1 0; 0 -1]

C1 = Vector{Clifford}(24)

# identity
C1[1]  = RI

# pi/2 and pi rotations about X,Y,Z
C1[2]  = expm(-im*pi/4*pX)
C1[3]  = im * pX
C1[4]  = expm(im*pi/4*pX)
C1[5]  = expm(-im*pi/4*pY)
C1[6]  = im * pY
C1[7]  = expm(im*pi/4*pY)
C1[8]  = expm(-im*pi/4*pZ)
C1[9]  = im * pZ
C1[10] = expm(im*pi/4*pZ)

# Hadamard class
C1[11] = expm(-im*pi/2/sqrt(2) * (pX+pY))
C1[12] = expm(-im*pi/2/sqrt(2) * (pX-pY))
C1[13] = expm(-im*pi/2/sqrt(2) * (pX+pZ)) # standard Hadamard
C1[14] = expm(-im*pi/2/sqrt(2) * (pX-pZ))
C1[15] = expm(-im*pi/2/sqrt(2) * (pY+pZ))
C1[16] = expm(-im*pi/2/sqrt(2) * (pY-pZ))

# Axis exchange class
C1[17] = expm(-1im*pi/3/sqrt(3) * (pX+pY+pZ))
C1[18] = expm(-2im*pi/3/sqrt(3) * (pX+pY+pZ))
C1[19] = expm(-1im*pi/3/sqrt(3) * (pX-pY+pZ))
C1[20] = expm(-2im*pi/3/sqrt(3) * (pX-pY+pZ))
C1[21] = expm(-1im*pi/3/sqrt(3) * (pX+pY-pZ))
C1[22] = expm(-2im*pi/3/sqrt(3) * (pX+pY-pZ))
C1[23] = expm(-1im*pi/3/sqrt(3) * (-pX+pY+pZ))
C1[24] = expm(-2im*pi/3/sqrt(3) * (-pX+pY+pZ))

rC1 = Dict{Clifford,UInt}()

for (k,v) in enumerate(C1)
    rC1[v] = k
end

localclifford(i::Int) = C1[i]
localclifford(v::Vector) = kron(C1[v]...)

localcliffordindex(c::Clifford) = rC1[c]


function K_eular_shear(beam::Beam)::Matrix{Float64}
    E,ν=beam.material.E,beam.material.ν
    A,I₂,I₃,J,l=beam.section.A,beam.section.I₂,beam.section.I₃,beam.section.J,beam.l
    As₂,As₃=beam.section.As₂,beam.section.As₃
    G=E/2/(1+ν)
    ϕ₂,ϕ₃=12E*I₃/(G*As₂*l^2),12E*I₂/(G*As₃*l^2)
    K=zeros(12,12)
    K[1,1]=E*A/l
    K[2,2]=12E*I₃/l^3/(1+ϕ₂)
    K[3,3]=12E*I₂/l^3/(1+ϕ₃)
    K[4,4]=G*J/l
    K[5,5]=(4+ϕ₃)*E*I₂/l/(1+ϕ₃)
    K[6,6]=(4+ϕ₂)*E*I₃/l/(1+ϕ₂)
    K[7,7]=E*A/l
    K[8,8]=12E*I₃/l^3/(1+ϕ₂)
    K[9,9]=12E*I₂/l^3/(1+ϕ₃)
    K[10,10]=G*J/l
    K[11,11]=(4+ϕ₃)*E*I₂/l/(1+ϕ₃)
    K[12,12]=(4+ϕ₂)*E*I₃/l/(1+ϕ₂)

    K[3,5]=K[5,3]=-6E*I₂/l^2/(1+ϕ₃)
    K[6,8]=K[8,6]=-6E*I₃/l^2/(1+ϕ₂)
    K[9,11]=K[11,9]=6E*I₂/l^2/(1+ϕ₃)

    K[2,6]=K[6,2]=6E*I₃/l^2/(1+ϕ₂)
    K[5,9]=K[9,5]=6E*I₂/l^2/(1+ϕ₃)
    K[8,12]=K[12,8]=-6E*I₃/l^2/(1+ϕ₂)

    K[7,1]=K[1,7]=-E*A/l
    K[8,2]=K[2,8]=-12E*I₃/l^3/(1+ϕ₂)
    K[9,3]=K[3,9]=-12E*I₂/l^3/(1+ϕ₃)
    K[10,4]=K[4,10]=-G*J/l
    K[11,5]=K[5,11]=(2-ϕ₃)*E*I₂/l/(1+ϕ₃)
    K[12,6]=K[6,12]=(2-ϕ₂)*E*I₃/l/(1+ϕ₂)

    K[3,11]=K[11,3]=-6E*I₂/l^2/(1+ϕ₃)

    K[2,12]=K[12,2]=6E*I₃/l^2/(1+ϕ₂)

    return K
end

function K2_eular_shear(beam::Beam)::Matrix{Float64}
    E,ν=beam.material.E,beam.material.ν
    A,I₂,I₃,J,l=beam.section.A,beam.section.I₂,beam.section.I₃,beam.section.J,beam.l
    As₂,As₃=beam.section.As₂,beam.section.As₃
    G=E/2/(1+ν)
    T=σ*A
    ϕ₂,ϕ₃=12E*I₃/(G*As₂*l^2),12E*I₂/(G*As₃*l^2)
    K=zeros(12,12)

    K[2,2]=(6/5+2ϕ₂+ϕ₂^2)/(1+ϕ₂)^2
    K[3,3]=(6/5+2ϕ₃+ϕ₃^2)/(1+ϕ₃)^2
    K[4,4]=J/A
    K[5,5]=(2*l^2/15+l^2*ϕ₃/6+l^2*ϕ₃^2/12)/(1+ϕ₃)^2
    K[6,6]=(2*l^2/15+l^2*ϕ₂/6+l^2*ϕ₂^2/12)/(1+ϕ₂)^2
    K[8,8]=(6/5+2ϕ₂+ϕ₂^2)/(1+ϕ₂)^2
    K[9,9]=(6/5+2ϕ₃+ϕ₃^2)/(1+ϕ₃)^2
    K[10,10]=J/A
    K[11,11]=(2*l^2/15+l^2*ϕ₃/6+l^2*ϕ₃^2/12)/(1+ϕ₃)^2
    K[12,12]=(2*l^2/15+l^2*ϕ₂/6+l^2*ϕ₂^2/12)/(1+ϕ₂)^2

    K[3,5]=K[5,3]=-(l/10)/(1+ϕ₃)^2
    K[6,8]=K[8,6]=-(l/10)/(1+ϕ₂)^2
    K[9,11]=K[11,9]=(l/10)/(1+ϕ₃)^2

    K[2,6]=K[6,2]=(l/10)/(1+ϕ₂)^2
    K[5,9]=K[9,5]=(l/10)/(1+ϕ₃)^2

    K[2,8]=K[8,2]=-(6/5+2ϕ₂+ϕ₂^2)/(1+ϕ₂)^2
    K[3,9]=K[9,3]=-(6/5+2ϕ₃+ϕ₃^2)/(1+ϕ₃)^2
    K[4,10]=K[10,4]=-J/A
    K[5,11]=K[11,5]=-(l^2/30+l^2*ϕ₃/6+l^2*ϕ₃^2/12)/(1+ϕ₃)^2
    K[6,12]=K[12,6]=-(l^2/30+l^2*ϕ₂/6+l^2*ϕ₂^2/12)/(1+ϕ₂)^2

    K[3,11]=K[11,3]=-(l/10)/(1+ϕ₃)^2

    K[2,12]=K[12,2]=(l/10)/(1+ϕ₂)^2

    K*=T/l
    return K
end


function giant_component_size(g, present)
    n = nv(g)
    Q = Queue{Int}()
	label = zeros(Int,n)
	cnt = zeros(Int,n)
	maxcnt = 0
    @inbounds for u in vertices(g)
		if present[u]
	        label[u] != 0 && continue
	        label[u] = u
			cnt[u] += 1
	        enqueue!(Q, u)
	        while !isempty(Q)
	            src = dequeue!(Q)
	            for vertex in all_neighbors(g, src)
					if present[vertex]
		                if label[vertex] == 0
		                    enqueue!(Q, vertex)
		                    label[vertex] = u
							cnt[u] += 1
		                end
					end
	            end
	        end
			if cnt[u] > maxcnt
				maxcnt = cnt[u]
			end
		end
    end
    return maxcnt
end

function recover_add_nodes(g, attack_nodes)
	n = nv(g)
	max_comp_sizes = Int[]
	ds = IntDisjointSets(n)
	present = fill(true, n)
	present[attack_nodes] .= false
	comp_sizes = ones(Int, n)
	max_comp_size = 1

	# add edges between present nodes and record component sizes
	for i in vertices(g)
		if present[i]
			for j in neighbors(g,i)
				if present[j]
					newroot = merge_nodes!(ds, comp_sizes, i, j)
					if comp_sizes[newroot] > max_comp_size
						max_comp_size = comp_sizes[newroot]
					end
				end
			end
		end
	end

	for i in reverse(attack_nodes)
		for j in neighbors(g, i)
			if present[j]
				newroot = merge_nodes!(ds, comp_sizes, i, j)
				if comp_sizes[newroot] > max_comp_size
					max_comp_size = comp_sizes[newroot]
				end
			end
		end
		present[i] = true
		push!(max_comp_sizes, max_comp_size)
	end
	reverse(max_comp_sizes)
end


#---------------------------------------------------------------------#
#This code computes the Chebyshev points & weights
#Written by F.X. Giraldo on 7/2014
#           Department of Applied Mathematics
#           Naval Postgraduate School
#           Monterey; CA 93943-5216
#---------------------------------------------------------------------#
function chebyshev_gauss(P::Integer)

    #Initialize arrays
    xgl=zeros(P)
    wgl=zeros(P)

    #Construct arrays
    for i=1:P
        xgl[P+1-i]=cos( (2*i-1)*pi/(2*P) );
        wgl[P+1-i]=pi/P
    end
    return (xgl,wgl)
end #function


using Plots

using ForwardDiff
eta_d = 1.0
eta_e = 0.5

pot(eta, e) = 1/eta * log(exp(eta*e) + exp(-eta*e))

e(y, y0) = abs(y0 - y)

yd = 0.5
ye = 1.0

potd(y) = pot(eta_d, e(y, yd))
pote(y) = pot(eta_e, e(y, ye))

pot2(y) = (pot(eta_d, e(y, yd)) + pot(eta_e, e(y, ye)))/2

# plot(potd, label="d")
# plot!(pote, label="e")
# plot!(pot2, label="sum")

nabla_pot2(y) = ForwardDiff.derivative(pot2, y)

y = -2.0:0.01:2.0
p2 = pot2.(y)
pd = potd.(y)
pe = pote.(y)

dp2 = [nabla_pot2(_y) for _y in y]

plot(y, p2, label="p2")
plot!(y, dp2,label="dp2")
plot!(y, pd, label="pd")
plot!(y, pe, label="pe")


abstract type AbstractSplineLayer <: Function end
Flux.trainable(m::AbstractSplineLayer) = (m.p,)

"""
Constructs a Spline Layer. At a high-level, it performs the following:
1. Takes as input a one-dimensional training dataset, a time span, a time step and
an interpolation method.
2. During training, adjusts the values of the function at multiples of the time-step
such that the curve interpolated through these points has minimum loss on the corresponding
one-dimensional dataset.

```julia
SplineLayer(time_span,time_step,spline_basis,saved_points=nothing)
```
Arguments:
- `time_span`: Tuple of real numbers corresponding to the time span.
- `time_step`: Real number corresponding to the time step.
- `spline_basis`: Interpolation method to be used yb the basis (current supported
  interpolation methods: ConstantInterpolation, LinearInterpolation, QuadraticInterpolation,
  QuadraticSpline, CubicSpline).
- 'saved_points': values of the function at multiples of the time step. Initialized by default
to a random vector sampled from the unit normal.
"""
struct SplineLayer{T<:Tuple{Real, Real},R<:Real,S1<:AbstractVector,S2<:UnionAll} <: AbstractSplineLayer
    time_span::T
    time_step::R
    saved_points::S1
    spline_basis::S2
    function SplineLayer(time_span,time_step,spline_basis,saved_points=nothing)
        saved_points = randn(length(time_span[1]:time_step:time_span[2]))
        new{typeof(time_span),typeof(time_step),typeof(saved_points),typeof(spline_basis)}(time_span,time_step,saved_points,spline_basis)
    end
end

function (layer::SplineLayer)(t::Real,p=layer.saved_points)
    return layer.spline_basis(p,layer.time_span[1]:layer.time_step:layer.time_span[2])(t)
end


# Tests the null of normality using the Jarque-Bera test statistic.
function jbtest(x::Vector)
    n = length(x)
    m1 = sum(x)/n
    m2 = sum((x - m1).^2)/n
    m3 = sum((x - m1).^3)/n
    m4 = sum((x - m1).^4)/n
    b1 = (m3/m2^(3/2))^2
    b2 = (m4/m2^2)
    statistic = n * b1/6 + n*(b2 - 3)^2/24
    d = Chisq(2.)
    pvalue = 1.-cdf(d,statistic)
    statistic, pvalue
end


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

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

Under the assumption that the image histogram is a mixture of two Gaussian
distributions the binarization threshold is chosen such that the expected
misclassification error rate is minimised.

# 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 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 minimum error thresholding method assumes that one can find a threshold
``T`` which partitions the data into two categories,  ``C_0`` and ``C_1``, such that
the data can be modelled by a mixture of two Gaussian distribution. 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.

Kittler and Illingworth proposed to use the minimum error criterion function
```math
J(T) = 1 + 2 \left[ P_0(T) \ln \sigma_0(T) + P_1(T) \ln \sigma_1(T) \right] - 2 \left[P_0(T) \ln P_0(T) + P_1(T) \ln P_1(T) \right]
```
to assess the discreprancy between the mixture of Gaussians implied by a particular threshold ``T``,
and the piecewise-constant probability density function represented by the histogram.
The discrete value ``T`` which minimizes the function ``J(T)`` produces
the sought-after threshold value (i.e. the bin which determines the threshold).

# 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, MinimumError())
```

# References

1. J. Kittler and J. Illingworth, “Minimum error thresholding,” Pattern Recognition, vol. 19, no. 1, pp. 41–47, Jan. 1986. [doi:10.1016/0031-3203(86)90030-0](https://doi.org/10.1016/0031-3203%2886%2990030-0)
2. Q.-Z. Ye and P.-E. Danielsson, “On minimum error thresholding and its implementations,” Pattern Recognition Letters, vol. 7, no. 4, pp. 201–206, Apr. 1988. [doi:10.1016/0167-8655(88)90103-1](https://doi.org/10.1016/0167-8655%2888%2990103-1)
"""
struct MinimumError <: AbstractImageBinarizationAlgorithm end

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

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


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ ad96c6b4-add2-11eb-1a72-8993033e1259
using Plots, LaTeXStrings, QuadGK, Roots, DifferentialEquations

# ╔═╡ 5168971b-2b60-494c-a8d8-cdfa8e1b8102
md"# 1D Bistable Example"

# ╔═╡ 261580d9-f1a5-442c-8e7f-3479f993f9b2
md"_Define separate production and degradation terms_"

# ╔═╡ 02d0fcf2-ca87-49ae-a8f3-f14310d0182a
begin
	n_max = 200.
	n_min = 1 # need to avoid 0 for log we'll evaluate later
	n = n_min:0.1:n_max
	γ = 1
	ks = 1e-4
	ν0 = 12.5
	ν1 = 200
	production(n) = (ν0+ ν1 * ks*n^2)/(1+ks*n^2)
	degradation(n) = γ*n
end

# ╔═╡ 13ddc727-e9bd-44f2-ba5e-c6c311b41723
begin

	plot(production, n, label=L"f(n)")
	p1 = plot!(degradation, n, label=L"g(n)", xlabel=L"n", title="Rate of Change", scale=:log10, xlimits=(15, n_max), ylimits=(15, 200), legend=:topleft)
# plot!(n->production(n)  - degradation(n), n, label=L"\frac{dn}{dt}",
# 		ylims=(-20, 150), xlims=(0, 150), legend=:topleft, title="Time evolution")
end

# ╔═╡ 868ba091-d0d2-42a2-bca7-f9e996b3870d
md"_Compute the fixed points to split up the 2 attractors into **Basins of Attractions**_"

# ╔═╡ 348718a0-8e19-4911-8998-5de9689bce1d
fixed_points = find_zeros(n-> production(n) - degradation(n), n_min, n_max)

# ╔═╡ 51d6af76-5ed3-41f8-9351-91e683c7bb57
md"_Integrate to get the quasi-potential_"

# ╔═╡ acb0adcf-fc19-45cf-bf0f-eb5532743b1a
begin
	f = production
	g = degradation

	ϕ(n) = quadgk(n -> -2*(f(n) -g(n))/(f(n) + g(n)), n_min, n)[1]
end

# ╔═╡ b4c511ea-fc96-475e-a5fc-1c9e282ca0e5
ϕ_vec = ϕ.(n)

# ╔═╡ 3373f26d-a6a4-4378-a8e1-c011c2e9f7a9
p2 = plot(n, ϕ_vec, ylabel=L"\phi(n)", xlabel=L"n", label="", title="Quasi-potential")

# ╔═╡ e147802a-2cc5-4610-821b-bf351ed62f9f
md"_Show the the probability due to the quasi-potential dominates compared to the other term in the equation_"

# ╔═╡ 178564ff-da53-4dd5-ba1f-b3e5db24b2ea
begin
	plot(n->log10.(f(n) + g(n)), n, label=L"f + g")
	plot!(n, log10.(ℯ.^(-ϕ_vec)), label=L"e^{-\phi}")
	p3 = plot!(xaxis=:log10, legend=:topleft)
end

# ╔═╡ 33cf16fa-679d-4447-aa59-b6d48e8fd15a
md"_Compute the normalized probability_"

# ╔═╡ a7e9be12-a1d4-4f1a-bbfc-d2642e5a9dd5
begin
	fn = f.(n)
	gn = g.(n)
	p_unnormal(n) =  ℯ^(-ϕ(n)) / (f(n)+g(n))
	norm = quadgk(p_unnormal, n_min, n_max)[1]
	p_normal(n) = 1/norm * p_unnormal(n)
	p4 = plot(p_normal,n,  title = "Normalized Probability", label="")
end

# ╔═╡ f5b8a081-222f-4dd6-ae4c-af7060aa4223
md"_Set up the jump process simulation_"

# ╔═╡ 467386e3-2ee8-46d4-9b40-831ae2778a57
random_count() = floor(Int, n_min + rand() * n_max)

# ╔═╡ 6ae4ed84-80c2-4c06-b69a-c81e6485b78b
begin
	affect1!(integrator) = integrator.u += 1
	affect2!(integrator) = integrator.u -= 1
	prodJump = ConstantRateJump((u, p, t) -> f(u), affect1!)
	degJump = ConstantRateJump((u, p, t) -> g(u), affect2!)
	u0 = random_count()
	t = (0.0, 5000.0)
	dprob = DiscreteProblem(u0, t)
	jprob = JumpProblem(dprob, Direct(), prodJump, degJump)
end

# ╔═╡ f32b2ab8-a941-463b-bb27-968c09e24777
begin
	sim_res = solve(jprob, SSAStepper())
	p5 = plot(sim_res, xlabel=L"t", label=L"n(t)", title="Simulation trace")
end

# ╔═╡ 50c911b1-57cf-4de4-9d65-054318c445fe
begin
	pf = plot(p1, p2, p4, p5, layout = grid(2,2), size=(750, 750))
	# savefig(pf, "ch3_bistable_prob.pdf")
	pf
end

# ╔═╡ 4547d20c-dab1-4df2-aa54-b8c52b1b5eeb
md"_Setup functions to do a longer time simulation and compare to integrated probability distribution_"

# ╔═╡ 456090be-e1c0-452e-b0ad-98fbc9f0576f
function calculate_time_fraction(sim_result, split)
	dts = diff(sim_result.t)
	dt_vals = collect(zip(dts, map(last, sim_result.u[1:end])))
	high = sum(map(first, filter(dtu -> last(dtu) > split, dt_vals)))
	low = sum(map(first, filter(dtu -> last(dtu) < split, dt_vals)))
	high, low
end

# ╔═╡ 23494ad1-82e3-4eee-86e3-695511de502e
function big_sim(num_sim = 25)
	function prob_func(prob, i, repeat)
        remake(prob, u0=random_count())
    end
	num_steps = 0
    ensemble_prob = EnsembleProblem(jprob, prob_func=prob_func)
	sim = solve(ensemble_prob, SSAStepper(), trajectories=num_sim)
	totals = [0., 0.]
	for j in 1:num_sim
		totals.+= calculate_time_fraction(sim[j], fixed_points[2])
		num_steps += length(sim[j])
	end
	totals, num_steps
end

# ╔═╡ 0e38a5c3-d970-4844-b560-fe4ed5452d85
totals, num_steps = big_sim()

# ╔═╡ 573a4914-b37e-4c2b-bd23-c10b58a6d4db
begin
	sim_frac = totals./sum(totals)
end

# ╔═╡ c62dd4a3-80ce-4b1f-b5a1-6142d30d0880
md"Time fractions from simulations is $(sim_frac)"

# ╔═╡ 86ee8326-51e1-453d-8d83-fd7021613c76
begin
	p_low = quadgk(p_normal, n_min, fixed_points[2])[1]
	p_high = quadgk(p_normal, fixed_points[2], n_max)[1]
	comp_frac = p_high, p_low
end

# ╔═╡ 5cabf565-b573-4446-b34d-8812986b25fe
md"Time fractions from integration of probability is $(comp_frac)"

# ╔═╡ Cell order:
# ╠═ad96c6b4-add2-11eb-1a72-8993033e1259
# ╟─5168971b-2b60-494c-a8d8-cdfa8e1b8102
# ╟─261580d9-f1a5-442c-8e7f-3479f993f9b2
# ╠═02d0fcf2-ca87-49ae-a8f3-f14310d0182a
# ╠═13ddc727-e9bd-44f2-ba5e-c6c311b41723
# ╟─868ba091-d0d2-42a2-bca7-f9e996b3870d
# ╠═348718a0-8e19-4911-8998-5de9689bce1d
# ╟─51d6af76-5ed3-41f8-9351-91e683c7bb57
# ╠═acb0adcf-fc19-45cf-bf0f-eb5532743b1a
# ╠═b4c511ea-fc96-475e-a5fc-1c9e282ca0e5
# ╠═3373f26d-a6a4-4378-a8e1-c011c2e9f7a9
# ╟─e147802a-2cc5-4610-821b-bf351ed62f9f
# ╠═178564ff-da53-4dd5-ba1f-b3e5db24b2ea
# ╟─33cf16fa-679d-4447-aa59-b6d48e8fd15a
# ╠═a7e9be12-a1d4-4f1a-bbfc-d2642e5a9dd5
# ╟─f5b8a081-222f-4dd6-ae4c-af7060aa4223
# ╠═467386e3-2ee8-46d4-9b40-831ae2778a57
# ╠═6ae4ed84-80c2-4c06-b69a-c81e6485b78b
# ╠═f32b2ab8-a941-463b-bb27-968c09e24777
# ╠═50c911b1-57cf-4de4-9d65-054318c445fe
# ╟─4547d20c-dab1-4df2-aa54-b8c52b1b5eeb
# ╠═456090be-e1c0-452e-b0ad-98fbc9f0576f
# ╠═23494ad1-82e3-4eee-86e3-695511de502e
# ╠═0e38a5c3-d970-4844-b560-fe4ed5452d85
# ╟─c62dd4a3-80ce-4b1f-b5a1-6142d30d0880
# ╟─5cabf565-b573-4446-b34d-8812986b25fe
# ╠═573a4914-b37e-4c2b-bd23-c10b58a6d4db
# ╠═86ee8326-51e1-453d-8d83-fd7021613c76


"""
The bistable model of mutual antagonism. To perform the bistability, the parameters are defined symmetrically, contributing to a stable system with steady-states determined by the initial state.

In this section, we are going to examine wheter this package can successfully find two steady states by grid searching and predefined ranges.

Reference
---------
1. p.82-p.83, Figure 4.8. Ingalls, B. P. (2013). Mathematical modeling in systems biology: an introduction. MIT press.
"""
function bistable_ode!(du, u, p ,t)
	s1, s2 = u
	K1, K2, k1, k2, k3, k4, n1 , n2  = p

	du[1] = k1 / (1 + (s2/K2)^n1) - k3*s1
	du[2] = k2/  (1 + (s1/K1)^n2) - k4*s2
end

p_ = LVector(K1=1., K2=1., k1=20., k2=20., k3=5., k4=5., n1= 4., n2=4.)

u_1 = LVector(s1=3., s2=1.)
u_2 = LVector(s1=1., s2=3.)


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


# # 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


###############################################
#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);


function findPi(num_samples, trials_per_sample)
  samples = Array{Float64}(undef, num_samples)
  mean = 0.0

  for sample = 1:num_samples
    samples[sample] = integrateCircle(trials_per_sample)
  end

  mean = sum(samples) / size(samples)[1]

  std_dev = getStdDev(mean, samples)

  return mean, std_dev
end

function getStdDev(mean, samples)
  variance = 0

  for sample in samples
    variance += (sample - mean) ^ 2
  end

  variance /= size(samples)[1]
  return sqrt(variance)
end

function integrateCircle(points_placed)
  counter = 0

  for trial = 1:points_placed
    y_pos = rand()
    x_pos = rand()

    dist = y_pos ^ 2 + x_pos ^ 2
    if dist <= 1
      counter += 1
    end
  end

  return 4.0 * counter / points_placed
end

function monte_carlo()
  res = [0.0, 0.0]
  for i in 1:20
    res1, res2 = findPi(1000, 50000)
    res[1] += res1 / 20.0
    res[2] += res2 / 20.0
  end
  print(res)
end


function ffd(f::Function, x=0.0, h=1e-8)
    val = (f(x+h)-f(x))/h
    return val
end

function bfd(f::Function, x=0.0, h=1e-8)
    val = (f(x)-f(x-h))/h
    return val
end

function cfd(f::Function, x=0.0, h=1e-5)
    val = (f(x+h)-f(x-h))/2/h
    return val
end


### wiener.jl ---  Wiener deconvolution
#
# Copyright (C) 2016 Mosè Giordano.
#
# Maintainer: Mosè Giordano <mose AT gnu DOT org>
# Keywords: wiener, deconvolution, signal processing
#
# This file is a part of Deconvolution.jl.
#
# License is MIT "Expat".
#
### Commentary:
#
# This file provides the `wiener' function to perform Wiener deconvolution.  It
# has two modes: one for the case of presence of a known blurring, the other one
# without blurring.
#
### Code:

export wiener

### The implementation of the deconvolution algorithm

function _wiener_no_blur(Y::AbstractArray, # Fourier transform of the input
                         S::AbstractArray, # Power spectrum of the signal
                         N::AbstractArray) # Power spectrum of the noise
    # Without blurring, the filter is:
    #   |S|² / (|S|² + |N|²)
    return real(ifft(Y .* S ./ (S .+ N)))
end

function _wiener_with_blur(Y::AbstractArray, # Fourier transform of the input
                           S::AbstractArray, # Power spectrum of the signal
                           N::AbstractArray, # Power spectrum of the noise
                           H::AbstractArray) # Fourier transform of the blurring
    # With blurring, the filter is:
    #   H* / (|H|² + |N/S|²)
    return real(ifft(Y .* conj(H) ./ (abs2.(H) .+ N ./ S)))
end

### User interface

## Without blurring
function wiener(input::AbstractArray, signal::AbstractArray,
                noise::AbstractArray)
    @assert size(input) == size(signal) == size(noise)
    input_ft = fft(input)
    signal_power_spectrum = abs2.(fft(signal))
    noise_power_spectrum = abs2.(fft(noise))
    return _wiener_no_blur(input_ft, signal_power_spectrum,
                           noise_power_spectrum)
end

# Noise as a real (it's converted to an array)
wiener(input::AbstractArray, signal::AbstractArray, noise::Real) =
    wiener(input, signal, fill!(similar(input), one(eltype(input)))*noise)

## With blurring
function wiener(input::AbstractArray, signal::AbstractArray,
                noise::AbstractArray, blurring::AbstractArray)
    @assert size(input) == size(signal) == size(noise) == size(blurring)
    input_ft = fft(input)
    signal_power_spectrum = abs2.(fft(signal))
    noise_power_spectrum = abs2.(fft(noise))
    blurring_ft = fft(blurring)
    return _wiener_with_blur(input_ft, signal_power_spectrum,
                             noise_power_spectrum, blurring_ft)
end

# Noise as a real (it's converted to an array)
wiener(input::AbstractArray, signal::AbstractArray,
       noise::Real, blurring::AbstractArray) =
           wiener(input, signal, fill!(similar(input), one(eltype(input)))*noise, blurring)

"""
    wiener(input, signal, noise[, blurring])

Return the Wiener deconvolution of `input`, using the power spectrum of `signal`
and `noise`.  If the `input` was blurred with a known blurring kernel, pass it
as fourth argument, `blurring`.

All arguments must be arrays in the time/space domain and all of the same size,
they will be converted into the frequency domain internally using the `fft`
function.
"""
wiener