using LinearAlgebra
using StaticArrays

include("./io.jl")

"""
    repulsion(npoints, dims, fileopts...; niter = 1e7, convergence_tol = 1e-10)

Generates a uniform grid of `npoints` points in `dims` dimensions, distributed uniformly on the surface of the unit sphere
using the REPULSION algorithm

The function iterates, up to a maximum of `niter` times, until the algorithm converges. The convergence criterion
`convergence_tol` is interpreted as the target for the maximum displacement between points in two successive iterations. The points
inherit their floating point precision from `convergence_tol`.
The points are returned as a vector of vectors of length `npoints`, with each element containing the angles `[α, β, γ]` and the
weight corresponding to each point (currently uniform weights are assigned)

The function optionally writes the resulting grid of points to a file  by calling `writedata`, which uses the contents of the optional
`fileopts` parameters to determine the file format and filename. The first fileopt parameter is a String containing the file format,
and the second is a String containing the filename. If the filename is not provided, a default is chosen using the current system time.
The file formats currently supported are: "CSV", "TSV", "MAT", "JLD", "JLD2". The [MAT](https://github.com/JuliaIO/MAT.jl) and
[JLD](https://github.com/JuliaIO/JLD.jl)/[JLD2](https://github.com/JuliaIO/JLD2.jl) file types require packages not in the standard library
that must be installed and loaded by the user.

This function is adapted from its counterpart in Spinach (https://spindynamics.org) for MATLAB. The Spinach
function, repulsion.m, is Copyright Ilya Kuprov and Frederic Mentink-Vigier
"""
function repulsion(npoints, dims, fileopts...; niter = 1e7, convergence_tol::T = 1e-10) where {T <: AbstractFloat}
    R_raw = rand(T, dims, npoints) .- one(T)/2 # generating the first set of random points
    R = [SVector{dims, T}(ntuple(i -> R_raw[n+i], Val(dims))) for n in 0:dims:(dims*npoints-1)]
    R_n = copy(R) # pre-allocate
    F = zeros(SVector{dims, T},npoints) # pre-allocate "forces"

    for i in 1:niter # starting the conjugate gradient optimization
        F .= zero(T) .* F # re-initializing F at each step
        for k in eachindex(R)
            F .-= dist_vector.(R,[R[k]]) # compute the distance between points, in place (speed-up)
        end
        R_n = normalize.(R_n - dims*F/npoints) # normalize
        max_diff = maximum(norm.(R-R_n)) # get the maximum distance
        R = R_n
	    if mod(i, 1000) == 0
	        @debug "Maximum difference " max_diff "At iteration " i # returns the convergence value every 1000 steps
	    end
        if max_diff <= convergence_tol
            @info "Converged at iteration " i # stops and return number of steps when converged
            break
        end
    end
    angles = @. vcat(toangles(R), one(T)/npoints) # computing the angles and weights

    if !isempty(fileopts)
        writedata(angles, fileopts...) # Optionally, write angles to file
    end

    return angles
end

"""
    dist_vector(a::SVector{N,T}, b::SVector{N,T}) where N where T

Computes the normalised distance between two vectors `a` and `b`, scaled by the cosine of the angle between them
"""
function dist_vector(a::SVector{N,T}, b::SVector{N,T}) where N where T
    dd = normalize(a - b)
    if isnan(first(dd))
        zeros(typeof(a))
    else
        dd*dot(a, b)
    end
end

"""
    toangles(p)

Converts a point `p` on the unit sphere to a vector of angles `[α, β, γ]`

The interpretation of `p` and the conversion depend on the number of dimensions:
- Two dimensions: `p` is treated as 2D cartesian co-ordinates, and the polar angle is returned as `β`
- Three dimensions: `p` is treated as 3D cartesian co-ordinates, and the spherical angles `θ` and `φ` are returned as `β` and `γ`
- Four dimensions: `p` is treated as a quaternion, and the ZYZ Euler angles are returned
All angles are shifted to lie in the ranges ``α ∈ [0,2π)``, ``β ∈ [0,π]``, ``γ ∈ [0,2π)``. Missing angles are returned as 0s
"""
function toangles(p::SVector{N,T}) where N where T
    angles = zeros(T, 3)
    if N == 2
	    angles[2] = atan(p[2], p[1])
    elseif N == 3
	    angles[3] = atan(p[2], p[1])
	    angles[2] = pi/(2*one(T)) + atan(p[3], hypot(p[1], p[2]))
    elseif N == 4
	    angles = quat_to_euler(p)
    end
    angles
end

"""
    quat_to_euler(q)

Converts a normalised quaternion to its corresponding Euler angles `[α, β, γ]` about the
ZYZ rotation axes.

An active rotation convention is assumed - invert the input quaternion
for passive rotations. This code is adapted from EasySpin's quat2euler implementation at:
https://github.com/StollLab/EasySpin/blob/main/easyspin/quat2euler.m
EasySpin is Copyright Stefan Stoll and other contributors
"""
function quat_to_euler(q::AbstractVector{T}) where T
    if length(q) != 4
        e = DomainError(q, "A Quaternion must have four elements")
        throw(e)
    end

    sy = 2*sqrt((q[3]*q[4] + q[2]*q[1])^2 + (q[2]*q[4] - q[3]*q[1])^2)
    beta = real(atan(sy, one(T) - 2*q[2]^2 - 2*q[3]^2))
    if sy < 1e-10 # Singularity check: 1e-10 ≈ 0 for normalised quantities
        alpha = zero(T)
        gamma = atan(-2*(q[4]*q[1] - q[2]*q[3]), one(T) - 2*(q[2]^2 - q[4]^2))
    else
        alpha = atan(2*(q[3]*q[4] - q[1]*q[2]), 2*(q[2]*q[4] + q[1]*q[3]))
        gamma = atan(2*(q[3]*q[4] + q[1]*q[2]), -2*(q[2]*q[4] - q[1]*q[3]))
        if alpha < 0
            alpha = alpha + pi + pi # Maintaining floating point precision
        end
        if gamma < 0
            gamma = gamma + pi + pi
        end
    end
    [alpha; beta; gamma]
end


###### beam splitter and special counter example matrices ######

function beam_splitter(transmission_amplitude = sqrt(0.5))

	"""2d beam_splitter matrix, follows the conventions of Leonardo"""

	# |t|^2 is the "transmission probability"

	t = transmission_amplitude
	r = sqrt(1-t^2)

	bs = [[t -r]; [r t]]

end

function beam_splitter_modes(;in_up, in_down, out_up, out_down, transmission_amplitude, n)

	"""beam splitter incorrporated for connecting specific input modes, output modes in a dimension n interferometer"""

	if in_up == in_down || out_up == out_down
		throw(ArgumentError())
	end

	bs = Matrix{ComplexF64}(I, n, n)

	sub_bs = beam_splitter(transmission_amplitude)

	bs[in_up, out_up] = sub_bs[1,1]
	bs[in_down, out_down] = sub_bs[2,2]
	bs[in_down, out_up] = sub_bs[2,1]
	bs[in_up, out_down] = sub_bs[1,2]
	#
	# # before convention checks:
	# bs[in_down, out_up] = sub_bs[1,2] ####### this change is a bit ad hoc and needs to be checked carefully with the conventions
	# bs[in_up, out_down] = sub_bs[2,1]
	#

	bs

end

function rotation_matrix(angle::Float64)
	[cos(angle) -sin(angle);
	 sin(angle) coss(angle)]
 end

 function rotation_matrix_modes(;in_up, in_dow, out_up, out_down, angle, n)

	 """beam splitter incorrporated for connecting specific input modes, output modes in a dimension n interferometer"""

 	if in_up == in_down || out_up == out_down
 		throw(ArgumentError())
 	end

 	r = Matrix{ComplexF64}(I, n, n)

 	sub_r = rotation_matrix(angle)

 	r[in_up, out_up] = sub_r[1,1]
 	r[in_down, out_down] = sub_r[2,2]
 	# bs[in_down, out_up] = sub_bs[2,1]
 	# bs[in_up, out_down] = sub_bs[1,2]
 	r[in_down, out_up] = sub_r[1,2] ####### this change is a bit ad hoc and needs to be checked carefully with the conventions
 	r[in_up, out_down] = sub_r[2,1]

 	r

 end

 function phase_shift(shifted_modes::Array, param_::Array)

	U = zeros(length(shifted_modes), length(shifted_modes))

 	for i in length(shifted_modes)
		for j in length(shifted_modes)
			i == j && i in shifted_modes ? U[i,j] = exp(1im * param_[i]) : continue
		end
	end

	  U

end

# function build_circuit(mode_occ::ModeOccupation, circuit_elements::Vector{Vector{T}}, layers::Vector) where T <: Interferometer
#
# 	m = mode_occ.m
# 	id = Matrix{ComplexF64}(LinearAlgebra.I, m, m)
# 	res = id
#
# 	for i in 1:length(circuit_elements)
# 		u = id
# 		for j in 1:length(circuit_elements[i])
# 			U = circuit_elements[i][j].U
# 			u[layers[i][j], layers[i][j]] = U
# 		end
#
# 		res *= u
# 	end
#
# 	res
# end
#
# mode_occ = first_modes(3,5)
#
# m = mode_occ.m
#
# b1 = BeamSplitter(1/sqrt(6))
#
# is_unitary(b1.U)
#
# U1 = RandHaar(1)
#
# is_unitary(U1.U)
#
# b2 = BeamSplitter(1/sqrt(2))
#
# is_unitary(b2.U)
#
# U2 = RandHaar(1)
#
# is_unitary(U2.U)
#
# circuit_elem = [[b1,U1],[b2,U2]]
#
# v_ = [[[1,2],[3]], [[1,3],[2]]]
#
# res = build_circuit(mode_occ, circuit_elem, v_)
#
# is_unitary(res)
#
# res * conj(transpose(res))
#


struct Spherical <: AbstractPredictorCorrectorAlgorithm{true} end

