# Practical Summary: Ch. 15, Maximum likelihood

# First, let's work with a simple one parameter model,
# for which we can get analytic results.

## Draw some data from the exponential distribution
using Distributions, Econometrics, Plots
θ⁰ = 4.0
n = 100
y = rand(Exponential(θ⁰),n)
histogram(y, normalize=true)
plot!(Exponential(θ⁰))

## Here are the ML results, for reference
using Econometrics, Distributions
logℒᵢ = θ -> logpdf.(Exponential(θ[1]),y)
θstart = [1.]
θhat = mleresults(logℒᵢ, θstart)[1]

##
# for this model, you should be able to work out, analytically,
# that the ML estimator is simply the sample mean. To confirm that:
mean(y)

## Let's get the results from basic theory: first the maximizer of log likelihood
using Optim
tol = 1e-08
s = θ -> -mean(logℒᵢ(θ))  # negative average logℒ
θhat = Optim.optimize(s, θstart, LBFGS(),
                            Optim.Options(
                            g_tol = tol,
                            x_tol=tol,
                            f_tol=tol);autodiff=:forward).minimizer

## Now get the t-stats, from Ihat and J hat.
using ForwardDiff
sc =  ForwardDiff.jacobian(logℒᵢ, θhat) # get the score contributions
Ihat = mean(sc.^2.)
Jhat = -ForwardDiff.hessian(s, θhat)
# when you set n large (above), you should see that Ihat ⩬ -Jhat
# also, note that Jhat = 1/θhat², which is a result we can get analytically

# three forms of estimated variance
V1 = inv(Ihat)/n
se1 = sqrt(V1)
V2 = inv(-Jhat)/n
se2 = sqrt(V2)
V3 =  inv(Jhat)*Ihat*inv(Jhat)/n
se3 = sqrt(V3)
[se1 se2 se3]   # note that the estimators are a little different from one another
                # the last one, sandwich, is what's reported in mleresults.
##

# now let's work with a model that has regressors
using Distributions, Econometrics, Plots
n = 1000
x = [ones(n) randn(n) rand(n)]
β⁰ = [-0.5, 1., 1.]
θ = β ->  exp.(x*β)
y = rand.(Exponential.(θ(β⁰))) # independent, non-identically distributed observations

