# # Pandemic Control
# A stochastic optimal control problem were we seek to control the spread of a
# contagion with uncertain parameters.

# ## Problem Statement and Model
# We wish to determine optimal social distancing policy to control spread of a
# contagion and minimize the economic impact of doing so. We'll model the spread
# of the virus through a given population using the SEIR Model which considers 4
# population subsets that follow:
# ```math
# \text{Susceptible} \rightarrow \text{Exposed} \rightarrow \text{Infectious} \rightarrow \text{Recovered}
# ```
# This model is formalized as:
# ```math
# \begin{gathered}
# \frac{ds(t)}{dt} = (u(t) - 1)\beta si(t) \\
# \frac{de(t)}{dt} = (1 - u(t))\beta si(t) - \xi e(t) \\
# \frac{di(t)}{dt} = \xi e(t) - \gamma i(t)\\
# \frac{dr(t)}{dt} = \gamma i(t) \\
# si(t) = s(t) i(t)
# \end{gathered}
# ```
# where ``s(t)`` is the susceptible population, ``e(t)`` is the exposed population,
# ``i(t)`` is the infectious population, ``r(t)`` is the recovered population,
# and ``u(t) \in [0, 1]`` is the enforced population isolation (social distancing).
# The other values denote model parameters that will be specific to the contagion
# in question.

# For our case study, we'll consider the incubation constant ``\xi`` to be an
# uncertain parameter ``\xi \sim \mathcal{U}(\underline{\xi}, \overline{\xi})``.
# This introduces to infinite dependencies into our model: time ``t`` and
# uncertain incubation ``\xi``. Moreover we'll seek to minimize the isolation
# measures ``u(t)`` that are implemented while limiting the amount of infected
# individuals ``i(t)`` to be below a threshold ``i_{max}``. Thus, the optimization
# problem becomes:
# ```math
# \begin{aligned}
# &&\min_{} &&& \int_{t \in \mathcal{D}_{t}} u(t) dt \\
# && \text{s.t.} &&& \frac{\partial s(t, \xi)}{\partial t} = (u(t) - 1)\beta si(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&&&& \frac{\partial e(t, \xi)}{\partial t} = (1 - u(t))\beta si(t, \xi) - \xi e(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&&&& \frac{\partial i(t, \xi)}{\partial t} = \xi e(t, \xi) - \gamma i(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&&&& \frac{\partial r(t, \xi)}{\partial t} = \gamma i(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&&&& si(t, \xi) = s(t, \xi) i(t, \xi), && \forall \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&&&& s(0, \xi) = s_0, e(0, \xi) = e_0, i(0, \xi) = i_0, r(0, \xi) = r_0, && \forall \xi \in \mathcal{D}_{\xi} \\
# &&&&& i(t, \xi) \leq i_{max}, && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&&&& u(t) \in [0, 0.8] \\
# &&&&& \xi \sim \mathcal{U}(\underline{\xi}, \overline{\xi})
# \end{aligned}
# ```
# Notice that the SEIR model variables now all depend on both ``t`` and ``\xi``,
# except ``u(t)`` which means we need to decide our policy before knowing the
# true value of ``\xi``.

# ## Modeling in InfiniteOpt
# Traditional modeling frameworks like JuMP cannot be used to solve this model
# directly because it is infinite, contains partial difference equations, and
# contains a time expectation. We would first have to transform it into a
# transcripted (discretized) variant by applying all the necessary techniques
# (e.g., orthogonal collocation over finite elements, trapezoid rule, etc.)
# which is very combersome and nontrivial in this case.

# However, we can directly model the above form in InfiniteOpt and it will take
# care of the rest! Let's get started by importing the needed packages and
# specifying the model parameters that we'll need.

using InfiniteOpt, Ipopt, Distributions, Plots

## Set the SEIR parameters
γ = 0.303
β = 0.727
N = 1e5
ξ_min = 0.1
ξ_max = 0.6

## Set the domain information
i_max = 0.02
ϵ = 0.005
t0 = 0
tf = 200
extra_ts = [0.001, 0.002, 0.004, 0.008, 0.02, 0.04, 0.08, 0.2, 0.4, 0.8]
num_samples = 5

## Set the intial condition values
e0 = 1 / N
i0 = 0
r0 = 0
s0 = 1 - 1 / N;

# ## Model Initialization
# Now let's setup our infinite model and select Ipopt as our optimizer that will
# be used to solve it. This is accomplished making an [`InfiniteModel`](@ref):
model = InfiniteModel(Ipopt.Optimizer);

# ## Infinite Parameter Definition
# We now need to define the infinite parameters ``t \in [t_0, t_f]`` and
# ``\xi \sim \mathcal{U}(\underline{\xi}, \overline{\xi})``. This is accomplished
# with [`@infinite_parameter`](@ref). We'll also include the following
# specifications:
# - use 51 equidistant time points
# - specify that orgothonal collocation using 2 nodes should be used for time derivatives
# - specify that the number of random scenarios should equal `num_samples`
# - add `extra_ts` as extra time points
@infinite_parameter(model, t ∈ [t0, tf], num_supports = 51,
                    derivative_method = OrthogonalCollocation(2))
@infinite_parameter(model, ξ ~ Uniform(ξ_min, ξ_max), num_supports = num_samples)
add_supports(t, extra_ts)

# ## Infinite Variable Definition
# With our infinite parameters defined, we can now define our infinite variables:
# - ``s(t, \xi) \geq 0``
# - ``e(t, \xi) \geq 0``
# - ``i(t, \xi) \geq 0``
# - ``r(t, \xi) \geq 0``
# - ``si(t, \xi)``
# - ``0 \leq u(t) \leq 0.8``
@variable(model, s ≥ 0, Infinite(t, ξ))
@variable(model, e ≥ 0, Infinite(t, ξ))
@variable(model, i ≥ 0, Infinite(t, ξ))
@variable(model, r ≥ 0, Infinite(t, ξ))
@variable(model, si, Infinite(t, ξ))
@variable(model, 0 ≤ u ≤ 0.8, Infinite(t), start = 0.2)

# ## Objective Definition
# Now its time to add the objective
# ``\min \ \int_{t \in \mathcal{D}_{t}} u(t) dt`` using `@objective`:
@objective(model, Min, ∫(u, t))

# ## Constraint Definition
# The last step now to defining our model is to define the constraints using
# `@constraint`. This will involve defining the initial conditions:
# ```math
# s(0, \xi) = s_0, e(0, \xi) = e_0, i(0, \xi) = i_0, r(0, \xi) = r_0, \ \forall \xi \in \mathcal{D}_{\xi}
# ```
# the model equations:
# ```math
# \begin{aligned}
# &&& \frac{\partial s(t, \xi)}{\partial t} = (u(t) - 1)\beta si(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&& \frac{\partial e(t, \xi)}{\partial t} = (1 - u(t))\beta si(t, \xi) - \xi e(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&& \frac{\partial i(t, \xi)}{\partial t} = \xi e(t, \xi) - \gamma i(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&& \frac{\partial r(t, \xi)}{\partial t} = \gamma i(t, \xi), && \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi} \\
# &&& si(t, \xi) = s(t, \xi) i(t, \xi), && \forall \forall t \in \mathcal{D}_{t}, \xi \in \mathcal{D}_{\xi}, \\
# \end{aligned}
# ```
# and the infection limit constraint:
# ```math
# i(t, \xi) \leq i_{max}, \ \forall t \in \mathcal{D}_t, \xi \in \mathcal{D}_{\xi}.
# ```
## Define the initial conditions
@constraint(model, s == s0, DomainRestrictions(t => 0))
@constraint(model, e == e0, DomainRestrictions(t => 0))
@constraint(model, i == i0, DomainRestrictions(t => 0))
@constraint(model, r == r0, DomainRestrictions(t => 0))

## Define the SEIR equations
@constraint(model, s_constr, ∂(s, t) == -(1 - u) * β * si)
@constraint(model, e_constr, ∂(e, t) == (1 - u) * β * si - ξ * e)
@constraint(model, i_constr, ∂(i, t) == ξ * e - γ * i)
@constraint(model, r_constr, ∂(r, t) == γ * i)
@constraint(model, si == s * i)

## Define the infection rate limit
@constraint(model, imax_constr, i ≤ i_max)

# ## Display the Infinite Model
# Let's display `model` now that it is fully defined:
print(model)

# ## Optimize the Model
# Let's solve our model and find the optimal policy. All we have to do is
# invoke `optimize!` and the model will automatically be transformed solved behind
# the scenes:
optimize!(model)

# ## Retrieve and Plot the Results
# Now we can retrieve the optimal results and plot them to visualize the optimal
# policy. Note that the values of infinite variables will be returned as arrays
# corresponding to how the supports were used to discretize our model. We can
# retrieve our values using `value`.

# Get the results:
r_opt = value(r, ndarray = true) * 100 # make the population fractions into percentages
s_opt = value(s, ndarray = true) * 100
i_opt = value(i, ndarray = true) * 100
e_opt = value(e, ndarray = true) * 100
u_opt = value(u)
obj_opt = objective_value(model)
ts = value(t)
ξs = value(ξ);

# Plot the values of ``r(t, \xi)`` and ``s(t, \xi)`` over time with
# confidence bounds:
r_mean = mean(r_opt, dims = 2)
r_std = std(r_opt, dims = 2)
plot(ts, r_mean, label = "r(t, ξ)", linecolor = :red, background_color = :transparent)
plot!(ts, r_mean + r_std, linecolor = :red, linestyle = :dash, linealpha = 0.4, label = "")
plot!(ts, r_mean - r_std, linecolor = :red, linestyle = :dash, linealpha = 0.4, label = "")

s_mean = mean(s_opt, dims = 2)
s_std = std(s_opt, dims = 2)
plot!(ts, s_mean, label = "s(t, ξ)", linecolor = :blue)
plot!(ts, s_mean + s_std, linecolor = :blue, linestyle = :dash, linealpha = 0.4, label = "")
plot!(ts, s_mean - s_std, linecolor = :blue, linestyle = :dash, linealpha = 0.4, label = "")
ylabel!("Pop. (%)")
xlabel!("Time (Days)")

# Plot the values of ``i(t, \xi)`` and ``e(t, \xi)`` over time with
# confidence bounds:
i_mean = mean(i_opt, dims = 2)
i_std = std(i_opt, dims = 2)
plot(ts, i_mean, label = "i(t, ξ)", linecolor = :green, background_color = :transparent)
plot!(ts, i_mean + i_std, linecolor = :green, linestyle = :dash, linealpha = 0.4, label = "")
plot!(ts, i_mean - i_std, linecolor = :green, linestyle = :dash, linealpha = 0.4, label = "")

e_mean = mean(e_opt, dims = 2)
e_std = std(e_opt, dims = 2)
plot!(ts, e_mean, label = "e(t, ξ)", linecolor = :purple)
plot!(ts, e_mean + e_std, linecolor = :purple, linestyle = :dash, linealpha = 0.4, label = "")
plot!(ts, e_mean - e_std, linecolor = :purple, linestyle = :dash, linealpha = 0.4, label = "")
ylabel!("Pop. (%)")
xlabel!("Time (Days)")

# Plot the values of ``u(t)`` over time:
plot(ts, u_opt, linecolor = :orange, label = "u(t)", ylims = (-0.02, 1.02),
     background_color = :transparent)
xlabel!("Time (Days)")
ylabel!("Distancing Ratio")

# ### Maintenance Tests
# These are here to ensure this example stays up to date.
using Test
@test termination_status(model) == MOI.LOCALLY_SOLVED
@test has_values(model)
@test u_opt isa Vector{<:Real}


using LinearAlgebra: pinv
using DSP: conv

export savitzky_golay_filter

