import Distributions.Categorical

immutable HiddenMarkovModel
    initial_state_prior::Array{Float64,1}
    transition_model::Array{Float64,2}
    observation_model::Array{Float64,2}
    log_initial_state_prior::Array{Float64,1}
    log_transition_model::Array{Float64,2}
    log_observation_model::Array{Float64,2}
    num_states::Int
    num_obs::Int

    function HiddenMarkovModel(initial_state_prior::Array{Float64,1},
                               transition_model::Array{Float64,2},
                               observation_model::Array{Float64,2})
        if !isapprox(sum(initial_state_prior), 1.)
            error("initial_state_prior is not normalized: $(initial_state_prior)")
        end
        num_states = size(initial_state_prior)[1]
        if size(transition_model) != (num_states, num_states)
            error("transition_model has wrong dimension: got $(size(transition_model)) expected $((num_states, num_states))")
        end
        if !isapprox(sum(transition_model, 2), ones(num_states, 1))
            error("transition_model rows are not all normalized: got $(sum(transition_model, 2))")
        end
        if size(observation_model)[1] != num_states
            error("observation_model has wrong number of rows: got $(size(observation_model)[1]) expected $num_states")
        end
        num_obs = size(observation_model)[2]
        if !isapprox(sum(observation_model, 2), ones(num_states, 1))
            error("observation_model rows are not all normalized: got $(sum(observation_model, 2))")
        end
        log_initial_state_prior = log.(initial_state_prior)
        log_transition_model = log.(transition_model)
        log_observation_model = log.(observation_model)
        new(initial_state_prior, transition_model, observation_model,
            log_initial_state_prior, log_transition_model, log_observation_model,
            num_states, num_obs)
    end
end

num_states(hmm::HiddenMarkovModel) = hmm.num_states
num_observations(hmm::HiddenMarkovModel) = hmm.num_obs

function hmm_simulate(hmm::HiddenMarkovModel, num_steps::Int)
    states = Array{Int, 1}(num_steps)
    observations = Array{Int, 1}(num_steps)
    states[1] = rand(Categorical(hmm.initial_state_prior))
    for i=2:num_steps
        prev = states[i-1]
        states[i] = rand(Categorical(hmm.transition_model[prev,:]))
    end
    for i=1:num_steps
        observations[i] = rand(Categorical(hmm.observation_model[states[i],:]))
    end
    (states, observations)
end

function hmm_forward_pass(hmm::HiddenMarkovModel, observations::Array{Int, 1})
    # compute the forward probabilities
    # p(x_t, y_{1:t-1}) for t=1,...,T
    # NOTE: we are assuming an observation for each time point
    num_steps = length(observations)
    num_states = hmm.num_states
    fprobs = Array{Float64,2}(num_steps, num_states)
    fprobs[1,:] = hmm.initial_state_prior
    for t=2:num_steps
        obs = observations[t-1]
        message = fprobs[t-1,:] .* hmm.observation_model[:,obs]
        # rows are x_{t-1}, cols are x_t
        product = message .* hmm.transition_model
        fprobs[t,:] = vec(sum(product, 1)) # sum over rows
    end
    return fprobs
end

function hmm_log_forward_pass(hmm::HiddenMarkovModel,
                              observations::Array{Int, 1})
    # compute the log forward probabilities
    # log p(x_t, y_{1:t-1}) for t=1,...,T
    # NOTE: we are assuming an observation for each time point
    num_steps = length(observations)
    num_states = hmm.num_states
    lfprobs = Array{Float64,2}(num_steps, num_states)
    lfprobs[1,:] = hmm.log_initial_state_prior
    for t=2:num_steps
        obs = observations[t-1]
        lmessage = lfprobs[t-1,:] .+ hmm.log_observation_model[:,obs]
        # rows are x_{t-1}, cols are x_t; sum over rows
        lproduct = lmessage .+ hmm.log_transition_model
        lfprobs[t,:] = logsumexp(lproduct, 1)
    end
    return lfprobs
end

function hmm_marginal_likelihood(hmm::HiddenMarkovModel,
                                 observations::Array{Int, 1})
    # use the forward probability p(x_T, y_{1:T-1})
    fprobs = hmm_forward_pass(hmm, observations)
    # p(y_{1:T}) = sum_{x_T} p(x_T, y_{1:T-1}) p(y_T | x_T)
    sum(fprobs[end,:] .* hmm.observation_model[:,observations[end]])
end

function hmm_log_marginal_likelihood(hmm::HiddenMarkovModel,
                                 observations::Array{Int, 1})
    lfprobs = hmm_log_forward_pass(hmm, observations)
    logsumexp(lfprobs[end,:] .+ hmm.log_observation_model[:,observations[end]])
end

function hmm_posterior_sample(hmm::HiddenMarkovModel,
                              observations::Array{Int, 1})
    num_steps = length(observations)
    lfprobs = hmm_log_forward_pass(hmm, observations)
    ldist = lfprobs[end,:] .+ hmm.log_observation_model[:,observations[end]]
    particle = Array{Int,1}(num_steps)
    particle[end] = rand(Categorical(exp.(ldist - logsumexp(ldist))))
    for t = num_steps-1:-1:1
        ldist = lfprobs[t,:]
        ldist = ldist .+ hmm.log_observation_model[:,observations[t]]
        ldist = ldist .+ hmm.log_transition_model[:,particle[t+1]]
        dist = exp.(ldist - logsumexp(ldist))
        particle[t] = rand(Categorical(dist))
    end
    particle
end

function hmm_log_joint_probability(hmm::HiddenMarkovModel,
                                   x::Array{Int,1},
                                   observations::Array{Int,1})
    if length(x) != length(observations)
        error("length(x) != length(observations)")
    end
    lj = hmm.log_initial_state_prior[x[1]]
    for t=2:length(x)
        lj += hmm.log_transition_model[x[t-1], x[t]]
    end
    for t=1:length(x)
        lj += hmm.log_observation_model[x[t], observations[t]]
    end
    lj
end

# -- Helper functions for SMC in HMM using prior proposals

immutable HMMPriorInitializer
    hmm::HiddenMarkovModel
    observation::Int # should be the first observation
end

function forward(init::HMMPriorInitializer)
    particle = Array{Int,1}([rand(Categorical(init.hmm.initial_state_prior))])
    log_weight = init.hmm.log_observation_model[particle[1], init.observation]
    (particle, log_weight)
end

function backward(init::HMMPriorInitializer, particle::Array{Int,1})
    @assert length(particle) == 1
    init.hmm.log_observation_model[particle[1], init.observation]
end

immutable HMMPriorIncrementer
    hmm::HiddenMarkovModel
    observation::Int
end

function forward(incr::HMMPriorIncrementer, particle::Array{Int,1})
    prev = particle[end]
    cur = rand(Categorical(incr.hmm.transition_model[prev,:]))
    log_weight = incr.hmm.log_observation_model[cur, incr.observation]
    new_particle = vcat(particle, [cur])
    (new_particle, log_weight)
end

function backward(incr::HMMPriorIncrementer, new_particle::Array{Int,1})
    cur = new_particle[end]
    log_weight = incr.hmm.log_observation_model[cur, incr.observation]
    particle = new_particle[1:end-1]
    (particle, log_weight)
end

function HMMPriorSMCScheme(hmm::HiddenMarkovModel,
                           observations::Array{Int,1},
                           num_particles::Int)
    initializer = HMMPriorInitializer(hmm, observations[1])
    incrementers = Array{Any,1}(length(observations) - 1)
    for i = 2:length(observations)
        incrementers[i-1] = HMMPriorIncrementer(hmm, observations[i])
    end
    SMCScheme(initializer, incrementers, num_particles)
end

# -- Helper functions for SMC in HMM using conditonal (optimal) proposals

immutable HMMConditionalInitializer
    hmm::HiddenMarkovModel
    observation::Int
    dist::Categorical
    log_weight::Float64
    function HMMConditionalInitializer(hmm::HiddenMarkovModel, observation::Int)
        lprior = hmm.log_initial_state_prior
        llikelihood = hmm.log_observation_model[:,observation]
        ldist = lprior .+ llikelihood
        log_weight = logsumexp(ldist)
        # log normalized distribution
        ldist = ldist - log_weight
        dist = Categorical(exp.(ldist))
        new(hmm, observation, dist, log_weight)
    end
end

function forward(init::HMMConditionalInitializer)
    # p(x_1 | y_1)
    particle = Array{Int,1}([rand(init.dist)])
    (particle, init.log_weight)
end

function backward(init::HMMConditionalInitializer, particle::Array{Int,1})
    @assert length(particle) == 1
    # p(y_1) = sum_{x_1} p(x_1) p(y_1 | x_1)
    init.log_weight
end

immutable HMMConditionalIncrementer
    hmm::HiddenMarkovModel
    observation::Int
    llikelihood::Array{Float64,1}
    function HMMConditionalIncrementer(hmm::HiddenMarkovModel, observation::Int)
        llikelihood = hmm.log_observation_model[:,observation]

        new(hmm, observation, llikelihood)
    end
end

function forward(incr::HMMConditionalIncrementer, particle::Array{Int,1})
    cur = particle[end]
    lprior = incr.hmm.log_transition_model[cur,:]
    ldist = lprior .+ incr.llikelihood
    # p(x_t | x_{t-1}, y_t)
    log_weight = logsumexp(ldist)
    new_component = rand(Categorical(exp.(ldist - log_weight)))
    new_particle = vcat(particle, [new_component])
    (new_particle, log_weight)
end

function backward(incr::HMMConditionalIncrementer, new_particle::Array{Int,1})
    cur = new_particle[end]
    lprior = incr.hmm.log_transition_model[cur,:]
    ldist = lprior .+ incr.llikelihood
    # p(y_t | x_{t-1}) = sum_{x_t} p(x_t | x_{t-1}) p(y_t | x_t)
    log_weight = logsumexp(ldist)
    particle = new_particle[1:end-1]
    (particle, log_weight)
end

function HMMConditionalSMCScheme(hmm::HiddenMarkovModel,
                                 observations::Array{Int,1},
                                 num_particles::Int)
    initializer = HMMConditionalInitializer(hmm, observations[1])
    incrementers = Array{Any,1}(length(observations) - 1)
    for i = 2:length(observations)
        incrementers[i-1] = HMMConditionalIncrementer(hmm, observations[i])
    end
    SMCScheme(initializer, incrementers, num_particles)
end


include("rw_1d.jl")

nsteps = length(ARGS)>0 ? parse(ARGS[1]) : 1000
prob = length(ARGS)>1 ? parse(ARGS[1]) : 0.5

result = randomwalk(nsteps,prob)

for (t,x) in enumerate(result)
    println("$(t-1) $x")
end


using Test
using QuantumOptics.sparsematrix
using SparseArrays, LinearAlgebra

# SparseMatrix = quantumoptics.sparsematrix.SparseMatrix
const SparseMatrix = SparseMatrixCSC{ComplexF64, Int}

@testset "sparsematrix" begin

# Set up test matrices
A = rand(ComplexF64, 5, 5)
A_sp = sparse(A)

B = Matrix{ComplexF64}(I, 5, 5)
B_sp = sparse(ComplexF64(1)*I, 5, 5)

C = rand(ComplexF64, 3, 3)
C[2,:] .= 0
C_sp = sparse(C)

R_sp = A_sp + B_sp
R = A + B

# Test arithmetic
@test 0 ≈ norm(Matrix(R_sp) - R)
@test 0 ≈ norm(Matrix(ComplexF64(0.5,0)*A_sp) - 0.5*A)
@test 0 ≈ norm(Matrix(A_sp/2) - A/2)
@test 0 ≈ norm(Matrix(A_sp*B_sp) - A*B)

# Test kronecker product
@test 0 ≈ norm(Matrix(kron(A_sp, C_sp)) - kron(A, C))
@test 0 ≈ norm(Matrix(kron(A_sp, B_sp)) - kron(A, B))

end # testset


using Spinach
using Test

@testset "Spinach.jl" begin
    ctx = LLVM.Context()
    sqrt_fun = "(define (square x) (* x x))
(define (average x y) (/ (+ x y) 2))
(define (abs x) (if (< x 0) (* -1 x) x))
(define (improve guess x) (average guess (/ x guess)))
(define (good-enough? guess x)
  (< (abs (- (square guess) x)) 0.001))

(define (sqrt-iter guess x) (if (good-enough? guess x)
      guess
      (sqrt-iter (improve guess x) x)))
(define (sqrt x) (sqrt-iter 1.00 x))
(sqrt 9)
"
    mod = generate_IR(sqrt_fun, ctx)
    @test run_module(mod,"0") == sqrt(9)

    abs_fun = "(define (abs x) (if (< x 0) (* -1 x) x))
    (abs -2)"
    mod = generate_IR(abs_fun, ctx)
    @test run_module(mod,"0") == abs(-2.)
    fact_fun = "(define (factorial x tot)
    (if (== x 1) tot (factorial (- x 1) (* x tot))))
    (factorial 5 1)"
    mod = generate_IR(fact_fun, ctx)
    @test run_module(mod,"0") == factorial(5)
# Write your tests here.
end


using StatsBase: countmap

export OneClassClassifier, regions, centroid, fit, predict

mutable struct OneClassClassifier{T}
    centers::Vector{T}
    freqs::Vector{Int}
    n::Int
    epsilon::Float64
end

function fit(::Type{OneClassClassifier}, dist::Function, X::AbstractVector{T}, m::Int; centroids=true) where T
    Q = fftclustering(dist, X, m)
    C = X[Q.irefs]
    P = Dict(Q.irefs[i] => i for i in eachindex(Q.irefs))
    freqs = zeros(Int, length(Q.irefs))
    for nn in Q.NN
        freqs[P[first(nn).objID]] += 1
    end

    if centroids
        CC = centroid_correction(dist, X, C)
        OneClassClassifier(CC, freqs, length(X), Q.dmax)
    else
        OneClassClassifier(C, freqs, length(X), Q.dmax)
    end

end

function regions(dist::Function, X, refs::Index)
    I = invindex(dist, X, refs, k=1)
    (freqs=[length(lst) for lst in I], regions=I)
end

function regions(dist::Function, X, refs)
    regions(dist, X, fit(Sequential, refs))
