using Distributions,StatsPlots, ForwardDiff, LinearAlgebra, PositiveFactorizations, Optim

n_samples = 2000

n = 100

x = randn(n,50)
theta = [repeat([1],10); repeat([0],40)]
y = x*theta + randn(n)

prior = Laplace(0,2)

ll(theta) = logpdf(MvNormal(I(n)),y-x*theta)

post(theta) = ll(theta) + sum(logpdf.(prior,theta))

init = rand(prior,50)

grad(theta) = ForwardDiff.gradient(post,theta)

ot = optimize(x->-1*ll(x),zeros(50),BFGS(), autodiff=:forward)

M = grad(ot.minimizer)*grad(ot.minimizer)'#mass matrix
M = inv(diagm(diag(M)))

mom_dist = MvNormal(M)

tol_psi = 1e18

#######
### start here
#

pars = zeros(n_samples,50)

uni = Uniform(0,1)

leap_step = 0.09

L = max(Int(ceil(1/leap_step)),20)

window_size = 10

window_dist = DiscreteUniform(1,window_size)

accept = 0

old_par = init#repeat([0],50)

leap_dist = Normal(leap_step,0.05)

for j = 1:500
    psi = rand(mom_dist)
    old_psi = psi
    new_par = old_par
    leap_step = rand(leap_dist)
    s = rand(window_dist)
    forward_leap = zeros(50,2)
    sum_p = 0
    for i = 1:(s*L) #forward leapfrog
        psi = psi + 1/2*leap_step*grad(new_par)
        new_par = new_par + inv(M)*leap_step*psi
        psi = psi + 1/2*leap_step*grad(new_par)
        pp = exp(logpdf(mom_dist,psi)+post(new_par))
        sum_p = sum_p + pp
        if pp/sum_p > rand(uni)
            forward_leap = [psi new_par]
        end
    end
    backward_leap = zeros(50,2)
    sum_p = 0
    psi = old_psi
    new_par = old_par
    for i = 1:((window_size-s)*L) #backward leapfrog
        psi = psi - 1/2*leap_step*grad(new_par)
        new_par = new_par - inv(M)*leap_step*psi
        psi = psi - 1/2*leap_step*grad(new_par)
        pp = exp(logpdf(mom_dist,psi)+post(new_par))
        sum_p = sum_p + pp
        if pp/sum_p > rand(uni)
            backward_leap = [psi new_par]
        end
    end
    #println(psi)
    if any(psi .> tol_psi)
        @info "Energy diverging. Quitting"
        break
    end
    r = post(forward_leap[:,2]) + logpdf(mom_dist,forward_leap[:,1]) - (post(backward_leap[:,2])+ logpdf(mom_dist,backward_leap[:,1]))
    a = rand(uni)
    r = min(1,exp(r))
    if r > a
        pars[j,:] = new_par
        global old_par = new_par
        global accept = accept + 1
    else
        pars[j,:] = old_par
    end
    println("Iteration ", j, " Accept ratio ", accept/j*100, "%")
end

leap_step = 0.15

L = max(Int(ceil(1/leap_step)),20)

accept = 0

old_par = reshape(mapslices(mode,pars[1:500,:],dims = 1),50)

for j = 1:2000
    psi = rand(mom_dist)
    old_psi = psi
    new_par = old_par
    leap_step = rand(leap_dist)
    s = rand(window_dist)
    forward_leap = zeros(50,2)
    sum_p = 0
    for i = 1:(s*L) #forward leapfrog
        psi = psi + 1/2*leap_step*grad(new_par)
        new_par = new_par + inv(M)*leap_step*psi
        psi = psi + 1/2*leap_step*grad(new_par)
        pp = exp(logpdf(mom_dist,psi)+post(new_par))
        sum_p = sum_p + pp
        if pp/sum_p > rand(uni)
            forward_leap = [psi new_par]
        end
    end
    backward_leap = zeros(50,2)
    sum_p = 0
    psi = old_psi
    new_par = old_par
    for i = 1:((window_size-s)*L) #backward leapfrog
        psi = psi - 1/2*leap_step*grad(new_par)
        new_par = new_par - inv(M)*leap_step*psi
        psi = psi - 1/2*leap_step*grad(new_par)
        pp = exp(logpdf(mom_dist,psi)+post(new_par))
        sum_p = sum_p + pp
        if pp/sum_p > rand(uni)
            backward_leap = [psi new_par]
        end
    end
    #println(psi)
    if any(psi .> tol_psi)
        @info "Energy diverging. Quitting"
        break
    end
    r = post(forward_leap[:,2]) + logpdf(mom_dist,forward_leap[:,1]) - (post(backward_leap[:,2])+ logpdf(mom_dist,backward_leap[:,1]))
    a = rand(uni)
    r = min(1,exp(r))
    if r > a
        pars[j,:] = new_par
        global old_par = new_par
        global accept = accept + 1
    else
        pars[j,:] = old_par
    end
    println("Iteration ", j, " Accept ratio ", accept/j*100, "%")
end

StatsPlots.density(pars[200:1700,1])

mm = mapslices(mode,pars[500:1126,:],dims = 1)

scatter(mm')

StatsPlots.density(pars[200:1700,1])
StatsPlots.density(pars[200:1700,2])
StatsPlots.density(pars[200:1700,3])
StatsPlots.density(pars[200:1700,4])
StatsPlots.density(pars[200:1700,5])
StatsPlots.density(pars[200:1700,6])
StatsPlots.density(pars[200:1700,7])
StatsPlots.density(pars[200:1700,8])
StatsPlots.density(pars[200:1700,9])
StatsPlots.density(pars[200:1700,10])
StatsPlots.density(pars[200:1700,11])
StatsPlots.density(pars[200:1700,12])
StatsPlots.density(pars[200:1700,13])
StatsPlots.density(pars[200:1700,14])
StatsPlots.density(pars[200:1700,15])
StatsPlots.density(pars[200:1700,16])


#=
compute the capacity of a classical-quantum channel
adapted from https://github.com/hfawzi/cvxquad/blob/master/examples/cq_channel_capacity.m
and listing 1 in "Efficient optimization of the quantum relative entropy" by H. Fawzi and O. Fawzi
=#

struct ClassicalQuantum{T <: Real} <: ExampleInstanceJuMP{T}
    d::Int
    complex::Bool
    use_EF::Bool
end

function build(inst::ClassicalQuantum{T}) where {T <: Float64}
    d = inst.d
    ext = (inst.use_EF ? MatNegEntropyEigOrd() : MatNegEntropy())
    rt2 = sqrt(T(2))
    R = (inst.complex ? Complex{T} : T)

    function hermtr1()
        P = randn(R, d, d)
        P = Hermitian(P * P', :U)
        P.data ./= tr(P)
        return Hermitian(P)
    end
    Ps = [hermtr1() for _ in 1:d]
    Hs = [dot(P, log(P)) for P in Ps]

    model = JuMP.Model()
    JuMP.@variable(model, ρ[1:d] >= 0)
    JuMP.@constraint(model, sum(ρ) == 1)
    JuMP.@variable(model, epi)
    JuMP.@objective(model, Max, -epi + dot(ρ, Hs))

    vec_dim = Cones.svec_length(R, d)
    entr = JuMP.AffExpr.(zeros(vec_dim))
    P_vec = zeros(T, vec_dim)
    for (Pi, ρi) in zip(Ps, ρ)
        Cones.smat_to_svec!(P_vec, Pi, rt2)
        JuMP.add_to_expression!.(entr, ρi, P_vec)
    end
    add_homog_spectral(ext, d, vcat(epi, entr), model)

    # save for use in tests
    model.ext[:epi] = epi
    model.ext[:entr] = entr

    return model
end

function test_extra(inst::ClassicalQuantum{T}, model::JuMP.Model) where T
    stat = JuMP.termination_status(model)
    @test stat == MOI.OPTIMAL
    (stat == MOI.OPTIMAL) || return

    # check constraint
    tol = eps(T)^0.2
    epi_opt = JuMP.value(model.ext[:epi])
    Entr_opt = zeros(inst.complex ? Complex{T} : T, inst.d, inst.d)
    Cones.svec_to_smat!(Entr_opt, JuMP.value.(model.ext[:entr]), sqrt(T(2)))
    λ = eigvals(Hermitian(Entr_opt, :U))
    @test minimum(λ) >= -tol
    qe_result = get_val(pos_only(λ), MatNegEntropy())
    @test epi_opt ≈ qe_result atol=tol rtol=tol
    return
end


# author: aliyyousuf@gmail.com

# assignment of week_03.

data_to_fit = readdlm("Week3_PR_Data.dat",',')

# Converting data into Float64

data_to_fit_2 = zeros(Float64, 15,2)
for r = 1:15
    for c = 1
        data_to_fit_2[r,c] = parse(Float64,split(data_to_fit[r,c])[1])
        data_to_fit_2[r,2] = parse(Float64,split(data_to_fit[r,c])[2])
    end
end

# "For loop" to show all the values of data:

for r = 1:15
    for c = 1:2
        #println(data_to_fit_2[r,c])
    end
end

# assign the values from first column to variable x
# and second column to variable y.

x = data_to_fit_2[:,1]
y = data_to_fit_2[:,2]

using Plots
pyplot()
plot(x,y, title = "Model versus Data",label = "Actual values from given data",
xlabel = " x values",
ylabel = "y values",
line =:scatter)
gui()

function parab_fit(x)
    return a * (x)^2 + b * x + c
end

a = -0.789
b = 2.988
c = -1.5
plot!(parab_fit,title = "Model versus Data",-10,10,
label = "Model values")
gui()


using JuMP, vOptGeneric, GLPK, Combinatorics
using Test, LinearAlgebra

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)

@variable(m, x[1:6,1:6], Bin)
@variable(m, y, Bin)
M = 50
p1 =   [M 4 5 M M M;
        M M 2 1 2 7;
        M M M 5 2 M;
        M M 5 M M 3;
        M M M M M 4;
        M M M M M M]
p2 =   [M 3 1 M M M;
        M M 1 4 2 2;
        M M M 1 7 M;
        M M 1 M M 2;
        M M M M M 2;
        M M M M M M]

@addobjective(m, Min, sum(x[i,j]*p1[i,j] for i=1:6, j=1:6) + 5)
@addobjective(m, Min, sum(x[i,j]*p2[i,j] for i=1:6, j=1:6) - 5)
@constraint(m, sum(x[1,j] for j = 1:6) == 1)
@constraint(m, sum(x[i,6] for i = 1:6) == 1)
@constraint(m, cstr[i=2:5], sum(x[i,j] for j = 1:6) - sum(x[j,i] for j =1:6) == 0)

vSolve(m, method=:epsilon)

Y_N = getY_N(m)
f1, f2 = first.(Y_N), last.(Y_N)
@test f1 == [13, 15, 16, 18] && f2 == [4, 2, 0, -1]
@test value.(x, 1) == [0.0 1.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 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]
printX_E(m)

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

# to test once https://github.com/JuliaOpt/JuMP.jl/issues/1956 is fixed
@test_broken begin
    m_copy = copy(m)
    # vSolve(m_copy, GLPK.Optimizer, method=:epsilon)
    # Y_N = getY_N(m)
    # f1 = map(x -> x[1], Y_N) ; f2 = map(x -> x[2], Y_N)
    # @test f1 == [8.0,10.0,11.0,13.0] && f2 == [9.0,7.0,5.0,4.0]
    # @test value.(x, 1) == [0.0 1.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 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]
    # printX_E(m_copy)
end

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

# tests DenseAxisArrays indexing
# http://www.juliaopt.org/JuMP.jl/v0.19.1/variables/#variable_jump_arrays-1

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)
# JuMP.set_optimizer_attribute(m, "logLevel", 0)
cities = ["Paris", "New-York", "Madrid"]
@variable(m, x[cities], Bin)
@constraint(m, constr, x["Paris"] >= 1)
@addobjective(m, Min, sum(x))
@addobjective(m, Max, sum(x))
vSolve(m, method=:epsilon)
@show value.(x, 1)
@test value.(x, 1)["Paris"] == 1.0
@test value.(x, 1)["New-York"] == 0.0
@test value.(x, 1)["Madrid"] == 0.0

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

# TODO : write a test for SparseAxisArray
# http://www.juliaopt.org/JuMP.jl/v0.19.1/variables/#variable_sparseaxisarrays-1

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

# TODO : re-implement writeMOP and parseMOP
# writeMOP(m, "test.txt")
# m = parseMOP("test.txt", solver = GLPKSolverMIP())
# solve(m, method=:dichotomy)
# Y_N = getY_N(m)
# f1 = map(x -> x[1], Y_N)
# f2 = map(x -> x[2], Y_N)
# @test f1 == [8.0, 11.0, 13.0]
# @test f2 == [9.0, 5.0, 4.0]

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

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)
# JuMP.set_optimizer_attribute(m, "logLevel", 0)

p1 = [77,94,71,63,96,82,85,75,72,91,99,63,84,87,79,94,90,60,69,62]
p2 = [65,90,90,77,95,84,70,94,66,92,74,97,60,60,65,97,93,60,69,74]
w = [80,87,68,72,66,77,99,85,70,93,98,72,100,89,67,86,91,79,71,99]
c = 900
size = length(p1)

@variable(m, x[1:size], Bin)
@addobjective(m, Max, dot(x, p1))
@addobjective(m, Max, dot(x, p2))
@constraint(m, dot(x, w) <= c)

vSolve(m, method=:dicho)

Y_N = getY_N(m)
f1, f2 = first.(Y_N), last.(Y_N)

@test f1 == [918.0, 934.0, 948.0, 955.0]
@test f2 == [984.0, 971.0, 939.0, 906.0]
@test value.(x, 1) ≈ [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0]

vSolve(m, method=:chalmet)

Y_N = getY_N(m)
f1, f2 = first.(Y_N), last.(Y_N)

@test f1 == [918., 924., 927., 934., 935., 940., 943., 948., 949., 952., 955.]
@test f2 == [984., 975., 972., 971., 947., 943., 940., 939., 915., 909., 906.]
@test value.(x, 1) ≈ [0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0]

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