"""

    predict(alg::Affine, H, J_H, ∂H∂t, x, t, Δt)

Spherical tangent predictor following the formulation of Chen[^1]. Denote by ``\hat{H}`` the lift of the homotopy ``H``
as a projective mapping.
Euler's method yields
```math
\begin{bmatrix}
D_x\hat{H}^(x,t) \\
x^{H}
\end{bmatrix}
\dot{x} =
\begin{bmatrix}
-D_t\hat{H}^(x,t) \\
0
\end{bmatrix}
```
and we can then obtain the spherical projection:
```math
\cos(\norm{\dot{x}}_2 Δt) x + \sin(\norm{\dot{x}}_2 Δt) \frac{\dot{x}}{\norm{\dot{x}}_2}
```

[^1]: Tianran Chen and Tien-Yien Li.
    “Spherical projective path tracking for homotopy continuation methods”.
    Communications in Information and Systems 12(3):195-220 (2012)
"""
function predict(alg::Spherical, H::AbstractHomotopy, J_H, ∂H∂t, x, t, Δt)
    dot_x = \([J_H(x,t); x'], [-∂H∂t(x,t); 0])
    norm_dot_x = norm(dot_x)
    # TODO: conditining? i.e. check singular values
    cos(-norm_dot_x*Δt) .* x .+ (sin(-norm_dot_x*Δt) / norm_dot_x) .* dot_x
end

function correct!(
    u::Vector{T},
    alg::Spherical,
    H::AbstractHomotopy,
    J_H::F,
    x::Vector{T},
    t,
    tol::Float64,
    max_iterations::Int
) where {T,F}
    N = length(x)
    res = zeros(x)
    dx_xt = zeros(T, N, N)
    Δx = zeros(x)
    u .= x

    for k = 1:max_iterations
        res[1:N-1] = -evaluate(H, u, t)
        # println("newton iteration: $(k), res: $(norm(res))")
        if norm(res) < tol
            return true
        end

        dx_xt[1:(N-1), :] .= J_H(u, t)
        dx_xt[N, :] = u'

        Δx .= \(dx_xt, res)
        norm_Δx = norm(Δx)
        u .= cos(norm_Δx) .* u .+ (sin(norm_Δx) / norm_Δx) .* Δx
    end

    norm(res) < tol
end


#=
void acb_elliptic_p(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes Weierstrass’s elliptic function

void acb_elliptic_inv_p(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes the inverse of the Weierstrass elliptic function

void acb_elliptic_zeta(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes the Weierstrass zeta function

void acb_elliptic_sigma(acb_t res, const acb_t z, const acb_t tau, slong prec)¶
Computes the Weierstrass sigma function
=#

"""
    weierstrass_p(z, tau)

- weierstrass_p(z, tau) == weierstrass_p(z+1, tau) == weierstrass_p(z+tau, tau)
"""
function weierstrass_p(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_p), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

"""
    weierstrass_invp(z, tau)

- weierstrass_p(weierstrass_invp(z, tau), tau) == z
"""
function weierstrass_invp(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_inv_p), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

"""
    weierstrass_zeta(z, tau)

- weierstrass_zeta(z+1, tau) == weierstrass_zeta(z, tau) + weierstrass_zeta(1/2, tau)
- weierstrass_zeta(z+tau, tau) == weierstrass_zeta(z, tau) + weierstrass_zeta(tau/2, tau)
"""
function weierstrass_zeta(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_zeta), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

"""
    weierstrass_sigma(z, tau)

- weierstrass_sigma(z+1, tau) == −exp(2*(z+τ/2)*weierstrass_zeta(τ/2,τ)) * weierstrass_sigma(z, tau)
"""
function weierstrass_sigma(z::ArbComplex{P}, tau::ArbComplex{P}) where {P}
    result = ArbComplex{P}()
    flag = 0
    ccall(@libarb(acb_elliptic_sigma), Cvoid, (Ref{ArbComplex}, Ref{ArbComplex}, Ref{ArbComplex}, Cint),
          result, z, tau, P)
    return result
end

weierstrass_p(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_p(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_p(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_p(ArbComplex{P}(z), ArbComplex{P}(tau))

weierstrass_pinv(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_pinv(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_pinv(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_pinv(ArbComplex{P}(z), ArbComplex{P}(tau))

weierstrass_zeta(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_zeta(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_zeta(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_zeta(ArbComplex{P}(z), ArbComplex{P}(tau))

weierstrass_sigma(z::ArbReal{P}, tau::ArbReal{P}) where {P} =
    weierstrass_sigma(ArbComplex{P}(z), ArbComplex{P}(tau))
weierstrass_sigma(z::ArbFloat{P}, tau::ArbFloat{P}) where {P} =
    weierstrass_sigma(ArbComplex{P}(z), ArbComplex{P}(tau))


using JuPOT

             m = 6 # No. of days data
             returns_data = ARGS[1]
			 returns = parse.(split(returns_data, ","))

             meanReturns_data =  ARGS[2] # Returns a matrix of size(n) with entries between 0-1s
			 meanReturns = parse.(split(meanReturns_data, ","))

             tickers_data = ARGS[3]
			 tickers = split(tickers_data, ",")

             n = length(tickers)

             matrix = reshape(returns,m,n)
             covariance = cov(matrix) #covariance matrix of returns

             # Assets data structure containing, names, expected returns, covarariance
             assets = AssetsCollection(tickers, meanReturns, covariance)

             target_ret_test = [0.10,0.15,0.20,0.25,0.30,0.35,0.40,0.45,0.50,0.55,0.60,0.65,0.70,0.75,0.80]

             function efficientFrontierMVO(assets, target_ret_test)

                 # Create the Arrays to be returned, and set them to be the length of the
                 # target_ret_test from the input.
                 n = length(target_ret_test)
                 risk = Array(Float64,n)
                 index = Array(Float64,n)

                 for i in 1:n
                     # Get the target return we want to test.
                     target_ret = target_ret_test[i]

                     # Setup and Optimize the SimpleMVO with the new Target Return.
                     mvo = SimpleMVO(assets, target_ret; short_sale=false)
                     result = optimize(mvo)

                     # Risk is the Objective Value. Index the Target Return.
                     risk[i] = result[1]
                     index[i] = target_ret
                 end

                 # Return as a Tuple of the Arrays
                 return risk, index
             end

             efd = efficientFrontierMVO(assets, target_ret_test)


export HPoint, getz, dist, midpoint, RandomHPoint, between, polar

"""
`HPoint(z::Complex)` creates a new point in the hyperbolic plane.
The argument `z` must have absolute value less than 1.

`HPoint(r,theta)` creates a new point with polar coordinates
`(r,theta)`. See also: `polar`.
"""
struct HPoint <: HObject
    z::Complex{Float64}
    attr::Dict{Symbol,Any}
    function HPoint(z::Complex)
        if _mag(z) >=  1
            throw(DomainError(z, "absolute value is too large"))
        end
        P = new(z,Dict{Symbol,Any}())
        set_color(P)
        set_radius(P)
        return P
    end
end

HPoint(P::HPoint) = HPoint(getz(P))  # copy constructor
HPoint(z::Number) = HPoint(Complex(z))
HPoint() = HPoint(0)

function show(io::IO,P::HPoint)
    r,theta = polar(P)
    print(io,"HPoint($r, $theta)")
end

"""
`RandomHPoint()` generates a point at random in the hyperbolic plane.
"""
RandomHPoint() = HPoint(randn(), 2*pi*rand())

"""
`getz(P::HPoint)` returns the point (complex number) in the
interior of the unit disc that represents `P`.
"""
getz(P::HPoint) = P.z

"""
`dist(P,Q)` gives the distance betwen two points in the
hyperbolic plane. If `Q` is omitted, give the distance
from `P` to `HPoint(0)`.
"""
function dist(P::HPoint, Q::HPoint)
    a = getz(P)
    b = getz(Q)
    delta = 2 * _mag(a-b)/(1-_mag(a))/(1-_mag(b))
    return acosh(1+delta)
end

dist(P::HPoint) = dist(P, HPoint(0))

HPoint(r::Real, theta::Real) = HPoint( solve_dist(r) * exp(im*theta) )

"""
`polar(P::HPoint)` gives the polar coordinates of `P`
"""
function polar(P::HPoint)
    r = dist(P)
    theta = angle(getz(P))
    return (r,theta)
end

(==)(P::HPoint,Q::HPoint) = _mag(getz(P)-getz(Q)) <= THRESHOLD*eps(1.0)

function _dist(t::Real)
    delta = 2*t*t/(1-t*t)
    return acosh(1+delta)
end

"""
`solve_dist(d)` is the inverse of `_dist()`
"""
function solve_dist(d::Real)
    ex = exp(d)
    return (ex-1)/(ex+1)
end

"""
`midpoint(p,q)` finds the mid point of the line segment
from `p` to `q`. Also `midpoint(L::HSegment)`.
"""
function midpoint(p::HPoint, q::HPoint)::HPoint
    if p==q
        return p
    end
    d = dist(p,q)
    f = move2xplus(p,q)

    t = solve_dist(d/2)

    r = HPoint(t)
    h = inv(f)
    return h(r)
end

"""
`between(a,b,c)` determines if the hyperbolic point `b`
lies on the segment from `a` to `c`.
"""
function between(a::HPoint, b::HPoint, c::HPoint)::Bool
    if a==b || b==c
        return true
    end
    f = move2xplus(a,c)
    aa = getz(f(a))
    bb = getz(f(b))
    cc = getz(f(c))

    if abs(imag(bb)) > THRESHOLD * eps(1.0)
        return false
    end

    if real(aa)-THRESHOLD*eps(1.0) <= real(bb) <= real(cc)+THRESHOLD*eps(1.0)
        return true
    end
    return false
end


# Conjugate

"""
  Conjugate(f::ProximableFunction)

Returns the conjugate function of `f`, that is `f*(x) = sup{y'x - f(y)}`.
"""

immutable Conjugate{T <: ProximableFunction} <: ProximableFunction
  f::T
end

fun_dom(f::Conjugate) = fun_dom(f.f)

# only prox! is provided here, call method would require being able to compute
# an element of the subdifferential of the conjugate

function prox!{R <: Real}(g::Conjugate, x::AbstractArray{R}, y::AbstractArray{R}, gamma::Real=1.0)
  # need to make a copy, doing this in place is probably not possible in general:
  # if object_id(x) == object_id(y), then prox! causes the loss of x (which is needed afterwards)
  x_copy = copy(x)
  # Moreau identity
  v = prox!(g.f, x/gamma, y, 1.0/gamma)
  v = vecdot(x_copy,y) - gamma*vecdot(y,y) - v
  for k in eachindex(y)
    y[k] = x_copy[k] - gamma*y[k]
  end
  return v
end

# complex case, need to cast inner products to real

function prox!{R <: Real}(g::Conjugate, x::AbstractArray{Complex{R}}, y::AbstractArray{Complex{R}}, gamma::Real=1.0)
  x_copy = copy(x)
  v = prox!(g.f, x/gamma, y, 1.0/gamma)
  v = real(vecdot(x_copy,y)) - gamma*real(vecdot(y,y)) - v
  for k in eachindex(y)
    y[k] = x_copy[k] - gamma*y[k]
  end
  return v
end

# special case for indicator of convex cones (the conjugate is also an indicator function)

prox!{R <: Real, C <: IndicatorConvexCone}(g::Conjugate{C}, x::AbstractArray{R}, y::AbstractArray{R}, gamma::Real=1.0) =
  prox_conjugate_convex_cone!(g.f, x, y, gamma)

prox!{R <: Real, C <: IndicatorConvexCone}(g::Conjugate{C}, x::AbstractArray{Complex{R}}, y::AbstractArray{Complex{R}}, gamma::Real=1.0) =
  prox_conjugate_convex_cone!(g.f, x, y, gamma)

function prox_conjugate_convex_cone!{T <: RealOrComplex, C <: IndicatorConvexCone}(f::C, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  x_copy = copy(x)
  prox!(f, x/gamma, y, 1.0/gamma)
  for k in eachindex(y)
    y[k] = x_copy[k] - gamma*y[k]
  end
  return 0.0
end

# naive implementation

function prox_naive{T <: RealOrComplex}(g::Conjugate, x::AbstractArray{T}, gamma::Real=1.0)
  y, v = prox_naive(g.f, x/gamma, 1.0/gamma)
  return x - gamma*y, real(vecdot(x,y)) - gamma*real(vecdot(y,y)) - v
end

is_prox_accurate(f::Conjugate) = is_prox_accurate(f.f)


function radixsort(tobesorted::Vector{Int64})
    arr = deepcopy(tobesorted)
    for shift in 63:-1:0
        tmp = Vector{Int64}(undef, length(arr))
        j = 0
        for i in 1:length(arr)
            if (shift == 0) == ((arr[i] << shift) >= 0)
                arr[i - j] = arr[i]
            else
                tmp[j + 1] = arr[i]
                j += 1
            end
        end
        tmp[j+1:end] .= arr[1:length(tmp)-j]
        arr = tmp
    end
    arr
end

function testradixsort()
    arrays = [[170, 45, 75, -90, -802, 24, 2, 66], [-4, 5, -26, 58, -990, 331, 331, 990, -1837, 2028]]
    for array in arrays
        println(radixsort(array))
    end
end

testradixsort()


#=
# 3D Poisson, Dirichlet bc
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("poisson3d");
useLog("poisson3dlog")

n = 10;
ord = 2;

domain(3)
functionSpace(order=ord)

mesh(HEXMESH, elsperdim=n)

u = variable("u")

testSymbol("v")

boundary(u, 1, DIRICHLET, 0)

# Write the weak form
coefficient("f", "-3*pi*pi*sin(pi*x)*sin(pi*y)*sin(pi*z)")
weakForm(u, "-dot(grad(u), grad(v)) - f*v")

solve(u);

outputValues(u, "p3dout", format="vtk", asci=false);

# exact solution is sin(pi*x)*sin(pi*y)*sin(pi*z)
# check error
maxerr = 0;
exact(x,y,z) = sin(pi*x)*sin(pi*y)*sin(pi*z);

for i=1:size(Finch.grid_data.allnodes,2)
    x = Finch.grid_data.allnodes[1,i];
    y = Finch.grid_data.allnodes[2,i];
    z = Finch.grid_data.allnodes[3,i];
    err = abs(u.values[i] - exact(x,y,z));
    global maxerr;
    maxerr = max(err,maxerr);
end
println("max error = "*string(maxerr));

finalize_finch()


using Random, LinearAlgebra, StatsBase, Distributions

##FUNCTION data_construction
"""Generate synthetic data
INPUTS
- ℓ             Loss function
- n             Number of samples
- k             Number of relevant features
- p             Number of features/covariates
- SNR           Squareroot of the Signal-to-Noise Ratio
- rho           Correlation coefficient between features
OUTPUTS
- X             Matrix of observed covariates
- Y             Vector of responses
- indices_true  Set of relevant features
- w_true        True regressor"""
function data_construction(ℓ::LossFunction, n::Int, k::Int, p::Int, SNR, ρ)
  indices_true, w_true = generate_estimator(p, k)
  #Auto-regresive AR(1) filter preserving noise power
  X = randn(n, p)
  for i in 2:p
      X[:, i] = X[:, i] + ρ*X[:, i-1]
  end
  # rho = ρ
  # R"""
  # X=matrix(NA,$n,$p)
  # for(i in 1:$n){
  #     X[i,]=arima.sim(n=$p,list(ar=$rho))
  # }
  # """
  # @rget X
  # X = √12 * rand(n, p) - (√12)/2

  Y = generate_output(ℓ, noisy_signal(X, indices_true, w_true, SNR))

  return X, Y, indices_true, w_true
end

function hard_data_construction(ℓ::LossFunction, n::Int, k::Int, p::Int, SNR, ρ)
  indices_true, w_true = generate_estimator(p, k)

  indices_true = collect(1:k)
  w_true = ones(k)./√k #sample(-1:2:1, k)

  #Covariance matrix
  θ = (1/√k + 1/k)/2
  Σ = SparseMatrixCSC{Float64}(I, p, p)
  for i in 1:k
    Σ[i,k+1] = θ
    Σ[k+1,i] = θ
  end

  # perm = shuffle(1:p)
  # indices_true = [findfirst( perm .== i) for i in indices_true]
  # Σ[:,:] = Σ[perm, perm]

  d = MvNormal(full(Σ))

  X = rand(d, n)'

  Y = generate_output(ℓ, noisy_signal(X, indices_true, w_true, SNR))

  # E = randn(n, p)  #Noise
  # for j in 1:p
  #   E[:,j] *= norm(X[:,j])/norm(E[:,j])/SNR*2
  # end

  return X, Y, indices_true, w_true
end

# FUNCTIONS Generate support and regressor w
function generate_estimator(p::Int, k::Int)
  weights = ones(p) ./ p
  indices = sort(sample(1:p, StatsBase.Weights(weights), k, replace=false))
  w = sample(-1:2:1, k)
  return indices, w
end

function noisy_signal(X, indices, w, SNR)
  n = size(X,1)
  S = X[:, indices]*w #Signal
  E = randn(n)  #Noise
  E *= (norm(S)/norm(E))/SNR #Scaling according to SNR
  return S+E
end

#FUNCTIONS Generate Output Y
function generate_output(ℓ::Regression, S)
  return S
end

function generate_output(ℓ::L1SVM, S)
  P = 1. ./( 1 .+ exp.(-S) )
  return 1.0*(P .> .5) - 1.0*(P .< .5)
end

function generate_output(ℓ::LogReg, S)
    n = size(S,1)
    P = 1. ./( 1 .+ exp.(-S) )
    return (-ones(n)).^(rand(n) .> P)
end


using DynamicEnergyBudgets, Unitful, Test
using DynamicEnergyBudgets: non_growth_flux

@testset "Half saturation curve" begin
    mx = 10.0
    half = 5.0
    x = 5.0
    @test half_saturation(mx, half, x) == 5.0

    # Maximum
    x = 100000000000000000000000000.0
    @test half_saturation(mx, half, x) == 10.0

    # Minimum
    x = 0.0
    @test half_saturation(mx, half, x) == 0.0
end

@testset "Flux" begin
    (a, b, c) = non_growth_flux.((0.1,0.2,0.3), (0.2,0.3,0.4), 0.1)
    @test a ≈ 0.01
    @test b ≈ 0.04
    @test c ≈ 0.09
end

@testset "Rate" begin
    j_E_mai    = 0.003
    y_E_CH_NO  = 1.0/1.5
    y_E_EN     = 1.0/0.7
    y_V_E      = 0.7
    kap_soma   = 0.6
    r          = 0.001
    reserve    = (1.0, 1.0, 1.0)
    turnover   = ((0.2, 0.2, 0.2) .* 0.05)
    su = ParallelComplementarySU()

    # rate_formula(r, su, rel_reserve::NTuple{3}, turnover::NTuple{3}, j_E_mai, y_E_Ea, y_E_Eb, y_V_E, κsoma) = begin
    rf = rate_formula(r, su, reserve, turnover, j_E_mai, y_E_CH_NO, y_E_EN, y_V_E, kap_soma)

    @test rf == 0.002874195250659631

    @testset "test rate with units" begin
        j_E_mai   = 0.003u"mol/mol*d^-1"
        turnover = ((0.2u"mol/mol*d^-1", 0.2u"mol/mol*d^-1", 0.2u"mol/mol*d^-1") .* 0.05)
        mass = (1.0u"mol/mol", 1.0u"mol/mol", 1.0u"mol/mol")
        r = 0.001u"mol/mol*d^-1"
        @test_nowarn non_growth_flux.(mass, turnover, r)
        rf = rate_formula(r, su, reserve, turnover, j_E_mai, y_E_CH_NO, y_E_EN, y_V_E, kap_soma)
        @test rf == 0.002874195250659631u"mol/mol*d^-1"
    end

end

nothing


include("backwards_activation_functions.jl")

"""
Partial derivates of the components of linear forward function using the linear output (δZ) and caches of these components (cache).
"""
function linear_backward(δZ, cache)
    # unpack cache
    A_prev, W, b = cache
    m = size(A_prev, 2)

    # partial derivates of each component
    δW = δZ * (A_prev') / m
    δb = sum(δZ, dims = 2) / m
    δA_prev = (W') * δZ

    @assert (size(δA_prev)[1] == size(A_prev)[1])
    @assert (size(δW) == size(W))
    @assert (size(δb)[1] == size(b)[1])

    return δW, δb, δA_prev
end # function

"""
Unpack the linear activated caches (cache) and compute their derivates from the applied activation function.
"""
function linear_activation_backwards(δA, cache; activation_function="relu")
    @assert activation_function ∈ ("sigmoid", "relu", "softmax", "tanh", "swish")

    linear_cache, cache_activation = cache

    if (activation_function == "relu")

        δZ = relu_backwards(δA, cache_activation)

    elseif (activation_function == "sigmoid")

        δZ = sigmoid_backwards(δA, cache_activation)

    elseif (activation_function ==  "tanh")

        δZ = tanh_backwards(δA, cache_activation)

    elseif (activation_function == "softmax")

        δZ = softmax_backwards(δA, cache_activation)

    elseif (activation_function == "swish")

        δZ = swish_backwards(δA, cache_activation)

    end # if

    δW, δb, δA_prev = linear_backward(δZ, linear_cache)

    return δW, δb, δA_prev
end # function


"""
    CircularArc3(c::Circle3)

Constructs an arc from a full circle.
"""
CircularArc3(c::Circle3)

"""
    CircularArc3(c::Circle3, pt::Point3)

Constructs an arc from a full circle, using `pt` as source and target.

!!! info "Precondition"

    `pt` lies on `c`
"""
CircularArc3(c::Circle3, pt::Point3)

"""
    CircularArc3(c::Circle3, p::Point3, q::Point3)

The circular arc constructed from a circle `c`, a source `p`, and a target `q`,
is defined as the set of points of the circle that lie between the source `p₁`
and the target `p₂`, when traversing the circle counterclockwise seen from the
side of the circle's supporting plane pointed by its *positive* normal vectors.

In this definition, we say that a normal vector ``(a,b,c)`` is positive if
``(a,b,c) > (0,0,0)`` (i.e. ``a > 0 ∨ (a = 0 ∧ b > 0) ∨ (b = 0 ∧ c > 0)``).

Constructs the circular arc supported by `c`, whose source and target are `p`
and `q`, respectively.

!!! info "Precondition"

    `p` and `q` lie on `c` and are different.
"""
CircularArc3(c::Circle3, p::Point3, q::Point3)

"""
    CircularArc3(p::Point3, q::Point3, r::Point3)

Constructs an arc that is supported by the circle of type [`Circle3`](@ref)
passing through points `p`, `q` and `r`.

The source and target are respectively `p` and `r`, when traversing the
supporting circle in the counterclockwise direction seen from the side of the
place containing the circle pointed by its *positive* normal vectors.  Note
that, depending on the orientation of the point triple `(p,q,r)`, `q` may not
lie on the arc.

!!! info "Precondition"

    `p`, `q`, and `r` are not collinear.
"""
CircularArc3(p::Point3, q::Point3, r::Point3)

"""
    supporting_circle(ca::CircularArc3)

Returns the supporting circle of `ca`.
"""
supporting_circle(ca::CircularArc3)

"""
    center(ca::CircularArc3)

Returns the center of `ca`'s supporting circle.
"""
center(ca::CircularArc3)

"""
    squared_radius(ca::CircularArc3)

Returns the squared radius of `ca`'s supporting circle.
"""
squared_radius(ca::CircularArc3)

"""
    supporting_plane(ca::CircularArc3)

Returns the supporting plane of `ca`'s supporting circle.
"""
supporting_plane(ca::CircularArc3)

"""
    diametral_sphere(ca::CircularArc3)

Returns the diametral_sphere of `ca`'s supporting circle.
"""
diametral_sphere(ca::CircularArc3)

"""
    source(ca::CircularArc3)

Returns the arc's smallest point, supposing its supporting circle is
traversed in counterclockwise direction.

When the arc was constructed from its (full) underlying circle, then source
returns the smallest ``x``-extremal point of the circle if the circle is not
contained in a plane ``x = A``, and the smallest ``y``-extremal point otherwise.
"""
source(ca::CircularArc3)

"""
    target(ca::CircularArc3)

Returns the arc's ending point, supposing its supporting circle is traversed
in counterclockwise direction.

When the arc was constructed from its (full) underlying circle, then target
returns the smallest ``x``-extremal point of the circle if the circle is not
contained in a plane ``x = A``, and the smallest ``y``-extremal point otherwise.
"""
target(ca::CircularArc3)

"""
    ==(ca₁::CircularArc3, ca₂::CircularArc3)

Test for equality.

Two arcs are equal, iff their non-oriented supporting planes are equal, the
centers and squared radii of their respective supporting circles are equal, and
their sources and targets are equal.
"""
==(ca₁::CircularArc3, ca₂::CircularArc3)


module DH32ParallelUtils

# package code goes here
export split_int, split_int_lengths

"""
Splits an integer number N into a set of n_procs contiguous, non-overlapping integer intervals [ix_start[i], ix_end[i]].

Input:

- N: number to be split.

- n_procs: number of processes.

Output: ix_start, ix_end

- ix_start: int array of start indexes.

- ix_end: int array of end indexes.
"""
function split_int(N::Int64, n_procs::Int64)
    ix_start = zeros(Int64,n_procs)
    ix_end = zeros(Int64,n_procs)
    float_bounds = range(1.0, stop=Float64(N), length=n_procs+1)
    ix_start[1] = 1
    ix_end[n_procs] = N
    for i in 2:n_procs
        ix_start[i] = round(Int64,float_bounds[i])
        ix_end[i-1] = ix_start[i]-1
    end
    ix_start, ix_end
end

"""
Split a integer N in n_procs integer pieces of approximately the same lengths.

Input: as for split_int.

Output:

- Array of n_procs integer lengths summing up to N

"""
function split_int_lengths(N::Int64, n_procs::Int64)
    ix_start, ix_end = split_int(N, n_procs)
    ix_end - ix_start .+ 1
end

end # module


# Sandpile64.jl
module Sandpile64

export sandpile_init64, run_sandpile64!, avalanche64!, add_grain!, is_unstable64

function sandpile_init64(x::Int64, y::Int64, setup::String="random")
    """
            sandpile_init(X, Y, setup)
    Initialize a 2d sandpile grid with size (`X`, `Y`). Can either have a zero or random
    initial condition.
    Arguments:
            x (Int64): length of grid along x
            y (Int64): length of grid along y
            setup (String): Type of inital setup for grid. Can either be `zero` or `random`.
    Returns:
            z (Matrix{Int64}): sandpile grid z with dimensions (x,y).
    """
    if setup == "random"
        return rand(0:3, x, y)
    elseif setup == "zero"
        return zeros(Int64, x, y)
    else
        error("`setup` must have value of `random` or `zero`")
    end
end

function add_grain64!(z::Array{Int64,2})
    """
            add_grain!(z)
    Adds a grain to a random location on the  2d array grid for the Abelian sandpile model, `z`.
    Arguments:
            z (Array{Int64, 2}): 2D sandpile grid.
    """
    return z[rand(1:size(z)[1]), rand(1:size(z)[2])] += 1
end

function is_unstable64(z::Array{Int64,2}, fᵪ::Int64=4)
    """
            is_unstable(z, fᵪ)
    Checks whether or not any location on the grid is unstable (has value above critical value).
    Arguments:
            z (Array{Int64, 2}): 2D sandpile grid.
            fₓ (Int64): Critical value for sandpile model (Must be > 1)
    Returns:
            (bool): Whether grid is stable or unstable.
    """
    if fᵪ <= 1
        error("`fᵪ` must be int greater than  1")
    end
    if length(findall(>=(fᵪ), z)) > 0
        return true
    else
        return false
    end
end

function avalanche64!(z::Array{Int64,2}, fᵪ=4)
    """
            avalanche(z, fᵪ)
    Update grid to follow rules on locations that are unstable.
    Follow Eq 3.2 of Self-organized criticality by Bak, Tang (1988).
    Returns how many avalanches (or slides s) occur for input state after collapsing all
    unstable locations.
    Arguments:
            z (Array{Int64, 2}): 2D sandpile grid.
            fₓ (Int64): Critical value for sandpile model (Must be > 1)
    Returns:
            (Int64): Number of unstable locations collapsed.
    """
    s::Int64 = 0
    unstable_locs = findall(>=(fᵪ), z)
    for crit_loc in unstable_locs
        z[crit_loc] -= 4
        s += 1
        neighbors = [
            CartesianIndex(0, 1),
            CartesianIndex(0, -1),
            CartesianIndex(1, 0),
            CartesianIndex(-1, 0),
        ]
        for neighbor in neighbors
            try
                z[crit_loc + neighbor] += 1
            catch
            end
        end
    end
    return s
end

function run_sandpile64!(z::Array{Int64,2}, N::Int64; N_crit::Bool=false, fᵪ::Int64=4)
    """
            run_sandpile(z, N, fᵪ)
    Run a sandpile model on an input grid `z` by adding `N` grains.
    Arguments:
            z (Array{Int64, 2}): 2D sandpile grid.
            N (Int64): Number of grains to add
            N_crit (boolean): Do you want to to all added `N` grains to result in collapses?
                              If set to `false`, `N` grains will be added to system but
                              returned sizes will only be those that result in avalanches.
                              Warning:  longer runtimes if set to `true`.
            fₓ (Int64): Critical value for sandpile model (Must be > 1)
    Returns
            (Vector{Int64}): Sizes of avalanches for each grain added.
    """
    s_list = zeros(Int64, N)
    if N_crit
            i = 1
            while i <= N
                add_grain64!(z)
                s = 0
                while is_unstable64(z, fᵪ)
                    s += avalanche64!(z)
                end
                if s > 0
                    s_list[i] = s
                    i += 1
                end
            end
    else
        for i in 1:N
            add_grain64!(z)
            s = 0
            while is_unstable64(z, fᵪ)
                s += avalanche64!(z)
            end
            s_list[i] = s
        end
    end
    return s_list[s_list .> 0]
end
end


"""
    eom_nonlinear3d(uᵢ, a₁, a₂, a₃,  b₁, b₂, b₃,
        c₁₂, c₂₃, c₁₃, σ₁, σ₂, σ₃) -> DiscreteDynamicalSystem

Equations of motion for a 3d nonlinear system with nonlinear couplings
``x_1 \\to x_2``, ``x_2 \\to x_3`` and ``x_1 \\to x_3``. Modified from [1].

## Equations of motion

The equations of motion are

```math
\\begin{aligned}
x_1(t+1) &= a_1 x_1 (1-x_1(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{1}(t) \\\\
x_2(t+1) &= a_1 x_2 (1-x_2(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{2}(t) + b x_1 x_2 \\\\
x_3(t+1) &= a_3 x_3 (1-x_3(t))^2  e^{-x_3(t)^2} + 0.4 \\xi_{3}(t) + c x_{2}(t) + d x_{1}(t)^2.
\\end{aligned}
```

## References

1. Gourévitch, B., Le Bouquin-Jeannès, R., & Faucon, G. (2006). Linear and nonlinear
    causality between signals: methods, examples and neurophysiological
    applications. Biological Cybernetics, 95(4), 349–369.
"""
function eom_nonlinear3d(x, p, n)
    x₁, x₂, x₃ = (x...,)
    a₁, a₂, a₃, b₁, b₂, b₃, c₁₂, c₂₃, c₁₃, σ₁, σ₂, σ₃ = (p...,)
    ξ₁ = rand(Normal(0, σ₁))
    ξ₂ = rand(Normal(0, σ₂))
    ξ₃ = rand(Normal(0, σ₃))

    dx₁ = a₁*x₁*(1-x₁)^2 * exp(-x₁^2) + b₁*ξ₁
    dx₂ = a₂*x₂*(1-x₂)^2 * exp(-x₂^2) + b₂*ξ₂ + c₁₂*x₁*x₂
    dx₃ = a₃*x₃*(1-x₃)^2 * exp(-x₃^2) + b₃*ξ₃ + c₂₃*x₂ + c₁₃*x₁^2

    return SVector{3}(dx₁, dx₂, dx₃)
end

function nonlinear3d(uᵢ, a₁, a₂, a₃,  b₁, b₂, b₃, c₁₂, c₂₃, c₁₃, σ₁, σ₂, σ₃)
    p = [a₁, a₂, a₃,  b₁, b₂, b₃, c₁₂, c₂₃, c₁₃, σ₁, σ₂, σ₃]
    s = DiscreteDynamicalSystem(eom_nonlinear3d, uᵢ, p)
    return s
end

"""
    nonlinear3d(;uᵢ = rand(3),
        σ₁ = 1.0, σ₂ = 1.0, σ₃ = 1.0,
        a₁ = 3.4, a₂ = 3.4, a₃ = 3.4,
        b₁ = 0.4, b₂ = 0.4, b₃ = 0.4,
        c₁₂ = 0.5, c₂₃ = 0.3, c₁₃ = 0.5) -> DiscreteDynamicalSystem

A 3d nonlinear system with nonlinear couplings ``x_1 \\to x_2``,
``x_2 \\to x_3`` and ``x_1 \\to x_3``. Modified from [1].

## Equations of motion

The equations of motion are

```math
\\begin{aligned}
x_1(t+1) &= a_1 x_1 (1-x_1(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{1}(t) \\\\
x_2(t+1) &= a_1 x_2 (1-x_2(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{2}(t) + b x_1 x_2 \\\\
x_3(t+1) &= a_3 x_3 (1-x_3(t))^2  e^{-x_3(t)^2} + 0.4 \\xi_{3}(t) + c x_{2}(t) + d x_{1}(t)^2.
\\end{aligned}
```

## References

1. Gourévitch, B., Le Bouquin-Jeannès, R., & Faucon, G. (2006). Linear and nonlinear
    causality between signals: methods, examples and neurophysiological
    applications. Biological Cybernetics, 95(4), 349–369.
"""
nonlinear3d(;uᵢ = rand(3),
        σ₁ = 1.0, σ₂ = 1.0, σ₃ = 1.0,
        a₁ = 3.4, a₂ = 3.4, a₃ = 3.4,
        b₁ = 0.4, b₂ = 0.4, b₃ = 0.4,
        c₁₂ = 0.5, c₂₃ = 0.3, c₁₃ = 0.5) =
    nonlinear3d(uᵢ, a₁, a₂, a₃,  b₁, b₂, b₃, c₁₂, c₂₃, c₁₃, σ₁, σ₂, σ₃)


using LightGraphs: AbstractGraph, adjacency_matrix

struct Graph
    neighbors::Vector{Set{Int}}
    active::BitSet
    peo::Vector{Int}
end

# # SIMPLE LINEAR-TIME ALGORITHMS TO TEST CHORDALITY OF GRAPHS, TEST ACYCLICITY OF HYPERGRAPHS, AND SELECTIVELY REDUCE ACYCLIC HYPERGRAPHS
# # ROBERT E. TARJAN AND MIHALIS YANNAKAKIS
# Algorithm 4.2 in VA
"""
    maximum_cardinality_search(A)

Compute the perfect elimination ordering for a chordal graph represented by
a symmetric sparse matrix `A`. This function can also be used with AbstractGraph
objects from `LightGraphs.jl`.
"""
function maximum_cardinality_search(A::SparseMatrixCSC)
    !issymmetric(A) && error(ArgumentError("Matrix must be symmetric"))
    n = size(A, 1)
    peo = zeros(Int, n)
    cache = zeros(Int, n)

    _maximum_cardinality_search!(peo, A; n = n, cache=cache)
    return peo
end
maximum_cardinality_search(G::AbstractGraph) = maximum_cardinality_search(adjacency_matrix(G))

function _maximum_cardinality_search!(peo, A; n=n, cache=cache)
    v = 1
    for i in n:-1:1
        max_val = 0
        for j in 1:n
            peo[j] != 0 && continue
            if cache[j] > max_val
                max_val = cache[j]
                v = j
            end
        end
        peo[v] = i
        @views cache[(rowvals(A)[nzrange(A, v)])] .+= 1
    end

    return peo
end

"""
    is_chordal(A)

Tests if the graph represented by symmetric sparse matrix `A` is chordal. This
function can also be used with AbstractGraph objects from `LightGraphs.jl`.

# References
* [Simple linear-time algorithms to test chordality of graphs, test acyclicity of hypergraphs, and selectively reduce acyclic hypergraphs](https://epubs.siam.org/doi/pdf/10.1137/0213035?casa_token=A22jkwrsrL0AAAAA:rx-G6F21ubTkMiJmRTH3IKqxmFTo_IVWDDZfJig5lsZxnQtNH2vUKWfZ3eZJKUv9CiKbIPt1VQs) by Robert Tarjan and Mihalis Yannakakis
"""
function is_chordal(A::SparseMatrixCSC; peo=nothing)
    if isnothing(peo)
        peo = maximum_cardinality_search(A)
    end
    i_peo = invperm(peo)

    n = length(peo)
    f = zeros(Int, n)
    ind = zeros(Int, n)

    for i in 1:n
        w = i_peo[i]
        f[w] = w
        ind[w] = i
        for v in rowvals(A)[nzrange(A, w)]
            peo[v] >= i && continue
            ind[v] = i
            if f[v] == v
                f[v] = w
            end
        end

        for v in rowvals(A)[nzrange(A, w)]
            peo[v] >= i && continue
            ind[f[v]] < i && return false
        end
    end
    return true
end

is_chordal(G::AbstractGraph; peo=nothing) = is_chordal(adjacency_matrix(G); peo=peo)


module LocalMethods

using SpecialFunctions: besselk, gamma

#M_vρσ is the matern covariance and G_σρ is Gaussian autocovariance model
export ℳ_ν, M_νρσ, G_σρ

tν𝒦t(t,ν) = t^ν * besselk(ν, t)

"""
    ℳ_ν(t, ν)

Compute (√(2ν)*t) ^ ν * 𝒦ν(√(2ν)*t) where 𝒦ν is the modified Bessel function of the second kind of order ν.

# Example
```julia-repl
julia> ℳ_ν(0.0, 0.5)
1.0
```
"""
function ℳ_ν(t, ν)
	pt, pν, p0, p1 = promote(t, ν, 0, 1)
	return (pt==p0) ? p1 : tν𝒦t(√(2pν)*pt,pν) * 2^(1-pν) / gamma(pν)
end

function M_νρσ(t; ν, ρ, σ)
	return σ * σ * ℳ_ν(t/ρ, ν)
end

function G_σρ(t; σ,ρ)
	return σ * σ *exp(-t^2/ρ^2)

end
end # end module


# Test rosenbrock pdf gradients
# Authors: Ali Siahkoohi, alisk@gatech.edu
# Date: September 2020

using DrWatson
@quickactivate :FastApproximateInference
import Pkg; Pkg.instantiate()

using LinearAlgebra
using Test

# Input
n_samples = 1

# Rosenbrock distribution
RB_dist = RosenbrockDistribution(0f0, 5f-1, ones(Float32, 2, 1))

# samples
X0 = rand(RB_dist, n_samples)
X = rand(RB_dist, n_samples)
dX = X - X0

function rosenbrock_nlpdf(RB_dist, x)
    return nl_pdf(x, RB_dist)[1], nl_pdf_grad(x, RB_dist)
end

# Gradient test
f0, ΔX = rosenbrock_nlpdf(RB_dist, X0)
h = 0.1f0
maxiter = 5
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test\n")
for j=1:maxiter
    f = rosenbrock_nlpdf(RB_dist, X0 + h*dX)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)


# Text provided under a Creative Commons Attribution license, CC-BY.  All code is made available under the FSF-approved BSD-3 license.  (c) Lorena A. Barba, Gilbert F. Forsyth 2017. Thanks to NSF for support via CAREER award #1149784.
# [@LorenaABarba](https://twitter.com/LorenaABarba)
# Translation to Julia by [@miguelraz_](https://twitter.com/miguelraz_)

# 12 steps to Navier–Stokes
# =====
# ***

# In the previous step, we solved the [2D Burgers' equation](./10_Step_8.ipynb): an important equation in the study of fluid mechanics because it contains the full convective nonlinearity of the flow equations. With that exercise, we also build the experience to incrementatlly code a Navier–Stokes solver.

# In the next two steps, we will solve Laplace and then Poisson equation. We will then put it all together!

# Step 9: 2D Laplace Equation
# ----
# ***

# Here is Laplace's equation in 2D:

# $$\frac{\partial ^2 p}{\partial x^2} + \frac{\partial ^2 p}{\partial y^2} = 0$$

# We know how to discretize a 2nd order derivative. But think about this for a minute — Laplace's equation has the features typical of diffusion phenomena. For this reason, it has to be discretized with *central differences*, so that the discretization is consistent with the physics we want to simulate.

# The discretized equation is:

# $$\frac{p_{i+1, j}^n - 2p_{i,j}^n + p_{i-1,j}^n}{\Delta x^2} + \frac{p_{i,j+1}^n - 2p_{i,j}^n + p_{i, j-1}^n}{\Delta y^2} = 0$$

# Notice that the Laplace Equation does not have a time dependence — there is no $p^{n+1}$.  Instead of tracking a wave through time (like in the previous steps), the Laplace equation calculates the equilibrium state of a system under the supplied boundary conditions.

# If you have taken coursework in Heat Transfer, you will recognize the Laplace Equation as the steady-state heat equation.

# Instead of calculating where the system will be at some time $t$, we will iteratively solve for $p_{i,j}^n$ until it meets a condition that we specify.  The system will reach equilibrium only as the number of iterations tends to $\infty$, but we can approximate the equilibrium state by iterating until the change between one iteration and the next is *very* small.

# Let's rearrange the discretized equation, solving for $p_{i,j}^n$:

# $$p_{i,j}^n = \frac{\Delta y^2(p_{i+1,j}^n+p_{i-1,j}^n)+\Delta x^2(p_{i,j+1}^n + p_{i,j-1}^n)}{2(\Delta x^2 + \Delta y^2)}$$

# Using second-order central-difference schemes in both directions is the most widely applied method for the Laplace operator. It is also known as the **five-point difference operator**, alluding to its stencil.

# We are going to solve Laplace's equation numerically by assuming an initial state of $p=0$ everywhere. Then we add boundary conditions as follows:

# $p=0$ at $x=0$

# $p=y$ at $x=2$

# $\frac{\partial p}{\partial y}=0$ at $y=0, \ 1$

# Under these conditions, there is an analytical solution for Laplace's equation:

# $$p(x,y)=\frac{x}{4}-4\sum_{n=1,odd}^{\infty}\frac{1}{(n\pi)^2\sinh2n\pi}\sinh n\pi x\cos n\pi y$$

# ### Exercise

# Write your own code to solve Poisson's equation using loops, in the style of coding used in our first lessons.
# Then, consider the demonstration of how to write it using functions (below) and modify your code in that style.
# Can you think of reasons to abandon the old style and adopt modular coding?

# Other tips:

# + Visualize each step of the iterative process
# + Think about what the boundary conditions are doing
# + Think about what the PDE is doing

# ### Using functions

# Remember the lesson on writing [functions with Julia](./11_Defining_Functions_in_Julia.jl)? We will use that style of code in this exercise.

# We're going to define two functions: one that plots our data in a 3D projection plot and the other that iterates to solve
# for $p$ until the change in the [L1 Norm][1] of $p$ is less than a specified value.

# [1]: http://en.wikipedia.org/wiki/Norm_(mathematics)#Taxicab_norm_or_Manhattan_norm

using Plots

# TODO PLOT
function  plot2D(x, y, p)
    fig = pyplot.figure(figsize=(11, 7), dpi=100)
    ax = fig.gca(projection='3d')
    X, Y = numpy.meshgrid(x, y)
    surf = ax.plot_surface(X, Y, p[:], rstride=1, cstride=1, cmap=cm.viridis,
            linewidth=0, antialiased=False)
    ax.set_xlim(0, 2)
    ax.set_ylim(0, 1)
    ax.view_init(30, 225)
    ax.set_xlabel('$x$')
    ax.set_ylabel('$y$')
end

# The function `plot2D` takes three arguments, an x-vector, a y-vector and our p matrix.
#     Given these three values, it produces a 3D projection plot,
#     sets the plot limits and gives us a nice viewing angle.

# $$p_{i,j}^n = \frac{\Delta y^2(p_{i+1,j}^n+p_{i-1,j}^n)+\Delta x^2(p_{i,j+1}^n + p_{i,j-1}^n)}{2(\Delta x^2 + \Delta y^2)}$$

function laplace2d(p, y, dx, dy, l1norm_target)
    l1norm = 1
    pn = empty(p)

    while l1norm > l1norm_target
        pn = copy(p)
        p[2:end, 2:end] = ((dy^2 * (pn[2:end, 2:] + pn[2:end, 0:-2]) +
                         dx^2 * (pn[2:, 2:end] + pn[0:-2, 2:end])) /
                        (2 * (dx^2 + dy^2)))

        p[:, begin] = 0  # p = 0 @ x = 0
        p[:,   end] = y  # p = y @ x = 2
        p[begin, :] = p[2, :]  # dp/dy = 0 @ y = 0
        p[end,   :] = p[end-1, :]  # dp/dy = 0 @ y = 1
        # TODO Port
        l1norm = (sum(abs.(p[:]) - abs.(pn[:])) /
                sum(abs.(pn[:]))
    end

    return p
end

# `laplace2d` takes five arguments, the `p` matrix, the `y`-vector, `dx`, `dy` and the value `l1norm_target`.
# This last value defines how close the `p` matrix should be in two consecutive iterations before the loop breaks and returns the calculated `p` value.

# Note that when executing the cells above in your own notebook, there will be no output.
# You have *defined* the function but you have not yet *called* the function.
# It is now available for you to use, the same as `numpy.linspace` or any other function in our namespace.

## Variable declarations
nx = 31
ny = 31
c = 1
dx = 2 / (nx - 1)
dy = 2 / (ny - 1)

# initial conditions
p = zeros((ny, nx))  # create a XxY vector of 0's

# plotting aids
x = LinRange(0, 2, nx)
y = LinRange(0, 1, ny)

##boundary conditions
p[:, begin] = 0  # p = 0 @ x = 0
p[:, end]   = y  # p = y @ x = 2
p[begin, :] = p[2, :]  # dp/dy = 0 @ y = 0
p[end, :]   = p[end-1, :]  # dp/dy = 0 @ y = 1

# Now let's try using our `plot2D` function to look at our initial conditions.
# If the function has been correctly defined,
# you should be able to begin typing `plot2D`
# and hit the **tab** key for auto-complete options.

# TODO PLOT
plot2D(x, y, p)

# It worked!  This is the initial state of our problem, where the value of `p`
# is zero everywhere except for along $x=2$ where $p=y$.
# Now let's try to run our `laplace2d` function with a specified L1 target of .01

# [Hint: if you are having trouble remembering the order in which variables are sent to a function,
# you can just type `laplace2d(` and the iPython Notebook will put up a little popup box to remind you]

p = laplace2d(p, y, dx, dy, 1e-4)

# Now try plotting this new value of `p` with our plot function.

# TODO PLOT
plot2D(x, y, p)

# ***

## Learn More

# The [next step](./13_Step_10.ipynb) will be to solve Poisson's equation. Watch **Video Lesson 11** on You Tube to understand why we need Poisson's equation in CFD.
#src="https://www.youtube.com/embed/ZjfxA3qq2Lg"

# And for a detailed walk-through of the discretization of Laplace and Poisson equations (steps 9 and 10), watch **Video Lesson 12** on You Tube:
#src="https://www.youtube.com/embed/iwL8ashXhWU"


module Probabilities
using Base:Integer, AbstractFloat

include("CheckArgument.jl")

"""
计算排列数 A_n^m,输入需要是整数类型，返回值是BigInt类型，因为排列数很容易变成大数
"""
function A(m::Integer, n::Integer)
    # basic param check
     if n < m
        throw(ArgumentError("n should be lager than m"))
    elseif n <= 0 || m < 0
        throw(ArgumentError("n or m should not be zero or negative"))
    elseif m == 0 return 1
            end

    function A_inner(m, n)
        if (m <= 1)
            return n
        end
        return A_inner(m - 1, n - 1) * n
    end
    return A_inner(BigInt(m), BigInt(n))
end

"""
计算组合数
"""
function C(m::Integer, n::Integer)
    # basic param check
    if n < m
        throw(ArgumentError("n should be lager than m"))
    elseif n <= 0 || m < 0
        throw(ArgumentError("n or m should not be zero or negative"))
    elseif m == 0 return 1
        end
    return A(BigInt(m), BigInt(n)) / factorial(BigInt(m))
end

"""
计算二项分布的列表
其中参数n为发生次数，p为单次独立事件的概率
Binomial_Distribution(n::Integer, p)
返回值是一个列表
"""
function Binomial_Distribution(n::Integer, p::AbstractFloat)
    check_probabilities(p)
    resList = []
	for k = 0:n
		# @info(k)
		push!(resList, C(k, n) * p^k * (1 - p)^(n - k))
	end
	return resList
end

"""
泊松分布
λᵏℯ^λ / k!,k=0,1,2
λ=n⋅p
function Poisson_Distribution(num::Integer, λ)
"""
function Poisson_Distribution(range::AbstractRange, λ::Real)
    check_is_nonnegative(λ, range[1])
    resList = []
    for k in range
        push!(resList, λ^k * ℯ^-λ / factorial(big(k)))
    end
    return resList
end

"""
    几何分布
他的每一项为 (1-p)^(k-1)*p
起始项为k=1zz
function Geometric_distribution(num::Integer, p::AbstractFloat)
"""
function Geometric_distribution(range::AbstractRange, p::AbstractFloat)
    check_probabilities(p)
    check_is_nonnegative(range[1])
    resList = []
    for k in range
        push!(resList, p * (1 - p)^(k - 1))
    end
    return resList
end

"""
    超几何分布
从N件产品（其中次品M件）中任取n件，k为取到的次品数

p{x=k}=C(k,M)C(n-k,n-M)/C(n,N)
"""
function Hypergeometric_Distribution(n::Integer, N::Integer, k::Integer, M::Integer)
    if !(0 <= k <= n <= N  && k <= M)
        throw(ArgumentError("0 <= k <= n <= N  &&k <= M"))
    end
    return C(k, M) * C(n - k, n - M) / C(n, N)
end

"""
    正态分布
function Normal_Distribution_pdf(range::AbstractRange, μ, σ)
"""
function Normal_Distribution_pdf(range::AbstractRange, μ::Real, σ::Real)
check_not_empty(range)
    check_is_nonnegative(μ)
    resList = []
    for x in range
        constant = 1 / (sqrt(2π) * σ)
        push!(resList, constant * ℯ^(-(x - μ)^2 / (2 * σ^2)))
    end
    return resList
end

"""
    指数分布 概率密度函数

function Exponential_Distribution(range::AbstractRange, θ::Real)
"""
function Exponential_Distribution_pdf(range::AbstractRange, θ::Real)
    check_not_empty(range)
    check_is_nonnegative(range[1])
    resList = []
    for x in range
        push!(resList, (1 / θ) * ℯ^(-x / θ))
    end
    return resList
end

"""
    指数分布
function Exponential_Distribution(range::AbstractRange, θ::Real)
"""
function Exponential_Distribution(range::AbstractRange, θ::Real)
    check_not_empty(range)
    check_is_nonnegative(range[1])
    resList = []
    for x in range
        push!(resList, 1 - ℯ^(-x / θ))
    end
    return resList
end
export A,C
end


export TimeScaleMODWT
import Wavelets
import Wavelets: wavelet, maxdyadiclevel, modwt

"""
    TimeScaleMODWT <: WaveletProbabilitiesEstimator
    TimeScaleMODWT(wl::Wavelets.WT.OrthoWaveletClass = Wavelets.WT.Daubechies{12}())

Apply the maximal overlap discrete wavelet transform (MODWT) to a
signal, then compute probabilities/entropy from the energies at different
wavelet scales. This implementation is based on Rosso et
al. (2001)[^Rosso2001].
Optionally specify a wavelet to be used.

The probability `p[i]` is the relative/total energy for the i-th wavelet scale.

## Example

Manually picking a wavelet is done as follows.

```julia
using Entropies, Wavelets
N = 200
a = 10
t = LinRange(0, 2*a*π, N)
x = sin.(t .+  cos.(t/0.1)) .- 0.1;

# Pick a wavelet (if no wavelet provided, defaults to Wavelets.WL.Daubechies{12}())
wl = Wavelets.WT.Daubechies{12}()

# Compute the probabilities (relative energies) at the different wavelet scales
probabilities(x, TimeScaleMODWT(wl))
```

[^Rosso2001]: Rosso, O. A., Blanco, S., Yordanova, J., Kolev, V., Figliola, A., Schürmann, M., & Başar, E. (2001). Wavelet entropy: a new tool for analysis of short duration brain electrical signals. Journal of neuroscience methods, 105(1), 65-75.
"""
struct TimeScaleMODWT <: WaveletProbabilitiesEstimator
    wl::Wavelets.WT.OrthoWaveletClass
    function TimeScaleMODWT(wl::Wavelets.WT.OrthoWaveletClass = Wavelets.WT.Daubechies{12}())
        new(wl)
    end
end

function get_modwt(x::AbstractVector{T}, wl::Wavelets.WT.OrthoWaveletClass = Wavelets.WT.Daubechies{12}()) where T<:Real
    orthofilter = wavelet(wl)
    nscales = maxdyadiclevel(x)
    W = modwt(x, orthofilter, nscales)
end

function energy_at_scale(W::AbstractArray{T, 2}, j::Int) where T<:Real
    1 <= j <= size(W, 2) || error("Scale j does not exist in wave coefficient matrix W. Available scales are j=1:$(size(W, 2))")
    Eⱼ = sum(W[:, j] .^ 2)
end

function energy_at_time(W::AbstractArray{T, 2}, t::Int) where T<:Real
    1 <= t <= size(W, 1) || error("Time t does not exist in wave coefficient matrix W. Available times are t=1:$(size(W, 1))")
    Eⱼ = sum(W[t, :] .^ 2)
end

function energy_total(W::AbstractArray{T, 2}) where T<:Real
    Etot = sum(W .^ 2)
end

function relative_wavelet_energy(W::AbstractArray{T, 2}, j::Int) where T<:Real
    energy_at_scale(W, j) / energy_total(W)
end

function relative_wavelet_energies(W::AbstractArray{T, 2}, js = 1:size(W, 2)) where T<:Real
    all(1 .<= js .<= size(W, 2)) || error(ArgumentError("scales $(js) contains scales not present in wavelet coefficient matrix with scales j=1:$(size(W, 2))"))
    [energy_at_scale(W, j) / energy_total(W) for j in js]
end

function time_scale_density(x::AbstractVector{T}, wl::Wavelets.WT.OrthoWaveletClass = WT.Daubechies{12}()) where T
    W = get_modwt(x, wl)
    Pⱼs = relative_wavelet_energies(W)
end

function probabilities(x::AbstractVector{T}, est::TimeScaleMODWT) where T<:Real
    Probabilities(time_scale_density(x, est.wl))
end


# CUR decomposition sketch
# LinearTimeCUR - Drineas, et  al. "Fast monte carlo algorithms for matrices iii: computing a compressed approximate matrix decomposition", SIAM Journal of Computing, 2006.
#
# Danny Perry (dperry@cs.utah.edu)
# April 2015

using Debug

# CUR decomposition (rank-k)
# A - the matrix to approximate
# fro_norm - frobenius norm of A
# num_cols - number of colums to sample for C (and U)
# num_rows - number of rows to sample for R (and U)
# k - rank k to approximate
#
# returns C,U,R matrices
@debug function LinearTimeCUR(A,fro_norm, num_cols,num_rows,k)
	n = size(A,1)
	d = size(A,2)

	c = num_cols
	r = num_rows

	C = zeros(n,c)
	U = zeros(c,r)
	R = zeros(r,d)

	for i=1:d
		pc = norm(A[:,i])^2 / fro_norm^2
		scale = 1/sqrt(c*pc)
		for j=1:c
			if rand() <= pc
				C[:,j] = scale * A[:,i]
			end
		end
	end
	CU,CS,CV = svd(C'*C)
	k = min(k, rank(diagm(CS)))
	for i=1:n
		pc = norm(A[i,:])^2 / fro_norm^2
		scale = 1/sqrt(r*pc)
		for j=1:r
			if rand() <= pc
				R[j,:] = scale * A[i,:]
				U[:,j] = scale * C[i,:]
			end
		end
	end
	U = CV[:,1:k] * pinv(diagm(CS[1:k])) * CU[:,1:k]' * U

	C,U,R
end

# This uses a decomposition to choose samples based on the leverage scores
# Uk,Vk - the first k columns of the matrix U and V, such that  A = USV'
@debug function LeverageCUR(A,fro_norm,Uk,Vk, num_cols,num_rows,k)
	n = size(A,1)
	d = size(A,2)

	c = num_cols
	r = num_rows

	C = zeros(n,c)
	U = zeros(c,r)
	R = zeros(r,d)

	col_lev = sum(Vk.^2,2) ./ k
	row_lev = sum(Uk.^2,2) ./ k

	for i=1:d
		pc = col_lev[i]
		scale = 1/sqrt(c*pc)
		for j=1:c
			if rand() <= pc
				C[:,j] = scale * A[:,i]
			end
		end
	end
	CU,CS,CV = svd(C'*C)
	k = min(k, rank(diagm(CS)))
	for i=1:n
		pc = row_lev[i]
		scale = 1/sqrt(r*pc)
		for j=1:r
			if rand() <= pc
				R[j,:] = scale * A[i,:]
				U[:,j] = scale * C[i,:]
			end
		end
	end
	U = CV[:,1:k] * pinv(diagm(CS[1:k])) * CU[:,1:k]' * U

	C,U,R
end


"""
`ExponentiatedKernel([ρ=1])`
The exponentiated kernel is a Mercer kernel given by:
```
    κ(x,y) = exp(ρ²xᵀy)
```
"""
struct ExponentiatedKernel{Tr} <: Kernel{Tr}
    transform::Tr
end
@inline kappa(κ::ExponentiatedKernel, xᵀy::T) where {T<:Real} = exp(xᵀy)

metric(::ExponentiatedKernel) = DotProduct()


"""
Struct for Geometric Brownian Motion

		gbmProcess=GeometricBrownianMotion(σ::num1,μ::num2) where {num1,num2 <: Number}

Where:\n
		σ	=	volatility of the process.
		μ	=	drift of the process.
		x0	=	initial value.
"""
mutable struct GeometricBrownianMotion{num <: Number, num1 <: Number, num2 <: Number, numtype <: Number} <: ItoProcess{numtype}
    σ::num
    μ::num1
    x0::num2
    function GeometricBrownianMotion(σ::num, μ::num1, x0::num2) where {num <: Number, num1 <: Number, num2 <: Number}
        if σ <= 0
            error("Volatility must be positive")
        else
            zero_typed = zero(num) + zero(num1)
            return new{num, num1, num2, typeof(zero_typed)}(σ, μ, x0)
        end
    end
end

export GeometricBrownianMotion;

function simulate!(X, mcProcess::GeometricBrownianMotion, mcBaseData::AbstractMonteCarloConfiguration, T::Number)
    @assert T > 0
    σ_gbm = mcProcess.σ
    mu_gbm = mcProcess.μ
    μ_bm = mu_gbm - σ_gbm^2 / 2
    simulate!(X, BrownianMotion(σ_gbm, μ_bm), mcBaseData, T)
    S0 = mcProcess.x0
    # @. X=S0*exp(X);
    broadcast!(x -> S0 * exp(x), X, X)

    nothing
end


function modeorthonorm(basis,N,err = 1e-9,ω=1.)

x,w,T = nfieldtrans(N,2,ω=ω)
Normtest = T'*(T.*(w.*ones(N)))

#check orthonormality is close to identity matrix for the basis
maxerr = maximum(abs.(Normtest-eye(Normtest)))
maxerr < err
end


#!/usr/bin/env julia

using SEM
using LinearAlgebra,Plots,UnPack

#----------------------------------#
# Set up grid
#----------------------------------#
Ex = 8; nr1 = 8;
Ey = 8; ns1 = 8;

function deform(x,y)
    x = @. 0.5*(x+1)
    y = @. 0.5*(y+1)
    return x,y
end

ifperiodic = [false,false] # [xdir, ydir]

m1 = Mesh(nr1,ns1,Ex,Ey,deform,ifperiodic)

#----------------------------------#
# case setup
#----------------------------------#
# solve with Newton Raphson approach
#
# ν∇²u + G(u,x) = f(x), x∈Ω
#     u = ub,         , x∈∂Ω
#
#----------------------------------#
x1 = m1.x
y1 = m1.y

function G(u,msh::Mesh) # G(u,x⃗)
    @unpack x,y = msh
#   Gu = @. u^3
    Gu = @. -exp.(-u)
    return Gu
end
function Gdu(u,msh::Mesh) # d/du G(u,x⃗)
    @unpack x,y = msh
#   Gdu = @. 3u^2
    Gdu = @. exp.(-u)
    return Gdu
end

ν  = @. 1+0*x1
u  = @. 0+0*x1 # initial guess (must agree with boundary data)
δu = @. 0+0*x1 # initialize
f  = @. 1+0*x1

# contrived example
kx = 3
ky = 3
ut = @. 0 + sin(kx*pi*x1)*sin(ky*pi*y1) # true solution
f  = @. ut*((kx^2+ky^2)*pi^2)         # forcing/RHS
f += G(ut,m1)

# boundary condition
# 'N': Neumann (homogeneous)
# 'D': Dirichlet (inhomogeniety applied through boundary data)
M1 = generateMask(['D','D','D','D'],m1) # [xmin,xmax,ymin,ymax]

#----------------------------------#
# solve linearized homogeneous equation for δu
#----------------------------------#
function residual(v,msh::Mesh)
    resi  = mass(f,msh)
    resi -= mass(G(v,msh),msh) .+ ν.*lapl(v,msh)
    resi .= mask(resi,M1)
    resi .= gatherScatter(resi,m1)
    return resi
end

function opHlmz(v,k,msh::Mesh) # LHS op
    Hu = hlmz(v,ν,k,msh)
    Hu = gatherScatter(Hu,msh)
    Hu = mask(Hu,M1)
    return Hu
end

for i=1:100
    local g = G(u,m1)
    local k = Gdu(u,m1)

    # LHS
    opLHS(v) = opHlmz(v,k,m1)

    # RHS
    b  = mass(f,m1)                  # forcing
    b -= ν.*lapl(u,m1) .+ mass(g,m1) # from LHS
    b .= mask(b,M1)
    b .= gatherScatter(b,m1)

    # solve
    pcg!(δu,b,opLHS,mult=m1.mult,ifv=false)
    u .+= δu

    r  = residual(u,m1);rr = norm(r,Inf);println("iter $i, res norm $rr")
    if(rr < 1e-8) break end
end
#----------------------------------#
er=ut-u; er=norm(er,Inf); println("Error ∞ norm: $er")
plt = meshplt(u,m1)
display(plt)
#----------------------------------#
nothing


# # Approximate Inference using ESS

# Loading the necessary packages and setting seed.

using AbstractGPs, Plots, Random
Random.seed!(1234);

# Loading toy regression
# [dataset](https://github.com/GPflow/docs/blob/master/doc/source/notebooks/basics/data/regression_1D.csv)
# taken from GPFlow examples.

x = [0.8658165855998895, 0.6661700880180962, 0.8049218148148531, 0.7714303440386239,
    0.14790478354654835, 0.8666105548197428, 0.007044577166530286, 0.026331737288148638,
    0.17188596617099916, 0.8897812990554013, 0.24323574561119998, 0.028590102134105955];
y = [1.5255314337144372, 3.6434202968230003, 3.010885733911661, 3.774442382979625,
    3.3687639483798324, 1.5506452040608503, 3.790447985799683, 3.8689707574953,
    3.4933565751758713, 1.4284538820635841, 3.8715350915692364, 3.7045949061144983];
scatter(x, y, xlabel="x", ylabel="y")

# Split the observations into train and test set.

(x_train, y_train) = (x[begin:8], y[begin:8]);
(x_test, y_test) = (x[9:end], y[9:end]);

# Instantiating the kernel.

k = Matern52Kernel()

# Instantiating a Gaussian Process with the given kernel `k`.

f = GP(k)

# Instantiating a `FiniteGP`, a finite dimentional projection at the inputs of the dataset
# observed under Gaussian Noise with $\sigma = 0.001$ .

fx = f(x_train, 0.001)

# Data's log-likelihood w.r.t prior `GP`.

logpdf(fx, y_train)

# Calculating the exact posterior over `f` given `y`. The GP's kernel currently has some
# arbitrary fixed parameters.

p_fx = posterior(fx, y_train)

# Data's log-likelihood under the posterior `GP`. We see that it drastically increases.

logpdf(p_fx(x_test), y_test)

# Plot the posterior `p_fx` along with the observations.

plt = scatter(x_train, y_train; label = "Train data")
scatter!(plt, x_test, y_test; label = "Test data")
plot!(plt, p_fx, 0:0.001:1; label="Posterior")

# # Elliptical Slice Sampler
# Previously we computed the log likelihood of the untuned kernel parameters
# of the GP, $-1.285$. We now also perform approximate inference over said
# kernel parameters using the
# [Elliptical Slice Sampling](http://proceedings.mlr.press/v9/murray10a/murray10a.pdf)
# provided by
# [EllipticalSliceSampling.jl](https://github.com/TuringLang/EllipticalSliceSampling.jl/).
# We start of by loading necessary packages.

using EllipticalSliceSampling, Distributions

# We define a function which returns log-probability of the data under the
# GP / log-likelihood of the parameters of the GP.

function logp(params; x=x_train, y=y_train)
    kernel = ScaledKernel(
        transform(
            Matern52Kernel(),
            ScaleTransform(exp(params[1]))
        ),
        exp(params[2])
    )
    f = GP(kernel)
    fx = f(x, 0.1)
    return logpdf(fx, y)
end

# We define a Gaussian prior over the joint distribution on kernel parameters
# space. Since we have only two parameters, we define a multi-variate Gaussian
# of dimension two.

prior = MvNormal(2, 1)

# Sanity check for the defined `logp` function and `prior` distribution.

logp(rand(prior))

# Generate 2,000 samples using `ESS_mcmc` provided by `EllipticalSliceSampling.jl`.

samples = sample(ESSModel(prior, logp), ESS(), 2_000; progress=false)
samples_mat = reduce(hcat, samples)';

# Mean of samples of both the parameters.

mean_params = mean(samples_mat, dims=1)

# Plot a histogram of the samples for the two parameters. The vertical line in each
# graph indicates the mean of the samples.

plt = histogram(samples_mat; layout=2, labels="Param")
vline!(plt, mean_params; layout=2, label="Mean")

# Average log-marginal-probability of data with posterior kernel parameter samples
# sampled using ESS. We can observe that there is significant improvement over
# exact posterior with default kernel parameters.

mean(logp(param; x=x_test, y=y_test) for param in samples)

# Plot sampled functions from posterior with tuned parameters

plt = scatter(x_train, y_train; label="Train data")
scatter!(plt, x_train, y_train; label="Test data")
for params in @view(samples[(end-100):end,:])
    opt_kernel = ScaledKernel(
        transform(
            Matern52Kernel(),
            ScaleTransform(exp(params[1]))
        ),
        exp(params[2])
    )
    f = GP(opt_kernel)
    p_fx = posterior(f(x, 0.1), y)
    sampleplot!(plt, p_fx(collect(0:0.02:1)), 1)
end
plt


# The function rref(), but with the modification that it also returns the pivot vector

export rref_with_pivots, rref_with_pivots!

function rref_with_pivots!(A::Matrix{T}, ɛ=T <: Union{Rational,Integer} ? 0 : eps(norm(A,Inf))) where T
    nr, nc = size(A)
    pivots = Vector{Int64}()
    i = j = 1
    while i <= nr && j <= nc
        (m, mi) = findmax(abs.(A[i:nr,j]))
        mi = mi+i - 1
        if m <= ɛ
            if ɛ > 0
                A[i:nr,j] .= zero(T)
            end
            j += 1
        else
            for k=j:nc
                A[i, k], A[mi, k] = A[mi, k], A[i, k]
            end
            d = A[i,j]
            for k = j:nc
                A[i,k] /= d
            end
            for k = 1:nr
                if k != i
                    d = A[k,j]
                    for l = j:nc
                        A[k,l] -= d*A[i,l]
                    end
                end
            end
            append!(pivots,j)
            i += 1
            j += 1
        end
    end
    return A, pivots
end

rref_with_pivots_conv(::Type{T}, A::Matrix) where {T} = rref_with_pivots!(copyto!(similar(A, T), A))

"""
    rref_with_pivots(A)
Compute the reduced row echelon form of the matrix A together with the
position of the pivots.
Since this algorithm is sensitive to numerical imprecision,
* Complex numbers are converted to ComplexF64
* Integer, Float16 and Float32 numbers are converted to Float64
* Rational are kept unchanged

```jldoctest
julia> rref_with_pivots([ 1  2 -1  -4;
              2  3 -1 -11;
             -2  0 -3  22])
3×4 Array{Float64,2}:
 1.0  0.0  0.0  -8.0
 0.0  1.0  0.0   1.0
 0.0  0.0  1.0  -2.0
 Int64[3]:
 1
 2
 3

julia> rref_with_pivots([16  2  3  13;
              5 11 10   8;
              9  7  6  12;
              4 14 15   1])
4×4 Array{Float64,2}:
 1.0  0.0  0.0   1.0
 0.0  1.0  0.0   3.0
 0.0  0.0  1.0  -3.0
 0.0  0.0  0.0   0.0
 Int64[3]:
 1
 2
 3

julia> rref_with_pivots([ 1  2  0   3;
              2  4  0   7])
2×4 Array{Float64,2}:
 1.0  2.0  0.0  0.0
 0.0  0.0  0.0  1.0
 Int64[3]:
 1
 4
```
"""
rref_with_pivots(A::Matrix{T}) where {T} = rref_with_pivots!(copy(A))
rref_with_pivots(A::Matrix{T}) where {T <: Complex} = rref_with_pivots_conv(ComplexF64, A)
rref_with_pivots(A::Matrix{ComplexF64}) = rref_with_pivots!(copy(A))
rref_with_pivots(A::Matrix{T}) where {T <: Union{Integer, Float16, Float32}} = rref_with_pivots_conv(Float64, A)

rref_with_pivots(A::AbstractMatrix) = rref_with_pivots(Matrix(A))


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

using Markdown
using InteractiveUtils

# ╔═╡ 6fd616f6-1ea1-11eb-3814-8bfb4a096c49
using Pkg, DrWatson

# ╔═╡ 6ffe2628-1ea1-11eb-24ea-57f985146a72
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Distributions, Random, StatsBase, Measures
	using DataFrames, Plots, LaTeXStrings
	Random.seed!(1)
end;

# ╔═╡ ed174bc4-1ea0-11eb-1e2f-a32874cec549
md"## Listing 3.31"

# ╔═╡ 70179b62-1ea1-11eb-13ce-7fb4d54bad4b
begin
	delta = 0.01
	grid = 0:delta:1
	f(x,y) = 9/8*(4x+y)*sqrt((1-x)*(1-y))
	z = [f(x,y) for y in grid, x in grid]

	densityIntegral = sum(z)*delta^2
	Text("2-dimensional Riemann sum over density: $(densityIntegral)")
end

# ╔═╡ f3fb9a0e-2294-11eb-227d-135a8b801a20
begin
	probB = sum([sum([f(x,y)*delta for y in x:delta:1])*delta for x in grid])
	Text("2-dimensional Riemann sum to evaluate probability: $(probB)")
end

# ╔═╡ f3fbcc54-2294-11eb-1a96-b1676ddb7772
begin
	p1 = surface(grid, grid, z,
		c=cgrad([:blue, :red]), la=1, camera=(60,50),
		ylabel="y", zlabel=L"f(x,y)", legend=:none)
	p2 = contourf(grid, grid, z,
		c=cgrad([:blue, :red]))
	p2 = contour!(grid, grid, z,
		c=:black, xlims=(0,1), ylims=(0,1), ylabel="y", ratio=:equal)

	plot(p1, p2, size=(800, 400), xlabel="x", margin=5mm)
end

# ╔═╡ 70182cd8-1ea1-11eb-094d-8d0d49cf15f3
md"## End of listing 3.31"

# ╔═╡ Cell order:
# ╟─ed174bc4-1ea0-11eb-1e2f-a32874cec549
# ╠═6fd616f6-1ea1-11eb-3814-8bfb4a096c49
# ╠═6ffe2628-1ea1-11eb-24ea-57f985146a72
# ╠═70179b62-1ea1-11eb-13ce-7fb4d54bad4b
# ╠═f3fb9a0e-2294-11eb-227d-135a8b801a20
# ╠═f3fbcc54-2294-11eb-1a96-b1676ddb7772
# ╟─70182cd8-1ea1-11eb-094d-8d0d49cf15f3