#############################
##  Numerical experiments associated with section 5.2.2
#############################
## This file contains the numerical experiments used to compare our method with
## Chebfun's Volterra integral equation implementation when high polynomial orders are required.
## We only include our own side of the implementation here. See the references in the paper for Chebfun.
##

using ApproxFun, MultivariateOrthogonalPolynomials, BandedMatrices, BlockBandedMatrices, SpecialFunctions, Plots, SparseArrays
using SparseVolterraExamples

#####
## First, we can check that the stated arctan function indeed approximates a step-like function.
## This is in Figure 6(a).
v(x,K) = atan(K*x)
    plot(x->v(x,10),0,1, legend=:true , xlabel = "x", label="k=10" ,ylabel = "u_2(x,K)",  legendfontsize=12, tickfontsize=10, thickness_scaling = 1.2 , grid=:none)
    plot!(x->v(x,50),0,1, legend=:true , xlabel = "x", label="k=50" ,ylabel = "u_2(x,K)",  legendfontsize=12, tickfontsize=10, thickness_scaling = 1.2 , grid=:none)
    plot!(x->v(x,100),0,1, legend=:true , xlabel = "x", label="k=100" , ylabel = "u_2(x,k)", legendfontsize=12, tickfontsize=10, thickness_scaling = 1.2 , grid=:none)
    plot!(x->v(x,200),0,1, legend=:bottomright , xlabel = "x", label="k=200" , ylabel = "u_2(x,k)", legendfontsize=12, tickfontsize=10, thickness_scaling = 1.2 , grid=:none)

#####
## First we define the Kernel and g(x,k) as in section 5.2.2.
gf(x,k) = k/(k^2*x^2+1)-(exp(x^2)*atan(k*x))/(2*k^2)+(exp(x^2)*x)/(2*k)-1/2*exp(x^2)*x^2*atan(k*x)
Kfun(x,y) = y*exp(x^2)

#####
## Solver function with step-by-step explanation for given k and with polynomial degree 'n'
function solveSec532(k,n,gf,Kfun)
    gF = Fun(x->gf(x,k),Jacobi(1,2, 0..1),n)             # Approximate g in the appropriate basis
    V = triVolterraFullKernelOpP01(Kfun,n,true,155)      # The following steps generate the appropriate Volterra operator
        V = reflectPabtoPba(n)*WLoweringP01P00(n)*V
        V = Conversion(Jacobi(0,0,0..1),Jacobi(1,2,0..1))[1:n,1:n]*V[1:n,1:n]
        V = Derivative(Jacobi(0,1,0..1),1)[1:n,1:n]-V
    coeff = [DirectEvalLHSP10at0(n);V[1:n-1,1:n]] \ [0;pad(gF.coefficients,n-1)] # Append evaluation and initial conditions and then solve the equation
    return Fun(Jacobi(0,1, 0..1),coeff) # returns solution Fun
end

#####
## Now we can compute and plot a specific example for k=100, plotting both the numerical approximation and the analytic solution.
## Computing errors for various k is a bit more work since automatic convergence testing is not currently implemented but a straightforward for loop can do it from here.
u = solveSec532(100,50,gf,Kfun)
plot(u, label="sparse method")
    plot!(x->v(x,100),0,1, legend=:bottomright , xlabel = "x", label="analytic" ,ylabel = "u_2(x,100)",  legendfontsize=12, tickfontsize=10, thickness_scaling = 1.2 , grid=:none)

#####
## We can also use BenchmarkTools to get a somewhat robust time estimation.
## Note however that this obviously strongly depends on the hardware you are using.
using BenchmarkTools
@benchmark solveSec532(100,50,gf,Kfun)
@benchmark solveSec532(200,300,gf,Kfun)

#####
## Now we plot the operator bandedness for the k=100 example, this is basically Figure 6(b).
function OperatorSec532(k,n,Kfun)
    V = triVolterraFullKernelOpP01(Kfun,n,true,155)    # The following steps generate the appropriate Volterra operator
        V = reflectPabtoPba(n)*WLoweringP01P00(n)*V
        V = Conversion(Jacobi(0,0,0..1),Jacobi(1,2,0..1))[1:n,1:n]*V[1:n,1:n]
        V = Derivative(Jacobi(0,1,0..1),1)[1:n,1:n]-V
        V = [DirectEvalLHSP10at0(n);V[1:n-1,1:n]] # Append evaluation and initial conditions
    return V # returns operator
end
#####
## Plots.jl's spy plot is not currently compatible with these types, so we instead convert to generic sparse for visualization.
## The exact bandedness properties depend on the chosen degrees and parameters but here is a standard example
spy(sparse(OperatorSec532(100,300,Kfun)),markersize=2.8,marker=:rect)

examples/Section 5-2-2 - High order solutions with step-like function.jl


using Jacobi
using Polynomials

# Script for writing the legerdre_RefLine file with Legendre polynomials up to
# a maximum degree of max_degree.
max_degree = 50

setprecision(BigFloat, 512)

function weights(points)
    N = length(points)
    return [ 2/((1-x^2)*(dlegendre(x, N)^2)) for x in points ]
end

open("legendre_line.jl", "w") do io
    for i = 1:max_degree
        # Legendre = Jacobi with α = β = 0
        x = jacobi_zeros(i, 0, 0, BigFloat)
        w = weights(x)/2
        println("Error in sum of the weights for degree $i: $(1-sum(w))")
        x = (x .+ 1)/2
        println(io, "@generated function gauss_quadrature(form::Val{:legendre},")
        println(io, "                                     shape::RefLine,")
        println(io, "                                     degree::Val{$i},")
        println(io, "                                     type::Type{T}) where {T}")
        println(io, """    weights = SVector(:(\$(T(big"$(w[1])"))),""")
        for j in 2:i
        println(io, """                      :(\$(T(big"$(w[j])"))),""")
        end
        println(io, "                       )")
        println(io, """    points = SVector(:(\$(NTuple{1,T}(big"$(x[1])"))),""")
        for j in 2:i
        println(io, """                     :(\$(NTuple{1,T}(big"$(x[j])"))),""")
        end
        println(io, "                      )")
        println(io, "    return weights, points")
        println(io, "end")
    end
end


#=
NacaThicknessFunction.jl
Jacob Child
February 3, 2022
Description:
NACA 4-series Thickness Formula The NACA 4-series formula for airfoil thickness
along the chord is given by the following polynomial function, where m is the value of maximum
thickness (as a percentage of the chord), and x is the x-position along the chord. (This formulation
is for airfoils with sharp trailing edges.)
t = 10m(.2969sqrt(x) - .1620x - .3537x^2 + .2843x^3 - .1015x^4)

Pseudo code- Have a function to ask for and generate x values in an array
Then have a function to calculate various thicknesses
Format as table? and output to screen
***use docstrings***
=#

"""
xinput(x, step)
Takes in the desired length, l, and step size and outputs an x array

"""
function x_input(L, step)
    x = [0:step:L;] #makes array stopping at L with step size step
    return x
end

"""
thickness_calculations(x)
Takes in the x coordinates, calculates thickness, and ouputs the array

"""
function thickness_calculations(x)
    m = 0.10
    t = 10*m.*(.2969*sqrt.(x) - .1260*x - .3537*x.^2 + .2843*x.^3 - .1015*x.^4)
    return t
end

#using Pkg  #Just for the first run through
#Pkg.add("DataFrames")  #just for the first run through

println("Hello! We are going to calculate the thickness along an airfoil \n")
println("\n Enter the Chord length")
L = readline()
L = parse(Int64, L)

println("\n Enter the desired step size to calculate the thickness at")
step = readline()
step = parse(Float64, step)

#println("Type of Inputs", typeof(L), typeof(step))

x = x_input(L, step)
t = thickness_calculations(x)

println(t) #If Running in VS Code Uncomment this as I couldn't get DataFrame to work

#using DataFrames #Only seems to work in Julia Terminal, not VS Code Terminal? ***research

#DataFrame()

#DataFrame("X" => x, "Thickness" => t)


function tournament_selection(P, a = rand(1:length(P)))
    # chose two different solutions at random
    b = rand(1:length(P))
    while a == b
        b = rand(1:length(P))
    end

    # perform selection
    P[a].rank < P[b].rank || (P[a].rank == P[b].rank && P[a].crowding > P[b].crowding ) ? P[a] : P[b]
end

function gen_β(β, η, D, R)
    α = 2.0 .- β .^ (-  η - 1.0 )
    mask = R .<= 1.0 ./ α
    s = 1.0 / (η + 1.0)
    βq = [ mask[i] ?  (R[i] * α[i])^s : (1.0 / (2.0 - R[i]*α[i]))^s for i in 1:D]
    βq
end

function SBX_crossover(vector1, vector2, bounds, η=15, p_variable = 0.9)
    xu = view(bounds, 2,:)
    xl = view(bounds, 1,:)
    D = length(vector1)

    do_crossover = ones(Bool, D)
    do_crossover[rand(D) .> p_variable] .= false
    do_crossover[ abs.( vector2 - vector1 ) .<= eps() ] .= false

    y1 = min.( vector1, vector2 )
    y2 = max.( vector1, vector2 )
    Δ = max.(eps(), y2 - y1)

    R = rand(D)

    β = @. 1.0 + (2.0 * (y1 - xl) / Δ)
    βq = gen_β(β, η, D, R)
    c1 = @. 0.5*(y1 + y2 -  βq*Δ)

    β = @. 1.0 + (2.0 * (y1 - xl) / Δ)
    βq = gen_β(β, η, D, R)
    c2 = @. 0.5*(y1 + y2 +  βq*Δ)

    # swap
    mask = rand(Bool, D)
    cc = copy(c1)
    c1[mask] = c2[mask]
    c2[mask] = cc[mask]

    cc1 = copy(vector1)
    cc1[do_crossover] = c1[do_crossover]
    cc2 = copy(vector2)
    cc2[do_crossover] = c2[do_crossover]

    reset_to_violated_bounds!(cc1, bounds)
    reset_to_violated_bounds!(cc2, bounds)

    return cc1, cc2
end

function polynomial_mutation!(vector, bounds, η=20, prob = 1 / length(vector))
    do_mutation = rand(length(vector)) .< prob

    xu = view(bounds, 2,do_mutation)
    xl = view(bounds, 1,do_mutation)
    x = view(vector, do_mutation)

    δ1 = (x - xl) ./ (xu - xl)
    δ2 = (xu - x) ./ (xu - xl)

    D = length(xu)
    R = rand(D)
    mask = R .< 0.5
    s = η+1.0
    mut_pow = 1.0 / (η + 1.0)
    δq = [ mask[i] ?
            ^(2.0R[i] + (1. - 2.0R[i]) * ^(1.0 - δ1[i], s), mut_pow) - 1.0 :
            1.0 - (2.0 * (1.0 - R[i]) + 2.0 * (R[i] - 0.5) * ^(1.0 - δ2[i], s))^mut_pow
            for i in 1:D
        ]

    vector[do_mutation] = x + δq .* ( xu - xl)
    # correct using reset to bound
    #
    vector

end


function bisectionMethod(f, a, b, tol; maxiter = 1000) #f=@(x)x^2-3; a=1; b=2; (ensure change of sign between a and b) error=1e-4
    if sign(f(a)) == sign(f(b))
        throw(DomainError("Function has the same sign on both ends in bisectionMethod."))
    end
    c = (a + b) / 2

    _fc = f(c)
    _fa = f(a)
    iter = 1
    while abs(_fc) > tol
        if _fc < 0 && _fa < 0 || _fc > 0 && _fa > 0
            a = c
            _fa = _fc
        else
            b = c
        end
        c = (a + b) / 2
        _fc = f(c)
        iter += 1
        iter < maxiter || throw(
            DomainError("Maximum number of iterations reached in bisectionMethod $iter."),
        )
    end
    return c
end

function bisectionMethodError(f, a, b, tol; maxiter = 1000) #f=@(x)x^2-3; a=1; b=2; (ensure change of sign between a and b) error=1e-4
    c = (a + b) / 2
    ftry(x) =
        try
            f(x)
        catch
            1.0
        end

    bisectionMethod(ftry, a, b, tol; maxiter = maxiter)
end

const vh2o_0 = m / rhoh2o
find_bh2o(the) = (bh2o) -> vh2o_0 - vh2o(the, bh2o)
bh2o0(the) = bisectionMethodError(find_bh2o(the), zmax - hb, zmax, 1e-12) # (zmax=>bh2o=>zmax-hb)

function get_bh2o(z, θ, p)
    bh2o = p.bh2o_0 - z / cos(θ)
    # if bh2o < zmax - hb || bh2o > zmax
    #     @warn "bh2o value in boatode! out of range" bh2o
    # end
    return bh2o
end


###############################################################################
#
# Collection of methods generating random projectors.
#
# Improvements:
#  * random_matrix_4
#  * include sparse projectors
#  * rank-1 projectors
#
###############################################################################

"""
Random projector of size k x m
with a normal distribution
"""
function random_matrix_1(k, m)
  P = zeros(k, m) #init

  #normal distribution
  P = rand(k, m)

  return 1 / sqrt(k) * P
end

"""
Random projector of size k x m
with -1 or 1 both with probability 1/2
"""
function random_matrix_2(k, m)
  P = zeros(k, m) #init

  #-1 or 1 both with probability 1/2
  P = rand([-1, 1], (k, m))

  return 1 / sqrt(k) * P
end

"""
Random projector of size k x m
with -1,0,1 respectively with probability 1/6,4/6,1/6
"""
function random_matrix_3(k, m)
  P = zeros(k, m) #init

  #-1,0,1 respectively with probability 1/6,4/6,1/6
  P = rand([-1, 0, 0, 1], (k, m))

  return 1 / sqrt(k) * P
end

"""
Random projector of size k x m
with orthogonal projection on a random k-dimensional linear subspace of R^m
"""
function random_matrix_4(k, m)
  P = zeros(k, m) #init

  #orthogonal projection on a random k-dimensional
  #linear subspace of R^m
  #P = rand(k, m)
  throw("NotImplemented")
  return 1 / sqrt(k) * P
end

"""
Check the average sparsity of random projector over N random matrices of size k*m
"""
function random_projector_sparse(projector::Function, N::Int64, k::Int64, m::Int64)
  tot = 0 #nb of non-zero elements

  for i = 1:N
    tot += SparseArrays.nnz(SparseArrays.sparse(projector(k, m))) / (k * m)
  end

  return tot / N
end