"""
	savitzky_golay_filter(y::AbstractVector, window_size::Integer, polynomial_order::Integer; deriv_order::Integer = 0, boundary_mode = :interpolation)

Apply Savitzky-Golay polynomial smoothing to input data `y` using a polynomial of order `polynomial_order` fit to a moving window `window_size` points wide. Optionally derivatives can be taken by specifying the `deriv_order` and the caller is responsible for the appropriate scaling by the point spacing. Handling of data within half the window size is specified by `boundary_mode`. When set to `:interpolation` the polynomial fit will be used; when set to `:nearest` the data will be padded using the edge values before convolution and the valid portion will then be returned.

# References
1. Savitzky, A., & Golay, M. J. E. (1964). Smoothing and Differentiation of Data by Simplified Least Squares Procedures. Analytical Chemistry, 36(8), 1627–1639. https://doi.org/10.1021/ac60214a047
2. Steinier, J., Termonia, Y., & Deltour, J. (1972). Comments on Smoothing and differentiation of data by simplified least square procedure. Analytical Chemistry, 44(11), 1906–1909. https://doi.org/10.1021/ac60319a045
3. Press, W. H., & Teukolsky, S. A. (1990). Savitzky-Golay Smoothing Filters. Computers in Physics, 4(6), 669. https://doi.org/10.1063/1.4822961
"""
function savitzky_golay_filter(y::AbstractVector, window_size::Integer, polynomial_order::Integer; deriv_order::Integer = 0, boundary_mode = :interpolation)

	# input validity checks
   	@assert isodd(window_size) "Window size must be an odd integer, i.e. fitting 2m + 1 points around the current value."
	@assert polynomial_order < window_size "Polynomial order must be less than the window size."
	@assert boundary_mode in (:interpolation, :nearest) "boundary_mode must be one of :interpolation, :nearest"

	# window size is 2m + 1 points
   	m = (window_size - 1) ÷ 2

	# build the Vandermonde design matrix A. Each row corresponds to a point in the fitting window -m:m
	# and each columns correspond to powers in the range 0:polynomial_order
   	fitting_points = -m:m
   	A = Matrix{Float64}(undef, window_size, polynomial_order + 1)
   	for i in 1:window_size, j in 1:polynomial_order + 1
        A[i,j] = fitting_points[i]^(j - 1)
    end

	if boundary_mode == :interpolation
		# for interpolation we'll want the full pseudo-inverse so we can calculate all the fit values at the edges
		# Ap = y
		C = pinv(A)

		# the filter coefficients are the rows of `C`
		filter_coeffs = C[deriv_order + 1,:] * factorial(deriv_order)

		# convolve with the filter coefficients with a couple extra steps:
		# 1. because of convolution will reverse coefficients we flip before
		# 2. c = conv(a,b) will return a vector of length(c) = length(a) + length(b) - 1 so we chop off the first and last m points
		smoothed = conv(reverse(filter_coeffs), y)[m+1:end-m]

		# for interpolation edge handling calculate the full fits
		if deriv_order == 0
			# if we are just smoothing then we can use the design and coefficient matrix as is
			AC = A*C
			smoothed[1:m] = (AC*y[1:window_size])[1:m]
			smoothed[end-m+1:end] = (AC*y[end-window_size+1:end])[end-m+1:end]
		else
			# otherwise we need to differentiate the polynomial coefficients
			# first m points
			p = C * y[1:window_size]
			for _ in 1:deriv_order
				p = [(i-1)*p[i] for i in 2:length(p)]
			end
			smoothed[1:m] = A[1:m, 1:size(A,2)-deriv_order]*p
			# last m points
			p = C * y[end-window_size+1:end]
			for _ in 1:deriv_order
				p = [(i-1)*p[i] for i in 2:length(p)]
			end
			smoothed[end-m+1:end] = A[m+2:end, 1:size(A,2)-deriv_order]*p

		end

		return smoothed

	elseif boundary_mode == :nearest

		# here we only need a single set of coefficients and so we can least-squares solve AᵀCᵀ = I for for a single row by picking a single column out of I
		Icol = zeros(Float64, polynomial_order+1, 1)
		Icol[deriv_order + 1] = 1.0
		filter_coeffs = transpose(A) \ Icol

		# pad the signal with the endpoints
		padded_y = [y[1] * ones(m); vec(y); y[end] * ones(m)]

		# convolve with filter
		smoothed = conv(filter_coeffs[end:-1:1], padded_y)

		# and return the valid midsection
		return smoothed[window_size:end-2*m]

	end

end


using KFEstimate
using LinearAlgebra, Plots, Zygote, Statistics, Revise
using Flux, Flux.Optimise
using ProgressBars
pathof(KFEstimate)

dt = 0.001
A = [1.0 dt 1/2*dt^2; 0.0 1.0 dt; 0.0 0.0 1.0]
B = [0.0; 0.0; 1.0]
B = reshape(B, length(B), 1)
Q = 1*Matrix{Float64}(I, 3, 3)

# observation model, assume we can noisily measure position
H = randn(3, 3)
# H = [1 0 0; 0 1 0; 0 0 1]
R = 10.0*Matrix{Float64}(I, 3, 3)

kf = KalmanFilter(A, B, Q, H, R)

# run simulation
time_step = 0:1000
x0 = [0.0; 0.0; 0.0]
action_sequence = [[1] for t in time_step]
sim_states, sim_measurements = run_simulation(kf, x0, action_sequence)
# run kalman filter
P0 = Matrix{Float64}(I, 3, 3)
s0 = State(x0, P0) # initial state belief
filtered_states = run_filter(kf, s0, action_sequence, sim_measurements)

# unpack sim and filtered states
μ, Σ = unpack(filtered_states)
p = [x[1] for x in sim_states]
v = [x[2] for x in sim_states]
a = [x[3] for x in sim_states]
x = hcat(p, v, a)
# plot
plot(time_step, [x[2:end, 1] x[2:end, 2] x[2:end, 3]], label = ["simulated p" "simulated v" "simulated a"])
plot!(time_step, [μ[2:end, 1] μ[2:end, 2] μ[2:end, 3]], label = ["filtred p" "filtered v" "filtered a"])
xlabel!("time step (t)")

##

# parametrized matrix estimates
Ahat(θ) = [θ[1] dt 1/2*dt^2; 0.0 θ[1] dt; 0.0 0.0 θ[1]]
Bhat(θ) = B
Qhat(θ) = Q
Hhat(θ) = H
Rhat(θ) = R
# define a parametrized kalman filter
param_kf = ParamKalmanFilter(Ahat, Bhat, Qhat, Hhat, Rhat)

θ0 = [1.1]
pkf = ParamKalmanFilter(Ahat, Bhat, Qhat, Hhat, Rhat)

loss = []
θ_range = 0.9:0.001:1.1
for i in θ_range
    θ_i = [i]
    states = run_param_kf(θ_i, pkf, s0, action_sequence, sim_measurements)
    l = kf_likelihood(θ_i, param_kf, states, action_sequence, sim_measurements)
    push!(loss, l)
end

plot(θ_range, loss)

##

# parametrized matrix estimates
Ahat(θ) = [θ[1] θ[2] θ[3]; 0 θ[1] θ[2]; 0 0 θ[1]]
# Ahat(θ) = [θ[1, 1] θ[1, 2] θ[1, 3]; 0.0 θ[2, 2] θ[2, 3]; 0.0 0.0 θ[3, 3]]
Bhat(θ) = B
Qhat(θ) = Q
Hhat(θ) = H
Rhat(θ) = R
# define a parametrized kalman filter
param_kf = ParamKalmanFilter(Ahat, Bhat, Qhat, Hhat, Rhat)

θ0 = [1.0 dt 1/2*dt^2] + 0.1*randn(1, 3)
opt = ADAM(0.01)
epochs = 200
newθ, loss = run_kf_gradient(θ0, param_kf, s0, action_sequence, sim_measurements, opt, epochs)

grad_states = run_param_filter(newθ, param_kf, s0, action_sequence, sim_measurements)
μgrad, Σgrad = unpack(grad_states)

l = @layout [a{0.7h};grid(1, 3)]
p1 = plot(time_step, [x[2:end, 1] x[2:end, 2] x[2:end, 3]], label = ["simulated p" "simulated v" "simulated a"], xlabel="time step (t)")
p1 = plot!(time_step, [μ[2:end, 1] μ[2:end, 2] μ[2:end, 3]], label = ["filtered p" "filtered v" "filtered a"], xlabel="time step (t)")
p1 = plot!(time_step, [μgrad[2:end, 1] μgrad[2:end, 2] μgrad[2:end, 3]], label = ["learned p" "learned v" "learned a"],  xlabel="time step (t)")
p2 = plot(1:epochs, loss, title="loss", xlabel="number of epochs")
p3 = plot(time_step, (x[2:end, :]-μ[2:end, :]).^2, title="KF vs. sim error", xlabel="time step (t)")
p4 = plot(time_step, (x[2:end, :]-μgrad[2:end, :]).^2, title="grad vs. sim error", xlabel="time step (t)")
plot(p1, p2, p3, p4, layout=l, titlefont = font(12), size=(1000, 700))


for f in [:chebyshev_basis, :sin_basis, :cos_basis, :fourier_basis, :polynomial_basis, :monomial_basis]
    @eval $f(x, c) = $f(scalarize(x), c)
end
function _generateBasis!(eqs, f, x, coeffs)
    n_x = size(x, 1)
    @assert length(eqs) == size(x, 1)*length(coeffs)
    @inbounds for (i, ti) in enumerate(coeffs)
        eqs[(i-1)*n_x+1:i*n_x] .= f(x, ti)
    end
    return
end

"""
$(SIGNATURES)

Constructs an array containing a Chebyshev basis in the variables `x` with coefficients `c`.
If `c` is an `Int` returns all coefficients from 1 to `c`.
"""
function chebyshev_basis(x::Array, coefficients::AbstractVector)
    eqs = Array{Num}(undef, size(x, 1)*length(coefficients))
    f(x, t) = cos.(t .* acos.(x))
    _generateBasis!(eqs, f, x, coefficients)
    eqs
end

chebyshev_basis(x::Array, terms::Int) = chebyshev_basis(x, 1:terms)

"""
$(SIGNATURES)

Constructs an array containing a Sine basis in the variables `x` with coefficients `c`.
If `c` is an `Int` returns all coefficients from 1 to `c`.
"""
function sin_basis(x::Array, coefficients::AbstractVector)
    eqs = Array{Num}(undef, size(x, 1)*length(coefficients))
    f(x, t) = sin.(t .* x)
    _generateBasis!(eqs, f, x, coefficients)
    eqs
end

sin_basis(x::Array, terms::Int) = sin_basis(x, 1:terms)

"""
$(SIGNATURES)

Constructs an array containing a Cosine basis in the variables `x` with coefficients `c`.
If `c` is an `Int` returns all coefficients from 1 to `c`.
"""
function cos_basis(x::Array, coefficients::AbstractVector)
    eqs = Array{Num}(undef, size(x, 1)*length(coefficients))
    f(x, t) = cos.(t .* x)
    _generateBasis!(eqs, f, x, coefficients)
    eqs
end

cos_basis(x::Array, terms::Int) = cos_basis(x, 1:terms)

"""
$(SIGNATURES)

Constructs an array containing a Fourier basis in the variables `x` with (integer) coefficients `c`.
If `c` is an `Int` returns all coefficients from 1 to `c`.
"""
function fourier_basis(x::Array, coefficients::AbstractVector{Int})
    eqs = Array{Num}(undef, size(x, 1)*length(coefficients))
    f(x, t) = iseven(t) ? cos.(t .* x ./ 2) : sin.(t .* x ./2)
    _generateBasis!(eqs, f, x, coefficients)
    eqs
end

fourier_basis(x::Array, terms::Int) = fourier_basis(x, 1:terms)

"""
$(SIGNATURES)

Constructs an array containing a polynomial basis in the variables `x` up to degree `c` of the form
`[x₁, x₂, x₃, ..., x₁^1 * x₂^(c-1)]`. Mixed terms are included.
"""
function polynomial_basis(x::Array, degree::Int = 1)
    @assert degree > 0
    n_x = length(x)
    n_c = binomial(n_x+degree, degree)
    eqs = Array{Num}(undef, n_c)
    _check_degree(x) = sum(x)<=degree ? true : false
    itr = Base.Iterators.product([0:degree for i in 1:n_x]...)
    itr_ = Base.Iterators.Stateful(Base.Iterators.filter(_check_degree, itr))
    filled = false
    @inbounds for i in 1:n_c
        eqs[i] = 1
        filled = true
        for (xi, ci) in zip(x, popfirst!(itr_))
            if !iszero(ci)
                filled ? eqs[i] = xi^ci : eqs[i] *= xi^ci
                filled = false
            end
        end
    end
    eqs
end

"""
$(SIGNATURES)

Constructs an array containing monomial basis in the variables `x` up to degree `c` of the form
`[x₁, x₁^2, ... , x₁^c, x₂, x₂^2, ...]`.
"""
function monomial_basis(x::AbstractArray, degree::Int = 1)
    @assert degree > 0
    n_x = length(x)
    exponents = 1:degree
    n_e = length(exponents)
    n_c = n_x * n_e + 1
    eqs = Array{Num}(undef, n_c)
    eqs[1] = Num(1)
    idx = 0
    for i in 1:n_x, j in 1:n_e
        idx = (i-1)*n_e+j+1
        eqs[idx] = x[i]^exponents[j]
    end
    eqs
end


# export fourier_nodes, fourier_wavenumbers, cheb
using LinearAlgebra, FFTW