end

function centroid(D)
    sum(D) ./ length(D)
end

function centroid_correction(dist::Function, X, C)
    [centroid(X[lst]) for lst in regions(dist, X, C).regions if length(lst) > 0]
end

function predict(occ::OneClassClassifier{T}, dist::Function, q::T) where T
    seq = fit(Sequential, occ.centers)
    res = search(seq, dist, q, KnnResult(1))
    #1.0 - first(res).dist  / occ.epsilon
    (similarity=max(0.0, 1.0 - first(res).dist  / occ.epsilon), freq=occ.freqs[first(res).objID])
    #occ.freqs[first(res).objID] / occ.n
end


using KNITRO, JuMP, Base.Test

m = Model(solver=KnitroSolver())
@variable(m, x)
setValue(x, 5)
@NLobjective(m, Min, -log(x)+x^2)
solve(m)
@test_approx_eq_eps getobjectivevalue(m) 0.84657 1e-5
@test_approx_eq_eps getvalue(x) 0.707107 1e-5

ktrmod = internalmodel(m)
MathProgBase.freemodel!(ktrmod)


using Distributed, StatsBase, ClusterManagers

addprocs(SlurmManager(parse(Int64, ARGS[1])))
@everywhere begin

    using TransitionalMCMC, Distributions

    # Prior Bounds
    lb, ub  = -5, 5

    # Prior Density and sampler
    logprior(x) = logpdf(Uniform(lb, ub), x[1]) + logpdf(Uniform(lb, ub), x[2])
    priorRnd(Nsamples) = rand(Uniform(lb, ub), Nsamples, 2)

    # Log Likelihood
    function logLik(x)
        sleep(1)
        return -1 * ((x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2)
    end

end

Nsamples = 2000

@time samps, acc = tmcmc(logLik, logprior, priorRnd, Nsamples, 5, 2)

rmprocs(workers())


"""
Squared distance of two vertices.
"""
@i @inline function sqdistance(dist!, x1::AbstractVector{T}, x2::AbstractVector) where T
    @invcheckoff @inbounds for i=1:length(x1)
        @routine begin
            d ← zero(T)
            d += x1[i] - x2[i]
        end
        dist! += d^2
        ~@routine
    end
end

"""
    var_and_mean_sq(var!, mean!, sqv)

the variance and mean value from squared values.
"""
@i function var_and_mean_sq(var!, var_sum, mean!, mean_sum, sqv::AbstractVector{T}) where T
    @inbounds for i=1:length(sqv)
        mean_sum += sqv[i] ^ 0.5
        var_sum += sqv[i]
    end
    mean! += mean_sum / (@const length(sqv))
    @routine begin
        @zeros T var_anc1 var_anc2 sqmean
        var_anc1 += var_sum / (@const length(sqv))
        sqmean += mean! ^ 2
        var_anc1 -= sqmean
        var_anc2 += var_anc1 * (@const length(sqv))
    end
    var! += var_anc2 / (@const length(sqv)-1)
    ~@routine
end

"""The loss of graph embedding problem."""
@i function embedding_loss(out!::T, x) where T
    @routine @invcheckoff begin
        d1 ← zeros(T, length(L1))
        d2 ← zeros(T, length(L2))
        for i=1:length(L1)
            @inbounds sqdistance(d1[i], x[:,L1[i][1]],x[:,L1[i][2]])
        end
        for i=1:length(L2)
            @inbounds sqdistance(d2[i], x[:,L2[i][1]],x[:,L2[i][2]])
        end
        @zeros T v1 m1 v2 m2 vacc1 vacc2 acc1 acc2
        var_and_mean_sq(v1, vacc1, m1, acc1, d1)
        var_and_mean_sq(v2, vacc2, m2, acc2, d2)
        m1 -= m2 - 0.1
    end
    out! += v1 + v2
    if m1 > 0
        # to ensure mean(v2) > mean(v1)
        # if mean(v1)+0.1 - mean(v2) > 0, punish it.
        out! += exp(m1)
        out! -= 1
    end
    ~@routine
end


import Base: Float64, Float32, Float16, Int

abstract type AbstractLogFixPoint <: AbstractFloat end

""" LogFixPoint16 is a logarithmic fixed-point number format with

1 sign bit
7 signed integer bits, with its signbit (=bit 15) flipped (i.e. excess-8)
8 fraction bits.

Exceptions
0x0000      = zero
0x8000      = Not-a-Real (NaR)."""
primitive type LogFixPoint16 <: AbstractLogFixPoint 16 end

Base.iszero(x::LogFixPoint16) = reinterpret(UInt16,x) == 0x0000
Base.isnan(x::LogFixPoint16) = reinterpret(UInt16,x) == 0x8000
Base.signbit(x::LogFixPoint16) = (reinterpret(UInt16,x) & 0x8000) == 0x8000

Base.zero(::Type{LogFixPoint16}) = reinterpret(LogFixPoint16,0x0000)
nan(::Type{LogFixPoint16}) = reinterpret(LogFixPoint16,0x8000)

Base.floatmin(::Type{LogFixPoint16}) = reinterpret(LogFixPoint16,0x0001)
Base.floatmax(::Type{LogFixPoint16}) = reinterpret(LogFixPoint16,0x7fff)

# In the absence of -Inf,Inf define typemin,typemax as floatmin,floatmax
Base.typemin(::Type{LogFixPoint16}) = floatmin(LogFixPoint16)
Base.typemax(::Type{LogFixPoint16}) = floatmax(LogFixPoint16)

Base.one(::Type{LogFixPoint16}) = reinterpret(LogFixPoint16,0x4000)

function Base.:(-)(x::LogFixPoint16)
    iszero(x) && return zero(LogFixPoint16)
    isnan(x) && return nan(LogFixPoint16)
    return reinterpret(LogFixPoint16,reinterpret(UInt16,x) ⊻ 0x8000)
end

function Base.inv(x::LogFixPoint16)
    iszero(x) && return nan(LogFixPoint16)
    ui = reinterpret(UInt16,x)
    sign = ui & 0x8000
    val = -ui & 0x7fff		# two's complement of all except sign bit
    return reinterpret(LogFixPoint16,sign | val)
end

# Default number of fraction bits
const nfrac = Ref{Int}(9)
const nint = Ref{Int}(15-nfrac[])
const scale = Ref{Int}(2^nfrac[])
const scale_over_logof2 = Ref{Float32}(scale[]/log(2f0))
const max_nfrac_supported = 11
const rounding_mode = Ref{Symbol}(:lin)

""" 32-bit value of floatmax/min are derived from

    Int32(round(scale[]*(-(2^i-2^-j)-1))) + bit15flip

where i=nint-1, j=nfrac."""
const logfixpoint16_max = Int32(2^15-1)             # val of floatmax
const logfixpoint16_halfmin = Int32(-2^nfrac[]+1)   # val of half floatmin
const bit15flip = Int32(2^14)    # flip the meaning of the sign bit of integers
const bit14flip = UInt16(2^13)
const uintbit15flip = UInt16(2^14)    # used for power2

"""Constant to convert from round-to-nearest in log to linear space.

Let
    1 = s*log2(f1)
    1.5 = s*log2(f2)
    2 = s*log2(f3)

With some scale s. Now f2 is not the arithmetic mean of f1,f3. Let f2* be that mean

    f2* = 2^(1/s-1) + 2^(2/s-1)

Then we find a constant c_b to be added before rounding by

    1.5 = c_b + 2*log2(f2*)
    => c_b = 1.5 - s*log2(f2*)
           = 1/2 - s*(log2(2^(1/s) + 1) - 1)"""
function rounding_correction(s::Int)
    return 1/2 - s*(log2(2^(1/s)+1)-1)
end

const c_b = Ref{Float32}(Float32(rounding_correction(scale[])))

"""Convert a Float32 to LogFixPoint16 via the base-2 logarithm and rounding.
Round to nearest is applied in lin-space due to the addition of c_b."""
function LogFixPoint16(f::Float32)
    iszero(f) && return zero(LogFixPoint16)
    ~isfinite(f) && return nan(LogFixPoint16)

    # scale_over_logof2 = scale/log(2) to scale by nfrac fraction bits
    # use log instead of log2 as it's faster and include 1/log(2) in the constant
    val = Int32(round(c_b[] + scale_over_logof2[]*log(abs(f)))) + bit15flip
    sign = ((reinterpret(UInt32,f) & 0x8000_0000) >> 16) % UInt16

    sign,val = under_overflow_check(sign,val)
    result = sign | (val % UInt16)
    return reinterpret(LogFixPoint16,result)
end

"""Check for under or overflow:

    - Any result beyond floatmax returns floatmax. (No-overflow rounding mode)
    - Any result smaller than floatmin/2 returns 0, any result between
        floatmin and floatmin/2 returns floatmin."""
function under_overflow_check(sign::UInt16,val::Int32)

    underflow = val < logfixpoint16_halfmin
    overflow = val > logfixpoint16_max

    val, sign = underflow ? (zero(Int32), 0x0000) : (max(val,one(Int32)), sign)
    val = overflow ? logfixpoint16_max : val

    return sign,val
end

"""Pre-calcualte look-up table for conversion from LogFixPoint16 to Float32."""
function createF32LookupTable(nint::Int,nfrac::Int)

    N = 2^(nint+nfrac)      # length of the table
    table = Array{UInt32,1}(undef,N)
    s = 2^nfrac             # scale
    bias = 2^(nint+nfrac-1) # exponent bias (=bit15flip)

    for i in 1:N
        # convert index to signed integer, with signbit flipped
        si = i-bias
        # store float32s as UInt32 for bitwise operations
        table[i] = reinterpret(UInt32,Float32(2.0 ^ (si/s)))
    end

    return table
end

const f32lookup = createF32LookupTable(nint[],nfrac[])

"""Conversion function from LogFixPoint16 to Float32 via table lookup
after branching off the special cases 0,NaR."""
function Float32(x::LogFixPoint16)
    iszero(x) && return 0f0
    isnan(x) && return NaN32

    ui = reinterpret(UInt16,x)
    val = ui & 0x7fff     # set the signbit to zero

    # set all except sign bit to zero
    # sign = either 0x8000_0000 or 0x0000_0000
    sign = ((ui & 0x8000) % UInt32) << 16

    # combine sign and table lookup
    @inbounds f = sign | f32lookup[val]
    return reinterpret(Float32,f)
end

# conversion between LogFixPoint16 and various floats
Float64(x::LogFixPoint16) = Float64(Float32(x))
Float16(x::LogFixPoint16) = Float16(Float32(x))
LogFixPoint16(x::Float64) = LogFixPoint16(Float32(x))
LogFixPoint16(x::Float16) = LogFixPoint16(Float32(x))
Int(x::LogFixPoint16) = Int(Float32(x))
LogFixPoint16(x::Int) = LogFixPoint16(Float32(x))

"""Multiplication for LogFixPoint16, equivalent to an addition of the exponents."""
function Base.:(*)(x::LogFixPoint16,y::LogFixPoint16)
    # special cases NaR*y = x*NaR = NaR
    isnan(x) | isnan(y) && return nan(LogFixPoint16)
    # if no NaN present: 0*y = x*0 = 0
    iszero(x) | iszero(y) && return zero(LogFixPoint16)

    uix = reinterpret(UInt16,x)
    uiy = reinterpret(UInt16,y)

    # mask exponent
    xsign = uix & 0x8000
    ysign = uiy & 0x8000

    # mask the sign and cast to Int32
    xval = (uix & 0x7fff) % Int32
    yval = (uiy & 0x7fff) % Int32

    # resulting sign
    sign = xsign ⊻ ysign

    # ADD EXPONENTS
    # xval and yval both contain the bias bit15flip
    # so subtract bit15flip, such that the result is again
    # biased with (a single) + bit15flip
    val = xval + yval - bit15flip

    # check for over or underflow
    sign,val = under_overflow_check(sign,val)

    # merge sign and exponent back together
    result = sign | (val % UInt16)
    return reinterpret(LogFixPoint16,result)
end

"""Division for LogFixPoint16, equivalent to a subtraction of the exponents."""
function Base.:(/)(x::LogFixPoint16,y::LogFixPoint16)
    # special case, if either x,y =NaN, or y=0 return NaR
    isnan(x) | isnan(y) | iszero(y) && return nan(LogFixPoint16)
    # else if x=0 return 0
    iszero(x) && return zero(LogFixPoint16)

    uix = reinterpret(UInt16,x)
    uiy = reinterpret(UInt16,y)

    # extract sign
    xsign = uix & 0x8000
    ysign = uiy & 0x8000

    # extract exponent
    xval = (uix & 0x7fff) % Int32
    yval = (uiy & 0x7fff) % Int32

    # combine to obtain the sign of the result
    sign = xsign ⊻ ysign

    # SUBTRACT EXPONENTS
    # xval and yval both contain the bias bit15flip
    # so add bit15flip, such that the result is again
    # biased with (a single) + bit15flip
    val = xval - yval + bit15flip

    sign,val = under_overflow_check(sign,val)

    # merge sign and exponent val
    result = sign | (val % UInt16)
    return reinterpret(LogFixPoint16,result)
end

function Base.sqrt(x::LogFixPoint16)
    iszero(x) && return zero(LogFixPoint16)
    signbit(x) && return nan(LogFixPoint16) #TODO throw DomainError?

    uix = reinterpret(UInt16,x)
    #TODO is this correct rounding?
    uix = (uix >> 1) + bit14flip
    return reinterpret(LogFixPoint16,uix)
end

function power2(x::LogFixPoint16)
    iszero(x) && return zero(LogFixPoint16)
    isnan(x) && return nan(LogFixPoint16)
    uix = reinterpret(UInt16,x)
    uix = (uix << 1) + uintbit15flip
    return reinterpret(LogFixPoint16,uix & 0x7fff)
end

""" Precomputes the Gaussian Logarithm as a table lookup for addition. Let

X = (-1)^signx * 2^x
Y = (-1)^signy * 2^y
Z = (-1)^signz * 2^z

then additions can be computed via the Gaussian logarithms

z = log2(||X|+|Y||) = x + log2(1+2^(y-x))

For fixed-point number exponents, interpreted as integers
including a bias c, x̂ = ax + c, where a = 2^n_frac the number of fraction bits:

ẑ = x̂ + a*log2(1+2^((ŷ-x̂)/a)     or
ẑ = ŷ - (ŷ-x̂) + a*log2(1+2^((ŷ-x̂)/a)

The last two terms are precomputed into a table lookup."""
function createAddLookup(scale::Int)

    tab = Array{Int32,1}(undef,max_table_size)

    for i in 0:max_table_size-1
        tab[i+1] = -i + Int(round(c_b[] + scale*log2(1+2^(i/scale))))
    end

    return tab
end

""" Precomputes the Gaussian Logarithm as a table lookup for subtraction.
Subtractions can be computed via the Gaussian logarithms

z = log2(||X|-|Y||) = x + log2(abs(1-2^(y-x)))

For fixed-point number exponents, interpreted as integers
including a bias c, x̂ = ax + c, where a = 2^n_frac the number of fraction bits:

ẑ = x̂ + a*log2(abs(1-2^((ŷ-x̂)/a))     or
ẑ = ŷ - (ŷ-x̂) + a*log2(abs(1-2^((ŷ-x̂)/a))

The last two terms are precomputed into a table lookup."""
function createSubLookup(scale::Int)
    tab = Array{Int32,1}(undef,max_table_size)

    # set the first entry manually to avoid -Inf due to the log2(0)
    tab[1] = -logfixpoint16_max-1

    for i in 1:max_table_size-1
        tab[i+1] = -i + Int(round(c_b[] + scale*log2(abs(1-2^(i/scale)))))
    end

    return tab
end

# table indices higher than that are 0
function find_max_diff_res(scale::Int)
    i = 1
    while (-i + round(c_b[] + scale*log2(1+2^(i/scale)))) > 0.0
        i += 1
    end
    return Int32(i)
end

const max_diff_resolvable = Ref{Int32}(find_max_diff_res(scale[]))
const max_table_size = find_max_diff_res(2^max_nfrac_supported)+1
const addTable = createAddLookup(scale[])
const subTable = createSubLookup(scale[])

function Base.:(+)(x::LogFixPoint16,y::LogFixPoint16)
    isnan(x) | isnan(y) && return nan(LogFixPoint16)

    # zero is an exceptional case and not reprsentable as (-1)^s * 2^x
    iszero(x) && return y
    iszero(y) && return x

    uix = reinterpret(UInt16,x)
    uiy = reinterpret(UInt16,y)

    xsign = uix & 0x8000
    ysign = uiy & 0x8000

    xval = (uix & 0x7fff) % Int32
    yval = (uiy & 0x7fff) % Int32

    # y is always the larger value
    # resulting sign is always that of y
    # in a-a sign is 1, but case is caught in underflow: sign set to 0
    xval,yval,xsign,ysign = xval > yval ? (yval,xval,ysign,xsign) : (xval,yval,xsign,ysign)
    diff = yval - xval     # diff >= 0

    # pull Gaussian logarithms from addTable, subTable
    # update yval with the increment
    @inbounds increment = diff < max_diff_resolvable[] ?
        (xsign == ysign ? addTable[diff+1] : subTable[diff+1]) : zero(Int32)
    yval += increment

    # check for over or underflow
    ysign,yval = under_overflow_check(ysign,yval)

    result = ysign | (yval % UInt16)
    return reinterpret(LogFixPoint16,result)
end

function diff_val(x::LogFixPoint16,y::LogFixPoint16)
    isnan(x) | isnan(y) && return nan(LogFixPoint16)

    uix = reinterpret(UInt16,x)
    uiy = reinterpret(UInt16,y)

    xsign = uix & 0x8000
    ysign = uiy & 0x8000

    xval = (uix & 0x7fff) % Int32
    yval = (uiy & 0x7fff) % Int32

    # y is always the larger value
    # resulting sign is always that of y
    # in a-a sign is 1, but case is caught in underflow: sign set to 0
    xval,yval,xsign,ysign = xval > yval ? (yval,xval,ysign,xsign) : (xval,yval,xsign,ysign)
    diff = yval - xval     # diff >= 0
end

"""Subtraction for LogFixPoint16 via sign switch of y as in x-y = x + (-y)."""
Base.:(-)(x::LogFixPoint16,y::LogFixPoint16) = x + (-y)

"""Rounding to integer via conversion to Float32."""
Base.round(x::LogFixPoint16, r::RoundingMode{:Up}) = Int(ceil(Float32(x)))
Base.round(x::LogFixPoint16, r::RoundingMode{:Down}) = Int(floor(Float32(x)))
Base.round(x::LogFixPoint16, r::RoundingMode{:Nearest}) = Int(round(Float32(x)))

for t in (Int8, Int16, Int32, Int64, Int128, UInt8, UInt16, UInt32, UInt64, UInt128)
	@eval Base.promote_rule(::Type{LogFixPoint16}, ::Type{$t}) = LogFixPoint16
end

function Base.nextfloat(x::LogFixPoint16)
    isnan(x) && return nan(LogFixPoint16)
    (x == -floatmin(LogFixPoint16)) && return zero(LogFixPoint16)

    sign = signbit(x)
    ui = reinterpret(UInt16,x)
    sign && return reinterpret(LogFixPoint16,ui - 0x0001)
    return reinterpret(LogFixPoint16,ui + 0x0001)
end

function Base.nextfloat(x::LogFixPoint16,n::Int)
    for i in 1:n
        x = nextfloat(x)
    end
    return x
end

function Base.prevfloat(x::LogFixPoint16)
	# -floatmax is 0xffff and would otherwise prevfloat to 0=0x0000
    isnan(x) | (x == -floatmax(LogFixPoint16)) && return nan(LogFixPoint16)

    sign = signbit(x)
    ui = reinterpret(UInt16,x)
    sign && return reinterpret(LogFixPoint16,ui + 0x0001)
    return reinterpret(LogFixPoint16,ui - 0x0001)
end

function Base.prevfloat(x::LogFixPoint16,n::Int)
    for i in 1:n
        x = prevfloat(x)
    end
    return x
end

function Base.log2(x::LogFixPoint16)
	signbit(x) | iszero(x) && return nan(LogFixPoint16)
	ui =  reinterpret(UInt16,x) % Int
	ui -= Int(bit15flip)
	ui /= scale
	return LogFixPoint16(ui)
end

function Base.:(==)(x::LogFixPoint16,y::LogFixPoint16)
    isnan(x) | isnan(y) && return false
    return reinterpret(UInt16,x) == reinterpret(UInt16,y)
end

function Base.:(>)(x::LogFixPoint16,y::LogFixPoint16)
    isnan(x) | isnan(y) && return nan(LogFixPoint16)

    uix = reinterpret(UInt16,x)
    uiy = reinterpret(UInt16,y)

    xsign = uix & 0x8000
    ysign = uiy & 0x8000

    xsign > ysign && return false
    xsign < ysign && return true

    return uix > uiy
end

function Base.:(<)(x::LogFixPoint16,y::LogFixPoint16)
    isnan(x) | isnan(y) && return nan(LogFixPoint16)

    uix = reinterpret(UInt16,x)
    uiy = reinterpret(UInt16,y)

    xsign = uix & 0x8000
    ysign = uiy & 0x8000

    xsign > ysign && return true
    xsign < ysign && return false

    return uix < uiy
end

function Base.:(<=)(x::LogFixPoint16,y::LogFixPoint16)
    isnan(x) | isnan(y) && return nan(LogFixPoint16)

    uix = reinterpret(UInt16,x)
    uiy = reinterpret(UInt16,y)

    xsign = uix & 0x8000
    ysign = uiy & 0x8000

    xsign > ysign && return true
    xsign < ysign && return false

    return uix <= uiy
end

# Showing
function Base.show(io::IO, x::LogFixPoint16)
    if isnan(x)
        print(io, "NaR")
    else
		io2 = IOBuffer()
        print(io2,Float32(x))
        f = String(take!(io2))
        print(io,"LogFixPoint16("*f*")")
    end
end

Base.bitstring(x::LogFixPoint16) = bitstring(reinterpret(UInt16,x))

function Base.bitstring(x::LogFixPoint16,mode::Symbol)
    if mode == :split	# split into sign, integer, fraction
        s = bitstring(x)
		return "$(s[1]) $(s[2:nint[]+1]) $(s[nint[]+2:end])"
    else
        return bitstring(x)
    end
end


#Julia Gadfly Histogram

using Gadfly
plot(x=randn(113), Geom.histogram(bincount=10))


module LDUFacts

using LinearAlgebra

using LinearAlgebra: checksquare
using LinearAlgebra: PivotingStrategy, NoPivot
import LinearAlgebra: rmul!, lmul!, ldiv!, rdiv!, adjoint, det, logdet, logabsdet
import Base: *, \, /, size, inv

export ldu, ldu!, ldu_update!
export NoPivot, DiagonalPivot, FullPivot, PivotLike

struct LDUPivoted{T,V<:AbstractMatrix{T},F<:Real} <: Factorization{T}
    factors::V
    piv::Vector{Int}
    rank::Int
    tol::F
    info::Int
    pam::Vector{Int}
    ame::Vector{T}
end

struct LDUPerm{T}
    piv::Vector{Int}
    pam::Vector{Int}
    ame::Vector{T}
    LDUPerm(piv, pam, ame::AbstractVector{T}) where T = new{T}(piv, pam, Vector{T}(ame))
end

struct DiagonalPivot <: PivotingStrategy end
struct FullPivot <: PivotingStrategy end
struct PivotLike{T} <: PivotingStrategy
    piv::Vector{Int}
    pam::Vector{Int}
    ame::Vector{T}
    PivotLike(la::LDUPivoted{T}) where T = new{T}(la.piv, la.pam, la.ame)
end

const SATMatrix{T} = Union{StridedMatrix{T}, Adjoint{T,<:StridedMatrix}, Transpose{T,<:StridedMatrix}}

include("ldufact.jl")
include("lduupdate.jl")

end # module


V, (VV,EV,FV,CV) = Lar.simplex(3, true);
#==
tetra = V, EV,FV,CV;
twotetra = Lar.Struct([ tetra, Lar.t(1.,1.,1.), tetra ]);
V,EV,FV,CV = Lar.struct2lar(twotetra);
==#
#==
cop_EV = Lar.coboundary_0(EV::Lar.Cells);
cop_EW = convert(Lar.ChainOp, cop_EV);
cop_FE = Lar.coboundary_1(V, FV::Lar.Cells, EV::Lar.Cells);

triangulated_faces = Lar.triangulate(convert(Lar.Points, V'), [cop_EV, cop_FE])

FE = Array{Int64,1}[
    [1, 2, 3],
    [1, 4, 5],
    [2, 4, 6],
    [3, 5, 6, 7, 8, 9],
    [7, 8, 9],
    [7, 10, 12],
    [7, 11, 13, 14],
    [8, 10, 15],
    [8, 11, 16, 17],
    [9, 12, 15],
    [9, 13, 16, 18],
    [14, 17, 18]
];

triFV = Array{Array{Int64,1},1}[
    [[2, 3, 1]], [[4, 2, 1]], [[4, 3, 1]],
    [[2, 3, 5]], [[6, 3, 7]], [[5, 3, 6]],
    [[7, 4, 2]], [[4, 7, 3]], [[7, 2, 5]],
    [[5, 9, 10]], [[10, 6, 5]], [[11, 7, 9]],
    [[5, 9, 7]], [[11, 7, 10]], [[6, 10, 7]],
    [[9, 10, 11]]
];
myV = [
    0.0 1.0 0.0 0.0 0.5 0.25 0.25 0.25 1.25 0.25 0.25;
    0.0 0.0 1.0 0.0 0.25 0.5 0.25 0.25 0.25 1.25 0.25;
    0.0 0.0 0.0 1.0 0.25 0.25 0.5 0.25 0.25 0.25 1.25
];
GL.VIEW(GL.GLExplode(myV,myFV,1.5,1.5,1.5,99));
==#

p = [.1;.1;.1]


function _density(data; xlims = (-Inf, Inf), trim = false, kwargs...)
    k = kde(data; kwargs...)
    x, y = k.x, k.density
    xmin, xmax = xlims
    xmin = max(xmin, minimum(data))
    xmax = min(xmax, maximum(data))
    if trim
        for i in eachindex(x, y)
            xmin ≤ x[i] ≤ xmax || (y[i] = NaN)
        end
    end
    return spec(:Lines) * style(x, y)
end

function _density(datax, datay; xlims = (-Inf, Inf), ylims = (-Inf, Inf), trim = false, kwargs...)
    k = kde((datax, datay); kwargs...)
    x, y, z = k.x, k.y, k.density
    xmin, xmax = xlims
    xmin = max(xmin, minimum(datax))
    xmax = min(xmax, maximum(datax))
    ymin, ymax = ylims
    ymin = max(ymin, minimum(datay))
    ymax = min(ymax, maximum(datay))
    if trim
        for i in eachindex(x, y)
            xmin ≤ x[i] ≤ xmax && ymin ≤ y[i] ≤ ymax || (z[i] = NaN)
        end
    end
    return spec(:Heatmap) * style(x, y, z)
end

const density = Analysis(_density)


@testset "Test Learn Offset" begin
    y = repmat([1; 2; 3; 4], 1, 10)
    u = repmat([5; 6; 7], 1, 10)

    yterms = [[1, 2, 3], [1, 2], [2], [3, 4]]
    uterms = [[1], [2, 3], [4, 5]]

    mdl = Linear(13, 4)

    new_mdl = learn_offset(mdl, yterms, uterms, IdData(y, u))

    @test new_mdl.mdls[1].a == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
    @test new_mdl.mdls[1].b == -[1, 1, 1, 2, 2, 3, 4, 4, 5, 6, 6, 7, 7]
    @test new_mdl.mdls[3].a == [1, 1, 1, 1]
    @test new_mdl.mdls[3].b == [1, 2, 3, 4]
end

@testset "Test Learn Normalization" begin
    @testset "Basic Usage" begin
        y = repmat([1 3; 1 3], 1, 1000)
        u = repmat([-1 1], 1, 1000)

        yterms = [[1, 2], [1, 2]]
        uterms = [[1, 2]]

        mdl = Linear(6, 2)

        new_mdl = learn_normalization(mdl, yterms, uterms, IdData(y, u); nσ=1)

        @test new_mdl.mdls[1].a ≈ 1./[1, 1, 1, 1, 1, 1] atol=1e-1
        @test new_mdl.mdls[1].b ≈  -[2, 2, 2, 2, 0, 0]./[1, 1, 1, 1, 1, 1] atol=1e-1
        @test new_mdl.mdls[3].a ≈ [1, 1] atol=1e-1
        @test new_mdl.mdls[3].b ≈ [2, 2] atol=1e-1
    end

    @testset "Zero Standard Deviation" begin
        y = repmat([1 1; -2 6], 1, 1000)
        u = repmat([-1 1], 1, 1000)

        yterms = [[1, 2], [1, 2]]
        uterms = [[1, 2]]

        mdl = Linear(6, 2)

        new_mdl = learn_normalization(mdl, yterms, uterms, IdData(y, u); nσ=1)

        @test new_mdl.mdls[1].a ≈ 1./[1, 1, 4, 4, 1, 1]  atol=1e-1
        @test new_mdl.mdls[1].b ≈ -[1, 1, 2, 2, 0, 0]./[1, 1, 4, 4, 1, 1]  atol=1e-1
        @test new_mdl.mdls[3].a ≈ [1, 4]  atol=1e-1
        @test new_mdl.mdls[3].b ≈ [1, 2]  atol=1e-1
    end
end


#
# Functions to find the appropriate ARMA model from data.
#

using FFTW
using Statistics
include("RandomMatrix.jl")
include("optimize_exponentials.jl")

function padded_length(N::Integer)
    powerof2 = round(Int, 2^ceil(log2(N)))
    if N == powerof2
        return N
    elseif N >= 0.75 * powerof2
        return powerof2
    elseif N >= 0.625 * powerof2
        return div(3*powerof2, 4)
    else
        return div(5*powerof2, 8)
    end
end

"""
    estimate_covariance(timeseries::Vector, nsamp::Int, chunklength::Int)

Take a noise sequence (as a 1D vector or as columns of a 2D matrix) and
estimate the sample-sample covariance vector from these.

`nsamp` and `chunklength` are both optional arguments.

`nsamp` is the number of samples to compute from the covariance. If not a positive
number, then compute all possible values.

`chunklength` is the amount of data to be considered at once. If omitted, then this
will be either the length of the timeseries or 15*nsamp, whichever is shorter.

The idea is that a full computation is very inefficient. Making chunklength
much smaller than the timeseries length allows computation more efficient, but
at the cost of losing a small amount of correlation between values in successive
chunks. The recommended value for long data sets is to let chunklength be at
least 10-20x nsamp.
"""
estimate_covariance(timeseries::AbstractVector) = estimate_covariance(timeseries, length(timeseries))

function estimate_covariance(timeseries::AbstractVector, nsamp::Int)
    N = length(timeseries)
    ideal_chunksize = 15*nsamp
    nchunks = div(N-1+ideal_chunksize, ideal_chunksize)
    chunklength = div(N, nchunks)
    estimate_covariance(timeseries, nsamp, chunklength)
end

function estimate_covariance(timeseries::Vector, nsamp::Int, chunklength::Int)
    N = length(timeseries)
    if nsamp > N
        error("Cannot compute $(nsamp) covariance values from a length-$N data.")
    end
    paddedsize = padded_length(chunklength+nsamp)
    padded_data = fill(0.0, paddedsize)
    result = fill(0.0, nsamp)

    i=0
    chunks_consumed = 0
    while i+chunklength <= N
        datamean = mean(timeseries[i+1:i+chunklength])
        padded_data[1:chunklength] = timeseries[i+1:i+chunklength] .- datamean
        chunks_consumed += 1

        power = abs2.(rfft(padded_data))
        acsum = irfft(power, length(padded_data))
        result += acsum[1:nsamp]
        i += chunklength
    end
    result ./ (chunks_consumed * collect(chunklength:-1:chunklength+1-nsamp))
end

"""
    main_exponentials(data, nexp; minexp=nothing)

Find the `nexp` "main exponentials" in the time stream `data`.

Specifically, follow the prescription of P. De Groen & B. De Moor (1987).
"The fit of a sum of exponentials to noisy data." J. Computational & Applied
Mathematics, vol. 20, pages 175–187.

We build a Hankel matrix H whose columns are contiguous segments of the time stream.
Perform a singular value decomposition to find the rank-`nexp` decomposition of H,
from which we can construct the square matrix A (of size `nexp`) which is known
to be similar to the "system matrix" (i.e., a time-step-advance matrix). The
latter has eigenvalues equal to the exponentials comprising the time stream, so
A (by similarity) has the same eigenvalues.

That's the best explanation I have. It might be more honest to say that we
find the exponentials by magic.

Note that the current implementation does not construct the exact SVD of H,
but rather uses a randomized matrix technique to compute an *approximate* SVD
very efficiently, containing only a specified number of leading singular vectors.

If `minexp==nothing` (the default), then only the length-`nexp` model is found.
If `minexp ≤ nexp`, then this function returns an array of arrays, one for each
value `p` with `minexp ≤ p ≤ nexp`, so that the length of the first is `minexp`
and the length of the last is `nexp`.

TODO: we need a better heuristic for constructing H out of a sufficient number
of segments from the data, without allowing it to get insanely large when `data`
is very, very long.
"""
function main_exponentials(data::Vector, nexp::Int; minexp=nothing)
    N = length(data)
    if 2nexp > N
        error("Cannot compute $(nexp) exponentials from data with fewer than twice as many data values.")
    end
    if (minexp != nothing) && (minexp > nexp)
        error("The value of `minexp` must be nothing or a number no larger than `nexp`.")
    end

    ncol = min(40 + 5nexp, div(N, 2))
    H = fill(0.0, N+1-ncol, ncol)
    for c=1:ncol
        H[:,c] = data[c:c+N-ncol]
    end
    U,s,V = find_svd_randomly(H[1:end-1,:], nexp)
    W = Diagonal(s .^ (-0.5))
    A = W*U'*H[2:end,:]*V*W

    # By default, return the exponential set of size nexp only.
    if minexp == nothing
        return sortbases!(eigvals(A))
    end

    # The minexp has been set, so return all sizes minexp...nexp, inclusive.
    exps = []
    for p = minexp:nexp
        A = W[1:p, 1:p]*U[:,1:p]'*H[2:end,:]*V[:,1:p]*W[1:p,1:p]
        push!(exps, sortbases!(eigvals(A)))
    end
    exps
end

"""
    sortbases!(b)

Re-order the elements of `b` in place, such that for every pair b[1],b[2] and so
on, either both are real, or they are complex conjugates of one another.
Elements of `b` with negative imaginary parts are simply assumed to be
conjugates of those with positive imaginary parts. The complex pairs will come
first, followed by the real elements.

Returns the sorted `b`.
"""
function sortbases!(b::Vector)
    real_ones = real(b[imag(b) .== 0])
    sort!(real_ones, rev=true)
    if length(real_ones) < length(b)
        imag_ones = b[imag(b) .> 0]
        if length(imag_ones) != sum(imag(b) .< 0)
            error("sortbases!(b) requires that b contain equal numbers of elements with imag(x) > 0 and < 0.")
        end
        for i=1:length(imag_ones)
            b[2i-1] = imag_ones[i]
            b[2i] = conj(imag_ones[i])
        end
        b[2*length(imag_ones)+1:end] = real_ones[:]
    else
        b[:] = real_ones[:]
    end
    b
end

"""
    B2C(B)

Converts an array `B` of exponential bases to the coefficients of the set of
monic quadratic polynomials.

It is assumed that B are arranged such that B[1] and B[2] are either both real
or are complex conjugates of one another, and similarly for elements (3,4), (5,6),
and so on. Use `sortbases!(B)` to achieve this.

If `B` violates these assumptions, then the computed `C` will make no sense.

Returns `C` such that B[1:2] are the roots of x^2+C[1]*X+C[2], and similarly for
[3:4] and all other pairs. If length(B) is odd, then B[end] is the root of x+C[end],
thus `-1 = B[end]*C[end]`.
"""
function B2C(B::AbstractVector{T}) where {T<:Number}
    C = fill(0.0, length(B))
    n = length(B)
    for i=1:2:n-1
        C[i] = -real(B[i]+B[i+1])
        C[i+1] = real(B[i]*B[i+1])
    end
    if n%2 > 0
        C[end] = -1/real(B[end])
    end
    C
end

"""
    C2B(C)

Converts a real array of polynomial coefficients `C` to their roots. See `B2C(B)`
for more.

Returns `B`, the possibly complex roots.
"""
function C2B(C::AbstractVector{T}) where {T<:Real}
    B = fill(zero(Complex{eltype(C)}), size(C))
    n = length(C)
    iscomplex = false
    for i = 1:2:n-1
        x = -0.5C[i]
        disc = x*x-C[i+1]
        if disc >= 0
            y::ComplexF64 = sqrt(disc)
        else
            y = 1im*sqrt(-disc)
            iscomplex = true
        end
        B[i] = x+y
        B[i+1] = x-y
    end
    if n%2 > 0
        B[end] = -1/C[end]
    end
    if !iscomplex
        return real(B)
    end
    B
end

"""
    findA(t, r, B, [w=weights])

Compute the amplitudes `A` for a linear model of the form

r_t ≈ f(t) = Sum_{i=1...p} A[i] B[i]^t

by minimizing

Sum_t [w_t (r_t - f(t))^2].

This is a linear problem, given `B`. If `w` is not given, then equal weights on
all data are assumed.

Returns the array of amplitudes `A`.
"""
function findA(t::AbstractVector, r::Vector, B::Vector{T}; w=nothing) where {T<:Number}
    @assert length(t) == length(r)
    if w==nothing
        w = ones(r)
    end
    wr = w.*r
    p = length(B)
    M = fill(zero(T), p, p)
    D = fill(zero(T), p)
    for i=1:p
        for j=1:i
            M[j,i] = M[i,j] = sum(w .* (B[i]*B[j]).^t)
        end
        D[i] = sum(wr .* B[i].^t)
    end

    # A = M\D is theoretically correct, but don't do this: can be ill-conditioned in certain cases
    A = pinv(M) * D
end

"""
    exponential_model(t, A, B)`

Computes and returns the sum-of-exponentials model with amplitudes `A` and
exponential bases `B` at time steps `t`.
"""
function exponential_model(t::AbstractVector, A::Vector, B::Vector)
    r = fill(0.0, length(t))
    for i=1:length(A)
        r .+= real(A[i]* B[i].^t)
    end
    r
end

"""
    squashgrowingexp!(bases::Vector, epsfrom1)

Replace any growing exponentials, that is, where  `abs(bases[j]) >= 1`, with
random decaying exponentials. These will be between `1-epsfrom1` and 1 in absolute
value (suggest `epsfrom1 <= 0.01`).
"""
function squashgrowingexp!(bases::Vector, epsfrom1::Float64)
    lastscale = 1 - 0.5*epsfrom1
    for (i,b) = enumerate(bases)
        if abs(b) >= 1 - 0.01*epsfrom1
            if i%2==1 || imag(b) == 0.0
                lastscale = (1-rand(1)[1]*epsfrom1) / abs(b)
            end
            bases[i] *= lastscale
        end
    end
end

"""
    fit_exponentials(data; pmin=0, pmax=6, w=nothing, deltar=nothing, good_enough=0.0)

Fit the time series `data` (regularly sampled) as a sum of possibly complex
exponentials numbering at least `pmin` and at most `pmax`. Uses
`main_exponentials` to determine a starting guess for the exponential bases.
Then a nonlinear fit is attempted, using `NLopt` with the `:LD_MMA` method.

Returns `A,B`, the amplitudes and exponential bases of the best-fit model. The
length of either gives the order of the best-fit model.

# Optional Arguments:
The penalty function is a weighted sum-of-square-errors. If the optional
argument `w` is given, it must be a vector of the same length as `data`, and
then the penalty is the sum over all samples of `w[i]*(model[i]-data[i])^2`.

If `w` is absent, then the weights are taken to be all equal to `deltar^(-2)`.

If `deltar` is also not given, then it is estimated from the standard deviation
of the successive-differences from the last 1/4 of the samples in `data`, on the
assumption that these are noise-dominated. If they aren't, then the caller must
offer a better figure for `w` or `deltar`.

If `pmin` < `pmax`, then models of more than one order will be tried, and the one
with the lowest penalty function will be returned.

If `good_enough` is assigned a positive value, and if any model order `p<pmax`
produces a penalty function less than this value, then this function returns
early with the lowest-order model that meets the criterion.

See also: [`fitARMA`](@ref)
"""
function fit_exponentials(data::Vector; pmin::Int=0, pmax::Int=6,
    w=nothing, deltar=nothing, good_enough::Float64=0.0)

    guess_exponentials = main_exponentials(data, pmax, minexp=pmin)

    N = length(data)
    if w == nothing
        # If deltar isn't given, use the heuristic that the std dev of the
        # diff of the last values in data estimates sqrt(2) times deltar.
        if deltar == nothing
            n = div(N, 4)
            deltar = std(diff(data[end-n:end]), mean=0) / sqrt(2)
        end

        w = ones(Float64, N) / float(deltar^2)
    end

    best_cost = +Inf
    best_fit = nothing
    for (guessB, p) = zip(guess_exponentials, pmin:pmax)
        squashgrowingexp!(guessB, 0.25/N)
        guessC = B2C(guessB)
        cost, fitC = optimize_exponentials(data, w, guessC)
        if (cost < best_cost)
            best_cost = cost
            best_fit = fitC
            # Stop early if cost is small enough.
            if cost <= good_enough
                break
            end
        end
    end

    bestB = C2B(best_fit)
    squashgrowingexp!(bestB, 0.25/N)

    t = 0:(length(data)-1)
    bestA = findA(t, data, bestB, w=w)
    bestA, bestB
end

"""
    fitARMA(covariance, p, q; kwargs)

Fit an `ARMAModel` of order `(p,q)` or lower to the `covariance` data. The model
order with the lowest cost function will be returned.

# Optional Arguments
- `pmin`: the minimum allowed ARMA order.
- `w`: the vector of weights to apply to the data (default: nothing, meaning equal weights)
- `deltar`: rms uncertainty on the covariance values.
- `good_enough`: if given, then the lowest-order model with a cost function at least
   this small will be used, rather than trying all orders up to `p`.

See also: [`fit_exponentials`](@ref)
"""
function fitARMA(covariance::Vector, p::Int, q::Int;
    w=nothing, deltar=nothing, good_enough::Float64=0.0, pmin::Int=0)
    nspecial = max(1+q-p, 0)
    amplitudes, bases = fit_exponentials(covariance[1+nspecial:end], pmin=pmin, pmax=p,
        w=w, deltar=deltar, good_enough=good_enough)
    if nspecial > 0
        amplitudes ./= (bases .^ nspecial)
    end
    ARMAModel(bases, amplitudes, covariance[1:nspecial])
end

# Absent further information, we find that q=p is a good choice most of the time,
# so let that be a default. Also, if we need a maximum order, we can supply
# the rather arbitrary guess of 6.
fitARMA(covariance::Vector, p::Int; kwargs...) = fitARMA(covariance, p, p; kwargs...)
fitARMA(covariance::Vector; kwargs...) = fitARMA(covariance, 6; kwargs...)


function HyperLouvain(H::hypergraph,kmax::Int64,Ω::IntensityFunction,maxits::Int64=100,bigInt::Bool=true;α,verbose=true,scan_order="lexical", Z0 = collect(1:length(H.D)))
    HyperLouvain(H,Ω;α=α,kmax=kmax,maxits=maxits,bigInt=bigInt,verbose=verbose,scan_order=scan_order, Z0 = Z0)
end

function HyperLouvain(H::hypergraph,Ω::IntensityFunction;α,kmax=maximum(keys(H.E)),maxits::Int64=100,bigInt::Bool=true,verbose=true,scan_order="lexical", Z0 = collect(1:length(H.D)))
    """
    Basic step Louvain algorithm: iterate through nodes and greedily move
    nodes to adjacent clusters. Does not form supernodes and does not recurse.

    H: hypergraph
    Ω: group interaction function, as constructed by ΩFromDict(D)
    bigInt::Bool: whether to convert the degree sequence to an array of BigInt when evaluating the volume term in second_term_eval(). Recommended.
    return: Z::array{Int64, 1}: array of group labels.
    """
    Hyp, w = hyperedge_formatting(H)    # hyperedge to node list
    node2edges = EdgeMap(H)             # node to hyperedge list
    n = length(H.D)

    Neighbs = NeighborList(node2edges, Hyp)  # Store neighbors of each node

    Ẑ = Z0                              # All nodes start in singleton clusters
    Clusters = Vector{Vector{Int64}}()  # Also store as cluster list
    for v = 1:n
        push!(Clusters, Vector{Int64}())
        push!(Clusters[v],v)
    end

    iter = 0
    improving = true
    changemade = false

    # Data structures for fast local changes in volumes
    r = kmax
    V, μ, M = evalSums(Ẑ, H.D, r;constants=false, bigInt=bigInt)
    C = evalConstants(r)

    cuts = evalCuts(H, Ẑ, Ω)

    # initialize memoization of edge penalties
    Pn           = [log(Ω.ω(Ω.P(collect(1:i)), α)) for i in 1:kmax]
    edge2penalty = [w[e]*Pn[length(Hyp[e])] for e in 1:length(Hyp)]

    # Main Greedy Local Move Loop
    while improving && iter < maxits

        iter += 1
        if mod(iter,1) == 0;
            if verbose println("Louvain Iteration $iter, Q = $(round(Float64(modularity(H, Ẑ, Ω;α = α)), digits = 2))") end
        end
        improving = false

        scan = scan_order == "lexical" ? (1:n) : Random.shuffle(1:n)

        for i = scan                     # visit each node in turn

            Ci_ind = Ẑ[i]               # Cluster index for node i
            Ci = Clusters[Ci_ind]       # Indices of nodes in i's cluster
            Ni = Neighbs[i]             # Get the indices of i's neighbors
            NC = unique(Ẑ[Ni])          # Clusters adjacent to i (candidate moves)

            BestZ = Ẑ[i]                # The default is do nothing
            BestImprove = 0
            V_best = 0
            μ_best = 0
            M_best = 0
            pen_best = 0

            # Check if it's better to move to a nearby cluster, Cj
            for j = 1:length(NC)
                Cj_ind = NC[j]

                # Check improvement for move from cluster i to to cluster j
                if Cj_ind == Ci_ind
                    change = 0
                else
                    ΔV, Δμ, ΔM = increments(V, μ, M, [i], Cj_ind, H.D, Ẑ)
                    voldiff = sum(Ω.ω(Ω.aggregator(p),α)*ΔM[p]*C[p] for p in keys(ΔM))

                    cdiff, Δpen = CutDiff_OneNode(Hyp,w,node2edges,Ẑ,i,Cj_ind,edge2penalty,Ω;α=α) # NOTE: need to check on this
                    change =  cdiff - voldiff
#                     println(cdiff, " ", voldiff)
                end

                # Check if this is currently the best possible greedy move
                if change - BestImprove > 1e-8
                    V_best = ΔV
                    μ_best = Δμ
                    M_best = ΔM
                    pen_best = Δpen
                    BestImprove = change
                    BestZ = Cj_ind
                    improving = true
                end
            end

            # Move i to the best new cluster, if it strictly improves modularity
            if BestImprove > 1e-8

                changemade = true
                improving = true    # we have a reason to keep iterating!

                # update increments for volume computation
                V, μ, M = addIncrements(V, μ, M, V_best, μ_best, M_best)

                # update clustering
                ci_old = Ẑ[i]
                Ẑ[i] = BestZ

                # Remove i from its old cluster and add it to its new cluster
                Clusters[ci_old] = setdiff(Clusters[ci_old],i)
                push!(Clusters[BestZ],i)

                # update map from edge ID to penalty at edge
                for eid in keys(pen_best)
                    newpenalty = pen_best[eid]
                    edge2penalty[eid] = newpenalty
                end
            end
        end
    end

    if ~changemade
        if verbose println("No nodes moved clusters") end
    end
    Ẑ, Clusters = renumber(Ẑ,Clusters)
    return Ẑ

end

function compute_moddiff(edge2part,Cts,Hyp,w,node2edges,V::Array, μ::Array, M::Dict,I_::T, t::Int64, D::Vector{Int64}, Z::Vector{Int64},C::Dict,Ω::IntensityFunction;α) where T <: Union{Int64, Vector{Int64}}
    """
    NOTE: I_ can now be a Vector{Int64} of nodes, all of which are assumed to belong in the same cluster.
    """
    # increments due to proposal
    ΔV, Δμ, ΔM = increments(V, μ, M, I_, t, D, Z)

    voldiff = sum(Ω.ω(Ω.aggregator(p),α)*ΔM[p]*C[p] for p in keys(M))

    ΔC, Δe2p = CutDiff_Many(Cts, I_, t, Z, Hyp, w, node2edges,edge2part)
    cdiff = sum(ΔC[p]*log(Ω.ω(Ω.aggregator(p), α)) for p in keys(ΔC))
    mdiff = cdiff-voldiff

    return mdiff, ΔV, Δμ, ΔM, ΔC, Δe2p
end

function SuperNodeStep(H::hypergraph,Z::Vector{Int64},kmax::Int64,Ω,maxits::Int64=100,bigInt::Bool=true;α,verbose=true)
    """
    A Louvain step, but starting with all nodes in an arbitrary initial cluster
    assignment Z. Louvain only considers moving an entire cluster at once.
    Running this code with Z = collect(1:n) is equivalent to HypergraphLouvain

    H: hypergraph
    Z: initial cluster assignment as a length n vector
    Ω: group interaction function, as constructed by ΩFromDict(D)
    bigInt::Bool: whether to convert the degree sequence to an array of BigInt when evaluating the volume term in second_term_eval(). Recommended.
    return: Z::array{Int64, 1}: array of group labels.
    """
    Hyp, w = hyperedge_formatting(H)    # hyperedge to node list
    node2edges = EdgeMap(H)             # node to hyperedge list
    n = length(H.D)

    Ẑ = renumber(Z)     # Ensure clusters go from 1 to number unique clusters
    sn = maximum(Ẑ)     # number of supernodes

    Clusters = Vector{Vector{Int64}}()  # Also store as cluster list
    for j = 1:sn
        push!(Clusters, Vector{Int64}())
    end
    for v = 1:n
        push!(Clusters[Ẑ[v]],v) # put node v in cluster Z[v]
    end
    SuperNodes = deepcopy(Clusters)

    # Store indices of supernodes that neighbor each supernode
    tic = time()
    Neighbs = SuperNeighborList(Hyp, SuperNodes,n)
    toc = time()-tic
    if verbose
        @show toc
    end

    # Note: There's a very subtle difference between the initial clustering,
    # which defines a set of supernodes, and the current clustering that
    # is being built by Louvain.
    # Supernodes don't change--they move as a unit. Clusters change.

    iter = 0
    improving = true
    changemade = false

    # Data structures for fast local changes in volumes
    r = kmax
    V, μ, M = evalSums(Ẑ, H.D, r;constants=false, bigInt=bigInt);
    C = evalConstants(r)

    # Data structures for fast local changes in cuts
    Cts = evalCuts(Ẑ,H)                    # edge partition -> no. times it appears in Z
    edge2part = Vector{Vector{Int64}}()    # edge id -> partition type
    for e = 1:length(Hyp)
        edge = Hyp[e]
        push!(edge2part,partitionize(Ẑ[edge]))
    end
    t1 = 0
    # Main Greedy Local Move Loop
    while improving && iter < maxits

        iter += 1
        if mod(iter,1) == 0;
            if verbose println("Louvain Iteration $iter") end
        end
        improving = false

        for i = 1:sn          # visit each supernode in turn

            S = SuperNodes[i]       # Indices associated with this supernode
            rep = S[1]              # A "representative" node from this supernode
            Ci_ind = Ẑ[rep]         # Cluster associated with this supernode
            Ci = Clusters[Ci_ind]   # Indices of nodes in this cluster

            Ni = Neighbs[i]         # All supernodes that neighbor this supernode

            # Get all the clusters of neighboring supernodes
            NC = Vector{Int64}()
            for ni in Ni
                push!(NC,Ẑ[SuperNodes[ni][1]])
            end
            NC = unique(NC)
            # println("$i = i, length(NC) = $(length(NC))")

            BestZ = Ci_ind          # The default is to not move the cluster
            BestImprove = 0
            V_best = 0
            μ_best = 0
            M_best = 0
            C_best = 0
            e2p_best = 0

            # Check if it's better to move to a nearby cluster, Cj
            for j = 1:length(NC)
                Cj_ind = NC[j]

                # Check improvement for move from cluster i to to cluster j
                if Cj_ind == Ci_ind
                    change = 0
                else
                    tic = time()
                    change, ΔV, Δμ, ΔM, ΔC, Δe2p = compute_moddiff(edge2part,Cts,Hyp, w, node2edges,V, μ, M, S, Cj_ind, H.D, Ẑ, C, Ω; α=α)
                    t1 += time()-tic
                end

                # Check if this is currently the best possible greedy move to make
                if change - BestImprove > 1e-8
                    V_best = ΔV
                    μ_best = Δμ
                    M_best = ΔM
                    C_best = ΔC
                    e2p_best = Δe2p
                    BestImprove = change
                    BestZ = Cj_ind
                    improving = true
                end
            end

            # Move supernode i to the best new cluster, only if it strictly improves modularity
            if BestImprove > 1e-8

                changemade = true
                improving = true # we have a reason to keep iterating!

                # update increments for volume computation
                V, μ, M = addIncrements(V, μ, M, V_best, μ_best, M_best)

                # update clustering
                ci_old = Ẑ[rep] # index of the old cluster
                Ẑ[S] .= BestZ

                # Remove all nodes from supernode i from their old cluster.
                Clusters[ci_old] = setdiff(Clusters[ci_old],S)

                # Add them to their new cluster
                append!(Clusters[BestZ],S)

                # Update map from partition type to number of those partitions
                for part in keys(C_best)
                    Δcount = C_best[part]
                    Cts[part] = get(Cts,part,0) + Δcount
                end

                # update map from edge ID to partition type
                for eid in keys(e2p_best)
                    newpart = e2p_best[eid]
                    edge2part[eid] = newpart
                end
            end
        end
    end
    Ẑ, Clusters = renumber(Ẑ,Clusters)
    if verbose
        @show t1
    end
    return Ẑ, changemade
end

function SuperNodeLouvain(H::hypergraph,kmax::Int64,Ω,maxits::Int64=100,bigInt::Bool=true;α,verbose=true,scan_order="lexical", Z0 = nothing)
    SuperNodeLouvain(H,Ω;α=α,kmax=kmax,maxits=maxits,bigInt=bigInt,verbose=verbose,scan_order=scan_order, Z0 = copy(Z0))
end

function SuperNodeLouvain(H::hypergraph,Ω::IntensityFunction;α,kmax=maximum(keys(H.E)),maxits::Int64=100,bigInt::Bool=true,verbose=true,scan_order="lexical", Z0 = collect(1:length(H.D)))
    """
    Running Louvain and then the super-node louvain steps until no more
    progress is possible

    H: hypergraph
    Ω: group interaction function, as constructed by ΩFromDict(D)
    bigInt::Bool: whether to convert the degree sequence to an array of BigInt when evaluating the volume term in second_term_eval(). Recommended.
    return: Z::array{Int64, 1}: array of group labels.
    """

    phase = 1

    if verbose println("Faster SuperNode Louvain: Phase $phase") end
    Ẑ = HyperLouvain(H,kmax,Ω;α=α,verbose=verbose,scan_order=scan_order, Z0 = copy(Z0))

    n = length(Ẑ)

    changed = true

    while changed
        phase += 1
        if verbose println("SuperNode Louvain: Phase $phase") end
        Ẑ, changed = SuperNodeStep(H,Ẑ,kmax,Ω,maxits;α=α,verbose=verbose)
    end

    return Ẑ
end

function evalCuts(Z::Array{Int64,1}, H::hypergraph)
    """
    Gets map from partition type to the number of times that partition type
    shows up in edges for the clustering Z
    """
    C = Dict{Vector{Int64}, Int64}()
    for k in keys(H.E)
        Ek = H.E[k]
        for e in keys(Ek)
            p = partitionize(Z[e])
            C[p] = get(C, p, 0) + Ek[e]
        end
    end
    return C
end

function evalCuts(H::hypergraph, Z::Array{Int64,1}, Ω::IntensityFunction)

    C = Dict(p => 0 for p in Ω.range)
    for k in keys(H.E)
        Ek = H.E[k]
        for e in keys(Ek)
            p = Ω.P(Z[e])
            C[p] = get(C, p, 0) + Ek[e]
        end
    end
    return C
end

function CutDiff_Many(C, I, t, Z, Hyp, w, node2edges,edge2part)
    """
    C: Dict{Vector{Int64}, Int64}, the dict of current cut values as would be produced by evalCuts().
        C is keyed by partition vectors p.
        C[p] is the # of edges with label partition p.
        Note: Ω is NOT used in the calculation of C
    I: Vector{Int64}, the list of nodes to move, assumed to be in the same group
    t: Int64, the new group to which to move the nodes I
    Z: Vector{Int64}, the current group assignments
    Hyp: list of nodes in each hyperedge
    w:  weight of the hyperedge (number of times this set of nodes is assigned a hyperedge)
    node2edges: node ID to list of edges the node is in.
    edge2part: for hyperedge index e, returns the current partitionize vector for e for the clustering Z
    return: ΔC, a Dict{Vector{Int64}, Int64}, where ΔC[p] is the change in C[p] caused by moving the nodes in I to group t
    """

    E_id = Vector{Int64}()  # edge IDs that invole some node from I
    for i in I
        append!(E_id,node2edges[i])
    end
    E = unique(E_id)

    # This stores proposed changes to edge partitions
    Δe2p = Dict{Int64,Vector{Int64}}()

    # Store changes in the number of partition types in the clustering
    ΔC = Dict(p => 0 for p in keys(C))

    Z_prop = copy(Z)    # Store the proposed new clustering
    Z_prop[I] .= t

    for eid in E
        e = Hyp[eid]    # get nodes in edge
        we = w[eid]     # get the edge weight

        p_old = edge2part[eid]  # instead of: pold = partitionize(Z[e])
        p_new = partitionize(Z_prop[e])
        Δe2p[eid] = p_new

        # Update number of partition types for p_old and p_new
        ΔC[p_old] -= we
        ΔC[p_new] = get(ΔC,p_new,0) + we

    end
    return ΔC, Δe2p
end

function CutDiff_OneNode(H::Vector{Vector{Int64}},w::Array{Float64,1},node2edges::Vector{Vector{Int64}},Z::Array{Int64,1}, I::Int64,J::Int64,edge2penalty::Union{Vector{Float64}, Vector{BigFloat}},Ω::IntensityFunction; α)
    """
    CutDiff: Compute change in the first term of the modularity function
    resulting from moving a node I to cluster J.
    """
    orig = Z[I]
    He   = node2edges[I]
    Δpen = Dict{Int64,Float64}()    # Change in penalty for moving a node

    obj  = 0
    Z[I] = J
    for i in He
        e        = H[i]
        z        = Z[e]                      # set of clusters
        new_pen = w[i]*log(Ω.ω(Ω.P(z),α))   # new penalty if this move is made
        Δpen[i]  = new_pen
        obj     += Δpen[i]-edge2penalty[i]
    end
    Z[I] = orig

    return obj, Δpen
end

function SuperNeighborList(Hyp, SuperNodes,n)
    """
    For an initial clustering into supernodes, get a list of other supernodes
    adjacenct to each supernode
    """
    H = elist2incidence(Hyp,n)
    m = size(H,1)
    sn = length(SuperNodes)
    I = Vector{Int64}()
    J = Vector{Int64}()

    for s = 1:sn
        # For supernode j, get indices of all other supernodes i such that
        # some node in j shares a node with some node in i
        S = SuperNodes[s]

        # this is silly and slow. But also not a bottleneck, so I'm
        # not worrying about making it faster right now
        Edges = findall(x->x>0,vec(sum(H[:,S],dims = 2)))  # all hyperedges nodes from S touch

        for e in Edges
            push!(I,s)
            push!(J,e)
        end
    end

    # binary (unweighted) supernodes-by-edges incidence matrix
    Hnew = SparseArrays.sparse(J,I,ones(length(I)),m,sn)
    Neighbs = NeighborList(Hnew, SparseArrays.sparse(Hnew'))
    return Neighbs
end


import Base: iterate
import Base.Iterators.take

"""
Structure to hold the iterates of the dual conditional gradient method.

DualCGIterable(M, b, A; optTol=√eps)

Construct a dual conditional gradient iterator for the problem

    minimize  ½‖Mx-b‖²  subject to gauge(A|x) ≤ 1.

The absolute tolerance `optTol` is used to determine the stopping condition:

    gap ≤ optTol(1+‖b‖₂)

"""
mutable struct DualCGIterable{T1<:Float64, T2<:Vector{T1}, T3<:AbstractArray{T1}, T4<:AbstractLinearOp, T5<:AbstractAtomicSet, T6<:AbstractAtom}
    M::T4           # linear operator
    Mt::T4          # adjoint of the linear operator
    b::T2           # observation
    A::T5           # atomic set
    r::T2           # residual
    z::T3           # negative gradient
    a::T6           # exposed atom
    Ma::T3          # image of the atom
    Mx::T3          # image of primal variable
    Δr::T3          # residual step
    τ::T1           # current scale of atomic set
    optTol::T1      # tolerance
    function DualCGIterable(M::AbstractLinearOp, b::Vector{Float64}, A::AbstractAtomicSet, optTol = sqrt(eps(1.0)))
        r = copy(b)
        Mt = M'
        z = Mt*r
        a = expose(A, z)
        Ma = M*a
        Mx = copy(Ma); fill!(Mx, 0.0)
        Δr = Ma - Mx
        τ = 0.0
        new{Float64, Vector{Float64}, AbstractArray{Float64}, AbstractLinearOp, AbstractAtomicSet, AbstractAtom}(M, Mt, b, A, r, z, a, Ma, Mx, Δr, τ, optTol)
    end
end

function oracleExit(p::DualCGIterable, gap::Float64)
    return gap < p.optTol
end

"""
Defines one complete iteration of the dual conditional gradient method.
"""
function iterate(p::DualCGIterable, k::Int=0)

    # Compute  M*a  where  a = expose(τA, z)
    expose!(p.A, p.z, p.a)
    # p.Ma .= p.M*p.a
    mul!(p.Ma, p.M, p.a)
    rmul!(p.Ma, p.τ)

    # Compute search direction Δr = M(a - x)
    @. p.Δr = p.Ma - p.Mx

    # Compute gap and possible exit
    gap = sumdot(p.Δr, p.r)
    oracleExit(p, gap) && return (gap, p.r), k+1

    if iszero(p.Mx)
        # Mx ← Ma
        p.Mx .= p.Ma
        # r ← b - Mx
        p.r .= p.b - sum(p.Mx, dims=2)[:]
    else
        # Linesearch
        α = linesearch(p.Δr, p.r)
        # r ← r - Δr*α
        mul!(p.r, p.Δr, α, -1, 1)
        # Mx ← Mx + Δr*Diagonal(α)
        broadcast!(*, p.Δr, p.Δr, α')
        BLAS.axpy!(1, p.Δr, p.Mx)
    end

    # Update gradient: z ← M'r
    mul!(p.z, p.Mt, p.r)

    return (gap, p.r), k+1
end

"""
    sumdot(X::Matrix, y::Vector)
Compute the sum of the dotproducts of each column of `X`
with the vector `y`.
"""
function sumdot(X::AbstractMatrix{Float64}, y::Vector{Float64})
    fdot = x -> dot(x, y)
    return sum(fdot, eachcol(X))
end

sumdot(x::AbstractVector{Float64}, y::Vector{Float64}) = dot(x, y)

"""
    setRadius!(p::DualCGIterable, τ::Float64)

Increase to `τ` the radius of the current Lasso problem. This routine
reflects the increase in the radius on various quantities of the dual
CG method.

If the current radius 0, then this routine only sets the new parameter,
but doesn't rescale any of the dual iterates.
"""
function setRadius!(p::DualCGIterable, τ::Float64)
    if p.τ > 0
        rescaleIterates!(p, τ/p.τ)
    end
    p.τ = τ
end

function rescaleIterates!(p::DualCGIterable, s::Float64)
    p.Mx .*= s
    Mxs = vec(sum(p.Mx, dims=2))
    p.r .= p.b .- Mxs
    p.z .= p.M'*p.r
end

function replaceIterates!(p::DualCGIterable, M::AbstractMatrix{Float64}, F::AbstractFace, c::Vector{Float64})
    MF = M*F
    p.Mx .= MF*c
    p.Mx .*= p.τ/sum(c)
    Mxs = vec(sum(p.Mx, dims=2))
    p.r .= p.b .- Mxs
    p.z .= p.M'*p.r
end

getNegGradient(p::DualCGIterable) = p.z
getResidual(p::DualCGIterable) = p.r

"""
    linesearch(Δr::Vector, r::Vector)

Solve univariate least-squares problem

    minimize_{θ∈[0,1]} ||θ⋅Δr-r||₂,

which has the solution

    θ = min(1, dot(Δr,Δr)/dot(Δr,Δr))
"""
function linesearch(Δr::AbstractVector{Float64}, r::Vector{Float64})
    return min(1.0, dot(Δr,r)/dot(Δr,Δr))
end

"""
    linesearch(Δr::Matrix, r::Vector)

Solve the box-constrained least squares problem

    minimize_{θ∈[0,1]ⁿ} ||Δr⋅θ - r||₂
"""
function linesearch(Δr::AbstractMatrix{Float64}, r::Vector{Float64})
    n = size(Δr, 2)
    bl = zeros(n)
    bu = ones(n)
    Q = Quadratic(Δr'*Δr, Δr'r)
    lbfgsq = LBFGSBQuad(Q, bl, bu, m=2)
    θ = boxquad(lbfgsq)
    return θ
end

"""
    primal recover
"""
function primalrecover!(p::DualCGIterable, sol::Solution, α::Float64, feaTol::Float64, exitFlag)
    flag = exitFlag
    s = support(p.A, p.z)
    face!(p.A, p.z/s, sol.Fnew)
    ϵ = copy(getResidual(p)); ϵ .*= sqrt(2*α)/norm(ϵ)
    face_project!(p.M, sol.Fnew, p.b-ϵ, sol.cnew, sol.r)
    sol.feasnew = norm(sol.r + ϵ)^2/2
    if sol.feasnew  ≤ sol.feas
        update!(sol)
    end
    if sol.feas - α ≤ feaTol
        flag = :feasible
    end
    return flag
end


using DSGE
using HDF5
path = dirname(@__FILE__)

# Test in model optimization
custom_settings = Dict{Symbol, Setting}(
    :date_forecast_start  => Setting(:date_forecast_start, quartertodate("2015-Q4")))
m = AnSchorfheide(custom_settings = custom_settings, testing = true)

file = "$path/../reference/optimize.h5"
x0 = h5read(file, "params")
data = h5read(file, "data")'
minimizer = h5read(file, "minimizer")
minimum = h5read(file, "minimum")
H_expected = h5read(file, "H")

# See src/estimate/estimate.jl
update!(m, x0)
n_iterations = 3

x0 = Float64[p.value for p in m.parameters]

out, H = optimize!(m, data; iterations=n_iterations)

@testset "Check optimize minimizers are the same" begin
    @test @test_matrix_approx_eq minimizer out.minimizer
    @test minimum ≈ out.minimum atol=5e-7
    @test @test_matrix_approx_eq H_expected H
end

nothing


using Pkg
Pkg.activate(".")
using RadMod2D

# model
include("./models2D.jl")
r1x = 0.8
r1y = 0.8
r2x = 1.2
r2y = 1.2
elemsize = 0.05
m = model_rect_in_rect(r1x, r1y, r2x, r2y, elemsize)

# view factors
vfmat = zeros(Float64, m.nelem, m.nelem)
@time existing_vf!(m, vfmat)
n = 15
dx, dy = create_tiles(m, n)
@time t_occ = check_tile_occupation(m, dx, dy, n)
@time blocking_vf_with_tiles!(m, vfmat, dx, dy, n, t_occ)
@time calculating_vf!(m, vfmat, normit = true)
vfmatp = vfmat_to_parts(m, vfmat, normit = true)
println("vf mat for parts:")
for i = 1:size(vfmatp,1)
    println(vfmatp[i,:])
end

# net radiation method
epsilon = zeros(m.nelem,1)
set_bc_part!(m, epsilon, 1, 0.3)
set_bc_part!(m, epsilon, 2, 0.6)
temp = zeros(m.nelem,1)
set_bc_part!(m, temp, 1:2, 300)
temp[79:106,1] .= 600 # bc for specific elements
@time Q, G = tempsolver(m, vfmat, temp, epsilon)
area = [m.elem[i].area for i = m.elem2par[1].first:m.elem2par[end].last]
q = Q[:] ./ area[:]

println("calculation finished")

# include("./test/run_example1.jl")


module RVM

using MLBase
using MLKernels

#__precompile__()

@inline logsig(x) = 1.0/(1.0+exp(-x))

export RVMspec, RVMfit, fit, predict

type RVMSpec
    kernel::MLKernels.Kernel
    n_iter::Integer
    normalizer::Function
    threshold::Real
    tol::Real
end

type RVMFit{R <: Real}
    kernel::MLKernels.Kernel
    normal::Function
    labelmap::MLBase.LabelMap
    converged::Bool
    steps::Integer
    w::Vector{R}
    RV::Matrix{R}
end

function _predicted_probability{R <: Real}(Model::RVMFit, Features::AbstractMatrix{R})
    X = Model.normal(Features)
    _ = MLKernels.kernelmatrix(Model.kernel, X, Model.RV)
    ϕ = [_ ones(size(_, 1))] # augment with bias

    P = _classify(Model.w, ϕ)
    P
end

function predict{R <: Real}(Model::RVMFit, Values::AbstractMatrix{R})
    P = _predicted_probability(Model, Values)
    _ = convert(Vector{Int64}, P .>= 0.5)
    Y = _ .+ 1 # booleans map to {0, 1} , {en,de}coders map to {1, 2}
    labeldecode(Model.labelmap, Y), P
end

function fit{R <: Real}(S::RVMSpec, Obs::AbstractMatrix{R}, ObsT::AbstractVector)
    # sample count, features
    n, p = size(Obs)

    # standardize normalizing process
    normalize = S.normalizer(Obs)
    X = normalize(Obs)

    # Relevance Vectors
    RV = X

    # standardize encoding process
    lm = labelmap(ObsT)
    t  = labelencode(lm, ObsT) - 1

    # design matrix / basis functions
    _ = MLKernels.kernelmatrix(S.kernel, X, X)
    ϕ = [_ ones(size(_, 1))] # augment with bias
    basis_count = size(ϕ, 2)

    # initialize priors N(0, 1)
    μ  = zeros(basis_count)
    α₀ = ones(basis_count) .* 1e0
    α₁ = α₀[1:end]
    #β  = 1e0

    # reporting information
    steps = S.n_iter
    succ  = false

    for i in 1:S.n_iter

        # update values
        μ, Σ = _solve_post_prob(μ, α₁, ϕ, t)
        γ  = 1 - (α₁ .* diag(Σ))
        α₁ = γ ./ (μ .^ 2)

        # identify informative vectors
        keep::Vector{Bool} = α₁ .< S.threshold
        if !any(keep)
            keep[1] = true # require at least one vecor
        end
        keep[end] = true # save bias

        # downselect uninformative vectors and weights
        α₀ = α₀[keep]
        α₁ = α₁[keep]
        γ  = γ[keep]
        ϕ  = ϕ[:, keep]
        Σ  = Σ[keep, keep]
        μ  = μ[keep]
        RV = RV[keep[1:end-1], :]

        # check for brreak
        Δα = maximum(abs(α₁ .- α₀))
        if (Δα < S.tol) & (steps > 2)
            steps = i
            succ  = true
            break
        end
        α₀ = α₁[1:end]
    end

    # Trained model as object
    modelfit::RVMFit{R} =
        RVMFit(S.kernel, normalize, lm, succ, steps, μ, RV)
end

function _classify(μ::AbstractVector, ϕ::AbstractMatrix)
    #@show "_classify"
    _ = ϕ * μ
    y = [logsig(x) for x in _]
    y
end

function _log_solve_post_prob(w::AbstractVector, α::AbstractVector,
                              ϕ::AbstractMatrix, t::AbstractVector)
    #@show "_log_post_probability"
    A = diagm(α)
    y = _classify(w, ϕ)
    pos = y[t .== 1]
    neg = y[t .!= 1]

    log_p = (0.5 * w' * A * w)[1] - sum(log(pos)) - sum(log(1 - neg))
end

function _hessian(w::AbstractVector, α::AbstractVector,
                  ϕ::AbstractMatrix, t::AbstractVector)
    #@show "_hessian"
    A = diagm(α)
    y = _classify(w, ϕ)
    B = diagm(y .* (1 - y))

    H = (ϕ' * B * ϕ + A)
    H
end

function _gradient(w::AbstractVector, α::AbstractVector,
                   ϕ::AbstractMatrix, t::AbstractVector)
    #@show "_gradient"
    A = diagm(α)
    y = _classify(w, ϕ)
    G = ϕ' * (t .- y) - (A * w)
    -G
end

function _newton_method(X₀::AbstractVector, F::Function, # Initial, Function
                        ∇::Function, ∇∇::Function)       # Gradient, Hessian
    #@show "_newton_method"
    for _ in 1:50
        #@show F(X₀) # uncomment to watch converge
        X₁ = X₀ - inv(∇∇(X₀)) * ∇(X₀)

        # convergence
        ΔX = maximum(abs(X₁ .- X₀))
        X₀ = X₁[1:end]
        if ΔX < 1e-6
            break
        end
    end
    X₀
end

function _solve_post_prob(w::AbstractVector, α::AbstractVector,
                          ϕ::AbstractMatrix, t::AbstractVector)
    w = _newton_method(w,
                       (w) -> _log_solve_post_prob(w, α, ϕ, t),
                       (w) -> _gradient(w, α, ϕ, t),
                       (w) -> _hessian(w, α, ϕ, t),
                       )

    Σ = inv(_hessian(w, α, ϕ, t))
    w, Σ
end

end


abstract type AbstractContinuation <: Function end

struct RationalContinuation{T}
    pmax::T
    kmax::Int
    b::T
end
function Base.getproperty(c::RationalContinuation, f::Symbol)
    f === :length && return length(c)
    return getfield(c, f)
end
function RationalContinuation(xmin, steps::Int)
    pmax = (1 - xmin) / xmin
    kmax = steps
    b = 1 / steps * (1 - sqrt(xmin)) / (1 + sqrt(xmin))
    return RationalContinuation(pmax, kmax, b)
end
function Base.iterate(s::RationalContinuation, k=1)
    (k > s.kmax + 1) && return nothing
    b = s.b
    return s(k), k + 1
end
Base.length(s::RationalContinuation) = s.kmax + 1
function (s::RationalContinuation{T})(k) where {T}
    k = clamp(k, 1, s.kmax + 1)
    b = s.b
    return 4 * b * (k - 1) / (1 - b * (k - 1))^2
end

struct FixedContinuation{T} <: AbstractContinuation
    param::T
    length::Int
end
function Base.iterate(s::FixedContinuation, x=1)
    (x > s.length) && return nothing
    return s.param, x + 1
end
Base.length(s::FixedContinuation) = s.length
(s::FixedContinuation{T})(x...) where {T} = s.param

"""
p(x) = 1/(a + b*ℯ^(-c*x))
"""
struct SigmoidContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
SigmoidContinuation(; kwargs...) = SigmoidContinuation{Float64}(; kwargs...)
function SigmoidContinuation{T}(;
    c::T=T(0.1), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = 1 - T(finish - start) / finish / (ℯ^(-c) - ℯ^(-steps * c)) * ℯ^(-c)
    b = T(finish - start) / finish / (ℯ^(-c) - ℯ^(-steps * c))
    return SigmoidContinuation(a, b, c, steps, min)
end
function Base.iterate(s::SigmoidContinuation, x=1)
    (x > s.length) && return nothing
    return max(1 / (s.a + s.b * ℯ^(-s.c * x)), s.min), x + 1
end
Base.length(s::SigmoidContinuation) = s.length
(s::SigmoidContinuation{T})(x) where {T} = max(1 / (s.a + s.b * ℯ^(-s.c * T(x))), s.min)

"""
p(x) = a(x-b*steps)^3 + c
"""
struct CubicSplineContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
CubicSplineContinuation(; kwargs...) = CubicSplineContinuation{Float64}(; kwargs...)
function CubicSplineContinuation{T}(;
    b::T=T(0.5), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = (finish - start) / (steps^3 * (1 - b)^3 - (1 - b * steps)^3)
    c = start - a * (1 - b * steps)^3
    return CubicSplineContinuation(a, b, c, steps, min)
end

function Base.iterate(s::CubicSplineContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * (x - s.b * s.length)^3 + s.c, s.min), x + 1
end
Base.length(s::CubicSplineContinuation) = s.length
function (s::CubicSplineContinuation{T})(x) where {T}
    return max(s.a * (T(x) - s.b * s.length)^3 + s.c, s.min)
end

"""
p(x) = a*x^b + c
"""
struct PowerContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
PowerContinuation(; kwargs...) = PowerContinuation{Float64}(; kwargs...)
function PowerContinuation{T}(;
    b=T(2), start=T(1), finish=T(5), steps::Int=30, min=-T(Inf)
) where {T}
    a = (finish - start) / max(T(1), steps^b - 1)
    c = start - a
    return PowerContinuation(a, b, c, steps, min)
end
function Base.iterate(s::PowerContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * x^s.b + s.c, s.min), x + 1
end
Base.length(s::PowerContinuation) = s.length
(s::PowerContinuation{T})(x) where {T} = max(s.a * T(x)^s.b + s.c, s.min)

"""
p(x) = a*ℯ^(b*x) + c
"""
struct ExponentialContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
ExponentialContinuation(; kwargs...) = ExponentialContinuation{Float64}(; kwargs...)
function ExponentialContinuation{T}(;
    b::T=T(0.1), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = (finish - start) / max(T(1), ℯ^(b * steps) - ℯ^(b))
    c = start - a * ℯ^(b)
    return ExponentialContinuation(a, b, c, steps, min)
end
function Base.iterate(s::ExponentialContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * ℯ^(s.b * x) + s.c, s.min), x + 1
end
Base.length(s::ExponentialContinuation) = s.length
(s::ExponentialContinuation{T})(x) where {T} = max(s.a * ℯ^(s.b * T(x)) + s.c, s.min)

"""
p(x) = a*log(b*x) + c
"""
struct LogarithmicContinuation{T} <: AbstractContinuation
    a::T
    b::T
    c::T
    length::Int
    min::T
end
LogarithmicContinuation(; kwargs...) = LogarithmicContinuation{Float64}(; kwargs...)
function LogarithmicContinuation{T}(;
    b::T=T(1), start::T=T(1), finish::T=T(5), steps::Int=30, min::T=-Inf
) where {T}
    a = (finish - start) / max(T(1), log(b * steps) - log(b))
    c = start - a * log(b)
    return LogarithmicContinuation(a, b, c, steps, min)
end
function Base.iterate(s::LogarithmicContinuation, x=1)
    x > s.length && return nothing
    return max(s.a * log(s.b * x) + s.c, s.min), x + 1
end
Base.length(s::LogarithmicContinuation) = s.length
(s::LogarithmicContinuation{T})(x) where {T} = max(s.a * log(s.b * T(x)) + s.c, s.min)

Continuation(::PowerPenalty; kwargs...) = Continuation(PowerPenalty; kwargs...)
Continuation(::RationalPenalty; kwargs...) = Continuation(RationalPenalty; kwargs...)
function Continuation(::Type{<:PowerPenalty}; steps, pmax=5.0, kwargs...)
    return PowerContinuation{Float64}(; b=1.0, start=1.0, steps=steps + 1, finish=pmax)
end
function Continuation(::Type{<:RationalPenalty}; steps, xmin=0.001, kwargs...)
    return RationalContinuation(xmin, steps)
end


using Random, Statistics

@testset "Assortativity" begin
    # Test definition of assortativity as Pearson correlation coefficient
    # between excess of degrees
    @testset "Small graphs" for n = 5:10
        @test @inferred assortativity(wheel_graph(n)) ≈ -1/3
    end
    @testset "Directed ($seed)" for seed in [1, 2, 3], (n, ne) in [(14, 18), (10, 22), (7, 16)]
        g = erdos_renyi(n, ne; is_directed=true, seed=seed)
        assort = assortativity(g)
        x = [outdegree(g, src(d))-1 for d in edges(g)]
        y = [indegree(g, dst(d))-1 for d in edges(g)]
        @test @inferred assort ≈ cor(x, y)
    end
    @testset "Undirected ($seed)" for seed in [1, 2, 3], (n, ne) in [(14, 18), (10, 22), (7, 16)]
        g = erdos_renyi(n, ne; is_directed=false, seed=seed)
        assort = assortativity(g)
        x = [outdegree(g, src(d))-1 for d in edges(g)]
        y = [indegree(g, dst(d))-1 for d in edges(g)]
        @test @inferred assort ≈ cor([x;y], [y;x])
    end
end


#== # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# Description
#
#   General functions to compute the gravity field potential and gravity
#   acceleration.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
# References
#
#   [1] Vallado, D. A (2013). Fundamentals of Astrodynamics and Applications.
#       Microcosm Press, Hawthorn, CA, USA.
#   [2] http://icgem.gfz-potsdam.de/home
#   [3] Delgado, M. R (2008). Aspherical Gravitational Field I: Modeling the
#       Space Environment. Universidad Poletécnica de Madrid, Madrid, Spain.
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # ==#

export create_gravity_model_coefs, compute_dU, compute_U, compute_g

################################################################################
#                               Public Functions
################################################################################

"""
    function create_gravity_model_coefs(icgem::ICGEM)

Return an instance of the structure `GravityModel_Coefs` based on the
information obtained from an ICGEM file in `icgem` (see `parse_icgem`).

"""
function create_gravity_model_coefs(icgem::ICGEM)

    legendre_norm = (icgem.norm == :fully_normalized) ? :full : :conv

    # Create and return the gravity model coefficients.
    gm_coefs::GravityModel_Coefs{Float64} =
        GravityModel_Coefs(name          = icgem.modelname,
                           μ             = icgem.gravity_constant,
                           R0            = icgem.radius,
                           n_max         = icgem.max_degree,
                           C             = copy(icgem.Clm),
                           S             = copy(icgem.Slm),
                           legendre_norm = legendre_norm)
end

"""
    function compute_dU(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector, n_max::Number = -1, m_max::Number = -1) where T<:Number

Compute the derivatives w.r.t. the spherical coordinates of the gravitational
field (`∂U/∂r`, `∂U/∂ϕ`, `∂U/∂λ`) defined by the coefficients `gm_coefs` (see
`GravityModel_Coefs`) at the position `r` [m] in ITRF. The maximum degree that
will be used while computing the spherical harmonics will be `n_max` and the
maximum order is `m_max`.

If `n_max` is negative, then the maximum available degree will be used. If
`n_max` is omitted, then it defaults to 0.

If `m_max` is negative or if it is greater than `n_max`, then it will be set to
`n_max`. If `m_max` is omitted, then it defaults to 0.

!!! info

    By convention, the result with `n_max` 0 and 1 will be the same.

# Returns

* The derivative of the gravitational field w.r.t. the radius (`∂U/∂r`).
* The derivative of the gravitational field w.r.t. the latitude (`∂U/∂ϕ`).
* The derivative of the gravitational field w.r.t. the longitude (`∂U/∂λ`).

# Remarks

In this case, `ϕ` is the geocentric latitude and `λ` is the geocentric
longitude.

"""
function compute_dU(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector,
                    n_max::Number = -1, m_max::Number = -1) where T<:Number

    # Unpack gravity model coefficients.
    #
    # We must not unpack all the variables because `n_max` will be overwritten.
    @unpack μ, R0, C, S, legendre_norm = gm_coefs

    # Check n_max value.
    if (n_max < 0) || (gm_coefs.n_max < n_max)
        n_max = gm_coefs.n_max
    end

    # Check m_max value.
    if (m_max < 0) || (m_max > n_max)
        m_max = n_max
    end

    # Get the geocentric latitude and longitude
    # =========================================
    r_gc = norm(r)
    ρ_gc = sqrt(r[1]^2 + r[2]^2)
    ϕ_gc = atan(r[3], ρ_gc)
    λ_gc = atan(r[2], r[1])

    # Auxiliary variables
    # ===================

    # Sine and cosine of the geocentric longitude.
    #
    # This values were be used in the algorithm to decrease the computational
    # burden.
    sin_λ,  cos_λ  = sincos(1λ_gc)
    sin_2λ, cos_2λ = sincos(2λ_gc)

    # First derivative of the non-spherical portion of the grav. field
    # ================================================================
    ∂Ur = T(1)  # Derivative w.r.t. the radius.
    ∂Uϕ = T(0)  # Derivative w.r.t. the geocentric latitude.
    ∂Uλ = T(0)  # Derivative w.r.t. the geocentric longitude.

    # Compute the associated Legendre functions `P_n,m[sin(ϕ_gc)]` with the
    # required normalization and its time derivative.
    #
    # Notice that cos(ϕ_gc-pi/2) = sin(ϕ_gc).
    dP, P = dlegendre(Val{legendre_norm}, ϕ_gc-pi/2, n_max, n_max, false)

    # Auxiliary variables.
    rn_fact = R0/r_gc
    rn      = rn_fact

    # Compute the derivatives.
    @inbounds for n = 2:n_max
        rn *= rn_fact # -> rn = (R0/r_gc)^n

        aux_∂Ur = T(0)
        aux_∂Uϕ = T(0)
        aux_∂Uλ = T(0)

        # Sine and cosine with m = 0.
        #
        # This values will be used to update recursively `sin(m*λ_gc)` and
        # `cos(m*λ_gc`), reducing the computational burden.
        sin_mλ   = T(0)      # sin( 0*λ_gc)
        sin_m_1λ = -sin_λ    # sin(-1*λ_gc)
        sin_m_2λ = -sin_2λ   # sin(-2*λ_gc)
        cos_mλ   = T(1)      # cos( 0*λ_gc)
        cos_m_1λ = cos_λ     # cos(-1*λ_gc)
        cos_m_2λ = cos_2λ    # cos(-2*λ_gc)

        for m = 0:n
            # Compute recursively `sin(m*λ_gc)` and `cos(m*λ_gc)`.
            sin_mλ = 2cos_λ*sin_m_1λ-sin_m_2λ
            cos_mλ = 2cos_λ*cos_m_1λ-cos_m_2λ

            CcSs_nm = C[n+1,m+1]*cos_mλ + S[n+1,m+1]*sin_mλ

            aux_∂Ur +=   P[n+1,m+1]*CcSs_nm
            aux_∂Uϕ +=  dP[n+1,m+1]*CcSs_nm
            aux_∂Uλ += m*P[n+1,m+1]*(S[n+1,m+1]*cos_mλ - C[n+1,m+1]*sin_mλ)

            # Check if we reached the maximum desired order.
            m >= m_max && break

            # Update the values for the next step.
            sin_m_2λ = sin_m_1λ
            sin_m_1λ = sin_mλ
            cos_m_2λ = cos_m_1λ
            cos_m_1λ = cos_mλ
        end

        ∂Ur += rn*(n+1)*aux_∂Ur
        ∂Uϕ += rn*aux_∂Uϕ
        ∂Uλ += rn*aux_∂Uλ
    end

    ∂Ur *= -μ/r_gc^2
    ∂Uϕ *= +μ/r_gc
    ∂Uλ *= +μ/r_gc

    ∂Ur, ∂Uϕ, ∂Uλ
end

"""
    function compute_g(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector, n_max::Number = -1, m_max::Number = -1) where T<:Number

Compute the gravitational acceleration (ITRF) \\[m/s²] at position `r` \\[m]
(ITRF) using the coefficients `gm_coefs` (see `GravityModel_Coefs`). The maximum
degree that will be used while computing the spherical harmonics will be
`n_max` and the maximum order it `m_max`.

If `n_max` is negative, then the maximum available degree will be used. If
`n_max` is omitted, then it defaults to 0.

If `m_max` is negative or if it is greater than `n_max`, then it will be set to
`n_max`. If `m_max` is omitted, then it defaults to 0.

!!! info

    By convention, the result with `n_max` 0 and 1 will be the same.

# Remarks

Notice that this function computes the **gravitational acceleration**. Hence,
the acceleration due to Earth rotation rate **is not** included.

"""
function compute_g(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector,
                   n_max::Number = -1, m_max::Number = -1) where T<:Number

    # Compute the partial derivatives of the gravitational field w.r.t. the
    # spherical coordinates.
    ∂Ur, ∂Uϕ, ∂Uλ = compute_dU(gm_coefs, r, n_max, m_max)

    # Auxiliary variables.
    r_gc = norm(r)
    ρ_gc = sqrt(r[1]^2 + r[2]^2)
    ϕ_gc = atan(r[3], ρ_gc)
    λ_gc = atan(r[2], r[1])

    # Compute the acceleration represented in the ITRF
    # ================================================

    # Compute the partial derivatives in spherical coordinate systems:
    #
    #     ∂U        1      ∂U     1   ∂U
    #    ---- , ---------.---- , ---.----
    #     ∂r     r.cos ϕ   ∂λ     r   ∂ϕ
    #
    # Notice that the singularity is not a problem here. When computing
    # `cos(pi/2)` a very small number will be returned and `∂U/∂λ` is 0. Hence,
    # the 2nd component will be 0.

    a_l = SVector{3,T}(∂Ur,
                       ∂Uλ/(r_gc*cos(ϕ_gc)),
                       ∂Uϕ/r_gc)

    # The vector `a_l` is represented in the local UEN (Up-Earth-North)
    # reference frame. Hence, we need to describe the unitary vectors of this
    # frame in the ECEF reference frame. This can be accomplished by the
    # following rotations matrix.

    Del = angle_to_dcm(0, ϕ_gc, -λ_gc, :XYZ)
    a_ITRF = Del*a_l

    a_ITRF
end

"""
    function compute_U(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector, n_max::Number = -1, m_max::Number = -1) where T<:Number

Compute the gravitational potential [J/kg] at `r` (ITRF) \\[m] using the
coefficients `gm_coefs` (see `GravityModel_Coefs`). The maximum degree that will
be used while computing the spherical harmonics will be `n_max` and the maximum
order is `m_max`.

If `n_max` is negative, then the maximum available degree will be used. If
`n_max` is omitted, then it defaults to 0.

If `m_max` is negative or if it is greater than `n_max`, then it will be set to
`n_max`. If `m_max` is omitted, then it defaults to 0.

!!! info

    By convention, the result with `n_max` 0 and 1 will be the same.

"""
function compute_U(gm_coefs::GravityModel_Coefs{T}, r::AbstractVector,
                   n_max::Number = -1, m_max::Number = -1) where T<:Number

    # Unpack gravity model coefficients.
    #
    # We must not unpack all the variables because `n_max` will be overwritten.
    @unpack μ, R0, C, S, legendre_norm = gm_coefs

    # Check n_max value.
    if (n_max < 0) || (gm_coefs.n_max < n_max)
        n_max = gm_coefs.n_max
    end

    # Check m_max value.
    if (m_max < 0) || (m_max > n_max)
        m_max = n_max
    end

    # Auxiliary variables.
    r_gc = norm(r)
    ρ_gc = sqrt(r[1]^2 + r[2]^2)
    ϕ_gc = atan(r[3], ρ_gc)
    λ_gc = atan(r[2], r[1])

    # Sine and cosine of the geocentric longitude.
    #
    # This values were be used in the algorithm to decrease the computational
    # burden.
    sin_λ,  cos_λ  = sincos(1λ_gc)
    sin_2λ, cos_2λ = sincos(2λ_gc)

    # Consider the zero degree term.
    U = T(1)

    # Compute the associated Legendre functions `P_n,m[sin(ϕ_gc)]` with the
    # required normalization.
    #
    # Notice that cos(ϕ_gc-pi/2) = sin(ϕ_gc).
    P = legendre(Val{legendre_norm}, ϕ_gc-pi/2, n_max, n_max, false)

    # Auxiliary variables.
    rn_fact = R0/r_gc
    rn      = rn_fact

    @inbounds for n = 2:n_max
        aux_U = T(0)

        # Sine and cosine with m = 0.
        #
        # This values will be used to update recursively `sin(m*λ_gc)` and
        # `cos(m*λ_gc`), reducing the computational burden.
        sin_mλ   = T(0)      # sin( 0*λ_gc)
        sin_m_1λ = -sin_λ    # sin(-1*λ_gc)
        sin_m_2λ = -sin_2λ   # sin(-2*λ_gc)
        cos_mλ   = T(1)      # cos( 0*λ_gc)
        cos_m_1λ = cos_λ     # cos(-1*λ_gc)
        cos_m_2λ = cos_2λ    # cos(-2*λ_gc)

        for m = 0:n
            # Compute recursively `sin(m*λ_gc)` and `cos(m*λ_gc)`.
            sin_mλ = 2cos_λ*sin_m_1λ-sin_m_2λ
            cos_mλ = 2cos_λ*cos_m_1λ-cos_m_2λ

            aux_U += P[n+1,m+1]*(C[n+1,m+1]*cos_mλ + S[n+1,m+1]*sin_mλ)

            # Check if we reached the maximum desired order.
            m >= m_max && break

            # Update the values for the next step.
            sin_m_2λ = sin_m_1λ
            sin_m_1λ = sin_mλ
            cos_m_2λ = cos_m_1λ
            cos_m_1λ = cos_mλ
        end

        rn *= rn_fact # -> rn = (R0/r_gc)^n
        U  += rn*aux_U
    end

    U *= μ/r_gc
end


using Base.Test
using Sigma
import Sigma: measure, logmeasure

# What should a domain be
b = HyperBox([0.0 0.0
              2.0 1.0])
@test ndims(b) == 2
@test measure(b) == 2
@test logmeasure(b) == log(2) + 0


using NumberIntervals, Turmeric, StaticArrays, ForwardDiff
using Test

all_unique(rts) = length(rts[2]) == 0

function test_newtonlike(f, X, method, nsol)
    rts = roots(f, X, method)
    @test length(rts[1]) == nsol
    @test all_unique(rts)
end

newtonlike_methods = (Newton(), Krawczyk())

@testset "1D roots" begin
    # Default
    rts = roots(sin, NumberInterval(-5, 5))
    @test length(rts[1]) == 3
    @test all_unique(rts)

    # Bisection()
    rts = roots(sin, NumberInterval(-5, 6), Bisection(), 1e-3)
    @test length(rts[2]) == 3

    for method in newtonlike_methods
        test_newtonlike(sin, NumberInterval(-5, 5), method, 3)
    end

    # Infinite interval
    rts = roots(x -> x^2 - 2, NumberInterval(-Inf, Inf), Newton())
    @test length(rts[1]) == 2
end

@testset "2D roots" begin
    f(x, y) = SVector(x^2 + y^2 - 1, y - 2x)
    f(X) = f(X...)
    X = SVector(
        NumberInterval(-6, 6),
        NumberInterval(-6, 6)
    )

    # Bisection()
    rts = roots(f, X, Bisection(), 1e-3)
    @test_broken length(rts[2]) == 2

    for method in newtonlike_methods
        test_newtonlike(f, X, method, 2)
    end

    # Infinite interval
    X = SVector(NumberInterval(-Inf, Inf), NumberInterval(-Inf, Inf))
    rts = roots(f, X, Newton())
    @test length(rts[1]) == 2
end

# From R docs:
# https://www.rdocumentation.org/packages/pracma/versions/1.9.9/topics/broyden

@testset "3D roots" begin
    function g(x)
        (x1, x2, x3) = x
        SVector(    x1^2 + x2^2 + x3^2 - 1,
                    x1^2 + x3^2 - 0.25,
                    x1^2 + x2^2 - 4x3
                )
    end

    X = NumberInterval(-5, 5)
    XX = SVector(X, X, X)

    for method in newtonlike_methods
        rts = roots(g, XX, method)
        @test length(rts[1]) == 4
        @test all_unique(rts)
    end
end

@testset "Out of domain" begin
    for method in newtonlike_methods
        @test_skip length(roots(log, NumberInterval(-100, 2), method)[1]) == 1
        @test length(roots(log, NumberInterval(-100, -1), method)[1]) == 0
    end
end

@testset "Infinite domain" begin
    for method in newtonlike_methods
        rts = roots(x -> x^2 - 2, NumberInterval(-Inf, Inf), method)
        @test length(rts[1]) == 2
        @test all_unique(rts)
    end
end

@testset "NaN return value" begin
    f(xx) = ( (x, y) = xx; SVector(log(y/x) + 3x, y - 2x) )
    X = SVector(NumberInterval(-100.0, 0.0), NumberInterval(-100.0, 0.0))
    for method in newtonlike_methods
        rts = roots(f, X, method)
        @test length(rts[1]) == 1
        @test_broken length(rts[2]) == 0 # failure due to singularity
    end
end

@testset "Stationary points" begin
    f(xx) = ( (x, y) = xx; sin(x) * sin(y) )
    gradf(x) = ForwardDiff.gradient(f, x)
    XX = SVector(NumberInterval(-5, 6), NumberInterval(-5, 6))
    tol = 1e-5

    for method in newtonlike_methods
        test_newtonlike(gradf, XX, method, 25)
    end
end

@testset "Full contraction of roots" begin
    region = NumberInterval(-5, 5)
    tol = 1e-7
    for method in newtonlike_methods
        rts = roots(sin, region, method, tol)
        contractor = GradientContractor(sin, method, region)
        contracted_regions = tfullcontract(contractor, rts[1], tol)
        for _region in contracted_regions
            @test diam(_region) < tol
        end
    end
end