module SimpleSolvers

    using Distances
    using ForwardDiff
    using LinearAlgebra
    using Printf

    import Base.minimum
    import Base.Callable

    include("utils.jl")

    export solve!
    export config, result, state, status
    export algorithm, objective
    export solution, minimizer, minimum

    export Newton, DFP, BFGS

    include("base/methods.jl")

    export UnivariateObjective,
           MultivariateObjective

    export value, value!, value!!,
           derivative, derivative!, derivative!!,
           gradient, gradient!, gradient!!,
           hessian, hessian!, hessian!!,
           d_calls, f_calls, g_calls, h_calls

    include("base/objectives.jl")

    export Options

    include("base/options.jl")

    export GradientParameters,
           GradientParametersAD,
           GradientParametersFD,
           GradientParametersUser

    export compute_gradient,
           compute_gradient!,
           compute_gradient_ad!,
           compute_gradient_fd!

    export check_gradient,
           print_gradient

    include("base/gradient.jl")

    export Hessian,
           HessianAD,
           HessianUser

    export compute_hessian,
           compute_hessian!,
           compute_hessian_ad!

    export check_hessian,
           print_hessian

    include("base/hessian.jl")

    export JacobianParameters,
           JacobianParametersAD,
           JacobianParametersFD,
           JacobianParametersUser

    export compute_jacobian!,
           compute_jacobian_ad!,
           compute_jacobian_fd!

    export check_jacobian,
           print_jacobian

    include("base/jacobian.jl")

    export LinearSolver, LUSolver, LUSolverLAPACK,
           factorize!

    include("linear/linear_solvers.jl")
    include("linear/lu_solver.jl")
    include("linear/lu_solver_lapack.jl")

    export bracket_minimum

    include("bracketing/bracketing.jl")
    include("bracketing/bracket_minimum.jl")

    export Linesearch, Static
    export Backtracking, backtracking,
           Bisection, bisection,
           Quadratic, quadratic

    include("linesearch/linesearch.jl")
    include("linesearch/static.jl")
    include("linesearch/backtracking.jl")
    include("linesearch/bisection.jl")
    include("linesearch/quadratic.jl")

    export NonlinearSolver, NonlinearSolverException,
           AbstractNewtonSolver, NLsolveNewton, NewtonSolver, QuasiNewtonSolver,
           residual_initial!, residual_absolute!, residual_relative!,
           assess_convergence, assess_convergence!,
           print_status, check_solver_status,
           get_solver_status, get_solver_status!,
           solve!

    include("nonlinear/nonlinear_solver_status.jl")
    include("nonlinear/nonlinear_solver.jl")
    include("nonlinear/abstract_newton_solver.jl")
    include("nonlinear/newton_solver.jl")
    include("nonlinear/quasi_newton_solver.jl")
    include("nonlinear/nlsolve_newton.jl")

    export Optimizer,
           NewtonOptimizer,
           BFGSOptimizer,
           DFPOptimizer,
           HessianBFGS,
           HessianDFP

    include("optimization/optimizer_status.jl")
    include("optimization/optimizer_result.jl")
    include("optimization/optimizer.jl")
    include("optimization/hessian_bfgs.jl")
    include("optimization/hessian_dfp.jl")
    include("optimization/newton_optimizer.jl")

end


export gradiente
export newton
export newton_rc

function gradiente(
  nlp :: AbstractNLPModel;
  max_time = 10.0,
  max_iter = 10_000,
  η₁ = 1e-2,
  atol = 1e-6,
  rtol = 1e-6,
)

t₀ = time()
Δt = time() - t₀
iter = 0

x = copy(nlp.meta.x0)
f(x) = obj(nlp, x)
∇f(x) = grad(nlp, x)
fx = f(x)
gx = ∇f(x)

x⁺ = similar(x)

ϵ = atol + rtol * norm(gx)

status = :unknown

# Cond. de parada
resolvido = norm(gx) < ϵ # true ou false
cansado = Δt > max_time || iter > max_iter
while !(resolvido || cansado) # Não tiver satisfeito as condições de parada
  # Calculo a direção e passo e xₖ₊₁
  d = -gx
  slope = dot(gx, d)

  α = 1.0
  x⁺ = x + α * d
  f⁺ = f(x⁺)
  while f⁺ ≥ fx + η₁ * α * slope
      α = α / 2
      x⁺ = x + α * d
      f⁺ = f(x⁺)
      if α < 1e-8
          status = :small_step
          break
      end
  end
  if status != :unknown
      break
  end
  x .= x⁺
  fx = f⁺
  gx = ∇f(x)

  resolvido = norm(gx) < ϵ
  Δt = time() - t₀
  iter += 1
  cansado = Δt > max_time || iter > max_iter
end

if resolvido
  status = :first_order
elseif cansado
  if Δt > max_time
      status = :max_time
  elseif iter > max_iter
      status = :max_iter
  end
end

return GenericExecutionStats(status, nlp, objective=fx,
  solution=x, dual_feas=norm(gx), iter=iter, elapsed_time=Δt)
end

##

function newton(
  nlp :: AbstractNLPModel;
  max_time = 10.0,
  max_iter = 10_000,
  η₁ = 1e-2,
  atol = 1e-6,
  rtol = 1e-6,
)

t₀ = time()
Δt = time() - t₀
iter = 0

x = copy(nlp.meta.x0)
f(x) = obj(nlp, x)
∇f(x) = grad(nlp, x)
H(x) = Symmetric(hess(nlp, x), :L)
fx = f(x)
gx = ∇f(x)

x⁺ = similar(x)

ϵ = atol + rtol * norm(gx)

status = :unknown

# Cond. de parada
resolvido = norm(gx) < ϵ # true ou false
cansado = Δt > max_time || iter > max_iter
while !(resolvido || cansado) # Não tiver satisfeito as condições de parada
  # Calculo a direção e passo e xₖ₊₁
  Hx = H(x)
  F = cholesky(Hx, check=false)
  if !issuccess(F)
      status = :not_desc
      break
  end
  d = -(F \ gx)
  slope = dot(gx, d)

  α = 1.0
  x⁺ = x + α * d
  f⁺ = f(x⁺)
  while f⁺ ≥ fx + η₁ * α * slope
      α = α / 2
      x⁺ = x + α * d
      f⁺ = f(x⁺)
      if α < 1e-8
          status = :small_step
          break
      end
  end
  if status != :unknown
      break
  end
  x .= x⁺
  fx = f⁺
  gx = ∇f(x)

  resolvido = norm(gx) < ϵ
  Δt = time() - t₀
  iter += 1
  cansado = Δt > max_time || iter > max_iter
end

if resolvido
  status = :first_order
elseif cansado
  if Δt > max_time
      status = :max_time
  elseif iter > max_iter
      status = :max_iter
  end
end

return GenericExecutionStats(status, nlp, objective=fx,
  solution=x, dual_feas=norm(gx), iter=iter, elapsed_time=Δt)
end

##

function newton_rc(
  nlp::AbstractNLPModel;
  atol::Real = 1e-6,
  rtol::Real = 1e-6,
  max_eval::Int = 1000,
  max_iter::Int = 0,
  max_time::Float64 = 10.0
)

  if !unconstrained(nlp)
    error("Problem is not unconstrained")
  end

  x = copy(nlp.meta.x0)

  f(x) = obj(nlp, x)
  ∇f(x) = grad(nlp, x)
  H(x) = Symmetric(hess(nlp, x), :L)

  fx = f(x)
  ∇fx = ∇f(x)

  ϵ = atol + rtol * norm(∇fx)
  t₀ = time()

  iter = 0
  Δt = time() - t₀
  solved = norm(∇fx) < ϵ # First order stationary
  tired = neval_obj(nlp) ≥ max_eval > 0|| iter ≥ max_iter > 0 || Δt ≥ max_time > 0 # Excess time, iteration, evaluations

  # status must be one of a few options found in SolverTools.show_statuses()
  # A good default value is :unknown.
  status = :unknown

  # log_header is up for some rewrite in the future. For now, it simply prints the column names with some spacing
  @info log_header(
    [:iter, :fx, :ngx, :nf, :Δt],
    [Int, Float64, Float64, Int, Float64],
    hdr_override=Dict(:fx => "f(x)", :ngx => "‖∇f(x)‖", :nf => "#f")
  )
  # log_row uses the type information of each value, thus we use `Any` here.
  @info log_row(
    Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
  )

  # This template implements a simple steepest descent method without any hopes of working.
  # This is where most of your change will happen
  η₁ = 1e-2
  η₂ = 0.75
  Δ = 1.0
  iter = 0

  while !(solved || tired)
    fx = f(x)
    gx = ∇f(x)
    Hx = H(x)
    model = Model(with_optimizer(Ipopt.Optimizer, print_level=0))
    @variable(model, d[1:2])
    @objective(model, Min, fx + dot(d, gx) + dot(d, Hx * d) / 2)
    @NLconstraint(model, d[1]^2 + d[2]^2 ≤ Δ^2)
    optimize!(model)
    d = value.(d)

    Ared = f(x) - f(x + d)
    Pred = f(x) - (fx + dot(d, gx) + dot(d, Hx * d) / 2)
    ρ = Ared / Pred
    if ρ < η₁
        Δ = Δ / 2
    elseif ρ < η₂
        x = x + d
    else
        x = x + d
        Δ = 2Δ
    end

    iter += 1
    Δt = time() - t₀
    solved = norm(∇fx) < ϵ # First order stationary
    tired = neval_obj(nlp) ≥ max_eval > 0|| #Evaluations
            iter ≥ max_iter > 0 || #iterations
            Δt ≥ max_time > 0 # Excess time

    @info log_row(
      Any[iter, fx, norm(∇fx), neval_obj(nlp), Δt]
    )
  end

  if solved
    status = :first_order
  elseif tired
    if neval_obj(nlp) ≥ max_eval > 0
      status = :max_eval
    elseif iter ≥ max_iter > 0
      status = :max_iter
    elseif Δt ≥ max_time > 0
      status = :max_time
    end
  end

  return GenericExecutionStats(
    status,
    nlp,
    solution=x,
    objective=f(x),
    dual_feas=norm(∇fx),
    elapsed_time=Δt,
    iter=iter
  )
end


module Tools
import Luna: Modes, PhysData, Capillary, RectModes, Maths
import Luna.PhysData: wlfreq
import Roots: find_zero
import Base: show
import Cubature: hquadrature
import Printf: @sprintf

"Calculate 'natural' pulse width from FWHM"
function τfw_to_τ0(τfw, shape)
    if shape == :sech
        τ0 = τfw/(2*log(1+sqrt(2)))
    elseif shape == :gauss
        τ0 = τfw/(2*sqrt(log(2)))
    else
        error("shape must be one of: :sech, :gauss")
    end
    τ0
end

"Get dispersion length"
function Ld(τfw, β2; shape=:sech)
    τ0 = τfw_to_τ0(τfw, shape)
    τ0^2/abs(β2)
end

"Get GVD coefficient"
function getβ2(ω, m::Modes.AbstractMode)
    Modes.dispersion(m, 2, ω)
end

"Get nonlinear length"
function Lnl(P0, γ)
    1/(γ*P0)
end

"Get fission length"
function Lfiss(P0, τfw, γ, β2; shape=:sech)
    Ld(τfw, β2, shape=shape)/getN(P0, τfw, γ, β2, shape=shape)
end

"Get nonlinear coefficient"
function getγ(ω, m::Modes.AbstractMode, n2)
    n2*ω/(PhysData.c*Modes.Aeff(m))
end

"Get linear and nonlinear refractive index and gas number density"
function getN0n0n2(ω, material; P=1.0, T=PhysData.roomtemp)
    N0 = PhysData.density(material, P, T)
    χ3 = PhysData.γ3_gas(material)*N0
    n0 = PhysData.ref_index(material, 2π*PhysData.c/ω, P, T)
    N0, n0, 3*χ3/(4*n0^2*PhysData.ε_0*PhysData.c)
end

"Get soliton order"
function getN(P0, τfw, γ, β2; shape=:sech)
    sqrt(Ld(τfw, β2, shape=:sech)/Lnl(P0, γ))
end

function E_to_P0(E, τfw; shape=:sech)
    τ0 = τfw_to_τ0(τfw, shape)
    if shape == :sech
        P0 = E/(2*τ0)
    else
        error("shape must be one of: :sech")
    end
    P0
end

function P0_to_I(P0, m)
    P0/Modes.Aeff(m)
end

function Pcr(ω, n0, n2)
    1.8962*(2π*PhysData.c/ω)^2/(4π*n0*n2)
end

paramfields = (:E, :τfw, :τ0, :ω, :λ, :material, :P, :T, :shape,:P0, :β2, :N0, :n0, :n2,
               :γ, :N, :I0, :Pcr, :Ld, :Lnl, :Lfiss, :zdw, :Lloss, :Aeff, :mode)

function show(io::IO, p::NamedTuple{paramfields, vT}) where vT
    mode = "MODE:\n  $(p.mode)"
    fill = @sprintf("FILL:\n  %.1f bar %s, Pcr = %.1e W, γ = %.1e (Wm)^-1, n2 = %.1e cm^2/W",
                     p.P, p.material, p.Pcr, p.γ, p.n2*1e4)
    wg = @sprintf("WAVEGUIDE:\n  Aeff = %.1e m^2, Lloss = %.1e m", p.Aeff, p.Lloss)
    pulse = @sprintf("PULSE:\n  %.2e J, %.2e s @ %.1f nm (shape: %s)",
                      p.E, p.τfw, p.λ*1e9, p.shape)
    dispersion = @sprintf("DISPERSION:\n  %.2e s^2/m @ %.1f nm, ZDW = %.1f nm",
                          p.β2, p.λ*1e9, p.zdw*1e9)
    intensity = @sprintf("INTENSITY:\n  %.1e W/cm^2", p.I0*1e-4)
    power = @sprintf("POWER:\n  %.1e W (%.4f of Pcr)", p.P0, p.P0/p.Pcr)
    sol = @sprintf("SOLITON:\n  Ld = %.1e m, Lnl = %.1e m, Lfiss = %.1e m, N = %.2f",
                   p.Ld, p.Lnl, p.Lfiss, p.N)
    out = join((mode, wg, fill, pulse, dispersion, power, intensity, sol), "\n")
    print(io, out)