m = vModel(GLPK.Optimizer) ; JuMP.set_silent(m)
# JuMP.set_optimizer_attribute(m, "logLevel", 0)

@variable(m, 0 <= x[1:4] <= 2, Int)
@addobjective(m, Max, sum(x))
@addobjective(m, Min, sum(x))
@addobjective(m, Max, x[3])
@addobjective(m, Max, x[4])
@constraint(m, sum(x) <= 3)
vSolve(m, method=:lex)
Y_N = getY_N(m)

for (ind, perm) in enumerate(permutations(1:4, 4))
    if perm[1] == 2
        @test Y_N[ind] == [0., 0., 0., 0.]
    elseif perm[2] == 2 && perm[1] != 1
        @test maximum(Y_N[ind]) == 2.
    end
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

using DifferentialEquations

function main()

function micro_1(du, u, parms, time)
    # PARAMETER DEFS
    # β transmition rate
    # r net population growth rate
    # μ hosts' natural mortality rate
    # Κ population size
    # α disease induced mortality rate

    β, r, μ, K, α = parms
    dS = r*(1-S/K)*S - β*S*I
    dI = β*S*I-(μ+α)*I
    du = [dS,dI]
end

# +
# PARAMETER DEFS
# w and m are used to define the other parameters allometrically

w = 1;
m = 10;
β = 0.0247*m*w^0.44;
r = 0.6*w^-0.27;
μ = 0.4*w^-0.26;
K = 16.2*w^-0.7;
α = (m-1)*μ;
# -

parms = [β,r,μ,K,α];
init = [K,1.];
tspan = (0.0,10.0);

sir_prob = ODEProblem(micro_1,init,tspan,parms)

sir_sol = solve(sir_prob);