"""
Check the average rank of random projector over N random matrices of size k*m
"""
function random_projector_rank(projector::Function, N::Int64, k::Int64, m::Int64)
  tot = 0 #average rank of the projector

  for i = 1:N
    tot += rank(projector(k, m))
  end

  return tot / N
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    laplacematrix(mesh; weights=:uniform, normalize=true)

The Laplace-Beltrami (a.k.a. Laplacian) matrix of the `mesh`.
Optionally specify the discretization `weights` as either
`:uniform` or `:cotangent` and `normalize` the rows by the
diagonal value.

## References

* Vallet, B & Lévy, B. 2008. [Spectral Geometry Processing with Manifold
  Harmonics](https://onlinelibrary.wiley.com/doi/10.1111/j.1467-8659.2008.01122.x)

* Zhang et al. 2007. [Spectral Methods for Mesh Processing and Analysis]
  (https://diglib.eg.org/handle/10.2312/egst.20071052.001-022)
"""
function laplacematrix(mesh; weights=:uniform, normalize=true)
  # convert to half-edge topology
  m = topoconvert(HalfEdgeTopology, mesh)

  # retrieve adjacency relation
  t = topology(m)
  𝒩 = Adjacency{0}(t)

  # initialize matrix
  n = nvertices(t)
  L = spzeros(n, n)

  # fill matrix with weights
  if weights == :uniform
    for i in 1:n
      js = 𝒩(i)
      for j in js
        L[i,j] = 1.0
      end
      L[i,i] = -1.0*length(js)
      if normalize
        for j in js
          L[i,j] /= -L[i,i]
        end
        L[i,i] /= -L[i,i]
      end
    end
  elseif weights == :cotangent
    v = vertices(m)
    for i in 1:n
      js = CircularVector(𝒩(i))
      for k in 1:length(js)
        j₋, j, j₊ = js[k-1], js[k], js[k+1]
        vᵢ, vⱼ =  v[i],  v[j]
        v₋, v₊ = v[j₋], v[j₊]
        αᵢⱼ = ∠(vⱼ, v₋, vᵢ)
        βᵢⱼ = ∠(vᵢ, v₊, vⱼ)
        L[i,j] = cot(αᵢⱼ) + cot(βᵢⱼ)
      end
      L[i,i] = -sum(L[i,js])
      if normalize
        for j in js
          L[i,j] /= -L[i,i]
        end
        L[i,i] /= -L[i,i]
      end
    end
  else
    throw(ArgumentError("invalid discretization weights"))
  end

  L
end


logistic(x::Float64) = 4*x*(1-x)

function bernoulli(x::Float64, l::Float64, u::Float64, nc::Float64)
  if l<=x && x<nc
    return ((u-l)*x+(nc-u)*l)/(nc-l)
  elseif nc<=x && x<=u
    return ((u-l)*x+(l-nc)*u)/(u-nc)
  else
    error("Input out of domain of Bernoulli map.")
  end
end

function nbernoulli(x::Float64, l::Float64, u::Float64, nc::Float64)
  if l<=x && x<nc
    return ((l-u)*x+(nc*u-l^2))/(nc-l)
  elseif nc<=x && x<=u
    return ((l-u)*x+(u^2-l*nc))/(u-nc)
  else
    error("Input out of domain of negative Bernoulli map.")
  end
end

function tent(x::Float64, l::Float64, u::Float64, nc::Float64)
  if l<=x && x<nc
    ((u-l)*x+(nc-u)*l)/(nc-l)
  elseif nc<=x && x<=u
    ((l-u)*x+(u^2-l*nc))/(u-nc)
  else
    error("Input out of domain of tent map.")
  end
end

function valley(x::Float64, l::Float64, u::Float64, nc::Float64)
  if l<=x && x<nc
    ((l-u)*x+(nc*u-l^2))/(nc-l)
  elseif nc<=x && x<=u
    ((u-l)*x+(l-nc)*u)/(u-nc)
  else
    error("Input out of domain of valley map.")
  end
end

function circular(x::Float64, nc::Float64)
  if -1.<=x && x<-sqrt(nc)
    y = -sqrt((1.-x^2)/(1.-nc))
  elseif -sqrt(nc)<=x && x<sqrt(nc)
    y = sqrt(1.-x^2/nc)
  elseif sqrt(nc)<=x && x<=1.
    y = -sqrt((1.-x^2)/(1.-nc))
  else
    error("Input out of domain of circular map.")
  end
end

immutable VDist <: ContinuousUnivariateDistribution
    nc::Float64

    function VDist(nc::Real)
      @assert -1.0<nc && nc<1.0 "Non-centrality parameter must be in (-1, 1) in the case circular map."
      new(float64(nc))
    end

    VDist() = VDist(0.0)
end

@continuous_distr_support VDist -1.0 1.0

function pdf(d::VDist, x::Real)
  if -1.0<=x && x<=0.0
    -2.0*(1.0-d.nc)*x
  elseif 0.0<x && x<=1.0
    2*d.nc*x
  else
    0.0
  end
end

function rand(d::VDist)
  u = rand()

  if 0.0<=u && u<1.0-d.nc
    -sqrt((u+d.nc-1.0)/(d.nc-1.0))
  elseif 1.0-d.nc<=u && u<=1.0
    sqrt((u+d.nc-1.0)/d.nc)
  end
end

mean(d::VDist) = (4.0*d.nc-2.0)/3.0

var(d::VDist) = 0.5-mean(d)^2


using LinearAlgebra

export fromroots,
       truncate!,
       chop!,
       coeffs,
       degree,
       domain,
       mapdomain,
       order,
       hasnan,
       roots,
       companion,
       vander,
       fit,
       integrate,
       derivative,
       variable,
       isintegral,
       ismonic

"""
    fromroots(::AbstractVector{<:Number}; var=:x)
    fromroots(::Type{<:AbstractPolynomial}, ::AbstractVector{<:Number}; var=:x)

Construct a polynomial of the given type given the roots. If no type is given, defaults to `Polynomial`.

# Examples
```jldoctest common
julia> using Polynomials

julia> r = [3, 2]; # (x - 3)(x - 2)

julia> fromroots(r)
Polynomial(6 - 5*x + x^2)
```
"""
function fromroots(P::Type{<:AbstractPolynomial}, roots::AbstractVector; var::SymbolLike = :x)
    x = variable(P, var)
    p =  prod(x - r for r in roots)
    return truncate!(p)
end
fromroots(r::AbstractVector{<:Number}; var::SymbolLike = :x) =
    fromroots(Polynomial, r, var = var)

"""
    fromroots(::AbstractMatrix{<:Number}; var=:x)
    fromroots(::Type{<:AbstractPolynomial}, ::AbstractMatrix{<:Number}; var=:x)

Construct a polynomial of the given type using the eigenvalues of the given matrix as the roots. If no type is given, defaults to `Polynomial`.

# Examples
```jldoctest common
julia> using Polynomials

julia> A = [1 2; 3 4]; # (x - 5.37228)(x + 0.37228)

julia> fromroots(A)
Polynomial(-1.9999999999999998 - 5.0*x + 1.0*x^2)
```
"""
fromroots(P::Type{<:AbstractPolynomial},
    A::AbstractMatrix{T};
    var::SymbolLike = :x,) where {T <: Number} = fromroots(P, eigvals(A), var = var)
fromroots(A::AbstractMatrix{T}; var::SymbolLike = :x) where {T <: Number} =
    fromroots(Polynomial, eigvals(A), var = var)

"""
    fit(x, y, deg=length(x) - 1; [weights], var=:x)
    fit(::Type{<:AbstractPolynomial}, x, y, deg=length(x)-1; [weights], var=:x)

Fit the given data as a polynomial type with the given degree. Uses linear least squares. When weights are given, as either a `Number`, `Vector` or `Matrix`, will use weighted linear least squares. The default polynomial type is [`Polynomial`](@ref). This will automatically scale your data to the [`domain`](@ref) of the polynomial type using [`mapdomain`](@ref)
"""
function fit(P::Type{<:AbstractPolynomial},
             x::AbstractVector{T},
             y::AbstractVector{T},
             deg::Integer = length(x) - 1;
    weights = nothing,
    var = :x,) where {T}
    x = mapdomain(P, x)
    vand = vander(P, x, deg)
    if weights !== nothing
        coeffs = _wlstsq(vand, y, weights)
    else
        coeffs = pinv(vand) * y
    end
    return P(T.(coeffs), var)
end

fit(P::Type{<:AbstractPolynomial},
    x,
    y,
    deg::Integer = length(x) - 1;
    weights = nothing,
    var = :x,) = fit′(P, promote(collect(x), collect(y))..., deg; weights = weights, var = var)

#  avoid issue  214
fit′(P::Type{<:AbstractPolynomial}, x, y, args...;kwargs...) = throw(MethodError("x and y do not produce abstract   vectors"))
fit′(P::Type{<:AbstractPolynomial},
     x::AbstractVector{T},
     y::AbstractVector{T},
     args...; kwargs...) where {T} = fit(P,x,y,args...;  kwargs...)

fit(x::AbstractVector,
    y::AbstractVector,
    deg::Integer = length(x) - 1;
    weights = nothing,
    var = :x,) = fit(Polynomial, x, y, deg; weights = weights, var = var)

# Weighted linear least squares
_wlstsq(vand, y, W::Number) = _wlstsq(vand, y, fill!(similar(y), W))
_wlstsq(vand, y, W::AbstractVector) = _wlstsq(vand, y, diagm(0 => W))
_wlstsq(vand, y, W::AbstractMatrix) = (vand' * W * vand) \ (vand' * W * y)

"""
    roots(::AbstractPolynomial; kwargs...)

Returns the roots of the given polynomial. This is calculated via the eigenvalues of the companion matrix. The `kwargs` are passed to the `LinearAlgeebra.eigvals` call.

!!! note

        The [PolynomialRoots.jl](https://github.com/giordano/PolynomialRoots.jl) package provides an alternative that is a bit faster and a bit more accurate; the [FastPolynomialRoots](https://github.com/andreasnoack/FastPolynomialRoots.jl) provides an interface to FORTRAN code implementing an algorithm that can handle very large polynomials (it is  `O(n^2)` not `O(n^3)`. the [AMRVW.jl](https://github.com/jverzani/AMRVW.jl) package implements the algorithm in Julia, allowing the use of other  number types.

"""
function roots(q::AbstractPolynomial{T}; kwargs...) where {T <: Number}

    p = convert(Polynomial{T},  q)
    roots(p; kwargs...)

end

"""
    companion(::AbstractPolynomial)

Return the companion matrix for the given polynomial.

# References
[Companion Matrix](https://en.wikipedia.org/wiki/Companion_matrix)
"""
companion(::AbstractPolynomial)

"""
    vander(::Type{AbstractPolynomial}, x::AbstractVector, deg::Integer)

Calculate the psuedo-Vandermonde matrix of the given polynomial type with the given degree.

# References
[Vandermonde Matrix](https://en.wikipedia.org/wiki/Vandermonde_matrix)
"""
vander(::Type{<:AbstractPolynomial}, x::AbstractVector, deg::Integer)

"""
    integrate(::AbstractPolynomial, C=0)

Returns the indefinite integral of the polynomial with constant `C`.
"""
integrate(p::AbstractPolynomial, C::Number = 0) = integrate(p, C)

"""
    integrate(::AbstractPolynomial, a, b)

Compute the definite integral of the given polynomial from `a` to `b`. Will throw an error if either `a` or `b` are out of the polynomial's domain.
"""
function integrate(p::AbstractPolynomial, a::Number, b::Number)
    P = integrate(p)
    return P(b) - P(a)
end

"""
    derivative(::AbstractPolynomial, order::Int = 1)

Returns a polynomial that is the `order`th derivative of the given polynomial. `order` must be non-negative.
"""
derivative(::AbstractPolynomial, ::Int)

"""
    truncate!(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0)

In-place version of [`truncate`](@ref)
"""
function truncate!(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
                   atol::Real = 0,) where {T}
    max_coeff = maximum(abs, coeffs(p))
    thresh = max_coeff * rtol + atol
    map!(c->abs(c) <= thresh ? zero(T) : c, coeffs(p), coeffs(p))
    return chop!(p, rtol = rtol, atol = atol)
end

"""
    truncate(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0)

Rounds off coefficients close to zero, as determined by `rtol` and `atol`, and then chops any leading zeros. Returns a new polynomial.
"""
function Base.truncate(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
    atol::Real = 0,) where {T}
    truncate!(deepcopy(p), rtol = rtol, atol = atol)
end

"""
    chop!(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0))

In-place version of [`chop`](@ref)
"""
function chop!(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
               atol::Real = 0,) where {T}
    isempty(coeffs(p)) && return p
    tol = norm(coeffs(p)) * rtol + atol
    for i = lastindex(p):-1:0
        val = p[i]
        if abs(val) > tol #!isapprox(val, zero(T); rtol = rtol, atol = atol)
            resize!(p.coeffs, i + 1);
            return p
        end
    end
    resize!(p.coeffs, 1)
    return p
end

"""
    chop(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0))

Removes any leading coefficients that are approximately 0 (using `rtol` and `atol`). Returns a polynomial whose degree will guaranteed to be equal to or less than the given polynomial's.
"""
function Base.chop(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
    atol::Real = 0,) where {T}
    chop!(deepcopy(p), rtol = rtol, atol = atol)
end

"""
    check_same_variable(p::AbstractPolynomial, q::AbstractPolynomial)

Check if either `p` or `q` is constant or if `p` and `q` share the same variable
"""
check_same_variable(p::AbstractPolynomial, q::AbstractPolynomial) =
    (Polynomials.isconstant(p) || Polynomials.isconstant(q)) || p.var ==  q.var

#=
Linear Algebra =#
"""
    norm(::AbstractPolynomial, p=2)

Calculates the p-norm of the polynomial's coefficients
"""
LinearAlgebra.norm(q::AbstractPolynomial, p::Real = 2) = norm(coeffs(q), p)

"""
    conj(::AbstractPolynomial)

Returns the complex conjugate of the polynomial
"""
LinearAlgebra.conj(p::P) where {P <: AbstractPolynomial} = ⟒(P)(conj(coeffs(p)), p.var)
LinearAlgebra.adjoint(p::P) where {P <: AbstractPolynomial} = ⟒(P)(adjoint.(coeffs(p)), p.var)
LinearAlgebra.transpose(p::AbstractPolynomial) = p
LinearAlgebra.transpose!(p::AbstractPolynomial) = p

#=
Conversions =#
Base.convert(::Type{P}, p::P) where {P <: AbstractPolynomial} = p
Base.convert(P::Type{<:AbstractPolynomial}, x) = P(x)
Base.promote_rule(::Type{<:AbstractPolynomial{T}},
    ::Type{<:AbstractPolynomial{S}},
) where {T,S} = Polynomial{promote_type(T, S)}

#=
Inspection =#
"""
    length(::AbstractPolynomial)

The length of the polynomial.
"""
Base.length(p::AbstractPolynomial) = length(coeffs(p))
"""
    size(::AbstractPolynomial, [i])

Returns the size of the polynomials coefficients, along axis `i` if provided.
"""
Base.size(p::AbstractPolynomial) = size(coeffs(p))
Base.size(p::AbstractPolynomial, i::Integer) = size(coeffs(p), i)
Base.eltype(p::AbstractPolynomial{T}) where {T} = T
# in  analogy  with  polynomial as a Vector{T} with different operations defined.
Base.eltype(::Type{<:AbstractPolynomial}) = Float64
Base.eltype(::Type{<:AbstractPolynomial{T}}) where {T} = T
#Base.eltype(::Type{P}) where {P <: AbstractPolynomial} = P # changed  in v1.1.0
function Base.iszero(p::AbstractPolynomial)
    if length(p) == 0
        return true
    end
    return all(iszero.(coeffs(p))) && p[0] == 0
end

# See discussions in https://github.com/JuliaMath/Polynomials.jl/issues/258
"""
    all(pred, poly::AbstractPolynomial)

Test whether all coefficients of an `AbstractPolynomial` satisfy predicate `pred`.

You can implement `isreal`, etc., to a `Polynomial` by using `all`.
"""
Base.all(pred, poly::AbstractPolynomial) = all(pred, poly[:])
"""
    any(pred, poly::AbstractPolynomial)

Test whether any coefficient of an `AbstractPolynomial` satisfies predicate `pred`.
"""
Base.any(pred, poly::AbstractPolynomial) = any(pred, poly[:])
"""
    map(fn, p::AbstractPolynomial)

Transform coefficients of `p` by applying a function (or other callables) `fn` to each of them.

You can implement `real`, etc., to a `Polynomial` by using `map`.
"""
Base.map(fn, p::P) where {P<:AbstractPolynomial} = ⟒(P)(map(fn, coeffs(p)), p.var)

"""
    isreal(p::AbstractPolynomial)

Determine whether a polynomial is a real polynomial, i.e., having only real numbers as coefficients.

See also: [`real`](@ref)
"""
Base.isreal(p::AbstractPolynomial) = all(isreal, p)
"""
    real(p::AbstractPolynomial)

Construct a real polynomial from the real parts of the coefficients of `p`.

See also: [`isreal`](@ref)

!!! note
    This could cause losing terms in `p`. This method is usually called on polynomials like `p = Polynomial([1, 2 + 0im, 3.0, 4.0 + 0.0im])` where you want to chop the imaginary parts of the coefficients of `p`.
"""
Base.real(p::AbstractPolynomial) = map(real, p)

"""
    isintegral(p::AbstractPolynomial)

Determine whether a polynomial is an integer polynomial, i.e., having only integers as coefficients.
"""
isintegral(p::AbstractPolynomial) = all(isinteger, p)

"""
    ismonic(p::AbstractPolynomial)

Determine whether a polynomial is a monic polynomial, i.e., its leading coefficient is one.
"""
ismonic(p::AbstractPolynomial) = isone(p[end])

"""
    coeffs(::AbstractPolynomial)

Return the coefficient vector `[a_0, a_1, ..., a_n]` of a polynomial.
"""
coeffs(p::AbstractPolynomial) = p.coeffs

"""
    degree(::AbstractPolynomial)

Return the degree of the polynomial, i.e. the highest exponent in the polynomial that
has a nonzero coefficient. The degree of the zero polynomial is defined to be -1.
"""
degree(p::AbstractPolynomial) = iszero(p) ? -1 : length(p) - 1

"""
    isconstant(::AbstractPolynomial)

Is the polynomial  `p` a constant.
"""
isconstant(p::AbstractPolynomial) = degree(p) <= 0

hasnan(p::AbstractPolynomial) = any(isnan.(coeffs(p)))

"""
    domain(::Type{<:AbstractPolynomial})

Returns the domain of the polynomial.
"""
domain(::Type{<:AbstractPolynomial})
domain(::P) where {P <: AbstractPolynomial} = domain(P)

"""
    mapdomain(::Type{<:AbstractPolynomial}, x::AbstractArray)
    mapdomain(::AbstractPolynomial, x::AbstractArray)

Given values of x that are assumed to be unbounded (-∞, ∞), return values rescaled to the domain of the given polynomial.

# Examples
```jldoctest  common
julia> using Polynomials

julia> x = -10:10
-10:10

julia> extrema(mapdomain(ChebyshevT, x))
(-1.0, 1.0)

```
"""
function mapdomain(P::Type{<:AbstractPolynomial}, x::AbstractArray)
    d = domain(P)
    x = collect(x)
    x_zerod = x .- minimum(x)
    x_scaled = x_zerod .* (last(d) - first(d)) ./ maximum(x_zerod)
    x_scaled .+= first(d)
    return x_scaled
end
mapdomain(::P, x::AbstractArray) where {P <: AbstractPolynomial} = mapdomain(P, x)
#=
indexing =#
Base.firstindex(p::AbstractPolynomial) = 0
Base.lastindex(p::AbstractPolynomial) = length(p) - 1
Base.eachindex(p::AbstractPolynomial) = 0:length(p) - 1
Base.broadcastable(p::AbstractPolynomial) = Ref(p)

# iteration
# iteration occurs over the basis polynomials
Base.iterate(p::AbstractPolynomial) = (p[0] * one(typeof(p)), 1)
function Base.iterate(p::AbstractPolynomial, state)
    state <= length(p) - 1 ? (p[state] * basis(p, state), state + 1) : nothing
end

Base.collect(p::P) where {P <: AbstractPolynomial} = collect(P, p)

# getindex
function Base.getindex(p::AbstractPolynomial{T}, idx::Int) where {T <: Number}
    idx < 0 && throw(BoundsError(p, idx))
    idx ≥ length(p) && return zero(T)
    return coeffs(p)[idx + 1]
end
Base.getindex(p::AbstractPolynomial, idx::Number) = getindex(p, convert(Int, idx))
Base.getindex(p::AbstractPolynomial, indices) = [getindex(p, i) for i in indices]
Base.getindex(p::AbstractPolynomial, ::Colon) = coeffs(p)

# setindex
function Base.setindex!(p::AbstractPolynomial, value::Number, idx::Int)
    n = length(coeffs(p))
    if n ≤ idx
        resize!(p.coeffs, idx + 1)
        p.coeffs[n + 1:idx] .= 0
    end
    p.coeffs[idx + 1] = value
    return p
end

Base.setindex!(p::AbstractPolynomial, value::Number, idx::Number) =
    setindex!(p, value, convert(Int, idx))
Base.setindex!(p::AbstractPolynomial, value::Number, indices) =
    [setindex!(p, value, i) for i in indices]
Base.setindex!(p::AbstractPolynomial, values, indices) =
    [setindex!(p, v, i) for (v, i) in zip(values, indices)]
Base.setindex!(p::AbstractPolynomial, value::Number, ::Colon) =
    setindex!(p, value, eachindex(p))
Base.setindex!(p::AbstractPolynomial, values, ::Colon) =
    [setindex!(p, v, i) for (v, i) in zip(values, eachindex(p))]

#=
identity =#
Base.copy(p::P) where {P <: AbstractPolynomial} = P(copy(coeffs(p)), p.var)
Base.hash(p::AbstractPolynomial, h::UInt) = hash(p.var, hash(coeffs(p), h))

#=
zero, one, variable, basis =#
"""
    zero(::Type{<:AbstractPolynomial})
    zero(::AbstractPolynomial)

Returns a representation of 0 as the given polynomial.
"""
Base.zero(::Type{P}, var=:x) where {P <: AbstractPolynomial} = ⟒(P)(zeros(eltype(P), 1), var)
Base.zero(p::P) where {P <: AbstractPolynomial} = zero(P, p.var)
"""
    one(::Type{<:AbstractPolynomial})
    one(::AbstractPolynomial)

Returns a representation of 1 as the given polynomial.
"""
Base.one(::Type{P}, var=:x) where {P <: AbstractPolynomial} = ⟒(P)(ones(eltype(P),1), var)  # assumes  p₀ = 1
Base.one(p::P) where {P <: AbstractPolynomial} = one(P, p.var)

Base.oneunit(::Type{P}, args...) where {P <: AbstractPolynomial} = one(P, args...)
Base.oneunit(p::P, args...) where {P <: AbstractPolynomial} = one(p, args...)

"""
    variable(var=:x)
    variable(::Type{<:AbstractPolynomial}, var=:x)
    variable(p::AbstractPolynomial, var=p.var)

Return the monomial `x` in the indicated polynomial basis.  If no type is give, will default to [`Polynomial`](@ref). Equivalent  to  `P(var)`.

# Examples
```jldoctest  common
julia> using Polynomials

julia> x = variable()
Polynomial(x)

julia> p = 100 + 24x - 3x^2
Polynomial(100 + 24*x - 3*x^2)

julia> roots((x - 3) * (x + 2))
2-element Array{Float64,1}:
 -2.0
  3.0

```
"""
variable(::Type{P}, var::SymbolLike = :x) where {P <: AbstractPolynomial} = MethodError()
variable(p::AbstractPolynomial, var::SymbolLike = p.var) = variable(typeof(p), var)
variable(var::SymbolLike = :x) = variable(Polynomial{Int}, var)

# basis
# var is a positional argument, not a keyword; can't deprecate so we do `_var; var=_var`
#@deprecate basis(p::P, k::Int; var=:x)  where {P<:AbstractPolynomial}  basis(p, k, var)
#@deprecate basis(::Type{P}, k::Int; var=:x) where {P <: AbstractPolynomial} basis(P, k,var)
# return the kth basis polynomial for the given polynomial type, e.g. x^k for Polynomial{T}
function basis(::Type{P}, k::Int, _var::SymbolLike=:x; var=_var) where {P <: AbstractPolynomial}
    zs = zeros(Int, k+1)
    zs[end] = 1
    ⟒(P){eltype(P)}(zs, var)
end
basis(p::P, k::Int, _var::SymbolLike=:x; var=_var) where {P<:AbstractPolynomial} = basis(P, k, var)

#=
arithmetic =#
Base.:-(p::P) where {P <: AbstractPolynomial} = P(-coeffs(p), p.var)
Base.:+(c::Number, p::AbstractPolynomial) = +(p, c)
Base.:-(p::AbstractPolynomial, c::Number) = +(p, -c)
Base.:-(c::Number, p::AbstractPolynomial) = +(-p, c)
Base.:*(c::Number, p::AbstractPolynomial) = *(p, c)

function Base.:*(p::P, c::S) where {P <: AbstractPolynomial,S}
    T = promote_type(P, S)
    return T(coeffs(p) .* c, p.var)
end

function Base.:/(p::P, c::S) where {T,P <: AbstractPolynomial{T},S}
    R = promote_type(P, eltype(one(T) / one(S)))
    return R(coeffs(p) ./ c, p.var)
end

Base.:-(p1::AbstractPolynomial, p2::AbstractPolynomial) = +(p1, -p2)

function Base.:+(p::P, n::Number) where {P <: AbstractPolynomial}
    p1, p2 = promote(p, n)
    return p1 + p2
end

function Base.:+(p1::P, p2::O) where {P <: AbstractPolynomial,O <: AbstractPolynomial}
    p1, p2 = promote(p1, p2)
    return p1 + p2
end

function Base.:*(p1::P, p2::O) where {P <: AbstractPolynomial,O <: AbstractPolynomial}
    p1, p2 = promote(p1, p2)
    return p1 * p2
end

Base.:^(p::AbstractPolynomial, n::Integer) = Base.power_by_squaring(p, n)

function Base.divrem(num::P, den::O) where {P <: AbstractPolynomial,O <: AbstractPolynomial}
    n, d = promote(num, den)
    return divrem(n, d)
end

"""
    gcd(a::AbstractPolynomial, b::AbstractPolynomial; atol::Real=0, rtol::Real=Base.rtoldefault)

Find the greatest common denominator of two polynomials recursively using
[Euclid's algorithm](http://en.wikipedia.org/wiki/Polynomial_greatest_common_divisor#Euclid.27s_algorithm).

# Examples

```jldoctest common
julia> using Polynomials

julia> gcd(fromroots([1, 1, 2]), fromroots([1, 2, 3]))
Polynomial(4.0 - 6.0*x + 2.0*x^2)

```
"""
function Base.gcd(p1::AbstractPolynomial{T}, p2::AbstractPolynomial{S}; kwargs...) where {T,S}
    gcd(promote(p1, p2)...; kwargs...)
end

function Base.gcd(p1::P, p2::Q;
                  atol::Real=zero(real(T)),
                  rtol::Real=Base.rtoldefault(real(T))
                  ) where {T, P <: AbstractPolynomial{T}, Q <: AbstractPolynomial{T}}

    r₀, r₁ = p1, p2
    iter = 1
    itermax = length(r₁)

    while !iszero(r₁) && iter ≤ itermax
        _, rtemp = divrem(r₀, r₁)
        r₀ = r₁
        r₁ = truncate(rtemp; atol=atol, rtol=rtol)
        iter += 1
    end
    return r₀
end

"""
    div(::AbstractPolynomial, ::AbstractPolynomial)
"""
Base.div(n::AbstractPolynomial, d::AbstractPolynomial) = divrem(n, d)[1]

"""
    rem(::AbstractPolynomial, ::AbstractPolynomial)
"""
Base.rem(n::AbstractPolynomial, d::AbstractPolynomial) = divrem(n, d)[2]

#=
Comparisons =#
Base.isequal(p1::P, p2::P) where {P <: AbstractPolynomial} = hash(p1) == hash(p2)
Base.:(==)(p1::AbstractPolynomial, p2::AbstractPolynomial) =
    check_same_variable(p1,p2) && (coeffs(p1) == coeffs(p2))
Base.:(==)(p::AbstractPolynomial, n::Number) = degree(p) <= 0 && p[0] == n
Base.:(==)(n::Number, p::AbstractPolynomial) = p == n

function Base.isapprox(p1::AbstractPolynomial{T},
    p2::AbstractPolynomial{S};
    rtol::Real = (Base.rtoldefault(T, S, 0)),
                       atol::Real = 0,) where {T,S}

    p1, p2 = promote(p1, p2)
    check_same_variable(p1, p2)  || error("p1 and p2 must have same var")

    # copy over from abstractarray.jl
    Δ  = norm(p1-p2)
    if isfinite(Δ)
        return Δ <= max(atol, rtol*max(norm(p1), norm(p2)))
    else
        for i in 0:max(degree(p1), degree(p2))
            isapprox(p1[i], p2[i]; rtol=rtol, atol=atol) || return false
        end
        return true
    end
end

function Base.isapprox(p1::P,
                       n::S;
                       rtol::Real = (Base.rtoldefault(T, S, 0)),
                       atol::Real = 0,) where {T,S, P<:AbstractPolynomial{T}}
    return isapprox(p1, ⟒(P){T}(n,p1.var))
end

Base.isapprox(n::S,
    p1::AbstractPolynomial{T};
    rtol::Real = (Base.rtoldefault(T, S, 0)),
    atol::Real = 0,) where {T,S} = isapprox(p1, n, rtol = rtol, atol = atol)


# Lions, Mercier, “Splitting algorithms for the sum of two nonlinear
# operators,” SIAM Journal on Numerical Analysis, vol. 16, pp. 964–979 (1979).

using Base.Iterators
using ProximalAlgorithms.IterationTools
using ProximalOperators: Zero
using LinearAlgebra
using Printf

"""
    ForwardBackwardIteration(; <keyword-arguments>)

Instantiate the forward-backward splitting algorithm (see [1]) for solving
optimization problems of the form

    minimize f(x) + g(x),

where `f` is smooth.

# Arguments
- `x0`: initial point.
- `f=Zero()`: smooth objective term.
- `g=Zero()`: proximable objective term.
- `Lf=nothing`: Lipschitz constant of the gradient of `f`.
- `gamma=nothing`: stepsize to use, defaults to `1/Lf` if not set (but `Lf` is).
- `adaptive=false`: forces the method stepsize to be adaptively adjusted.
- `minimum_gamma=1e-7`: lower bound to `gamma` in case `adaptive == true`.

# References
1. Lions, Mercier, “Splitting algorithms for the sum of two nonlinear operators,” SIAM Journal on Numerical Analysis, vol. 16, pp. 964–979 (1979).
"""
Base.@kwdef struct ForwardBackwardIteration{R,Tx,Tf,Tg,TLf,Tgamma}
    f::Tf = Zero()
    g::Tg = Zero()
    x0::Tx
    Lf::TLf = nothing
    gamma::Tgamma = Lf === nothing ? nothing : (1 / Lf)
    adaptive::Bool = gamma === nothing
    minimum_gamma::R = real(eltype(x0))(1e-7)
end

Base.IteratorSize(::Type{<:ForwardBackwardIteration}) = Base.IsInfinite()

Base.@kwdef mutable struct ForwardBackwardState{R,Tx}
    x::Tx             # iterate
    f_x::R            # value of f at x
    grad_f_x::Tx      # gradient of f at x
    gamma::R          # stepsize parameter of forward and backward steps
    y::Tx             # forward point
    z::Tx             # forward-backward point
    g_z::R            # value of g at z
    res::Tx           # fixed-point residual at iterate (= z - x)
    Az::Tx=similar(x) # TODO not needed
    grad_f_z::Tx=similar(x)
end

function Base.iterate(iter::ForwardBackwardIteration)
    x = copy(iter.x0)
    grad_f_x, f_x = gradient(iter.f, x)
    gamma = iter.gamma === nothing ? 1 / lower_bound_smoothness_constant(iter.f, I, x, grad_f_x) : iter.gamma
    y = x - gamma .* grad_f_x
    z, g_z = prox(iter.g, y, gamma)
    state = ForwardBackwardState(
        x=x, f_x=f_x, grad_f_x=grad_f_x,
        gamma=gamma, y=y, z=z, g_z=g_z, res=x - z,
    )
    return state, state
end

function Base.iterate(iter::ForwardBackwardIteration{R}, state::ForwardBackwardState{R,Tx}) where {R,Tx}
    if iter.adaptive == true
        state.gamma, state.g_z, state.f_x = backtrack_stepsize!(
            state.gamma, iter.f, nothing, iter.g,
            state.x, state.f_x, state.grad_f_x, state.y, state.z, state.g_z, state.res, state.z, state.grad_f_z,
            minimum_gamma = iter.minimum_gamma,
        )
        state.x, state.z = state.z, state.x
        state.grad_f_x, state.grad_f_z = state.grad_f_z, state.grad_f_x
    else
        state.x, state.z = state.z, state.x
        state.f_x = gradient!(state.grad_f_x, iter.f, state.x)
    end

    state.y .= state.x .- state.gamma .* state.grad_f_x
    state.g_z = prox!(state.z, iter.g, state.y, state.gamma)

    state.res .= state.x .- state.z

    return state, state
end

# Solver

struct ForwardBackward{R, K}
    maxit::Int
    tol::R
    verbose::Bool
    freq::Int
    kwargs::K
end

function (solver::ForwardBackward)(x0; kwargs...)
    stop(state::ForwardBackwardState) = norm(state.res, Inf) / state.gamma <= solver.tol
    disp((it, state)) =
        @printf("%5d | %.3e | %.3e\n", it, state.gamma, norm(state.res, Inf) / state.gamma)
    iter = ForwardBackwardIteration(; x0=x0, solver.kwargs..., kwargs...)
    iter = take(halt(iter, stop), solver.maxit)
    iter = enumerate(iter)
    if solver.verbose
        iter = tee(sample(iter, solver.freq), disp)
    end
    num_iters, state_final = loop(iter)
    return state_final.z, num_iters
end

ForwardBackward(; maxit=10_000, tol=1e-8, verbose=false, freq=100, kwargs...) =
    ForwardBackward(maxit, tol, verbose, freq, kwargs)

# Aliases

const ProximalGradientIteration = ForwardBackwardIteration
const ProximalGradient = ForwardBackward


################################################################################
## Linear Regression
##   y ~ N(b0 + b1 * x, s2)
##   b0, b1 ~ N(0, 1000)
##   s2 ~ invgamma(0.001, 0.001)
################################################################################

using Mamba

## Data
data = Dict(
  :x => [1, 2, 3, 4, 5],
  :y => [1, 3, 3, 3, 5]
)

## Log-transformed Posterior(b0, b1, log(s2)) + Constant and Gradient Vector
logfgrad = function(x::DenseVector)
  b0 = x[1]
  b1 = x[2]
  logs2 = x[3]
  r = data[:y] - b0 - b1 * data[:x]
  logf = (-0.5 * length(data[:y]) - 0.001) * logs2 -
           (0.5 * dot(r, r) + 0.001) / exp(logs2) -
           0.5 * b0^2 / 1000 - 0.5 * b1^2 / 1000
  grad = [
    sum(r) / exp(logs2) - b0 / 1000,
    sum(data[:x] .* r) / exp(logs2) - b1 / 1000,
    -0.5 * length(data[:y]) - 0.001 + (0.5 * dot(r, r) + 0.001) / exp(logs2)
  ]
  logf, grad
end

## MCMC Simulation with Metropolis-Adjusted Langevin Algorithm
## Without (1) and with (2) a user-specified proposal covariance matrix
n = 5000
sim1 = Chains(n, 3, names = ["b0", "b1", "s2"])
sim2 = Chains(n, 3, names = ["b0", "b1", "s2"])
epsilon = 0.1
Sigma = eye(3)
theta1 = MALAVariate([0.0, 0.0, 0.0], epsilon, logfgrad)
theta2 = MALAVariate([0.0, 0.0, 0.0], epsilon, Sigma, logfgrad)
for i in 1:n
  sample!(theta1)
  sample!(theta2)
  sim1[i, :, 1] = [theta1[1:2]; exp(theta1[3])]
  sim2[i, :, 1] = [theta2[1:2]; exp(theta2[3])]
end
describe(sim1)
describe(sim2)


"""
    internal_rate_of_return(cashflows::vector)::Yields.Rate
    internal_rate_of_return(cashflows::Vector, timepoints::Vector)::Yields.Rate

Calculate the internal_rate_of_return with given timepoints. If no timepoints given, will assume that a series of equally spaced cashflows, assuming the first cashflow occurring at time zero and subsequent elements at time 1, 2, 3, ..., n.

Returns a Yields.Rate type with periodic compounding once per period (e.g. annual effective if the `timepoints` given represent years). Get the scalar rate by calling `Yields.rate()` on the result.

# Example
```julia-repl
julia> internal_rate_of_return([-100,110],[0,1]) # e.g. cashflows at time 0 and 1
0.10000000001652906
julia> internal_rate_of_return([-100,110]) # implied the same as above
0.10000000001652906
```

# Solver notes
Will try to return a root within the range [-2,2]. If the fast solver does not find one matching this condition, then a more robust search will be performed over the [.99,2] range.

The solution returned will be in the range [-2,2], but may not be the one nearest zero. For a slightly slower, but more robust version, call `ActuaryUtilities.irr_robust(cashflows,timepoints)` directly.
"""
function internal_rate_of_return(cashflows)

    return internal_rate_of_return(cashflows, 0:length(cashflows)-1)

end

function internal_rate_of_return(cashflows,times)
    # first try to quickly solve with newton's method, otherwise
    # revert to a more robust method
    lower,upper = -2.,2.

    v = try
        return irr_newton(cashflows,times)
    catch e
        if isa(e,Roots.ConvergenceFailed) || sprint(showerror, e) =="No convergence"
            return irr_robust(cashflows,times)
        else
            throw(e)
        end
    end

    if v <= upper && v >= lower
        return v
    else
        return irr_robust(cashflows,times)
    end
end

irr_robust(cashflows) = irr_robust(cashflows,0:length(cashflows)-1)

function irr_robust(cashflows, times)
    f(i) =  sum(cf / (1+i)^t for (cf,t) in zip(cashflows,times))
    # lower bound at -.99 because otherwise we can start taking the root of a negative number
    # when a time is fractional.
    roots = Roots.find_zeros(f, -0.99, 2)

    # short circuit and return nothing if no roots found
    isempty(roots) && return nothing
    # find and return the one nearest zero
    min_i = argmin(roots)
    return Yields.Periodic(roots[min_i],1)

end

irr_newton(cashflows) = irr_newton(cashflows,0:length(cashflows)-1)

function irr_newton(cashflows, times)
    # use newton's method with hand-coded derivative
    f(r) =  sum(cf * exp(-r*t) for (cf,t) in zip(cashflows,times))
    f′(r) = sum(-t*cf * exp(-r*t) for (cf,t) in zip(cashflows,times) if t > 0)
    # r = Roots.solve(Roots.ZeroProblem((f,f′), 0.0), Roots.Newton())
    r = Roots.newton(x->(f(x),f(x)/f′(x)),0.0)
    return Yields.Periodic(exp(r)-1,1)

end

"""
    irr(cashflows::vector)
    irr(cashflows::Vector, timepoints::Vector)

    An alias for `internal_rate_of_return`.
"""
irr = internal_rate_of_return

"""
    present_value(interest, cashflows::Vector, timepoints)
    present_value(interest, cashflows::Vector)

Discount the `cashflows` vector at the given `interest_interestrate`,  with the cashflows occurring
at the times specified in `timepoints`. If no `timepoints` given, assumes that cashflows happen at times 1,2,...,n.

The `interest` can be an `InterestCurve`, a single scalar, or a vector wrapped in an `InterestCurve`.

# Examples
```julia-repl
julia> present_value(0.1, [10,20],[0,1])
28.18181818181818
julia> present_value(Yields.Forward([0.1,0.2]), [10,20],[0,1])
28.18181818181818 # same as above, because first cashflow is at time zero
```

Example on how to use real dates using the [DayCounts.jl](https://github.com/JuliaFinance/DayCounts.jl) package
```jldoctest

using DayCounts
dates = Date(2012,12,31):Year(1):Date(2013,12,31)
times = map(d -> yearfrac(dates[1], d, DayCounts.Actual365Fixed()),dates) # [0.0,1.0]
present_value(0.1, [10,20],times)

# output
28.18181818181818

```

"""
function present_value(yc::T, cashflows, timepoints) where {T <: Yields.AbstractYield}
    sum(discount(yc,t) * cf for (t,cf) in zip(timepoints, cashflows))
end

function present_value(yc::T, cashflows) where {T <: Yields.AbstractYield}
    present_value(yc,cashflows,1:length(cashflows))
end

function present_value(i, x)

    v = 1.0
    v_factor = discount(i,0,1)
    pv = 0.0

    for (t,cf) in zip(1:length(x),x)
        v *= v_factor
        pv += v * cf
    end
    return pv
end

function present_value(i, v, times)
    return present_value(Yields.Constant(i), v, times)
end

# Interest Given is an array, assume forwards.
function present_value(i::AbstractArray, v)
    yc = Yields.Forward(i)
    return sum(discount(yc, t) * cf for (t,cf) in zip(1:length(v),v))
end

# Interest Given is an array, assume forwards.
function present_value(i::AbstractArray, v, times)
    yc = Yields.Forward(i, times)
    return sum(discount(yc, t) * cf for (cf, t) in zip(v,times))
end

"""
    pv()

    An alias for `present_value`.
"""
pv = present_value

"""
    present_value(interest, cashflows::Vector, timepoints)
    present_value(interest, cashflows::Vector)

Efficiently calculate a vector representing the present value of the given cashflows at each period prior to the given timepoint.

# Examples
```julia-repl
julia> present_values(0.00, [1,1,1])
[3,2,1]

julia> present_values(Yields.Forward([0.1,0.2]), [10,20],[0,1])
2-element Vector{Float64}:
 28.18181818181818
 18.18181818181818
```

"""
function present_values(interest, cashflows)
    pvs = Vector{Float64}(undef,length(cashflows))
    pvs[end] = Yields.discount(interest, lastindex(cashflows) - 1, lastindex(cashflows)) * cashflows[end]
    for (t, cf) in Iterators.reverse(enumerate(cashflows[1:end - 1]))
        pvs[t] = Yields.discount(interest, t - 1, t) * (cf + pvs[t + 1])
    end

    return pvs
end

function present_values(interest,cashflows,times)
    present_values_accumulator(interest,cashflows,times)
end

function present_values_accumulator(interest,cashflows,times,pvs=[0.0])
    from_time = length(times) == 1 ? 0. : times[end-1]
    pv = discount(interest,from_time,last(times)) *(first(pvs) + last(cashflows))
    pvs = pushfirst!(pvs,pv)

    if length(cashflows) > 1

        new_cfs = @view cashflows[1:end-1]
        new_times = @view times[1:end-1]
        return present_values_accumulator(interest,new_cfs,new_times,pvs)
    else
        # last discount and return
        return pvs[1:end-1] # end-1 get rid of trailing 0.0
    end
end

# if given a vector of rates, assume that it should be a forward discount yield
function present_values(y::Vector{T}, cfs, times) where {T <: Real}
    return present_values(Yields.Forward(y), cfs, times)
end

"""
    price(...)

The absolute value of the `present_value(...)`.

# Extended help

Using `price` can be helpful if the directionality of the value doesn't matter. For example, in the common usage, duration is more interested in the change in price than present value, so `price` is used there.
"""
price(x1,x2) = present_value(x1, x2) |> abs
price(x1,x2,x3) = present_value(x1, x2, x3) |> abs

"""
    breakeven(yield, cashflows::Vector)
    breakeven(yield, cashflows::Vector,times::Vector)

Calculate the time when the accumulated cashflows breakeven given the yield.

Assumptions:

- cashflows occur at the end of the period
- cashflows evenly spaced with the first one occuring at time zero if `times` not given

Returns `nothing` if cashflow stream never breaks even.

```jldoctest
julia> breakeven(0.10, [-10,1,2,3,4,8])
5

julia> breakeven(0.10, [-10,15,2,3,4,8])
1

julia> breakeven(0.10, [-10,-15,2,3,4,8]) # returns the `nothing` value

```
"""
function breakeven(y::T, cashflows::Vector, timepoints::Vector) where {T <: Yields.AbstractYield}
    accum = zero(eltype(cashflows))
    last_neg = nothing

    accum += cashflows[1]
    if accum >= 0 && isnothing(last_neg)
        last_neg = timepoints[1]
    end

    for i in 2:length(cashflows)
        # accumulate the flow from each timepoint to the next
        accum *= Yields.accumulation(y, timepoints[i - 1], timepoints[i])
        accum += cashflows[i]

        if accum >= 0 && isnothing(last_neg)
            last_neg = timepoints[i]
        elseif accum < 0
            last_neg = nothing
        end
    end

    return last_neg

end

function breakeven(y::T, cfs, times) where {T <: Real}
    return breakeven(Yields.Constant(y), cfs, times)
end

function breakeven(y::Vector{T}, cfs, times) where {T <: Real}
    return breakeven(Yields.Forward(y), cfs, times)
end

function breakeven(i, cashflows::Vector)
    return breakeven(i, cashflows, [t for t in 0:length(cashflows) - 1])
end

abstract type Duration end

struct Macaulay <: Duration end
struct Modified <: Duration end
struct DV01 <: Duration end

abstract type KeyRateDuration <: Duration end

"""
    KeyRatePar(timepoint,shift=0.001) <: KeyRateDuration

Shift the par curve by the given amount at the given timepoint. Use in conjunction with `duration` to calculate the key rate duration.

Unlike other duration statistics which are computed using analytic derivatives, `KeyRateDuration`s are computed via a shift-and-compute the yield curve approach.

`KeyRatePar` is more commonly reported (than [`KayRateZero`](@ref)) in the fixed income markets, even though the latter has more analytically attractive properties. See the discussion of KeyRateDuration in the Yields.jl docs.

"""
struct KeyRatePar{T,R} <: KeyRateDuration
    timepoint::T
    shift::R
    KeyRatePar(timepoint, shift=.001) = new{typeof(timepoint),typeof(shift)}(timepoint,shift)
end

"""
    KeyRateZero(timepoint,shift=0.001) <: KeyRateDuration

Shift the par curve by the given amount at the given timepoint. Use in conjunction with `duration` to calculate the key rate duration.

Unlike other duration statistics which are computed using analytic derivatives, `KeyRateDuration` is computed via a shift-and-compute the yield curve approach.

`KeyRateZero` is less commonly reported (than [`KayRatePar`](@ref)) in the fixed income markets, even though the latter has more analytically attractive properties. See the discussion of KeyRateDuration in the Yields.jl docs.
"""
struct KeyRateZero{T,R} <: KeyRateDuration
    timepoint::T
    shift::R
    KeyRateZero(timepoint, shift=.001) = new{typeof(timepoint),typeof(shift)}(timepoint,shift)
end

"""
    KeyRate(timepoints,shift=0.001)

A convenience constructor for [`KeyRateZero`](@ref).

## Extended Help
[`KeyRateZero`](@ref) is chosen as the default constructor because it has more attractive properties than [`KeyRatePar`](@ref):

- rates after the key `timepoint` remain unaffected by the `shift`
  - e.g. this causes a 6-year zero coupon bond would have a negative duration if the 5-year par rate was used

"""
KeyRate = KeyRateZero

"""
    duration(Macaulay(),interest_rate,cfs,times)
    duration(Modified(),interest_rate,cfs,times)
    duration(DV01(),interest_rate,cfs,times)
    duration(interest_rate,cfs,times)             # Modified Duration
    duration(interest_rate,valuation_function)    # Modified Duration

Calculates the Macaulay, Modified, or DV01 duration. `times` may be ommitted and the valuation will assume evenly spaced cashflows starting at the end of the first period.
- `interest_rate` should be a fixed effective yield (e.g. `0.05`).

When not given `Modified()` or `Macaulay()` as an argument, will default to `Modified()`.

# Examples

Using vectors of cashflows and times
```julia-repl
julia> times = 1:5
julia> cfs = [0,0,0,0,100]
julia> duration(0.03,cfs,times)
4.854368932038834
julia> duration(Macaulay(),0.03,cfs,times)
5.0
julia> duration(Modified(),0.03,cfs,times)
4.854368932038835
julia> convexity(0.03,cfs,times)
28.277877274012614

```

Using any given value function:

```julia-repl
julia> lump_sum_value(amount,years,i) = amount / (1 + i ) ^ years
julia> my_lump_sum_value(i) = lump_sum_value(100,5,i)
julia> duration(0.03,my_lump_sum_value)
4.854368932038835
julia> convexity(0.03,my_lump_sum_value)
28.277877274012617

```
"""
function duration(::Macaulay, yield, cfs, times)
    return sum(times .* price.(yield, vec(cfs), times) / price(yield, vec(cfs), times))
end

function duration(::Modified, yield, cfs, times)
    D(i) = price(i, vec(cfs), times)
    return duration(yield, D)
end

function duration(yield, valuation_function)
    D(i) = log(valuation_function(i + yield))
    δV =  - ForwardDiff.derivative(D, 0.0)
end

function duration(yield::Y, valuation_function) where {Y <: Yields.AbstractYield}
    D(i) = log(valuation_function(i + yield))
    δV =  - ForwardDiff.derivative(D, 0.0)
end

function duration(yield, cfs, times)
    return duration(Modified(), yield, vec(cfs), times)
end
function duration(yield::Y, cfs::A) where {Y <: Yields.AbstractYield,A <: AbstractArray}
    times = 1:length(cfs)
    return duration(Modified(), yield, vec(cfs), times)
end

function duration(yield::R, cfs) where {R <: Real}
    return duration(Yields.Constant(yield), cfs)
end

function duration(::DV01, yield, cfs, times)
    return duration(DV01(), yield, i -> price(i, vec(cfs), times))
end
function duration(d::Duration, yield, cfs)
    times = 1:length(cfs)
    return duration(d, yield, vec(cfs), times)
end

function duration(::DV01, yield, valuation_function)
    return duration(yield, valuation_function) * valuation_function(yield) / 100
end

"""
    convexity(yield,cfs,times)
    convexity(yield,valuation_function)

Calculates the convexity.
    - `yield` should be a fixed effective yield (e.g. `0.05`).
    - `times` may be omitted and it will assume `cfs` are evenly spaced beginning at the end of the first period.

# Examples

Using vectors of cashflows and times
```julia-repl
julia> times = 1:5
julia> cfs = [0,0,0,0,100]
julia> duration(0.03,cfs,times)
4.854368932038834
julia> duration(Macaulay(),0.03,cfs,times)
5.0
julia> duration(Modified(),0.03,cfs,times)
4.854368932038835
julia> convexity(0.03,cfs,times)
28.277877274012614

```

Using any given value function:

```julia-repl
julia> lump_sum_value(amount,years,i) = amount / (1 + i ) ^ years
julia> my_lump_sum_value(i) = lump_sum_value(100,5,i)
julia> duration(0.03,my_lump_sum_value)
4.854368932038835
julia> convexity(0.03,my_lump_sum_value)
28.277877274012617

```

"""
function convexity(yield, cfs, times)
    return convexity(yield, i -> price(i, vec(cfs), times))
end

function convexity(yield, cfs::A) where {A <: AbstractArray}
    times = 1:length(cfs)
    return convexity(yield, i -> price(i, vec(cfs), times))
end

function convexity(yield, valuation_function)
    v(x) = abs(valuation_function(yield + x[1]))
    ∂²P = ForwardDiff.hessian(v, [0.0])
    return ∂²P[1] / v([0.0])
end

"""
    duration(keyrate::KeyRateDuration,curve,cashflows)
    duration(keyrate::KeyRateDuration,curve,cashflows,timepoints)
    duration(keyrate::KeyRateDuration,curve,cashflows,timepoints,krd_points)

Calculate the key rate duration by shifting the **zero** (not par) curve by the kwarg `shift` at the timepoint specified by a KeyRateDuration(time).

The approach is to carve up the curve into `krd_points` (default is the unit steps between `1` and  the last timepoint of the casfhlows). The
zero rate corresponding to the timepoint within the `KeyRateDuration` is shifted by `shift` (specified by the `KeyRateZero` or `KeyRatePar` constructors. A new curve is created from the shifted rates. This means that the
"width" of the shifted section is ± 1 time period, unless specific points are specified via `krd_points`.

The `curve` may be any Yields.jl curve (e.g. does not have to be a curve constructed via `Yields.Zero(...)`).

!!! Experimental: Due to the paucity of examples in the literature, this feature does not have unit tests like the rest of JuliaActuary functionality. Additionally, the API may change in a future major/minor version update.

# Examples

```julia-repl
julia> riskfree_maturities = [0.5, 1.0, 1.5, 2.0];

julia> riskfree    = [0.05, 0.058, 0.064,0.068];

julia> rf_curve = Yields.Zero(riskfree,riskfree_maturities);

julia> cfs = [10,10,10,10,10];

julia> duration(KeyRate(1),rf_curve,cfs)
8.932800152336995

```

# Extended Help

Key Rate Duration is not a well specified topic in the literature and in practice. The reference below suggest that shocking the par curve is more common
in practice, but that the zero curve produces more consistent results. Future versions may support shifting the par curve.

References:
- [Quant Finance Stack Exchange: To compute key rate duration, shall I use par curve or zero curve?](https://quant.stackexchange.com/questions/33891/to-compute-key-rate-duration-shall-i-use-par-curve-or-zero-curve)
- (Financial Exam Help 123](http://www.financialexamhelp123.com/key-rate-duration/)

"""
function duration(keyrate::KeyRateDuration, curve, cashflows, timepoints, krd_points)
    shift = keyrate.shift
    curve_up = _krd_new_curve(keyrate,curve,krd_points)
    curve_down = _krd_new_curve(opposite(keyrate),curve,krd_points)
    price = pv(curve, cashflows, timepoints)
    price_up = pv(curve_up, cashflows, timepoints)
    price_down = pv(curve_down, cashflows, timepoints)

    return (price_down - price_up) / (2*shift*price)

end

opposite(kr::KeyRateZero) = KeyRateZero(kr.timepoint,-kr.shift)
opposite(kr::KeyRatePar) = KeyRatePar(kr.timepoint,-kr.shift)

function _krd_new_curve(keyrate::KeyRateZero,curve,krd_points)
    curve_times = krd_points
    shift = keyrate.shift

    zeros = Yields.zero.(curve, curve_times)

    zero_index = findfirst(==(keyrate.timepoint), curve_times)

    target_rate = zeros[zero_index]

    zeros[zero_index] += Yields.Rate(shift,target_rate.compounding)

    new_curve = Yields.Zero(zeros, curve_times)

    return new_curve
end

function _krd_new_curve(keyrate::KeyRatePar,curve,krd_points)
    curve_times = krd_points
    shift = keyrate.shift

    pars = Yields.par.(curve, curve_times)

    zero_index = findfirst(==(keyrate.timepoint), curve_times)

    target_rate = pars[zero_index]
    pars[zero_index] += Yields.Rate(shift,target_rate.compounding)

    new_curve = Yields.Par(pars, curve_times)

    return new_curve
end

function duration(keyrate::KeyRateDuration, curve, cashflows, timepoints)
    krd_points = 1:maximum(timepoints)
    return duration(keyrate, curve, cashflows, timepoints, krd_points)

end

function duration(keyrate::KeyRateDuration, curve, cashflows)
    timepoints = eachindex(cashflows)
    krd_points = 1:maximum(timepoints)
    return duration(keyrate, curve, cashflows, timepoints, krd_points)

end

"""
    moic(cashflows<:AbstractArray)

The multiple on invested capital ("moic") is the un-discounted sum of distributions divided by the sum of the contributions. The function assumes that negative numbers in the array represent contributions and positive numbers represent distributions.

# Examples

```julia-repl
julia> moic([-10,20,30])
5.0
```

"""
function moic(cfs::T) where {T<:AbstractArray}
    invested = zero(eltype(cfs))
    returned = zero(eltype(cfs))
    for i = 1:length(cfs)
        @inbounds cf = cfs[i]
        if cf > 0
            returned += cf
        else
            invested += -cf
        end
    end

    return returned / invested
end


function CellRef(n::AbstractString)
    @assert is_valid_cellname(n) "$n is not a valid CellRef."
    column_name, row_number = split_cellname(n)
    return CellRef(n, row_number, decode_column_number(column_name))
end

@inline CellRef(row::Int, col::Int) = CellRef(encode_column_number(col) * string(row))
@inline CellPosition(ref::CellRef) = CellPosition(row_number(ref), column_number(ref))
@inline row_number(p::CellPosition) = p.row
@inline column_number(p::CellPosition) = p.column
@inline CellRef(p::CellPosition) = CellRef(row_number(p), column_number(p))

"""
    decode_column_number(column_name::AbstractString) :: Int

Converts column name to a column number.

```julia
julia> XLSX.decode_column_number("D")
4
```

See also: `encode_column_number`.
"""
function decode_column_number(column_name::AbstractString) :: Int
    local result::Int = 0

    num_characters = length(column_name)

    iteration = 1
    for i in num_characters:-1:1
        column_char_as_int = Int(column_name[i])
        result += (26^(iteration-1)) * (column_char_as_int - 64) # From A to Z we have 26 values. 'A' Char is ASCII 65.
        iteration += 1
    end

    return result
end

"""
    encode_column_number(column_number::Int) :: String

Converts column number to a column name.

# Example

```julia
julia> XLSX.encode_column_number(4)
"D"
```

See also: `decode_column_number`.
"""
function encode_column_number(column_number::Int) :: String
    @assert column_number > 0 && column_number <= 16384 "Column number should be in the range from 1 to 16384."

    third_letter_sequence = div(column_number - 26 - 1, 26*26)
    column_number = column_number - third_letter_sequence*(26*26)

    second_letter_sequence = div(column_number - 1, 26) # 26^1
    column_number = column_number - second_letter_sequence*(26)

    first_letter_sequence = column_number # 26^0

    if third_letter_sequence > 0
        # result will have 3 letters
        return String([ Char(third_letter_sequence+64), Char(second_letter_sequence+64), Char(first_letter_sequence+64) ])

    elseif second_letter_sequence > 0
        # result will have 2 letters
        return String([ Char(second_letter_sequence+64), Char(first_letter_sequence+64) ])

    else
        # result will have 1 letter
        return String([ Char(first_letter_sequence+64) ])
    end
end

Base.string(c::CellRef) = c.name
Base.show(io::IO, c::CellRef) = print(io, string(c))
Base.:(==)(c1::CellRef, c2::CellRef) = c1.name == c2.name
Base.hash(c::CellRef) = hash(c.name)

const RGX_COLUMN_NAME = r"^[A-Z]?[A-Z]?[A-Z]$"
const RGX_CELLNAME = r"^[A-Z]+[0-9]+$"
const RGX_CELLRANGE = r"^[A-Z]+[0-9]+:[A-Z]+[0-9]+$"

function is_valid_column_name(n::AbstractString) :: Bool
    if !occursin(RGX_COLUMN_NAME, n)
        return false
    end

    column_number = decode_column_number(n)
    if column_number < 1 || column_number > 16384
        return false
    end

    return true
end

const RGX_CELLNAME_LEFT = r"^[A-Z]+"
const RGX_CELLNAME_RIGHT = r"[0-9]+$"

"""
    split_cellname(n::AbstractString) -> column_name, row_number

Splits a string representing a cell name to its column name and row number.

# Example

```julia
julia> XLSX.split_cellname("AB:12")
("AB:", 12)
```
"""
@inline function split_cellname(n::AbstractString)
    for (i, c) in enumerate(n)
        if isdigit(c)
            column_name = SubString(n, 1, i-1)
            row = parse(Int, SubString(n, i, length(n)))

            return column_name, row
        end
    end

    error("Couldn't split (column_name, row) for cellname $n.")
end

"""
    is_valid_cellname(n::AbstractString) :: Bool

Checks wether `n` is a valid name for a cell.

Cell names are bounded by `A1 : XFD1048576`.
"""
function is_valid_cellname(n::AbstractString) :: Bool

    if !occursin(RGX_CELLNAME, n)
        return false
    end

    column_name, row = split_cellname(n)

    if row < 1 || row > 1048576
        return false
    end

    if !is_valid_column_name(column_name)
        return false
    end

    return true
end

const RGX_CELLRANGE_START = r"^[A-Z]+[0-9]+"
const RGX_CELLRANGE_STOP = r"[A-Z]+[0-9]+$"

"""
    split_cellrange(n::AbstractString) -> start_name, stop_name

Splits a string representing a cell range into its cell names.

# Example

```julia
julia> XLSX.split_cellrange("AB12:CD24")
("AB12", "CD24")
```
"""
@inline function split_cellrange(n::AbstractString)
    s = split(n, ":")
    @assert length(s) == 2 "$n is not a valid cell range."
    return s[1], s[2]
end

function is_valid_cellrange(n::AbstractString) :: Bool

    if !occursin(RGX_CELLRANGE, n)
        return false
    end

    start_name, stop_name = split_cellrange(n)

    if !is_valid_cellname(start_name)
        return false
    end

    if !is_valid_cellname(stop_name)
        return false
    end

    return true
end

macro ref_str(ref)
    CellRef(ref)
end

function CellRange(r::AbstractString)
    @assert occursin(RGX_CELLRANGE, r) "Invalid cell range: $r."
    start_name, stop_name = split_cellrange(r)
    return CellRange(CellRef(start_name), CellRef(stop_name))
end

CellRange(start_row::Integer, start_column::Integer, stop_row::Integer, stop_column::Integer) = CellRange(CellRef(start_row, start_column), CellRef(stop_row, stop_column))

Base.string(cr::CellRange) = "$(string(cr.start)):$(string(cr.stop))"
Base.show(io::IO, cr::CellRange) = print(io, string(cr))
Base.:(==)(cr1::CellRange, cr2::CellRange) = cr1.start == cr2.start && cr2.stop == cr2.stop
Base.hash(cr::CellRange) = hash(cr.start) + hash(cr.stop)

macro range_str(cellrange)
    CellRange(cellrange)
end

"""
    Base.in(ref::CellRef, rng::CellRange) :: Bool

Checks wether `ref` is a cell reference inside a range given by `rng`.
"""
function Base.in(ref::CellRef, rng::CellRange) :: Bool
    top = row_number(rng.start)
    bottom = row_number(rng.stop)
    r = row_number(ref)

    if top <= r && r <= bottom
        left = column_number(rng.start)
        right = column_number(rng.stop)
        c = column_number(ref)

        if left <= c && c <= right
            return true
        end
    end

    return false
end

"""
    Base.issubset(subrng::CellRange, rng::CellRange)

Checks wether `subrng` is a cell range contained in `rng`.
"""
Base.issubset(subrng::CellRange, rng::CellRange) :: Bool = in(subrng.start, rng) && in(subrng.stop, rng)

function Base.size(rng::CellRange)
    top = row_number(rng.start)
    bottom = row_number(rng.stop)
    left = column_number(rng.start)
    right = column_number(rng.stop)

    return ( bottom - top + 1, right - left + 1 )
end

"""
    row_number(c::CellRef) :: Int

Returns the row number of a given cell reference.
"""
row_number(c::CellRef) :: Int = c.row_number

"""
    column_number(c::CellRef) :: Int

Returns the column number of a given cell reference.
"""
column_number(c::CellRef) :: Int = c.column_number

column_name(c::CellRef) :: String = encode_column_number(column_number(c))

"""
Returns (row, column) representing a `ref` position relative to `rng`.

For example, for a range "B2:D4", we have:

* "C3" relative position is (2, 2)

* "B2" relative position is (1, 1)

* "C4" relative position is (3, 2)

* "D4" relative position is (3, 3)

"""
function relative_cell_position(ref::CellRef, rng::CellRange)
    @assert ref ∈ rng "$ref is outside range $rng."

    top = row_number(rng.start)
    left = column_number(rng.start)

    r, c = row_number(ref), column_number(ref)

    return ( r - top + 1 , c - left + 1 )
end

#
# ColumnRange
#

Base.string(cr::ColumnRange) = "$(encode_column_number(cr.start)):$(encode_column_number(cr.stop))"
Base.show(io::IO, cr::ColumnRange) = print(io, string(cr))
Base.:(==)(cr1::ColumnRange, cr2::ColumnRange) = cr1.start == cr2.start && cr2.stop == cr2.stop
Base.hash(cr::ColumnRange) = hash(cr.start) + hash(cr.stop)
Base.in(column_number::Integer, rng::ColumnRange) = rng.start <= column_number && column_number <= rng.stop

function relative_column_position(column_number::Integer, rng::ColumnRange)
    @assert column_number ∈ rng "Column $column_number is outside range $rng."
    return column_number - rng.start + 1
end

@inline relative_column_position(ref::CellRef, rng::ColumnRange) = relative_column_position(column_number(ref), rng)

const RGX_COLUMN_RANGE = r"^[A-Z]?[A-Z]?[A-Z]:[A-Z]?[A-Z]?[A-Z]$"
const RGX_COLUMN_RANGE_START = r"^[A-Z]+"
const RGX_COLUMN_RANGE_STOP = r"[A-Z]+$"

"""
Returns tuple (column_name_start, column_name_stop).
"""
@inline function split_column_range(n::AbstractString)
    s = split(n, ":")
    return s[1], s[2]
end

function is_valid_column_range(r::AbstractString) :: Bool
    if !occursin(RGX_COLUMN_RANGE, r)
        return false
    end

    start_name, stop_name = split_column_range(r)

    if !is_valid_column_name(start_name) || !is_valid_column_name(stop_name)
        return false
    end

    return true
end

function ColumnRange(r::AbstractString)
    @assert is_valid_column_range(r) "Invalid column range: $r."
    start_name, stop_name = split_column_range(r)
    return ColumnRange(decode_column_number(start_name), decode_column_number(stop_name))
end

convert(::Type{ColumnRange}, str::AbstractString) = ColumnRange(str)
convert(::Type{ColumnRange}, column_range::ColumnRange) = column_range

column_bounds(r::ColumnRange) = (r.start, r.stop)
Base.length(r::ColumnRange) = r.stop - r.start + 1

<<<<<<< HEAD
# ColumnRange iterator
start(itr::ColumnRange) = itr.start
done(itr::ColumnRange, column_index::Int) = column_index > itr.stop
next(itr::ColumnRange, column_index::Int) = (encode_column_number(column_index), column_index + 1)

# CellRange iterator
start(rng::CellRange) = CellRefIteratorState(row_number(rng.start), column_number(rng.start))
done(rng::CellRange, state::CellRefIteratorState) = state.row > row_number(rng.stop)
=======
# ColumnRange iterator: element is a String with the column name, the state is the column number.
function Base.iterate(itr::ColumnRange, state::Int=itr.start)
    if state > itr.stop
        return nothing
    end
>>>>>>> upstream/master

    return encode_column_number(state), state + 1
end

<<<<<<< HEAD
# (i, state) = next(I, state)
function next(rng::CellRange, state::CellRefIteratorState)
    local next_state::CellRefIteratorState
    if state.col == column_number(rng.stop)
=======
# CellRange iterator: element is a CellRef, the state is a CellPosition.
function Base.iterate(rng::CellRange, state::CellPosition=CellPosition(rng.start))

    if row_number(state) > row_number(rng.stop)
        return nothing
    elseif column_number(state) == column_number(rng.stop)
>>>>>>> upstream/master
        # reached last column. Go to the next row.
        next_state = CellPosition(row_number(state) + 1, column_number(rng.start))
    else
        # go to the next column
        next_state = CellPosition(row_number(state), column_number(state) + 1)
    end

    return CellRef(state), next_state
end

function Base.length(rng::CellRange)
    (r, c) = size(rng)
    return r * c
end

#
# SheetCellRef, SheetCellRange, SheetColumnRange
#

Base.string(cr::SheetCellRef) = string(cr.sheet, "!", cr.cellref)
Base.show(io::IO, cr::SheetCellRef) = print(io, string(cr))
Base.:(==)(cr1::SheetCellRef, cr2::SheetCellRef) = cr1.sheet == cr2.sheet && cr2.cellref == cr2.cellref
Base.hash(cr::SheetCellRef) = hash(cr.sheet) + hash(cr.cellref)

Base.string(cr::SheetCellRange) = string(cr.sheet, "!", cr.rng)
Base.show(io::IO, cr::SheetCellRange) = print(io, string(cr))
Base.:(==)(cr1::SheetCellRange, cr2::SheetCellRange) = cr1.sheet == cr2.sheet && cr2.rng == cr2.rng
Base.hash(cr::SheetCellRange) = hash(cr.sheet) + hash(cr.rng)

Base.string(cr::SheetColumnRange) = string(cr.sheet, "!", cr.colrng)
Base.show(io::IO, cr::SheetColumnRange) = print(io, string(cr))
Base.:(==)(cr1::SheetColumnRange, cr2::SheetColumnRange) = cr1.sheet == cr2.sheet && cr2.colrng == cr2.colrng
Base.hash(cr::SheetColumnRange) = hash(cr.sheet) + hash(cr.colrng)

const RGX_SHEET_CELLNAME = r"^.+![A-Z]+[0-9]+$"
const RGX_SHEET_CELLRANGE = r"^.+![A-Z]+[0-9]+:[A-Z]+[0-9]+$"
const RGX_SHEET_COLUMN_RANGE = r"^.+![A-Z]?[A-Z]?[A-Z]:[A-Z]?[A-Z]?[A-Z]$"

const RGX_SHEET_CELLNAME_RIGHT = r"[A-Z]+[0-9]+$"
const RGX_SHEET_CELLRANGE_RIGHT = r"[A-Z]+[0-9]+:[A-Z]+[0-9]+$"
const RGX_SHEET_COLUMN_RANGE_RIGHT = r"[A-Z]?[A-Z]?[A-Z]:[A-Z]?[A-Z]?[A-Z]$"

function is_valid_sheet_cellname(n::AbstractString) :: Bool
    if !occursin(RGX_SHEET_CELLNAME, n)
        return false
    end

    cellname = match(RGX_SHEET_CELLNAME_RIGHT, n).match
    if !is_valid_cellname(cellname)
        return false
    end

    return true
end

function is_valid_sheet_cellrange(n::AbstractString) :: Bool
    if !occursin(RGX_SHEET_CELLRANGE, n)
        return false
    end

    cellrange = match(RGX_SHEET_CELLRANGE_RIGHT, n).match
    if !is_valid_cellrange(cellrange)
        return false
    end

    return true
end

function is_valid_sheet_column_range(n::AbstractString) :: Bool
    if !occursin(RGX_SHEET_COLUMN_RANGE, n)
        return false
    end

    column_range = match(RGX_SHEET_COLUMN_RANGE_RIGHT, n).match
    if !is_valid_column_range(column_range)
        return false
    end

    return true
end

const RGX_CELLNAME_RIGHT_FIXED = r"\$[A-Z]+\$[0-9]+$"
const RGX_SHEET_CELNAME_RIGHT_FIXED = r"\$[A-Z]+\$[0-9]+:\$[A-Z]+\$[0-9]+$"

function SheetCellRef(n::AbstractString)
    if is_valid_fixed_sheet_cellname(n)
        fixed_cellname = match(RGX_CELLNAME_RIGHT_FIXED, n).match
        cellname = replace(fixed_cellname, "\$" => "")
        sheetname = SubString(n, 1, length(n) - length(fixed_cellname) - 1)
        return SheetCellRef(sheetname, CellRef(cellname))
    else
        @assert is_valid_sheet_cellname(n) "$n is not a valid SheetCellRef."
        cellname = match(RGX_SHEET_CELLNAME_RIGHT, n).match
        sheetname = SubString(n, 1, length(n) - length(cellname) - 1)
        return SheetCellRef(sheetname, CellRef(cellname))
    end
end

function SheetCellRange(n::AbstractString)
    if is_valid_fixed_sheet_cellrange(n)
        fixed_cellrange = match(RGX_SHEET_CELNAME_RIGHT_FIXED, n).match
        cellrange = replace(fixed_cellrange, "\$" => "")
        sheetname = SubString(n, 1, length(n) - length(fixed_cellrange) - 1)
        return SheetCellRange(sheetname, CellRange(cellrange))
    else
        @assert is_valid_sheet_cellrange(n) "$n is not a valid SheetCellRange."
        cellrange = match(RGX_SHEET_CELLRANGE_RIGHT, n).match
        sheetname = SubString(n, 1, length(n) - length(cellrange) - 1)
        return SheetCellRange(sheetname, CellRange(cellrange))
    end
end

function SheetColumnRange(n::AbstractString)
    @assert is_valid_sheet_column_range(n) "$n is not a valid SheetColumnRange."
    column_range = match(RGX_SHEET_COLUMN_RANGE_RIGHT, n).match
    sheetname = SubString(n, 1, length(n) - length(column_range) - 1)
    return SheetColumnRange(sheetname, ColumnRange(column_range))
end

# Named ranges
const RGX_FIXED_SHEET_CELLNAME = r"^.+!\$[A-Z]+\$[0-9]+$"
const RGX_FIXED_SHEET_CELLRANGE = r"^.+!\$[A-Z]+\$[0-9]+:\$[A-Z]+\$[0-9]+$"

is_valid_fixed_sheet_cellname(s::AbstractString) = occursin(RGX_FIXED_SHEET_CELLNAME, s)
is_valid_fixed_sheet_cellrange(s::AbstractString) = occursin(RGX_FIXED_SHEET_CELLRANGE, s)


## Whittle.jl : implementation of Whittle covariance function

## Whittle ##
struct Whittle{T} <: IsotropicCovarianceStructure{T}
    λ::T
    σ::T
    p::T

    function Whittle{T}(λ::T, σ::T, p::T) where T
        λ > 0 || throw(DomainError(λ, "correlation length λ of Whittle covariance cannot be negative or zero"))
        σ > 0 || throw(DomainError(σ, "marginal standard deviation σ of Whittle covariance cannot be negative or zero"))
        p >= 1 || throw(DomainError(p, "in p-norm, p must be greater than or equal to 1"))
        isinf(p) && throw(DomainError(p, "in p-norm, p cannot be infinity"))

        new{T}(λ, σ, p)
    end
end

"""
    Whittle(λ, [σ = 1], [p = 2])

Whittle covariance structure with correlation length `λ`, (optional) marginal standard deviation `σ` and (optional) `p`-norm, defined as

``C(x, y) = σ \\displaystyle\\frac{ρ}{λ} K₁\\left(\\frac{ρ}{λ}\\right)``

with ``ρ = ||x-y||_p``.

# Examples
```jldoctest
julia> Whittle(0.1)
Whittle (λ=0.1, σ=1.0, p=2.0)

julia> Whittle(1.0, σ=2)
Whittle (λ=1.0, σ=2.0, p=2.0)

```
See also: [`Exponential`](@ref), [`Linear`](@ref), [`Spherical`](@ref), [`Gaussian`](@ref), [`SquaredExponential`](@ref), [`Matern`](@ref)
"""
Whittle(λ::Real; σ::Real=1.0, p::Real=2) = Whittle{promote_type(typeof(λ),typeof(σ),typeof(p))}(promote(λ, σ, p)...)

# evaluate Whittle covariance
function apply(w::Whittle, x::Real)
    if iszero(x)
        float(one(x))
    else
        ρ = x / w.λ
        ρ * besselk(1, ρ)
    end
end

# short name
shortname(::Whittle) = "Whittle"


# this illustrates effect and detection of influential observations
using Plots, LinearAlgebra
n = 20
x = (1:n-1)/(n-1)
x = [x; 3] # the last observation is an outlying value of x

x = [ones(n,1) x]
P = x*inv(x'*x)*x'

beta = [10, -1]
e = 2*randn(n,1)
y = x*beta + e

# The fit
yhat = P*y

# calculate leverage and influence
leverage = diag(P)
e = y - yhat
influence = (leverage ./ (1.0 .-leverage)) .* e

xlabel = "X"
x = x[:,2]
scatter(x, y, label = "Data points")
plot!(x, yhat, label = "fitted")
plot!(x, leverage, label = "Leverage")
plot!(x, influence, label = "Influence")
gui()
#savefig("InfluentialObservation.svg")


struct SimpleCurve
    start::Tuple{Number, Number}
    stop::Tuple{Number, Number}
    control::Tuple{Number, Number}

    SimpleCurve(start::Tuple{Number, Number}, stop::Tuple{Number, Number}, control::Tuple{Number, Number}=(start .+ stop) ./ 2) = new(start, stop, control)
end

function getPoint(curve::SimpleCurve, percent::Number)
    x = (1 - percent)^2 * curve.start[1] + 2 * (1 - percent) * percent * curve.control[1] + percent^2 * curve.stop[1]
    y = (1 - percent)^2 * curve.start[2] + 2 * (1 - percent) * percent * curve.control[2] + percent^2 * curve.stop[2]

    return x, y
end


###
#   Simple implmentation of the Metropolis Hasting algorithm in julia
#
#            Institute for Risk and Uncertainty, Uni of Liverpool
#
#                       Authors: Ander Gray, Adolphus Lye
#
#                       Email: Ander.Gray@liverpool.ac.uk,
#                              Adolphus.Lye@liverpool.ac.uk
#
#
#
#       W. K. Hastings (1970). Monte Carlo sampling methods using
#       Markov chains and their applications. Biometrika,
#       57(1), 97-109. doi:10.1093/biomet/57.1.97
#
###

###
#   Add description of function and inputs
###

function MHsample(Target :: Function, Prop, start :: Vector{<:Real}, Nsamples :: Integer, burnin :: Integer = 50, thin ::Integer = 3; islogged :: Bool = true)

    dims = length(start)                                    # Dimensions of input/ prior

    PropRnd = mu -> rand(Prop(mu))                          # Generates a sample from the proposal given mean mu

    chain = zeros( Nsamples * thin + burnin, dims)
    chain[1,:] = start
    accRate = 0

    islogged ? evalDen =  x -> exp(Target(x)) : evalDen = x -> Target(x)

    for i = 2:( Nsamples*thin +burnin)

        next = PropRnd(chain[i-1,:])              # Draw candidate

        targDen = evalDen(next)                  # Target Density at next sample
        targPrevious = evalDen(chain[i-1,:])     # Target Density at current sample

        propDen = pdf(Prop(chain[i-1,:]), next)         # Proposal at next centred at current
        propPrevious = pdf(Prop(next), chain[i-1,:])    # Propsoal at current centred at next

        α = targDen/targPrevious * propPrevious/propDen     # General formula for acceptance probability

        accepted = α >= rand()

        if accepted
            chain[i,:] = next
            accRate = accRate +1
        else
            chain[i,:] = chain[i-1, :]
        end

    end
    accRate = accRate/(Nsamples*thin+burnin)
    return chain[burnin+1:thin:end,:], accRate
end

##
#   Symetric proposal and logged target
##
function MHsampleSimple(Target :: Function, PropRnd, start :: Vector{<:Real}, Nsamples :: Integer, burnin :: Integer = 50, thin ::Integer = 3)

    dims = length(start)                                    # Dimensions of input/ prior
    chain = zeros( Nsamples * thin + burnin, dims)
    chain[1,:] = start[:]
    accRate = 0

    for i = 2:( Nsamples*thin+burnin)

        next = PropRnd(chain[i-1,:])                 # Draw candidate

        targDen = Target(next)[1]                    # Target Density at next sample
        targPrevious = Target(chain[i-1,:])[1]       # Target Density at current sample

        α =  min(0, targDen - targPrevious)

        accepted = α >= log(rand())

        if accepted
            chain[i,:] = next
            accRate = accRate +1
        else
            chain[i,:] = chain[i-1, :]
        end

    end
    accRate = accRate/(Nsamples*thin+burnin)
    return chain[burnin+1:thin:end,:], accRate
end


export Beta

import Distributions: Beta, params
import SpecialFunctions: digamma

vague(::Type{ <: Beta }) = Beta(1.0, 1.0)

prod_analytical_rule(::Type{ <: Beta }, ::Type{ <: Beta }) = ProdAnalyticalRuleAvailable()

function prod(::ProdAnalytical, left::Beta, right::Beta)
    left_a, left_b   = params(left)
    right_a, right_b = params(right)
    T                = promote_type(eltype(left), eltype(right))
    return Beta(left_a + right_a - one(T), left_b + right_b - one(T))
end

function logmean(dist::Beta)
    a, b = params(dist)
    return digamma(a) - digamma(a + b)
end

function mirroredlogmean(dist::Beta)
    a, b = params(dist)
    return digamma(b) - digamma(a + b)
end


export getAnalyticalConstGrad2D,getAnalyticalConstGrad3D,getAnalyticalConstGradInv2D,getAnalyticalConstGradInv3D,getSmoothGaussianMedium,getSmoothFactoredModel,getSmoothFactoredModel3D

function getAnalyticalConstGrad2D(n::Array{Int64,1},h::Array{Float64,1})
src = [1,div(n[2],2)];

source1 = (src[1]-1)*h[1];
source2 = (src[2]-1)*h[2];

(X1,X2) = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2]);

a  = -0.4;
s0 = 2.0;

kappaSquared = s0^2 .+ 2.0*a*(X1.-source1);

# matshow(kappaSquared); colorbar();
SBarSquared = s0^2 .+ a*(X1.-source1);
rSquared = ((X1.-source1).^2 .+ (X2.-source2).^2);
sigmaSquarred = 2*rSquared./(SBarSquared + sqrt.(SBarSquared.^2 .- (a^2).*rSquared));
sigma = sqrt.(sigmaSquarred);
T_exact =  SBarSquared.*sigma .- (a^2).*(sigma.^3)./6;

return kappaSquared,src,T_exact;
end

function getAnalyticalConstGrad3D(n::Array{Int64,1},h::Array{Float64,1})
src = [div(n[1],2),div(n[2],2),1];

source1 = (src[1]-1)*h[1];
source2 = (src[2]-1)*h[2];
source3 = (src[3]-1)*h[3];

(X1,X2,X3) = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2],(0:(n[3]-1))*h[3]);

# a  = -1.75;
a  = -1.65;
s0 = 2.0;

kappaSquared = s0^2 .+ 2*a*(X3.-source3);
SBarSquared = s0^2 .+ a*(X3.-source3);
rSquared = (X1.-source1).^2 .+ (X2.-source2).^2 .+ (X3.-source3).^2;
sigmaSquarred = 2*rSquared./(SBarSquared .+ sqrt.(SBarSquared.^2 .- (a^2).*rSquared));
sigma = sqrt.(sigmaSquarred);
T_exact =  SBarSquared.*sigma .- (a^2).*(sigma.^3)./6;

return kappaSquared,src,T_exact;
end

function getAnalyticalConstGradInv2D(n::Array{Int64,1},h::Array{Float64,1})

src = [1,div(n[2],2)];

source1 = (src[1]-1)*h[1];
source2 = (src[2]-1)*h[2];

(X1,X2) = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2]);

a  = 1.0;
ainv = (1.0./a);
s0 = 2.0;

Kappa = 1.0./(1.0./s0 .+ a*(X1.-source1));
T_exact = ainv*acosh.(1.0.+(0.5*s0*a*a).*Kappa.*((X1.-source1).^2 .+ (X2.-source2).^2));
G2_exact = ainv*(1.0./sqrt.((1.0.+(0.5*s0*a*a).*Kappa.*((X1.-source1).^2 .+ (X2 .- source2).^2)).^2 .- 1)).*(0.5*s0*a*a).*Kappa.*2.0.*(X2.-source2);
G1_exact = ainv*(1.0./sqrt.((1.0.+(0.5*s0*a*a).*Kappa.*((X1.-source1).^2 .+ (X2 .- source2).^2)).^2 .- 1)).*((0.5*s0*a*a).*Kappa.*2.0.*(X1.-source1)
			+ (0.5*s0*a*a).*((X1.-source1).^2 + (X2.-source2).^2).*(-(Kappa.^2)).*a);
kappaSquared = Kappa.^2;

return kappaSquared,src,T_exact;

end

function getAnalyticalConstGradInv3D(n::Array{Int64,1},h::Array{Float64,1})

src = [div(n[1],2),div(n[1],2),1];

source1 = (src[1]-1)*h[1];
source2 = (src[2]-1)*h[2];
source3 = (src[3]-1)*h[3];

(X1,X2,X3) = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2],(0:(n[3]-1))*h[3]);

a  = 1.0;
ainv = (1.0./a);
s0 = 2.0;
Kappa = 1.0./(1.0./s0 .+ a*(X3.-source3));
Rsquared = (X1.-source1).^2 .+ (X2.-source2).^2 .+ (X3.-source3).^2;
T_exact = ainv*acosh.(1.0.+(0.5*s0*a*a).*Kappa.*(Rsquared));
G1_exact = ainv*(1.0./sqrt.((1.0.+(0.5*s0*a*a).*Kappa.*Rsquared).^2 .- 1.0)).*(0.5*s0*a*a).*Kappa.*2.0.*(X1.-source1);
G2_exact = ainv*(1.0./sqrt.((1.0.+(0.5*s0*a*a).*Kappa.*Rsquared).^2 .- 1.0)).*(0.5*s0*a*a).*Kappa.*2.0.*(X2.-source2);
G3_exact = ainv*(1.0./sqrt.((1.0.+(0.5*s0*a*a).*Kappa.*Rsquared).^2 .- 1.0)).*((0.5*s0*a*a).*Kappa.*2.0.*(X3.-source3)
			.+ (0.5*s0*a*a).*(Rsquared).*(-(Kappa.^2)).*a);
kappaSquared = Kappa.^2;

return kappaSquared,src,T_exact;

end

function getSmoothGaussianMedium(n::Array{Int64,1},h::Array{Float64,1})
src = div.(n,4);
kappaSquared = [];
T_exact = [];
if length(n)==2
	(T1_exact,G11_exact,G12_exact) = getSmoothFactoredModel(n,h);
	(T0,G01,G02,L0) = getAnalytic2DeikonalSolutionAll(n,h,src);
	G1_exact = T0.*G11_exact + G01.*T1_exact;
	G2_exact = T0.*G12_exact + G02.*T1_exact;
	kappaSquared = G1_exact.*G1_exact + G2_exact.*G2_exact;
	# matshow(sqrt(kappaSquared)); colorbar();
	# xlabel("y");
	# ylabel("x");
	# xticks(0:div(n[2],8):n[2],0:1:8);
	# yticks(0:div(n[1],4):n[1],0:1:4);

	#figure()
	#CS = contour(kappaSquared,50);colorbar();title("kappaSquared");
	#clabel(CS, inline=1, fontsize=10)
	#figure()
	#contour(G1_exact);colorbar();title("G1_exact");
	T_exact = T0.*T1_exact;
	#figure()
	#contour(T_exact,100);colorbar();title("T exact");

else

	(T1_exact,G11_exact,G12_exact,G13_exact) = getSmoothFactoredModel3D(n,h);

	(T0,G01,G02,G03) = getAnalytic3DeikonalSolutionAll(n,h,src);

	G1_exact = T0.*G11_exact + G01.*T1_exact;
	G2_exact = T0.*G12_exact + G02.*T1_exact;
	G3_exact = T0.*G13_exact + G03.*T1_exact;
	kappaSquared = G1_exact.*G1_exact + G2_exact.*G2_exact + G3_exact.*G3_exact;

	# Er = kappaSquared;
	# for k=1:n[3]
		# matshow(reshape(Er[:,:,k],n[1],n[2]));colorbar();title("kappaSQ");
	# end

	T_exact = T0.*T1_exact;

end

return kappaSquared,src,T_exact;
end

function getSmoothFactoredModel(n::Array{Int64,1},h::Array{Float64,1})
xsrc = zeros(2);
src_kappa = zeros(Int64,2);
src_kappa[1] = div(n[1],3);
src_kappa[2] = div(n[2],4);

xsrc[1] = (src_kappa[1]-1)*h[1];
xsrc[2] = (src_kappa[2]-1)*h[2];

X1,X2 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2]);

sigma = 0.1;
T1_exact = (exp.( - (sigma*((X1 .- xsrc[1]).^2) + 4*sigma*((X2.-xsrc[2]).^2))) .+ 1)/2;
G11_exact = -2*sigma*(X1 .- xsrc[1]).*exp.( - (sigma*((X1 .- xsrc[1]).^2) + 4*sigma*((X2.-xsrc[2]).^2)))/2;
G12_exact = -8*sigma*(X2 .- xsrc[2]).*exp.( - (sigma*((X1 .- xsrc[1]).^2) + 4*sigma*((X2.-xsrc[2]).^2)))/2;

return T1_exact,G11_exact,G12_exact;
end

function getSmoothFactoredModel3D(n::Array{Int64,1},h::Array{Float64,1})
xsrc = zeros(3);
src_kappa = zeros(Int64,3);
src_kappa[1] = div(n[1],3);
src_kappa[2] = div(n[2],4);
src_kappa[3] = div(n[3],2);

xsrc[1] = (src_kappa[1]-1)*h[1];
xsrc[2] = (src_kappa[2]-1)*h[2];
xsrc[3] = (src_kappa[3]-1)*h[3];

X1,X2,X3 = ndgrid((0:(n[1]-1))*h[1],(0:(n[2]-1))*h[2],(0:(n[3]-1))*h[3]);

Sigma = [0.1,0.4,0.2];
EXPRSkewedSquared = exp.(-Sigma[1]*((X1 .- xsrc[1]).^2) .- Sigma[2]*((X2.-xsrc[2]).^2) .- Sigma[3]*((X3.-xsrc[3]).^2));

T1_exact = EXPRSkewedSquared/2.0 .+ 0.5;
G11_exact = -Sigma[1]*(X1 .- xsrc[1]).*EXPRSkewedSquared;
G12_exact = -Sigma[2]*(X2 .- xsrc[2]).*EXPRSkewedSquared;
G13_exact = -Sigma[3]*(X3 .- xsrc[3]).*EXPRSkewedSquared;

return T1_exact,G11_exact,G12_exact,G13_exact;
end


# Miscellaneous functions

"""
    logmvgamma(p::Int, a::Real)

Return the logarithm of [multivariate gamma function](https://en.wikipedia.org/wiki/Multivariate_gamma_function) ([DLMF 35.3.1](https://dlmf.nist.gov/35.3.1)).
"""
function logmvgamma(p::Int, a::Real)
    # NOTE: one(a) factors are here to prevent unnecessary promotion of Float32
    res = p * (p - 1) * log(pi * one(a)) / 4
    for ii in 1:p
        res += lgamma(a + (1 - ii) * one(a)/ 2)
    end
    return res
end

"""
    lstirling_asym(x)

The remainder term after
[Stirling's approximation](https://en.wikipedia.org/wiki/Stirling%27s_approximation)
to [`lgamma`](@ref):

```math
\\log \\Gamma(x) \\approx x \\log(x) - x + \\log(2\\pi/x)/2 = \\log(x)*(x-1/2) + \\log(2\\pi)/2 - x
```

In Julia syntax, this means:

    lstirling_asym(x) = lgamma(x) + x - (x-0.5)*log(x) - 0.5*log(2π)

For sufficiently large `x`, this can be approximated using the asymptotic
_Stirling's series_ ([DLMF 5.11.1](https://dlmf.nist.gov/5.11.1)):

```math
\\frac{1}{12x} - \\frac{1}{360x^3} + \\frac{1}{1260x^5} - \\frac{1}{1680x^7} + \\ldots
```

The truncation error is bounded by the first omitted term, and is of the same sign.

Relative error of approximation is bounded by
    (174611/125400 x^-19) / (1/12 x^-1 - 1/360 x^-3)
which is < 1/2 ulp for x >= 10.0, and total numeric error appears to be < 2 ulps

# References

* Temme, N. (1996) Special functions: An introduction to the classical functions of
   mathematical physics, Wiley, New York, ISBN: 0-471-11313-1, Chapter 3.6, pp 61-65.
* Weisstein, Eric W. ["Stirling's Series."](http://mathworld.wolfram.com/StirlingsSeries.html).
  MathWorld.
* [OEIS A046968](http://oeis.org/A046968) and [OEIS A046969](http://oeis.org/A046969)
  for the series coefficients
"""
function lstirling_asym end

lstirling_asym(x::BigFloat) = lgamma(x) + x - log(x)*(x - big(0.5)) - log2π/big(2)

lstirling_asym(x::Integer) = lstirling_asym(float(x))

const lstirlingF64 = Float64[lstirling_asym(k) for k in big(1):big(64)]
const lstirlingF32 = Float64[lstirling_asym(k) for k in big(1):big(40)]

function lstirling_asym(x::Float64)
    isinteger(x) && (0 < x ≤ length(lstirlingF64)) && return lstirlingF64[Int(x)]
    t = inv(abs2(x))
    @horner(t,
             8.33333333333333333e-2, #  1/12 x^-1
            -2.77777777777777778e-3, # -1/360 x^-3
             7.93650793650793651e-4, #  1/1260 x^-5
            -5.95238095238095238e-4, # -1/1680 x^-7
             8.41750841750841751e-4, #  1/1188 x^-9
            -1.91752691752691753e-3, # -691/360360 x^-11
             6.41025641025641026e-3, #  1/156 x^-13
            -2.95506535947712418e-2, # -3617/122400 x^-15
             1.79644372368830573e-1)/x #  43867/244188 x^-17
end

function lstirling_asym(x::Float32)
    isinteger(x) && (0 < x ≤ length(lstirlingF32)) && return lstirlingF32[Int(x)]
    t = inv(abs2(x))
    @horner(t,
             8.333333333333f-2, #  1/12 x^-1
            -2.777777777777f-3, # -1/360 x^-3
             7.936507936508f-4, #  1/1260 x^-5
            -5.952380952381f-4, # -1/1680 x^-7
             8.417508417508f-4)/x #  1/1188 x^-9
end


__precompile__()

"""
Multi-dimensional integration routines: Monte-Carlo, stratified Monte-Carlo,
multi-armed bandit UCB.

See the
[README file](https://github.com/florianLepretre/ucbature/blob/master/README.md)
for documentation.
"""

module ucbature

export mc, str, ucb

rand_min_max(xmin, xmax) = xmin .+ rand(length(xmin)).*(xmax.-xmin)

volume(xmin, xmax) = prod(xmax .- xmin)

"""
    mc(f, xmin, xmax, nb_evals)

Compute the n-dimensional integral f(x), where
`n == length(xmin) == length(xmax)`, over the hypercube whose corners are given
by the vectors (or tuples) `xmin` and `xmax`. `f` should be a function `f(x)`
that takes an n-dimensional vector `x` and returns the integrand at `x`.

Use the classic Monte-Carlo integration method using `nb_evals` samples.
"""
function mc(integrand, xmin, xmax, nb_evals)
    f_val(i) = integrand(rand_min_max(xmin, xmax))
    mean_val = mapreduce(f_val, +, 1:nb_evals) / nb_evals
    volume(xmin, xmax) * mean_val
end

function make_x_cube(i_zone, nb_zones_per_dim, dim)
    # compute position i_zone in [0, nb_zones_per_dim]^dim
    x_cube = Array{Int}(dim)
    q = nb_zones_per_dim^(dim-1)
    for d = 1:dim
        x_cube[d] = div(i_zone, q)
        i_zone -= x_cube[d] * q
        q = div(q, nb_zones_per_dim)
    end
    x_cube
end

function make_zones(xmin, xmax, nb_zones_per_dim)
    dim = length(xmin)
    nb_zones = nb_zones_per_dim ^ dim
    xstep = (xmax .- xmin) / nb_zones_per_dim
    zones = [xmin .+ xstep.*make_x_cube(i, nb_zones_per_dim, dim)
             for i = 0:(nb_zones-1)]
    dim, nb_zones, xstep, zones
end

"""
    str(f, xmin, xmax, nb_evals, nb_zones_per_dim=3)

Compute the n-dimensional integral f(x), where
`n == length(xmin) == length(xmax)`, over the hypercube whose corners are given
by the vectors (or tuples) `xmin` and `xmax`. `f` should be a function `f(x)`
that takes an n-dimensional vector `x` and returns the integrand at `x`.

Use the stratified monte-carlo integration method using `nb_evals` samples and
`nb_zones_per_dim` uniform strata in each dimension.
"""
function str(integrand, xmin, xmax, nb_evals, nb_zones_per_dim=3)
    dim, nb_zones, xstep, zones = make_zones(xmin, xmax, nb_zones_per_dim)
    if nb_evals < nb_zones
        warn("Integrator.str, nb_evals < $(nb_zones)")
        return mc(integrand, xmin, xmax, nb_evals)
    end
    nb_evals_per_zone = div(nb_evals, nb_zones)
    f_mc(z) = mc(integrand, z, z.+xstep, nb_evals_per_zone)
    mapreduce(f_mc, +, zones)
end

ucb1(R, k, kn, Vn) = Vn ./ sqrt.(kn) .+ R .* sqrt.(log(k) ./ kn)

ucb_vars(s1, s2, k) = (s2./k) .- (s1./k).^2

function ucb_strat(integrand, xstep, zones, nb_init_evals_per_zone)
    nb_zones = length(zones)
    zones_sum1 = zeros(nb_zones)
    zones_sum2 = zeros(nb_zones)
    for (n, zmin) = enumerate(zones)
        zmax = zmin .+ xstep
        for _ = 1:nb_init_evals_per_zone
            y = integrand(rand_min_max(zmin, zmax))
            zones_sum1[n] += y
            zones_sum2[n] += y^2
        end
    end
    zones_evals = fill(nb_init_evals_per_zone, nb_zones)
    zones_sum1, zones_sum2, zones_evals
end

function ucb_mab(integrand, nb_init_evals, nb_evals, k_exploration, zones,
                 xstep, zones_sum1, zones_sum2, zones_evals, zones_vars)
    for k=nb_init_evals:nb_evals
        scores = ucb1(k_exploration, k, zones_evals, zones_vars)
        n = indmax(scores)
        zmin = zones[n]
        zmax = zmin .+ xstep
        y = integrand(rand_min_max(zmin, zmax))
        zones_sum1[n] += y
        zones_sum2[n] += y^2
        zones_evals[n] += 1
        zones_vars[n] = ucb_vars(zones_sum1[n], zones_sum2[n], zones_evals[n])
    end
    zones_sum1, zones_evals
end

"""
    ucb(integrand, xmin, xmax, nb_evals, nb_zones_per_dim=3, k_exploration=0.01, nb_init_evals_per_zone=10)

Compute the n-dimensional integral f(x), where
`n == length(xmin) == length(xmax)`, over the hypercube whose corners are given
by the vectors (or tuples) `xmin` and `xmax`. `f` should be a function `f(x)`
that takes an n-dimensional vector `x` and returns the integrand at `x`.

Use the multi-armed bandit UCB method using `nb_evals` samples,
`nb_zones_per_dim` uniform strata in each dimension, `nb_init_evals_per_zone`
initial samples in each stratum and an exploration coefficient `k_exploration`.
"""
function ucb(integrand, xmin, xmax, nb_evals, nb_zones_per_dim=3,
             k_exploration=0.01, nb_init_evals_per_zone=10)
    # init
    dim, nb_zones, xstep, zones = make_zones(xmin, xmax, nb_zones_per_dim)
    nb_init_evals = nb_zones * nb_init_evals_per_zone
    if nb_evals < nb_init_evals
        warn("nb_evals < $(nb_init_evals)")
        return mc(integrand, xmin, xmax, nb_evals)
    end
    # stratified sampling
    zones_sum1, zones_sum2, zones_evals = ucb_strat(integrand, xstep, zones,
                                                    nb_init_evals_per_zone)
    zones_vars = ucb_vars(zones_sum1, zones_sum2, zones_evals)
    # multi-armed bandit
    zones_sum1, zones_evals =
    ucb_mab(integrand, nb_init_evals, nb_evals, k_exploration, zones,
            xstep, zones_sum1, zones_sum2, zones_evals, zones_vars)
    # final result
    prod(xstep) * sum(zones_sum1 ./ zones_evals)
end

end # module ucbature