end

"Soliton parameter collection"
function params(E, τfw, λ, mode, material; shape=:sech, P=1.0, T=PhysData.roomtemp)
    ω = 2π*PhysData.c/λ
    P0 = E_to_P0(E, τfw, shape=shape)
    τ0 = τfw_to_τ0(τfw, shape)
    β2 = getβ2(ω, mode)
    N0, n0, n2 = getN0n0n2(ω, material, P=P, T=T)
    γ = getγ(ω, mode, n2)
    N = getN(P0, τfw, γ, β2, shape=shape)
    zdw = Modes.zdw(mode)
    if ismissing(zdw)
        zdw = Modes.zdw(mode, λ)
    end
    p = (E=E, τfw=τfw, τ0=τ0, ω=ω, λ=λ, material=material, P=P, T=T, shape=shape,
         P0=P0, β2=β2, N0=N0, n0=n0, n2=n2, γ=γ, N=N,
         I0=P0_to_I(P0, mode), Pcr=Pcr(ω, n0, n2),
         Ld=Ld(τfw, β2, shape=shape),
         Lnl=Lnl(P0, γ),
         Lfiss=Lfiss(P0, τfw, γ, β2, shape=shape),
         zdw=zdw,
         Lloss=Modes.losslength(mode, ω),
         Aeff=Modes.Aeff(mode),
         mode=mode)
end

function capillary_params(E, τfw, λ, a, material;
                          shape=:sech, P=1.0, T=PhysData.roomtemp, clad=:SiO2, n=1, m=1,kind=:HE, ϕ=0.0)
    mode = Capillary.MarcatiliMode(a, material, P, n=n, m=m, kind=kind, ϕ=ϕ, T=T, clad=clad)
    params(E, τfw, λ, mode, material, shape=shape, P=P, T=T)
end

function rectangular_params(E, τfw, λ, a, b, material;
                            shape=:sech, P=1.0, T=PhysData.roomtemp, clad=:SiO2, n=1, m=1,
                            pol=:x)
    mode = RectModes.RectMode(a, b, material, P, clad, T=T, n=n, m=m, pol=pol)
    params(E, τfw, λ, mode, material, shape=shape, P=P, T=T)
end

function gas_ratio(gas1, gas2, λ)
    χ3r = PhysData.χ3_gas(gas1, 1) / PhysData.χ3_gas(gas2, 1)
    β2r = PhysData.dispersion(2, gas1, λ) / PhysData.dispersion(2, gas2, λ)
    β2r, χ3r
end

field_to_intensity(E) = 0.5*PhysData.ε_0*PhysData.c*E^2
intensity_to_field(I) = sqrt(2I/PhysData.ε_0/PhysData.c)

"""
    λRDW(m::Modes.AbstractMode, λ0; z=0, λlims=(100e-9, 0.9λ0))

Calculate the phase-matching wavelength for resonant dispersive wave (RDW) emission in the
mode `m` when pumping at `λ0`.

This neglects the nonlinear contribution to the phase mismatch.
"""
function λRDW(m::Modes.AbstractMode, λ0; z=0, λlims=(100e-9, 0.9λ0))
    ω0 = wlfreq(λ0)
    β1 = Modes.dispersion(m, 1, ω0; z=z)
    β0 = Modes.β(m, ω0; z=z)
    Δβ(ω) = Modes.β(m, ω; z=z) - β1*(ω.-ω0) - β0
    try
        ωRDW = find_zero(Δβ, extrema(wlfreq.(λlims)))
        wlfreq(ωRDW)
    catch
        missing
    end
end

"""
    λRDW(a::Number, gas::Symbol, pressure, λ0; λlims=(100e-9, 0.9λ0), kwargs...)

Calculate the phase-matching wavelength for resonant dispersive wave (RDW) emission in a
capillary with core radius `a` filled with `gas` at a certain `pressure`
when pumping at `λ0`. Additional `kwargs` are passed onto `Capillary.MarcatiliMode`.

This neglects the nonlinear contribution to the phase mismatch.
"""
function λRDW(a::Number, gas::Symbol, pressure, λ0; λlims=(100e-9, 0.9λ0), kwargs...)
    m = Capillary.MarcatiliMode(a, gas, pressure; kwargs...)
    λRDW(m, λ0; λlims=λlims)
end

"""
    pressureRDW(a::Number, gas::Symbol, λ_target, λ0; Pmax=100, clad=:SiO2, kwargs...)

Calculate the phase-matching pressure for resonant dispersive wave (RDW) emission at
`λ_target` in a capillary with core radius `a` filled with `gas` when pumping at `λ0`.
"""
function pressureRDW(a::Number, gas::Symbol, λ_target, λ0; Pmax=100, clad=:SiO2, kwargs...)
    # cladn is likely based on interpolation so requires creating the BSpline.
    # By creating the function here we only have to do that once
    rfc = PhysData.ref_index_fun(clad)
    cladn = (ω; z) -> rfc(wlfreq(ω))
    ω0 = wlfreq(λ0)
    ω_target = wlfreq(λ_target)
    function Δβ(P)
        m = Capillary.MarcatiliMode(a, gas, P, cladn; kwargs...)
        β1 = Modes.dispersion(m, 1, ω0)
        β0 = Modes.β(m, ω0)
        Modes.β(m, ω_target) - β1*(ω_target.-ω0) - β0
    end

    try
        find_zero(Δβ, (1e-6, Pmax))
    catch
        missing
    end
end

function aperture_filter(a, dist, radius)
    w0 = 0.64*a
    function filter(λ)
        w1 = dist*λ/(π*w0)
        I(r) = Maths.gauss(r, w1/2) / Maths.gaussnorm(w1/2)
        2*hquadrature(I, 0, radius)[1]
    end
end

end


using Combinatorics

# Part 1
ratings = sort(parse.(Int32, readlines("day10_input.txt")))
diffs = vcat(ratings, [ratings[end]+3]) - vcat([0], ratings)
println("Part 1: product = $(count(x -> x==1, diffs) * count(x -> x==3, diffs))")

# Part 2
valid(chain) = maximum([chain; chain[end]+3] - [chain[1]-3; chain]) <= 3
innercombinations(chain) = combinations(chain[2:(lastindex(chain)-1)])
swapguts(chain, inner) = [first(chain); inner; last(chain)]
arrangements(chain) = count([[valid(swapguts(chain, c)) for c in innercombinations(chain)]; valid(swapguts(chain, []))])
function subchains(chain)
    results = []
    current = [first(chain)]
    for i in 2:lastindex(chain)
        if last(current) + 3 == chain[i]
            push!(results, current)
            current = []
        end
        push!(current, chain[i])
    end
    push!(results, current)
    results
end
println("Part 2: arrangements = $(prod(arrangements.(subchains([0; ratings; ratings[end]+3]))))")


"""
    MSD!(data::Dict, alpha=10)
This will add a variable "MSD" to `data` which is created by `read_dump`. "MSD" is the abbreviation of Mean Square Displacement (MSD), which could be used to calculate diffusion coefficients with the Einstein relationship. `alpha` is the ration between simulation time and the maximum of MSD calculating times.
"""
function MSD!(data::Dict, alpha=10)
    try
        data["coord_corr"]
    catch
        PBC!(data)
    end
    # Reading Input
    num_steps = data["num_steps"]
    num_atoms = data["num_atoms"]
    coord_corr = copy_array(data["coord_corr"])

    # MSD
    num_msd_steps = convert(Int64, round(num_steps / alpha))
    msd = zeros(num_msd_steps, num_atoms, 4)
    for step = 0 : num_msd_steps-1
        msd[step+1, :, 1:3] = mean((coord_corr[1+step:end, :, :] .- coord_corr[1:end-step, :, :]).^2, dims=1)
        msd[step+1, :, 4] = sum(msd[step+1, :, 1:3], dims=2)
    end
    data["MSD"] = msd
    println("\"MSD\"\t\thas been added to `data`")
end

"""
    diffusion(data::Dict, alpha=10)
This will return the diffusion coefficients of each atoms in `data`.

Notes: time_step shoulb be added by `time_step!()` before calling this function.
"""
function diffusion(data::Dict, alpha=10)
    try
        data["MSD"]
    catch
        MSD!(data, alpha)
    end

    try
        data["time_step"]
    catch
        error("\"time_step\" is not contained in `data`. Please call `time_step!()` before calling this function")
    end

    # Reading Input
    MSD = copy_array(data["MSD"])
    time_step = prod(data["time_step"])
    num_atoms = data["num_atoms"]

    # Diffusion
    num_step_msd = size(MSD, 1)
    t = [i*time_step for i=0:num_step_msd-1]
    res = zeros(num_atoms)

    for atom = 1:num_atoms
        res[atom] = leastsq(t, MSD[:, atom, 4])[2]
    end
    return res
end


using Test
using DiscreteFunctions, Permutations

f = DiscreteFunction(2,3,4,1)
ff = DiscreteFunction([2,3,4,1])
@test f==ff
@test hash(f) == hash(ff)
@test f(1) == 2
g = f^2
@test g(1) == 3
@test has_inv(g)
g = inv(f)
@test g(1) == 4
@test g[1] == g(1)
@test f*g == IdentityFunction(4)
g[1]=2
@test g(1)==2

@test length(f) == 4

@test f*f*f == f^3
@test inv(f) == f^-1

f = RandomFunction(4)
@test f*IdentityFunction(4) == f

f = IdentityFunction(4)
@test fixed_points(f) == collect(1:4)
@test has_inv(f)
@test is_permutation(f)
@test image(f) == Set{Int}(1:4)
@test length(cycles(f)) == length(f)
@test length(sources(f)) == 0

p = RandomPermutation(10)
f = DiscreteFunction(p)
@test f.data == p.data

g = f*f
A = Matrix(f)
B = Matrix(g)
@test A*A == B

f = sqrt(g)
@test g == f*f

f = RandomFunction(9)
A = Matrix(f)
@test f == DiscreteFunction(A)


"""
DIJKSTRA
--------
    compute shortest paths using Dijkstra's algorithm.
    d = dijkstra(A,u) computes the shortest path from vertex u to all nodes
    reachable from vertex u using Dijkstra's algorithm for the problem.
    The graph is given by the weighted sparse matrix A, where A(i,j) is
    the distance between vertex i and j.  In the output vector d,
    the entry d[v] is the minimum distance between vertex u and vertex v.
    A vertex w unreachable from u has d(w)=Inf.
    pred is the predecessor tree to generate the actual shorest paths.
    In the predecessor tree pred[v] is the vertex preceeding v in the
    shortest path and pred[u]=0. Any unreachable vertex has pred[w]=0 as well.
    If your network is unweighted, then use bfs instead.

Functions
---------
- (d,pred) = dijkstra(A::MatrixNetwork,u::Int64)
- (d,pred) = dijkstra{F}(A::SparseMatrixCSC{F,Int64},u::Int64)

Example
-------
~~~
# Find the minimum travel time between Los Angeles (LAX) and
# Rochester Minnesota (RST).

(A,xy,labels) = load_matrix_network_metadata("airports")
A = -A; # fix funny encoding of airport data
lax = 247; rst = 355
(d,pred) = dijkstra(A,lax)
@printf("Minimum time: %d",d[rst]); #Print the path
@printf("Path:")
u = rst
while(u != lax)
    @printf("%s <-- ", labels[u])
    u = pred[u];
    if (u == lax)
        @printf("%s", labels[lax])
    end
end
~~~
"""

function dijkstra(A::MatrixNetwork,u::Int64)
    (rp,ci,ai) = (A.rp, A.ci, A.vals)
    return dijkstra_internal(rp,ci,ai,u)
end

function dijkstra{F}(A::SparseMatrixCSC{F,Int64},u::Int64)
    (rp,ci,ai) = sparse_to_csr(A)
    return dijkstra_internal(rp,ci,ai,u)
end