"""
fourier_nodes(n; a = 0, b = 2π)
# Description
- Create a uniform grid of points for periodic functions
# Arguments
- `N`: integer | number of evenly spaced points
# Keyword Arguments
- `a`: number | starting point of interval [a, b)
- `b`: number | ending point of interval [a, b)
# Return
- `g`: array | an array of points of evenly spaced points from [a, b)
"""
function fourier_nodes(N; a = 0, b = 2π)
    return (b - a) .* collect(0:(N - 1)) / N .+ a
end

"""
fourier_wavenumbers(N; L = 2π)
# Description
- Create wavenumbers associated with the domain of length L
# Arguments
- `N`: integer | number of wavevectors
# Keyword Arguments
- `L`: number | length of interval [a, b), L = b-a
# Return
- `wavenumbers`: array | an array of wavevectors
"""
function fourier_wavenumbers(N; L = 2π)
    up = collect(0:1:(N - 1))
    down = collect((-N):1:-1)
    indices = up
    indices[(div(N, 2) + 1):end] = down[(div(N, 2) + 1):end]
    wavenumbers = 2π / L .* indices
    return wavenumbers
end

"""
# Description
Spectral differentiation matrix and nodes for periodic domains
# Argument
- 'N': number of gridpoints
# Keyword Argument
a and b specify the interval [a, b)
default assumes x ∈ [0, 2π)
# Return
- 'D': Fourier differentiation matrix
- 'x': Fourier grid points
"""
function fourier(N; a = 0, b = 2π)
    if N == 0
        return [0], [0]
    else
        k = fourier_wavenumbers(N, L = b - a)
        x = fourier_nodes(N, a = a, b = b)
        ℱ = fft(I + zeros(N, N), 1)
        ℱ⁻¹ = ifft(I + zeros(N, N), 1)
        D = real.(ℱ⁻¹ * Diagonal(im .* k) * ℱ)
        return D, x
    end
end