end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -3.1276627592672863 <= q <= 3.1323678679599185)

                     add_NL_constraint(m, :(1/(1 + exp(-(0.7015310798579062 + -0.563498653088923*$(x[1]) + -0.9918531495953298*$(x[2]) + 0.03337339686246876*$(x[3])))) + 1/(1 + exp(-(-0.6991785255115901 + -0.8453063373771403*$(x[1]) + 0.408983421114415*$(x[2]) + 0.28700035557532555*$(x[3])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


__precompile__()

module NNLS
using LinearAlgebra
using Statistics: mean

export nnls,
       solve!,
       NNLSWorkspace,
       QPWorkspace,
       QP,
       primal_infeasibility,
       dual_infeasibility,
       stationarity_violation,
       slackness_violation,
       check_optimality_conditions,
       load!,
       NNLSSolver

"""
CONSTRUCTION AND/OR APPLICATION OF A SINGLE
HOUSEHOLDER TRANSFORMATION..     Q = I + U*(U**T)/B

The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 12, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.
"""
function construct_householder!(u::AbstractVector{T}, up::T)::T where T
    m = length(u)
    if m <= 1
        return up
    end

    cl = maximum(abs, u)
    @assert cl > 0
    clinv = 1 / cl
    sm = zero(T)
    for ui in u
        sm += (ui * clinv)^2
    end
    cl *= sqrt(sm)
    if u[1] > 0
        cl = -cl
    end
    result = u[1] - cl
    u[1] = cl

    return result
end

"""
CONSTRUCTION AND/OR APPLICATION OF A SINGLE
HOUSEHOLDER TRANSFORMATION..     Q = I + U*(U**T)/B

The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 12, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.
"""
function apply_householder!(u::AbstractVector{T}, up::T, c::AbstractVector{T}) where T
    m = length(u)
    if m > 1
        cl = abs(u[1])
        @assert cl > 0
        b = up * u[1]
        if b >= 0
            return
        end
        b = 1 / b

        sm = c[1] * up
        for i in 2:m
            sm += c[i] * u[i]
        end
        if sm != 0
            sm *= b
            c[1] += sm * up
            for i in 2:m
                c[i] += sm * u[i]
            end
        end
    end
end

"""
   COMPUTE ORTHOGONAL ROTATION MATRIX..
The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 12, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.

   COMPUTE.. MATRIX   (C, S) SO THAT (C, S)(A) = (SQRT(A**2+B**2))
                      (-S,C)         (-S,C)(B)   (   0          )
   COMPUTE SIG = SQRT(A**2+B**2)
      SIG IS COMPUTED LAST TO ALLOW FOR THE POSSIBILITY THAT
      SIG MAY BE IN THE SAME LOCATION AS A OR B .
"""
function orthogonal_rotmat(a::T, b::T)::Tuple{T, T, T} where T
    if abs(a) > abs(b)
        xr = b / a
        yr = sqrt(1 + xr^2)
        c = (1 / yr) * sign(a)
        s = c * xr
        sig = abs(a) * yr
    elseif b != 0
        xr = a / b
        yr = sqrt(1 + xr^2)
        s = (1 / yr) * sign(b)
        c = s * xr
        sig = abs(b) * yr
    else
        sig = zero(T)
        c = zero(T)
        s = one(T)
    end
    return c, s, sig
end

"""
The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 15, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.
"""
function solve_triangular_system!(zz, A, idx, nsetp, jj)
    for l in 1:nsetp
        ip = nsetp + 1 - l
        if (l != 1)
            for ii in 1:ip
                zz[ii] -= A[ii, jj] * zz[ip + 1]
            end
        end
        jj = idx[ip]
        zz[ip] /= A[ip, jj]
    end
    return jj
end

mutable struct NNLSWorkspace{T, I <: Integer}
    QA::Matrix{T}
    Qb::Vector{T}
    x::Vector{T}
    w::Vector{T}
    zz::Vector{T}
    idx::Vector{I}
    rnorm::T
    mode::I
    nsetp::I

    function NNLSWorkspace{T, I}(m, n) where {T, I <: Integer}
        new{T, I}(Matrix{T}(undef, m, n), # A
            Vector{T}(undef, m),    # b
            Vector{T}(undef, n),    # x
            Vector{T}(undef, n),    # w
            Vector{T}(undef, m),    # zz
            Vector{I}(undef, n),    # idx
            zero(T), # rnorm
            zero(I), # mode
            zero(I)  # nsetp
        )
    end
end

function Base.resize!(work::NNLSWorkspace{T}, m::Integer, n::Integer) where T
    work.QA = Matrix{T}(undef, m, n)
    work.Qb = Vector{T}(undef, m)
    resize!(work.x, n)
    resize!(work.w, n)
    resize!(work.zz, m)
    resize!(work.idx, n)
end

function load!(work::NNLSWorkspace{T}, A::AbstractMatrix{T}, b::AbstractVector{T}) where T
    m, n = size(A)
    @assert size(b) == (m,)
    if size(work.QA, 1) != m || size(work.QA, 2) != n
        resize!(work, m, n)
    end
    work.QA .= A
    work.Qb .= b
    work
end

NNLSWorkspace(m::Integer, n::Integer,
              eltype::Type{T}=Float64,
              indextype::Type{I}=Int) where {T, I} = NNLSWorkspace{T, I}(m, n)

function NNLSWorkspace(A::Matrix{T}, b::Vector{T}, indextype::Type{I}=Int) where {T, I}
    m, n = size(A)
    @assert size(b) == (m,)
    work = NNLSWorkspace{T, I}(m, n)
    load!(work, A, b)
    work
end

"""
Views in Julia still allocate some memory (since they need to keep
a reference to the original array). This type allocates no memory
and does no bounds checking. Use it with caution.
"""
struct UnsafeVectorView{T} <: AbstractVector{T}
    offset::Int
    len::Int
    ptr::Ptr{T}
end

UnsafeVectorView(parent::DenseArray{T}, start_ind::Integer, len::Integer) where {T} = UnsafeVectorView{T}(start_ind - 1, len, pointer(parent))
Base.size(v::UnsafeVectorView) = (v.len,)
Base.getindex(v::UnsafeVectorView, idx) = unsafe_load(v.ptr, idx + v.offset)
Base.setindex!(v::UnsafeVectorView, value, idx) = unsafe_store!(v.ptr, value, idx + v.offset)
Base.length(v::UnsafeVectorView) = v.len
Base.IndexStyle(::Type{V}) where {V <: UnsafeVectorView} = Base.IndexLinear()

"""
UnsafeVectorView only works for isbitstype types. For other types, we're already
allocating lots of memory elsewhere, so creating a new View is fine.

This function looks type-unstable, but the isbitstype(T) test can be evaluated
by the compiler, so the result is actually type-stable.
"""
function fastview(parent::Array{T}, start_ind::Integer, len::Integer) where T
    if isbitstype(T)
        UnsafeVectorView(parent, start_ind, len)
    else
        @view(parent[start_ind:(start_ind + len - 1)])
    end
end

"""
Fallback for non-contiguous arrays, for which UnsafeVectorView does not make
sense.
"""
fastview(parent::AbstractArray, start_ind::Integer, len::Integer) = @view(parent[start_ind:(start_ind + len - 1)])

@noinline function checkargs(work::NNLSWorkspace)
    m, n = size(work.QA)
    @assert size(work.Qb) == (m,)
    @assert size(work.x) == (n,)
    @assert size(work.w) == (n,)
    @assert size(work.zz) == (m,)
    @assert size(work.idx) == (n,)
end

function largest_positive_dual(w::AbstractVector{T}, idx::AbstractVector{TI}, range) where {T, TI}
    wmax = zero(T)
    izmax = zero(TI)
    for i in range
        j = idx[i]
        if w[j] > wmax
            wmax = w[j]
            izmax = i
        end
    end
    wmax, izmax
end

"""
Algorithm NNLS: NONNEGATIVE LEAST SQUARES

The original version of this code was developed by
Charles L. Lawson and Richard J. Hanson at Jet Propulsion Laboratory
1973 JUN 15, and published in the book
"SOLVING LEAST SQUARES PROBLEMS", Prentice-HalL, 1974.
Revised FEB 1995 to accompany reprinting of the book by SIAM.

GIVEN AN M BY N MATRIX, A, AND AN M-VECTOR, B,  COMPUTE AN
N-VECTOR, X, THAT SOLVES THE LEAST SQUARES PROBLEM
                 A * X = B  SUBJECT TO X .GE. 0
"""
function solve!(work::NNLSWorkspace{T, TI}, max_iter::Integer=(3 * size(work.QA, 2))) where {T, TI}
    checkargs(work)

    A = work.QA
    b = work.Qb
    x = work.x
    w = work.w
    zz = work.zz
    idx = work.idx
    factor = 0.01
    work.mode = 1

    m = convert(TI, size(A, 1))
    n = convert(TI, size(A, 2))

    iter = 0
    x .= 0
    idx .= 1:n

    iz2 = n
    iz1 = one(TI)
    iz = zero(TI)
    j = zero(TI)
    jj = zero(TI)
    nsetp = zero(TI)
    up = zero(T)

    terminated = false

    # ******  MAIN LOOP BEGINS HERE  ******
    while true
        # println("jl main loop")
        # QUIT IF ALL COEFFICIENTS ARE ALREADY IN THE SOLUTION.
        # OR IF M COLS OF A HAVE BEEN TRIANGULARIZED.
        if (iz1 > iz2 || nsetp >= m)
            terminated = true
            break
        end

        # COMPUTE COMPONENTS OF THE DUAL (NEGATIVE GRADIENT) VECTOR W().
        for i in iz1:iz2
            idxi = idx[i]
            sm = zero(T)
            for l in (nsetp + 1):m
                sm += A[l, idxi] * b[l]
            end
            w[idxi] = sm
        end

        while true
            # FIND LARGEST POSITIVE W(J).
            wmax, izmax = largest_positive_dual(w, idx, iz1:iz2)

            # IF WMAX .LE. 0. GO TO TERMINATION.
            # THIS INDICATES SATISFACTION OF THE KUHN-TUCKER CONDITIONS.
            if wmax <= 0
                terminated = true
                break
            end

            iz = izmax
            j = idx[iz]

            # THE SIGN OF W(J) IS OK FOR J TO BE MOVED TO SET P.
            # BEGIN THE TRANSFORMATION AND CHECK NEW DIAGONAL ELEMENT TO AVOID
            # NEAR LINEAR DEPENDENCE.
            Asave = A[nsetp + 1, j]
            up = construct_householder!(
                fastview(A, LinearIndices(A)[nsetp + 1, j], m - nsetp),
                up)
            unorm::T = zero(T)
            for l in 1:nsetp
                unorm += A[l, j]^2
            end
            unorm = sqrt(unorm)

            if ((unorm + abs(A[nsetp + 1, j]) * factor) - unorm) > 0
                # COL J IS SUFFICIENTLY INDEPENDENT.  COPY B INTO ZZ, UPDATE ZZ
                # AND SOLVE FOR ZTEST ( = PROPOSED NEW VALUE FOR X(J) ).
                # println("copying b into zz")
                zz .= b
                apply_householder!(
                    fastview(A, LinearIndices(A)[nsetp + 1, j], m - nsetp),
                    up,
                    fastview(zz, nsetp + 1, m - nsetp))
                ztest = zz[nsetp + 1] / A[nsetp + 1, j]

                # SEE IF ZTEST IS POSITIVE
                if ztest > 0
                    break
                end
            end

            # REJECT J AS A CANDIDATE TO BE MOVED FROM SET Z TO SET P.
            # RESTORE A(NPP1,J), SET W(J)=0., AND LOOP BACK TO TEST DUAL
            # COEFFS AGAIN.
            A[nsetp + 1, j] = Asave
            w[j] = 0
        end
        if terminated
            break
        end

        # THE INDEX  J=INDEX(IZ)  HAS BEEN SELECTED TO BE MOVED FROM
        # SET Z TO SET P.    UPDATE B,  UPDATE INDICES,  APPLY HOUSEHOLDER
        # TRANSFORMATIONS TO COLS IN NEW SET Z,  ZERO SUBDIAGONAL ELTS IN
        # COL J,  SET W(J)=0.
        b .= zz

        idx[iz] = idx[iz1]
        idx[iz1] = j
        iz1 += one(TI)
        nsetp += one(TI)

        if iz1 <= iz2
            for jz in iz1:iz2
                jj = idx[jz]
                apply_householder!(
                    fastview(A, LinearIndices(A)[nsetp, j], m - nsetp + 1),
                    up,
                    fastview(A, LinearIndices(A)[nsetp, jj], m - nsetp + 1))
            end
        end

        if nsetp != m
            for l in (nsetp + 1):m
                A[l, j] = 0
            end
        end

        w[j] = 0

        # SOLVE THE TRIANGULAR SYSTEM.
        # STORE THE SOLUTION TEMPORARILY IN ZZ().
        jj = solve_triangular_system!(zz, A, idx, nsetp, jj)

        # ******  SECONDARY LOOP BEGINS HERE ******
        #
        # ITERATION COUNTER.
        while true
            iter += 1
            if iter > max_iter
                work.mode = 3
                terminated = true
                println("NNLS quitting on iteration count")
                break
            end

            # SEE IF ALL NEW CONSTRAINED COEFFS ARE FEASIBLE.
            # IF NOT COMPUTE ALPHA.
            alpha = convert(T, 2)
            for ip in Base.OneTo(nsetp)
                l = idx[ip]
                if zz[ip] <= 0
                    t = -x[l] / (zz[ip] - x[l])
                    if alpha > t
                        alpha = t
                        jj = ip
                    end
                end
            end

            # IF ALL NEW CONSTRAINED COEFFS ARE FEASIBLE THEN ALPHA WILL
            # STILL = 2.    IF SO EXIT FROM SECONDARY LOOP TO MAIN LOOP.
            if alpha == 2
                break
            end

            # OTHERWISE USE ALPHA WHICH WILL BE BETWEEN 0 AND 1 TO
            # INTERPOLATE BETWEEN THE OLD X AND THE NEW ZZ.
            for ip in Base.OneTo(nsetp)
                l = idx[ip]
                x[l] = x[l] + alpha * (zz[ip] - x[l])
            end

            # MODIFY A AND B AND THE INDEX ARRAYS TO MOVE COEFFICIENT I
            # FROM SET P TO SET Z.
            i = idx[jj]

            while true
                x[i] = 0

                if jj != nsetp
                    jj += one(TI)
                    for j in jj:nsetp
                        ii = idx[j]
                        idx[j - 1] = ii
                        cc, ss, sig = orthogonal_rotmat(A[j - 1, ii], A[j, ii])
                        A[j - 1, ii] = sig
                        A[j, ii] = 0
                        for l in Base.OneTo(n)
                            if l != ii
                                # Apply procedure G2 (CC,SS,A(J-1,L),A(J,L))
                                temp = A[j - 1, l]
                                A[j - 1, l] = cc * temp + ss * A[j, l]
                                A[j, l] = -ss * temp + cc * A[j, l]
                            end
                        end

                        # Apply procedure G2 (CC,SS,B(J-1),B(J))
                        temp = b[j - 1]
                        b[j - 1] = cc * temp + ss * b[j]
                        b[j] = -ss * temp + cc * b[j]
                    end
                end

                nsetp -= one(TI)
                iz1 -= one(TI)
                idx[iz1] = i

                # SEE IF THE REMAINING COEFFS IN SET P ARE FEASIBLE.  THEY SHOULD
                # BE BECAUSE OF THE WAY ALPHA WAS DETERMINED.
                # IF ANY ARE INFEASIBLE IT IS DUE TO ROUND-OFF ERROR.  ANY
                # THAT ARE NONPOSITIVE WILL BE SET TO ZERO
                # AND MOVED FROM SET P TO SET Z.
                allfeasible = true
                for jji in Base.OneTo(nsetp)
                    jj = jji
                    i = idx[jj]
                    if x[i] <= 0
                        allfeasible = false
                        break
                    end
                end
                if allfeasible
                    break
                end
            end

            # COPY B( ) INTO ZZ( ).  THEN SOLVE AGAIN AND LOOP BACK.
            zz .= b
            jj = solve_triangular_system!(zz, A, idx, nsetp, jj)
        end
        if terminated
            break
        end
        # ******  END OF SECONDARY LOOP  ******

        for i in 1:nsetp
            x[idx[i]] = zz[i]
        end
        # ALL NEW COEFFS ARE POSITIVE.  LOOP BACK TO BEGINNING.
    end

    # ******  END OF MAIN LOOP  ******
    # COME TO HERE FOR TERMINATION.
    # COMPUTE THE NORM OF THE FINAL RESIDUAL VECTOR.

    sm = zero(T)
    if nsetp < m
        for i in (nsetp + 1):m
            sm += b[i]^2
        end
    else
        w .= 0
    end
    work.rnorm = sqrt(sm)
    work.nsetp = nsetp
    return work.x
end

function solve!(work::NNLSWorkspace{T}, A::AbstractMatrix{T}, b::AbstractVector{T}, max_iter=(3 * size(A, 2))) where T
    load!(work, A, b)
    solve!(work, max_iter)
    work.x
end

function nnls(A::DenseMatrix{T}, b::DenseVector{T}, max_iter=(3 * size(A, 2))) where T
    work = NNLSWorkspace(A, b)
    solve!(work, max_iter)
    work.x
end

# Implementation of an NNLS-based QP solver, based on section II of  Bemporad,
# "A quadratic programming algorithm based on nonnegative least squares with
# applications to embedded model predictive control", IEEE Transactions on
# Automatic Control, 2016.
# Variable names match the paper wherever possible

const AllColsSubArray{T} = SubArray{T,2,Array{T,2},Tuple{UnitRange{Int},Base.Slice{Base.OneTo{Int}}},false}

"""
Structure describing the QP:

Minimize ``\\frac{1}{2} z' Q z + c' z``
Subject to ``G z \\leq g``
"""
struct QP{T}
    Q::Matrix{T}
    c::Vector{T}
    G::Matrix{T}
    g::Vector{T}

    function QP{T}(Q::Matrix, c::Vector, G::Matrix, g::Vector) where T
        @boundscheck begin
            LinearAlgebra.checksquare(Q)
            length(c) == size(Q, 1) || throw(DimensionMismatch())
            size(G, 1) == length(g) || throw(DimensionMismatch())
            size(G, 2) == size(Q, 1) || throw(DimensionMismatch())
        end
        new{T}(Q, c, G, g)
    end

    QP(Q::Matrix{T}, c::Vector{T}, G::Matrix{T}, g::Vector{T}) where {T} = QP{T}(Q, c, G, g)
    QP{T}(qp::QP) where {T} = QP{T}(qp.Q, qp.c, qp.G, qp.g)
end

mutable struct QPWorkspace{T, I} # TODO: consider not having `I` as a parameter
    # Variables from paper:
    L::Matrix{T}
    c::Vector{T}
    G::Matrix{T}
    g::Vector{T}
    M::Matrix{T}
    d::Vector{T}
    r::Vector{T}

    # Additional variables:
    e::Vector{T} # intermediate variable
    A::Matrix{T} # 'A'-matrix of NNLS problem
    AM::AllColsSubArray{T} # upper block of A
    Ad::AllColsSubArray{T} # last row of A
    b::Vector{T} # 'b'-vector of NNLS problem
    nnlswork::NNLSWorkspace{T, I}
    status::Symbol

    function QPWorkspace{T, I}(q::Integer, n::Integer) where {T, I}
        work = new{T, I}()
        resize!(work, q, n)
    end
end

QPWorkspace(q::Integer, n::Integer) = QPWorkspace{Float64, Int}(q, n)

function Base.convert(::Type{QP{T1}}, qp::QP{T2}) where {T1, T2}
    QP(convert(Matrix{T1}, qp.Q),
       convert(Vector{T1}, qp.c),
       convert(Matrix{T1}, qp.G),
       convert(Vector{T1}, qp.g))
end

"""
    QPWorkspace(qp::QP)

Construct a workspace and load problem data for the QP

Minimize ``\\frac{1}{2} z' Q z + c' z``
Subject to ``G z \\leq g``
"""
function QPWorkspace(qp::QP{T}) where T
    work = QPWorkspace{T, Int}(size(qp.G, 1), size(qp.G, 2))
    load!(work, qp.Q, qp.c, qp.G, qp.g)
    work
end

function Base.resize!(work::QPWorkspace{T}, q::Integer, n::Integer) where T
    work.L = Matrix{T}(undef, n, n)
    work.c = Vector{T}(undef, n)
    work.G = Matrix{T}(undef, q, n)
    work.g = Vector{T}(undef, q)
    work.M = Matrix{T}(undef, q, n)
    work.d = Vector{T}(undef, q)
    work.r = Vector{T}(undef, n + 1)
    work.e = Vector{T}(undef, n)
    work.A = Matrix{T}(undef, n + 1, q)
    work.AM = view(work.A, 1 : n, :)
    work.Ad = view(work.A, n + 1 : n + 1, :)
    work.b = Vector{T}(undef, n + 1)
    work.nnlswork = NNLSWorkspace{T, Int}(size(work.A)...)
    work.status = :Unsolved
    work
end

"""
    load!(work::QPWorkspace, Q, c, G, g)

Load problem data for the QP

Minimize ``\\frac{1}{2} z' Q z + c' z``
Subject to ``G z \\leq g``
"""
function load!(work::QPWorkspace{T}, Q::AbstractMatrix{T}, c::AbstractVector{T}, G::AbstractMatrix{T}, g::AbstractVector{T}) where T
    work.L .= Q
    work.c .= c
    work.G .= G
    work.g .= g
    work.status = :Unsolved
    nothing
end

load!(work::QPWorkspace{T}, qp::QP{T}) where {T} = load!(work, qp.Q, qp.c, qp.G, qp.g)
checkunsolved(work::QPWorkspace) = work.status == :Unsolved || error("Problem was already solved.")

"""
    z, λ = solve!(work::QPWorkspace)

Solve the QP that was loaded into `work` using `load!`. Returns the primal
solution ``z`` and the dual solution ``λ``.
"""
function solve!(work::QPWorkspace{T}, eps_infeasible = 1e-4) where T<:LinearAlgebra.BlasFloat
    checkunsolved(work)

    L = work.L
    c = work.c
    G = work.G
    g = work.g
    M = work.M
    d = work.d
    e = work.e
    A = work.A
    AM = work.AM
    Ad = work.Ad
    b = work.b
    r = work.r
    nnlswork = work.nnlswork

    # Compute M
    LinearAlgebra.LAPACK.potrf!('U', L) # L <- upper cholesky factor of Q
    M .= G
    LinearAlgebra.BLAS.trsm!('R', 'U', 'N', 'N', one(T), L, M) # M <- G L⁻¹

    # Compute d
    e .= c
    LinearAlgebra.LAPACK.potrs!('U', L, e) # e <- Q⁻¹ c

    d .= g
    LinearAlgebra.BLAS.gemv!('N', one(T), G, e, one(T), d) # d <- g + G Q⁻¹ c

    # Populate A
    transpose!(AM, M)
    transpose!(Ad, d)
    rmul!(A, -1)

    # Populate b
    γ = one(T)
    b .= 0
    b[end] = γ

    # Solve the NNLS
    load!(nnlswork, A, b)
    solve!(nnlswork)
    y = nnlswork.x

    # Compute the residual
    r = b
    LinearAlgebra.BLAS.gemv!('N', one(T), A, y, -one(T), r) # r <- A * y - b

    # Check for feasibility
    work.status = sum(abs, r) < eps_infeasible ? :Infeasible : :Optimal

    # Back out solution and dual
    z = c
    λ = y
    if work.status == :Optimal
        # Note: r[end] == -(γ + d ⋅ y)
        LinearAlgebra.BLAS.gemv!('T', 1 / r[end], G, y, -one(T), c) # z <- -1 / (γ + d ⋅ y) G^ᵀ y - c
        LinearAlgebra.LAPACK.potrs!('U', L, z)
        rmul!(λ, -1 / sqrt(-r[end])) # the sqrt appears to be missing in (12) in the paper
    else
        fill!(z, NaN)
        fill!(λ, NaN)
    end

    z, λ
end

function solve!(work::QPWorkspace{T}, eps_infeasible = 1e-4) where T
    checkunsolved(work)

    Q = work.L # is set to Q initially; modified to be L
    c = work.c
    G = work.G
    g = work.g
    M = work.M
    d = work.d
    e = work.e
    A = work.A
    AM = work.AM
    Ad = work.Ad
    b = work.b
    r = work.r
    nnlswork = work.nnlswork

    # Compute M
    cholQ = cholesky!(Hermitian(Q, :U))
    L = cholQ.U
    M[:] = G
    rdiv!(M, L)

    # Compute d
    e[:] = c
    ldiv!(cholQ, e) # e <- Q⁻¹ c
    mul!(d, G, e)
    d .+= g # d <- g + G Q⁻¹ c

    # Populate A
    transpose!(AM, M)
    transpose!(Ad, d)
    rmul!(A, -1)

    # Populate b
    γ = one(T)
    b .= 0
    b[end] = γ

    # Solve the NNLS
    load!(nnlswork, A, b)
    solve!(nnlswork)
    y = nnlswork.x

    # Compute the residual
    mul!(r, A, y)
    r .-= b # r <- A * y - b

    # Check for feasibility
    work.status = sum(abs, r) < eps_infeasible ? :Infeasible : :Optimal

    # Back out solution and dual
    z = similar(c) # TODO
    λ = y
    if work.status == :Optimal
        # Note: r[end] == -(γ + d ⋅ y)
        mul!(z, transpose(G), y)
        z .= z ./ r[end] .- c # z <- -1 / (γ + d ⋅ y) G^ᵀ y - c
        ldiv!(cholQ, z) # z <- -Q⁻¹ (1 / (γ + d ⋅ y) G^ᵀ y + c)
        rmul!(λ, -1 / sqrt(-r[end])) # the sqrt appears to be missing in (12) in the paper
    else
        fill!(z, NaN)
        fill!(λ, NaN)
    end

    z, λ
end

"""
    primal_infeasibility(qp::QP, z)

Measure of primal infeasibility of the solution to the QP

Minimize ``\\frac{1}{2} z' Q z + c' z``
Subject to ``G z \\leq g`

where ``z`` is the primal solution and ``λ`` is the dual solution.

For a primal-feasible solution, this will return a value <= 0.
"""
primal_infeasibility(qp::QP, z, λ) = maximum(qp.G * z .- qp.g)

"""
    dual_infeasibility(qp::QP, λ)

Measure of primal infeasibility of the solution to the QP

Minimize ``\\frac{1}{2} z' Q z + c' z``
Subject to ``G z \\leq g`

where ``z`` is the primal solution and ``λ`` is the dual solution.

For a dual-feasible solution, this will return a value <= 0.
"""
dual_infeasibility(qp::QP, z, λ) = maximum(λ)

"""
    stationarity_violation(qp::QP, z, λ)

Measure of the violation of the stationarity KKT condition for the QP.

For an optimal solution, this should return a value near 0.
"""
function stationarity_violation(qp::QP, z, λ)
    slack_grad = qp.G' * λ
    cost_grad = qp.Q * z + qp.c
    scale = mean(cost_grad ./ slack_grad)
    maximum(abs, cost_grad .- scale .* slack_grad)
end

"""
    slackness_violation(qp::QP, z, λ)

Measure of the violation of the complementary slackness condition for the QP.

For an optimal solution, this should return a value near 0.
"""
slackness_violation(qp::QP, z, λ) = maximum(abs, (qp.G * z .- qp.g) .* λ)

"""
    check_optimality_conditions(qp::QP, z, λ)

Provies a scalar indicator of the violation of the KKT optimality conditions for
a given QP, primal solution ``z``, and dual solution ``λ``. A feasible & optimal
solution should return a value close to 0, while an infeasible or suboptimal
solution should return a value greater than zero.
"""
function check_optimality_conditions(qp::QP, z, λ)
    return max(
               primal_infeasibility(qp, z, λ),
               dual_infeasibility(qp, z, λ),
               stationarity_violation(qp, z, λ),
               slackness_violation(qp, z, λ)
               )
end

include("NNLSSolverInterface.jl")
using .NNLSSolverInterface: NNLSSolver

end # module


using QSM: Poisson2
using QSM: wrapped_laplacian
using QSM: wrapped_laplacian_boundary_neumann!
using QSM: wrapped_laplacian_boundary_periodic!
using QSM: solve_poisson_dct, solve_poisson_fft, solve_poisson_mgpcg

L2norm(x) = norm(x) / sqrt(length(x))

function test_problem(::Type{T}, sz::NTuple{3, Integer}) where {T<:AbstractFloat}
    f = sinpi

    R = CartesianIndices(map(n -> 2:n-1, sz))
    m = zeros(Bool, sz)
    u = zeros(T, sz)
    d2u = zeros(T, sz)
    d2uh = zeros(T, sz)

    X = range(0, 1, length=sz[1])
    Y = range(0, 1, length=sz[2])
    Z = range(0, 1, length=sz[3])

    dx = (X.step.hi, Y.step.hi, Z.step.hi)
    _u = [f(x).*f(y).*f(z) for x in X, y in Y, z in Z]

    m[R] .= true
    u[R] .= _u[R]
    d2u[R] .= -3*π^2*_u[R]

    A = Poisson2(T, m, dx)
    d2uh = mul!(d2uh, A, u)

    return A, u, d2u, d2uh
end

@testset "Multigrid" begin
    include("multigrid/multigrid.jl")
end

@testset "DCT with wrapped Laplacian" for T in (Float32, Float64)
    Random.seed!(501)
    for sz in [(11, 12, 13), (12, 13, 14, 3)]
        x, dx = rand(T, sz), Tuple(rand(T, 3))
        x .-= mean(x, dims=(1,2,3))
        d2x = wrapped_laplacian(x, dx)
        d2x = wrapped_laplacian_boundary_neumann!(d2x, x, dx)
        @test solve_poisson_dct(d2x, dx) ≈ x
    end
end

@testset "FFT with wrapped Laplacian" for T in (Float32, Float64)
    Random.seed!(501)
    for sz in [(11, 12, 13), (12, 13, 14, 3)]
        x, dx = rand(T, sz), Tuple(rand(T, 3))
        x .-= mean(x, dims=(1,2,3))
        d2x = wrapped_laplacian(x, dx)
        d2x = wrapped_laplacian_boundary_periodic!(d2x, x, dx)
        @test solve_poisson_fft(d2x, dx) ≈ x
    end

end

@testset "MGPCG with wrapped_laplacian" for T in (Float32, Float64)
    Random.seed!(501)
    for sz in [(127, 128, 129), (130, 131, 132)]
        A, u, d2u, d2uh = test_problem(T, sz)
        @test wrapped_laplacian(u, A.dx) ≈ d2uh
        @test solve_poisson_mgpcg(d2uh, A.interior, A.dx; rtol=eps(T)) ≈ u
    end
end


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   Functions related with True Equator Mean Equinox (TEME) reference frame.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Remarks
#
# As mentioned in [1, p. 233], there is not an official definition of the TEME
# frame. Hence, in this package, it is considered the definition presented in
# [1, p. 233] in which the complete form of the Equation of Equinoxes is used.
# This seems to be the case when comparing the values shown in Table 3-6 [1, p.
# 232].
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export rTEMEtoTOD,  rTODtoTEME
export rTEMEtoMOD,  rMODtoTEME
export rTEMEtoGCRF, rGCRFtoTEME
export rTEMEtoPEF,  rPEFtoTEME

#                                 TEME <=> TOD
# ==============================================================================

"""
    rTEMEtoTOD([T,] JD_TT::Number [, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0])

Compute the rotation that aligns the True Equator Mean Equinox (TEME) frame with
the True of Date (TOD) frame at the Julian Day `JD_TT` [Terrestrial Time]. This
algorithm uses the IAU-76/FK5 theory and TEME definition in [1, p. 233]. Notice
that one can provide corrections for the nutation in obliquity (`δΔϵ_1980`)
\\[rad] and in longitude (`δΔψ_1980`) \\[rad] that are usually obtained from
IERS EOP Data (see `get_iers_eop`).

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the TEME frame with the TOD frame. The rotation
representation is selected by the optional parameter `T`.

"""
rTEMEtoTOD(JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    rTEMEtoTOD(DCM, JD_TT, δΔϵ_1980, δΔψ_1980)

function rTEMEtoTOD(T::Type, JD_TT::Number, δΔϵ_1980::Number = 0,
                    δΔψ_1980::Number = 0)

    # Compute the nutation in the Julian Day (Terrestrial Time) `JD_TT`.
    (mϵ_1980, Δϵ_1980, Δψ_1980) = nutation_fk5(JD_TT)

    # Add the corrections to the nutation in obliquity and longitude.
    Δϵ_1980 += δΔϵ_1980
    Δψ_1980 += δΔψ_1980

    # Compute the obliquity.
    ϵ_1980 = mϵ_1980 + Δϵ_1980

    # Evaluate the Delaunay parameters associated with the Moon in the interval
    # [0,2π]°.
    #
    # The parameters here were updated as stated in the errata [2].
    T_TT = (JD_TT - JD_J2000)/36525
    r    = 360
    Ω_m  = 125.04452222 - (5r + 134.1362608)*T_TT +
                          0.0020708*T_TT^2 +
                          2.2e-6*T_TT^3
    Ω_m  = mod(Ω_m, 360)*pi/180

    # Compute the equation of Equinoxes.
    #
    # According to [2], the constant unit before `sin(2Ω_m)` is also in [rad].
    Eq_equinox1982 = Δψ_1980*cos(mϵ_1980) +
        ( 0.002640*sin(1Ω_m) + 0.000063*sin(2Ω_m) )*pi/648000

    # Compute the rotation.
    return angle_to_rot(T, -Eq_equinox1982, 0, 0, :ZYX)
end

"""
    rTODtoTEME([T,] JD_TT::Number [, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0])

Compute the rotation that aligns the True of Date (TOD) frame with the True
Equator Mean Equinox (TEME) frame at the Julian Day `JD_TT` [Terrestrial Time].
This algorithm uses the IAU-76/FK5 theory and TEME definition in [1, p. 233].
Notice that one can provide corrections for the nutation in obliquity
(`δΔϵ_1980`) \\[rad] and in longitude (`δΔψ_1980`) \\[rad] that are usually
obtained from IERS EOP Data (see `get_iers_eop`).

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the TOD frame with the TEME frame. The rotation
representation is selected by the optional parameter `T`.

"""
rTODtoTEME(JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    rTODtoTEME(DCM, JD_TT, δΔϵ_1980, δΔψ_1980)

rTODtoTEME(T::Type, JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    inv_rotation(rTEMEtoTOD(T, JD_TT, δΔϵ_1980, δΔψ_1980))

#                                 TEME <=> MOD
# ==============================================================================

"""
    rTEMEtoMOD([T,] JD_TT::Number [, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0])

Compute the rotation that aligns the True Equator Mean Equinox (TEME) frame with
the Mean of Date (MOD) frame at the Julian Day `JD_TT` [Terrestrial Time]. This
algorithm uses the IAU-76/FK5 theory and TEME definition in [1, p. 233]. Notice
that one can provide corrections for the nutation in obliquity (`δΔϵ_1980`)
\\[rad] and in longitude (`δΔψ_1980`) \\[rad] that are usually obtained from
IERS EOP Data (see `get_iers_eop`).

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the TEME frame with the MOD frame. The rotation
representation is selected by the optional parameter `T`.

"""
rTEMEtoMOD(JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    rTEMEtoMOD(DCM, JD_TT, δΔϵ_1980, δΔψ_1980)

function rTEMEtoMOD(T::Type, JD_TT::Number, δΔϵ_1980::Number = 0,
                    δΔψ_1980::Number = 0)

    # Notice that, in this case, we will not use `rTEMEtoTOD`, and `rTODtoMOD`
    # because this would call the function `nutation` twice, leading to a huge
    # performance drop. Hence, the code of those two functions is almost
    # entirely rewritten here.

    # Compute the nutation in the Julian Day (Terrestrial Time) `JD_TT`.
    (mϵ_1980, Δϵ_1980, Δψ_1980) = nutation_fk5(JD_TT)

    # Add the corrections to the nutation in obliquity and longitude.
    Δϵ_1980 += δΔϵ_1980
    Δψ_1980 += δΔψ_1980

    # Compute the obliquity.
    ϵ_1980 = mϵ_1980 + Δϵ_1980

    # Evaluate the Delaunay parameters associated with the Moon in the interval
    # [0,2π]°.
    #
    # The parameters here were updated as stated in the errata [2].
    T_TT = (JD_TT - JD_J2000)/36525
    r    = 360
    Ω_m  = 125.04452222 - (5r + 134.1362608)*T_TT +
                          0.0020708*T_TT^2 +
                          2.2e-6*T_TT^3
    Ω_m  = mod(Ω_m, 360)*pi/180

    # Compute the equation of Equinoxes.
    #
    # According to [2], the constant unit before `sin(2Ω_m)` is also in [rad].
    Eq_equinox1982 = Δψ_1980*cos(mϵ_1980) +
        ( 0.002640*sin(1Ω_m) + 0.000063*sin(2Ω_m) )*pi/648000

    # Compute the rotation TEME => TOD.
    r_TOD_TEME = angle_to_rot(T, -Eq_equinox1982, 0, 0, :ZYX)

    # Compute the rotation TOD => MOD.
    r_MOD_TOD = angle_to_rot(T, ϵ_1980, Δψ_1980, -mϵ_1980, :XZX)

    # Return the full rotation.
    return compose_rotation(r_TOD_TEME, r_MOD_TOD)
end

"""
    rMODtoTEME([T,] JD_TT::Number [, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0])

Compute the rotation that aligns the Mean of Date (MOD) frame with the True
Equator Mean Equinox (TEME) frame at the Julian Day `JD_TT` [Terrestrial Time].
This algorithm uses the IAU-76/FK5 theory and TEME definition in [1, p. 233].
Notice that one can provide corrections for the nutation in obliquity
(`δΔϵ_1980`) \\[rad] and in longitude (`δΔψ_1980`) \\[rad] that are usually
obtained from IERS EOP Data (see `get_iers_eop`).  .

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the MOD frame with the TEME frame. The rotation
representation is selected by the optional parameter `T`.

"""
rMODtoTEME(JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    rMODtoTEME(DCM, JD_TT, δΔϵ_1980, δΔψ_1980)

rMODtoTEME(T::Type, JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    inv_rotation(rTEMEtoMOD(T, JD_TT, δΔϵ_1980, δΔψ_1980))

#                                TEME <=> GCRF
# ==============================================================================

"""
    rTEMEtoGCRF([T,] JD_TT::Number [, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0])

Compute the rotation that aligns the True Equator Mean Equinox (TEME) frame with
the Geocentric Celestial Reference Frame (GCRF) at the Julian Day `JD_TT`
[Terrestrial Time]. This algorithm uses the IAU-76/FK5 theory and TEME
definition in [1, p. 233]. Notice that one can provide corrections for the
nutation in obliquity (`δΔϵ_1980`) \\[rad] and in longitude (`δΔψ_1980`) \\[rad]
that are usually obtained from IERS EOP Data (see `get_iers_eop`).

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the TEME frame with the GCRF frame. The rotation
representation is selected by the optional parameter `T`.

# Remarks

The EOP data related to the nutation of the obliquity (`δΔϵ_1980`) and the
nutation of the longitude (`δΔψ_1980`) can be omitted. In this case, the GCRF
frame is what is usually called J2000 reference frame.

"""
rTEMEtoGCRF(JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    rTEMEtoGCRF(DCM, JD_TT, δΔϵ_1980, δΔψ_1980)

function rTEMEtoGCRF(T::Type, JD_TT::Number, δΔϵ_1980::Number = 0,
                     δΔψ_1980::Number = 0)

    # Compute the rotation TEME => MOD.
    r_MOD_TEME  = rTEMEtoMOD(T, JD_TT, δΔϵ_1980, δΔψ_1980)

    # Compute the rotation MOD => GCRF.
    r_GCRF_MOD = rMODtoGCRF_fk5(T, JD_TT)

    # Return the full rotation.
    return compose_rotation(r_MOD_TEME, r_GCRF_MOD)
end

"""
    rGCRFtoTEME([T,] JD_TT::Number [, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0])

Compute the rotation that aligns the GCRF frame with the True Equator Mean
Equinox (TEME) frame at the Julian Day `JD_TT` [Terrestrial Time]. This
algorithm uses the IAU-76/FK5 theory and TEME definition in [1, p. 233]. Notice
that one can provide corrections for the nutation in obliquity (`δΔϵ_1980`)
\\[rad] and in longitude (`δΔψ_1980`) \\[rad] that are usually obtained from
IERS EOP Data (see `get_iers_eop`).

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the GCRF frame with the TEME frame. The rotation
representation is selected by the optional parameter `T`.

# Remarks

The EOP data related to the nutation of the obliquity (`δΔϵ_1980`) and the
nutation of the longitude (`δΔψ_1980`) can be omitted. In this case, the GCRF
frame is what is usually called J2000 reference frame.

"""
rGCRFtoTEME(JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
    rGCRFtoTEME(DCM, JD_TT, δΔϵ_1980, δΔψ_1980)

rGCRFtoTEME(T::Type, JD_TT::Number, δΔϵ_1980::Number = 0, δΔψ_1980::Number = 0) =
   inv_rotation(rTEMEtoGCRF(T, JD_TT, δΔϵ_1980, δΔψ_1980))

#                                 TEME <=> PEF
# ==============================================================================

"""
    rTEMEtoPEF([T,] JD_TT::Number)

Compute the rotation that aligns the True Equator Mean Equinox (TEME) frame with
the Pseudo-Earth Fixed (PEF) frame at the Julian Day `JD_TT` [Terrestrial Time].
This algorithm uses the IAU-76/FK5 theory and TEME definition in [1, p. 233].

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the TEME frame with the PEF frame. The rotation
representation is selected by the optional parameter `T`.

"""
rTEMEtoPEF(JD_UT1::Number) = rTEMEtoPEF(DCM, JD_UT1)

function rTEMEtoPEF(T::Type, JD_UT1::Number)
    # Compute the Greenwich Mean Sidereal Time.
    θ_gmst = JDtoGMST(JD_UT1)

    # Compute the rotation.
    return angle_to_rot(T, θ_gmst, 0, 0, :ZYX)
end

"""
    rPEFtoTEME([T,] JD_TT::Number)

Compute the rotation that aligns the Pseudo-Earth Fixed (PEF) frame with the
True Equator Mean Equinox (TEME) frame at the Julian Day `JD_TT` [Terrestrial
Time]. This algorithm uses the IAU-76/FK5 theory and TEME definition in [1, p.
233].

The rotation type is described by the optional variable `T`. If it is `DCM`,
then a DCM will be returned. Otherwise, if it is `Quaternion`, then a Quaternion
will be returned. In case this parameter is omitted, then it falls back to
`DCM`.

# Returns

The rotation that aligns the PEF frame with the TEME frame. The rotation
representation is selected by the optional parameter `T`.

"""
rPEFtoTEME(JD_UT1::Number)          = rPEFtoTEME(DCM, JD_UT1)
rPEFtoTEME(T::Type, JD_UT1::Number) = inv_rotation(rTEMEtoPEF(T, JD_UT1))


# calculate required particle num for kld sampling

module ParticleNumKld
  using Plots, Distributions
  pyplot()

  function num(epsilon, delta, bin_num)
    return ceil(quantile(Chisq(bin_num-1), 1-delta)/(2*epsilon))
  end

  function main(is_test=false)
    bs = 2:1:9
    n = [num(0.1, 0.01, b) for b in bs]
    p1 = plot(bs, n, title="bin: 2-10")

    bs = 2:1:9999
    n = [num(0.1, 0.01, b) for b in bs]
    p2 = plot(bs, n, title="bin: 2-100000")

    plot(p1, p2, layout=(1, 2), size=(750, 300), legend=false)

    if is_test == false
      save_path = joinpath(split(@__FILE__, "src")[1], "src/localization/particle_filter/kld_sampling/particle_num_kld.png")
      savefig(save_path)
    end
  end
end


using .Curve: BN, FP, EP, EP2, LIB, Limb
using .Util: Point
using .Config: PRIME, ORDER, @ID

evalpoly(poly, x) = evalpoly(poly, BN(x))
function evalpoly(poly, x::BN)
    @assert length(poly) > 0
    xpow, y = x, first(poly)
    for coeff in Iterators.drop(poly, 1)
        y += xpow * coeff
        xpow = mod(x * xpow, ORDER)
    end
    return isa(y, Point) ? y : mod(y, ORDER)
end

function lagrange_interpolate_c0(coeffs::LagrangeCoeffGenerator, shares::AbstractDict{@ID,V})  where {V<:Union{BN,EP,EP2}}
    length(shares) >= 2 || error("need at least two shares")
    c0 = zero(V)
    for (x, y) in shares
        c0 += coeffs[x] * y
    end
    return isa(c0, Point) ? c0 : mod(c0, ORDER)
end

function lagrange_interpolate_c0(weights::BarycentricWeightGenerator, shares::AbstractDict{@ID,V}) where {V<:Union{BN,EP,EP2}}
    length(shares) >= 2 || error("need at least two shares")
    num, denom = zero(V), zero(BN)
    for (x, y) in shares
        # optimisation: normally inv is done when calculating weights
        #w = invmod(weights[j], ORDER)
        w = invmod(weights[x], ORDER)
        #@info weights[j], w
        num += w * y
        # just addition in practice no risk of overflow
        denom += w
    end
    c0 = invmod(mod(denom, ORDER), ORDER) * num
    return isa(c0, Point) ? c0 : mod(c0, ORDER)
end

lagrange_interpolate_c0(shares::AbstractDict) =
    lagrange_interpolate_c0(LagrangeCoeffGenerator(keys(shares)), shares)


module IonizationRates
using SpecialFunctions
using HCubature

# * Common

Γ(x) = SpecialFunctions.gamma(x)

@doc raw"""
```math
\begin{equation}
\tag{ADK19}
\abs{C_{n^*\ell^*}}^2 =
\frac{2^{2n^*}}{n^*\Gamma(n^*+\ell^*+1)\Gamma(n^*-\ell^*)}
\end{equation}
```
"""
C2(n, ℓ) = 2.0^(2n)/(n*Γ(n+ℓ+1)*Γ(n-ℓ))

@doc raw"""
    Etilde(E0, Iₚ)
```math
\tilde{E} = \frac{E_0}{2(2I_p)^{3/2}} \equiv
\frac{\sqrt{I}}{2(2I_p)^{3/2}}.
```
"""
Etilde(E0, Iₚ) = E0/(2(2Iₚ)^1.5)

@doc raw"""
    f(ℓ, m)

Common prefactor appearing in e.g Eqs. (PPT54,59), (ADK12)

```math
f(\ell, m) =
\frac{(2\ell+1)(\ell+\abs{m})!}
{2^{\abs{m}}(\abs{m})!(\ell-\abs{m})!}
```
"""
function f(ℓ::Integer,m::Integer)
    ℓ >= 0 || error("Invalid value of ℓ ($(ℓ) < 0)")
    abs(m) <= ℓ || error("Invalid value of m ($(abs(m)) > $(ℓ))")
    (2ℓ+1)*(factorial(ℓ+abs(m)) /
            (2^abs(m)*factorial(abs(m))*factorial(ℓ-abs(m))))
end

Uₚ(I,ω) = I/4ω^2

keldysh(Iₚ, Uₚ) = √(Iₚ/2Uₚ)
keldysh(Iₚ, I, ω) = keldysh(Iₚ, Uₚ(I, ω))

@doc raw"""
    effective_n(Iₚ, Z)

Effective principal quantum number for ionization potential `Iₚ` and
residual nuclear charge `Z`.

```math
\begin{equation}
\tag{ADK2}
n^* = \frac{Z}{\sqrt{2I_p}}
\end{equation}
```
"""
effective_n(Iₚ, Z) =
    Z/√(2Iₚ)

# * PPT

@doc raw"""
    β(γ)

```math
\begin{equation}
\tag{PPT56ff}
\beta(\gamma) =
\frac{2\gamma}{\sqrt{1+\gamma^2}}
\end{equation}
```

Defined just below Eq. (PPT56).
"""
β(γ) = 2γ/√(1+γ^2)

@doc raw"""
    α(γ)

```math
\begin{equation}
\tag{PPT34}
\alpha(\gamma) =
2\left(
\arcsinh\gamma -
\frac{\gamma}{\sqrt{1+\gamma^2}}
\right) \equiv
2\arcsinh\gamma -
\beta(\gamma)
\end{equation}
```
"""
α(γ) = 2asinh(γ) - β(γ)

@doc raw"""
    g(γ)

```math
\begin{equation}
\tag{PPT33}
g(\gamma) =
\frac{3}{2\gamma}
\left[
\left(
1+\frac{1}{2\gamma^2}
\right)
\arcsinh\gamma-
\frac{\sqrt{1+\gamma^2}}{2\gamma}
\right] \equiv
\frac{3}{2\gamma}
\left[
\left(
1+\frac{1}{2\gamma^2}
\right)
\arcsinh\gamma-
\frac{1}{\beta(\gamma)}
\right]
\end{equation}
```
"""
g(γ) = 3/(2γ)*((1+1/(2γ^2))*asinh(γ) - 1/β(γ))

@doc raw"""
    ν(Iₚ, I, ω)

Number of photons of energy ``\omega`` needed to reach above the
ionization potential ``I_p`` in the presence of a strong dressing
field of intensity ``I``.

```math
\begin{equation}
\tag{PPT26}
\nu = \frac{I_p}{\omega}
\left(1+\frac{1}{2\gamma^2}\right)
\end{equation}
```
"""
ν(Iₚ, I, ω) = Iₚ/ω * (1.0 + 1/2keldysh(Iₚ,I,ω)^2)

@doc raw"""
    w(m,x)

```math
\begin{equation}
\tag{PPT56}
w_m(x) =
\frac{x^{2\abs{m}+1}}{2}
\int_0^1\diff{t}
\frac{\exp(-x^2t)t^{\abs{m}}}{\sqrt{1-t}}
\end{equation}
```
"""
function w(m,x)
    val, = hquadrature(t -> exp(-x^2*t)*t^abs(m)/√(1-t), 0, 1)
    x^(2abs(m)+1)/2*val
end

@doc raw"""
    A(m, Iₚ, I, ω[; tol = √(eps()), maxterms=typemax(Int)])

```math
\begin{equation}
\tag{PPT55}
A_m(\omega,\gamma) =
\frac{4}{\sqrt{3\pi}}
\frac{1}{\abs{m}!}
\frac{\gamma^2}{1+\gamma^2}
\sum_{n\geq\nu}^\infty
\exp[-\alpha(n-\nu)]
w_m[\sqrt{\beta(n-\nu)}],
\end{equation}
```
where ``\nu`` is computed using [`ν`](@ref).
"""
function A(m, Iₚ, I, ω; tol = √(eps()), maxterms=typemax(Int))
    γ = keldysh(Iₚ,I,ω)
    nν = ν(Iₚ, I, ω)
    S = 0
    dS = Inf
    kappa = ceil(Int, nν)
    ii = 0
    α_γ = α(γ)
    β_γ = β(γ)
    while dS/S > tol && ii < maxterms
        dk = kappa - nν
        dS = exp(-α_γ*dk)*w(m,√(β_γ*dk))
        S += dS
        kappa += 1
        ii += 1
    end
    (4/√(3π))/factorial(abs(m)) * γ^2/(1+γ^2) * S
end

@doc raw"""
    PPT(Iₚ, I, ω, ℓ, m[, Z=1])

Compute the strong-field photoionization rate from an initial state
with ionization potential `Iₚ`, angular quantum numbers `ℓ` & `m`, and
a binding potential of charge `Z`, when subjected to monochromatic
light of intensity `I` and angular frequency `ω`, using the PPT
formalism (slightly generalized to effective principal quantum numbers
``n^*``):

```math
\begin{equation}
\tag{PPT54,ADK1}
w_{\ell m}(F, \omega) =
I_p
\abs{C_{n^*\ell^*}}^2
\sqrt{\frac{3}{2\pi}}
f(\ell, m)
\tilde{E}^{-(2n^* - \abs{m} - 3/2)}
(1+\gamma^2)^{-n^* + \abs{m}/2 + 3/4}
A_m(\omega,\gamma)
\exp\left[
-\frac{g(\gamma)}{3\tilde{E}}
\right],
\end{equation}
```
where ``\gamma`` is the Keldysh parameter, and ``n^*`` the [`effective_n`](@ref).
"""
function PPT(Iₚ, I, ω, ℓ, m, Z=1)
    γ = keldysh(Iₚ, I, ω)
    n⭑ = effective_n(Iₚ, Z)

    pre = √(3/2π)*C2(n⭑,n⭑-1)*f(ℓ,m)*Iₚ
    Ẽ = Etilde(√(I), Iₚ)

    Am = A(0, Iₚ, I, ω)

    w = pre*Ẽ^(-(2n⭑-abs(m)-3/2))*(1+γ^2)^(-n⭑+abs(m)/2+3/4)*Am*exp(-g(γ)/(3Ẽ))

    isnan(w) ? zero(w) : w
end

end


function e002(N)
    return sum_evens(N)
end

function fib(channel::Channel, num0::Integer, num1::Integer)
    put!(channel, num0)
    put!(channel, num1)
    while true
        sum = num0 + num1
        num0 = num1
        num1 = sum
        put!(channel, sum)
    end
end

function sum_evens(N)
    channel = Channel(c -> fib(c,1,2))
    i = take!(channel)
    sum = 0
    while i < N
        if i % 2 == 0
            sum += i
        end
        i = take!(channel)
    end
    close(channel)
    return sum
end

function solution002()
    N = 4000000
    answer = sum_evens(N)
    println("The answer to problem 002 is $answer")
end


using IntrinsicPersistence
using Test

using StaticArrays
using LinearAlgebra

# sanitycheck
@testset "Sanity checks" begin
    @testset "regular circle" begin
        n = 12
        pts = [SVector(cos(t), sin(t)) for t in range(0, stop = 2π, length = n+1)[1:end-1]]
        landmarks = 1:2:n

        gc = GeodesicComplex(pts, 0.6, landmarks = landmarks)
        res = persistence(gc)
        @test length(res) == 1
        α = first(res).indices
        @test sort(α) == landmarks
    end

    @testset "circle with contraction" begin
        n = 24
        pts = vcat(
            [SVector(1.0cos(t), 1.0sin(t)) for t in range(0, stop = 2π, length = n+1)[1:end-1]],
            [SVector(0.9cos(t), 0.9sin(t)) for t in range(0, stop = 2π, length = n+1)[1:end-1]],
            [SVector(0.8cos(t), 0.8sin(t)) for t in range(0, stop = 2π, length = n+1)[1:end-1]],
            [SVector(0.7cos(t), 0.7sin(t)) for t in range(0, stop = 2π, length = n+1)[1:end-1]],
            [SVector(0.6cos(t), 0.6sin(t)) for t in range(0, stop = 2π, length = n+1)[1:end-1]],
            [SVector(0.5cos(t), 0.5sin(t)) for t in range(0, stop = 2π, length = n+1)[1:end-1]],
        )
        landmarks = 1:2:n
        gc = GeodesicComplex(pts, 0.3, landmarks = landmarks)
        res = persistence(gc)
        @test length(res) == 1
        @test diameter(first(res)) ≈ 1
    end

    @testset "landmarks form a circle, but the underlying space is contractible" begin
        n = 12
        pts = vcat([SVector(0.5cos(t), 0.5sin(t))
                    for t in range(0, stop = 2π, length = n÷2+1)[1:end-1]],
                   [SVector(cos(t), sin(t))
                    for t in range(0, stop = 2π, length = n+1)[1:end-1]],
                   [SVector(0.0, 0.0)])
        landmarks = n÷2+1:2:n+n÷2
        gc = GeodesicComplex(pts, 0.6, landmarks = landmarks)
        res = persistence(gc)
        @test isempty(res)

        res = persistence(gc, shrinked = false)
        @test !isempty(res)
    end

    @testset "cylinder" begin
        n = 12
        pts = [SVector(cos(t), sin(t), h)
               for t in range(0, stop = 2π, length = n+1)[1:end-1]
               for h in 0:0.5:10]
        landmarks = 1:2:length(pts)
        gc = GeodesicComplex(pts, 0.6, landmarks = landmarks)
        res = persistence(gc)
        @test length(res) == 1
    end

    @testset "random full graph" begin
        n = 100
        pts = [SVector(rand(), rand(), rand()) for _ in 1:n]
        landmarks = 1:n
        gc = GeodesicComplex(pts, 2, landmarks = landmarks)
        res = persistence(gc)
        @test isempty(res)
    end

    @testset "regular cone" begin
        n = 12
        pts = [SVector(h/10*cos(t), h/10*sin(t), h)
               for t in range(0, stop = 2π, length = n+1)[1:end-1]
               for h in 0:0.5:10]
        landmarks = 1:2:length(pts)
        gc = GeodesicComplex(pts, 0.6, landmarks = landmarks)
        res = persistence(gc)
        @test isempty(res)
    end

    @testset "random cone" begin
        n = 100_000
        pts = map(rand() for _ in 1:n) do h
            t = rand()
            SVector(h*sinpi(2t), h*cospi(2t), h)
        end
        gc = GeodesicComplex(pts, 0.2)
        res = persistence(gc)
        @test isempty(res)
    end

    @testset "random truncated cone" begin
        n = 100_000
        pts = map(rand() for _ in 1:n) do h
            t = rand()
            SVector((h+1)*sinpi(2t), (h+1)*cospi(2t), h)
        end
        gc = GeodesicComplex(pts, 0.2)
        res = persistence(gc)
        @test length(res) == 1
        @test diameter(first(res)) ≈ 2 atol = 0.2
        @test perimeter(first(res)) ≈ 2π atol = 0.2
    end
end


# # Continuous Bag of Words (CBOW) Model

include("bow.jl")

# Instead of mapping words directly to 5 dimensions for 5 classes, we map them to 64
# dimensions to represent features, then go to 5 dimensions after adding up these features.

pred(words) = b_sm .+ W_sm * sum(W_emb[:,words], dims=2)

# Only the prediction function and parameters need to change, the rest of the code is the
# same.

EMB_SIZE = 64
W_emb = param(EMB_SIZE, nwords)
W_sm = param(ntags, EMB_SIZE)
b_sm = param(ntags)

# Bonus question: can this model represent a function that the BOW model cannot represent?

train()


using Test, FluxReconstruction, KitBase

cd(@__DIR__)
include("test_triangle.jl")

let u = rand(3), u0 = rand(3)
    FR.L1_error(u, u0, 0.1)
    FR.L2_error(u, u0, 0.1)
    FR.L∞_error(u, u0, 0.1)
end

shock_detector(-Inf, 3)
shock_detector(log10(0.1), 3)
shock_detector(log10(1e5), 3)

deg = 5
ps2 = FRPSpace2D(0.0, 1.0, 20, 0.0, 1.0, 20, deg, 1, 1)

rs_jacobi(ps2.xpl, [0 0; √3 -1; √3+1 √3-1; 1 √3])
rs_jacobi(ps2.xpl, rand(3, 3, 4, 2))

ps1 = TriFRPSpace("../assets/linesource.msh", 2)
ps = FRPSpace1D(0.0, 1.0, 20, deg)

positive_limiter(ones(6), 1 / 6, ps.ll, ps.lr)
positive_limiter(ones(6, 3), 5 / 3, 1 / 6, ps.ll, ps.lr)

let u = rand(deg + 1)
    ℓ = FR.basis_norm(deg)

    modal_filter!(u, 1e-6; filter = :l2)
    modal_filter!(u, 1e-6; filter = :l2opt)
    modal_filter!(u, 1e-6, ℓ; filter = :l1)
    modal_filter!(u, ℓ; filter = :lasso)
    modal_filter!(u, 10; filter = :exp)
    modal_filter!(u, 10; filter = :houli)

    # 2D exponential filter
    FR.filter_exp(2, 10, Array(ps.V))
end

let u = rand(deg+1, deg+1)
    modal_filter!(u, 1e-6, 1e-6; filter = :l2)
    modal_filter!(u, 1e-6, 1e-6; filter = :l2opt)
end

f = randn(5, deg + 1)
fδ = randn(5, 2)
FR.interp_face!(fδ, f, ps.ll, ps.lr)

FR.standard_lagrange(ps.xpl)


#
# morphology.jl --
#
# Implementation of non-linear morphological filters.
#
#------------------------------------------------------------------------------
#
# This file is part of the `LocalFilters.jl` package licensed under the MIT
# "Expat" License.
#
# Copyright (C) 2017-2018, Éric Thiébaut.
#

"""
Basic operations of mathematical morphology are:

```julia
erode(A, B) -> Amin
dilate(A, B) -> Amax
```

which respectively return the local minima `Amin` and the local maxima `Amax`
of argument `A` in a neighborhood defined by `B`.  The returned result is
similar to `A` (same size and type).

The two operations can be combined in one call:

```julia
localextrema(A, B) -> Amin, Amax
```

The in-place versions:

```julia
erode!(Amin, A, B) -> Amin
dilate!(Amax, A, B) -> Amax
localextrema!(Amin, Amax, A, B) -> Amin, Amax
```

apply the operation to `A` with structuring element `B` and store the
result in the provided arrays `Amin` and/or `Amax`.

See also [`localmean`](@ref), [`opening`](@ref), [`closing`](@ref),
[`top_hat`](@ref) and [`bottom_hat`](@ref).

"""
erode(A::AbstractArray, args...) = erode!(similar(A), A, args...)

erode!(dst, src::AbstractArray{T,N}, B=3) where {T,N} =
    erode!(dst, src, Neighborhood{N}(B))

@doc @doc(erode) erode!

function erode!(dst::AbstractArray{T,N},
                A::AbstractArray{T,N},
                B::RectangularBox{N}) where {T,N}
    localfilter!(dst, A, :, min, axes(B))
end

function erode!(A::AbstractArray{T,N},
                B::RectangularBox{N}) where {T,N}
    localfilter!(A, :, min, axes(B))
end

function erode!(dst::AbstractArray{T,N},
                A::AbstractArray{T,N},
                B::Kernel{Bool,N}) where {T,N}
    if ismmbox(B)
        localfilter!(dst, A, :, min, axes(B))
    else
        @assert axes(dst) == axes(A)
        localfilter!(dst, A, B,
                     (a)     -> typemax(T),
                     (v,a,b) -> b && a < v ? a : v,
                     (d,i,v) -> d[i] = v)
    end
    return dst
end

function erode!(dst::AbstractArray{T,N},
                A::AbstractArray{T,N},
                B::Kernel{K,N}) where {T<:AbstractFloat,
                                       K<:AbstractFloat,N}
    if ismmbox(B)
        localfilter!(dst, A, :, min, axes(B))
    else
        @assert axes(dst) == axes(A)
        localfilter!(dst, A, Kernel{T}(B),
                     (a)     -> typemax(T),
                     (v,a,b) -> min(v, a - b),
                     (d,i,v) -> d[i] = v)
    end
    return dst
end

dilate(A::AbstractArray, args...) = dilate!(similar(A), A, args...)

@doc @doc(erode) dilate

dilate!(dst, src::AbstractArray{T,N}, B=3) where {T,N} =
    dilate!(dst, src, Neighborhood{N}(B))

@doc @doc(dilate) dilate!

function dilate!(dst::AbstractArray{T,N},
                 A::AbstractArray{T,N},
                 B::RectangularBox{N}) where {T,N}
    localfilter!(dst, A, :, max, axes(B))
end

function dilate!(A::AbstractArray{T,N},
                 B::RectangularBox{N}) where {T,N}
    localfilter!(A, :, max, axes(B))
end

function dilate!(dst::AbstractArray{T,N},
                 A::AbstractArray{T,N},
                 B::Kernel{Bool,N}) where {T,N}
    if ismmbox(B)
        localfilter!(dst, A, :, max, axes(B))
    else
        @assert axes(dst) == axes(A)
        localfilter!(dst, A, B,
                     (a)     -> typemin(T),
                     (v,a,b) -> b && a > v ? a : v,
                     (d,i,v) -> d[i] = v)
    end
    return dst
end

function dilate!(dst::AbstractArray{T,N},
                 A::AbstractArray{T,N},
                 B::Kernel{K,N}) where {T<:AbstractFloat,
                                        K<:AbstractFloat,N}
    if ismmbox(B)
        localfilter!(dst, A, :, max, axes(B))
    else
        @assert axes(dst) == axes(A)
        localfilter!(dst, A, Kernel{T}(B),
                     (a)     -> typemin(T),
                     (v,a,b) -> max(v, a + b),
                     (d,i,v) -> d[i] = v)
    end
    return dst
end

localextrema(A::AbstractArray, args...) =
    localextrema!(similar(A), similar(A), A, args...)

@doc @doc(erode) localextrema

function localextrema!(Amin::AbstractArray{T,N},
                       Amax::AbstractArray{T,N},
                       A::AbstractArray{T,N}, B=3) where {T,N}
    localextrema!(Amin, Amax, A, Neighborhood{N}(B))
end

@doc @doc(localextrema) localextrema!

function localextrema!(Amin::AbstractArray{T,N},
                       Amax::AbstractArray{T,N},
                       A::AbstractArray{T,N},
                       B::RectangularBox{N}) where {T,N}
    @assert axes(Amin) == axes(Amax) == axes(A)
    localfilter!((Amin, Amax), A, B,
                 (a)     -> (typemax(T),
                             typemin(T)),
                 (v,a,b) -> (min(v[1], a),
                             max(v[2], a)),
                 (d,i,v) -> (Amin[i], Amax[i]) = v)
end

function localextrema!(Amin::AbstractArray{T,N},
                       Amax::AbstractArray{T,N},
                       A::AbstractArray{T,N},
                       B::Kernel{Bool,N}) where {T,N}
    @assert axes(Amin) == axes(Amax) == axes(A)
    localfilter!((Amin, Amax), A, B,
                 (a)     -> (typemax(T),
                             typemin(T)),
                 (v,a,b) -> (b && a < v[1] ? a : v[1],
                             b && a > v[2] ? a : v[2]),
                 (d,i,v) -> (Amin[i], Amax[i]) = v)
end

function localextrema!(Amin::AbstractArray{T,N},
                       Amax::AbstractArray{T,N},
                       A::AbstractArray{T,N},
                       B::Kernel{K,N}) where {T<:AbstractFloat,
                                              K<:AbstractFloat,N}
    @assert axes(Amin) == axes(Amax) == axes(A)
    localfilter!((Amin, Amax), A, Kernel{T}(B),
                 (a)     -> (typemax(T),
                             typemin(T)),
                 (v,a,b) -> (min(v[1], a - b),
                             max(v[2], a + b)),
                 (d,i,v) -> (Amin[i], Amax[i]) = v)
end

#------------------------------------------------------------------------------
# Higher level operators.

"""
```julia
closing(A, R)
opening(A, R)
```

respectively perform a closing or an opening of array `A` by the structuring
element `R`.  If not specified, `R` is a box of size 3 along all the dimensions
of `A`.  A closing is a dilation followed by an erosion, whereas an opening is
an erosion followed by a dilation.

The in-place versions are:

```julia
closing!(dst, wrk, src, R)
opening!(dst, wrk, src, R)
```

which perform the operation on the source `src` and store the result in
destination `dst` using `wrk` as a workspace array.  These 3 arguments must be
similar arrays, `dst` and `src` may be identical, but `wrk` must not be the
same array as `src` or `dst`.  The destination `dst` is returned.

See [`erode`](@ref) or [`dilate`](@ref) for the meaning of the arguments.

"""
closing(A::AbstractArray, args...) =
    closing!(similar(A), similar(A), A, args...)

function closing!(dst::AbstractArray{T,N},
                  wrk::AbstractArray{T,N},
                  src::AbstractArray{T,N}, B=3) where {T,N}
    closing!(dst, wrk, src, Neighborhood{N}(B))
end

function closing!(dst::AbstractArray{T,N},
                  wrk::AbstractArray{T,N},
                  src::AbstractArray{T,N},
                  B::Neighborhood{N}) where {T,N}
    erode!(dst, dilate!(wrk, src, B), B)
end

opening(A::AbstractArray, args...) =
    opening!(similar(A), similar(A), A, args...)

function opening!(dst::AbstractArray{T,N},
                  wrk::AbstractArray{T,N},
                  src::AbstractArray{T,N}, B=3) where {T,N}
    opening!(dst, wrk, src, Neighborhood{N}(B))
end

function opening!(dst::AbstractArray{T,N},
                  wrk::AbstractArray{T,N},
                  src::AbstractArray{T,N},
                  B::Neighborhood{N}) where {T,N}
    dilate!(dst, erode!(wrk, src, B), B)
end

@doc @doc(closing) closing!
@doc @doc(closing) opening
@doc @doc(closing) opening!

# Out-of-place top hat filter requires 2 allocations without a
# pre-filtering, 3 allocations with a pre-filtering.

"""
```julia
top_hat(A, R)
top_hat(A, R, S)
bottom_hat(A, R)
bottom_hat(A, R, S)
```

Perform A summit/valley detection by applying a top-hat filter to array
`A`.  Argument `R` defines the structuring element for the feature
detection.  Optional argument `S` specifies the structuring element used to
apply a smoothing to `A` prior to the top-hat filter.  If `R` and `S` are
specified as the radii of the structuring elements, then `S` should be
smaller than `R`.  For instance:

```julia
top_hat(bitmap, 3, 1)
```

may be used to detect text or lines in a bimap image.

The in-place versions:

```julia
top_hat!(dst, wrk, src, R)
bottom_hat!(dst, wrk, src, R)
```

apply the top-hat filter on the source `src` and store the result in the
destination `dst` using `wrk` as a workspace array.  These 3 arguments must
be similar but different arrays.  The destination `dst` is returned.

See also [`dilate`](@ref), [`closing`](@ref) and [`morph_enhance`](@ref).

"""
top_hat(a, r=3) = top_hat!(similar(a), similar(a), a, r)

function top_hat(a, r, s)
    wrk = similar(a)
    top_hat!(similar(a), wrk, closing!(similar(a), wrk, a, s), r)
end

function top_hat!(dst::AbstractArray{T,N},
                  wrk::AbstractArray{T,N},
                  src::AbstractArray{T,N}, r=3) where {T,N}
    opening!(dst, wrk, src, r)
    @inbounds for i in eachindex(dst, src)
        dst[i] = src[i] - dst[i]
    end
    return dst
end

bottom_hat(a, r=3) = bottom_hat!(similar(a), similar(a), a, r)

function bottom_hat(a, r, s)
    wrk = similar(a)
    bottom_hat!(similar(a), wrk, opening!(similar(a), wrk, a, s), r)
end

function bottom_hat!(dst::AbstractArray{T,N},
                     wrk::AbstractArray{T,N},
                     src::AbstractArray{T,N}, r=3) where {T,N}
    closing!(dst, wrk, src, r)
    @inbounds for i in eachindex(dst, src)
        dst[i] -= src[i]
    end
    return dst
end

@doc @doc(top_hat)    top_hat!
@doc @doc(top_hat) bottom_hat
@doc @doc(top_hat) bottom_hat!


module MnistTest

# this re-exports Transformations, StochasticOptimization, PenaltyFunctions, and ObjectiveFunctions
using Learn

# my version of ML iteration.  Hopefully will be replaced with what's currently in MLDataUtils dev branch
using StochasticOptimization.Iteration

import MLDataUtils: rescale!

# for loading the data
import MNIST

# for plotting
using StatPlots, MLPlots
gr(leg=false, linealpha=0.5)

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

# create a one-hot matrix given class labels
# TODO: this should be added as a utility in MLDataUtils
function to_one_hot(y::AbstractVector)
    yint = map(yi->round(Int,yi)+1, y)
    nclasses = maximum(yint)
    hot = zeros(Float64, nclasses, length(y))
    for (i,yi) in enumerate(yint)
        hot[yi,i] = 1.0
    end
    hot
end

# randomly pick a subset of testdata (size = totcount) and compute the total loss
function my_test_loss(obj, testdata, totcount = 500)
    totloss = 0.0
    totcorrect = 0
    for (x,y) in each_obs(rand(each_obs(testdata), totcount))
        totloss += transform!(obj,y,x)

        # logistic version:
        # ŷ = output_value(obj.transformation)[1]
        # correct = (ŷ > 0.5 && y > 0.5) || (ŷ <= 0.5 && y < 0.5)

        # softmax version:
        ŷ = output_value(obj.transformation)
        chosen_idx = indmax(ŷ)
        correct = y[chosen_idx] > 0

        totcorrect += correct
    end
    totloss, totcorrect/totcount
end

function doit()

    # our data:
    x_train, y_train = MNIST.traindata()
    x_test, y_test = MNIST.testdata()

    # normalize the input data given μ/σ for the input training data
    μ, σ = rescale!(x_train)
    rescale!(x_test, μ, σ)

    # note: needs my update to MLDataUtils
    y_train, y_test = map(to_one_hot, (y_train, y_test))

    # to_isone(y::AbstractVector) = (z = Array(eltype(y), 1, length(y)); map!(yi->float(yi==1.0), z, y))
    # y_train, y_test = map(to_isone, (y_train, y_test))

    # # keep only 0's and 1's
    # train = filterobs(i -> y_train[i] < 1.5, x_train, y_train)
    # test = filterobs(i -> y_test[i] < 1.5, x_test, y_test)

    train = (x_train, y_train)
    test = (x_test, y_test)

    # At this point we have train and test where each column of x is length-784 corresponding to pixel intensities,
    # and each column of y is length-10 corresponding to output class.

    nin, nh, nout = 784, [100], 10

    # build our objective from a neural net with nh hidden nodes and cross-entropy loss
    # NOTE: cross entropy loss is assumed from the softmax output
    t = nnet(nin, nout, nh, :relu, :softmax)
    # penalty = ElasticNetPenalty(1e-5)
    penalty = L2Penalty(1e-5)
    obj = objective(t, penalty)
    @show obj

    # ---------------------------------------
    # this section is ONLY FOR PLOTTING
    # it can be skipped completely if you only care about learning the model

    # the parts of the plot
    chainplt = ChainPlot(t, maxn=100)
    lossplt = TracePlot(title="Test Loss", ylim=(0,Inf))
    accuracyplt = TracePlot(title="Accuracy", ylim=(0.6,1))
    hmplt = heatmap(rand(28,28), ratio=1)

    # put together the full plot... a ChainPlot with loss, accuracy, and the heatmap
    plot(
        chainplt.plt,
        lossplt.plt,
        accuracyplt.plt,
        hmplt,
        size = (1200,800),
        layout=@layout([a; grid(1,3){0.2h}])
    )

    doanim = false
    # anim = Animation()

    tracer = IterFunction((obj, i) -> begin
        # sample points from the test set and compute/save the loss
        @show i
        if mod1(i,500)==500
            totloss, accuracy = my_test_loss(obj, test, 200)
            @show totloss, accuracy
            push!(lossplt, i, totloss)
            push!(accuracyplt, i, accuracy)
        end

        # add transformation data
        update!(chainplt)

        # update the heatmap of the total outgoing weight from each pixel
        pixel_importance = reshape(sum(t[1].params.views[1],1), 28, 28)
        # pixel_importance = reshape(abs(input_grad(t)),28,28)
        hmplt[1][1][:z].surf[:] = pixel_importance

        # handle animation frames/output
        if doanim
            lastframe = 5000
            if i < lastframe
                frame(anim)
            elseif i == lastframe
                gif(anim, fps=10)
            end
        end

        # display the plot
        gui()
        # if i>0
        #     @profile gui()
        # else
        #     gui()
        # end
    end, every=100)

    # trace once before we start learning to see initial values
    tracer.f(obj, 0)

    # end of plotting section
    # ---------------------------------------

    # create a gradient descent learner
    learner = make_learner(
        # averages the gradient over minibatches, updating params using Adadelta method
        GradientLearner(1e-2, Adam()),

        # our custom iteration method
        tracer,

        # shorthand to add a MaxIter(50000)
        maxiter = 50000
    )

    # do the learning... average over minibatches of size 5
    learn!(obj, learner, infinite_batches(train, size=5))

    # return these in case we want to analyze afterwards
    obj, learner
end

obj, learner = doit()

end #module


export normcols, normrows, rank, qrfactors, isorth
export tolabsdefault, tolreldefault, adjustsize
export QRFactorization, QRWrapper

const QRFactorization = Union{QRPivoted,
                              LinearAlgebra.QRCompactWY,
                              SuiteSparse.SPQR.QRSparse}

struct QRWrapper{F<:Union{QRPivoted,LinearAlgebra.QRCompactWY,SuiteSparse.SPQR.QRSparse}}
    parent::F
end

"""
    calculate 2-norms all columns of a sparse matrix
"""
function normcols(A::SparseMatrixCSC{T}) where T<:Number
    i = 1
    TT = typeof(sqrt(one(eltype(A))))
    z = zero(T)
    sum = z
    r = Vector{TT}(undef, A.n)
    for j = 1:A.n
        k = A.colptr[j+1]
        while i < k
            sum += abs2(A.nzval[i])
            i += 1
        end
        r[j] = sum
        sum = z
        end
        sqrt.(real(r))
end

function normcols(A::AbstractMatrix{T}) where T<:Number
    [norm(view(A, :, i)) for i in 1:size(A, 2)]
end

"""
    calculate 2-norms of all rows of a sparse matrix
"""
function normrows(A::SparseMatrixCSC{T}) where T<:Number
    TT = typeof(sqrt(one(eltype(A))))
    r = zeros(TT, size(A, 1))
    for i = 1:nnz(A)
        j = A.rowval[i]
        r[j] += abs2(A.nzval[i])
    end
    sqrt.(r)
end

function normrows(A::AbstractMatrix{T}) where T<:Number
    [norm(view(A, i, :)) for i in 1:size(A, 1)]
end

"""
    default tolerance absolute for qr-calculations
"""
function tolabsdefault(A::AbstractMatrix{T}) where T<:Number
    20maximum(normcols(A)) * eps(real(T)) * max(size(A)...)
end

"""
    default tolerance relative for rank-calculations
"""
function tolreldefault(A)
    T = eltype(A)
    sqrt(size(A,1)) * eps(real(T)) * 20
end

"""
    absolute tolerance as a combination of tolrel and tolabs
"""
function tolerance(A::AbstractMatrix, pivot::Bool,
                   tolrel::AbstractFloat, tolabs::AbstractFloat)

    tolabs >= 0 || throw(ArgumentError("absolute tolerance must not be negative"))
    tolrel >= 0 || throw(ArgumentError("relative tolerance must not be negative"))
    if tolrel == 0 == tolabs && pivot
        tolrel = tolreldefault(A)
    end
    tolrel == 0 ? tolabs : max(tolrel * maximum(normcols(A)), tolabs)
end

"""
    rank of a QR-factorization of a matrix
"""
function LinearAlgebra.rank(F::Union{QRFactorization,QRWrapper};
                            tolrel::AbstractFloat=0.0, tolabs::AbstractFloat=0.0)

    FR = F.R
    tol = tolerance(FR, true, tolrel, tolabs)
    count(abs2.(diag(FR)) .> tol^2)
end

"""
    `adjustsize(A::Matrix, m:Integer, n::Integer)`

Return a Matrix with `m` rows and `n` columns. If `m > size(A, 1)` new zero rows are added,
if `n > size(A, 2)` zero columns are added. The left upper corner of the resulting
matrix is filled with the corresponding elements of `A`.
"""
function adjustsize(A::AbstractMatrix, m::Integer=size(A,1), n::Integer=size(A,2))

    m0, n0 = size(A)
    if n == n0 && m == m0
        A
    elseif A isa SparseMatrixCSC &&
        ( n0 <= n && ( m > m0 || maximum(view(A.rowval, 1:nnz(A))) <= m < m0 ))
        if n == n0
            acolptr = A.colptr
        else
            acolptr = resize!(copy(A.colptr), n+1)
            fill!(view(acolptr, n0+2:n+1), acolptr[n0+1])
        end
        SparseMatrixCSC(m, n, acolptr, A.rowval, A.nzval)
    elseif m <= m0 && n <= n0
        A[1:m,1:n]
    else
        B = similar(A, m, n)
        fill!(B, 0)
        m0 = min(m, m0)
        n0 = min(n, n0)
        B[1:m0,1:n0] = A[1:m0,1:n0]
        B
    end
end

Base.show(io::IO, qr::QRWrapper) = show(io, qr.parent)
Base.propertynames(qr::QRWrapper, private::Bool) = (:R, :Q, :pcol, :prow)
import Base.getproperty

function getproperty(qr::QRWrapper{<:SuiteSparse.SPQR.QRSparse}, d::Symbol)
    if d == :R || d == :Q || d == :pcol || d == :prow
        getproperty(qr.parent, d)
    elseif d == :p
        qr.parent.pcol
    else
        getfield(qr, d)
    end
end
function getproperty(qr::QRWrapper{<:QRPivoted}, d::Symbol)
    if d == :R || d == :Q || d == :p || d == :P
        getproperty(qr.parent, d)
    elseif d == :pcol
        getproperty(qr.parent, :p)
    elseif d == :prow
        collect(1:size(qr.parent.Q, 1))
    else
        getfield(qr, d)
    end
end
function getproperty(qr::QRWrapper{<:LinearAlgebra.QRCompactWY}, d::Symbol)
    if d == :R || d == :Q
        getproperty(qr.parent, d)
    elseif d == :pcol
        collect(1:size(qr.parent.R, 2))
    elseif d == :prow
        collect(1:size(qr.parent.Q, 1))
    else
        getfield(qr, d)
    end
end

Base.size(qrw::QRWrapper, arg...) = size(qrw.parent, arg...)

"""
    qrfactors(A::AbstractMatrix; tolrel, tolabs, pivot::Bool=true)

Perform QR-factorization of matrix `A` with various matrix types (dense and sparse).
Allow to specify if pivoting is done or not and a relative - and absolute tolerance
bounds on the diagonal elements of R.
Return a common abstract result type `q::QRWrapper` in all cases, which allows to
access the elements of the factorization `q.Q q.R q.pcol q.prow` in all cases.
In the case of dense `A` with `pivot=true` and `tolerance != 0`, the tiny elements
in `R` are set to zero and the corresponding Householder-factors to unity.
"""
function qrfactors(A::SparseMatrixCSC{<:Union{ComplexF64,Float64}};
                 tolrel::AbstractFloat=0.0, tolabs::AbstractFloat=0.0, pivot=true)

    tol = tolerance(A, pivot, tolrel, tolabs)
    QRWrapper(qrfact(A, tol=tol))
end

function qrfactors(A::AbstractSparseMatrix;
                 tolrel::AbstractFloat=0.0, tolabs::AbstractFloat=0.0, pivot=true)

    throw(ArgumentError("qrfact does not support $(typeof(A))"))
end

function qrfactors(A::AbstractMatrix;
                 tolrel::AbstractFloat=0.0, tolabs::AbstractFloat=0.0, pivot=true)

    tol = pivot ? tolerance(A, pivot, tolrel, tolabs) : 0.0
    if tol > 0 && !pivot
        throw(ArgumentError("tol > 0 requires pivot=true"))
    end
    qr = qrfact(A, Val(pivot))
    if tol > 0
        m = size(qr.factors, 1)
        k = rank(qr, tolrel=tolrel, tolabs=tolabs)
        for i = k+1:length(qr.τ)
            qr.τ[i] = 0
            qr.factors[i:m,i] = 0
        end
    end
    QRWrapper(qr)
end

function qrfactors(A::Union{Adjoint, Transpose};
                 tolrel::AbstractFloat=0.0, tolabs::AbstractFloat=0.0, pivot=true)

    qrfactors(copy(A), tolrel=tolrel, tolabs=tolabs)
end

"""
    isorth(A) -> Bool

Check if matrix `A` is orthogonal/unitary.
If `A` is not square, check if all columns/rows are mutually orthogonal.
"""
function isorth(A::AbstractMatrix)
    m, n = size(A)
    tol = eps(eltype(A)) * max(size(A)...) * maximum(abs.(collect(extrema(A))))
    if m >= n
        maximum(abs.(A'A - I)) < tol
    else
        maximum(abs.(A*A' - I)) < tol
    end
end


using Catlab.WiringDiagrams

f = singleton_diagram(Box(:f, [:A], [:A]))
g = singleton_diagram(Box(:g, [:B], [:B]))
h = singleton_diagram(Box(:h, [:A,:B], [:C]))
diagram = compose(otimes(f,g),h)


function tokenize(str)
    tokens = String[]
    token = ""
    for c in str
        if c == ' '
            token != "" && push!(tokens, token)
            token = ""
        elseif c == '+' || c == '*'
            token != "" && push!(tokens, token)
            push!(tokens, string(c))
        elseif c == '(' || c == ')'
            token != "" && push!(tokens, token)
            push!(tokens, string(c))
            token = ""
        else
            token *= c
        end
    end
    token != "" && push!(tokens, token)

    tokens
end

function part1(tokens::AbstractVector{<:AbstractString})
    n = 0
    op = +
    i = 1
    while i ≤ length(tokens)
        if tokens[i] == "("
            j, m = part1((@view tokens[i+1:end]))
            i += j + 1
            n = op(n, m)
        elseif tokens[i] == ")"
            break
        elseif tokens[i] == "+"
            op = +
            i += 1
        elseif tokens[i] == "*"
            op = *
            i += 1
        else
            n = op(n, parse(Int, tokens[i]))
            i += 1
        end
    end
    i, n
end
part1(str::AbstractString) = last(part1(tokenize(str)))

abstract type AbstractExpression end

mutable struct Expression <: AbstractExpression
    op::Function
    left::AbstractExpression
    right::AbstractExpression
end

mutable struct Value <: AbstractExpression
    val::Int
end

function parseexpr(tokens::AbstractVector{<:AbstractString})
    i, op, exp = 1, +, Value(0)
    while i ≤ length(tokens)
        if tokens[i] == "("
            j, right = parseexpr(@view tokens[i+1:end])
            i += j + 1
            if exp isa Value
                exp = Expression(op, exp, right)
            elseif op === +
                exp.right = Expression(op, exp.right, right)
            else
                exp = Expression(op, exp, right)
            end
        elseif tokens[i] == ")"
            break
        elseif tokens[i] == "+"
            op = +
            i += 1
        elseif tokens[i] == "*"
            op = *
            i += 1
        else
            leaf = Value(parse(Int, tokens[i]))
            if exp isa Value
                exp = Expression(op, exp, leaf)
            elseif op === +
                exp.right = Expression(op, exp.right, leaf)
            else
                exp = Expression(op, exp, leaf)
            end
            i += 1
        end
    end
    i, exp
end
parseexpr(str::AbstractString) = last(parseexpr(tokenize(str)))

evaluate(expr::Expression) = expr.op(evaluate(expr.left), evaluate(expr.right))
evaluate(expr::Value) = expr.val

part2(str::AbstractString) = evaluate(parseexpr(str))

let lines = readlines("input.txt")
    @info "Solution" sum(map(part1, lines)) sum(map(part2, lines))
end


"""
    FiniteDimOperator{F,W,B}

A type representing a map from [0,1] to ``Rⁿ``, that is

    z[i] = σ(w[i](x) + b[i])

where `w` is a vector of linear forms on the Hilbert space of ``[0,1] → R``-functions.
"""
struct FiniteDimOperator{F,W,B}
    w::Vector{W} # linear forms
    b::Vector{B} # biases
    σ::F # activation function
end

#=
function FiniteDimOperator(out::Integer, σ=identity; init=Flux.glorot_uniform)
    b = init(out)
    FiniteDimOperator(w,b,σ)
end
=#

function Base.show(io::IO, ::MIME"text/plain", x::FiniteDimOperator{F,W,B}) where {F,W,B}
    println(io, "$(length(x.b))-dimensional FiniteDimOperator with\n functional type $W\n bias type $B\n activation function $F")
end

Flux.@functor FiniteDimOperator

function (fdo::FiniteDimOperator)(x)
    w, b = fdo.w, fdo.b
    fdo.σ.([w[i](x) + b[i] for i in 1:length(b)])
end


# This file inludes fractal integration tools.

export integrate

"""
    $SIGNATURES

Calculates numerical integration of `pts` using the fractal integration method `method` that may be

* [`Interp1D`](@ref)
* [`HInterp1D`](@ref)
* [`Interp2D`](@ref)
* [`HInterp2D`](@ref)

"""
integrate(pts::AbstractVector{<:AbstractVector{<:Real}}, method::AbstractInterp) =
    integrate(map(pnt -> Point(pnt...), pts), method)

function integrate(pts::PointVector, method::AbstractInterp)
    tess = tessellate(pts, method)
    transforms = gettransforms(pts, method)
    val = evaluate(transforms, method, pts)
end

# Interp1D
function evaluate(transforms::AbstractVector{<:Transformation}, method::Interp1D, pts::PointVector{2})
    x  = getindex.(pts, 1)
    (a11, a21, _, a22), (b1, b2) = extract(transforms)

    # J   = abs.(a11)  # Taking abs is not true, |⋅| denotes determinant, not abs
    J   = a11
    K1  = (x[end]^2 - x[1]^2) / 2
    K2  = x[end] - x[1]

    num = sum(J .* (a21 * K1 + b2 * K2))
    denum = 1 - sum(J .* a22)
    num / denum
end

# HInterp1D
function evaluate(transforms::AbstractVector{<:Transformation}, method::HInterp1D, pts::PointVector{3})
    x = getindex.(pts, 1)
    (a11, a21, a31, _, a22, a32, _, a23, a33), (b1, b2, b3) = extract(transforms)

    # J   = abs.(a11) # Taking abs is not true, |⋅| denotes determinant, not abs
    J   = abs.(a11)
    K1  = (x[end]^2 - x[1]^2) / 2
    K2  = (x[end] - x[1])

    W11 = sum(a22 .* J);    W12 = sum(a23 .* J);    Λ1  = sum((a21 * K1 + b2 * K2) .* J);
    W21 = sum(a32 .* J);    W22 = sum(a33 .* J);    Λ2  = sum((a31 * K1 + b3 * K2) .* J);

    A = [1 - W11    -W12;
         -W21       1 - W22]
    b = [Λ1, Λ2]
    I = A \ b
    I[1]    # Hidden integration value I[2] is not returned.
end

# Interp2D
function evaluate(transforms::AbstractVector{<:Transformation}, method::Interp2D, pts::PointVector{3})
    (x1, y1), (x2, y2), (x3, y3) = getboundary(pts, method)

    k11 = x2 - x1;      k12 = x3 - x1;      l1 = x1
    k21 = y2 - y1;      k22 = y3 - y1;      l2 = y1

    # JT = abs(k11 * k22 - k21 * k12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JT = k11 * k22 - k21 * k12
    Δ1 = k11 + k12 + 3 * l1
    Δ2 = k21 + k22 + 3 * l2

    (a11, a21, a31, a12, a22, a32, _, _, a33), (b1, b2, b3) = extract(transforms)
    # JL = abs.(a11 .* a22 - a21 .* a12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JL = a11 .* a22 - a21 .* a12

    num = JT / 6 * sum(JL .* (a31 * Δ1 + a32 * Δ2 + 3 * b3))
    denum = 1 - sum(JL .* a33)
    num / denum
end

# HInterp2D
function evaluate(transforms::AbstractVector{<:Transformation}, method::HInterp2D, pts::PointVector{4})
    (x1, y1), (x2, y2), (x3, y3) = getboundary(pts, method)
    k11 = x2 - x1;      k12 = x3 - x1;      l1 = x1
    k21 = y2 - y1;      k22 = y3 - y1;      l2 = y1

    # JT = abs(k11 * k22 - k21 * k12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JT = k11 * k22 - k21 * k12
    Δ1 = k11 + k12 + 3 * l1
    Δ2 = k21 + k22 + 3 * l2

    (a11, a21, a31, a41, a12, a22, a32, a42, _, _, a33, a43, _, _, a34, a44), (b1, b2, b3, b4) = extract(transforms)
    # JL = abs.(a11 .* a22 - a21 .* a12) # Taking abs is not true, |⋅| denotes determinant, not abs
    JL = a11 .* a22 - a21 .* a12

    W11 = sum(a33 .* JL);    W12 = sum(a34 .* JL);    Λ1  = JT / 6 * sum((a31 * Δ1 + a32 * Δ2 + 3 * b3) .* JL);
    W21 = sum(a43 .* JL);    W22 = sum(a44 .* JL);    Λ2  = JT / 6 * sum((a41 * Δ1 + a42 * Δ2 + 3 * b4) .* JL);

    A = [1 - W11    -W12;
         -W12       1 - W22]
    b = [Λ1, Λ2]
    I = A \ b
    I[1]    # Hidden integration value I[2] is not returned.
end

function extract(transforms::AbstractVector{<:Transformation})
    As = getfield.(transforms, :A)
    bs = getfield.(transforms, :b)
    A = map(idx -> getindex.(As, idx...), [(i, j) for i in 1 : size(As[1], 1), j in 1 : size(As[1], 2)])
    b = map(idx -> getindex.(bs, idx), 1 : length(bs[1]))
    (A, b)
end


module RegERMs

using StatsBase
using Optim

import StatsBase: predict
import Optim: optimize

## Exports

export
    # loss
    Loss,                       # abstract type for all kinds of loss
    LogisticLoss,               # Type to represent logistic loss
    SquaredLoss,                # Type to represent squared loss
    HingeLoss,                  # Type to represent hinge loss
    MultinomialLogisticLoss,    # Type to represent multinomial logistic loss (soft max)

    value,                      # evaluate a single value
    values,                     # evaluate multiple values
    deriv,                      # evaluate a single derivative value
    derivs,                     # evaluate multiple derivative values
    gradient,                   # evaluate a signle gradient
    value_and_deriv,            # jointly evaluate function value and derivative
    tloss,                      # the total loss

    # regularizer
    Regularizer,                # abstract type for all kinds of regularizers
    L2reg,                      # Type to represent L2 regularizer

    # model
    Model,                      # automatic selection of model depending on regression and model type
    RegressionModel,            # abstract type for all regression models
    OrdinalModel,               # regression model for ordinal regression
    BinomialModel,              # regression model for binomial classification
    MultinomialModel,           # regression model for multinomial classification
    predict,                    # evaluate a prediction

    # regression function
    RegressionFunction,         # abstract type for possible model functions
    LinearRegressionFunction,   # linear model for regression

    # models
    SVM,                        # SVM model (binomial regression)
    RidgeReg,                   # Ridge regression (ordinal regression)
    BinomialLogReg,             # Logistic Regression (binomial regression)
    MultinomialLogReg,          # Logistic Regression (multinomial regression)

    # solvers
    LBFGSSolver,                # LBFGS solver
    SGDSolver,                  # SGD solver

    # optim
    RegERM,                     # abstract type for learning models
    RegressionSolver,           # solver for regression problems
    objective,                  # evaluate the objective
    optimize,                   # optimize the objective

    # hyperparameters
    L1RDAParameters,
    RegularizationParameters

## Source files

# include
include("loss.jl")
include("regularizer.jl")
include("mercer_map.jl")   ## TODO: need some discussion about the API for this
include("models/hyperparameters.jl")
include("regression_function.jl")
include("model.jl")
include("optim.jl")

# classification methods
include("models/svm.jl")
include("models/logistic_regression.jl")
# regression methods
include("models/ridge_regression.jl")
# solver
include("solvers/sgd.jl")
include("solvers/l1rda.jl")
include("solvers/lbfgs.jl")

end # module