function dijkstra_internal{F}(rp::Vector{Int64},ci::Vector{Int64},ai::Vector{F},u::Int64)

    if any(ai.<0)
        error("dijkstra''s algorithm cannot handle negative edge weights")
    end

    n = length(rp) - 1
    d = Inf*ones(Float64,n)
    T = zeros(Int64,n)
    L = zeros(Int64,n)
    pred = zeros(Int64,length(rp)-1)

    n = 1
    T[n] = u
    L[u] = n # oops, n is now the size of the heap

    # enter the main dijkstra loop
    d[u] = 0
    while n > 0
        v = T[1]
        ntop = T[n]
        T[1] = ntop
        L[ntop] = 1
        n = n - 1 # pop the head off the heap
        k = 1
        kt = ntop # move element T[1] down the heap
        while true
            i = 2*k
            if i > n
                break
            end      # end of heap
            if i == n
                it = T[i]        # only one child, so skip
            else               # pick the smallest child
                lc = T[i]
                rc = T[i+1]
                it = lc
                if d[rc] < d[lc]
                    i = i+1
                    it = rc
                end # right child is smaller
            end
            if d[kt] < d[it]
                break   # at correct place, so end
            else
                T[k] = it
                L[it] = k
                T[i] = kt
                L[kt] = i
                k=i # swap
            end
        end       # end heap down
        # for each vertex adjacent to v, relax it
        for ei = rp[v]:rp[v+1]-1       # ei is the edge index
            w = ci[ei]
            ew = ai[ei]          # w is the target, ew is the edge weight
            # relax edge (v,w,ew)
            if d[w] > d[v] + ew
                d[w] = d[v] + ew
                pred[w] = v
                # check if w is in the heap
                k = L[w]
                onlyup=false
                if k == 0
                    # element not in heap, only move the element up the heap
                    n = n+1
                    T[n] = w
                    L[w] = n
                    k = n
                    kt = w
                    onlyup = true
                else
                    kt = T[k]
                end
                # update the heap, move the element down in the heap
                while !onlyup
                    i = 2 * k
                    if i > n
                        break
                    end          # end of heap
                    if i == n
                        it = T[i]    # only one child, so skip
                    else            # pick the smallest child
                        lc = T[i]
                        rc = T[i+1]
                        it = lc
                        if d[rc] < d[lc]
                            i = i+1
                            it = rc
                        end # right child is smaller
                    end
                    if d[kt] < d[it]
                        break      # at correct place, so end
                    else
                        T[k] = it
                        L[it] = k
                        T[i] = kt
                        L[kt] = i
                        k = i # swap
                    end
                end
                # move the element up the heap
                j = k
                tj = T[j]
                while j > 1                      # j==1 => element at top of heap
                    j2 = round(Int,floor(j/2))
                    tj2 = T[j2]    # parent element
                    if d[tj2] < d[tj]
                        break      # parent is smaller, so done
                    else                         # parent is larger, so swap
                        T[j2] = tj
                        L[tj] = j2
                        T[j] = tj2
                        L[tj2] = j
                        j=j2
                    end
                end
            end
        end
    end
    return (d,pred)
end


using Gadfly

"""
Precision-recall curve for two hypothetical classifications systems.
A is better than B.
"""
function PRhandDemo()
    domain = 0:0.01:1

    fA(x) = 1-x.^(3)
    fB(x) = 1-x.^(3/2)

    lA = layer(x=domain, y=fA(domain), Geom.line, Theme(default_color=colorant"red",  line_width=3px, line_style=Gadfly.get_stroke_vector(:dot)))
    lB = layer(x=domain, y=fB(domain), Geom.line, Theme(default_color=colorant"blue", line_width=3px))

    plot(lA, lB, Guide.xlabel("recall"), Guide.ylabel("precision"))
end

"""
ROC curves for two hypothetical classification systems.
A is better than B.
Plots true positive rate, (tpr) vs false positive (fpr).
"""
function ROChandDemo()
    domain = 0:0.01:1

    fA(x) = x.^(1/3)
    fB(x) = x.^(2/3)

    lA = layer(x=domain, y=fA(domain), Geom.line, Theme(default_color=colorant"red",  line_width=3px, line_style=Gadfly.get_stroke_vector(:dot)))
    lB = layer(x=domain, y=fB(domain), Geom.line, Theme(default_color=colorant"blue", line_width=3px))
    shade = layer(x=domain, y=zeros(domain), ymin=zeros(domain), ymax=fB(domain), Geom.ribbon)

    plot(lA, shade, lB, Guide.xlabel("fpr"), Guide.ylabel("tpr"))
end


@testset "colorfulness" begin

    @info "test: colorfulness"

    # Test against a simple result calculated by hand

    x = [RGB(1,0,0), RGB(0,1,0), RGB(0,0,1)]
    @test trunc(colorfulness(x)) == 337

    # Black is not colorful

    x = [RGB(0,0,0), RGB(0,0,0), RGB(0,0,0)]
    @test trunc(colorfulness(x)) == 0

    # White is not colorful

    x = [RGB(1,1,1), RGB(1,1,1), RGB(1,1,1)]
    @test trunc(colorfulness(x)) == 0

    # A grayscale image is not colorful

    cameraman = testimage("cameraman")
    @test colorfulness(cameraman) == 0

    # A color image with only grays is not colorful

    x = convert(Array{Float64}, cameraman)
    img = RGB.(x, x, x)
    @test colorfulness(img) == 0

    # Lena 256 is a reduced color image and so should be less colorful
    # than the original

    imga = testimage("lena_color_256")
    imgb = testimage("lena_color_512")

    @test  colorfulness(imga) < colorfulness(imgb)

    # Test all invocation styles are equivalent

    c1 = colorfulness(imga)
    c2 = hasler_and_susstrunk_m3(imga)
    c3 = colorfulness(HASLER_AND_SUSSTRUNK_M3(), imga)
    c4 = assess(HASLER_AND_SUSSTRUNK_M3(), imga)

    @test c1 == c2 == c3 == c4

end


type EpsilonGreedy <: BanditAlgorithm
  epsilon::Float64
  counts::Vector{Int64}
  values::Vector{Float64}
end

function EpsilonGreedy(epsilon::Float64, n_arms::Int64)
  EpsilonGreedy(epsilon, zeros(Int64, n_arms), zeros(n_arms))
end

function initialize(algo::EpsilonGreedy, n_arms::Int64)
  algo.counts = zeros(Int64, n_arms)
  algo.values = zeros(n_arms)
end

function initialize(algo::EpsilonGreedy, arms::Vector)
  initialize(algo, length(arms))
end

function select_arm(algo::EpsilonGreedy)
  if rand() > algo.epsilon
    return ind_max(algo.values)
  else
    return randi(length(algo.values))
  end
end

function update(algo::EpsilonGreedy, chosen_arm::Int64, reward::Real)
  algo.counts[chosen_arm] += 1
  n = algo.counts[chosen_arm]

  value = algo.values[chosen_arm]
  algo.values[chosen_arm] = ((n - 1) / n) * value + (1 / n) * reward
end


"""
Make a cartesian product equispaced grid on `[-1,1]^N`.
"""
function equispaced_grid(::Val{N}, n::Int) where {N}
    grid1 = PeriodicEquispacedGrid{T}(n)
    grid = ProductGrid( ntuple(t->grid1, Val{N}())... )
end

make_centers(grid::AbstractGrid) = collect(grid)

restrict_grid(grid, domain) = MaskedGrid(grid, domain)

function rbfapprox(rbf::RBF, centers, points, f; verbose = false, method = :svd)
    basis = RBFDictionary(centers, rbf)
    A = evaluation_matrix(basis, points)
    B = f.(points)
    if method == :LAPACK
        coef, rank = Base.LinAlg.LAPACK.gelsy!(A,B)
        if verbose
            println("Effective rank: ", rank)
        end
    elseif method == :svd
        threshold = 1e-14
        usv = LAPACK.gesdd!('S',A)
        s = usv[2]
        maxindex = findlast(s.>threshold)
        if verbose
            println("Effective svd-rank: ", maxindex)
        end
        sreg = 1 ./ s[1:maxindex]
        ureg = usv[1][:,1:maxindex]
        vreg = usv[3][1:maxindex,:]
        coef = vreg' * diagm(sreg) * (ureg'*B)
    else
        error("Please supply a valid method for rbfapprox")
    end

    Expansion(basis, coef)
end


cd(@__DIR__); include("setups/simpletree.jl")
gr(dpi = 200)