"""
# Description
Julia version of Spectral Methods in Matlab
# Argument
- 'N': polynomial order
# Keyword Argument
a and b specify the interval [a, b]
default assumes x ∈ [-1, 1]
# Return
- 'D': Chebyshev differentiation matrix
- 'x': Guass-Lobatto points
"""
function chebyshev(N; a = -1, b = 1)
    if N == 0
        return [0], [(a + b) / 2]
    else
        x = @. cos(pi * (0:N) / N)
        c = [2; ones(N - 1); 2] .* (-1) .^ (0:N)
        dX = x .- x'
        D = (c ./ c') ./ (dX + I)                # off-diagonal entries
        D = D - Diagonal(sum(D', dims = 1)[1:(N + 1)]) # diagonal entries
        return 2 / (b - a) * D, (b - a) * (x .+ 1) / 2 .+ a
    end
end


# This file includes the MonteCarlo simulation of FSK modulation scheme and compares
# the numerical results with the theoretical results.

using DigitalCommunications
using Plots

# Settings
pulse = RectangularPulse()          # Modulation pulse
fs = 10                             # Samling frequency
ts = 1 / fs                         # Sampling period
tb = pulse.duration                 # Pulse duration
k = 3                               # Bits per symbol
M = 2^k                             # Costellation size
nsymbols = Int(1e6)                 # Number of symbols
ebno = collect(-2 : 8)              # Energy per bit to noise power spectral density ratio
esno = ebno .+ 10 * log10(k)        # Energy per symbol to noise power spectral densit ratio.

# Communcation system components
gen = SymbolGenerator(nsymbols, M)
modulator = Modulator(FSK(M), RectangularPulse(), ts)
channel = AWGNChannel(1., ts, tb)
detector = Detector(modulator(1:M))

# Monte Carlo simulation
message = gen.symbols
symerr = zeros(length(esno))
for i in 1 : length(symerr)
    channel.esno = esno[i]  # Update channel noise level
    mbar = message |> modulator |> channel |> detector  # Run communication system
    symerr[i] = sum(mbar .!= message) / length(message)  # Symbol error rate
end

# Plots
plt = plot(title="$M-PSK", xlabel="ebno [dB]", ylabel="Pe")
plot!(ebno, berfsk.(esno, M), marker=:circle, yscale=:log10, label="theoretical")
plot!(ebno, symerr, marker=:circle, yscale=:log10, label="montecarlo")


#===============================================================================

The most naive way of computing n15 requires fourteen multiplications:

n × n × ... × n = n15

But using a "binary" method you can compute it in six multiplications:

n × n = n^2
n^2 × n^2 = n^4
n^4 × n^4 = n^8
n^8 × n^4 = n^12
n^12 × n^2 = n^14
n^14 × n = n^15

However it is yet possible to compute it in only five multiplications:

n × n = n^2
n^2 × n = n^3
n^3 × n^3 = n^6
n^6 × n^6 = n^12
n^12 × n^3 = n^15

We shall define m(k) to be the minimum number of multiplications to compute n^k;
for example m(15) = 5.

For 1 ≤ k ≤ 200, find ∑ m(k).

--------------------------------------------------------------------------------

To start with, this problem can be reformulated by asking how many additions are
necessary to reach a number starting from 1 and using only additions.

The general approach is bottom up. At every iteration we build a new set of sets
of numbers. An inner set contains all the numbers that can be together obtained
by 'iter' calculations.

An inner set is generated from a previous inner set by taking a copy of the
later, then by selecting one of its elements, by adding it to the initial
elements and adding the numbers so obtained to the new set. This must be done
for all elements of the initial set.

At the next iteration, it is only necessary to consider the new elements that
have been added when selecting the element to add to the other.

Because a same inner set can be obtained from different paths, each new inner
set is associated with the numbers selected to produce it.

Each time a new number is generated, the corresponding iteration is recorded
with it.

===============================================================================#

using Base.Test

function compute_nbr_steps(limit)
    nbr_steps = Array{Int}(limit)
    # nbr_steps[i] is the number of addition needed to compute i (i.e. it is
    # m(i)).

    fill!(nbr_steps, -1)
    sets = Dict{Set{Int}, Set{Int}}()
    # Key is all numbers in the set, Value is the set of the numbers added
    # during the previous iteration.

    const empty_value = Set{Int}()
    # empty_value is used as an indicator

    completion = 1
    iter = 0
    nbr_steps[1] = 0
    sets[Set{Int}(1)] = Set{Int}(1)

    while completion<limit
        iter += 1
        updated_sets = typeof(sets)()

        for (set, added_numbers) in sets, l in set, m in added_numbers
            n=l+m
            if n ≤ limit && n ∉ set
                updated_set = copy(set)
                push!(updated_set, n)

                nums = get(updated_sets, updated_set, empty_value)

                if nums == empty_value
                    nums = Set{Int}(n)
                else
                    push!(nums, n)
                end

                updated_sets[updated_set] = nums

                if nbr_steps[n]<0
                    nbr_steps[n] = iter
                    completion += 1
                    println("\t completion = $(completion)")
                    if completion == limit
                        return nbr_steps
                    end
                end
            end
        end
        sets = updated_sets
    end
end

result = compute_nbr_steps(200)

@test result[15] == 5
@test length(result) == 200

print("Euler 122: $(sum(result))")


using LinearAlgebra

@doc """
Gaussian kernel function.

# Examples
```jldoctest
y = k(x)
```
where the 'x' is Vector{Float64} sample and 'y' is density in point 'x'.
""" ->
k(x::Vector{Float64}) = √(2*π)*ℯ^(-(x'*x)/2)

@doc """
Function to return Parzen window estimation.

# Examples
```jldoctest
fun = Σ(h, X, f, x)
```
where the 'fun' is parzen window estimation, 'h' is window size,
'X' is set of train samples (size(dim, N) ~ (dimension of data, data count),
'f' is kernel function and 'x' is given Vector{Float64} sample.
""" ->
create_parzen_window(h::Float64, X::Matrix{Float64},
                        f, x::Vector{Float64}) = (1/(h*size(X)[1]))*mapreduce(a->f((x-a)/h), +, eachcol(X))


export choice_probabilities, choice_probabilities!

"""
Compute the choice probabilities using an overflow safe algorithm.

This is just a wrapper around [`choice_probabilities!`](@ref) which computes them in-place.
"""
function choice_probabilities(utilities::AbstractMatrix{<:AbstractFloat})
    probabilities = similar(utilities)
    choice_probabilities!(probabilities, utilities)
    return probabilities
end

"""
Compute the choice probabilities in-place using an overflow safe algorithm.

Based on the SoftmaxThreePassReload function from
Dukhan and Ablavatski (2020), The Two-Pass Softmax Algorithm (https://arxiv.org/abs/2001.04438).
"""
function choice_probabilities!(output::AbstractMatrix{T}, utilities::AbstractMatrix{T}) where T<:AbstractFloat
    @assert size(output) == size(utilities) "Output matrix must have same size as utilities matrix."
    J, I = size(utilities)

    @inbounds for i = 1:I
        max_u = T(0)
        for j = 1:J
            if utilities[j, i] > max_u
                max_u = utilities[j, i]
            end
        end

        sigma = @fastmath exp(-max_u)

        @simd ivdep for j = 1:J
            output[j, i] = @fastmath exp(utilities[j, i] - max_u)
            # SIMD makes left-to-right summation both faster and more accurate
            # https://discourse.julialang.org/t/when-shouldnt-we-use-simd/18276/14
            sigma += output[j, i]
        end

        sigma_inv = 1 / sigma
        @simd ivdep for j = 1:J
            output[j, i] *= sigma_inv
        end
    end
end


"""
    expected_shortfall(returns, confidence, method; multiplier=1.0)

Computes the Expected Shortfall (ES), also known as Conditional Value-at-Risk (CVaR), Average Value-at-Risk (AVaR) or Expected Tail Loss (ETL). The ES is the expected return on the asset in the worst `α%` of cases, therefore quantifies the tail-risk of an asset. It is calculated by averaging all of the returns in the distribution that are worse than the VaR of the portfolio at a given significance level `α`. For instance, for a 5% significance level, the expected shortfall is calculated by taking the average of returns in the worst 5% of cases.

Expected Shortfall puts emphasis on the tail of the loss distribution, whereas Value-at-risk neglects this aspect.

# Arguments
- `returns`:        Vector of asset returns.
- `α`:              Significance level, e.g. use `0.05` for 95% confidence, or `0.01` for 99% confidence.
- `method`:         Distribution estimation method: `:historical`, `:gaussian` or `:cornish_fisher`.
- `multiplier`:     Optional scalar multiplier, i.e. use `12` to annualize monthly returns, and use `252` to annualize daily returns.

# Methods
- `:historical`:        Historical based on empirical distribution of returns.
- `:gaussian`:          Gaussian distribution based on parametric fit (mean, variance).
- `:cornish_fisher`:    Cornish-Fisher based on Gaussian parametric distribution fit adjusted for third and fourth moments (skewness, kurtosis). Cornish-Fisher expansion aims to approximate the quantile of a true distribution by using higher moments (skewness and kurtosis) of that distribution to adjust for its non-normality. See https://thema.u-cergy.fr/IMG/pdf/2017-21.pdf for details.

# Sources
- Amédée-Manesme, Charles-Olivier and Barthélémy, Fabrice and Maillard, Didier (2017). Computation of the Corrected Cornish–Fisher Expansion using the Response Surface Methodology: Application to VaR and CVaR. THEMA Working Paper n°2017-21, Université de Cergy-Pontoise, France.
"""
function expected_shortfall(returns, α, method::Symbol; multiplier=1.0)
    if method == :historical
        # average return below significance level (quantile)
        sorted = sort(returns)
        idx = floor(Int64, length(sorted) * α)
        return mean(sorted[1:idx]) * sqrt(multiplier)
    elseif method == :gaussian
        # derivation: http://blog.smaga.ch/expected-shortfall-closed-form-for-normal-distribution/
        q = quantile(Normal(), α)
        μ = mean(returns)
        σ = std(returns; corrected=false)
        return (μ - σ*pdf(Normal(), q)/α) * sqrt(multiplier)
    elseif method == :cornish_fisher
        # third/fourth moment adjusted Gaussian distribution fit
        # https://papers.ssrn.com/sol3/papers.cfm?abstract_id=1024151
        q = quantile(Normal(), α)
        S = skewness(returns)
        K = kurtosis(returns; method=:excess)
        g = q + 1/6*(q^2-1)S + 1/24*(q^3-3q)*K - 1/36*(2q^3-5q)*S^2
        ϕ = pdf(Normal(), g)
        EG2 = -1/α*ϕ * (1 + 1/6*(g^3)*S + 1/72*(g^6 - 9g^4 + 9g^2 + 3)*S^2 + 1/24*(g^4 - 2g^2 - 1)*K)
        μ = mean(returns)
        σ = std(returns; corrected=false)
        return (μ + σ*EG2) * sqrt(multiplier)
    end

    throw(ArgumentError("Passed method parameter '$(method)' is invalid, must be one of :historical, :gaussian, :cornish_fisher."))
end


function project(vec1::Array{Float64}, vec2::Array{Float64})::Array{Float64}
    return (dot(vec1,vec2)/dot(vec2,vec2)) * vec2
end

function unitvec(vec::Array{Float64})::Array{Float64}
    return (1/norm(vec)) * vec
end

function gram_schmidt(vecs...)
    u = [unitvec(vecs[1])]

    for i in vecs[2:end]
        new_u = i
        for j in u
            p = project(j, i)
            new_u -= p
        end
        push!(u, unitvec(new_u))
    end

    return u
end


module DH32StatsHelpers

# package code goes here
export cMDS,
Cohen_d,
DataFrame_to_distmat,
zscore_log10_cols

#cMDS_of_DataFrame,
#kmeans_of_cMDS,

using DataFrames
using Distances
using Distributions
using MultivariateStats
using Clustering

using Distributions
using DataFrames

"""
Compute effect size measure Cohen's d

Input:
  - float array x1
  - float array x2
  - confidence level (default 0.95)

Output:
  - DataFrame with 1 row and 3 columns:
  d, lower_ci, upper_ci
"""
function Cohen_d(x1::Array{Float64,1}, x2::Array{Float64,1}, conf_level::Float64=0.95)

    n1 = length(x1)
    n2 = length(x2)
    dof = n1 + n2 - 2.
    d = (mean(x1) - mean(x2))/sqrt((((n1-1.)*var(x1)+(n2-1.)*var(x2))/(n1+n2-2.)))
    S_d = sqrt(((n1 + n2)/(n1 * n2) + 0.5 * d^2/dof) * ((n1 + n2)/dof))
    Z = -quantile(TDist(dof), (1.0-conf_level)/2.0)
    lower_ci = d - Z * S_d
    upper_ci = d + Z * S_d

    return DataFrame(d=d, lower_ci=lower_ci, upper_ci=upper_ci)
end

"""
Modify columns of a data frame of count data or other positive data by:

- Replacing 0 by 1 (to avoid NA after log)
- log10-transformation
- z-scores of log10-transformed

Output: transformed data frame
"""
function zscore_log10_cols(df::DataFrame)

    df_new = copy(df)

    for i in 1:length(df)

        #replace all zeros by 1 to avoid log(0)
        df_new[df_new[:,i] .== 0,i] = 1

        #take the log10 of all columns
        df_new[:,i] = map(log10,df[:,i])

        #take z-scores
        df_new[:,i] = zscore(convert(Array{Float64},df_new[:,i]))
    end

    return df_new
end

"""
Translates a data frame into a distance matrix.

Input:

   - data frame

   - distance (from Distances.jl), e.g. HellingerDist()

   - row_instances: true or false. Are rows of data frame the instances for which distances should be evaluated (true), or the columns (false)

Output: distance matrix (symmetric matrix with zero diagonal)

"""
function DataFrame_to_distmat(df::DataFrame, dist::Any, row_instances::Bool=true)

    x = convert(Array, df)

    if (row_instances == true)
        return pairwise(dist, x')
    end

    return pairwise(dist, x)

end

"""
Carry out a classical multi-dimensional scaling (cMDS).

Input:

    - distances (symmetric matrix)

    - dimension of cMDS output (default: 2)

    - should eigenvalues be computed? (true or false)

Output:

    - cMDS array with instances as columns

    - if eigenvalues == true: also 1-dim array of eigenvalues

"""
function cMDS(d::Array{Float64,2}, dim::Int64=2, eigenvalues::Bool=true)

    #do classical MDS:
    cMDS = classical_mds(d, dim)

    if eigenvalues == true
        #compute eigenvalues
        G = dmat2gram(d)
        E = eigfact!(Symmetric(G))
        #return cMDS array and eigenvalue vector descending from largest
        return cMDS, sort(E.values[:,1],rev=true)
    end

    return cMDS

end

end # module


function wave_equation(m,M,y)

# evaluates the discrete matrix equation for the waves
#
# a1 b'' + a2 b' + (a3 + a4 lambda + a5 I lambda^2)b = 0
#
# where a1-a2 are tridiagonal matrices resulting from derivatives
# and a3-a5 are diagonal matrices. The eigenvalue is lambda = C
# (see eq. 37) B&M 2019. We recover the value of frequency from
# the definition of C
#
# input
#  m -  angular order of wave
#  C -  Coriolis parameter  (dimensionless)
#  M -  magnetic parameter  (dimensionless)
#  y -  meridional coordinate
#
# output
# A0,A1,A2 - matrices used in nonlinear eigenvalue problem
#
# Form of eigenvalue problem
# (A0 + A1 lambda + A2 lambda^2 ) b = 0
#

# size of matrix system
n = length(y);

# matrices defining first and second derivatives
d1,d2 = derivatives(y);

# assemble matrix a1 (second derivative)
y2 = y.*y;
y2c = -(y2.-1);
a1 = Diagonal(y2c) * d2;

# assemble matrix a2   (first derivative)
a2 = -2*Diagonal(y) * d1;

# assemble matrix a3   (constant term)
y2c_inv = 1.0./y2c;
a3 = -m^2*Diagonal(y2c_inv);

# now combine into A0, A1 and A2
A0 = convert(Array{Complex{Float64},2},a1+a2+a3);
A1 = convert(Array{Complex{Float64},2},m*Diagonal(ones(n))/M);
A2 = convert(Array{Complex{Float64},2},Diagonal(y2)/M);

return A0,A1,A2

end


"""
```jldoctest
julia> floor(ArithematicAnalysis.bisection(x->x^4-4x^3+x+4,0,3);digits=2)
1.23
```
"""
function bisection(custom_func,starting::Real,ending::Real)
    if custom_func(starting)==0
        return starting
    elseif custom_func(ending)==0
        return ending
    elseif  !_root_finder(starting,ending,custom_func)
        print("Cannot find roots between a and b")
        return Nothing
    else
        mid_point = _get_mid_point(starting,ending)
        while !_root_evaluate(mid_point,custom_func)
            if _root_finder(starting,mid_point,custom_func)
                ending = mid_point
            elseif _root_finder(mid_point,ending,custom_func)
                starting = mid_point
            end
            mid_point = _get_mid_point(starting,ending)
        end
        return mid_point
    end
end

function _get_mid_point(x::Real, y::Real)
    return x+(y-x)/2
end

function _root_evaluate(x::Real,func)
    value = func(x)
    abs_value = abs(value)
    return abs_value<10^-7
end

function _root_finder(x::Real, y::Real,func)
    x,y = func(x), func(y)
    return x * y < 0
end

# Should be Monotonic Function
function custom_func(x::Real)
    return x^4-4x^3+x+4
end


function prob(q)
    p = zeros(Float64, 21, 50)

    p[1, 50] = 50/q
    p[2, 50] = 1 - 50/q

    for k in 49:-1:1
        for n in 0:min(51-k, 20)
            if n == 0
                p[n+1, k] = (k/q) * p[n+1, k+1]
            elseif n + k == 51
                p[n+1, k] = (1 - k/q) * p[n, k+1]
            else
                p[n+1, k] = (k/q) * p[n+1, k+1] + (1 - k/q) * p[n, k+1]
            end
        end
    end

    return p[21, 1]
end

function main()
    l = 50
    h = 55
    e = 10^-12

    while h - l > e
        p = prob((h + l)/2)
        if p > 0.02
            l = (h + l)/2
        else
            h = (h + l)/2
        end
    end

    println(round((h + l)/2, digits=10))
end

main()


#----------------------------#
# see page 57 in Sarkka book #
#----------------------------#

function simulatedata(N = 1_000; seed = 1)

    rg = MersenneTwister(seed)

    A, H = 0.5*randn(rg, 2, 2),   0.5*randn(rg, 2, 2)
    Q, R = 0.005*Matrix(I, 2, 2), 0.05*Matrix(I, 2, 2)
    m₀   = 0.5*randn(rg, 2)
    P₀   = 0.05*Matrix(I, 2, 2)

    yclean = Array{Vector{Float64}, 1}(undef, N)
    y = Array{Vector{Float64}, 1}(undef, N)
    x = Array{Vector{Float64}, 1}(undef, N)

    x[1] = rand(rg, MvNormal(m₀, P₀))
    yclean[1] = H*x[1]
    y[1] = rand(rg, MvNormal(H * x[1], R))

    for k in 2:N

        # prediction step

        x[k] = rand(rg, MvNormal(A * x[k-1], Q))

        yclean[k] = H * x[k]

        y[k] = rand(rg, MvNormal(H * x[k], R))

    end

    return yclean, y, x, A, H, Q, R, m₀, P₀

end


#===============================================================================

A bag contains one red disc and one blue disc. In a game of chance a player
takes a disc at random and its colour is noted. After each turn the disc is
returned to the bag, an extra red disc is added, and another disc is taken at
random.

The player pays £1 to play and wins if they have taken more blue discs than red
discs at the end of the game.

If the game is played for four turns, the probability of a player winning is
exactly 11/120, and so the maximum prize fund the banker should allocate for
winning in this game would be £10 before they would expect to incur a loss. Note
that any payout will be a whole number of pounds and also includes the original
£1 paid to play the game, so in the example given the player actually wins £9.

Find the maximum prize fund that should be allocated to a single game in which
fifteen turns are played.

===============================================================================#

using Memoize
using Base.Test

#=------------------------------------------------------------------------------
Let p(b, n) the probability to take b blue disks in n attempts, and pb(n) the
probability to take a blue disk at the nth attempt.

It comes:
    p(b, n) = | pb(n) p(b-1, n-1) + (1-pb(n)) p(b, n-1) if b≤n
              | 0 if b>n

    pb(n) = 1 / (1+n)

In addition p(1, 1) = pb(1) p(0,0) + (1-pb(1)) p(0, -1)
                    = 1/2 p(0, 0)
                    = 1/2
            => p(0, 0) = 1
------------------------------------------------------------------------------=#

@memoize
function p(b::Int, n::Int)::Rational{Int}
    if     n<b  0
    elseif b<0  0
    elseif n==0 1  # In that case, b must also be equal to 0
    else        p(b-1, n-1)//(1+n) + n*p(b, n-1)//(1+n)
    end
end

#=------------------------------------------------------------------------------
The probability to win it then given by
    p(n) = sum(p(b,n) for b ∈ [div(n,2)+1, n])
------------------------------------------------------------------------------=#

p(n::Int)::Rational{Int} = sum([p(b,n) for b in div(n,2)+1:n])

@test p(4) == 11//120

#=------------------------------------------------------------------------------
If p(n) = N//D, then the maximal prize fund is div(D/N), i.e. floor(1//p(n))
------------------------------------------------------------------------------=#

maxfund(n) = let p = p(n); div(p.den, p.num) end

@test maxfund(4) == 10

#=----------------------------------------------------------------------------=#
println("Euler 121: $(maxfund(15))")


# Based on the introductory discussion on
# https://www.maa.org/press/periodicals/loci/joma/the-sir-model-for-spread-of-disease-the-differential-equation-model

## Uncomment the following lines if used as a standalone
# using DifferentialEquations
# using Plots

function SIR(du, u, p, t)
    s, i, r = u
    b, k = p
    du[1] = -b * s * i
    du[2] = b * s * i - k * i
    du[3] = k * i
end

# TODO: make the following a doctest example
    # #initian conditions
    # p = [0.5/7900000.0,0.33]
    # u0 = [7900000.0,10.0,0.0]
    # tspan = (0.0,140.0)

    # #solve
    # sir = ODEProblem(SIR,u0,tspan,p)
    # sol = solve(sir)

    # #plot
    # plot(sol)


function exponential_fg!(∇f, x)

    if !(∇f==nothing)
        ∇f[1] = -2.0 * (2.0 - x[1]) * exp((2.0 - x[1])^2)
        ∇f[2] = -2.0 * (3.0 - x[2]) * exp((3.0 - x[2])^2)
    end

    fx = exp((2.0 - x[1])^2) + exp((3.0 - x[2])^2)

    return fx
end

function exponential_hessian!(storage, x)
    storage[1, 1] = 2.0 * exp((2.0 - x[1])^2) * (2.0 * x[1]^2 - 8.0 * x[1] + 9)
    storage[1, 2] = 0.0
    storage[2, 1] = 0.0
    storage[2, 2] = 2.0 * exp((3.0 - x[2])^2) * (2.0 * x[2]^2 - 12.0 * x[2] + 19)
end

tp_fletch_powell_fg! = TestProblem(exponential_fg!, [0.0, 0.0], I, NearlyNewton.OptOptions())
tp_fletch_powell_fg!_alt = TestProblem(exponential_fg!, [0.0, 0.0], I, NearlyNewton.OptOptions())


"""
    laplacian(n::Int, h::Real)
    laplacian(x::AbstractRange{T}) where {T}

Returns ``n \\times n`` differentiation matrix ``\\mathrm{D^{(2)}}`` using sinc interpolants.
See Weideman, J.A., and Reddy, S.C. (2000). A MATLAB differentiation matrix suite.
ACM Transactions on Mathematical Software (TOMS), 26(4), 465-519.
Equation (20) on page 485.

On input: `n` - size of a uniform grid with step size `h`, or `x` - `AbstractRange` object.

# Example
Second derivative of the gaussian function
```math
    \\partial^2 \\exp(-x^2)/\\partial x^2 = 2(2x^2 - 1)\\exp(-x^2)
```
```jldoctest
julia> x = -6:0.01:6
-6.0:0.01:6.0

julia> d²f_exact = 2(2x.^2 .- 1).*exp.(-x.^2);

julia> d²f_approx = laplacian(length(x), step(x)) * exp.(-x.^2);

julia> isapprox(d²f_exact, d²f_approx, atol = 1e-10)
true
```
"""
function laplacian(n, h)
  Δ = zeros(n, n)
  Δ[diagind(Δ, 0)] .= -1/3*π^2 / h^2
  for i=2:n
    @inbounds Δ[diagind(Δ, i-1)] =
      Δ[diagind(Δ, 1-i)] .= 2*(-1)^i / (i-1)^2 / h^2
  end
  return Δ
end
laplacian(x::AbstractRange) = laplacian(length(x), step(x))

function FGH(x::AbstractRange)
  N, Δx = length(x), step(x)
  iseven(N) && throw(DomainError(N, "number of grid points must be odd"))
  K = zeros(N, N)
  for d=0:N-1
    # See Eq. 26 // C. Clay Marston et al. The Fourier grid Hamiltonian method
    # for bound state eigenvalues and eigenfunctions. JCP 91, 3571 (1989)
    # d ≝ i - j
    K[diagind(K, d)] = K[diagind(K, -d)] .=
    -2/N*(2π/N)^2 / Δx^2 * sum(l -> cos(2π*l*d/N)*l^2, 1:(N >> 1))
  end
  return K
end

function sincdiff(x::AbstractRange)
  n, dx = length(x), step(x)
  ∇ = zeros(n, n)
  ∇[diagind(∇, 0)] .= 0
  for i = 2:n
      @inbounds ∇[diagind(∇, i - 1)] .= (-1)^i / (i - 1) / dx
      @inbounds ∇[diagind(∇, 1 - i)] .= (-1)^(i - 1) / (i - 1) / dx
  end
  return ∇
end

function ∂r(f, x, r = exp.(x), D = sincdiff(x) - I, rₘᵢₙ = 5e-5)
  df = D * (f .* r)
  @. df /= r^2
  extrapolate_left!(df, r, rₘᵢₙ)
  return df
end

#=
 ρ = similar(r)
∇ρ = similar(r)

@. ρ = (2/27*r^2 - 2/3*r + 1)^2 * exp(-2/3*r) / 27π
@. ∇ρ = -(8*r^4 - 192*r^3 + 1512*r^2 - 4536r + 4374)*exp(-2/3*r) / 59049π

norm(∂r(ρ, x) .- ∇ρ)
=#


using JuMP
using Clp
using LinearAlgebra
using Plots

n = 100
grid = range(-1, 1, length=n)
f = ℯ.^grid
coeff = zeros(n, 5)
for degree in 0:4
    coeff[:, degree + 1] = grid .^ degree
end
errors = zeros(5)
polynomials = zeros(5, 5)
taylors = [1 0 0 0 0; 1 1 0 0 0; 1 1 1/2 0 0; 1 1 1/2 1/6 0; 1 1 1/2 1/6 1/24]

for d in 0:4

    model = Model(Clp.Optimizer)

    @variables(model, begin
        m >= 0
        a[1:5]
    end
    )

    @objective(model, Min, m)

    @constraints(model, begin
        [i=d+2:5], a[i] == 0
        m .>= f - coeff * a
        m .>= coeff * a - f
    end
    )

    optimize!(model)
    errors[d + 1] = objective_value(model)
    polynomials[d + 1, :] = value.(a)
end

plot(grid, f, label="e^x")
for d in 0:4
    plot!(grid, coeff * (polynomials[d + 1, :]), label="cheb $d")
end
savefig("polynomial.png")

plot(0:4, errors)
savefig("error.png")

for d in 0:4
    plot(grid, f, label="e^x")
    plot!(grid, coeff * polynomials[d + 1, :], label="cheb $d")
    plot!(grid, coeff * taylors[d + 1, :], label="tayl $d")
    savefig("compare-$d.png")
end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _lsf_rotate{T<:AbstractFloat}(Δ::T, vsini::T, ɛ::T)
    n_2 = floor(ceil(2*vsini/Δ)/2)
    vel = collect(-n_2:n_2)*Δ
    e1 = 2*(1.0 - ɛ)
    e2 = pi*ɛ/2.0
    e3 = pi*(1.0 - ɛ/3.0)
    x = vel/vsini
    x1 = abs(1.0 - x.^2)
    return vel, (e1*sqrt(x1) + e2*x1)/e3
end

"""
    lsf_rotate(delta_v, v_sin_i[, epsilon = 0.3]) -> velocity_grid, lsf

### Purpose ###

Create a 1-d convolution kernel to broaden a spectrum from a rotating star.

### Explanation ###

Can be used to derive the broadening effect (LSF, line spread function) due to
rotation on a synthetic stellar spectrum.  Assumes constant limb darkening
across the disk.

### Arguments ###

* `delta_v`: numeric scalar giving the step increment (in km/s) in the output
  rotation kernel
* `v_sin_i`: the rotational velocity projected along the line of sight (km/s)
* `epsilon` (optional numeric argument): numeric scalar giving the
  limb-darkening coefficient, default = 0.6 which is typical for photospheric
  lines.  The specific intensity \$I\$ at any angle \$\\theta\$ from the
  specific intensity \$I_{\\text{cen}}\$ at the center of the disk is given by:

\$\$ I = I_{\\text{cen}}\\cdot(1 - \\varepsilon\\cdot(1 - \\cos(\\theta))) \$\$

### Output ###

The 2-tuple (`velocity_grid`, `lsf`):

* `velocity_grid`: vector of velocity grid with the same number of elements as
  `lsf` (see below)
* `lsf`: the convolution kernel vector for the specified rotational velocity.
  The number of points in `lsf` will be always be odd (the kernel is symmetric)
  and equal to either `ceil(2*v_sin_i/delta_v)` or `ceil(2*v_sin_i/delta_v) + 1`,
  whichever number is odd.  Elements of `lsf` will always be of type
  `AbstractFloat`.  To actually compute the broadening, the spectrum should be
  convolved with the rotational `lsf`

### Example ###

Plot the line spread function for a star rotating at 90 km/s in velocity space
every 3 km/s.  Use [PyPlot.jl](https://github.com/stevengj/PyPlot.jl) for
plotting.

``` julia
using PyPlot
plot(lsf_rotate(3, 90)...)
```

### Notes ###

Code of this function is based on IDL Astronomy User's Library.
"""
lsf_rotate(Δ::Real, v::Real, ɛ::Real=0.6) =
    _lsf_rotate(promote(float(Δ), float(v), float(ɛ))...)


"""
Imputes gradients based on a vector of x and y coordinates.
### Takes
 * x - A Float64 vector of x coordinates
 * y - A Float64 vector of y coordinates

### Returns
 * A Float64 vector of gradients for each input point
"""
function imputeGradients(x,y)
     n = length(x)
     # Judd (1998), page 233, second last equation
     L = sqrt( (x[2:n]-x[1:(n-1)]).^2 + (y[2:n]-y[1:(n-1)]).^2)
     # Judd (1998), page 233, last equation
     d = (y[2:n]-y[1:(n-1)])./(x[2:n]-x[1:(n-1)])
     # Judd (1998), page 234, Eqn 6.11.6
     Conditionsi = d[1:(n-2)].*d[2:(n-1)] .> 0
     MiddleSiwithoutApplyingCondition = (L[1:(n-2)].*d[1:(n-2)]+L[2:(n-1)].* d[2:(n-1)]) ./ (L[1:(n-2)]+L[2:(n-1)])
     sb = Conditionsi .* MiddleSiwithoutApplyingCondition
     # Judd (1998), page 234, Second Equation line plus 6.11.6 gives this array of slopes.
     ff = [((-sb[1]+3*d[1])/2);  sb ;  ((3*d[n-1]-sb[n-2])/2)]
     return ff
 end

 """
 Creates a spline defined by interval starts IntStarts and quadratic coefficients SpCoefs which evaluates an input point.
### Takes
  * IntStarts - A Float64 vector that gives the starting points of intervals (in the x plane)
  * SpCoefs - A 3 column matrix with the same number of rows as the length of the IntStarts vector. The first column is the coefficient of the quadratic term, the second column for the linear term. The third column is the constant.

### Returns
  * A spline function that takes a single Float64 input and returns the spline value at that point.
 """
function ppmak(IntStarts,SpCoefs)
  function sp(PointToExamine)
    IntervalNum = searchsortedlast(IntStarts, PointToExamine)
    IntervalNum = max(IntervalNum, 1)
    xmt = PointToExamine - IntStarts[IntervalNum]
    Coefs = SpCoefs[ IntervalNum , :]
    return reshape(Coefs' * [xmt^2 xmt 1]', 1)[1]
  end
  function Vsp(PointToExamine)
    return map(x -> sp(x), PointToExamine)
  end
  return Vsp
end

 """
 Creates the derivative function of the spline defined by interval starts IntStarts and quadratic coefficients SpCoefs which evaluates an input point.
### Takes
  * IntStarts - A Float64 vector that gives the starting points of intervals (in the x plane)
  * SpCoefs - A 3 column matrix with the same number of rows as the length of the IntStarts vector. The first column is the coefficient of the quadratic term, the second column for the linear term. The third column is the constant.

### Returns
 * The derivative function that takes a single Float64 input and returns the derivative at that point.
 """
function ppmakDeriv(IntStarts,SpCoefs)
  function sp(PointToExamine)
    IntervalNum = searchsortedlast(IntStarts, PointToExamine)
    IntervalNum = max(IntervalNum, 1)
    xmt = PointToExamine - IntStarts[IntervalNum]
    Coefs = SpCoefs[ IntervalNum , :]
    return reshape(Coefs' * [2*xmt 1 0]', 1)[1]
  end
  function Vsp(PointToExamine)
    return map(x -> sp(x), PointToExamine)
  end
  return Vsp
end

"""
Creates the second derivative function of the spline defined by interval starts IntStarts and quadratic coefficients SpCoefs which evaluates an input point.
### Takes
 * IntStarts - A Float64 vector that gives the starting points of intervals (in the x plane)
 * SpCoefs - A 3 column matrix with the same number of rows as the length of the IntStarts vector. The first column is the coefficient of the quadratic term, the second column for the linear term. The third column is the constant.

### Returns
 * The second derivative function that takes a single Float64 input and returns the second derivative at that point.
"""
function ppmak2Deriv(IntStarts,SpCoefs)
  function sp(PointToExamine)
    IntervalNum = searchsortedlast(IntStarts, PointToExamine)
    IntervalNum = max(IntervalNum, 1)
    xmt = PointToExamine - IntStarts[IntervalNum]
    Coefs = SpCoefs[ IntervalNum , :]
    return reshape(Coefs' * [2 0 0]', 1)[1]
  end
  function Vsp(PointToExamine)
    return map(x -> sp(x), PointToExamine)
  end
  return Vsp
end

"""
Splits an interval into 2 subintervals and creates the quadratic coefficients
### Takes
 * s - A 2 entry Float64 vector with gradients at either end of the interval
 * z - A 2 entry Float64 vector with y values at either end of the interval
 * Smallt - A 2 entry Float64 vector with x values at either end of the interval

### Returns
 * A 2 x 5 matrix. The first column is the x values of start of the two subintervals. The second column is the ends. The last 3 columns are quadratic coefficients in two subintervals.
"""
function schumakerIndInterval(s,z,Smallt)
   # The SchumakerIndInterval function takes in each interval individually
   # and returns the location of the knot as well as the quadratic coefficients in each subinterval.

   # Judd (1998), page 232, Lemma 6.11.1 provides this if condition:
   if (sum(s)*(Smallt[2]-Smallt[1]) == 2*(z[2]-z[1]))
     tsi = Smallt[2]
   else
     # Judd (1998), page 233, Algorithm 6.3 along with equations 6.11.4 and 6.11.5 provide this whole section
     delta = (z[2] -z[1])/(Smallt[2]-Smallt[1])
     Condition = ((s[1]-delta)*(s[2]-delta) >= 0)
     Condition2 = abs(s[2]-delta) < abs(s[1]-delta)
     if (Condition)
       tsi = sum(Smallt)/2
     elseif (Condition2)
       tsi = (Smallt[1] + (Smallt[2]-Smallt[1])*(s[2]-delta)/(s[2]-s[1]))
     else
       tsi = (Smallt[2] + (Smallt[2]-Smallt[1])*(s[1]-delta)/(s[2]-s[1]))
     end
   end

   # Judd (1998), page 232, 3rd last equation of page.
   alpha = tsi-Smallt[1]
   beta = Smallt[2]-tsi
   # Judd (1998), page 232, 4th last equation of page.
   sbar = (2*(z[2]-z[1])-(alpha*s[1]+beta*s[2]))/(Smallt[2]-Smallt[1])
   # Judd (1998), page 232, 3rd equation of page. (C1, B1, A1)
   Coeffs1 = [ (sbar-s[1])/(2*alpha)  s[1]  z[1] ]
   if (beta == 0)
     Coeffs2 = Coeffs1
   else
     # Judd (1998), page 232, 4th equation of page. (C2, B2, A2)
     Coeffs2 = [ (s[2]-sbar)/(2*beta)  sbar  Coeffs1 * [alpha^2, alpha, 1] ]
   end
   Machine4Epsilon = 4*eps()
     if (tsi  <  Smallt[1] + Machine4Epsilon )
         return [Smallt[1] Smallt[2] Coeffs2]
     elseif (tsi + Machine4Epsilon > Smallt[2] )
         return [Smallt[1] Smallt[2] Coeffs1]
     else
         return [Smallt[1] tsi Coeffs1 ; tsi Smallt[2] Coeffs2]
     end
 end

 """
 Calls SchumakerIndInterval many times to get full set of spline intervals and coefficients. Then calls extrapolation for out of sample behaviour
### Takes
 * gradients - A Float64 vector of gradients at each point
 * x - A Float64 vector of x coordinates
 * y - A Float64 vector of y coordinates
 * extrapolation - A string in ("Curve", "Linear", "Constant") that gives behaviour outside of interpolation range.

### Returns
 * A vector of interval starts
 * A vector of interval ends
 * A matrix of all coefficients
  """
 function getCoefficientMatrix(gradients,y,x, extrapolation)
   n = length(x)
   fullMatrix = schumakerIndInterval([gradients[1] gradients[2]], [y[1] y[2]], [x[1] x[2]] )
    for intrval = 2:(n-1)
      Smallt = [ x[intrval] , x[intrval + 1] ]
      s = [ y[intrval], y[intrval + 1] ]
      z = [ gradients[intrval], gradients[intrval + 1] ]
      intMatrix = schumakerIndInterval(z,s,Smallt)
      fullMatrix = vcat(fullMatrix,intMatrix)
    end
    fullMatrix = extrapolate(fullMatrix, extrapolation, x, y)
   return fullMatrix[:,1], fullMatrix[:,2], fullMatrix[:,3:5]
 end

"""
 Adds a row on top and bottom of coefficient matrix to give out of sample prediction.
### Takes
 * fullMatrix - output from GetCoefficientMatrix first few lines
 * extrapolation - A string in ("Curve", "Linear", "Constant") that gives behaviour outside of interpolation range.
 * x - A Float64 vector of x coordinates
 * y - A Float64 vector of y coordinates

### Returns
  * A new version of fullMatrix with out of sample prediction built into it.
"""
function extrapolate(fullMatrix, extrapolation, x, y)

  if (extrapolation == "Curve")
    return fullMatrix
  end

  dim = size(fullMatrix)[1]

  Botx   = fullMatrix[1,1]
  Boty   = y[1]

  if (extrapolation == "Linear")
    BotB = fullMatrix[1 , 4]
    BotC   = Boty - BotB
  else
    BotB = 0.0
    BotC = Boty
  end

  BotRow = [ Botx-1, Botx, 0.0, BotB, BotC]

  Topx = fullMatrix[dim,2]
  Topy = y[length(y)]

  if (extrapolation == "Linear")
    TopB = fullMatrix[dim ,4]
    TopC = Topy
  else
    TopB = 0.0
    TopC = Topy
  end

  TopRow = [ Topx, Topx + 1, 0.0 ,TopB ,TopC]

  fullMatrix = vcat(BotRow' , fullMatrix,  TopRow')

  return fullMatrix
end

"""
Creates splines for a given set of x and y values (and optionally gradients) and the first and second derivatives of this spline.
### Takes
* x - A Float64 vector of x coordinates.
* y - A Float64 vector of y coordinates.
* gradients (optional)- A Float64 vector of gradients at each point. If not supplied these are imputed from x and y.
* extrapolation (optional) - This should be a string in ("Curve", "Linear", "Constant") specifying how to interpolate outside of the sample domain. By default it is "curve" which extends out the first and last quadratic curves. The other options are "Linear" which extends the line (from first and last curve) out from the first and last point and "Constant" which extends out the y value at the first and last point.

### Returns
* A spline which takes and input value and returns the spline y value.
* The derivative of this spline.
* The second derivative of this spline.
 """
function schumaker(x,y,gradients = "Not-Supplied", extrapolation = ("Curve", "Linear", "Constant"))
  # This is the main function of the package that creates and returns a schumaker spline.
  # Inputs : An x array, the corresponding y array and optionally the gradients at each point.
  #           If gradients are not input then they are estimated.
  # Outputs: A schumaker spline function. Its derivative and its second derivative.

  if (extrapolation == ("Curve", "Linear", "Constant"))
    extrapolation = "Curve"
  end

  if (gradients == "Not-Supplied")
     gradients = imputeGradients(x,y)
  end

  IntStarts, IntEnds, SpCoefs = getCoefficientMatrix(gradients,y,x, extrapolation)

  Sp      = ppmak(IntStarts,SpCoefs)
  SpDeriv = ppmakDeriv(IntStarts,SpCoefs)
  SpDeriv2= ppmak2Deriv(IntStarts,SpCoefs)
  return Sp, SpDeriv, SpDeriv2
end


"""
    Beta
* Parametrization
* Score
* Fisher Information
* `time_varying_params` map.
* Default link
"""
Beta

function score!(score_til::Matrix{T}, y::T, ::Type{Beta}, param::Matrix{T}, t::Int) where T
    score_til[t, 1] = log(y) + digamma(param[t, 1] + param[t, 2]) - digamma(param[t, 1])
    score_til[t, 2] = log(1 - y) + digamma(param[t, 1] + param[t, 2]) - digamma(param[t, 2])
    return
end

function fisher_information!(aux::AuxiliaryLinAlg{T}, ::Type{Beta}, param::Matrix{T}, t::Int) where T
    minus_trigamma_a_b = -trigamma(param[t, 1] + param[t, 2])
    aux.fisher[1, 1] = trigamma(param[t, 1]) + minus_trigamma_a_b
    aux.fisher[2, 2] = trigamma(param[t, 2]) + minus_trigamma_a_b
    aux.fisher[2, 1] = minus_trigamma_a_b
    aux.fisher[1, 2] = minus_trigamma_a_b
    return
end

function log_likelihood(::Type{Beta}, y::Vector{T}, param::Matrix{T}, n::Int) where T
    loglik = 0.0
    for t in 1:n
        loglik += (param[t, 1] - 1)*log(y[t]) + (param[t, 2] - 1)*log(1 - y[t]) - logbeta(param[t, 1], param[t, 2])
    end
    return -loglik
end

# Links
function link!(param_tilde::Matrix{T}, ::Type{Beta}, param::Matrix{T}, t::Int) where T
    param_tilde[t, 1] = link(LogLink, param[t, 1], zero(T))
    param_tilde[t, 2] = link(LogLink, param[t, 2], zero(T))
    return
end
function unlink!(param::Matrix{T}, ::Type{Beta}, param_tilde::Matrix{T}, t::Int) where T
    param[t, 1] = unlink(LogLink, param_tilde[t, 1], zero(T))
    param[t, 2] = unlink(LogLink, param_tilde[t, 2], zero(T))
    return
end
function jacobian_link!(aux::AuxiliaryLinAlg{T}, ::Type{Beta}, param::Matrix{T}, t::Int) where T
    aux.jac[1] = jacobian_link(LogLink, param[t, 1], zero(T))
    aux.jac[2] = jacobian_link(LogLink, param[t, 2], zero(T))
    return
end

# utils
function update_dist(::Type{Beta}, param::Matrix{T}, t::Int) where T
    small_threshold!(param, SMALL_NUM, t)
    return Beta(param[t, 1], param[t, 2])
end

function params_sdm(d::Beta)
    return Distributions.params(d)
end

function num_params(::Type{Beta})
    return 2
end


using Calculus

################################################################################
# Function 1
################################################################################
f = function(x)
    n = length(x)
    total = 0
    for i in 1:n-1
        total += (x[i+1]-x[i])^2 + (exp(x[i]) - 1)^2
    end
    return(total + (exp(x[n]) - 1)^2)
end

grad_f = function(x)
    n = length(x)
    J = zeros(n)
    for i in 1:n
        if i == 1
            J[1] = -2 * (x[2]-x[1]) + 2*(exp(2*x[1]) - exp(x[1]))
        elseif i == n
            J[n] = 2 * (x[n]-x[n-1]) + 2*(exp(2*x[n]) - exp(x[n]))
        else
            J[i] = -2 * (x[i+1]-x[i]) + 2 * (x[i]-x[i-1]) + 2*(exp(2*x[i])-exp(x[i]))
        end
    end
    return(J)
end

hess_f = function(x)
    n = length(x)
    H = zeros(n,n)
    for i in 1:n
        for j in 1:n
            if i == j
                H[i,j] = 2 + 4*exp(2*x[i]) - 2*exp(x[i])
                if (i != n) && (i != 1)
                    H[i,j] += 2
                end

            elseif abs(i-j) == 1
                H[i,j] = -2
            else
                H[i,j] = 0
            end
        end
    end
    return(H)
end

# test that the numerical gradients equal the analytically computed gradients
#x = [1.;2.;3.]
#println(sum(grad_f(x)-Calculus.gradient(f,x)))
#println(sum(hess_f(x)-Calculus.hessian(f,x)))

f_g = Calculus.gradient(f)
f_h = Calculus.hessian(f)

################################################################################
# Rosenbrock
################################################################################

function rosenbrock(x)
    """
    Assumes x has an even number of elemnts.
    """
    n = length(x)
    @assert n%2 == 0
    total = 0
    for i in 1:Integer(n/2)
        total += (1-x[2*i-1])^2 + 10*(x[2*i] - x[2*i-1]^2)^2
    end
    return total
end

rosenbrock_g = Calculus.gradient(rosenbrock)
rosenbrock_h = Calculus.hessian(rosenbrock)

################################################################################
# Cute functions
################################################################################

function cute(x)
    n = length(x)
    total = 0
    for i in 1:n-4
        total += (-4x[i]+3)^2 + (x[i]^2 + 2*x[i+1]^2 + 3*x[i+2]^2 + 4*x[i+3]^2 + 5*x[n]^2)^2
    end
	return total
end

cute_g = Calculus.gradient(cute)
cute_h = Calculus.hessian(cute)

function cute2(x)
    """
    the cute2 function
    """
    n = length(x)
    inds = 2:n-3
    return sum((-4*x[inds]+3.0).^2) +
        sum((x[inds].^2 + 2*x[inds+1].^2 +
		3*x[inds+2].^2 + 4*x[inds+3].^2 + 5*x[1].^2).^2)
end

cute2_g = Calculus.gradient(cute2)
cute2_h = Calculus.hessian(cute2)

################################################################################
# Fenton
################################################################################

function fenton(x)
    """
    Using Newton's method, the starting point [3;2] converges but [3;4] diverges
    """
    x1 = x[1]; x2 = x[2];
    return((12 + x1*x1 + (1+x2*x2)/(x1*x1) + (x1*x1*x2*x2+100)/((x1*x2)^4)) / 10)
end

fenton_g = Calculus.gradient(fenton)
fenton_h = Calculus.hessian(fenton)


function gausslobatto(n)
    # Gauss-Legendre-Lobatto Quadrature Nodes and Weights
    if n == 1
        error("Lobatto undefined for n = 1.")
    elseif n == 2
        [-1.0, 1.0], [1.0, 1.0]
    elseif n == 3
        [-1.0, 0.0, 1.0], [1.0 / 3, 4.0 / 3, 1.0 / 3]
    else
        # Compute via GaussJacobi:
        x, w = gaussjacobi(n - 2, 1.0, 1.0)
        @inbounds for i in 1:length(x)
            w[i] = w[i] / (1 - x[i]^2)
        end
        pushfirst!(x, -1.0)
        push!(x, 1.0)
        pushfirst!(w, 2 / (n * (n - 1)))
        push!(w, 2 / (n * (n - 1)))
        x, w
    end
end


module LP

using JuMP
using Clp

function example()
    # inspired by Linear Algebra and its applications Ch. 9

    m = Model(Clp.Optimizer)

    # unknowns
    @variable(m, 0 <= x1)
    @variable(m, 0 <= x2)
    @variable(m, 0 <= x3)

    # dot(C, x)
    @objective(m, Max, 2x1 + 3x2 + 4x3)

    # A ∈ R(mxn), with m constraints and n unknowns
    @constraint(m, x1 + x2 + x3 <= 50)
    @constraint(m, x1 + 2x2 + 4x3 <= 80)

    optimize!(m)

    println("Optimal X: x1: $(value(x1)) x2: $(value(x2)) x3: $(value(x3))")
    println("Optimal Objective value: $(objective_value(m))")

end # example

function practice_problem_one()
    # inspired by Linear Algebra and its applications Ch. 9

    m = Model(Clp.Optimizer)

    # unknowns
    @variable(m, 0 <= x1)
    @variable(m, 0 <= x2)

    # dot(C, x)
    @objective(m, Max, 2x1 + x2)

    @constraint(m, -x1 + x2 <= 8)
    @constraint(m, 3x1 + 2x2  <= 24)

    optimize!(m)

    println("Optimal X: x1: $(value(x1)) x2: $(value(x2))")
    println("Optimal Objective value: $(objective_value(m))")

end # practice problem one

function exercise_seven()
    # inspired by Linear Algebra and its applications Ch. 9

    m = Model(Clp.Optimizer)

    # unknowns
    @variable(m, 0 <= x1)
    @variable(m, 0 <= x2)

    # dot(C, x)
    @objective(m, Max, 80x1 + 65x2)

    @constraint(m, 2x1 + x2 <= 32)
    @constraint(m, x1 + x2  <= 18)
    @constraint(m, x1 + 3x2  <= 24)

    optimize!(m)

    println("Optimal X: x1: $(value(x1)) x2: $(value(x2))")
    println("Optimal Objective value: $(objective_value(m))")

end # exercise seven

end # module


module CubicHermiteSpline

using ArgCheck

export CubicHermiteSplineInterpolation
export interp, grad

struct CubicHermiteSplineInterpolation
    x
    y
    gradient
end

"""
    basis(t)

Compute the basis functions of cubic Hermite spline interpolation. They are:
    H_{00} = 2t^3 - 3t^2 + 1 = (1+2t)(t-1)^2
    H_{10} = t^3 - 2t^2 + t = t(t-1)^2
    H_{01} = -2t^3 + 3t^2 = t^2(3-2t)
    H_{11} = t^3 - t^2 = t^2(t - 1)
"""
function basis(t)
    t2 = t * t
    it = t - 1
    it2 = it * it
    tt = 2 * t
    h00 = (1 + tt) * it2
    h10 = t * it2
    h01 = t2 * (3 - tt)
    h11 = t2 * it
    return h00, h10, h01, h11
end

"""
    basis_derivative(t)

Compute the basis functions of cubic Hermite spline interpolation. They are:
    H_{00} = 6t^2 - 6t = 6t(t-1)
    H_{10} = 3t^2 - 4t + 1 = (3t-1)(t-1)
    H_{01} = -6t^2 + 6t = -6t(t-1)
    H_{11} = 3t^2 - 2t = t(3t - 2)
"""
function basis_derivative(t)
    t2 = t * t
    h00 = 6*t2 - 6*t
    h10 = 3*t2 - 4*t + 1
    h01 = -6*t2 + 6*t
    h11 = 3*t2 - 2*t
    return h00, h10, h01, h11
end

function findinterval(v, x)
    for i in eachindex(x)
        if x[i] == v
            return i, x[i], nothing
        elseif x[i] > v
            return i-1, x[i-1], x[i]
        end
    end
end

function _interp(spl::CubicHermiteSplineInterpolation, v; grad=false)
    x = spl.x
    y = spl.y
    gradient = spl.gradient
    @argcheck v >= x[1]
    @argcheck v <= x[end]

    idx, x1, x2 = findinterval(v, x)
    if x2 === nothing
       return grad ? gradient[idx] : y[idx]
    end

    # mapping (x1, x2) to (0, 1), h is the scaling constant
    t = (v - x1) / (x2 - x1)
    h = x2 - x1
    y1, y2 = y[idx], y[idx+1]
    k1, k2 = gradient[idx], gradient[idx+1]

    h00, h10, h01, h11 = grad ? basis_derivative(t) : basis(t)
    r = y1*h00 + h*k1*h10 + y2*h01 + h*k2*h11

    # Need to rescale the interpolated gradient before return
    return grad ? (r/h) : r
end

(spl::CubicHermiteSplineInterpolation)(v::Real; grad=false) = _interp(spl, v; grad=grad)
(spl::CubicHermiteSplineInterpolation)(x::AbstractVector; grad=false) = spl.(x; grad=grad)

# handy methods
interp(spl::CubicHermiteSplineInterpolation, p) = spl(p)
grad(spl::CubicHermiteSplineInterpolation, p) = spl(p; grad=true)

end # module


# Locally Linear Embedding (LLE)
# ------------------------
# Nonlinear dimensionality reduction by locally linear embedding,
# Roweis, S. & Saul, L., Science 290:2323 (2000)

using Logging

#### LLE type
struct LLE{T <: Real} <: AbstractDimensionalityReduction
    k::Int
    λ::AbstractVector{T}
    proj::Projection{T}

    LLE{T}(k::Int, λ::AbstractVector{T}, proj::Projection{T})  where T = new(k, λ, proj)
end

## properties
outdim(R::LLE) = size(R.proj, 1)
eigvals(R::LLE) = R.λ
neighbors(R::LLE) = R.k

## show
summary(io::IO, R::LLE) = print(io, "LLE(outdim = $(outdim(R)), neighbors = $(neighbors(R)))")

## interface functions
function fit(::Type{LLE}, X::AbstractMatrix{T};
             maxoutdim::Int=2, k::Int=12, tol::Real=1e-5,
             use_naive=false, use_eigen=false) where {T<:Real}
    # Construct NN graph
    D, E = find_nn(X, k)
    _, C = largest_component(SimpleWeightedGraph(adjmat(D,E)))
    X = X[:, C]
    n = length(C)

    # Correct indexes of neighbors if more then one connected component
    Ec = E
    if size(E,2) != n
        R = Dict(zip(C, collect(1:n)))
        Ec = zeros(Int,k,n)
        for i in 1 : n
            Ec[:,i] = map(j->get(R,j,C[i]), E[:,C[i]])
        end
    end

    if k > maxoutdim
        @warn("k > maxoutdim: regularization will be used")
    end

    # Reconstruct weights and compute embedding:
    # M = (I - w)'(I - w) = I - w'I - Iw + w'w
    M = spdiagm(0 => fill(one(T), n))
    Ones = fill(one(T), k, 1)
    for i in 1 : n
        J = Ec[:,i]
        Z = view(X, :, J) .- view(X, :, i)
        G = transpose(Z)*Z
        G += I * tol # regularize
        w = vec(G \ Ones)
        w ./= sum(w)
        ww = w*transpose(w)
        for (l, j) in enumerate(J)
            M[i,j] -= w[l]
            M[j,i] -= w[l]
            for (m, jj) in enumerate(J)
                M[j,jj] = ww[l,m]
            end
        end
    end

    if use_eigen
        λ, V = decompose(M, maxoutdim)
    else
        λ, V = specDecompose(M, tol, maxoutdim, use_naive)
    end
    return LLE{T}(k, λ, rmul!(transpose(V), sqrt(n)))
end

transform(R::LLE) = R.proj


"""
 Convert cartesian coordinates to spherical coordinates (θ,ϕ,r)
 θ, ϕ are in radians
 ϕ is the latitude, i.e., the angle from the equator.
"""
function xyz2θϕr(x)
  r = sqrt(x[1]^2 + x[2]^2 + x[3]^2)
  θ = atan(x[2], x[1])
  ϕ = asin(x[3]/r)
  VectorValue(θ,ϕ,r)
end

"""
  Map point on the unit sphere on ambientspace to the parametric space of the unit sphere
  (x,y,z) ↦ (θ,ϕ)
"""
function xyz2θϕ(x)
  θ = atan(x[2], x[1])
  ϕ = asin(x[3])
  VectorValue(θ,ϕ)
end

"""
  Map point on the parametric space of the unit sphere to the ambient space
  (θ,ϕ) ↦ (x,y,z)
"""
function θϕ2xyz(θϕ)
  θ,ϕ = θϕ
  x = cos(θ)*cos(ϕ)
  y = sin(θ)*cos(ϕ)
  z = sin(ϕ)
  VectorValue(x,y,z)
end

"""
Matrix transformation from spherical vector field to Cartesian vector field.
  θ∈(0,2π)
  ϕ∈(-π/2,π/2)
  r=constant for the sphere
"""
function spherical_to_cartesian_matrix(θϕr)
  θ,ϕ,r = θϕr
  TensorValue(-sin(θ)       , cos(θ)       ,      0,
              -sin(ϕ)*cos(θ),-sin(ϕ)*sin(θ), cos(ϕ),
               cos(ϕ)*cos(θ), cos(ϕ)*sin(θ), sin(ϕ))
end

"""
Matrix transformation from cartesian vector field to spherical vector field
"""
function cartesian_to_spherical_matrix(xyz)
  x,y,z = xyz
  sr = sqrt(x^2+y^2+z^2)
  cr = sqrt(x^2+y^2)
  TensorValue(-y/cr, (x*z)/(sr*cr), x/sr,
               x/cr, (y*z)/(sr*cr), y/sr,
                  0,   -cr/(sr*cr), z/sr)
end


using Plots

function bracket_minimum(f, x=0; s=1e-2, k=2.0)
    a, ya = x, f(x)
    b, yb = a + s, f(a + s)
    if yb > ya
        a, b = b, a
        ya, yb = yb, ya
        s = -s
    end
    while true
        c, yc = b + s, f(b + s)
        if yc > yb
            return a < c ? (a,b, c) : (c, b, a)
        end
        a, ya, b, yb = b, yb, c, yc
        s *= k
    end
end

function quadratic_fit_search(f, a, b, c, n)
    ya, yb, yc = f(a), f(b), f(c)
    for i in 1:n-3
        x = 0.5*(ya*(b^2-c^2)+yb*(c^2-a^2)+yc*(a^2-b^2)) /
            (ya*(b-c) +yb*(c-a) +yc*(a-b))
        yx = f(x)
        if x > b
            if yx > yb
                c, yc = x, yx
            else
                a, ya, b, yb = b, yb, x, yx
            end
        elseif x < b
            if yx > yb
                a, ya = x, yx
            else
                c, yc, b, yb = b, yb, x, yx
            end
        end
    end
    return (a, b, c)
end

f = x->sin(10x)+cos(3x)
x0 = 3.4
#busca do intervalo
a,b,c = bracket_minimum(f,x0)

#numero de iteracoes
φ = MathConstants.φ
n=(c-a)/(10^-3*log(φ))

am,bm,cm = quadratic_fit_search(f,a,b,c,n)
(am,f(am))
am,bm,cm

plot(f,0.,4.,lw=2,draw_arrow="true")
scatter!([a,b,c],[f(a),f(b),f(c)])
scatter!([am,bm,cm],[f(am),f(bm),f(cm)])


# This file investigates mean square error (MSE) with respect to free variables in hidden 1D interpolation

using FractalTools
using Makie

# Construct interpolation data
# f(x) = [FractalTools.sinusoid(x), FractalTools.sinusoid(x)]  # For regular data
f(x) = [FractalTools.wen(x), FractalTools.weierstrass(x)]  # For irregular data
line = [[0.], [1.]]
pts = getdata(f, line, 11)

# Construct test data
tpts = getdata(line, 101)
npts = length(tpts)

# Compute errors
fvals = getindex.(map(pnt -> f(pnt[1]), tpts), 1)
freevars = 0.001 : 0.01 : 0.5
mse = map(freevars) do freevar
    interp = interpolate(pts, HInterp1D(fill(freevar, 2, 2)))
    ivals = getindex.(map(pnt -> interp(pnt[1]), tpts), 1)
    sum((fvals - ivals).^2) / npts
end

# Plot mse
fig = Figure()
ax = fig[1, 1] = Axis(fig, xlabel="Free Variable", ylabel="MSE", title="1D Hidden Interpolation MSE")
stem!(ax, freevars, mse, color=:black)
save(joinpath(@__DIR__, "hinterp1d_mse.png"), fig)
display(fig)


export accuracy, bce, logistic, nll, zeroone

"""
    nll(scores, labels; dims=1, average=true)

Return the negative log likelihood for a single batch of data given an unnormalized `scores`
matrix and an `Integer` array of correct `labels`. The `scores` matrix should have size
`(classes,instances)` if `dims=1` or `(instances,classes)` if `dims=2`. `labels[i]` should be
in `1:classes` to indicate the correct class for instance i, or 0 to skip instance i.

The return value is `(total/count)` if `average=true` and `(total,count)` if `average=false`
where `count` is the number of instances not skipped (i.e. `label != 0`) and `total` is their
total negative log likelihood.

## Example

Let's assume that there are three classes (cat, dog, ostrich) and just 2 instances with
the unnormalized score `scores[:,1]` and `scores[:,2]` respectively. The first instance
is actually a cat and the second instance a dog:

    scores = [12.2    0.3;
               2.0   21.5;
               0.0  -21.0]
    labels = [1, 2]
    nll(scores,labels)
    # returns 2.1657e-5

The probabilites are derived from the scores and the negative log-probabilities corresponding
to the labels are averaged:

    probabilites = exp.(scores) ./ sum(exp.(scores),dims=1)
    -(log(probabilites[labels[1],1]) + log(probabilites[labels[2],2]))/2
    # returns 2.1657e-5

"""
function nll(scores,labels::AbstractArray{<:Integer}; dims=1, average=true)
    indices = findindices(scores,labels,dims=dims)
    lp = logsoftmax(scores,dims=dims)[indices]
    average ? (-sum(lp) / length(lp)) : (-sum(lp), length(lp))
end

"""
    accuracy(scores, labels; dims=1, average=true)

Given an unnormalized `scores` matrix and an `Integer` array of correct `labels`, return the
ratio of instances where the correct label has the maximum score. `dims=1` means instances are
in columns, `dims=2` means instances are in rows. Use `average=false` to return the pair
(ncorrect,count) instead of the ratio (ncorrect/count). The valid labels should be integers in
the range `1:numclasses`, if `labels[i] == 0`, instance i is skipped.

"""
function accuracy(scores,labels::AbstractArray{<:Integer}; dims=1, average=true)
    indices = findindices(scores,labels,dims=dims)
    ycpu = convert(Array,value(scores))
    (maxval,maxind) = findmax(ycpu,dims=dims)
    maxind = LinearIndices(ycpu)[maxind]
    maxind = vec(maxind)[vec(labels) .!= 0]
    correct = (maxind .== indices)
    average ? (sum(correct) / length(correct)) : (sum(correct), length(correct))
end

"""
    logistic(scores, labels; average=true)

Computes logistic loss given predicted unnormalized scores and answer labels for a binary
prediction task.

    log.(1 .+ exp.(-labels .* scores))

Label values should be {-1,1}. Scores are unrestricted.  The return value is `(total/count)`
if `average=true` and `(total,count)` if `average=false` where `count` is the number of
instances and `total` is their total loss.

See also `bce` which computes the same loss with {0,1} labels.

Reference: https://towardsdatascience.com/nothing-but-numpy-understanding-creating-binary-classification-neural-networks-with-e746423c8d5c
"""
function logistic(scores, labels::AbstractVector{<:Integer}; average=true)
    labels = ((labels .+ 1) .÷ 2) # (-1,1)->(0,1)
    bce(scores,labels; average=average)
end

"""
    bce(scores, labels; average=true)

Computes binary cross entropy loss given predicted unnormalized scores and answer labels for a
binary prediction task. Label values should be in {0,1}. Scores are unrestricted and will be
converted to probabilities using

    probs = 1 ./ (1 .+ exp.(-scores))

The loss calculated is

    -(labels .* log.(probs) .+ (1 .- labels) .* log.(1 .- probs))

The return value is `(total/count)` if `average=true` and `(total,count)` if `average=false`
where `count` is the number of instances and `total` is their total loss.

See also `logistic` which computes the same loss with {-1,1} labels.

Reference: https://towardsdatascience.com/nothing-but-numpy-understanding-creating-binary-classification-neural-networks-with-e746423c8d5c
"""
function bce(scores, labels::AbstractVector{<:Integer}; average=true)
    labels = oftype(scores,labels)
    l = max.(0, scores) .- labels .* scores .+ log.(1 .+ exp.(-abs.(scores)))
    average ? (sum(l)/length(l)) : (sum(l),length(l))
end

# Indexing help for nll and accuracy

function findindices(scores, labels::AbstractArray{<:Integer}; dims=1)
    ninstances = length(labels)
    nindices = 0
    indices = Vector{Int}(undef,ninstances)
    if dims == 1                   # instances in first dimension
        y1 = size(scores,1)
        y2 = div(length(scores),y1)
        if ninstances != y2; throw(DimensionMismatch()); end
        @inbounds for j=1:ninstances
            if labels[j] == 0; continue; end
            indices[nindices+=1] = (j-1)*y1 + labels[j]
        end
    elseif dims == 2               # instances in last dimension
        y2 = size(scores,ndims(scores))
        y1 = div(length(scores),y2)
        if ninstances != y1; throw(DimensionMismatch()); end
        @inbounds for j=1:ninstances
            if labels[j] == 0; continue; end
            indices[nindices+=1] = (labels[j]-1)*y1 + j
        end
    else
        error("findindices only supports dims = 1 or 2")
    end
    return (nindices == ninstances ? indices : view(indices,1:nindices))
end

"""
    nll(model; data, dims=1, average=true, o...)

Compute the negative log likelihood for a model over a dataset:

    nll(model(inputs; kwargs...), labels; dims) for (inputs,labels) in data

and return `(total/count)` if `average=true` or `(total,count)` if `average=false` where
`count` is the number of instances not skipped (instances with `label==0` are skipped) and
`total` is their total negative log likelihood.

The `model` should be a function returning scores given inputs, and data should be an iterable
of `(inputs,labels)` pairs. The valid labels should be integers in the range `1:numclasses`,
if `labels[i] == 0`, instance i is skipped.

"""
function nll(model; data, dims=1, average=true, o...)
    sum = cnt = 0
    for (x,y) in data
        (z,n) = nll(model(x; o...), y; dims=dims, average=false)
        sum += z; cnt += n
    end
    average ? sum / cnt : (sum, cnt)
end

"""
    accuracy(model; data, dims=1, average=true, o...)

Compute the number of correct predictions of a model over a dataset:

    accuracy(model(inputs; kwargs...), labels; dims) for (inputs,labels) in data

and return `(ncorrect/count)` if `average=true` or `(ncorrect,count)` if `average=false` where
`count` is the number instances not skipped (instances with `label==0` are skipped) and
`ncorrect` is the number of them correctly labeled by the model.

The `model` should be a function returning scores given inputs, and data should be an iterable
of `(inputs,labels)` pairs. The valid labels should be integers in the range `1:numclasses`,
if `labels[i] == 0`, instance i is skipped.

"""
function accuracy(model; data, dims=1, average=true, o...)
    sum = cnt = 0
    for (x,y) in data
        (z,n) = accuracy(model(x; o...), y; dims=dims, average=false)
        sum += z
        cnt += n
    end
    average ? sum / cnt : (sum, cnt)
end

# DEPRECATE:

# The two-arg model,data calls cause mix-ups with the scores,labels calls:
function nll(model,data; o...)
    @warn "nll(model,data; o...) is deprecated, please use nll(model; data=data, o...)" maxlog=1
    nll(model; data=data, o...)
end

function accuracy(model,data; o...)
    @warn "accuracy(model,data; o...) is deprecated, please use accuracy(model; data=data, o...)" maxlog=1
    accuracy(model; data=data, o...)
end

# We need the (model,x,y) interface to implement regularization:
function nll(f, x, y; dims=1, average=true, o...)
    @warn "nll(f,x,y; o...) is deprecated, please use nll(f(x),y; o...) instead." maxlog=1
    nll(f(x; o...), y; dims=dims, average=average)
end

function accuracy(f, x, y; dims=1, average=true, o...)
    @warn "accuracy(f,x,y; o...) is deprecated, please use accuracy(f(x),y; o...) instead." maxlog=1
    accuracy(f(x; o...), y; dims=dims, average=average)
end

# We need the (weights,data,predict) interface to support the old interface:
function nll(w, data, f::Function; dims=1, average=true, o...)
    @warn "nll(weights,data,func; o...) is deprecated, please use nll(x->func(weights,x); data=data, o...) instead." maxlog=1
    nll(x->f(w,x;o...), data; dims=dims, average=average)
end

function accuracy(w, data, f::Function; dims=1, average=true, o...)
    @warn "accuracy(weights,data,func; o...) is deprecated, please use accuracy(x->func(weights,x); data=data, o...) instead." maxlog=1
    accuracy(x->f(w,x;o...), data; dims=dims, average=average)
end

"zeroone loss is equal to 1 - accuracy"
function zeroone(x...; o...)
    @warn "zeroone() is deprecated, please use 1-accuracy()" maxlog=1
    1 - accuracy(x...; o...)
end


using NumericalMethodsforEngineers, Plots

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

  f(x::Float64, y::Vector{Float64}) = (x + y[1])/x
  steps = 3
  h = 0.333

  x = 2.0
  y = [2.0]

  # y(3.0) = 4.2165

  r = Array{Float64,2}[]
  push!(r, euler(f, x, y, steps, h))
  push!(r, modified_euler(f, x, y, steps, h))
  push!(r, mid_point_euler(f, x, y, steps, h))
  push!(r, runga_kutta_4(f, x, y, steps, h))
  println()
  r |> display
  println()

  xs = x + Float64[(i-1)*h for i in 1:steps+1]
  titles = ["Euler", "Modified_Euler", "Mid_Point_Euler", "Runga_Kutta_4"]

  p1 = plot(xs, r[1][:, 2])
  plot!(p1, xs, r[2][:, 2])
  plot!(p1, xs, r[3][:, 2])
  plot!(p1, xs, r[4][:, 2])
  show(p1)
  savefig("Ex7.2.png")
  gui()

end


using LinearAlgebra

function swap_rows(i::T, nlinha::T) where {T<:Integer}
    for n ∈ (i+1):nlinha        # iterate over lines above to check if could be swap
        if A[n,i] ≠ 0.0         # condition to swap row
            L = copy(A[i,:])    # copy line to swap
            A[i,:] = A[n,:]     # swap occur
            A[n,:] = L
            break
        end
    end
end

function gauss_jordan(A::Matrix{T}) where {T<:Number}

    # convert to float to avoid InexactError: Int64()
    (T <: Integer) && (A = convert.(Float64, A))

    # check if matrix is singular
    m, n = size(A)
    if m == n
        @assert det(A) ≠ 0.0 "Must insert a non-singular matrix"
    else
        @assert det(A[:,1:end-1]) ≠ 0.0 "Must insert a non-singular matrix or a system matrix [A b]"
    end

    for i ∈ axes(A, 1)
        if A[i,i] == 0.0                            # check if need swap rows
            swap_rows(i, m)
        end

        @. A[i,:] = A[i,:] / A[i,i]                 # divide pivot line by pivot element

        for j ∈ axes(A, 1)                          # iterate each line for each pivot column, except pivot line
            if j ≠ i                                # jump pivot line
                @. A[j,:] = A[j,:] - A[i,:]*A[j,i]  # apply gauss jordan in each line
            end
        end
    end

    return A
end

## logic in the for loops
# A[1,:] = A[1,:]./ A[1,1]
# A[2,:] = A[2,:] - A[1,:]*A[2,1]
# A[3,:] = A[3,:] - A[1,:]*A[3,1]

# A[2,:] = A[2,:]./ A[2,2]
# A[1,:] = A[1,:] - A[2,:].*A[1,2]
# A[3,:] = A[3,:] - A[2,:].*A[3,2]

# A[3,:] = A[3,:]./ A[3,3]
# A[1,:] = A[1,:] - A[3,:].*A[1,3]
# A[2,:] = A[2,:] - A[3,:].*A[2,3]


isharshad(x)   = x % sum(digits(x)) == 0
nextharshad(x) = begin while !isharshad(x+1) x += 1 end; return x + 1 end

function harshads(n::Integer)
	h = Vector{typeof(n)}(n)
	h[1] = 1
	for j in 2:n
		h[j] = nextharshad(h[j-1])
	end
	return h
end

println("First 20 harshad numbers: ", join(harshads(20), ", "))
println("First harshad number after 1001: ", nextharshad(1000))