## Here's the exponential (correct) model
using Econometrics, Distributions, ForwardDiff, Calculus
logℒᵢ = β -> [logpdf(Exponential(exp(x[i,:]'*β)),y[i]) for i = 1:n]
s = β -> mean(logℒᵢ(β))
βstart = zeros(3)
βhat, objval, junk = mleresults(logℒᵢ, βstart, "exponential model (correct)")
## get the score contributions
sc =  ForwardDiff.jacobian(logℒᵢ, βhat)
# compute information matrix estimate
Ihat = zeros(3,3)
for i = 1:n
	Ihat .+= sc[i,:]*sc[i,:]'
end
Ihat ./= n
## now the Hessian
Jhat = Calculus.hessian(s, βhat, :central)
## check that J=-I
Ihat+Jhat # pretty close to zeros, as information matrix equality tells us

## Likelihood ratio test
using Distributions, Econometrics
# a true restriction
R = [1 1 1]
r = 1.5
obj(b) = -s(b) # need to minimize
βhatR, objvalR, junk = fmincon(obj, βhat, R, r)
objvalR = -objvalR # back to maximization
LR = 2*n*(objval - objvalR)
pval = 1. - cdf(Chisq(1), LR)
println("LR test: $R*β=$r (true restriction)  \nLR stat: $LR. p-value: $pval")
# a false restriction
R = [0 1 -1]
r = 1.
obj(b) = -s(b)
βhatR, objvalR, junk = fmincon(obj, βhat, R, r)
objvalR = -objvalR # back to maximization
LR = 2*n*(objval - objvalR)
pval = 1. - cdf(Chisq(1), LR)
println("LR test: $R*β=$r (false restriction)  \nLR stat: $LR. p-value: $pval")

## What about trying a χ² model? If we don't know the true density, we don't know which to use
logℒᵢ = β -> [logpdf(Chisq(exp(x[i,:]'*β)),y[i]) for i = 1:n]
βstart = zeros(3)
βhat, junk = mleresults(logℒᵢ, βstart, "χ² model (incorrect)")
sc =  ForwardDiff.jacobian(logℒᵢ, βhat); # get the score contributions
Ihat = zeros(3,3)
for i = 1:n
	Ihat .+= sc[i,:]*sc[i,:]'
end
Ihat ./= n
Jhat = Calculus.hessian(s, βhat, :central)
# check J=-I
println("check Ihat+Jhat approx zero,  χ² model\n")
Ihat+Jhat # a lot farther away from zeros, at least when n is large.

# How can we tell which model is best? The information criteria
# favor the exponential model, clearly. Comparing I and J,
# the information matrix equality holds much better for the
# exponential model than for the Chi-square. The IM test formalizes
# this, but here, we can see the intuitive basis for the test.


module FourierTransforms

using Compat: cispi
using Primes

################################################################################

function direct_ft!(X::AbstractVector{T}, x::AbstractVector{T}) where {T<:Complex}
    N = length(x)
    @assert length(X) == N
    RT = typeof(real(zero(T)))
    for k in 1:N
        s = zero(T)
        for n in 1:N
            s += cispi(-2 * ((n - 1) * (k - 1) % N) / RT(N)) * x[n]
        end
        X[k] = s
    end
    return X
end

direct_ft(x::AbstractVector) = direct_ft!(similar(x), x)

################################################################################

function direct_ft_1!(X::AbstractVector{T}, x::AbstractVector{T}) where {T<:Complex}
    @assert length(X) == length(x) == 1
    X[1] = x[1]
    return X
end

direct_ft_1(x::AbstractVector) = direct_ft_1!(similar(x), x)

################################################################################

function direct_ft_2!(X::AbstractVector{T}, x::AbstractVector{T}) where {T<:Complex}
    @assert length(X) == length(x) == 2
    X[1] = x[1] + x[2]
    X[2] = x[1] - x[2]
    return X
end

direct_ft_2(x::AbstractVector) = direct_ft_2!(similar(x), x)

################################################################################

function direct_ft_4!(X::AbstractVector{T}, x::AbstractVector{T}) where {T<:Complex}
    @assert length(X) == length(x) == 4
    X[1] = (x[1] + x[3]) + (x[2] + x[4])
    X[2] = (x[1] - x[3]) - im * (x[2] - x[4])
    X[3] = (x[1] + x[3]) - (x[2] + x[4])
    X[4] = (x[1] - x[3]) + im * (x[2] - x[4])
    return X
end

direct_ft_4(x::AbstractVector) = direct_ft_4!(similar(x), x)

################################################################################

# <https://en.wikipedia.org/w/index.php?title=Cooley%E2%80%93Tukey_FFTur_algorithm&oldid=1056450535>
function ditfft2!(X::AbstractVector{T}, x::AbstractVector{T}) where {T<:Complex}
    N = length(x)
    @assert length(X) == N
    RT = typeof(real(zero(T)))
    if N == 0
        # do nothing
    elseif N == 1
        X[1] = x[1]
    else
        # Ensure N is even
        @assert N % 2 == 0
        N2 = N ÷ 2
        ditfft2!((@view X[(0 * N2 + 1):(1 * N2)]), (@view x[1:2:end]))
        ditfft2!((@view X[(1 * N2 + 1):(2 * N2)]), (@view x[2:2:end]))
        for k in 1:N2
            ϕ = cispi(-2 * (k - 1) / RT(N))
            p = X[0 * N2 + k]
            q = ϕ * X[1 * N2 + k]
            X[0 * N2 + k] = p + q
            X[1 * N2 + k] = p - q
        end
    end
    return X
end

ditfft2(x::AbstractVector) = ditfft2!(similar(x), x)

################################################################################

function choose_radix_sqrt(n::Integer)
    @assert n ≥ 0
    n ≤ 1 && return n
    # Handle special cases efficiently
    if ispow2(n)
        # Power of 2
        return 2
    end
    # Power of 2 times a factor
    p2 = trailing_zeros(n)
    factor = n >> p2
    if factor ∈ (3, 5, 7, 11, 13, 17, 19)
        return factor
    end
    # Use a greedy algorithm to find the largest factor not larger than sqrt(n)
    prime_factor_counts = factor(n)
    prime_factors = sort!(collect(keys(prime_factor_counts)))
    sqrt_n = isqrt(n)
    radix = 1
    for prime_factor in Iterators.reverse(Iterators.filter(≤(sqrt_n), prime_factors))
        for count in 1:prime_factor_counts[prime_factor]
            if radix * prime_factor ≤ sqrt_n
                radix *= prime_factor
            else
                break
            end
        end
    end
    # Always return a radix larger than 1
    radix == 1 && return n
    return radix
end

# X: output
# Y: workspace
# x: input (will be destroyed)
function radix_fft!(X::AbstractVector{T}, Y::AbstractVector{T}, x::AbstractVector{T};
                    choose_radix=choose_radix_sqrt) where {T<:Complex}
    N = length(x)
    @assert length(X) == length(Y) == N
    RT = typeof(real(zero(T)))
    radix = choose_radix(N)
    if N == 0
        # do nothing
    elseif N ≤ radix
        direct_ft!(X, x)
    else
        N₁ = radix              # aka "decimation in time"
        @assert N % N₁ == 0
        N₂ = N ÷ N₁
        x2 = reshape(x, (N₁, N₂))
        X2 = reshape(X, (N₂, N₁))
        Y2 = reshape(Y, (N₂, N₁))
        Z2 = reshape(x, (N₂, N₁))
        for n₁ in 1:N₁
            radix_fft!((@view Y2[:, n₁]), (@view X2[:, n₁]), (@view x2[n₁, :]); choose_radix=choose_radix)
            for n₂ in 1:N₂
                Y2[n₂, n₁] *= cispi(-2 * (n₁ - 1) * (n₂ - 1) / RT(N))
            end
        end
        for n₂ in 1:N₂
            radix_fft!((@view X2[n₂, :]), (@view Z2[n₂, :]), (@view Y2[n₂, :]); choose_radix=choose_radix)
        end
    end
    return X
end

"x will contain output"
radix_fft!(x::AbstractVector; kws...) = radix_fft!(x, similar(x), copy(x); kws...)
"X will contain output, x will be preserved"
radix_fft!(X::AbstractVector, x::AbstractVector; kws...) = radix_fft!(X, similar(x), copy(x); kws...)
"x will be preserved"
radix_fft(x::AbstractVector; kws...) = radix_fft!(similar(x), x; kws...)

################################################################################

export fft!, fft, inv_fft!, inv_fft

"""
    fft!(X::AbstractVector, x::AbstractVector)

Calculate the Fourier transform of `x` and store it into `X`. `x` is
not modified. Both vectors must have the same length and must have
complex element types.

See also: [`fft`](@ref), [`inv_fft!`](@ref).
"""
fft!(X::AbstractVector, x::AbstractVector) = radix_fft!(X, x)

"""
    inv_fft!(x::AbstractVector, X::AbstractVector)

Calculate the inverse Fourier transform of `X` and store it into `x`.
`X` is not modified. Both vectors must have the same length and must
have complex element types.

See also: [`inv_fft`](@ref), [`fft!`](@ref).
"""
function inv_fft!(X::AbstractVector{T}, x::AbstractVector{T}) where {T<:Complex}
    # TODO: Don't modify `x`
    x .= conj(x)
    fft!(X, x)
    x .= conj(x)
    X .= conj(X) / length(X)
    return X
end

"""
    X = fft(x::AbstractVector)
    X::AbstractVector

Calculate the Fourier transform of `x` and return it in a newly
allocate vector. `x` is not modified. The element type of `x` must be
a complex number type.

See also: [`inv_fft`](@ref), [`fft!`](@ref).
"""
fft(x::AbstractVector) = fft!(similar(x), x)

"""
    x = inv_fft(X::AbstractVector)
    x::AbstractVector

Calculate the inverse Fourier transform of `X` and return it in a
newly allocate vector. `X` is not modified. The element type of `X`
must be a complex number type.

See also: [`fft`](@ref), [`inv_fft!`](@ref).
"""
inv_fft(x::AbstractVector) = inv_fft!(similar(x), x)

end


using LinearAlgebra, Statistics, Plots
gr(fmt=:png);

x_vals = [0 0 0; 2 -3 -4]
y_vals = [0 0 0; 4 3 -3.5]

plot(x_vals, y_vals, arrow = true, color = :blue,
     legend = :none, xlims = (-5, 5), ylims = (-5, 5),
     annotations = [(2.2, 4.4, "[2, 4]"),
                    (-3.3, 3.3, "[-3, 3]"),
                    (-4.4, -3.85, "[-4, -3.5]")],
     xticks = -5:1:5, yticks = -5:1:5,
     framestyle = :origin)

x = ones(3)
y = [2,4,6]

dot(x,y)
norm(x)

# fixed linear function, to generate a plane
f(x, y) = 0.2x + 0.1y

# lines to vectors
x_vec = [0 0; 3 3]
y_vec = [0 0; 4 -4]
z_vec = [0 0; f(3, 4) f(3, -4)]

# draw the plane
n = 20
grid1 = range(-5, 5, length = n)
z2 = [ f(grid1[row], grid1[col]) for row in 1:n, col in 1:n ]
wireframe(grid1, grid1, z2, fill = :blues, gridalpha =1 )
plot!(x_vec, y_vec, z_vec, color = [:blue :green], linewidth = 3, labels = "",
      colorbar = false)


# https://i.stack.imgur.com/egDbs.png
"""
Does a math thing?
"""
function extendedeuclidean(a::Int64, b::Int64)::Tuple{Int64,String}
    d0 = (a, 1, 0)
    d1 = (b, 0, 1)

    while d1[1] != 0
        d0, d1 = d1, d0 .- (d0[1] ÷ d1[1]) .* d1
    end

    (d0, "$(d0[1]) = $(a) * $(d0[2]) + $(b) * $(d0[3])")
end


### Logistic function in its general form

# See https://en.wikipedia.org/wiki/Logistic_function
# l scales the curve. For l>1, the curve is stretched, whereas for l<1 it is shrinked
# The curve's maximum value coincides with l
# k gives the curve's steepness. For larger k, the curve becomes more steep
# x0 is the x value of the curve's midpoint
# y0 is the y value of the curve's midpoint
# The logistic function has been defined for tuning purposes

logistic(x::Real, l::Real=1., k::Real=1., x0::Real=0., y0::Real=0.) = l/(1+exp(-k*(x-x0)))+y0


"""
`LinearKernel([ρ=1.0,[c=0.0]])`
The linear kernel is a Mercer kernel given by
```
    κ(x,y) = ρ²xᵀy + c
```
Where `c` is a real number
"""
struct LinearKernel{T,Tr,Tc<:Real} <: Kernel{T,Tr}
    transform::Tr
    metric::DotProduct
    c::Tc
    function LinearKernel{T,Tr,Tc}(transform::Tr,c::Tc) where {T,Tr<:Transform,Tc<:Real}
        return new{T,Tr,Tc}(transform,DotProduct(),c)
    end
end

function LinearKernel(ρ::T₁=1.0,c::T₂=zero(T₁)) where {T₁<:Real,T₂<:Real}
    LinearKernel{T₁,ScaleTransform{T₁},T₂}(ScaleTransform(ρ),c)
end

function LinearKernel(ρ::AbstractVector{T₁},c::T₂=zero(T₁)) where {T₁<:Real,T₂<:Real}
    LinearKernel{T₁,ARDTransform{T₁,length(ρ)},T₂}(ARDTransform(ρ),c)
end

function LinearKernel(t::Tr,c::T=zero(Float64)) where {Tr<:Transform,T<:Real}
    LinearKernel{eltype(t),Tr,T}(t,c)
end

params(k::LinearKernel) = (params(transform(k)),k.c)
opt_params(k::LinearKernel) = (opt_params(transform(k)),k.c)

@inline kappa(κ::LinearKernel, xᵀy::T) where {T<:Real} = xᵀy + κ.c

"""
`PolynomialKernel([ρ=1.0[,d=2.0[,c=0.0]]])`
The polynomial kernel is a Mercer kernel given by
```
    κ(x,y) = (ρ²xᵀy + c)^d
```
Where `c` is a real number, and `d` is a shape parameter bigger than 1
"""
struct PolynomialKernel{T,Tr,Tc<:Real,Td<:Real} <: Kernel{T,Tr}
    transform::Tr
    metric::DotProduct
    c::Tc
    d::Td
    function PolynomialKernel{T,Tr,Tc,Td}(transform::Tr,c::Tc,d::Td) where {T,Tr<:Transform,Tc<:Real,Td<:Real}
        return new{T,Tr,Tc,Td}(transform,DotProduct(),c,d)
    end
end

function PolynomialKernel(ρ::T₁=1.0,d::T₂=2.0,c::T₃=zero(T₁)) where {T₁<:Real,T₂<:Real,T₃<:Real}
    @check_args(PolynomialKernel, d, d >= one(T₁), "d >= 1")
    PolynomialKernel{T₁,ScaleTransform{T₁},T₂,T₃}(ScaleTransform(ρ),c,d)
end

function PolynomialKernel(ρ::AbstractVector{T₁},d::T₂=2.0,c::T₃=zero(T₁)) where {T₁<:Real,T₂<:Real,T₃<:Real}
    @check_args(PolynomialKernel, d, d >= one(T₂), "d >= 1")
    PolynomialKernel{T₁,ARDTransform{T₁,length(ρ)},T₂,T₃}(ARDTransform(ρ),c,d)
end

function PolynomialKernel(t::Tr,d::T₁=2.0,c::T₂=zero(eltype(T₁))) where {Tr<:Transform,T₁<:Real,T₂<:Real}
    @check_args(PolynomialKernel, d, d >= one(T₁), "d >= 1")
    PolynomialKernel{eltype(Tr),Tr,T₁,T₂}(t,c,d)
end

params(k::PolynomialKernel) = (params(transform(k)),k.d,k.c)
opt_params(k::PolynomialKernel) = (opt_params(transform(k)),k.d,k.c)

@inline kappa(κ::PolynomialKernel, xᵀy::T) where {T<:Real} = (xᵀy + κ.c)^(κ.d)


## -------------  Logistic function

"""
	$(SIGNATURES)

Generalized logistic function object.

``
lb + (ub - lb) / (1 + lrShifter * exp(-slope * x)) ^ twister
``

Note that slope can be set to 1 if the input is `x / x0`. But it makes sense to have a slope when x0 depends on the college.

# Parameters
- `lb` and `ub` are the bounds of `f(x)`.
- `lrShifter >= 0` roughly shifts the curve left/right. It really sets the grad prob at h = 0. Higher `lrShifter` implies lower grad probs.
- `slope > 0` is the slope.
- `twister > 0` produces asymmetry (`twister = 1` is symmetric).
"""
@with_kw mutable struct GeneralizedLogistic
    lb :: Double = 0.0
    ub :: Double = 1.0
    lrShifter :: Double = 1.0
    slope :: Double = 1.0
    twister :: Double = 1.0
end

"""
	$(SIGNATURES)

Generalized logistic function.
"""
function logistic(x; lb = 0.0, ub = 1.0, lrShifter = 1.0, slope = 1.0, twister = 1.0)
    return lb .+ (ub - lb) ./ ((1 .+ lrShifter .* exp.(-slope .* x)) .^ twister);
end

function logistic(glf :: GeneralizedLogistic, x)
    return logistic(x; lb = glf.lb, ub = glf.ub, lrShifter = glf.lrShifter,
        slope = glf.slope, twister = glf.twister)
end

# -----------


using LinearAlgebra

# k_local(Elastic Modulus,thickness,height,length,number of nodes,row/column one,row/column two, ==1 adjusts for fixed node)

#creates local stiffness matrix and maps to global
function k_local(E,t,h,l,i,n1,n2,fixed=0)
    k_global = zeros(i,i)
    A = t*h
    k = ((E*A)/l)
    k_global[n1,n1] = k
    k_global[n1,n2] = -k
    k_global[n2,n1] = -k
    k_global[n2,n2] = k
    if fixed == 1
      adj = zeros(1,i)
      adj[n1] = 1
      k_global[n1,:] = adj
      k_global[:,n1] = adj
    end
    return k_global
end

# constants from problem statement
E = 68.9*10^9
t = 0.006
h13 = 0.08
h2 = 0.02
l = 0.025
l2 = 0.1

# combining all globally mapped matrices
kG = k_local(E,t,h13,l,4,1,2,1)+
     k_local(E,t,h2,l2,4,2,3)+
     k_local(E,t,h13,l,4,3,4)

# Force Vector
f = Float64[0 0 0 1800]
F = transpose(f)

# Displacement vector
u = kG\F
println("Nodal Displacement Vector (meters)")
display(u)

# Strain Function
ϵ(x2,x1,l) = (x2 - x1)/l

#Strain Vector
ϵM = [ϵ(u[2],u[1],l)
      ϵ(u[3],u[2],l2)
      ϵ(u[4],u[3],l)]

# Stress Function
σ(E,ϵ) = E*ϵ

# Stress Vector
σM = [σ(E,ϵM[1])
      σ(E,ϵM[2])
      σ(E,ϵM[3])]
println("Element Stress Vector (Pa)")
display(σM)


function fourfilt(x,delt,tmax,tmin)
#  FOURFILT  Fourier low, high, or bandpass filter.
#
#     [filtdat]=fourfilt(x,delt,tmax,tmin)
#
#     where:   x:      data series to be filtered
#              delt:   sampling interval
#              tmax:   maximum period filter cutoff
#              tmin:   minimum period filter cutoff
#
#     usage:  lowpassdata=fourfilt(data,0.5,2000,20)
#
#               gives lowpass filter with cutoff at 20.0 sec
#               tmax set > (length(x)*delt) for no cutoff at low freq end
#
#     usage:  highpassdata=fourfilt(x,0.5,20,0.9)
#
#               gives highpass filter with cutoff at 20.0 sec
#               tmin set < (2*delt) for no cutoff at high freq end
#
#     usage:  bandpassdata=fourfilt(x,0.5,20,10)
#
#               gives bandpass filter passing 10-20 sec. band
#
#     REQUIREMENTS: Julia "FFTW" package
#
# Reference:
# Walters, R. A.  and Heston, C., 1982. Removing the tidal-period variations from time-series
# data using low-pass digital filters. Journal of Physical Oeanography, 12 112-115 .
#
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# Julia version that has been translated from MATLAB version
# by William Pringle on 08/21/2019
# MATLAB verion is available from Rich Signell's "RPSstuff" github
#
# Version 1.0  (12/4/96)  Jeff List (jlist@usgs.gov)
# Version 1.1  (1/8/97)  Rich Signell (rsignell@usgs.gov)
#     removed argument for number of points and add error trapping for matrices
# Version 1.1b (12/1/2005) Rich Signell (rsignell@usgs.gov)
#     added reference
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

global f

if ndims(x) != 1
  display("fourfilt cannot handle matrices");
  return
end
if sum(isnan.(x)) > 0
  display("fourfilt will not work with nans in timeseries")
  return
end

npts = length(x);
nby2 = convert(Int,floor(npts/2));
tfund = npts*delt;
ffund = 1.0/tfund;

#  remove the mean from data:
datamean = mean(x);
x = x .- datamean;

# fourier transform data:
coeffs = fft(x);

#  filter coefficients:
f = ffund
for i = 2:nby2+1
  t = 1/f
  if t > tmax || t < tmin
     coeffs[i] = 0.0*coeffs[i];
  end
  f += ffund
end

#  calculate the remaining coefficients:
for i = 2:nby2
   coeffs[npts+2-i] = conj(coeffs[i])
end

#  backtransform data and take real part:
backcoeffs = ifft(coeffs);
filtdat    = real(backcoeffs);

# add back the mean:
return filtdat .+ datamean

end


using Distributions

# Use Monte Carlo methods to estimate the divergences between two
# theoretical distributions. Take random samples of source distribution.
# Estimate cost at those points and weight by probability.

function mc_divergence_estimate(d1::Distribution, d2::Distribution, samples::Integer, sample_divergence_func::Function)
    divergence = 0.0
    for i in 1:samples
        x = rand(d1)
        divergence += sample_divergence_func(pdf(d1, x), pdf(d2, x))
    end
    return divergence
end

# non-symmetric (i.e. conventional) divergence for a single sample

function kl_sample_divergence(dist1_probab::Float64, dist2_probab::Float64)
    return -dist1_probab * log(dist2_probab / dist1_probab)
end

# symmetrized divergence for a single sample

function skl_sample_divergence(dist1_probab::Float64, dist2_probab::Float64)
    return kl_sample_divergence(dist1_probab,dist2_probab) + kl_sample_divergence(dist2_probab,dist1_probab)
end

# Jensen-Shannon divergence for a single sample (symmetric)

function js_sample_divergence(dist1_probab::Float64, dist2_probab::Float64)
    const mean_probab::Float64 = (dist1_probab + dist2_probab) / 2
    return (kl_sample_divergence(dist1_probab, mean_probab) + kl_sample_divergence(dist2_probab, mean_probab)) / 2
end

kl_divergence(d1::Distribution, d2::Distribution) = mc_divergence_estimate(d1, d2, 10_000, kl_sample_divergence)
skl_divergence(d1::Distribution, d2::Distribution) = mc_divergence_estimate(d1, d2, 10_000, skl_sample_divergence)
js_divergence(d1::Distribution, d2::Distribution) = mc_divergence_estimate(d1, d2, 10_000, js_sample_divergence)


# (weighted) sum of L2 norm and L1 norm
# for Group Lasso, this can be used together with src/calculus/slicedSeparableSum
export NormL1plusL2

"""
**L1 norm plus L2 norm**

    NormL1plusL2(λ_1=1, λ_2=1)

With two nonegative scalars λ_1 and λ_2, returns the function
```math
f(x) = λ_1 ∑_{i=1}^{n} |x_i| + λ_2 \\sqrt{x_1^2 + … + x_n^2}.
```
With nonnegative array λ_1 and nonnegative scalar λ_2, returns the function
```math
f(x) =  ∑_{i=1}^{n} {λ_1}_i |x_i| + λ_2 \\sqrt{x_1^2 + … + x_n^2}.
```
"""
struct NormL1plusL2{L1<:NormL1, L2 <: NormL2} <: ProximableFunction
    l1::L1
    l2::L2
end

is_separable(f::NormL1plusL2) = false
is_convex(f::NormL1plusL2) = true
is_positively_homogeneous(f::NormL1plusL2) = true

function NormL1plusL2(lambda1::L=1, lambda2::M=1) where {L <: Union{Real, AbstractArray}, M <: Real}
    NormL1plusL2(NormL1(lambda1), NormL2(lambda2))
end

function (f::NormL1plusL2)(x::AbstractArray{T}) where T <: RealOrComplex
    return f.l1(x) + f.l2(x)
end

function prox!(y::AbstractArray{T}, f::NormL1plusL2, x::AbstractArray{T}, gamma::Real=1) where T <: RealOrComplex
    prox!(y, f.l1, x, gamma)
    prox!(y, f.l2, y, gamma)
    return f(y)
end

fun_name(f::NormL1plusL2) = "L2-norm + L1-norm"
fun_dom(f::NormL1plusL2) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::NormL1plusL2{L1,L2}) where {L1<:NormL1{<:Real}, L2} = "x ↦ λ_1 ||x||_1 + λ_2||x||_2"
fun_expr(f::NormL1plusL2{L1,L2}) where {L1<:NormL1{<:AbstractArray}, L2} = "x ↦ sum( (λ_1)_i |x_i| ) + λ_2||x||_2"
fun_params(f::NormL1plusL2{L1,L2}) where {L1<:NormL1{<:Real}, L2}  = "λ_1 = $(f.l1.lambda), λ_2 = $(f.l2.lambda)"
fun_params(f::NormL1plusL2{L1,L2}) where {L1<:NormL1{<:AbstractArray}, L2}  = "λ_1 = $(typeof(f.lambda)) of size $(size(f.lambda)), λ_2 = $(f.l2.lambda)"

function prox_naive(f::NormL1plusL2, x::AbstractArray{T}, gamma::Real=1) where T <: RealOrComplex
    y1, v1 = prox_naive(f.l1, x, gamma)
    y2, v2 = prox_naive(f.l2, y1, gamma)
    return y2, f(y2)
end


abstract type Kernel end

struct PolynomialKernel <: Kernel
    # return (x^T @ y)^M
    _degree::Int64 # degree of polynomial feature
    _intercept::Float64 # intercept
    function PolynomialKernel(degree::Int64, intercept = 0.0)
        new(degree, intercept)
    end
end

"""
    kernel(kernel::PolynomialKernel, x::AbstractArray{Float64, 2}, y::AbstractArray{Float64, 2})

For x = [x_1,... x_N] of size (n_feature, n_samples1) and y = [y_1,... y_N] of size (n_feature, n_samples2),
 compute the gram matrix K of size(n_samples1, n_samples2), where K_ij =(x_i * y_j)^M.
"""
function kernel(
    kernel::PolynomialKernel,
    x::AbstractArray{Float64,2},
    y::AbstractArray{Float64,2},
)
    n_samples1 = size(x)[2]
    n_samples2 = size(y)[2]
    n_feature = size(x)[1]

    intercept = kernel._intercept
    degree = kernel._degree

    n_feature == size(y)[1] || error("feature dimension does not match")
    ret = zeros(n_samples1, n_samples2)
    for i = 1:n_samples1
        for j = 1:n_samples2
            x_i = view(x, :, i)
            y_j = view(y, :, j)
            ret[i, j] = (sum(x_i .* y_j) + intercept)^degree
        end
    end

    return ret
end

"Kernel to compute c0 * exp(-0.5 * c1^T * (x1 - x2)^2) where c0 is scalar and c1 is Array{Float64, 1}"
mutable struct RBFKernel <: Kernel
    _params::Array{Float64,1}
    _n_params::Int64
    function RBFKernel(params::Array{Float64,1})
        new(copy(params), size(params)[1])
    end
end

function rbf(
    c0::Float64,
    c1::AbstractArray{Float64,1},
    x1::AbstractArray{Float64,1},
    x2::AbstractArray{Float64,1},
)
    d = (x1 - x2) .^ 2
    return c0 * exp(-0.5 * transpose(c1) * d)
end

"""
    kernel(kernel::PolynomialKernel, x::AbstractArray{Float64, 2}, y::AbstractArray{Float64, 2})

For x = [x_1,... x_N] of size (n_feature, n_samples1) and y = [y_1,... y_N] of size (n_feature, n_samples2),
 compute the gram matrix K of size(n_samples1, n_samples2), where K_ij = rbf(x_i, x_j).
"""
function kernel(kernel::RBFKernel, x::AbstractArray{Float64,2}, y::AbstractArray{Float64,2})
    n_feature = size(x)[1]
    n_samples1 = size(x)[2]
    n_samples2 = size(y)[2]
    n_feature == size(y)[1] || error("size does not match")
    n_feature == kernel._n_params - 1 || error("kernel size does not match")

    ret = zeros(n_samples1, n_samples2)
    for i = 1:n_samples1
        x_i = x[:, i]
        for j = 1:n_samples2
            y_j = y[:, j]
            ret[i, j] = rbf(kernel._params[1], kernel._params[2:end], x_i, y_j)
        end
    end

    return ret
end

"""
    derivatives(kernel::RBFKernel, x::AbstractArray{Float64, 2}, y::AbstractArray{Float64, 2})

Return the derivative of K(gram matrix) with respect to c0 and c1 as Array{Float64, 3}
"""
function derivatives(
    kernel::RBFKernel,
    x::AbstractArray{Float64,2},
    y::AbstractArray{Float64,2},
)
    n_feature = size(x)[1]
    n_samples1 = size(x)[2]
    n_samples2 = size(y)[2]
    n_feature == size(y)[1] || error("size does not match")
    n_feature == kernel._n_params - 1 || error("kernel size does not match")
    n_params = kernel._n_params

    ret = zeros(n_params, n_samples1, n_samples2)
    c0 = kernel._params[1]
    c1 = kernel._params[2:end]
    for i = 1:n_samples1
        for j = 1:n_samples2
            x_i = x[:, i]
            y_j = y[:, j]
            d = (x_i - y_j) .^ 2
            delta_c0 = exp(-0.5 * sum(c0 * d))
            delta_c1 = -0.5 * d * c0 * delta_c0
            ret[:, i, j] = vcat([delta_c0], delta_c1)
        end
    end

    return ret
end

function update_parameters(kernel::RBFKernel, updates_to_add::AbstractArray{Float64,1})
    kernel._params += updates_to_add
end


#-------------monte_carlo.jl---------------------------------------------------#
#
# Purpose: A simple monte carlo script to integrate x^2 in a given range
#
#------------------------------------------------------------------------------#

# A function to integrate x^2 in a given range with num number of points
function monte_carlo(range::Float64, num::Int64)
    count = 0
    for i = 1:num
        x_pos = rand() * range - range * 0.5
        y_pos = rand() * 0.25 * range * range

        if (y_pos < x_pos*x_pos)
            count += 1
        end
    end

    #println("The integral of x^2 between ", -range*0.5, " and ",
            #range*0.5, " is ", (count / num) *0.25*range*range*range)
end

monte_carlo(6.0, 10000)


module MadMax

using PrettyTables
using Flux
using LinearAlgebra

export BisectionEval
export RegulaFalsiEval
export NewtonEval
export GaussJacobiEval
export GaussSeidelEval
export SOREval

function getLowerTriangular(X)
    R = zeros(size(X)...)
    for i in 1:size(X)[1]
        for j in 1:size(X)[2]
            if i>j
                R[i,j] = X[i, j]
            end
        end
    end
    return R
end

function getUpperTriangular(X)
    R = zeros(size(X)...)
    for i in 1:size(X)[1]
        for j in 1:size(X)[2]
            if i<j
                R[i,j] = X[i, j]
            end
        end
    end
    return R
end

function SOREval(X, x0, ω=1; tol=0.005, N=5, verbose=true)
    A = X[:, 1:end-1]
    b = -X[:, end]

    D = Diagonal(A)
    U = getUpperTriangular(A)
    L = getLowerTriangular(A)

    H = inv(D+ω*L) * ((1-ω)D - ω*U)
    c = ω*inv(D+ω*L)*b

    x = x0
    approximations = [x0]

    n = 0
    while n != N
        x = H*x + c
        approximations = [approximations; [x]]
        n += 1
    end

    return x, approximations
end

function GaussJacobiEval(X, x0; tol=0.005, N=5, verbose=true)
    A = X[:, 1:end-1]
    b = -X[:, end]

    D = Diagonal(A)
    U = getUpperTriangular(A)
    L = getLowerTriangular(A)

    H = -inv(D) * (L+U)
    c = inv(D)*b

    x = x0
    approximations = [x0]

    n = 0
    while n != N
        x = H*x + c
        approximations = [approximations; [x]]
        n += 1
    end

    return x, approximations
end

function GaussSeidelEval(X, x0; tol=0.005, N=5, verbose=true)
    A = X[:, 1:end-1]
    b = -X[:, end]

    D = Diagonal(A)
    U = getUpperTriangular(A)
    L = getLowerTriangular(A)

    H = -inv(L+D)*U
    c = inv(L+D)*b

    x = x0
    approximations = [x0]

    n = 0
    while n != N
        x = H*x + c
        approximations = [approximations; [x]]
        n += 1
    end

    return x, approximations
end

function NewtonEval(f::Function, x0; tol = 0.005, N=100, verbose=true)
    error = Base.Inf64

    # find gradient
    x = 0
    df(x) = gradient(f, x)[1]

    approximations = Array{Float64}(undef, 1)
    approximations = [x0]
    errors = [error]
    n = 0

    # iterate
    while error > tol && n != N
        x_last = last(approximations)
        delta = f(x_last) / df(x_last)
        x_n = x_last - delta

        approximations = [approximations; x_n]

        # get error
        error = abs(x_last - x_n)
        errors = [errors; error]

        n += 1
    end

    if verbose
        headers = ["S.N" "x_n" "Error"]
        data = hcat(0:n, approximations, errors)

        pretty_table(data, headers)
        println("$n iterations")
    end

    return approximations, error
end

function BisectionEval(f::Function, a, b; tol = 0.005, N = 100, verbose=true)
    # make expr
    error = Base.Inf64

    # find starting interval
    interval = (a, b)

    # iterate
    approximations = []
    errors = []
    n = 0

    if verbose
        as = []
        bs = []
        prod_s =[]
    end

    while error > tol && n!=N
        # get m_k
        m_k = sum(interval)/2

        # get error
        if length(approximations) > 0
            error = abs(last(approximations) - m_k)
        end
        push!(errors, error)

        if verbose
            push!(as, interval[1])
            push!(bs, interval[2])
        end

        # append the appproximation
        push!(approximations, m_k)

        # get new interval
        interval = (f(interval[1])*f(m_k)) < 0 ? (interval[1], m_k) : (m_k, interval[2])

        prod = (f(interval[1])*f(m_k)) < 0 ? "<0" : ">0"
        verbose && push!(prod_s, prod)

        n += 1
    end

    if verbose
        data = hcat(1:n, as, bs, approximations, prod_s, errors)
        header = ["S.No" "a_k-1" "b_k-1" "m_k" "f(a_k-1)*f(m_k)" "error"]
        pretty_table(data, header)
        println("$n iterations")
    end

    return approximations, errors
end

function RegulaFalsiEval(f::Function, a, b; tol = 0.005, N = 100, verbose=true)
    # make expr
    error = Base.Inf64

    # find starting interval
    interval = (a, b)

    # iterate
    approximations = []
    errors = []
    n = 0

    if verbose
        as = []
        bs = []
        f_as = []
        f_bs = []
        prod_s =[]
    end

    while error > tol && n<=N
        # get m_k

        m_k = (interval[1]*f(interval[2]) - interval[2]*f(interval[1])) / (f(interval[2]) - f(interval[1]))

        if verbose
            push!(as, interval[1])
            push!(bs, interval[2])
            push!(f_as, f(interval[1]))
            push!(f_bs, f(interval[2]))
        end

        # get error
        if length(approximations) > 0
            error = abs(last(approximations) - m_k)
        end
        push!(errors, error)

        # append the appproximation
        push!(approximations, m_k)

        # get new interval
        interval = (f(interval[1])*f(m_k)) < 0 ? (interval[1], m_k) : (m_k, interval[2])

        prod = (f(interval[1])*f(m_k)) < 0 ? "<0" : ">0"
        verbose && push!(prod_s, prod)

        n += 1
    end

    if verbose
        data = hcat(1:n, as, bs, f_as, f_bs, approximations, prod_s, errors)
        header = ["S.No" "a_k-1" "b_k-1" "f(a)" "f(b)" "m_k" "f(a_k-1)*f(m_k)" "error"]
        pretty_table(data, header)
        println("$n iterations")
    end

    return approximations, errors
end

end


# # Subspace angles
# This example considers the angles between neighbouring Laguerre polynomials with a perturbed measure:
# ```math
# \cos\theta_n = \frac{\langle L_n, L_{n+k}\rangle}{\|L_n\|_2 \|L_{n+k}\|_2},\quad{\rm for}\quad 0\le n < N-k,
# ```
# where the inner product is defined by $\langle f, g\rangle = \int_0^\infty f(x) g(x) x^\beta e^{-x}{\rm\,d}x$.
#
# We do so by connecting Laguerre polynomials to the normalized generalized Laguerre polynomials associated with the perturbed measure. It follows by the inner product of the connection coefficients that:
# ```math
# \cos\theta_n = \frac{(V^\top V)_{n, n+k}}{\sqrt{(V^\top V)_{n, n}(V^\top V)_{n+k, n+k}}}.
# ```
#
using FastTransforms, LinearAlgebra

# The neighbouring index `k` and the maximum degree `N-1`:
k, N = 1, 11

# The Laguerre connection parameters:
α, β = 0.0, 0.125

# We precompute a Laguerre--Laguerre plan:
P = plan_lag2lag(Float64, N, α, β; norm2=true)

# We apply the plan to the identity, followed by the adjoint plan:
VtV = P*I
lmul!(P', VtV)

# From this matrix, the angles are recovered from:
θ = [acos(VtV[n, n+k]/sqrt(VtV[n, n]*VtV[n+k, n+k])) for n in 1:N-k]


using LinearAlgebra
using SparseArrays
using Plots
using ForwardDiff
using FastGaussQuadrature # more accurate quadrature

"This routine solves Poisson's equation -u''(x) = f(x) in weak form using piecewise
linear finite element methods. It's used to illustrate convergence in different norms."

m = 200 # number of elements

# Manufactured solution
uexact(x) = log(2+sin(4*pi*x))
dudx_exact(x) = ForwardDiff.derivative(uexact,x)
f(x) = -ForwardDiff.derivative(dudx_exact,x) # -d/dx (κ(x)*du/dx) = f
α = uexact(-1)
β = uexact(1)

# define spatial grid
x = LinRange(-1,1,m+2) # x_0, x_1, ..., x_m, x_{m+1} = x_0
# x = @. x + randn() / (2*m) # randomly perturb nodes
# x[1] = -1; x[end] = 1 # reset endpoints

# define local FEM basis
rq,wq = gausslegendre(100) # define overkill-accurate Gauss quadrature
λ(r) = [(1 .-r)./2 (1 .+r)./2]
#dλ(r) = [-.5*ones(length(r)) .5*ones(length(r))]
dλ(r) = [-.5 .5]
map_point(x,a,b) = a + (b-a) * (1+x)/2 # maps x ∈ [-1,1] to interval [a,b]

"∫u' * ϕ_i' = ∫f(x)*ϕ_i(x)"
A = spzeros(m+2,m+2)
A_local = [1 -1;-1 1]
b = zeros(m+2)
for e = 1:m+1 # loop over the intervals
    ids = e:e+1
    h_e = x[e+1]-x[e]

    x_midpt = .5*(x[e+1]+x[e])
    A[ids,ids] .+= A_local * 1/h_e

    # # one point midpoint quadrature
    # b[ids] .+= h_e / 2.0 * f(.5*(x[e]+x[e+1])) # f(.5*(x_i+x_i+1))

    # higher accuracy Gauss quadrature
    xq = map_point.(rq,x[e],x[e+1])
    b[ids] .+= h_e / 2.0 * λ(rq)'*(wq.*f.(xq)) # ∑_k f(x_k) * w_k * ϕ_i(x_k)
end

# impose Dirichlet BCs on both sides
b̃ = b - vec(A[:,1]*α)
b̃[1] = α
A[1,:] .= 0
A[:,1] .= 0
A[1,1] = 1

b̃ = b̃ - vec(A[:,m+2]*β)
b̃[m+2] = β
A[m+2,:] .= 0
A[:,m+2] .= 0
A[m+2,m+2] = 1

# convert b̃ from a SparseVector to a Vector
# (Julia doesn't currently allow "\" to be applied to SparseVectors)
u = A\Vector(b̃)

nodal_err = abs.(u - uexact.(x))
plot(x,u,mark=:dot,ms=3,legend=false,title="Max error at nodes = $(maximum(nodal_err))")
xfine = LinRange(-1,1,1000)
plot!(xfine,uexact.(xfine),legend=false)

# compute H1 and L2 errors
function compute_error(u,m,x,rq,wq)
    H1_err2 = 0.0
    L2_err2 = 0.0
    for e = 1:m+1
        ids = e:e+1
        h_e = x[e+1]-x[e]
        rx = 2 / h_e
        xq = map_point.(rq,x[e],x[e+1])
        dudx_err = rx * dλ(rq)*u[ids] .- dudx_exact.(xq)
        H1_err2 += h_e / 2 * dot(wq, dudx_err.^2)

        u_err = λ(rq)*u[ids] - uexact.(xq) # ∫(u-uexact)^2
        L2_err2 += h_e / 2 * dot(wq, u_err.^2)
    end
    return sqrt(H1_err2), sqrt(L2_err2)
end
H1_err,L2_err = compute_error(u,m,x,rq,wq)
@show H1_err # will be proportional to O(h)
@show L2_err


export clenshaw_curtis

"""
    clenshaw_curtis(N)

Compute the nodes `x` and weights `w` for integrating a
continuous functions from [-1,1] using the Clenshaw-Curtis
integration rule with order `N`.
"""
function clenshaw_curtis(N::Int64)

    @assert N>0 "Order must be greater than 0"

    x = zeros(N)
    w = ones(N)

    if N==1
        #x[1]=0.0
        w[1] = 2.0
        return x, w
    end

    # Compute x values
    @inbounds for i=1:N
        x[i] = cos((N-i)*π/(N-1))
    end

    n = N
    x[1] = -1.0

    if mod(n,2)==1
        x[ceil(Int64, (n+1)/2)] = 0.0
    end

    x[n] = 1.0

    @inbounds for i=1:n
        θ = (i-1)*π/(n-1)
        for j=1:ceil(Int64,(n-1)/2)
            if 2*j == n-1
                b = 1.0
            else
                b = 2.0
            end

            w[i] -= b*cos(2*j*θ)/(4*j*j-1)
        end
    end

    w[1] *= 1/(n-1)
    w[2:n] .*= 2.0/(n-1)
    w[n] *= 1/(n-1)

    return x, w
end


# Description:
# Dimensions: 2

# The plot on the left shows the three-hump Camel function on its recommended input domain, and the plot on the right shows only a portion of this domain, to allow for easier viewing of the function's key characteristics. The function has three local minima.

# Input Domain:
# The function is usually evaluated on the square xi ∈ [-5, 5], for all i = 1, 2.

# Global Minimum:
# https://www.sfu.ca/~ssurjano/camel33.png

function threeHumpCamel(xx)

    x1 = xx[1];
    x2 = xx[2];

    term1 = 2 * x1^2;
    term2 = -1.05 * x1^4;
    term3 = x1^6 / 6;
    term4 = x1 * x2;
    term5 = x2^2;

    y = term1 + term2 + term3 + term4 + term5;

    return [y]
end


require("bspline.jl")

function exercise_zero()
    controlPoints = [
                     [-0.2 2];
                     [-0.3 6.2];
                     [-1.2 4.8];
                     [-2.8 8.8];
                     [-0.7 14];
                     [1.4 14.7];
                     [3.6 10.2];
                     [3.2 5.1];
                     [1.5 6.2];
                     [1.4 2];
                     ]

    ## continuityVector = ones(length(controlPoints[:,1]))
    k = 4
    extendedPartition = extendedPartitionAttempt(k, ones(6), Clumped())
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(controlPoints,
                         "exercise-zero-clumped-control-poly.coordinates")

    writeArrayForGnuplot(bspline,
                         "exercise-zero-clumped-bspline.coordinates")

    extendedPartition = extendedPartitionAttempt(k, ones(6), Uniformed())
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(controlPoints,
                         "exercise-zero-uniformed-control-poly.coordinates")

    writeArrayForGnuplot(bspline,
                         "exercise-zero-uniformed-bspline.coordinates")

    closedControlPoints = [controlPoints; controlPoints[1,:];]
    closedBSpline = drawBSpline(k, closedControlPoints, 200)
    writeArrayForGnuplot(closedControlPoints,
                         "exercise-zero-closed-control-poly.coordinates")

    writeArrayForGnuplot(closedBSpline,
                         "exercise-zero-closed-bspline.coordinates")

end

function exercise_one()
    controlPoints = [
                     [-2.5 -1];
                     [-3 2];
                     [0.3 2.0];
                     [-0.3 -2.0];
                     [4 -2.0];
                     [3 1.0];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-one-control-poly.coordinates")

    fittingPoints = 200

    k = 2
    continuityVector = ones(4)
    k_two_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_two_extendedPartition)
    k_two_bspline = drawCurve(k_two_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_two_bspline,
                         "exercise-one-clumped-bspline-k-two.coordinates")

    k = 3
    continuityVector = ones(3)
    k_three_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_three_extendedPartition)
    k_three_bspline = drawCurve(k_three_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_three_bspline,
                         "exercise-one-clumped-bspline-k-three.coordinates")

    k = 4
    continuityVector = ones(2)
    k_four_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_four_extendedPartition)
    k_four_bspline = drawCurve(k_four_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_four_bspline,
                         "exercise-one-clumped-bspline-k-four.coordinates")

    k = 5
    continuityVector = ones(1)
    k_five_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_five_extendedPartition)
    k_five_bspline = drawCurve(k_five_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_five_bspline,
                         "exercise-one-clumped-bspline-k-five.coordinates")

    k = 6
    continuityVector = ones(0)
    k_six_extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, k_six_extendedPartition)
    k_six_bspline = drawCurve(k_six_extendedPartition,
                              k,
                              M,
                              controlPoints,
                              matrix,
                              fittingPoints)
    writeArrayForGnuplot(k_six_bspline,
                         "exercise-one-clumped-bspline-k-six.coordinates")

end

function exercise_two()
    controlPoints = [
                     [-0.7 -0.4];
                     [-1 -1];
                     [-1.3 -0.3];
                     [-1 1.2];
                     [-0.3 1.4];
                     [-0.15 0];
                     [0.15 0];
                     [0.3 1.6];
                     [2 0.6];
                     [2.6 0.2];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-two-control-poly.coordinates")

    k = 4
    extendedPartition = expandPartition(0:10, [4 [1 for _ in 1:10]'])
    extendedPartition = normalizeRespect(extendedPartition, Maximum(), k)
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-first.coordinates")

    k = 4
    extendedPartition = expandPartition(0:4, [4 2 2 2 4])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(6))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-second.coordinates")

    k = 6
    extendedPartition = expandPartition(0:5, [6 1 1 1 1 6])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(4))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-three.coordinates")

    k = 6
    extendedPartition = expandPartition(0:4, [6 1 1 2 6])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(4))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-four.coordinates")

    k = 8
    extendedPartition = expandPartition(0:3, [8 1 1 8])
    extendedPartition = normalizeRespect(extendedPartition, Maximum())
    M = sum(ones(2))
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    fittingPoints = 200
    bspline = drawCurve(extendedPartition, k, M, controlPoints, matrix, fittingPoints)

    writeArrayForGnuplot(bspline,
                         "exercise-two-bspline-five.coordinates")

end

function exercise_three()
    controlPoints = [
                     [-3.0 2];
                     [-1.0 -2];
                     [1 1.0];
                     [1 1.0];
                     [3.0 -2];
                     [5 3.0];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-three-control-poly.coordinates")

    fittingPoints = 200

    k = 3
    continuityVector = ones(3)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-three-bspline.coordinates")
end

function exercise_four()
    controlPoints = [
                     [-2.0 2];
                     [-3.0 4];
                     [0 6];
                     [0 6];
                     [3 4];
                     [2.0 2];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-four-control-poly.coordinates")

    fittingPoints = 200

    k = 4
    ## continuityVector = ones(3)
    ## extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    extendedPartition = [0 0 0 0 1/4 3/4 1 1 1 1]
    M = 2
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-four-bspline-first-partition.coordinates")

    extendedPartition = [0 0 0 0 1/2 1/2 1 1 1 1]
    M = 2
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-four-bspline-second-partition.coordinates")

end

function exercise_five()
    controlPoints = [
                     [-2.0 1];
                     [-1.0 -1];
                     [0 1];
                     [1 -1];
                     [2.0 1];
                     ]

    writeArrayForGnuplot(controlPoints,
                         "exercise-five-control-poly.coordinates")

    fittingPoints = 200

    k = 4
    continuityVector = ones(1)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-five-bspline-single.coordinates")

    controlPoints = [
                     [-2.0 1];
                     [-1.0 -1];
                     [0 1];
                     [0 1];
                     [1 -1];
                     [2.0 1];
                     ]

    fittingPoints = 200

    k = 4
    continuityVector = ones(2)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-five-bspline-doubled.coordinates")

    controlPoints = [
                     [-2.0 1];
                     [-1.0 -1];
                     [0 1];
                     [0 1];
                     [0 1];
                     [1 -1];
                     [2.0 1];
                     ]

    fittingPoints = 200

    k = 4
    continuityVector = ones(3)
    extendedPartition = extendedPartitionAttempt(k, continuityVector, Clumped())
    M = sum(continuityVector)
    matrix = buildFunctionsMatrix(k, M, extendedPartition)
    bspline = drawCurve(extendedPartition,
                        k,
                        M,
                        controlPoints,
                        matrix,
                        fittingPoints)
    writeArrayForGnuplot(bspline,
                         "exercise-five-bspline-tripled.coordinates")

end

function exercise_six()

    k = 4

    firstClosedControlPoints = [
                                [1.0 4];
                                [-0.8 4];
                                [-1.9 8.2];
                                [0.3 13.2];
                                [2.8 4];
                                [1.0 4];
                                ]

    firstClosedBSpline = drawBSpline(k, firstClosedControlPoints, 200)
    writeArrayForGnuplot(firstClosedControlPoints,
                         "exercise-six-first-closed-control-poly.coordinates")

    writeArrayForGnuplot(firstClosedBSpline,
                         "exercise-six-first-closed-bspline.coordinates")

    secondClosedControlPoints = [
                                 [0.4 6];
                                 [-1.8 3];
                                 [-2 14.8];
                                 [0.4 11];
                                 [3.0 14];
                                 [3.4 2.2];
                                 [0.4 6];
                                 ]

    secondClosedBSpline = drawBSpline(k, secondClosedControlPoints, 200)
    writeArrayForGnuplot(secondClosedControlPoints,
                         "exercise-six-second-closed-control-poly.coordinates")

    writeArrayForGnuplot(secondClosedBSpline,
                         "exercise-six-second-closed-bspline.coordinates")

end

function exercise_knotsInsertion()
    controlPoints = [
                     [0.0 0.0];
                     [1.0 2.0];
                     [2.0 2.5];
                     [3.0 2.2];
                     [4.0 0.5];
                     ]

    partition = [-1; 0.5; 1.8; 2.8; 3.2; 4;]

    k = 3
    refined_controlPoints =
        knotsInsertion_head(controlPoints, partition, k, 1)

    writeArrayForGnuplot(controlPoints,
                         "exercise-knots-insertion-original-control-points.coordinates")

    writeArrayForGnuplot(refined_controlPoints,
                         "exercise-knots-insertion.coordinates")

    ## deBoorBSpline = drawBSpline(k, controlPoints, 200)
    ## writeArrayForGnuplot(deBoorBSpline,
    ##                      "exercise-knots-insertion-deBoor.coordinates")

end


using QuadGK

function main()

    f(x) = x^2
    g(x) = x^4
    h(x) = (x+1)^(-1)

    a = 0
    b = 2
    n = 10

    println(simpsons(f, a, b, n))
    println(quadgk(f, a, b, rtol=1e-5))

    println(simpsons(g, a, b, n))
    println(quadgk(g, a, b, rtol=1e-5))

    println(simpsons(h, a, b, n))
    println(quadgk(h, a, b, rtol=1e-5))
end

function simpsons(df, a, b, n)

    # n should be even
    h = (b - a) / n

    xi0 = df(a) + df(b)         # end points
    xi1 = 0                     # odd points
    xi2 = 0                     # even points

    # performing integration
    for i in 1:n-1

        x = a + i*h

        # if even add to xi2
        if i%2 == 0
            xi2 += df(x)
        else
            # else odd add to xi1
            xi1 += df(x)
        end
    end

    # add everything togeteher
    xi = h*(xi0 + 2*xi2 + 4*xi1)/3

    return xi

end

main()


function eval(membership_function::TriangularMF, x::T) where {T <: Number}
    return maximum([minimum([((x - membership_function.l_vertex) / (membership_function.center - membership_function.l_vertex)), ((membership_function.r_vertex - x) / (membership_function.r_vertex - membership_function.center))]), 0])
end

function mean_at(membership_function::TriangularMF, firing_strength::T) where {T <: Number}
    if firing_strength != 1
        p1 = (membership_function.center - membership_function.l_vertex) * firing_strength + membership_function.l_vertex
        p2 = (membership_function.center - membership_function.r_vertex) * firing_strength + membership_function.r_vertex
        (p1 + p2) / 2
    elseif firing_strength == 1
        return membership_function.center
    end
end

function eval(menbership_function::GaussianMF, x::T) where {T <: Number}
    return exp(- 0.5 * ((x - menbership_function.center) / menbership_function.sigma)^2)
end

function mean_at(membership_function::GaussianMF, firing_strength::T) where {T <: Number}
    return membership_function.center
end

function eval(membership_function::BellMF, x::T) where {T <: Number}
    return (1 / (1 + abs((x - membership_function.c) / membership_function.a)^(2 * membership_function.b)))
end

function mean_at(membership_function::BellMF, firing_strength::T) where {T <: Number}
    return membership_function.c
end

function eval(membership_function::TrapezoidalMF, x::T) where {T <: Number}
    return  maximum([minimum([((x - membership_function.l_bottom_vertex) / (membership_function.l_top_vertex - membership_function.l_bottom_vertex)), 1, ((membership_function.r_bottom_vertex - x) / (membership_function.r_bottom_vertex - membership_function.r_top_vertex))]), 0])
end

function mean_at(membership_function::TrapezoidalMF, firing_strength::T) where {T <: Number}
    p1 = (membership_function.l_top_vertex - membership_function.l_bottom_vertex) * firing_strength + membership_function.l_bottom_vertex
    p2 = (membership_function.r_top_vertex - membership_function.r_bottom_vertex) * firing_strength + membership_function.r_bottom_vertex
    (p1 + p2) / 2
end

function eval(membership_function::SigmoidMF, x::T) where {T <: Number}
    return  1 / (1 + exp(-membership_function.a * (x - membership_function.c)))
end

function mean_at(membership_function::SigmoidMF, firing_strength::T) where {T <: Number}
    if firing_strength == 1
        p_firing_strength = 0.999
    elseif firing_strength == 0
        p_firing_strength = 0.001
    else
        p_firing_strength = firing_strength
    end

    p1 = -log((1 / p_firing_strength) - 1) / membership_function.a + membership_function.c
    p2 = membership_function.limit
    (p1 + p2) / 2
end


"""
   function lagdiff(x::AbstractVector, d=1)

Returns `out[i] = x[i] - x[i-d]` with out padded with missings so that `length(out)==length(x)`.
"""
function lagdiff(x::AbstractVector, d=1)
  if d>0
    dx = vcat(fill(missing, d), x[(d+1):end] .- x[1:(end-d)])
  elseif d<0
    dx = vcat(-x[(d+1):end] .+ x[1:(end-d)], fill(missing,d))
  else
    dx = zero(x)
  end
  return(dx)
end

"""
    function smooth(x::AbstractVector; w=pdf(Normal(), range(-3, 3, length=7)))

Returns the running weighted mean of x*w. When `n=length(w)`, then
`out[i] = sum(x[i - n÷2 + j-1] * w[j] for j in 1:n)/sum(w)`.
"""
function smooth(x::AbstractVector;  w=pdf(Normal(), range(-3, 3, length=7)))
  sx = Vector{Union{Missing, Float64}}(undef, length(x))
  n = length(w)
  if n % 2 != 1
    error("only odd length windows allowed")
  end
  shift = -(n÷2):(n÷2)
  for i in 1:length(sx)
    s = findfirst((i .+ shift) .> 0)
    l = findlast((i .+ shift) .< length(x))
    sx[i] = sum(w[s:l].*x[i .+ shift[s:l]])./sum(w[s:l])
    if isnan(sx[i])
      sx[i] = 0
    end
  end
  return(sx)
end


## Julia program for Atkinson's Poisson Simulation
## author: weiya <szcfweiya@gmail.com>
## date: 2018-08-19

function AtkinsonPois(lambda)
    # parameters
    beta = pi/sqrt(3*lambda)
    alpha = lambda*beta
    c = 0.767 - 3.36/lambda
    k = log(c) - lambda - log(beta)
    # step 1: propose new x
    u1 = rand()
    while true
        global x
        x = (alpha - log((1-u1)/u1))/beta
        x > -0.5 && break
    end
    while true
        # step 2: transform to N
        N = floor(Int, x)
        u2 = rand()
        # step 3: accept or not
        lhs = alpha - beta*x + log(u2/(1+exp(alpha-beta*x))^2)
        rhs = k + N*log(lambda) - log(factorial(lambda))
        if lhs <= rhs
            return(N)
        end
    end
end

## example
N = 100;
res = ones(Int, N);
for i = 1:N
    res[i] = AtkinsonPois(10)
end
# ans: 8 9 13 10 12 .......

# alternative method
function SimplePois(lambda)
    s = 0
    k = 0
    while true
        u = rand()
        x = -log(u)/lambda
        s = s + x
        if s > 1
            return(k)
        end
        k = k + 1
    end
end

## example for simple poisson
res2 = ones(Int, N);
for i = 1:N
    res2[i] = SimplePois(10)
end
# ans: 5 7 16 7 10 .......


using Plots;gr()

f(x) = x*x - 2

function bisection(a, b, max_iter)
    c = 1e+9
    error = zeros(max_iter)
    for i in 1:max_iter
        c = (a + b) / 2
        error[i] = abs(c - sqrt(2))
        f_c = f(c)

        if (f_c == 0)
            break
        elseif (f_c < 0)
            a = c
        else
            b = c
        end
    end

    (c, error)
end

a = 1.0
b = 2.0
max_iter = 20
c, error = bisection(a, b, max_iter)
t = 1:max_iter

plot(t, error, color=:cornflowerblue, label="bisection",xlabel="t",ylabel="Error")
savefig("bisection.png")


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

          Solving the HJB with a state constraint using
                the explicit method

               Translated from matlab code on Ben Moll's website:
                              http://www.princeton.edu/~moll/HACTproject.htm

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

using LinearAlgebra, SparseArrays, Plots

γ = 2 # parameter from CRRA utility
r = 0.03 # the interest rate
ρ = 0.05 # the discount rate

z1=.1
z2 =.2
z= [z1 z2]

λ1 = 0.02
λ2 = 0.03

λ = [λ1 λ2]

H = 500 # number of points in the grid space
amin=-0.02
amax=2

a = LinRange(amin,amax,H)
da = (amax-amin)/(H-1)

aa = [a a]
zz = ones(H,1)*z

maxit = 100 # the maximum number of iterations we allow the finite differencing algorithm
crit = 10^(-6) # our critical value
Δ = 1000

dVf=zeros(H,2)
dVb = zeros(H,2)
c= zeros(H,2)

# Create the matrix that captures the systems dynamic
Aswitch = [-sparse(I,H,H)*λ[1] sparse(I,H,H)*λ[1]; sparse(I,H,H)*λ[2] -sparse(I,H,H)*λ[2]]

# Inital Guess
v0=zeros(H,2)
v0[:,1] = (z[1] .+ r.*a).^(1-γ)/(1-γ)/ρ
v0[:,2] = (z[2] .+ r.*a).^(1-γ)/(1-γ)/ρ
global v=v0

dist=[]
V_n=[]

# The finite differeing loop

for n in 1:maxit
    global V=v
    push!(V_n,V)
    # Forward differencing
    dVf[1:H-1,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVf[H,:] = (z .+ r.*amax).^(-γ) # impose state constraint at the max, just in case
    # Backward differencing
    dVb[2:H,:]=(V[2:H,:]-V[1:H-1,:])/da
    dVb[1,:]= (z .+r.*amin).^(-γ)

    # Find consumption and savings with the forward difference
    cf = dVf.^(-1/γ)
    sf = zz +r.*aa -cf

    # Find consumption and savings with the backward difference
    cb= dVb.^(-1/γ)
    sb = zz +r.*aa -cb

    # Find consumption and savings at steady state
    c0 = zz + r.*aa
    dV0 = c0.^(-γ)

    # Now implement the upwind scheme in order to select the best differencing method
    If = sf .> 0
    Ib = sb .< 0
    I0 = (1 .- If - Ib)

    # State constraint at amin is automatically implemented

    global dV_Upwind = dVf.*If + dVb.*Ib + dV0.*I0

    global c = dV_Upwind.^(-1/γ)
    global u = c.^(1-γ)/(1-γ)
    # Construct matrix for the evolution of the system
    X = -min.(sb,0)/da
    Y= -max.(sf,0)/da + min.(sb,0)/da
    Z = max.(sf,0)/da

    A1 = spdiagm(0=>Y[:,1], -1 => X[2:H,1], 1=> Z[1:H-1,1])
    A2 = spdiagm(0=>Y[:,2], -1 => X[2:H,2], 1=> Z[1:H-1,2])
    AA = [A1 spzeros(H,H); spzeros(H,H) A2]

    A = AA + Aswitch

    B = (ρ +1/Δ)*sparse(I,2*H,2*H) -A
    u_stacked = [u[:,1]; u[:,2]]
    V_stacked = [V[:,1]; V[:,2]]

    b = u_stacked + V_stacked/Δ
    V_stacked = B\b # Solves the system of equations

    V = [V_stacked[1:H] V_stacked[H+1:2*H]]

    V_change = V - v

    global v = V

    push!(dist, findmax(abs.(V_change))[1])

    if dist[n] < crit
        println("Value Function converged, Iteration=")
        println(n)
        break
    end
end

# Graphs
adot = zz + r.*aa -c

plot(dist[:], xlabel="Iteration", ylabel="\$\\lvert V^{n+1}-V^{n}\\rvert\$")

plot(a,V, xlabel="a", ylabel="\$V_{i}(a)\$", xlims=(amin,amax),legend=false)

plot(a,c,xlabel="a", ylabel="\$c_{i}(a)\$", xlims=(amin,amax), legend=false)

plot(a, adot, xlabel="a", ylabel="\$s_{i}(a)\$", xlims=(amin,amax))


# # [Sparse linear regression](@id sparse_linreg)
#
# Let's look at a least squares regression problem with L1 regularization:
# we will use the "diabetes dataset" (see [here](https://www4.stat.ncsu.edu/~boos/var.select/)), so let's start by loading the data.

using HTTP

splitlines(s) = split(s, "\n")
splitfields(s) = split(s, "\t")
parsefloat64(s) = parse(Float64, s)

function load_diabetes_dataset()
    res = HTTP.request("GET", "https://www4.stat.ncsu.edu/~boos/var.select/diabetes.tab.txt")
    lines = res.body |> String |> strip |> splitlines
    return hcat((line |> splitfields .|> parsefloat64 for line in lines[2:end])...)'
end

data = load_diabetes_dataset()

training_input = data[1:end-100, 1:end-1]
training_label = data[1:end-100, end]

test_input = data[end-99:end, 1:end-1]
test_label = data[end-99:end, end]

n_training, n_features = size(training_input)

# Now we can set up the optimization problem we want to solve: we will minimize the mean squared error
# for a linear model, appropriately scaled so that the features in the training data are normally distributed
# ("standardization", this is known to help the optimization process).
#
# After some simple manipulation, this standardized linear model can be implemented as follows:

using LinearAlgebra
using Statistics

input_loc = mean(training_input, dims=1) |> vec
input_scale = std(training_input, dims=1) |> vec

linear_model(wb, input) = input * wb[1:end-1] .+ wb[end]

function standardized_linear_model(wb, input)
    w_scaled = wb[1:end-1] ./ input_scale
    wb_scaled = vcat(w_scaled, wb[end] - dot(w_scaled, input_loc))
    return linear_model(wb_scaled, input)
end

# The loss term in the cost is then the following. Note that this is a regular Julia function:
# since the algorithm we will apply requires its gradient, automatic differentiation will
# do the work for us.

mean_squared_error(label, output) = mean((output .- label) .^ 2) / 2

training_loss(wb) = mean_squared_error(training_label, standardized_linear_model(wb, training_input))

# As regularization we will use the L1 norm, implemented in [ProximalOperators](https://github.com/JuliaFirstOrder/ProximalOperators.jl):

using ProximalOperators

reg = ProximalOperators.NormL1(1)

# We want to minimize the sum of `training_loss` and `reg`, and for this task we can use `FastForwardBackward`,
# which implements the fast proximal gradient method (also known as fast forward-backward splitting, or FISTA).
# Therefore we construct the algorithm, then apply it to our problem by providing a starting point,
# and the objective terms `f=training_loss` (smooth) and `g=reg` (non smooth).

using ProximalAlgorithms

ffb = ProximalAlgorithms.FastForwardBackward()
solution, iterations = ffb(x0=zeros(n_features + 1), f=training_loss, g=reg)

# We can now check how well the trained model performs on the test portion of our data.

test_output = standardized_linear_model(solution, test_input)
mean_squared_error(test_label, test_output)


#=

Plot the price-dividend ratio in the risk neutral case, for the Markov asset pricing
lecture.

=#
using QuantEcon
using Plots
using LaTeXStrings

pyplot()

n = 25  # size of state space
beta = 0.9
mc = tauchen(n, 0.96, 0.02)

K = mc.p .* reshape(exp(mc.state_values), 1, n)

I = eye(n)
v = (I - beta * K) \  (beta * K * ones(n, 1))

plot(mc.state_values,
    v,
    lw=2,
    ylabel="price-dividend ratio",
    xlabel="state",
    alpha=0.7,
    label=L"$v$")


function agm(x, y, e::Real = 5)
    (x ≤ 0 || y ≤ 0 || e ≤ 0) && throw(DomainError("x, y must be strictly positive"))
    g, a = minmax(x, y)
    while e * eps(x) < a - g
        a, g = (a + g) / 2, sqrt(a * g)
    end
    a
end

x, y = 1.0, 1 / √2
println("# Using literal-precision float numbers:")
@show agm(x, y)

println("# Using half-precision float numbers:")
x, y = Float32(x), Float32(y)
@show agm(x, y)

println("# Using ", precision(BigFloat), "-bit float numbers:")
x, y = big(1.0), 1 / √big(2.0)
@show agm(x, y)


"""
$(SIGNATURES)
"""
average(m::AbstractMatrix) = (sum(m) / length(m))

"""
$(SIGNATURES)
"""
average_absolute(m::AbstractMatrix) = (sum(abs.(m)) / length(m))

"""
$(SIGNATURES)
"""
test_average_absolute(m::AbstractMatrix) = (average_absolute(m)) < 1e6

"""
$(SIGNATURES)
"""
function smallest_positive(x1::Float64, x2::Float64)
    (x1 > 0.0 && x2 < 0.0) && return x1
    (x1 < 0.0 && x2 > 0.0) && return x2
    (x1 > 0.0 && x2 > 0.0) && return min(x1, x2)
    return nothing
end

"""
$(SIGNATURES)

Approximates the gradient of a scalar field.
"""
function ∇(ϕ::AbstractMatrix{Float64})
    height, width = size(ϕ)

    # x, y only represents the first and the second variables. No reference to graphical representations.
    ∇ϕx = zeros(Float64, height, width)   # divergence on the right edge will be filled with zeros
    ∇ϕy = zeros(Float64, height, width)   # divergence on bottom edge will be filled with zeros

    ∇ϕx[begin:end-1, :] = ϕ[begin+1:end, :] - ϕ[begin:end-1, :]
    ∇ϕy[:, begin:end-1] = ϕ[:, begin+1:end] - ϕ[:, begin:end-1]

    fill_borders!(∇ϕx, 0.0)
    fill_borders!(∇ϕy, 0.0)

    return ∇ϕx, ∇ϕy
end

"""
$(SIGNATURES)

Calculate the second-order Laplace operator by convolution of a kernel.
"""
function laplacian(ϕ::AbstractMatrix{Float64})
    height, width = size(ϕ)

    # Embed matrix within a larger matrix for better vectorization and avoid duplicated code
    # Padded matrix adds 1 row/col after the size of ϕ.
    # ϕ is inserted in padded matrix within  1:1+height+1 x 1:1+width+1.
    # The rest of the padded matrix (its borders) are set at 0.0.
    padded_ϕ = zeros(Float64, 1 + height + 1, 1 + width + 1)
    padded_ϕ[begin+1:end-1, begin+1:end-1] .= ϕ[1:end, 1:end]

    # Convolution = up + down + left + right
    ∇²ϕ =
        padded_ϕ[begin+1-1:end-1-1, begin+1:end-1] +
        padded_ϕ[begin+1+1:end-1+1, begin+1:end-1] +
        padded_ϕ[begin+1:end-1, begin+1-1:end-1-1] +
        padded_ϕ[begin+1:end-1, begin+1+1:end-1+1] -
        4.0 * padded_ϕ[begin+1:end-1, begin+1:end-1]

    return ∇²ϕ[1:end-1, 1:end-1]
end


using Downhill

"""
    rosenbrock(x; b=2)

Return the value of Rosenbrock function
    ```math
    \\sum_{i=1}^n \\right[(1 - x_{i-1})^2 + b (x_i - x_{i-1})^2 \\left]
    ```
    in `n = length(x)` dimensions.
"""
function rosenbrock(x::AbstractVector{T}; b=2) where {T}
    f = zero(T) * zero(b) # type-generic zero
    for i in firstindex(x)+1:lastindex(x)
        f += (one(T) - x[i-1])^2 + b * (x[i] - x[i-1]^2)^2
    end
    return f
end

"""
    drosenbrock(x; b=2)

Return the gradient of Rosenbrock function
    ```math
    \\sum_{i=1}^n \\right[(1 - x_{i-1})^2 + b (x_i - x_{i-1})^2 \\left]
    ```
    in `n = length(x)` dimensions.
"""
function drosenbrock(x::AbstractVector{T}; b=2) where {T}
    g = fill!(similar(x, promote_type(T, typeof(b))), 0)
    for i in firstindex(x)+1:lastindex(x)
        g[i-1] += 2 * (x[i-1] - 1) + 4 * b * x[i-1] * (x[i-1]^2 - x[i])
        g[i] += 2 * b * (x[i] - x[i-1]^2)
    end
    return g
end

"""
    rosenbrock!(x, g; b=2)

Return the value of Rosenbrock function in `length(x)` dimensions.
"""
function rosenbrock!(x::AbstractVector, g::AbstractVector; b=2)
    f = zero(eltype(g))
    fill!(g, 0)
    inds = eachindex(x, g)
    for i in 2:last(inds)
        f += (1 - x[i-1])^2 + b * (x[i] - x[i-1]^2)^2
        g[i-1] += 2 * (x[i-1] - 1) + 4 * b * x[i-1] * (x[i-1]^2 - x[i])
        g[i] += 2 * b * (x[i] - x[i-1]^2)
    end
    return f, g
end

ros(x, g) = (rosenbrock(x), drosenbrock(x))

ros!(x, g) = rosenbrock!(x, g)

ans_nonmutating = let x0 = zeros(2)
    opt = BFGS(x0)
    optresult = optimize!(ros, opt, x0; maxiter=1000)
    println("""
        ==Nonmutating gradient evaluation==

        Optimization converged: $(optresult.converged)
        Final argument: $(optresult.argument)
        Final gradient: $(optresult.gradient)
        Iteration count: $(optresult.iterations)
        Call count: $(optresult.calls)
    """)
    optresult.argument
end

ans_mutating = let x0 = zeros(2)
    opt = BFGS(x0)
    optresult = optimize!(ros!, opt, x0; maxiter=1000)
    println("""
        ==Mutating gradient evaluation==

        Optimization converged: $(optresult.converged)
        Final argument: $(optresult.argument)
        Final gradient: $(optresult.gradient)
        Iteration count: $(optresult.iterations)
        Call count: $(optresult.calls)
    """)
    optresult.argument
end

@assert ans_nonmutating ≈ ans_mutating


"""
   HermiteWeight()

is a quasi-vector representing `exp(-x^2)` on ℝ.
"""
struct HermiteWeight{T} <: Weight{T} end

HermiteWeight() = HermiteWeight{Float64}()
axes(::HermiteWeight{T}) where T = (Inclusion(ℝ),)
function getindex(w::HermiteWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    exp(-x^2)
end

sum(::HermiteWeight{T}) where T = sqrt(convert(T, π))

struct Hermite{T} <: OrthogonalPolynomial{T} end
Hermite() = Hermite{Float64}()
orthogonalityweight(::Hermite{T}) where T = HermiteWeight{T}()

==(::Hermite, ::Hermite) = true
axes(::Hermite{T}) where T = (Inclusion(ℝ), oneto(∞))

"""
     hermiteh(n, z)

computes the `n`-th Hermite polynomial, orthogonal with
respec to `exp(-x^2)`, at `z`.
"""
hermiteh(n::Integer, z::Number) = Base.unsafe_getindex(Hermite{typeof(z)}(), z, n+1)

broadcasted(::LazyQuasiArrayStyle{2}, ::typeof(*), ::HermiteWeight{T}, ::Hermite{V}) where {T,V} = Weighted(Hermite{promote_type(T,V)}())

# H_{n+1} = 2x H_n - 2n H_{n-1}
# 1/2 * H_{n+1} + n H_{n-1} = x H_n
# x*[H_0 H_1 H_2 …] = [H_0 H_1 H_2 …] * [0    1; 1/2  0     2; 1/2   0  3; …]
jacobimatrix(H::Hermite{T}) where T = Tridiagonal(Fill(one(T)/2,∞), Zeros{T}(∞), one(T):∞)
recurrencecoefficients(H::Hermite{T}) where T = Fill{T}(2,∞), Zeros{T}(∞), zero(T):2:∞

massmatrix(::Hermite{T}) where T = Diagonal(sqrt(convert(T,π)) .* convert(T,2) .^ (0:∞) .* factorial.(convert(T,0):∞))

@simplify *(A::QuasiAdjoint{<:Any,<:Weighted{<:Any,<:Hermite}}, B::Hermite) = massmatrix(Hermite{promote_type(eltype(A),eltype(B))}())

##########
# Derivatives
##########

@simplify function *(D::Derivative, H::Hermite)
    T = promote_type(eltype(D),eltype(H))
    D = _BandedMatrix((zero(T):2:∞)', ℵ₀, -1,1)
    H*D
end

@simplify function *(D::Derivative, Q::OrthonormalWeighted{<:Any,<:Hermite})
    X = jacobimatrix(Q.P)
    Q * Tridiagonal(-X.ev, X.dv, X.ev)
end


#nb # %% A slide [markdown] {"slideshow": {"slide_type": "subslide"}}
# ## Symmetric Strip Grating
# This example consists of a symmetric strip grating, i.e. a grating where the strip width
# is half the unit cell period ``P``:
#
# ![diagram](./assets/symmetric_strip_diagram.png)
#
# Only three of the infinite number of strips in the grating are shown, and they extend infinitely to the left and right.
# The grating lies in the ``z=0`` plane with free space on both sides. The shaded areas represent metalization.
# The dashed lines show two possible choices for the unit cell location: "J" for a formulation in terms of electric
# surface currents, and "M" for magnetic surface currents.
#
# For normal incidence there is a closed-form solution due to Weinstein,
# but for a more recent reference one can find the solution in Problem 10.6 of R. E. Collin,
# *Field Theory of Guided Waves, Second Ed.*,
# IEEE Press, 1991.  Here is the code for computing the exact solution:
"""
    grating(kP, nterms=30) -> (Γ, T)

Compute the normal incidence reflecton and transmission coefficients of a symmetric grid of
zero-thickness conducting strips.  The product of the period of the strips and the incident
electric field wavenumber is `kP` (dimensionless).  The incident electric field is
perpendicular to the direction along the axis of the strips.  The series have been
accelerated by applying a Kummer's transformation, using the first two terms in the Maclaurin
series for the inverse sin function.  `kP` must be in the half-open interval [0,1). The
default number of summed terms `nterms` yields better than 10 digits of accuracy over the
interval [0.01,0.99].
"""
function grating(kP; nterms=30)
    sum1 = 1.3862943611198906 # \sum_{n=1}^{\infty} 1/(n-1/2) - 1/n = log(4)
    sum3 = 7.2123414189575710 # \sum_{n=1}^{\infty} (n-1/2)^{-3} - n^{-3} = 6 * \zeta(3)
    x = kP/(4π)
    θ = x*sum1 + x^3/6 * sum3
    for n = 1:nterms
        xonmhalf = x/(n - 0.5)
        xon = x/n
        term = asin(xonmhalf) - (xonmhalf + (xonmhalf)^3/6) -
              (asin(xon) - (xon + xon^3/6))
        θ += term
    end
    Γ = sin(θ) * cis(-π/2 - θ)
    T = 1 + Γ
    return (Γ, T)
end

#  Note that using the extension of
# [Babinet's Principle for electromagnetic fields](http://kirkmcd.princeton.edu/examples/babinet.pdf)
# this also provides the solution (upon appropriate interchange and sign change of the coefficients) for
# the case where the incident wave polarization is parallel to the direction of the strips.

# Here is the PSSFSS code to analyze this structure using electric currents as the unknowns.  We will
# scale the geometry so that the frequency in GHz is numerically equal to the period of the strips
# measured in wavelengths.

using Plots, PSSFSS
c = 11.802852677165355 # light speed [inch*GHz]
period = c  # so the period/wavelength = freq in GHz
Py = period
Ly = period/2
Px = Lx = Ly/10 # Infinite in x direction so this can be anything
Ny = 60
Nx = round(Int, Ny*Lx/Ly)
sheet = rectstrip(;Px, Py, Lx, Ly, Nx, Ny, units=inch)
flist = 0.02:0.02:0.98
steering = (θ=0, ϕ=0)
strata = [Layer()
          sheet
          Layer()]
results_j = analyze(strata, flist, steering, showprogress=false,
                    resultfile=devnull, logfile=devnull);
p1 = plot(sheet)
p2 = plot(sheet, unitcell=true)
ptitle = plot(title = "Symmetric Strip Triangulation",
             grid = false, showaxis = false, xtick=[], ytick=[],
             bottom_margin = -50Plots.px)
plot(ptitle, p1, p2, layout = @layout([A{0.09h}; [B C]]))
#md savefig("symstrip1.png"); nothing  # hide
#-
#md # ![](symstrip1.png)

# Note that setting `Lx = Px` causes the strip to fully occupy the x-extent
# of the unit cell.  PSSFSS automatically ensures that the triangle edges at these unit
# cell boundaries define basis functions that satisfy the Floquet (phase shift) boundary
# conditions, so that currents are free to flow across these unit cell boundaries.

# We can also analyze the same structure using magnetic currents in the areas free of
# metalization as the unknowns:
sheet = rectstrip(;class='M', Px, Py, Lx, Ly, Nx, Ny, units=inch)
strata = [Layer()
          sheet
          Layer()]
results_m = analyze(strata, flist, steering, showprogress=false,
                    resultfile=devnull, logfile=devnull);

#
# Each 50-frequency run of `analyze` takes about 14 seconds
# for this geometry of 720 triangles on my machine.
# More detailed timing information is available in the log file
# (which is omitted for generating this documentation).

# We will compare the PSSFSS results to the analytic solution:
## Generate exact results:
rt = grating.(2π*flist)
rperp_exact = first.(rt)
tperp_exact = last.(rt)
rpar_exact = -tperp_exact
tpar_exact = -rperp_exact;

# Obtain PSSFSS results for electric and magnetic currents:
outrequest = @outputs s11(v,v) s21(v,v) s11(h,h) s21(h,h)
rperp_j, tperp_j, rpar_j, tpar_j =
      collect.(eachcol(extract_result(results_j, outrequest)))
rperp_m, tperp_m, rpar_m, tpar_m =
      collect.(eachcol(extract_result(results_m, outrequest)));

# Generate the comparison plots:
angdeg(z) = rad2deg(angle(z)) # Convenience function

p1 = plot(title = "Perpendicular Reflection Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude",
          legend=:topleft)
plot!(p1, flist, abs.(rperp_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(rperp_j), label="PSSFSS J")
plot!(p1, flist, abs.(rperp_m), label="PSSFSS M")
#md savefig("symstrip2.png"); nothing  # hide
#-
#md # ![](symstrip2.png)
#-
p2 = plot(title = "Perpendicular Reflection Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(rperp_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(rperp_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(rperp_m), label="PSSFSS M")
#md savefig("symstrip3.png"); nothing  # hide
#-
#md # ![](symstrip3.png)
#-
p1 = plot(title = "Parallel Reflection Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude")
plot!(p1, flist, abs.(rpar_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(rpar_j), label="PSSFSS J")
plot!(p1, flist, abs.(rpar_m), label="PSSFSS M")
#md savefig("symstrip4.png"); nothing  # hide
#-
#md # ![](symstrip4.png)
#-
p2 = plot(title = "Parallel Reflection Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(rpar_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(rpar_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(rpar_m), label="PSSFSS M")
#md savefig("symstrip5.png"); nothing  # hide
#-
#md # ![](symstrip5.png)

# Now look at the transmission coefficients:
p1 = plot(title = "Perpendicular Transmission Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude")
plot!(p1, flist, abs.(tperp_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(tperp_j), label="PSSFSS J")
plot!(p1, flist, abs.(tperp_m), label="PSSFSS M")
#md savefig("symstrip6.png"); nothing  # hide
#-
#md # ![](symstrip6.png)
#-
p2 = plot(title = "Perpendicular Transmission Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(tperp_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(tperp_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(tperp_m), label="PSSFSS M")
#md savefig("symstrip7.png"); nothing  # hide
#-
#md # ![](symstrip7.png)
#-
p1 = plot(title = "Parallel Transmission Magnitude",
          xlabel = "Period (wavelengths)",
          ylabel = "Coefficient Magnitude", legend=:topleft)
plot!(p1, flist, abs.(tpar_exact), ls=:dash, label="Exact")
plot!(p1, flist, abs.(tpar_j), label="PSSFSS J")
plot!(p1, flist, abs.(tpar_m), label="PSSFSS M")
#md savefig("symstrip8.png"); nothing  # hide
#-
#md # ![](symstrip8.png)
#-
p2 = plot(title = "Parallel Transmission Phase",
          xlabel = "Period (wavelengths)",
          ylabel = "Phase (deg)")
plot!(p2, flist, angdeg.(tpar_exact), ls=:dash, label="Exact")
plot!(p2, flist, angdeg.(tpar_j), label="PSSFSS J")
plot!(p2, flist, angdeg.(tpar_m), label="PSSFSS M")
#md savefig("symstrip9.png"); nothing  # hide
#-
#md # ![](symstrip9.png)

# ### Conclusion
# Although good agreement is obtained, as expected the best agreement between
# all three results occurs for the lowest frequencies, where the triangles are
# smallest in terms of wavelength.  This emphasizes the fact that it is necessary for the
# user to check that enough triangles have been requested for good convergence
# over the frequency band of interest.  This example is an extremely demanding case
# in terms of bandwidth, as the ratio of maximum to minimum frequency here
# is ``0.98/0.02 = 49:1``


using LabelledArrays

export verdes

"""
    eom_verdes(u, p, t) → SVector{3}

Equations of motion for a 3D system where the response X is a highly
nonlinear combination of Y and Z [1]. The forcings Y and Z involve
sines and cosines, respectively, and have different periods.

The implementation here allows for tuning the period of the
forcing signals.

# References
Verdes, P. F. "Assessing causality from multivariate time series." Physical Review E 72.2 (2005): 026222.
"""
function eom_verdes(u, p, t)
    x, y, z = (u...,)
    ωy, ωz, σx, σy, σz = (p...,)

    ηx = σx == 0 ? 0 : rand(Normal(0, σx))
    ηy = σy == 0 ? 0 : rand(Normal(0, σy))
    ηz = σz == 0 ? 0 : rand(Normal(0, σz))

    dx = y*(18y - 27y^2 + 10)/2 + z*(1-z) + ηx
    dy = (1 - cos((2*pi/ωy) * t))/2 + ηy
    dz = (1 - sin((2*pi/ωz) * t))/2 + ηz
    return SVector{3}(dx, dy, dz)
end

function verdes(u₀, ωy, ωz, σx, σy, σz)
    p = LVector(ωy = ωy, ωz = ωz, σx = σx, σy = σy, σz = σz)

    DiscreteDynamicalSystem(eom_verdes, u₀, p)
end

"""
    verdes(;u₀ = rand(3), ωy = 315, ωz = 80,
        σx = 0.0, σy = 0.0, σz = 0.0) → DiscreteDynamicalSystem

Intitialise a 3D system where the response X is a highly nonlinear combination
of Y and Z. The forcings Y and Z involve sines and cosines, and
have different periods, which controlled by `ωy` and `ωz`.

The equations of motion are

```math
\\begin{aligned}
x(t+1) &= \\dfrac{y(t)(18y(t) - 27y(t)^2 + 10)}{2} + z(t)(1-z(t)) + ηx \\
y(t+1) &= \\dfrac{(1 - \\dfrac{\\cos(2\\pi)}{\\omega y}t)}{2} + ηy \\
z(t+1) &= \\dfrac{(1 - \\dfrac{\\sin(2\\pi)}{\\omega z}t)}{2} + ηz
\\end{aligned}
```
where ηx, ηy, ηz is gaussian noise with mean 0 and standard deviation `σx`, `σy`
and `σz`.

# References
Verdes, P. F. "Assessing causality from multivariate time series." Physical
Review E 72.2 (2005): 026222.
"""
verdes(;u₀ = rand(3),
    ωy = 315, ωz = 80,
    σx = 0.01, σy = 0.01, σz = 0.01) =
    verdes(u₀, ωy, ωz, σx, σy, σz)


using Plots
using DifferentialEquations

@userplot LorenzPlot
@recipe function f(lp::LorenzPlot)
    x, y, z, i = lp.args
    n = length(x)
    inds = circshift(1:n, 1-i)
    seriesalpha --> range(0, 1, length=n)
    aspect_ratio --> 1
    label --> false
    x[inds], y[inds], z[inds]
end

function anim()

    function lorenz(du, u, p, t)

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

    end

    u0 = [1, 5, 10]
    tspan = (0, 100)
    p = (10, 28, 8/3)

    prob = ODEProblem(lorenz, u0, tspan, p, reltol=1e-10, abstol=1e-10)
    sol = solve(prob)
    n = length(sol.t)

    plt = plot3d(1, xaxis=("x", (-30, 30)), yaxis=("y", (-30, 30)), zaxis=("z", (0, 60)),
                 title="Lorenz Attractor", marker=2, markercolor="black", dpi=300)

    # anim = @animate for i in 1:n
    #     push!(plt, sol[1,i], sol[2,i], sol[3,i])
    # end
    # gif(anim, "de_anim.gif", fps=30)

    anim = @animate for i in 1:10:n
        lorenzplot(sol[1,:], sol[2,:], sol[3,:], i, c = :thermal)
    end
    gif(anim, "de_anim_60.gif", fps=60)

end

anim()

    # anim = @animate for i in 1:10:n
    #     lorenzplot(sol[1,:], sol[2,:], sol[3,:], i)
    # end
    # gif(anim, "de_anim60.gif", fps=60)


module test

using Roots,Interpolations,PyPlot,Optim
using CompEcon, NLsolve

alpha     = 0.65
beta      = 0.95
grid_max  = 2  # upper bound of capital grid
n         = 150  # number of grid points
N_iter    = 3000  # number of iterations
kgrid     = 1e-6:(grid_max-1e-6)/(n-1):grid_max  # equispaced grid
f(x) = x.^alpha  # defines the production function f(k)
tol = 1e-9

ab        = alpha * beta
c1        = (log(1 - ab) + log(ab) * ab / (1 - ab)) / (1 - beta)
c2        = alpha / (1 - ab)
v_star(k) = c1 .+ c2 .* log(k)
k_star(k) = ab * k.^alpha
c_star(k) = (1-ab) * k.^alpha
ufun(x) = log(x)

# Bellman Operator
# inputs
# `grid`: grid of values of state variable
# `v0`: current guess of value function

# output
# `v1`: next guess of value function
# `pol`: corresponding policy function

#takes a grid of state variables and computes the next iterate of the value function.
function bellman_operator(grid,v0)

    v1  = zeros(n)     # next guess
    pol = zeros(Int,n)     # policy function
    w   = zeros(n)   # temporary vector

    # loop over current states
    # current capital
    for (i,k) in enumerate(grid)

        # loop over all possible kprime choices
        for (iprime,kprime) in enumerate(grid)
            if f(k) - kprime < 0   #check for negative consumption
                w[iprime] = -Inf
            else
                w[iprime] = ufun(f(k) - kprime) + beta * v0[iprime]
            end
        end
        # find maximal choice
        v1[i], pol[i] = findmax(w)     # stores Value und policy (index of optimal choice)
    end
    return (v1,pol)   # return both value and policy function
end

# VFI iterator
#
## input
# `n`: number of grid points
# output
# `v_next`: tuple with value and policy functions after `n` iterations.
function VFI()
    v_init = zeros(n)     # initial guess
    for iter in 1:N_iter
        v_next = bellman_operator(kgrid,v_init)  # returns a tuple: (v1,pol)
        # check convergence
        if maxabs(v_init.-v_next[1]) < tol
            verrors = maxabs(v_next[1].-v_star(kgrid))
            perrors = maxabs(kgrid[v_next[2]].-k_star(kgrid))
            println("discrete VFI:")
            println("Found solution after $iter iterations")
            println("maximal value function error = $verrors")
            println("maximal policy function error = $perrors")
            return v_next
        elseif iter==N_iter
            warn("No solution found after $iter iterations")
            return v_next
        end
        v_init = v_next[1]  # update guess
    end
end

# plot
function plotVFI()
    v = VFI()
    figure("discrete VFI",figsize=(10,5))
    subplot(131)
    plot(kgrid,v[1],color="blue")
    plot(kgrid,v_star(kgrid),color="black")
    xlim(-0.1,grid_max)
    ylim(-50,-30)
    xlabel("k")
    ylabel("value")
    title("value function")

    subplot(132)
    plot(kgrid,kgrid[v[2]])
    plot(kgrid,k_star(kgrid),color="black")
    xlabel("k")
    title("policy function")

    subplot(133)
    plot(kgrid,kgrid[v[2]].-k_star(kgrid))
    title("policy function error")
end

function bellman_operator2(grid,v0)

    v1  = zeros(n)     # next guess
    pol = zeros(n)     # consumption policy function

    Interp = interpolate((collect(grid),), v0, Gridded(Linear()) )

    # loop over current states
    # of current capital
    for (i,k) in enumerate(grid)

        objective(c) = - (log(c) + beta * Interp[f(k) - c])
        # find max of ojbective between [0,k^alpha]
        res = optimize(objective, 1e-6, f(k))
        pol[i] = f(k) - res.minimizer
        v1[i] = -res.minimum
    end
    return (v1,pol)   # return both value and policy function
end

function VFI2()
    v_init = zeros(n)     # initial guess
    for iter in 1:N_iter
        v_next = bellman_operator2(kgrid,v_init)  # returns a tuple: (v1,pol)
        # check convergence
        if maxabs(v_init.-v_next[1]) < tol
            verrors = maxabs(v_next[1].-v_star(kgrid))
            perrors = maxabs(v_next[2].-k_star(kgrid))
            println("continuous VFI:")
            println("Found solution after $iter iterations")
            println("maximal value function error = $verrors")
            println("maximal policy function error = $perrors")
            return v_next
        elseif iter==N_iter
            warn("No solution found after $iter iterations")
            return v_next
        end
        v_init = v_next[1]  # update guess
    end
    return nothing
end

function plotVFI2()
    v = VFI2()
    figure("discrete VFI - continuous control",figsize=(10,5))
    subplot(131)
    plot(kgrid,v[1],color="blue")
    plot(kgrid,v_star(kgrid),color="black")
    xlim(-0.1,grid_max)
    ylim(-50,-30)
    xlabel("k")
    ylabel("value")
    title("value function")

    subplot(132)
    plot(kgrid,v[2])
    plot(kgrid,k_star(kgrid),color="black")
    xlabel("k")
    title("policy function")

    subplot(133)
    plot(kgrid,v[2].-k_star(kgrid))
    xlabel("k")
    title("policy function error")
end

function policy_iter(grid,c0,u_prime,f_prime)

    c1  = zeros(length(grid))     # next guess
    pol_fun = interpolate((collect(grid),), c0, Gridded(Linear()) )

    # loop over current states
    # of current capital
    for (i,k) in enumerate(grid)
        objective(c) = u_prime(c) - beta * u_prime(pol_fun[f(k)-c]) * f_prime(f(k)-c)
        c1[i] = fzero(objective, 1e-10, f(k)-1e-10)
    end
    return c1
end

uprime(x::Float64) = 1.0 / x
fprime(x::Float64) = alpha* x ^ (alpha-1)
function fprime(x::Vector)
    r = similar(x)
    for ix in eachindex(x)
        if x[ix] > 0
            r[ix] = alpha* x[ix]^(alpha-1)
        else
            r[ix] = 10000.0
        end
    end
    return r
end

function PFI()
    c_init = kgrid
    for iter in 1:N_iter
        c_next = policy_iter(kgrid,c_init,uprime,fprime)
        # check convergence
        if maxabs(c_init.-c_next) < tol
            perrors = maxabs(c_next.-c_star(kgrid))
            println("PFI:")
            println("Found solution after $iter iterations")
            println("max policy function error = $perrors")
            return c_next
        elseif iter==N_iter
            warn("No solution found after $iter iterations")
            return c_next
        end
        c_init = c_next  # update guess
    end
end
function plotPFI()
    v = PFI()
    figure("PFI")
    subplot(121)
    plot(kgrid,v)
    plot(kgrid,c_star(kgrid),color="black")
    xlabel("k")
    title("policy function")

    subplot(122)
    plot(kgrid,v.-c_star(kgrid))
    xlabel("k")
    title("policy function error")
end
#plotVFI()
#plotVFI2()
#plotPFI()

function proj2(n=25)

    a = 1e-6
    aa = 1e-2
    b = 1.0
    bb = 0.95
    basis = fundefn(:cheb,n,a,b)
    k = funnode(basis)[1]   # collocation points

    fs = f.(k)

    # system of equations: for each k, one line
    function ff(coef::Vector,result::Vector,b,coll_points,fs)
        # get cons function
        cons = funeval(coef,b,coll_points)[1]
        # println("coef= $coef")
        # println("cons = $cons")

        # put resulting residuals into result[:]
        # put into euler equation
        nextk = fs.-cons
        for i in eachindex(nextk)
            if nextk[i] < 0
                result[i] = -8000.0
            else
                result[i] = uprime(cons[i]) .- beta * uprime(nextk[i]) .* fprime(nextk[i])
            end
        end
        # println(result)
    end

    f_closure(x::Vector,r::Vector) = ff(x,r,basis,k,fs)

    res = nlsolve(f_closure,ones(n)*0.4)
    x = linspace(aa,bb,501)
    y = funeval(res.zero,basis,x)[1]
    plot(x,y)
    return res
end

function proj(n=25)

    alpha = 1.0
    eta   = 1.5
    a     = 0.1
    b     = 3.0
    basis = fundefn(:cheb,n,a,b)
    p     = funnode(basis)[1]   # collocation points

    c0 = ones(n)*0.3
    function resid!(c::Vector,result::Vector,p,basis,alpha,eta)
        # your turn!
        q = funeval(c,basis,p)[1]
        q2 = zeros(q)
        for i in eachindex(q2)
            if q[i] < 0
                q2[i] = -20.0
            else
                q2[i] = sqrt(q[i])
            end
        end
        result[:] = p.+ q .*((-1/eta)*p.^(eta+1)) .- alpha*q2 .- q.^2
    end
    f_closure(r::Vector,x::Vector) = resid!(x,r,p,basis,alpha,eta)
    res = nlsolve(f_closure,c0)
    println(res)

    # plot residual function
    x = collect(linspace(a,b,501))
    y = similar(x)
    resid!(res.zero,y,x,basis,alpha,eta);
    y = funeval(res.zero,basis,x)[1]
    pl = Any[]
    push!(pl,plot(x,y,title="residual function"))

    # plot supply functions at levels 1,10,20

    # plot demand function

    y = funeval(res.zero,basis,x)[1]
    p2 = plot(y,x,label="supply 1")
    plot!(10*y,x,label="supply 10")
    plot!(20*y,x,label="supply 20")
    d = x.^(-eta)
    plot!(d,x,label="Demand")

    push!(pl,p2)

    plot(pl...,layout=2)
end

end # module


# Refined Löwner-John ellipsoid example with Mosek, MosekTools, and JuMP

# Computes the inner and outer Löwner-John ellipsoids around the feasible
# space of the JuMP model.

# Uses Polyhedra and CDDLib packages to express the polyhedral set of the JuMP
# model by inequalities and a convex hull of a finite set of points

using JuMP, Mosek, MosekTools
using LinearAlgebra, Polyhedra, CDDLib
using Plots

function lowner_john_inner_ellipsoid(A,b)
    """
    Löwner-John inner ellipsoid:
        the function takes H-representation (A,b) of a polytope and
        outputs the parameters of the maximum volume inscribed ellipsoid.

        polyhedral set S:
        S = { x ∈ Rn | ai' x ⩽ bi, i= 1,...,m }

        Proposition 3.7.1 [BenTalN01]
            ... Then the largest volume ellipsoid contained in S is
            E = { x = Z∗u + z∗ | uTu ≤1 }, where (Z∗,z∗) are from

            maximize    t                                           (1.0)
                s.t.    t ⩽ (DetZ)^1/n                              (1.1)
                        Z - PSD                                     (1.2)
                        || Z * ai ||_2 ≤ bi −ai' * z, i = 1,...,m   (1.3)
                        Z ∈ S^n, z ∈ Rn, t ∈ R                      (1.4)

        To model (1.1), refer to [Mosek, 6.2.3 Log-determinant]

            t ⩽ (DetZ)^1/n <==> [Z B; B' diag(B)] - PSD             (2.1)
                                B is lower triangular               (2.2)
                                t ⩽ (Π_i={1,...,n} Bii)^1/n         (2.3)

            where (2.3) maximizes the geometric mean of the Bii variables. It
            can be modelled using power cone [Mosek, 4.2.4 Geometric mean] and
            MOI.GeometricMeanCone available with the MathOptInterface.jl

    [BenTalN01]	    A. Ben-Tal and A. Nemirovski. Lectures on Modern Convex
                    Optimization: Analysis, Algorithms, and Engineering Applications.
                    MPS/SIAM Series on Optimization. SIAM, 2001.
    [Mosek]         https://docs.mosek.com/modeling-cookbook/index.html
    """
    # model size
    m,n = size(A)
    # model declaraton
    model = Model(optimizer_with_attributes(Mosek.Optimizer, "LOG" => 0))
    # variables
    @variable(model, t)
    @variable(model, Z[1:n,1:n], PSD)
    @variable(model, z[1:n])
    @variable(model, B[1:n,1:n])
    # maximize the volume
    @objective(model, Max, t)
    # SOC constraint (1.3)
    @constraint(model, SOC[i=1:m], [ b[i] - (A[i,:]' * z) ; (A[i,:]'*Z)'] in SecondOrderCone())
    # PSD constraint (2.1)
    @constraint(model, [Z LowerTriangular(B); LowerTriangular(B)' diagm(diag(B))] in PSDCone())
    # Power cone costraint (2.3)
    @constraint(model, [t; vec([B[i,i] for i in 1:n])] in MOI.GeometricMeanCone(n+1))
    # optimize
    optimize!(model)
    status=termination_status(model)
    CPU_time = MOI.get(model, MOI.SolveTime())
    # return solution
    return JuMP.value.(Z), JuMP.value.(z), CPU_time
end

function lownerjohn_outer(x)
    """
    Löwner-John outer ellipsoid:
        the function takes V-representation (x) of a polytope and
        outputs the parameters of the maximum volume inscribed ellipsoid.

        Polyhedral set S = Conv{x_1,...,x_n}

        Uses the same Mosek features as the inner ellipsoid function,
        refer to Proposition 3.7.2 [BenTalN01] for further details.
    """

    m,n = size(x)
    model = Model(optimizer_with_attributes(Mosek.Optimizer, "LOG" => 0))

    @variable(model,t >= 0)
    @variable(model,P[1:n,1:n],PSD)
    @variable(model,B[1:n,1:n])
    @variable(model,c[1:n])

    # (1, Px-c) in cone
    for i in 1:m
        @constraint(model,[1 ; (x[i,:]' * P)' - c] in MOI.SecondOrderCone(n+1))
    end

    @constraint(model, [P LowerTriangular(B); LowerTriangular(B)' diagm(diag(B))] in PSDCone())
    @constraint(model, [t; vec([B[i,i] for i in 1:n])] in MOI.GeometricMeanCone(n+1))

    @objective(model,Max,t)

    optimize!(model)
    status=termination_status(model)
    CPU_time = MOI.get(model, MOI.SolveTime())

    JuMP.value.(P), JuMP.value.(c), CPU_time
end

# create a JuMP model
function JuMP_model(n)
    model = Model(optimizer_with_attributes(Mosek.Optimizer, "LOG" => 0))
    @variable(model, x[1:n])
    @objective(model, Min, x'diagm(ones(n))'x)
    @constraint(model, con1[i=1:n], 0 <= x[i] <= 1)
    @constraint(model, ones(n)'x >= n/3)
    optimize!(model)
    return model
end

# model dimension
n = 2
# get JuMP model
model = JuMP_model(n)
@info("done exporting JuMP model")
# create a polyhedron from the feasible set of a JuMP model
poly = polyhedron(model, CDDLib.Library(:exact))
@info("done exporting the polyhedron")

# obtain H-representation of the polyhedron
hr = MixedMatHRep(hrep(poly))
A = hr.A
b = hr.b
@info("done obtaining H-representation")
# solve ellipsoid problem
Ci, di, CPU_time = lowner_john_inner_ellipsoid(A,b)
@info("done solving inner ellipsoid problem in $(CPU_time) seconds")

# obtain V-representation of the polyhedron
vr = vrep(poly)
vr = MixedMatVRep(vr)
p = vr.V
@info("done obtaining V-representation")
Po, co, CPU_time = lownerjohn_outer(p)
@info("done solving outter ellipsoid problem in $(CPU_time) seconds")
Poinv = Po^-1

# make a plot
if n == 2
    println("Plotting solution...")
    Plots.gr()
    Plots.plot([p[:,1] ; p[1,1]], [p[:,2] ; p[1,2]])
    Plots.plot!(t -> Ci[1,:]' * [ cos(t),sin(t) ] + di[1],
                t -> Ci[2,:]' * [ cos(t),sin(t) ] + di[2],
                0,2*pi)
    Plots.plot!(t -> Poinv[1,:]' * ( [cos(t),sin(t)] .+ co[1] ),
                t -> Poinv[2,:]' * ( [cos(t),sin(t)] .+ co[2] ),
                0,2*pi)
end