## build frames
SGWT = sgwt_frame(Matrix(W); nf = 6)
SGWT = reshape(SGWT, (N, :))
SGWT_dual = (SGWT * SGWT') \ SGWT

distROT = natural_eigdist(𝚽, 𝛌, Q; α = 1.0, input_format = :pmf1, distance = :ROT)
rNGWF, dic_l2x = rngwf_all_vectors(distROT, 𝚽; σ = 0.1 * maximum(distROT), thres = 0.15)
rNGWF_dual = (rNGWF * rNGWF') \ rNGWF
Γ = rngwf_lx(dic_l2x)

## (a)
plt = plot(size = (600, 500))
    gplot!(W, X, width = 1)
    scatter_gplot!(X; marker = f, ms = 4)
    plot!(frame = :none, cbar = true, xlim = [-10, 14])
savefig(plt, "../figs/simpletree_f.png")

## (b)
rel_approx_sgwt, f_approx_sgwt = frame_approx(f, SGWT, SGWT_dual; num_kept = 3 * N)
rel_approx_rngwf, f_approx_rngwf = frame_approx(f, rNGWF, rNGWF_dual; num_kept = 3 * N)
plt = plot(0:length(rel_approx_rngwf)-1, [rel_approx_sgwt rel_approx_rngwf],
    grid = false, lw = 2, c = [:blue :green],
    xlab = "Number of Coefficients Retained",
    ylab = "Relative Approximation Error", yaxis = :log,
    lab = ["SGWT" "rNGWF"])
    plot!(xguidefontsize = 14, yguidefontsize = 14, legendfontsize = 11, size = (600, 500))
savefig(plt, "../figs/simpletree_approx_f.png")


## JL-ML SPGL Comparison
export scale_cs

"""
An example of using implicit JOLI operators to solve a system using GenSPGL

Use: x, r, g, info, SNR = ex_cs()
"""
function scale_cs(n,k,i)

    #Create a sparse vector
    #inds = randperm(n)
    #x0 = zeros(n)
    #keep = Int(floor(n/1000))
    #x0[inds[1:keep]] = randn(keep)
    #r_inds = inds[1:k]

    #= Load sparse vector x and restriction mask
    fid = matopen(Pkg.dir("GenSPGL")*"/src/Examples/scaletest/x0.mat")
    x0 = vec(read(fid, "x0"))
    r_inds = Int.(vec(read(fid, "r_inds")))
    n = length(vec(x0))
    close(fid)
    =#

    # Load sparse vector x and restriction mask
    fid = matopen(Pkg.dir("GenSPGL")*"/src/Examples/scaletest/x0.mat")
    x0 = vec(read(fid, "x0"))
    r_inds = Int.(vec(read(fid, "r_inds")))
    n = length(vec(x0))
    close(fid)

    # Sparsity Promoting Transform
    S = joDCT(n)

    # Restriction Operator
    R = joRestriction(n, r_inds)

    # Modelling op
    A = R*S

    # Create data
    b = A*x0

    # Set BLAS Threads
    BLAS.set_num_threads(8)

    # Solve
    opts = spgOptions(verbosity = 1, iterations = 100)
    t = [@elapsed spgl1(A, b, options = opts) for ii = 1:i]
    return t

end


function dataenvelop(input::Array{Float64,2}, output::Array{Float64,1}; verbose::Bool=false)::DataEnvelopResult

    nrow, ncol = size(input)

    efficiencies = zeros(Float64, nrow)
    resultdf = DataFrame()
    casenames::Array{Symbol,1} = []

    for objectnum in 1:nrow

        model = Model(Cbc.Optimizer);
        MOI.set(model, MOI.Silent(), !verbose)

        @variable(model, x[1:nrow])
        @variable(model, theta)

        @objective(model, Min, theta)

        for varnum in 1:ncol
            @constraint(model, sum(input[:, varnum] .* x[1:nrow]) - input[objectnum, varnum] * theta <= 0)
        end
        @constraint(model, sum(output .* x[1:nrow]) - output[objectnum] >= 0)

        @constraint(model, theta >= 0)
        for i in 1:nrow
            @constraint(model, x[i] >= 0)
        end

        optimize!(model)

        values = JuMP.value.(x)
        theta = objective_value(model)

        efficiencies[objectnum] = theta
        push!(casenames,  Symbol(string("Case", objectnum)))
        resultdf[:, casenames[objectnum]] = values[1:nrow]
    end

    orderedEfficiencyIndex = sortperm(efficiencies, ) |> reverse
    orderedCases = casenames[orderedEfficiencyIndex]

    result = DataEnvelopResult(
        efficiencies,
        resultdf,
        orderedCases
    )

    return result
end


# Different methods of generating PSF

using FFTW

field(amplitude, phase) = amplitude .* exp.( 1im * phase)

"""
psf(amplitude, phase) -> psfimage
psf(pupilfield) -> psfimage

Calculate psf for given amplitude and phase.

Examples
====
```jldoctest

julia> ap,_ = PhaseRetrieval.aperture(-1:.2:1,-1:.2:1,.8)
([0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0 1.0 1.0 1.0 1.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0], [NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN; NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN; NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN; NaN NaN NaN NaN NaN 1.0 NaN NaN NaN NaN NaN; NaN NaN NaN NaN 1.0 1.0 1.0 NaN NaN NaN NaN; NaN NaN NaN 1.0 1.0 1.0 1.0 1.0 NaN NaN NaN; NaN NaN NaN NaN 1.0 1.0 1.0 NaN NaN NaN NaN; NaN NaN NaN NaN NaN 1.0 NaN NaN NaN NaN NaN; NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN; NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN; NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN NaN])

julia> psfimage = psf(ap)
[17.720305330423212 17.720305330423212 2.830830026003772 0.7990467619253477 1.7153703234534299 0.002318499057199975 0.8566413660014238 0.0023184990571999695 1.7153703234534299 0.7990467619253475 2.830830026003772; 17.720305330423212 17.720305330423212 2.830830026003772 0.7990467619253475 1.7153703234534299 0.0023184990571999695 0.8566413660014238 0.002318499057199975 1.7153703234534299 0.7990467619253477 2.830830026003772; 2.830830026003773 2.830830026003773 0.2240431494891658 5.881503709461158 4.671643508435733 0.0810140527710051 0.6181197482998197 0.0810140527710051 4.671643508435733 5.881503709461155 0.2240431494891658; 0.7990467619253475 0.7990467619253477 5.881503709461158 11.063720826850961 3.6825070656623606 0.6902785321094302 4.96008586865757 0.6902785321094302 3.6825070656623606 11.063720826850961 5.881503709461158; 1.7153703234534294 1.7153703234534299 4.671643508435735 3.6825070656623615 0.5365498748309357 19.64548752112056 38.22662768629444 19.64548752112056 0.5365498748309356 3.682507065662362 4.671643508435733; 0.00231849905719993 0.002318499057199919 0.08101405277100532 0.6902785321094301 19.64548752112056 78.45538081840574 118.33852533074673 78.45538081840574 19.64548752112056 0.6902785321094301 0.08101405277100517; 0.8566413660014239 0.8566413660014239 0.6181197482998197 4.96008586865757 38.22662768629444 118.33852533074673 169.0 118.33852533074673 38.22662768629444 4.96008586865757 0.6181197482998197; 0.002318499057199919 0.00231849905719993 0.08101405277100517 0.6902785321094301 19.64548752112056 78.45538081840574 118.33852533074673 78.45538081840574 19.64548752112056 0.6902785321094301 0.08101405277100532; 1.7153703234534299 1.7153703234534294 4.671643508435733 3.682507065662362 0.5365498748309356 19.64548752112056 38.22662768629444 19.64548752112056 0.5365498748309357 3.6825070656623615 4.671643508435735; 0.7990467619253477 0.7990467619253475 5.881503709461158 11.063720826850961 3.6825070656623606 0.6902785321094302 4.96008586865757 0.6902785321094302 3.6825070656623606 11.063720826850961 5.881503709461158; 2.830830026003773 2.830830026003773 0.2240431494891658 5.881503709461155 4.671643508435733 0.0810140527710051 0.6181197482998197 0.0810140527710051 4.671643508435733 5.881503709461158 0.2240431494891658]

```
"""
psf(field::AbstractArray) = abs.(toimageplane(field)) .^2

psf(amplitude, phase) = psf(field(amplitude, phase))

subpsf(field::AbstractArray,Q::Integer) = psf(subdivide_sum(field,Q))
subpsf(amplitude, phase,Q::Integer) = psf(subdivide_sum(field(amplitude, phase),Q))

toimageplane(field) = ifftshift(fft(fftshift(field)))


function michaelis_menten(u, p, t)
    [0.6 - 1.5u[1]/(0.3+u[1])]
end

u0 = [0.5]
tspan = (0.0, 4.0)

problem_1 = ODEProblem(michaelis_menten, u0, tspan)
solution_1 = solve(problem_1, Tsit5(), saveat = 0.1)
problem_2 = ODEProblem(michaelis_menten, 2*u0, tspan)
solution_2 = solve(problem_2, Tsit5(), saveat = 0.1)
X = [solution_1[:,:] solution_2[:,:]]
ts = [solution_1.t; solution_2.t]

DX = similar(X)
for (i, xi) in enumerate(eachcol(X))
    DX[:, i] = michaelis_menten(xi, [], ts[i])
end

@parameters t
@variables u[1:2]
h = [monomial_basis(u[1:1], 4)...]
basis = Basis([h; h .* u[2]], u)

@testset "Ideal data" begin

    prob = ContinuousDataDrivenProblem(X, ts, DX)

    opts = [ImplicitOptimizer(5e-1);ImplicitOptimizer(0.4:0.1:0.7)]
    for opt in opts
        res = solve(prob, basis, opt, normalize = false, denoise = false, maxiter = 10000)
        m = metrics(res)
        @test m.Error < 5e-1
        @test m.AICC < 23.0
        @test m.Sparsity == 4
    end

    for opt in [ADM(5e-1); ADM(4e-1:1e-3:5e-1)]
        res = solve(prob, basis, opt, normalize = false, denoise = false, maxiter = 1000)
        m = metrics(res)
        @test m.Error < 8e-1
        @test m.AICC < 23.0
        @test m.Sparsity == 4
    end
end

Random.seed!(2345)
X = X .+ 1e-3*randn(size(X))

@testset "Noisy data" begin

    prob = ContinuousDataDrivenProblem(X, ts, GaussianKernel())

    for opt in [ImplicitOptimizer(4e-1);ImplicitOptimizer([0.3; 0.4; 0.5])]
        res = solve(prob, basis, opt, normalize = true, denoise = true)
        m = metrics(res)
        @test m.Error < 3e-1
        @test m.AICC < 35.0
        @test m.Sparsity == 4
    end

    for opt in [ADM(0.01:0.01:4e-1)]
        res = solve(prob, basis, opt, normalize = false, denoise = false)
        m = metrics(res)
        @test m.Error < 5e-1
        @test m.AICC < 12.0
        @test m.Sparsity == 4
    end
end


ENV["JULIA_CUDA_SILENT"] = true
# Change the working directory to the file location
cd(@__DIR__)

using DrMZ
using Parameters: @with_kw
using Flux.NNlib
using LaTeXStrings
using ColorSchemes
using Flux
using Plots; pyplot()
using Plots.PlotMeasures
using Random
using Printf

default(fontfamily="serif",frame=:box,grid=:hide,palette=:viridis,markeralpha=0.4,dpi=200,legendfontsize=6);
PyPlot.matplotlib.rc("mathtext",fontset="cm")

@with_kw mutable struct Args
    num_sensors::Int = 128;
    L1::Float64 = 0.0;
    L2::Float64 = 2*pi;
    n_epoch::Int = 50000;
    N::Int = 128;
end

function generate_basis_function_results(pde_function;kws...)

    args = Args(;);

    M = args.num_sensors;

    # Load the trained trunk
    trunk = load_trunk(args.n_epoch,pde_function)

    # Generate quadrature grid
    x, w = gauss_legendre(args.N,args.L1,args.L2);

    # Generate custom basis functions
    basis, S = build_basis(trunk,args.L1,args.L2,M,x,w);

    # Plot singular values
    pltsing = plot((1:1:length(S)),S,label=false,seriestype=:scatter,yaxis = :log10,xlabel = L"k", ylabel = L"$\sigma_k$",xlims=(1,length(S)))
    display(pltsing)
    savefig(pltsing, @sprintf("singular_values_%s_%i_functions.png",pde_function,M))

    # Plot custom basis functions
    x_plot = (args.L1:0.01:args.L2);
    pltbasis = plot(x_plot,basis[1].(x_plot),legend=false,label=L"$\phi_1$",xlabel=L"x",xlims=(x[1],x[end]),linewidth=1.5)
    for i in 2:M
      plot!(pltbasis,x_plot,basis[i].(x_plot),legend=false,label=L"$\phi_{$i}$",linewidth=1.5)
    end
    display(pltbasis)
    savefig(pltbasis, @sprintf("basis_functions_%s_%i_functions.png",pde_function,M))

    # Compute the expansion coefficients
    ak = expansion_coefficients(basis,exp.(sin.(x)),x,w);

    # Plot the expansion coefficients
    pltak = plot((1:1:length(basis)),abs.(ak),label=false,seriestype=:scatter,yaxis = :log10,xlabel = L"k", ylabel = L"|$a_k$|")
    display(pltak)
    savefig(pltak, @sprintf("expansion_coefficients_%s_%i_functions_exp_sinx.png",pde_function,M))

end

# Generate the custom basis functions from the trained trunk functions
generate_basis_function_results("advection_equation")


using Test

using MarketData
using TimeSeries

using MarketTechnicals

@testset "Volatility" begin

@testset "bollinger_bands" begin
    # TTR default uses sample=FALSE and value 117.3251
    @test isapprox(values(bollingerbands(cl)[:up])[1]  , 117.3251, atol=.01)
    # TTR default uses sample=FALSE and value 89.39392
    @test isapprox(values(bollingerbands(cl)[:down])[1], 89.39392, atol=.01)
    # TTR 103.3595
    @test isapprox(values(bollingerbands(cl)[:mean])[1], 103.3595, atol=.01)
    @test timestamp(bollingerbands(cl))[1]   == Date(2000,1,31)
    @test timestamp(bollingerbands(cl))[end] == Date(2001,12,31)
end

@testset "truerange" begin
    @test isapprox(values(truerange(ohlc))[end], 0.83, atol=.01) # TTR 0.83
    @test timestamp(truerange(ohlc))[end] == Date(2001,12,31)
end

@testset "donchianchannels" begin
    ta = donchianchannels(ohlc)

    @test meta(ta)      == meta(ohlc)
    @test timestamp(ta) == timestamp(ohlc[20:end])
    @test isapprox(values(ta[:up])[1]   , 121.5, atol=.01)
    @test isapprox(values(ta[:down])[1] , 86.5, atol=.01)
    @test values(ta[:mid])[1] == (values(ta[:up])[1] + values(ta[:down])[1]) / 2
end

@testset "atr" begin
    @test isapprox(values(atr(ohlc))[1]  , 8.343571428571428)   # TTR value 8.343571
    @test isapprox(values(atr(ohlc))[end], 0.9664561242651976)  # TTR value 0.9664561
    @test timestamp(atr(ohlc))[1] == Date(2000,1,24)
end

@testset "keltner_bands" begin
    ta = keltnerbands(ohlc)

    @test values(ta[:kup])[end] > values(ta[:kma])[end]
    @test values(ta[:kma])[end] > values(ta[:kdn])[end]

    @test isapprox(values(ta[:kup])[end], 23.3156, atol=.01)  # needs confirmation
    @test isapprox(values(ta[:kma])[end], 21.3705, atol=.01)  # needs confirmation
    @test isapprox(values(ta[:kdn])[end], 19.4254, atol=.01)  # needs confirmation
    @test timestamp(ta)[1]   == Date(2000, 2, 1)
    @test timestamp(ta)[end] == Date(2001, 12, 31)
end

@testset "chaikin_volatility" begin
    ta = chaikinvolatility(ohlc)

    @test timestamp(ta) == timestamp(ohlc[10+10:end])
    @test colnames(ta)  == [:chaikinvolatility]
    @test meta(ta)      == meta(ohlc)
    @test isapprox(values(ta)[1], -2.2401, atol=.01)  # needs confirmation
    @test isapprox(values(ta)[2], -3.2901, atol=.01)  # needs confirmation
end

end  # @testset "Volatility"


function petk04()
#
# M. Petkovic testing polynomials, p139
#
    z = [3.0; -1*[1;1;1]; 2*im*[1;1;1]; (-2+im)*[1;1]; (-2-im)*[1;1];
         (2+im)*[1;1]; (2-im)*[1;1]];
    p = poly(z);
    z = [3 1; -1 3;  2*im 3; -2+im 2; -2-im 2; 2+im 2; 2-im 2];

    p, PolyZeros(z)
end


"""
    AbstractInverseRetractionMethod

Abstract type for methods for inverting a retraction (see [`inverse_retract`](@ref)).
"""
abstract type AbstractInverseRetractionMethod end

"""
    AbstractRetractionMethod

Abstract type for methods for [`retract`](@ref)ing a tangent vector to a manifold.
"""
abstract type AbstractRetractionMethod end

"""
    ApproximateInverseRetraction <: AbstractInverseRetractionMethod

An abstract type for representing approximate inverse retraction methods.
"""
abstract type ApproximateInverseRetraction <: AbstractInverseRetractionMethod end

"""
    ApproximateRetraction <: AbstractInverseRetractionMethod

An abstract type for representing approximate retraction methods.
"""
abstract type ApproximateRetraction <: AbstractRetractionMethod end

"""
    ExponentialRetraction

Retraction using the exponential map.
"""
struct ExponentialRetraction <: AbstractRetractionMethod end

"""
    PolarRetraction <: AbstractRetractionMethod

Retractions that are based on singular value decompositions of the matrix / matrices
for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct PolarRetraction <: AbstractRetractionMethod end

"""
    ProjectionRetraction <: AbstractRetractionMethod

Retractions that are based on projection and usually addition in the embedding.
"""
struct ProjectionRetraction <: AbstractRetractionMethod end

"""
    QRRetraction <: AbstractRetractionMethod

Retractions that are based on a QR decomposition of the
matrix / matrices for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct QRRetraction <: AbstractRetractionMethod end

"""
    LogarithmicInverseRetraction

Inverse retraction using the [`log`](@ref)arithmic map.
"""
struct LogarithmicInverseRetraction <: AbstractInverseRetractionMethod end

"""
    PolarInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a singular value decomposition of the
matrix / matrices for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct PolarInverseRetraction <: AbstractInverseRetractionMethod end

"""
    ProjectionInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a projection (or its inversion).
"""
struct ProjectionInverseRetraction <: AbstractInverseRetractionMethod end

"""
    QRInverseRetraction <: AbstractInverseRetractionMethod

Inverse retractions that are based on a QR decomposition of the
matrix / matrices for point and tangent vector on a [`AbstractManifold`](@ref)
"""
struct QRInverseRetraction <: AbstractInverseRetractionMethod end

"""
    NLsolveInverseRetraction{T<:AbstractRetractionMethod,TV,TK} <:
        ApproximateInverseRetraction

An inverse retraction method for approximating the inverse of a retraction using `NLsolve`.

# Constructor

    NLsolveInverseRetraction(
        method::AbstractRetractionMethod[, X0];
        project_tangent=false,
        project_point=false,
        nlsolve_kwargs...,
    )

Constructs an approximate inverse retraction for the retraction `method` with initial guess
`X0`, defaulting to the zero vector. If `project_tangent` is `true`, then the tangent
vector is projected before the retraction using `project`. If `project_point` is `true`,
then the resulting point is projected after the retraction. `nlsolve_kwargs` are keyword
arguments passed to `NLsolve.nlsolve`.
"""
struct NLsolveInverseRetraction{TR<:AbstractRetractionMethod,TV,TK} <:
       ApproximateInverseRetraction
    retraction::TR
    X0::TV
    project_tangent::Bool
    project_point::Bool
    nlsolve_kwargs::TK
    function NLsolveInverseRetraction(m, X0, project_point, project_tangent, nlsolve_kwargs)
        return new{typeof(m),typeof(X0),typeof(nlsolve_kwargs)}(
            m,
            X0,
            project_point,
            project_tangent,
            nlsolve_kwargs,
        )
    end
end
function NLsolveInverseRetraction(
    m,
    X0 = nothing;
    project_tangent::Bool = false,
    project_point::Bool = false,
    nlsolve_kwargs...,
)
    return NLsolveInverseRetraction(m, X0, project_point, project_tangent, nlsolve_kwargs)
end

"""
    default_inverse_retraction_method(M::AbstractManifold)

The [`AbstractInverseRetractionMethod`](@ref) that is used when calling
[`inverse_retract`](@ref) without specifying the inverse retraction method.
By default, this is the [`LogarithmicInverseRetraction`](@ref).
"""
default_inverse_retraction_method(::AbstractManifold) = LogarithmicInverseRetraction()

"""
    default_retraction_method(M::AbstractManifold)

The [`AbstractRetractionMethod`](@ref) that is used when calling [`retract`](@ref) without
specifying the retraction method. By default, this is the [`ExponentialRetraction`](@ref).
"""
default_retraction_method(::AbstractManifold) = ExponentialRetraction()

"""
    inverse_retract!(M::AbstractManifold, X, p, q[, method::AbstractInverseRetractionMethod])

Compute the inverse retraction, a cheaper, approximate version of the
[`log`](@ref)arithmic map), of points `p` and `q` on the [`AbstractManifold`](@ref) `M`.
Result is saved to `X`.

Inverse retraction method can be specified by the last argument, defaulting to
[`default_inverse_retraction_method`](@ref)`(M)`. See the documentation of respective manifolds for
available methods.

See also [`retract!`](@ref).
"""
function inverse_retract!(M::AbstractManifold, X, p, q)
    return inverse_retract!(M, X, p, q, default_inverse_retraction_method(M))
end
function inverse_retract!(
    M::AbstractManifold,
    X,
    p,
    q,
    method::LogarithmicInverseRetraction,
)
    return log!(M, X, p, q)
end
function inverse_retract!(
    M::AbstractManifold,
    X,
    p,
    q,
    method::AbstractInverseRetractionMethod,
)
    return error(
        manifold_function_not_implemented_message(M, inverse_retract!, X, p, q, method),
    )
end

"""
    inverse_retract(M::AbstractManifold, p, q)
    inverse_retract(M::AbstractManifold, p, q, method::AbstractInverseRetractionMethod

Compute the inverse retraction, a cheaper, approximate version of the
[`log`](@ref)arithmic map), of points `p` and `q` on the [`AbstractManifold`](@ref) `M`.

Inverse retraction method can be specified by the last argument, defaulting to
[`default_inverse_retraction_method`](@ref)`(M)`.
For available inverse retractions on certain manifolds see the documentation on the
corresponding manifold.

See also [`retract`](@ref).
"""
function inverse_retract(M::AbstractManifold, p, q)
    X = allocate_result(M, inverse_retract, p, q)
    inverse_retract!(M, X, p, q)
    return X
end
function inverse_retract(M::AbstractManifold, p, q, method::AbstractInverseRetractionMethod)
    X = allocate_result(M, inverse_retract, p, q)
    inverse_retract!(M, X, p, q, method)
    return X
end

@doc raw"""
    retract(M::AbstractManifold, p, X)
    retract(M::AbstractManifold, p, X, t::Real=1)
    retract(M::AbstractManifold, p, X, method::AbstractRetractionMethod)
    retract(M::AbstractManifold, p, X, t::Real=1, method::AbstractRetractionMethod)

Compute a retraction, a cheaper, approximate version of the [`exp`](@ref)onential map,
from `p` into direction `X`, scaled by `t`, on the [`AbstractManifold`](@ref) `M`.

A retraction ``\operatorname{retr}_p: T_p\mathcal M → \mathcal M`` is a smooth map that fulfills

1. ``\operatorname{retr}_p(0) = p``
2. ``D\operatorname{retr}_p(0): T_p\mathcal M \to T_p\mathcal M`` is the identity map, i.e. ``D\operatorname{retr}_p(0)[X]=X``,

where ``D\operatorname{retr}_p`` denotes the differential of the retraction

The retraction is called of second order if for all ``X`` the curves ``c(t) = R_p(tX)``
have a zero acceleration at ``t=0``, i.e. ``c''(0) = 0``.

Retraction method can be specified by the last argument, defaulting to
[`default_retraction_method`](@ref)`(M)`. For further available retractions see the documentation of respective manifolds.

Locally, the retraction is invertible. For the inverse operation, see [`inverse_retract`](@ref).
"""
function retract(M::AbstractManifold, p, X)
    q = allocate_result(M, retract, p, X)
    retract!(M, q, p, X)
    return q
end
retract(M::AbstractManifold, p, X, t::Real) = retract(M, p, t * X)
function retract(M::AbstractManifold, p, X, method::AbstractRetractionMethod)
    q = allocate_result(M, retract, p, X)
    retract!(M, q, p, X, method)
    return q
end
function retract(M::AbstractManifold, p, X, t::Real, method::AbstractRetractionMethod)
    return retract(M, p, t * X, method)
end

"""
    retract!(M::AbstractManifold, q, p, X)
    retract!(M::AbstractManifold, q, p, X, t::Real=1)
    retract!(M::AbstractManifold, q, p, X, method::AbstractRetractionMethod)
    retract!(M::AbstractManifold, q, p, X, t::Real=1, method::AbstractRetractionMethod)

Compute a retraction, a cheaper, approximate version of the [`exp`](@ref)onential map,
from `p` into direction `X`, scaled by `t`, on the [`AbstractManifold`](@ref) manifold `M`.
Result is saved to `q`.

Retraction method can be specified by the last argument, defaulting to
[`default_retraction_method`](@ref)`(M)`. See the documentation of respective manifolds for available
methods.

See [`retract`](@ref) for more details.
"""
retract!(M::AbstractManifold, q, p, X) = retract!(M, q, p, X, default_retraction_method(M))
retract!(M::AbstractManifold, q, p, X, t::Real) = retract!(M, q, p, t * X)
retract!(M::AbstractManifold, q, p, X, ::ExponentialRetraction) = exp!(M, q, p, X)
function retract!(M::AbstractManifold, q, p, X, t::Real, method::AbstractRetractionMethod)
    return retract!(M, q, p, t * X, method)
end
function retract!(M::AbstractManifold, q, p, X, method::AbstractRetractionMethod)
    return error(manifold_function_not_implemented_message(M, retract!, q, p, method))
end


# These examples are documented by the LCP notebooks from
# https://github.com/QuantEcon/SimpleDifferentialOperators.jl
# In general, these are optimal stopping problems solved
# as linear complementarity problems.

using Parameters, DiffEqOperators, Random, Suppressor, Test, PATHSolver

# Seed
Random.seed!(1793)

function LCP_split(S_)
  # setup
    u1(x) = x^2
    ρ = 0.75
    μ1(x) = -x
    σ = 0.1
    M = 300
    x = range(-5., 5., length = M)

    # construct operator and non-operator LCP objects
    BC = RobinBC((0., 1., 0.), (0., 1., 0.), x[2] - x[1], 1)
    L1 = UpwindDifference(1, 1, x[2] - x[1], M, map(t -> μ1(t), x)) * BC
    L2 = σ^2 / 2 * CenteredDifference(2, 2, x[2] - x[1], M) * BC
    u_vec = u1.(x)

    return (L1 = Array(L1)[1], L2 = Array(L2)[1], S = S_ .* ones(M), u = u_vec, ρ = ρ)
end

@testset "Split Case, Nonzero S" begin
  # setup
  @unpack L1, L2, S, u, ρ = LCP_split(0.125)
  g = z -> ρ .* z - L1 * z - L2 * z -
      u + ρ .* S - L1 * S - L2 * S

  n = 300
  lb = zeros(n)
  ub = 300. .* ones(n) # a reasonable guess?
  options(convergence_tolerance = 1e-14, output = :no,
          time_limit = 600) # 10 minute budget
  exit_code, sol_z, sol_g = @suppress solveLCP(g, lb, ub);
  # tests
  @test exit_code == :StationaryPointFound
  @test sol_z[3] ≈ 8.774937148286833 atol = 1e-5
  @test sol_z[120] ≈ 0.3021168981772892 atol = 1e-5
  @test sol_z[270] ≈ 5.729541001488482 atol = 1e-5
  @test sol_g[end-1] ≈ 3.197442310920451e-14 atol = 1e-5
end

@testset "Split Case, Zero S" begin
  # setup
  @unpack L1, L2, S, u, ρ = LCP_split(0.)
  f = z -> ρ .* z - L1 * z - L2 * z -
      u + ρ .* S - L1 * S - L2 * S

  # @unpack L, B, S, q = LCP_split(0.)
  # f = z -> B*z + q
  n = 300
  lb = zeros(n)
  ub = 300*ones(n) # a reasonable guess?
  options(convergence_tolerance = 1e-12, output = :no,
          time_limit = 600) # 10 minute budget
  exit_code, sol_z, sol_f = @suppress solveLCP(f, lb, ub);

  # tests
  @test exit_code == :StationaryPointFound
  @test sol_z[3] ≈ 8.888461348456772 atol = 1e-5
  @test sol_z[76] ≈ 2.279767804635279 atol = 1e-5
  @test sol_z[150] ≈ 0.005770703117189083 atol = 1e-5
  @test sol_z[269] ≈ 5.744079221450249 atol = 1e-5
  @test sol_f[123] ≈ 1.7763568394002505e-15 atol = 1e-5
end

nothing


"""
Attitude objects contains Euler angles (rad) and quaternions.

# Constructors
    Attitude(psi, theta, phi)
    Attitude(q0, q1, q2, q3)
"""
struct Attitude
    # Euler angles
    psi::Number
    theta::Number
    phi::Number
    # Quaternions
    q0::Number
    q1::Number
    q2::Number
    q3::Number
end

# Constructor from Euler angles
Attitude(psi, theta, phi) = Attitude(psi, theta, phi,
                                     euler2quaternion(psi, theta, phi)...)

# Constructor from quaternions
Attitude(q0, q1, q2, q3) = Attitude(quaternion2euler(q0, q1, q2, q3)...,
                                    q0, q1, q2, q3)

# Getters
"""
    get_euler_angles(att::Attitude)

Return Euler angles ψ, θ, ϕ (yaw, pith, roll).
"""
get_euler_angles(att::Attitude) = [att.psi, att.theta, att.phi]

"""
    get_quaternions(att::Attitude)

Return quaternions.
"""
get_quaternions(att::Attitude) = [att.q0, att.q1, att.q2, att.q3]


module HyperbolicCrossApprox

import SmolyakApprox: normalize_node,
                      scale_nodes,
                      combine_nodes

import ChebyshevApprox: chebyshev_polynomial,
                        derivative_of_chebyshev_polynomial,
                        chebyshev_nodes,
                        chebyshev_extrema,
                        chebyshev_extended,
                        vertesi_nodes

using ThreadPools

include("hyperbolic_cross_weights.jl")
include("hyperbolic_cross_evaluate.jl")
include("hyperbolic_cross_grid.jl")
include("generate_multi_index.jl")
include("hyperbolic_cross_derivative.jl")
include("hyperbolic_cross_gradient.jl")
include("determine_grid_size.jl")
include("hyperbolic_cross_inverse_interpolation_matrix.jl")

export chebyshev_nodes,
       chebyshev_extrema,
       chebyshev_extended,
       vertesi_nodes,
       hyperbolic_cross_grid,
       hyperbolic_cross_weights,
       hyperbolic_cross_weights_threaded,
       hyperbolic_cross_inverse_interpolation_matrix,
       hyperbolic_cross_inverse_interpolation_matrix_threaded,
       hyperbolic_cross_polynomial,
       hyperbolic_cross_evaluate,
       hyperbolic_cross_derivative,
       hyperbolic_cross_gradient

end


using CompScienceMeshes
using Test

m = meshrectangle(1.0,1.0,1.0,3)
p = chart(m, cells(m,i))

# bouding box of an array of points
c, s = boundingbox(p.vertices)
@test isa(c, eltype(p.vertices))
@test isa(s, Number)

# bouding box of a point
c, s = boundingbox(p.vertices[1])
@test c == p.vertices[1]
@test s == 0


#=
    The file contains rational reconstrction and CRT functions
=#

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

# Rational number reconstruction implementation borrowed from CLUE
# and modified a bit to suit the 'Modern Computer Algebra' definitions
# Returns a rational r // h of QQ field in a canonical form such that
#   r // h ≡ a (mod m)
#
# let n = max( λ(a), λ(m) ) , where λ(x) is a number of bits for x
# O(n^2)
function rational_reconstruction(a::I, m::I) where {I<:Union{Int, BigInt}}
    a = mod(a, m)
    if a == 0 || m == 0
        return QQ(0, 1)
    end
    if m < 0
        m = -m
    end
    if a < 0
        a = m - a
    end
    if a == 1
        return QQ(1, 1)
    end
    bnd = sqrt(float(m) / 2)

    @debug "" bnd
    U = (I(1), I(0), m)
    V = (I(0), I(1), a)
    while abs(V[3]) >= bnd
        q = div(U[3], V[3])
        T = U .- q .* V
        U = V
        V = T
    end

    t = abs(V[2])
    r = V[3] * sign(V[2])

    @debug "" r t

    # changed from `<= bnd` to `<= m / bnd`
    # we can speed up this !
    # if t <= bnd && gcd(r, t) == 1
    #    return QQ(r, t)
    # end

    return QQ(r, t)

    # throw(DomainError(
    #    :($a//$m), "rational reconstruction of $a (mod $m) does not exist"
    # ))

    # return QQ(0, 1)
end

# returns the bound for rational reconstruction
# based on the current modulo
#
# The bound for rational reconstrction:
# as soon as the numerator in rational reconstruction
# exceeds this bound, the gcd iteration is stopped
function rational_reconstruction_bound(modulo::BigInt)
    setprecision(2*Base.GMP.MPZ.sizeinbase(modulo, 2)) do
        ceil(BigInt, sqrt(BigFloat(modulo) / 2))
    end
end

# 0 allocations!
#=
    Computes the rational reconstruction of `a` mod `m`.
    Namely, a pair of numbers num, den , such that
        num//den ≡ a (mod m)

    Additional params:
        bnd  = stores stopping criterion threshold
        buf, buf1, buf2, buf3  = buffers
        u1, u2, u3 = buffers
        v1, v2, v3 = buffers
=#
function rational_reconstruction!(
            num::BigInt, den::BigInt, bnd::BigInt, buf::BigInt,
            buf1::BigInt, buf2::BigInt, buf3::BigInt,
            u1::BigInt, u2::BigInt, u3::BigInt,
            v1::BigInt, v2::BigInt, v3::BigInt,
            a::BigInt, m::BigInt)

    if Base.GMP.MPZ.cmp_ui(a, 0) == 0
        Base.GMP.MPZ.set_ui!(num, 0)
        Base.GMP.MPZ.set_ui!(den, 1)
        return true
    end

    # assumes input is nonnegative
    @assert Base.GMP.MPZ.cmp_ui(a, 0) > 0

    if Base.GMP.MPZ.cmp_ui(a, 1) == 0
        Base.GMP.MPZ.set_ui!(num, 1)
        Base.GMP.MPZ.set_ui!(den, 1)
        return true
    end

    Base.GMP.MPZ.set_ui!(u1, 1)
    Base.GMP.MPZ.set_ui!(u2, 0)
    Base.GMP.MPZ.set!(u3, m)
    Base.GMP.MPZ.set_ui!(v1, 0)
    Base.GMP.MPZ.set_ui!(v2, 1)
    Base.GMP.MPZ.set!(v3, a)

    while true
        if Base.GMP.MPZ.cmp(v2, bnd) > 0
            # @error "" v2 bnd
            return false
        end

        Base.GMP.MPZ.set!(buf, v3)
        if Base.GMP.MPZ.cmp_ui(buf, 0) < 0
            Base.GMP.MPZ.neg!(buf)
        end

        if Base.GMP.MPZ.cmp(buf, bnd) < 0
            break
        end

        Base.GMP.MPZ.tdiv_q!(buf, u3, v3)

        Base.GMP.MPZ.mul!(buf1, buf, v1)
        Base.GMP.MPZ.mul!(buf2, buf, v2)
        Base.GMP.MPZ.mul!(buf3, buf, v3)

        Base.GMP.MPZ.sub!(buf1, u1, buf1)
        Base.GMP.MPZ.sub!(buf2, u2, buf2)
        Base.GMP.MPZ.sub!(buf3, u3, buf3)

        Base.GMP.MPZ.set!(u1, v1)
        Base.GMP.MPZ.set!(u2, v2)
        Base.GMP.MPZ.set!(u3, v3)

        Base.GMP.MPZ.set!(v1, buf1)
        Base.GMP.MPZ.set!(v2, buf2)
        Base.GMP.MPZ.set!(v3, buf3)
    end

    Base.GMP.MPZ.set!(den, v2)
    Base.GMP.MPZ.set!(num, v3)

    #=
    num and den are obviously coprime,
    gcd is not nessesary

    Base.GMP.MPZ.gcd!(buf, den, num)
    Base.GMP.MPZ.tdiv_q!(den, buf)
    Base.GMP.MPZ.tdiv_q!(num, buf)
    =#

    if Base.GMP.MPZ.cmp_ui(den, 0) < 0
        Base.GMP.MPZ.neg!(den)
        Base.GMP.MPZ.neg!(num)
    end

    true
end

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

# 0 allocations!
#=
    Computes the unique x s.t
        x ≡ a1 mod m1
        x ≡ a2 mod m2

    Additional params:
        M   = m1*m2
        buf = buffer
        n1  = buffer
        n2  = buffer
=#
function CRT!(
            M::BigInt, buf::BigInt, n1::BigInt, n2::BigInt,
            a1::BigInt, minv1::BigInt, a2::UInt, minv2::BigInt,
            m1::BigInt, m2::BigInt)

    Base.GMP.MPZ.mul!(buf, m1, minv1)
    Base.GMP.MPZ.mul_ui!(n1, buf, a2)

    Base.GMP.MPZ.mul!(buf, m2, minv2)
    Base.GMP.MPZ.mul!(n2, buf, a1)

    Base.GMP.MPZ.add!(buf, n1, n2)
    Base.GMP.MPZ.fdiv_r!(buf, M)
end

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


using Test

using CompScienceMeshes

e0 = point(0,0,0)
e1 = point(1,0,0)
e2 = point(0,1,0)
e3 = point(0,0,1)

p = simplex(e0, e1)
q = simplex((e0+e1)/2, (-e0+3*e1)/2)

pq = intersection(p,q)
@test length(pq) == 1
pq = pq[1]
@test volume(pq) == 1/2
@test pq.vertices[1] == q.vertices[1]
@test pq.vertices[2] == p.vertices[2]
@test 2*pq.tangents[1] == p.tangents[1] == q.tangents[1]

c = (e0+e1+e2)/3
p = simplex(e0, e1, e2)
q = simplex(e0, e1, c)

qq = intersection(q,q)
pq = intersection(p,q)
qp = intersection(q,p)

# For these case where points of the subject are on the clipper or
# the other way around, the precise computationof the intersection
# is subject to numerical rounding errors. For efficiency reasons
# this something that needs fixing. It should not affect the computation
# of the interaction elements, however.
@test sum([volume(r) for r in qq]) ≈ volume(q)
@test sum([volume(r) for r in pq]) ≈ volume(q)
@test sum([volume(r) for r in qp]) ≈ volume(q)

# This case caused NaNs in the RWG BC Gram matrix. Turned out to be caused
# by computing the intersection between two parallel lines. This needs a
# more robust solution in the long run.
p = point(-0.819550109062379,-0.394674531365384,-0.415415013001886)
q = point(-0.66680743687093,-0.618706884644307,-0.415415013001886)
r = point(-0.832690784854388,-0.535137873473586,-0.142314838273285)

a = point(-0.66680743687093,-0.618706884644307,-0.415415013001886)
b = point(-0.749749110862659,-0.5769223790589465,-0.27886492563758547)
c = point(-0.7730161102625658,-0.5161730964944257,-0.32438162142568566)

isct1 = CompScienceMeshes.sutherlandhodgman([p,q,r],[a,b,c])

@test isct1 == [
    point(-0.7730161102625658,-0.5161730964944257,-0.32438162142568566),
    point(-0.66680743687093,-0.618706884644307,-0.415415013001886),
    point(-0.66680743687093,-0.618706884644307,-0.415415013001886),
    point(-0.749749110862659,-0.5769223790589465,-0.27886492563758547)]

splx1 = simplex(
    point(0,0,0),
    point(1,0,0),
    point(0,1,0))
splx2 = simplex(
    point(0,-1,0),
    point(1,1,0),
    point(0.5,1,0))

# Verify that all parts of the intersection are oriented consistently
# and equal to the orientation of the inputs
isct2 = intersection(splx1, splx2)
@test length(isct2) == 2

for splx in isct2
    @test normal(splx) ≈ point(0,0,1)
end

# Repeat this test but now with the orientation of the second
# operand flipped
splx1 = simplex(
    point(0,0,0),
    point(1,0,0),
    point(0,1,0))
splx2 = simplex(
    point(1,1,0),
    point(0,-1,0),
    point(0.5,1,0))

isct2 = intersection(splx1, splx2)
@test length(isct2) == 2

for splx in isct2
    @test normal(splx) ≈ point(0,0,1)
end

# Repeat this test but now with the orientation of the first operand flipped
splx1 = simplex(
    point(1,0,0),
    point(0,0,0),
    point(0,1,0))
splx2 = simplex(
    point(0,-1,0),
    point(1,1,0),
    point(0.5,1,0))

isct2 = intersection(splx1, splx2)
@test length(isct2) == 2

for splx in isct2
    @test normal(splx) ≈ point(0,0,-1)
end


function n_depth_increments(depths)
  difference = depths[2:1:end] - depths[1:1:end-1]

  n_increments = 0
  for diff in difference
    if diff > 0
      n_increments += 1
    end
  end
  return n_increments
end

function three_measurement_window(depths)
  window = depths[1:1:end-2] + depths[2:1:end-1] + depths[3:1:end]
  difference = window[2:1:end] - window[1:1:end-1]

  n_increments = 0
  for diff in difference
    if diff > 0
      n_increments += 1
    end
  end
  return n_increments
end

# Test case
# Each entry is the sea floor depth in increasing distance from sub.
test_input = [199, 200, 208, 210, 200, 207, 240, 269, 260, 263]
n_test = n_depth_increments(test_input)
println("Number of increments in the test case: $n_test")

n_test_three_window = three_measurement_window(test_input)
println("Number of increments in the three window test case: $n_test_three_window")

# Application to actual input.
open("input.txt", "r") do f
  data = [parse(Int, s) for s in readlines(f)]
  n_increments = n_depth_increments(data)
  println("Number of increments: $n_increments")

  n_three_window = three_measurement_window(data)
  println("Number of increments in the three window case: $n_three_window")
end


export create_curl

# Wrapper to create the discrete curl by default
create_curl(isfwd::AbsVecBool,  # isfwd[w] = true|false: ∂w is forward|backward difference
            N::AbsVecInteger,  # size of grid
            ∆l⁻¹::NTuple{K,Number}=ntuple(k->1.0,length(isfwd)),  # ∆l⁻¹[w]: inverse of uniform distance between grid planes in x-direction
            isbloch::AbsVecBool=fill(true,K),  # boundary conditions in x, y, z
            e⁻ⁱᵏᴸ::AbsVecNumber=ones(K);  # Bloch phase factor in x, y, z
            cmp_shp::AbsVecInteger=1:K,
            cmp_out::AbsVecInteger=1:K,
            cmp_in::AbsVecInteger=1:K,
            order_cmpfirst::Bool=true  # true to use Cartesian-component-major ordering for more tightly banded matrix
            ) where {K} =
    # I should not cast e⁻ⁱᵏᴸ into a complex vector, because then the entire curl matrix
    # becomes a complex matrix.  Sometimes I want to keep it real (e.g., when no PML and
    # Bloch phase factors are used).  In fact, this is the reason why I accept e⁻ⁱᵏᴸ instead
    # of constructing it from k and L as exp.(-im .* k .* L), which is always complex even
    # if k = 0.
    #
    # I should not cast ∆l⁻¹ to a vector of any specific type (e.g., Float, ComplexF), either,
    # because sometimes I would want to even create an integral curl operator.
    create_curl(isfwd, fill.(∆l⁻¹,(N...,)), isbloch, e⁻ⁱᵏᴸ; cmp_shp, cmp_out, cmp_in, order_cmpfirst)

# Wrapper to convert AbstractVector's to SVec's
create_curl(isfwd::AbsVecBool,  # isfwd[w] = true|false: ∂w is forward|backward difference
            ∆l⁻¹::NTuple{K,AbsVecNumber},  # ∆l⁻¹[w]: inverse of uniform distance between grid planes in x-direction
            isbloch::AbsVecBool=fill(true,K),  # boundary conditions in x, y, z
            e⁻ⁱᵏᴸ::AbsVecNumber=ones(K);  # Bloch phase factor in x, y, z
            cmp_shp::AbsVecInteger=1:K,
            cmp_out::AbsVecInteger=1:K,
            cmp_in::AbsVecInteger=1:K,
            order_cmpfirst::Bool=true  # true to use Cartesian-component-major ordering for more tightly banded matrix
            ) where {K} =
    (Kout = length(cmp_out); Kin = length(cmp_in);
     create_curl(SVec{K}(isfwd), ∆l⁻¹, SVec{K}(isbloch), SVec{K}(e⁻ⁱᵏᴸ);
                 cmp_shp=SInt{K}(cmp_shp), cmp_out=SInt{Kout}(cmp_out), cmp_in=SInt{Kin}(cmp_in), order_cmpfirst))

# ∇×F = x̂ (∂y Fz - ∂z Fy) + ŷ (∂z Fx - ∂x Fz) + ẑ (∂x Fy - ∂y Fx)
#
# In matrix form,
# ⎡ 0  -∂z  ∂y⎤
# ⎢ ∂z  0  -∂x⎥
# ⎣-∂y  ∂x  0 ⎦
#
# CURL_BLK takes only the signs of the blocks:
const CURL_BLK = S²Int{3,9}(0,  1, -1,  # 1st column (not row)
                           -1,  0,  1,  # 2nd column (not row)
                            1, -1,  0)  # 3rd column (not row)

function create_curl(isfwd::SBool{K},  # isfwd[w] = true|false: ∂w is forward|backward difference
                     ∆l⁻¹::NTuple{K,AbsVecNumber},  # ∆l⁻¹[w]: inverse of distances between grid planes in w-direction
                     isbloch::SBool{K},  # isblock[w]: boundary conditions in w-direction
                     e⁻ⁱᵏᴸ::SNumber{K};  # e⁻ⁱᵏᴸ[w]: Bloch phase factor in w-direction
                     cmp_shp::SInt{K}=SVec(ntuple(identity, Val(K))),
                     cmp_out::SInt{Kout}=SVec(ntuple(identity, Val(K))),
                     cmp_in::SInt{Kin}=SVec(ntuple(identity, Val(K))),
                     order_cmpfirst::Bool=true  # true to use Cartesian-component-major ordering for more tightly banded matrix
                     ) where {K,Kout,Kin}  # {space dimension, output field dimension, input field dimension}
    T = promote_type(eltype.(∆l⁻¹)..., eltype(e⁻ⁱᵏᴸ))  # eltype(eltype(∆l⁻¹)) can be Any if ∆l⁻¹ is inhomogeneous
    N = SVec(length.(∆l⁻¹))
    M = prod(N)

    num_blk = sum(abs, CURL_BLK[cmp_out,cmp_in])  # no allocations

    # Below, create_∂info() is called num_blk times in the double for loops, and each call
    # returns I, J, V of length 2M, so we preallocate Itot, Jtot, Vtot of length num_blk⋅2M.
    Itot = VecInt(undef, num_blk * 2M)
    Jtot = VecInt(undef, num_blk * 2M)
    Vtot = Vector{T}(undef, num_blk * 2M)

    ind_blk = 0  # index of matrix block
    for ind_nv = 1:Kout
        nv = cmp_out[ind_nv]  # Cartesian compotent of output field
        istr, ioff = order_cmpfirst ? (Kout, ind_nv-Kout) : (1, M*(ind_nv-1))  # (row stride, row offset)
        for ind_nu = 1:Kin
            nu = cmp_in[ind_nu]  # Cantesian component of input field

            parity = CURL_BLK[nv,nu]
            iszero(parity) && continue  # if zero block, skip to next nu

            jstr, joff = order_cmpfirst ? (Kin, ind_nu-Kin) : (1, M*(ind_nu-1))  # (column stride, column offset)

            nw = 6 - nv - nu  # direction of differentiation; 6 = 1 + 2 + 3
            is_nw = cmp_shp.==nw
            sum(is_nw)==1 || @error "cmp_shp = $cmp_shp does not have one and only one nw = $nw"
            ind_nw = findfirst(is_nw)
            I, J, V = create_∂info(ind_nw, isfwd[ind_nw], N, ∆l⁻¹[ind_nw], isbloch[ind_nw], e⁻ⁱᵏᴸ[ind_nw])

            @. I = istr * I + ioff
            @. J = jstr * J + joff
            V .*= parity

            # For some reason, using .= below is slower because it uses 1 allocatiotn.  On
            # the other hand, using = does not use allocation and therefore faster.
            indₛ, indₑ = ind_blk*2M + 1, (ind_blk+1)*2M
            Itot[indₛ:indₑ] = I
            Jtot[indₛ:indₑ] = J
            Vtot[indₛ:indₑ] = V
            ind_blk += 1
        end
    end

    return dropzeros!(sparse(Itot, Jtot, Vtot, Kout*M, Kin*M))
end


const TensorSubGrid = ProductGrid{<:NTuple{N,GRID}} where N where {GRID<:AbstractSubGrid}
function tensorproductbitarray(vectors::Union{BitVector,Vector{Bool}}...)
	N = length(vectors)
	R = falses(map(length, vectors))
	for i in CartesianIndices(size(R))
		R[i] = reduce(&,map(k->vectors[k][i.I[k]],1:length(vectors)))
	end
	R
end

mask(grid::TensorSubGrid) = tensorproductbitarray(map(mask, elements(grid))...)
subindices(grid::TensorSubGrid) = findall(mask(grid))
supergrid(grid::TensorSubGrid) = ProductGrid(map(supergrid, elements(grid))...)
issubindex(i, g::TensorSubGrid) = all(map(issubindex, i, elements(g)))
issubindex(i::CartesianIndex, g::TensorSubGrid) = issubindex(i.I, g)


# Copyright 2021 Gravifer
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#=
  Author: Gravifer, X. H. Zhu, H. J. Yan, H. Y. Sun
  Date: 2021-04-10
  Description: TBD
=#

module transfinite_interpolate

using StaticArrays
using OffsetArrays
include("../../src/misc-util.jl")
using .__CFD2021__misc_util__: tuplejoin

"""
    transfinite_interpolate_2d(f_lo, f_hi[, v_cr])

Return a interpolating function defined on ``[0, 1]²``.

**Notice**: does not check or warn when extrapolating.

See also: https://en.wikipedia.org/wiki/Transfinite_interpolation

TODO: generalize to higher dimensions.

# Arguments
- `f_lo::Tuple{Function, Function}`: a pair of functions.
- `f_hi::Tuple{Function, Function}`: a pair of functions.
- `v_cr::SMatrix{2, 2} = nothing`: (optional) values to take at the corners;
    use average of edge functions if not provided.
    Still, in practical use you should make sure the ends of the edges meet.

# Output
- `interpolating_function(u::NTuple{2, T} where T <: Number):: Number`:
    Take in a length-2 static vector of numbers and return a number.

# Examples
```jldoctest
julia> interp_func = transfinite_interpolate_2d( (x->2x, y->3y), (x->4x+3, y->2+5y) )
(::var"#interpolating_function#1"{Tuple{var"#2#6", var"#3#7"}, Tuple{var"#4#8", var"#5#9"}}) (generic function with 1 method)

julia> interp_func((0.5,0.5))
3.0
```
"""
function transfinite_interpolate_2d(f_lo::Tuple{Function,Function},
                                    f_hi::Tuple{Function,Function};
                                    v_cr::Union{Nothing,SMatrix{2,2}}=nothing)::Function
    #= handling of optionals suggested by
       https://stackoverflow.com/questions/42499528/julia-convention-for-optional-arguments;
       though Base.Void is deprecated and now called Nothing instead, as pointed out by
       https://discourse.julialang.org/t/void-not-defined-version-problem/14183 =#
    if typeof(v_cr) <: Nothing
        v_cr = SA[(f_lo[1](0) + f_lo[2](0))/2 (f_lo[1](1) + f_hi[2](0))/2
                  (f_hi[1](0) + f_lo[2](1))/2 (f_hi[1](1) + f_hi[2](1))/2]
    end

    function interpolating_function(u::NTuple{2,T} where {T<:Number})::Number
        #= linear-part from edges   bilinear-part from vertices =#
        (1-u[2])*f_lo[1](u[1]) - (1-u[1])*(1-u[2])*v_cr[1,1] +
        (  u[2])*f_hi[1](u[1]) - (  u[1])*(1-u[2])*v_cr[1,2] +
        (1-u[1])*f_lo[2](u[2]) - (1-u[1])*(  u[2])*v_cr[2,1] +
        (  u[1])*f_hi[2](u[2]) - (  u[1])*(  u[2])*v_cr[2,2]
    end
    return interpolating_function
end

"""
    transfinite_interpolate_2d(v_lo, v_hi[, v_cr])
    transfinite_interpolate_2d!(interpolated_array, v_lo, v_hi[, v_cr])

Return a interpolated array defined on a discretization of ``[0, 1]²``.

**Notice**: the in-place version does not check dimension-compatibility of its inputs.

TODO: generalize to higher dimensions.

# Arguments
- `v_lo::Tuple{Vector{<: Real}, Vector{<: Real}}`: a pair of vectors.
- `v_hi::Tuple{Vector{<: Real}, Vector{<: Real}}`: a pair of vectors.
- `v_cr::SMatrix{2, 2, <: Real} = nothing`: (optional) values to take at the corners;
    use average of edge values if not provided.
    Still, in practical use you should make sure the ends of the edges meet.

# Output
- `interpolated_array::Array{Real, 2}`

# Examples
```jldoctest
julia> transfinite_interpolate_2d(([0,0.5,1], [0,0.5,1]), ([1,1.5,2], [1,1.5,2]))
3×3 Matrix{Real}:
 0.0  0.5  1.0
 0.5  1.0  1.5
 1.0  1.5  2.0
```
"""
function transfinite_interpolate_2d( # for regular Arrays
    v_lo::Tuple{Vector{T}, Vector{S}} where {T <: Real, S <: Real},
    v_hi::Tuple{Vector{T}, Vector{S}} where {T <: Real, S <: Real},
    v_cr::Union{Nothing, SMatrix{2, 2, T} where T <: Real} = nothing)::Array{Real, 2}

    #= check dimension-compatibility before allocating anything =#
    length(v_lo) == length(v_hi) || throw(
        DimensionMismatch(" incompatible grid specification ($(length(v_lo))-d at one end but $(length(v_hi))-d at the other)."))
    N = length(v_lo); errmsg = "" #= In this function, dimensions are always indexed from 1, but not so for each axis of them =#
    for dim in eachindex(v_lo)  if ! (axes(v_lo[dim]) == axes(v_hi[dim]))
        errmsg *= "grid mismatch in dimension $(dim) ($(axes(v_lo[dim])) at one end but $(axes(v_hi[dim])) at the other)."
    end end; errmsg == "" || throw(DimensionMismatch(errmsg))

    interpolated_array = Array{Float64}(undef, (length(v) for v in v_lo)...)
    transfinite_interpolate_2d!(interpolated_array, v_lo, v_hi, v_cr)
    return interpolated_array
end
function transfinite_interpolate_2d( # for OffsetArrays
    v_lo::Tuple{OffsetArray, OffsetArray},
    v_hi::Tuple{OffsetArray, OffsetArray},
    v_cr::Union{Nothing, SMatrix{2, 2, T} where T <: Real} = nothing)::OffsetArray

    #= check dimension-compatibility before allocating anything =#
    length(v_lo) == length(v_hi) || throw(
        DimensionMismatch(" incompatible grid specification ($(length(v_lo))-d at one end but $(length(v_hi))-d at the other)."))
    N = length(v_lo); errmsg = "" #= In this function, dimensions are always indexed from 1, but not so for each axis of them =#
    for dim in eachindex(v_lo)  if ! (axes(v_lo[dim]) == axes(v_hi[dim]))
        errmsg *= "grid mismatch in dimension $(dim) ($(axes(v_lo[dim])) at one end but $(axes(v_hi[dim])) at the other)."
    end end; errmsg == "" || throw(DimensionMismatch(errmsg))

    interpolated_array = Array{Float64}(undef, (length(v) for v in v_lo)...)
    interpolated_array = OffsetArray(interpolated_array, tuplejoin((axes(v) for v in v_lo)...))
    transfinite_interpolate_2d!(interpolated_array, v_lo, v_hi, v_cr)
    return interpolated_array
end
function transfinite_interpolate_2d!( # for regular Arrays
    interpolated_array::Array{T, 2}   where  T <: Real,
    v_lo::Tuple{Vector{T}, Vector{S}} where {T <: Real, S <: Real},
    v_hi::Tuple{Vector{T}, Vector{S}} where {T <: Real, S <: Real},
    v_cr::Union{Nothing, SMatrix{2, 2, T} where T <: Real} = nothing)

    if typeof(v_cr) <: Nothing # e.g, the following lines still assumes range ``1:2``.
        v_cr = SA[(v_lo[1][begin]+v_lo[2][begin])/2 (v_lo[1][ end ]+v_hi[2][begin])/2
                  (v_hi[1][begin]+v_lo[2][ end ])/2 (v_hi[1][ end ]+v_hi[2][ end ])/2]
    end; M = [length(v)-1 for v in v_lo] # this is only for scaling, and never for iterating

    # loop Version
    for u in CartesianIndices(interpolated_array) #= #= Deprecated =# Iterators.product(IndexCartesian(), [eachindex(v) for v in v_lo]...) =#
        interpolated_array[u] =
        #= linear-part   from   edges      bilinear-part   from   vertices    =#
        (1-(u[2]-1)/M[2])*v_lo[1][u[1]] - (1-(u[1]-1)/M[1])*(1-(u[2]-1)/M[2])*v_cr[1,1] +
        (  (u[2]-1)/M[2])*v_hi[1][u[1]] - (  (u[1]-1)/M[1])*(1-(u[2]-1)/M[2])*v_cr[1,2] +
        (1-(u[1]-1)/M[1])*v_lo[2][u[2]] - (1-(u[1]-1)/M[1])*(  (u[2]-1)/M[2])*v_cr[2,1] +
        (  (u[1]-1)/M[1])*v_hi[2][u[2]] - (  (u[1]-1)/M[1])*(  (u[2]-1)/M[2])*v_cr[2,2]
    end

    return interpolated_array
end
function transfinite_interpolate_2d!(interpolated_array::OffsetArray,
                                     v_lo::Tuple{OffsetArray,OffsetArray},
                                     v_hi::Tuple{OffsetArray,OffsetArray},
                                     v_cr::Union{Nothing,SMatrix{2,2,T} where T<:Real}=nothing)
    if typeof(v_cr) <: Nothing # e.g, the following lines still assumes range ``1:2``.
        v_cr = SA[(v_lo[1][begin]+v_lo[2][begin])/2 (v_lo[1][ end ]+v_hi[2][begin])/2
                  (v_hi[1][begin]+v_lo[2][ end ])/2 (v_hi[1][ end ]+v_hi[2][ end ])/2]
    end; M = [length(v)-1 for v in v_lo] # this is only for scaling, and never for iterating

    # loop Version
    for u in CartesianIndices(interpolated_array) #= #= Deprecated =# Iterators.product(IndexCartesian(), [eachindex(v) for v in v_lo]...) =#
        interpolated_array[u] =
        #= linear-part   from   edges      bilinear-part   from   vertices    =#
        (1-u[2]/M[2])*v_lo[1][u[1]] - (1-u[1]/M[1])*(1-u[2]/M[2])*v_cr[1,1] +
        (  u[2]/M[2])*v_hi[1][u[1]] - (  u[1]/M[1])*(1-u[2]/M[2])*v_cr[1,2] +
        (1-u[1]/M[1])*v_lo[2][u[2]] - (1-u[1]/M[1])*(  u[2]/M[2])*v_cr[2,1] +
        (  u[1]/M[1])*v_hi[2][u[2]] - (  u[1]/M[1])*(  u[2]/M[2])*v_cr[2,2]
    end
    return interpolated_array
end

end # module transfinite_interpolate


using ParallelAccelerator

@acc function kernel_gesummv(n::Int64, alpha::Float64, beta::Float64,
                             A::Matrix{Float64}, B::Matrix{Float64},
                             x::Vector{Float64})
    y = Array{Float64}(n)
    for i = 1:n
        t = 0.0
        y[i] = 0.0
        for j = 1:n
            t = A[i, j] * x[j] + t
            y[i] = B[i, j] * x[j] + y[i]
        end
        y[i] = alpha * t + beta * y[i]
    end
    return y
 end

@inbounds function kernel_gesummv_o(n::Int64, alpha::Float64, beta::Float64,
                             A::Matrix{Float64}, B::Matrix{Float64},
                             x::Vector{Float64})
    tic()
    y = Array{Float64}(n)
    for i = 1:n
        t = 0.0
        y[i] = 0.0
        for j = 1:n
            t = A[i, j] * x[j] + t
            y[i] = B[i, j] * x[j] + y[i]
        end
        y[i] = alpha * t + beta * y[i]
    end
    println("Julia ", toq())
    return y
 end

 n = 15000
 alpha = 1.5
 beta = 1.2

 A = rand(n, n)
 B = rand(n, n)
 x = rand(n)

res_tiramisu = kernel_gesummv(n, alpha, beta, A, B, x)
res_julia = kernel_gesummv_o(n, alpha, beta, A, B, x)

function compare(A, B)
    for i in 1:n
        if A[i] != B[i]
            println("different $((i)) : $(A[i]) != $(B[i])")
            return false
        end
    end
    return true
end

println(compare(res_tiramisu, res_julia))


export qromp
# Based on the algorithm developed by
# Baptista, R., Stolbunov, V., & Nair, P. B. (2019).
# Some greedy algorithms for sparse polynomial chaos expansions.
# Journal of Computational Physics, 387, 303-325.

# Solve ψc = u using the orthogonal matching pursuit with QR formulation

function qromp(ψ::AbstractMatrix{T}, u::AbstractVector{T}; invert::Bool=true, verbose::Bool = true, ϵrel::Float64 = 1e-1, maxterms::Int64=typemax(Int64)) where {T}
    m, n = size(ψ)

    residue = copy(u)
    idxset = Int64[]
    dict = collect(1:n)
    ϵu = norm(u)
    ϵhist = [ϵu]
    ϵrel *= ϵu
    ϕ = zeros(T, n)

    @inbounds for j=1:n
        for i=1:m
        ϕ[j] += ψ[i,j]^2
        end
    end

    # F = qrfactUnblocked(zeros(0,0))
    # F = elasticqrfact(zeros(m,0))
    cache = zeros(T, m)
    F = elasticqrfact(reshape(cache, length(cache),1))
    ek = zeros(T, n)
    # ektest = zeros(T, n)
    q = zeros(T, m)
    new_idx = 1

    @inbounds for k=1:n
        # Step 6 & 7 update denominators and choose best candidate

        entry = 0.0
        @fastmath @inbounds for (j, idx) in enumerate(dict)
            ψj = view(ψ,:,idx)
            ratio = dot(residue, ψj)
            ratio = ratio^2
            ratio /= ϕ[idx]

            if ratio >= entry
                entry = copy(ratio)
                new_idx = copy(idx)
            end
        end

        # if k>1
        #     for (j, idx) in enumerate(dict)
        #         ψj = view(ψ,:,idx)
        #         @show norm(ϕ[idx] - norm(ψj-Matrix(F.Q)[:,1:k-1]*(Matrix(F.Q)[:,1:k-1]'*ψj))^2)
        #     end
        # end

        # Step 8 Update set of selected basis
        push!(idxset, new_idx)

        # Step 9 Update candidate dictionary
        filter!(x-> x != new_idx, dict)

        # Step 10 Update Q(k-1) and R(k-1) with ψi(k)
        if k>1
            # F = updateqrfactUnblocked!(F, view(ψ,:,new_idx))
            updateelasticqrfact!(F, view(ψ,:,new_idx), cache)

            ek[k-1] = 0.0
        else
            # F = qrfactUnblocked(ψ[:,new_idx:new_idx])
            F = elasticqrfact(view(ψ,:,new_idx:new_idx))
        end

        # Step 11 Update residual
        ek[k] = 1.0
        mul!(q, F.Q, view(ek,1:k))

        @inbounds for (j, idx) in enumerate(dict)
            ψj = view(ψ,:,idx)
            ϕj = dot(q, ψj)
            ϕ[idx] -= ϕj^2
        end
        factor = dot(q, u)
        residue -= factor*q

        # Step 12 Calculate stopping critera
        ϵ = norm(residue)

        if verbose == true
            push!(ϵhist, copy(ϵ))
        end
        if ϵ < ϵrel || k == maxterms
            break
        end
    end

    # Solve the system with the set of indices
    if verbose == true
        if invert == true
            c = zeros(length(idxset))
            ldiv!(c, F, u)
            return idxset, c, ϵhist
        else
            return idxset, ϵhist
        end
    else
        if invert == true
            c = zeros(length(idxset))
            ldiv!(c, F, u)
            return idxset, c
        else
            return idxset
        end
    end
end


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: Julia 1.5.1
#     language: julia
#     name: julia-1.5
# ---

using Plots

x = range(-pi,pi,length=100)
y = sin.(x)
plot(aspect_ratio=:equal,xlim=[-pi,pi])
plot!(sin,color=:red)
plot!(x, y, fill=y .+ 1,color=:green)
plot!(x, y, fill=y .- 1,color=:blue)

x = range(-pi,pi,length=100)
y = sin.(x)
l = -ones(length(y))
u = ones(length(y))
plot(aspect_ratio=:equal,xlim=[-pi,pi])
plot!(sin, ribbons=[l, u], color=:blue,label=:none)
plot!(sin,color=:red,label=:none)

x=range(-pi,pi,length=200)
plot(x,sin.(x),fill=(x-> sin(x)>0 ? sin(x) : 0),alpha=0.5)


"""
# plotbounds

Plot regression line and intervals based on Stan samples of coeffficients.

$(SIGNATURES)

### Required arguments
```julia
* `df::DataFrame`                      : DataFrame with observed variables and scaled variables
* `xvar::Symbol`                       : X variable in df
* `yvar::Symbol`                       : Y variable in df
* `dfs::DataFrame`                     : DataFrame with Stan samples
* `linkvars::Vector{Symbol}`           : Initial 2 Symbols are regression coefficients,
                                         3rd - only for :predict bounds indicates σ.
```
### Optional arguments
```julia
* `fnc = link`                         : Link function
* `fig::AbstractString=""`             : File to store the plot. If "", a plot is returned
* `bounds::Vector{Symbol}`             : Bounds to display, see below
* `title::AbstractString=""`           : Title for plot
* `xlab::AbstractString=String(xvar)`  : X axis variable label
* `ylab::AbstractString=String(yvar)`  : Y axis variable label
* `alpha::Float64=0.11`                : Interval value, defaults to [0.045, 0.945]
* `colors::Vector{Symbol}`             : Colors for regions, defaults to [:lightgrey, :grey]
* `stepsize::Float64=0.01`             : Stepsize for boundary accuracy
* `rescale_axis=true`                  : Display using un-standardized scale
```

This method is primarily intended to display 2 regions, typically predicted values and
the quantile or hpdi region around the mean line. This is specified using the `bounds`
keyword argument, e.g. `bounds = [:predicted, :hpdi]`.

For the prediction interval, a 3rd parameter needs to be present in the Stan sample
DataFrame (dfs) containing the σ value. This symbol needs to be added to `linkvars`, e.g.
```julia
linkvars = [:a, :bM, :sigma]
```

For other options, :quantile and :hpdi, two parameters suffice (typically the itercept
and slope parameters).

"""
function plotbounds(
  df::DataFrame,
  xvar::Symbol,
  yvar::Symbol,
  dfs::DataFrame,
  linkvars::Vector{Symbol};
  fnc = link,
  fig::AbstractString="",
  bounds::Vector{Symbol}=[:predicted, :hpdi],
  title::AbstractString="",
  xlab::AbstractString=String(xvar),
  ylab::AbstractString=String(yvar),
  alpha::Float64=0.11,
  colors::Vector{Symbol}=[:lightgrey, :grey],
  stepsize::Float64=0.01,
  rescale_axis=true,
  kwargs...
)

  xbar = mean(df[:, xvar])
  xstd = std(df[:, xvar])
  ybar = mean(df[:, yvar])
  ystd = std(df[:, yvar])

  xvar_s = Symbol(String(xvar)*"_s")
  yvar_s = Symbol(String(yvar)*"_s")

  x_s = minimum(df[:, xvar_s]):stepsize:maximum(df[:, xvar_s])
  y_s = fnc(dfs, linkvars, x_s);

  if rescale_axis
    x = rescale(x_s, xbar, xstd)
    y = [rescale(y_s[i], ybar, ystd) for i in 1:length(x)]
  else
    x = x_s
    y = y_s
    xbar = 0.0
    xstd = 1.0
    ybar = 0.0
    ystd = 1.0
  end

  mu = [mean(y[i]) for i in 1:length(x)]

  p = plot(;xlab=xlab, ylab=ylab, title=title, kwargs...)

  if :predicted in bounds
    predictions = zeros(length(x), nrow(dfs))
    for i in 1:length(x)
      for j in 1:nrow(dfs)
        predictions[i, j] =
          ybar .+ rand(Normal(dfs[j, linkvars[1]] + ystd/xstd * dfs[j, linkvars[2]] * (x[i] - xbar), ystd*dfs[j, linkvars[3]]))
      end
    end
    pred_hpd = [hpdi(predictions[i, :], alpha=alpha) for i in 1:length(x)]
    for i in 1:length(x)
      plot!([x[i], x[i]], pred_hpd[i], color=colors[1], leg=false)
    end
  end

  if :quantile in bounds
    bnds_quantile = [quantile(y[i], [alpha/2, 1-alpha/2]) for i in 1:length(x)]
    for i in 1:length(x)
      plot!([x[i], x[i]], bnds_quantile[i], color=colors[2], leg=false)
    end
  end

  if :hpdi in bounds
    bnds_hpd = [hpdi(y[i], alpha=alpha) for i in 1:length(x)]
    for i in 1:length(x)
      plot!([x[i], x[i]], bnds_hpd[i], color=colors[2], leg=false)
    end
  end

  plot!(x, mu, color=:black)
  if rescale_axis
    scatter!(df[:, xvar], df[:, yvar], leg=false, color=:darkblue)
  else
    scatter!(df[:, xvar_s], df[:, yvar_s], leg=false, color=:darkblue)
  end

  if fig == ""
    return(p)
  else
    savefig(p, fig)
  end

end

"""
# plotbounds

Plot regression line and PI interval with alpha=0.11.

$(SIGNATURES)

### Required arguments
```julia
* `df::DataFrame`                      : DataFrame with observed variables and scaled variables
* `xvar::Symbol`                       : X variable in df
* `yvar::Symbol`                       : Y variable in df
* `nt::NamedTuple`                     : NamedTuple with Stan samples
* `linkvars::Vector{Symbol}`           : Initial 2 Symbols are regression coefficients,
                                         3rd - only for :predict bounds indicates σ.
```
### Optional arguments
```julia
* `fnc = link`                         : Link function
* `fig::AbstractString=""`             : File to store the plot. If "", a plot is returned
* `stepsize::Float64=0.01`             : Stepsize for boundary accuracy
* `rescale_axis=true`                  : Display using un-standardized scale
* `title::AbstractString=""`           : Title for plot
* `lab::AbstractString=""`             : X axis variable label
```

This method is intended to the PI region around the mean line.

The 2nd symbol in `linkvars` needs to match the matrix
```julia
linkvars = [:a, :bM]
```

For other options, :quantile and :hpdi, two parameters suffice (typically the itercept
and slope parameters).

"""
function plotbounds(
    df::DataFrame,
    xvar::Symbol,
    yvar::Symbol,
    nt::NamedTuple,
    linkvars::Vector{Symbol};
    fnc::Function=link,
    fig::AbstractString="",
    stepsize=0.01,
    rescale_axis=true,
    lab::AbstractString="",
    title::AbstractString="",
    kwargs...)

    xbar = mean(df[:, xvar])
    xstd = std(df[:, xvar])
    ybar = mean(df[:, yvar])
    ystd = std(df[:, yvar])

    xvar_s = Symbol(String(xvar)*"_s")
    yvar_s = Symbol(String(yvar)*"_s")

    minx = minimum(df[!, xvar_s]) - 0.1
    maxx = maximum(df[!, xvar_s]) + 0.1
    x_s = minx:stepsize:maxx

    k = size(nt[Symbol(linkvars[2])], 1)
    m = fnc(collect(x_s), k)
    y_s = m * nt[Symbol(linkvars[2])]
    y_s .+= nt[Symbol(linkvars[1])]'
    mul = meanlowerupper(y_s)

    p = plot(;title, kwargs...)
    if rescale_axis
        xrange = rescale(x_s, xbar, xstd)
        mm = rescale(mul.mean, ybar, ystd)
        lm = rescale(mul.lower, ybar, ystd)
        um = rescale(mul.upper, ybar, ystd)
        plot!(p, xrange, mm;
            xlab=String(xvar), ylab=String(yvar),
            ribbon=(mm-lm, um-mm),
            lab)
    else
        plot!(p, xrange, mul.mean;
            xlab=String(xvar), ylab=String(yvar),
            ribbon=(mul.mean-mul.lower, mul.upper-mul.mean),
            lab)
    end

    if fig == ""
        return(p)
    else
        savefig(p, fig)
    end
end

export
  plotbounds