###### MLE ######

"""
Shrinkage of small eigen values.
"""
function eigen_truncation(S::AbstractVector{<:Real}, ρ::Real; ptrunc::Real=2., kwargs...)
    @assert 0<=ρ<=1  "Relative threshold must be in [0,1]"
    @assert Inf > ptrunc > 0

    Sp = abs.(S).^ptrunc
    sidx = sortperm(Sp, lt=≥)  # index in decreasing order
    Sn = (cumsum(view(Sp, sidx))/sum(Sp)).^(1/ptrunc)  # decreasing order
    k = findfirst(Sn .>= 1-ρ)
    return isnothing(k) ? (1:length(S)) : sidx[1:k]
end

"""
    xiAx(A, X; ρ=0)

Safe evaluation of the inverse quadratic form
    trace(X' * inv(A) * X)
where the matrix `A` is symmetric and positive definite. Small eigen values of `A` are truncated.
"""
function xiAx(A::AbstractMatrix{<:Real}, X::AbstractVecOrMat{<:Real}; ρ::Real=0, kwargs...)
    # # Sanity check
    # @assert issymmetric(A)
    # @assert size(X, 1) == size(A, 1)

    # # Direct inversion
    # return tr(X' * pinv(A + 1e-3*I) * X)

    # Eigen decomposition, S are the eigen values in increasing order
    S, U = eigen(A)  # so that U * Diagonal(S) * inv(U) == A, in particular, U' == inv(U)
    # On covariance matrix this is equivalent to
    # U, S, V = svd(A)

    # shrinkage of small eigen values for stability
    idx = eigen_truncation(S, ρ; kwargs...)

    return sum((U[:,idx]'*X).^2 ./ S[idx])  # this may raise error if idx is empty
end

"""
    log_likelihood(A, X)

Log-likelihood of a general Gaussian vector.

# Args
- A: covariance matrix
- X: sample vector or matrix. For matrix each column is a sample.
"""
function log_likelihood(A::AbstractMatrix{<:Real}, X::AbstractMatrix{<:Real}; kwargs...)
    # # Sanity check
    # @assert issymmetric(A)
    # @assert size(X, 1) == size(A, 1)

    N = size(X,2) # number of i.i.d. samples in data
    return -1/2 * (xiAx(A,X, kwargs...) + N*logdet(A) + length(X)*log(2π))
end

log_likelihood(A::AbstractMatrix, X::AbstractVector; kwargs...) = log_likelihood(A, reshape(X,:,1), kwargs...)

"""
    log_likelihood_H(A, X)

Safe evaluation of the log-likelihood of a fBm model with the implicite σ (optimal in the MLE sense).

# Args
- A: covariance matrix
- X: sample vector or matrix. For matrix each column is a sample.

# Notes
- This function is common to all MLEs with the covariance matrix of form `σ²A(h)`, where `{σ, h}` are unknown parameters. This kind of MLE can be carried out in `h` uniquely and `σ` is obtained from `h`.
- `log_likelihood_H(Σ, X)` and `log_likelihood(σ^2*Σ, X)` are equivalent with the optimal `σ = sqrt(xiAx(Σ, X) / length(X))`
"""
function log_likelihood_H(A::AbstractMatrix{<:Real}, X::AbstractMatrix{<:Real}; ρ::Real=0, kwargs...)
    # Sanity check
    # @assert issymmetric(A)
    # @assert size(X, 1) == size(A, 1)

    N = size(X,2) # number of i.i.d. samples in data

    # Compute the log-likelihood: log_likelihood(A,X)
    S, U = eigen(A)  # so that U * Diagonal(S) * inv(U) == A, in particular, U' == inv(U)

    # shrinkage of small eigen values for stability
    idx = eigen_truncation(S, ρ; kwargs...)

    val = -1/2 * (length(X)*log(sum((U[:,idx]'*X).^2 ./ S[idx])) + N*sum(log.(S[idx])))  # non-constant part of log-likelihood

    return val - length(X) * log(2π*exp(1)/length(X))/2  # with the constant part
end

# function log_likelihood_H_partial(A::AbstractMatrix{<:Real}, X::AbstractMatrix{<:Real}; kwargs...)
#     log_likelihood_H_full(diagm(0=>diag(A)), X; kwargs...)
# end

# function log_likelihood_H(A::AbstractMatrix, X::AbstractMatrix; partial::Bool=false, kwargs...)
#     return if partial
#         log_likelihood_H_partial(A, X; kwargs...)
#     else
#         log_likelihood_H_full(A, X; kwargs...)
#     end
# end

log_likelihood_H(A::AbstractMatrix, X::AbstractVector; kwargs...)  = log_likelihood_H(A, reshape(X,:,1); kwargs...)

#### fWn (bank)-MLE ####

"""
    fWn_log_likelihood_H(X, H, F, G)

Log-likelihood of a fWn bank model with the optimal volatility.
"""
function fWn_log_likelihood_H(X::AbstractVecOrMat{<:Real}, H::Real, F::AbstractVector{<:AbstractVector{<:Real}}, G::AbstractVector{<:Integer}; mode::Symbol=:causal, kwargs...)
    # # Sanity check
    # @assert size(X,1) == length(G) * length(F)  "Mismatched dimensions."

    # the process and the number of filters in the filter bank
    proc = FractionalWaveletNoiseBank(H, F, mode)  # mode of wavelet transform

    return log_likelihood_H(covmat(proc, G; kwargs...), X; kwargs...)
end

"""
Efficient evaluation of `fWn_log_likelihood_H`. This function reuses the precomputed summation coefficients `A` and their supports `K` that are independent of the Hurst exponent `H` hence avoids the re-evaluation. Note that `A` and `K` depend on the mode of the wavelet transform.
"""
function fWn_log_likelihood_H(X::AbstractVecOrMat{<:Real}, H::Real, A::AbstractMatrix{<:AbstractVector{<:Real}}, K::AbstractMatrix{<:AbstractVector{<:Real}}, G::AbstractVector{<:Integer}; kwargs...)
    # # Sanity check
    # @assert size(X,1) == length(G) * length(F)  "Mismatched dimensions."

    # the process and the number of filters in the filter bank
    return log_likelihood_H(fWn_covmat(A, K, H, G; kwargs...), X; kwargs...)
end

"""
    fWn_MLE_estim(X, F, G; method)

Maximum likelihood estimation of Hurst exponent and volatility for fWn bank.

# Args
- X: sample matrix.
- F: array of filters used for computing `X`
- G: integer time grid of `X`.
- method: :optim for optimization based or :table for look-up table based solution.

# Returns
- hurst, σ: estimation
- loglikelihood: log-likelihood of estimation
- optimizer: object of optimizer, for method==:optim only

# Notes
- The fWnb process is multivariate. An observation at time `t` is a vector that the dimension equals to the number of filters used in fWnb. A sample vector is the concatenation of observations made on some time grid `G` and the row dimension of `X` must be `length(F) * length(G)`. Columns of `X` are treated as i.i.d. sample vectors hence the horizontal axis (i.e. last dimension) should NOT be interpreted as time axis (even it is physically the case, e.g. wavelet coefficients of a time series). In fact the time axis is "in" the vertical axis (i.e. the first dimension) that the grid `G` must be conformal with.
"""
function fWn_MLE_estim(X::AbstractMatrix{<:Real}, F::AbstractVector{<:AbstractVector{<:Real}}, G::AbstractVector{<:Integer}; method::Symbol=:optim, mode::Symbol=:causal, kwargs...)
    if length(X) == 0 || any(isnan.(X))  # for empty input or input containing nans
        return (hurst=NaN, σ=NaN, loglikelihood=NaN, optimizer=nothing)
    else
        # Sanity check
        @assert size(X,1) == length(F) * length(G)  "Mismatched dimensions."
        @assert length(G) == length(unique(G))  "All elements of the grid must be distinct."

        # X = X .- mean(X, dims=2)  # force the zero-mean condition. Attention: this can cause the estimation of H to fail!

        proc = FractionalWaveletNoiseBank(0.5, F, mode)
        # func = h -> -fWn_log_likelihood_H(X, h, F, G; kwargs...)
        func = h -> -fWn_log_likelihood_H(X, h, proc.coeffs, proc.supps, G; kwargs...)

        opm = nothing
        hurst = nothing
        ε::Real=1e-2  # search Hurst in the range [ɛ, 1-ɛ]

        if method == :optim
            # Gradient-free constrained optimization
            opm = Optim.optimize(func, ɛ, 1-ɛ, Optim.Brent())
            # # Gradient-based optimization
            # optimizer = Optim.GradientDescent()  # e.g. Optim.BFGS(), Optim.GradientDescent()
            # opm = Optim.optimize(func, ε, 1-ε, [0.5], Optim.Fminbox(optimizer))
            hurst = Optim.minimizer(opm)[1]
        elseif method == :table
            Hs = collect(ε:ε:1-ε)
            hurst = Hs[argmin([func(h) for h in Hs])]
        else
            throw("Unknown method: ", method)
        end

        # proc = FractionalWaveletNoiseBank(hurst, F, mode)
        # Σ = covmat(proc, G)
        Σ = fWn_covmat(proc.coeffs, proc.supps, hurst, G; kwargs...)

        # Estimation of volatility
        σ = sqrt(xiAx(Σ, X) / length(X))

        # Log-likelihood
        L = log_likelihood_H(Σ, X)
        # # or equivalently
        # L = log_likelihood(σ^2*Σ, X)
        L /= length(X)  # normalization

        return (hurst=hurst, σ=σ, loglikelihood=L, optimizer=opm)
    end
end

"""
Accelerated MLE by dividing a large vector of samples into smaller ones.

# Args
- X, F: wavelet coefficient matrix and filters
- s: sub window size
- u: downsampling factor
- l: length of decorrelation
"""
function fWn_MLE_estim(X::AbstractMatrix{<:Real}, F::AbstractVector{<:AbstractVector{<:Real}}, s::Integer, u::Integer, l::Integer; kwargs...)
    t, V = rolling_vectorize(X, s, u, l; mode=:causal)  # TODO: case s = size(X,2)
    fWn_MLE_estim(V, F, u*(1:s); kwargs...)  # regular grid is implicitely used here.
end

#### B-Spline MLE ####

"""
Filter of a fWn.

# Note
- The extra 1/sqrt(s) factor is due to the definition of B-Spline wavelet filter and DCWT.
"""
bspline_fWn_filter = (s,v) -> intscale_bspline_filter(s,v) / sqrt(s)

"""
    bspline_MLE_estim(X::AbstractMatrix{<:Real}, sclrng::AbstractVector{<:Integer}, vm::Integer, G::AbstractVector{<:Integer}; kwargs...)

Maximum likelihood estimation of Hurst exponent and volatility for fractional Wavelet noise bank.

# Args
- X: matrix of B-Spline wavelet coefficients. Each row corresponds to a scale.
- sclrng: scales of the rows in `X`.
- vm: vanishing moment of the B-Spline wavelet.
- G: integer grid correpsonding to the columns of `X`.

# Notes
- In rolling window estimation, to avoid recomputing the B-Spline filters one can use directly `fWn_MLE_estim`.
"""
function bspline_MLE_estim(X::AbstractMatrix{<:Real}, sclrng::AbstractVector{<:Integer}, vm::Integer, G::AbstractVector{<:Integer}; kwargs...)
    # @assert size(X,1) == length(sclrng) * length(G)  "Mismatched dimensions."

    F = [bspline_fWn_filter(s, vm) for s in sclrng]
    fWn_MLE_estim(X, F, G; kwargs...)
end

"""
    bspline_MLE_estim(X::AbstractMatrix{<:Real}, sclrng::AbstractVector{<:Integer}, vm::Integer, s::Integer, u::Integer, l::Integer; kwargs...)

Cross-scale and cross-time B-Spline MLE, by concatenating adjacent columns into a larger vector and using both the cross-scale and the cross-time correlations.
"""
function bspline_MLE_estim(X::AbstractMatrix{<:Real}, sclrng::AbstractVector{<:Integer}, vm::Integer, s::Integer, u::Integer, l::Integer; kwargs...)
    t, V = rolling_vectorize(X, s, u, l; mode=:causal)
    bspline_MLE_estim(V, sclrng, vm, u*(1:s); kwargs...)
end

"""
    bspline_MLE_estim(X::AbstractMatrix{<:Real}, sclrng::AbstractVector{<:Integer}, vm::Integer; kwargs...)

Cross-scale only B-Spline MLE, by treating the columns of `X` as independent observations and using only the cross-scale correleation.

# Notes
- Setting the keyword argument `partial=true` to use the local Whittle wavelet estimator.
"""
function bspline_MLE_estim(X::AbstractMatrix{<:Real}, sclrng::AbstractVector{<:Integer}, vm::Integer; kwargs...)
    bspline_MLE_estim(X, sclrng, vm, 1:1; kwargs...)
end

"""
    bspline_MLE_estim(X::AbstractVector{<:Real}, scl::Integer, vm::Integer, args...;  kwargs...)

Cross-time only B-Spline MLE (with acceleration).
"""
function bspline_MLE_estim(X::AbstractVector{<:Real}, scl::Integer, vm::Integer, args...; kwargs...)
    # Call two different functions depending on whether the arguments `s, u, l` are passed or not in `args...`.
    if isempty(args)
        # reshape to a column vector with a regular grid
        bspline_MLE_estim(reshape(X,:,1), [scl], vm, 1:length(X); kwargs...)
    else
        bspline_MLE_estim(reshape(X,1,:), [scl], vm, args...; kwargs...)
    end
end

#### fGn-MLE ####

"""
Filter of a fGn.
"""
fGn_filter = d -> vcat(1, zeros(d-1), -1)

"""
Maximum likelihood estimation of Hurst exponent and volatility for fractional Gaussian noise bank.

# Args
- X: sample matrix.
- lags: time lags of the finite difference operator used for computing `X`.
- G: integer time grid of `X`.
"""
function fGn_MLE_estim(X::AbstractMatrix{<:Real}, lags::AbstractVector{<:Integer}, G::AbstractVector{<:Integer}; kwargs...)
    # @assert size(X,1) == length(lags) * length(G)  "Mismatched dimensions."

    F = [fGn_filter(l) for l in lags]
    fWn_MLE_estim(X, F, G; kwargs...)
end

"""
Cross-scale and cross-time fGn MLE.
"""
function fGn_MLE_estim(X::AbstractMatrix{<:Real}, lags::AbstractVector{<:Integer}, s::Integer, u::Integer, l::Integer; kwargs...)
    t, V = rolling_vectorize(X, s, u, l; mode=:causal)
    fGn_MLE_estim(V, lags, u*(1:s); kwargs...)
end

"""
Cross-scale only fGn MLE.
"""

function fGn_MLE_estim(X::AbstractMatrix{<:Real}, lags::AbstractVector{<:Integer}; kwargs...)
    fGn_MLE_estim(X, lags, 1:1; kwargs...)
end

"""
Cross-time only fGn MLE (with acceleration).
"""
function fGn_MLE_estim(X::AbstractVector{<:Real}, lag::Integer, args...; kwargs...)
    # Call two different functions depending on whether the arguments `s, u, l` are passed or not in `args...`.
    if isempty(args)
        # reshape to a column vector with a regular grid
        fGn_MLE_estim(reshape(X,:,1), [lag], 1:length(X); kwargs...)
    else
        fGn_MLE_estim(reshape(X,1,:), [lag], args...; kwargs...)
    end
end


##  Fast Fourier transform of a real-valued signal
##  + simple power spectrum
##  + Hilbert transform
##  Author Bernd Blasius (bernd.blasius@gmail.com)

# this files will need the Julia Package FFTW.jl

function power(x,fs)
    # simple power spectrum
    # x: time series
    # f: sampling frequency
    n = length(x)
    xbar = fftshift(fft(x))
    f = (-n/2:n/2-1)*(fs/n)    # zero-centered frequency range
    power= abs.(xbar).^2/n     # zero-centered power

    f, power
end

function power_real(x,fs)
    # simple power spectrum using rfft
    # x: time series
    # f: sampling frequency
    n = length(x)
    xbar = rfft(x)
    n1 = length(xbar)
    f = (0:n1-1) * (fs/n)     #frequency range
    power = abs.(xbar).^2/n    #power

    f, power
end

function hilbert(x,fs)
    # calculate the Hilbert transform of a real valued signal
    n = length(x)
    xbar = rfft(x)
    irfft(im*xbar,n) # Hilbert transform
end


k, a = quadratic_field(10)
@test degree(k) == 2
k, a = cyclotomic_field(11)
G, mG = automorphism_group(k)
@test order(G) == 10
kt, t = k["t"]
factor(t^2-a)
factor(t^5-a)
k, a = wildanger_field(3, 13)
h = hilbert_class_field(k)
@test degree(h) == 9
K = number_field(h)
@test degree(h) == 9
L = simple_extension(K)[1]
@test degree(L) == 9
absolute_simple_field(L)
discriminant(maximal_order(K))
norm_equation(k, 27)


"""
DatingEstimate is a simple datatype to store the results of the coaltime function.
coaltime estimates date range that may be considered a 95% confidence range in
which the true coalsescence time between two sequences lies.
The type stores the 5%, 50%, and 95% values for the range.
"""
struct DatingEstimate
    lower::Float64
    middle::Float64
    upper::Float64
end

DatingEstimate() = DatingEstimate(.0, .0, .0)

lower(x::DatingEstimate) = x.lower
middle(x::DatingEstimate) = x.middle
upper(x::DatingEstimate) = x.upper

lower(x::Vector{DatingEstimate}) = lower.(x)
middle(x::Vector{DatingEstimate}) = middle.(x)
upper(x::Vector{DatingEstimate}) = upper.(x)

function Base.show(io::IO, de::DatingEstimate)
    println(io, "Coalescence time estimate:\n5%: $(de.lower), 95%: $(de.upper)")
end

@inline function binomzero(p0::Float64, N::Int, B::Int)
    f(p::Float64) = cdf(Binomial(N, p), B) - p0
    return fzero(f, 0, 1)
end

"""
    estimate_time(N::Int, K::Int, µ::Float64)

Compute the coalescence time between two sequences by modelling the process of
mutation accumulation between two sequences as a bernoulli process.

`N` is the length of the two aligned sequences.
`K` is the estimated number of mutations.
`µ` is the assumed mutation rate.
"""
@inline function estimate_time(N::Int, K::Int, µ::Float64)
    if !(N ≥ K ≥ 0)
        throw(ArgumentError("Condition must be satisfied: N ≥ K ≥ 0"))
    end
    div = 2µ
    ninetyfive = ceil(binomzero(0.05, N, K) / div)
    fifty = ceil(binomzero(0.5, N, K) / div)
    five = ceil(binomzero(0.95, N, K) / div)
    return DatingEstimate(five, fifty, ninetyfive)
end

"""
    estimate_time(N::Int, K::Int, µ::Float64)

Compute the coalescence time between two sequences by modelling the process of
mutation accumulation between two sequences as a bernoulli process.

`N` is the length of the two aligned sequences.
`d` is the a value of estimated genetic distance.
`µ` is the assumed mutation rate.
"""
@inline function estimate_time(N::Int, d::Float64, µ::Float64)
    if !(1.0 ≥ d ≥ 0.0)
        throw(ArgumentError("Genetic distance `d` must be a value between 1 and 0"))
    end
    K = convert(Int, ceil(d * N))
    return estimate_time(N, K, µ)
end

@inline function estimate_time(m::MutationCount, μ::Float64)
    return estimate_time(m.n_sites, m.n_mutations, μ)
end

@inline estimate_time(V::Vector{MutationCount}, μ::Float64) = estimate_time.(V, μ)
@inline estimate_time(V::Vector{Vector{MutationCount}}, μ::Float64) = estimate_time.(V, μ)

@inline function estimate_time(M::PairwiseListMatrix{T,false}, μ::Float64) where {T}
    V = estimate_time(getlist(M), μ)
    return PairwiseListMatrix(V, false, eltype(V)())
end

@inline function estimate_time(M::N, μ::Float64) where {N<:NamedArray}
    V = estimate_time(M.array, μ)
    return setlabels(V, names(M, 1))
end


function WrappedGaussian(x, μ, σ)
    p = 0
    K = 10
    for k = -K:K
        p += exp(-0.5 * ((x - μ - 2pi * k) / σ)^2)
    end
    return p / (sqrt(2π) * σ)
end

function ω2k(ω, h)
    g = 9.81
    α = ω^2*(h/g)
    out = α*(tanh(α))^(-0.5)/h
    return out .* !isnan(out)
end


cd(@__DIR__); include("setups/grid7x3.jl");

## ROT1 + pmf2
D = natural_eigdist(𝚽, 𝛌, Q; α = 0.1, input_format = :pmf2, distance = :ROT)
E = transform(fit(MDS, D, maxoutdim=2, distances=true))
E[1, :] .*= sign(E[1, 2]); E[2, :] .*= -sign(E[1, 2])
plt = grid7x3_mds_heatmaps(E, 𝚽; backend = :pyplot)
savefig(plt, "../figs/Grid7x3_MDS_ROT1_pmf2_alpha01.png")
display(plt)


function IMPES!(sim::Simulation, res::Reservoir, kr::RelPerms, fluid::Fluids, grid::Grid,
                bc::BoundaryConditions, post::FlowResults)

    # Define flow functions
    # krw(s) = water_rel_perm(s, res.kr)
    # kro(s) = oil_rel_perm(s, res.kr)

    # λw(s) = krw(s) / fluid.μw
    # λo(s) = kro(s) / fluid.μo
    # λ(s) = λo(s) + λw(s)

    # fw(s) = λw(s) ./ (λw(s) .+ λo(s))

    # Get data
    n = grid.Nx
    A = res.A
    k = res.k
    ϕ = res.phi
    Δx = grid.dx
    Δt = sim.dt
    s = sim.s
    θ = sim.θ

    # Initialize velocity vectors
    v = zeros(n+1)
    v[1] = bc.value[1] / A

    ∇v = zeros(n)

    # Calculate all mobilities
    krw, kro = kr(s, θ)
    Λw = krw ./ fluid.μw
    Λo = kro ./ fluid.μo
    Λ = Λw + Λo

    # Create T Vector
    T = zeros(n,n)
    for i in 2:n-1
        T[i,i-1] = Λ[i-1]
        T[i,i+1] = Λ[i]
        T[i,i] = - (T[i, i-1] + T[i,i+1])
    end

    # Left - Neumann
    T[1,1] = -Λ[1]
    T[1,2] = Λ[1]
    # Right - Dirichlet
    T[n,n] = - 9*Λ[n] - 3*Λ[n-1]
    T[n,n-1] = 3*Λ[n-1] + Λ[n]

    ## Get Q vector
    Q = zeros(n)
    # Left - Neumann
    Q[1] = - bc.value[1] / A / k * Δx
    # Right - Dirichlet
    Q[n] = -8 * (Λw[n] + Λo[n]) * bc.value[2]

    ## Solve for Pressure
    p = T \ Q

    ## Solve for Saturation
    # Calculate water velocity
    v[2:end-1] = [Λw[l-1] * (p[l-1] - p[l]) for l in 2:n] / Δx * k
    v[end] = 2 * Λw[n]*(p[n]-bc.value[2]) / Δx * k
    # Calculate gradient of water velocity
    ∇v = [v[l-1] - v[l] for l in 2:n+1] / Δx

    s = s .+ ∇v / ϕ * Δt

    # Update Simulation
    sim.p = p
    sim.s = s
    sim.v = v
    sim.i += 1
    sim.t += sim.dt

    # Save Results
    post.p[:,sim.i] = p
    post.s[:,sim.i] = s
    post.θ[:,sim.i] = θ

end

function IMPEC!(ctrans::CompTransport, sim, res, grid, post)
    n = grid.Nx
    Δx = grid.dx
    ϕ = res.phi
    Δt = sim.dt
    v = sim.v
    s = sim.s
    i = sim.i
    Δs = sim.s .- post.s[:,i-1]

    ncomps = ctrans.ncomps

    for j in 1:ncomps
        c = ctrans.c[:,j,i-1]

        cinj = ctrans.cinj[j]
        ∇cuw  = zeros(n)
        ∇cuw[2:n] = [c[l] * v[l+1] - c[l-1] * v[l]  for l in 2:n] / Δx
        ∇cuw[1] = (c[1] * v[2] - cinj * v[1])  / Δx

        c = c  .- Δt * ∇cuw / ϕ ./ s - c .* Δs ./ s

        ctrans.c[:,j,i] = c
    end

end

# function IMPEC!(ctrans::CompTransport, sim::Simulation,  res::Reservoir, fluid::Fluids, grid::Grid, bc::BoundaryConditions, wells::Wells)
#     Nx = grid.Nx
#     Ny = grid.Ny
#     p = sim.p
#     s = sim.s
#     ϕ = res.phi
#     k = res.k
#     Δx = grid.dx
#     Δt = sim.dt
#     μw = fluid.μw
#     Vi = res.h * res.W * Δx
#     qinj = wells.Q[1]
#     kr = res.kr
#     c = ctrans.c
#     cinj = ctrans.cinj
#     sold = ctrans.s_old

#     η = k * Δt / (ϕ * μw * Δx^2) * 6894757 * 1.06e-14 * 3600 * 24
#     A = zeros(Nx * Ny, Nx * Ny)
#     krw = [water_rel_perm(s[l], kr) for l in 1:Nx]

#     # Upwinding
#     for l in 2:Nx-1
#         # krw1 = water_rel_perm(s[l], kr)
#         # krw0 = water_rel_perm(s[l-1], kr)

#         A[l, l-1] =  krw[l-1] * (p[l-1] - p[l]) * η / s[l]
#         A[l, l] = krw[l] * (p[l+1] - p[l])  * η / s[l]
#     end

#     A[Nx, Nx-1] =  krw[Nx-1] * (p[Nx-1] - p[Nx]) * η / s[Nx]
#     A[Nx, Nx] =  - krw[Nx] * (2* p[Nx]) * η  / s[Nx]

#     A[1, 1] =  krw[1] * (p[2] - p[1]) * η  / s[1]

#     Q = zeros(Nx * Ny)
#     Q[1] = qinj * Δt / (s[1] * Vi * ϕ) * cinj

#     Δs = (s .- sold) ./ sold

#     ctrans.c = c + A * c + Q - Δs .* c
#     ctrans.s_old = sim.s
# end

# function IMPEC2!(ctrans::CompTransport, sim::Simulation,  res::Reservoir, fluid::Fluids, grid::Grid, bc::BoundaryConditions, wells::Wells)
#     Nx = grid.Nx
#     Ny = grid.Ny
#     p = sim.p
#     s = sim.s
#     ϕ = res.phi
#     k = res.k
#     Δx = grid.dx
#     Δt = sim.dt
#     μw = fluid.μw
#     Vi = res.h * res.W * Δx
#     qinj = wells.Q[1]
#     kr = res.kr
#     c = ctrans.c
#     cinj = ctrans.cinj
#     sold = ctrans.s_old

#     η = k * Δt / (ϕ * μw * Δx^2) * 6894757 * 1.06e-14 * 3600 * 24
#     A = zeros(Nx * Ny, Nx * Ny)

#     krw = [water_rel_perm(s[l], kr) for l in 1:Nx]

#     u = zeros(Nx)
#     for l in 2:Nx-1
#         u[l] = krw[l] * (p[l] - p[l+1])
#     end
#     u[1] = krw[1] * (p[2] - p[1])
#     u[Nx] = krw[Nx] * (p[Nx] - p[Nx-1])

#     # Upwinding
#     for l in 3:Nx - 1
#         A[l, l] = -3 * u[l] * η / 2 / s[l]
#         A[l, l-1] =  4 * u[l-1] * η / 2 / s[l]
#         A[l, l-2] =  -u[l-2] * η / 2 / s[l]
#     end

#     A[2, 1] = krw[1] * (p[1] - p[2]) * η / s[2]
#     A[2, 2] = krw[2] * (p[3] - p[2])  * η / s[2]

#     A[1, 1] =  krw[1] * (p[2] - p[1]) * η  / s[1]

#     A[Nx, Nx-1] =  krw[Nx-1] * (p[Nx-1] - p[Nx]) * η / s[Nx]
#     A[Nx, Nx] =  - krw[Nx] * (2* p[Nx]) * η  / s[Nx]

#     Q = zeros(Nx * Ny)
#     Q[1] = qinj * Δt / (s[1] * Vi * ϕ) * cinj

#     Δs = (s .- sold) ./ sold

#     ctrans.s_old = sim.s
#     ctrans.c = c + A * c + Q - Δs .* c
# end


using InterpretableModels: zerooneconvert, optimize

@testset "LIMScoringSystem via Lasso" begin

N, P = 1_000, 500
realcoefs = zeros(Int, P)
realcoefs[1] = 4
realcoefs[42] = -2
realcoefs[244] = -8
X = rand(N, P) .< 0.5
yraw = X * realcoefs
y = zerooneconvert.(yraw) # class is 0 or 1
models = optimize(LIMScoringSystem, LassoLIMOptimizer(), X, y)

@test models isa AbstractVector
@test models[1] isa LIMScoringSystem

end


# Bayesian Logistic Regression using the Polya-Gamma augmentation.
#

using LinearAlgebra

export BinaryLogisticRegression
export LogisticRegression
export predict

#######################################################################
# Binary logistic regression
#
# NOTE:
# We discourage the use of BinaryLogisticRegression model directly.
# Instead, you should use the general LogisticRegression model (which
# uses the BinaryLogisticRegression itself).

struct BinaryLogisticRegression <: Model
    β::ConjugateParameter{<:Normal}
    hasbias::Bool

    function BinaryLogisticRegression(β::ConjugateParameter{<:Normal}, hasbias::Bool)
        model = new(β, hasbias)
        β.stats = data -> stats_β(model, data...)
        return model
    end
end

function BinaryLogisticRegression(μ₀::Vector{T}, Σ₀::Matrix{T},
                                  μ::Vector{T}, Σ::Matrix{T},
                                  hasbias::Bool) where T <: AbstractFloat
    BinaryLogisticRegression(ConjugateParameter(Normal(μ₀, Σ₀), Normal(μ, Σ)),
                             hasbias)
end

function BinaryLogisticRegression(T; inputdim::Integer, initσ::Real = 0.0,
                                  pseudocounts::Real = 1.0, hasbias::Bool = true)
    dim = hasbias ? inputdim + 1 : inputdim

    # Prior parameters
    μ₀ = zeros(T, dim)
    Σ₀ = T(1/pseudocounts) * Matrix{T}(I, dim, dim)

    # Initial paremeterization of the posterior
    μ = μ₀ .+ T(initσ) .* randn(T)
    Σ = deepcopy(Σ₀)

    BinaryLogisticRegression(μ₀, Σ₀, μ, Σ, hasbias)
end

function BinaryLogisticRegression(;inputdim::Integer, initσ::Real = 0.0,
                                  pseudocounts::Real = 1.0, hasbias::Bool = true)
    BinaryLogisticRegression(Float64, inputdim = inputdim, initσ = initσ,
                             pseudocounts = pseudocounts, hasbias = hasbias)
end

function regressors(model::BinaryLogisticRegression, X::Matrix{T}) where T <: AbstractFloat
    if model.hasbias
        # We add an extra dimension always set to "1" to incorporate a bias
        # parameter to the model.
        return vcat(X, ones(T, 1, size(X, 2)))
    end
    return X
end

# Return the Polya-Gamma posterior of the augmented
# logistic function
function _augment_posterior(ψ²::Vector{T}) where T <: AbstractFloat
    b = ones(T, length(ψ²))
    qω = PolyaGamma(b)
    update!(qω, -.5 * ψ²)
    return qω
end

# Return the sufficient statistics of the log likelihood and the
# the expected statistics of β.
function _precompute(model::BinaryLogisticRegression, X::Matrix{T1},
                     z::Vector{T2},
                     Nₖ::Union{Vector{T2}, Nothing} = nothing) where {T1 <: AbstractFloat,
                                                                      T2 <: Real}
    X̂ = regressors(model, X)

    # Quadratic expansion of the regressors
    D, N = size(X̂)
    X̂X̂ᵀ = reshape(X̂, D, 1, N) .* reshape(X̂, 1, D, N)
    vec_X̂X̂ᵀ = reshape(X̂X̂ᵀ, :, N)

    # Expectation of the regressors: E[ T(β) ] = E[ ( β, vec(ββᵀ) )]
    E_Tβ = gradlognorm(model.β.posterior)
    vec_E_ββᵀ = E_Tβ[D+1:end]

    # Augmentation of the logistic function
    pω = PolyaGamma{T1, size(X̂, 2)}()
    ψ² = vec_X̂X̂ᵀ' * vec_E_ββᵀ
    qω = _augment_posterior(ψ²)
    E_ω = mean(qω)

    zstats = isnothing(Nₖ) ? (z .- 0.5) : z .- Nₖ * 0.5

    stats = vcat(reshape(zstats, 1, N) .* X̂,
                 -.5 .* reshape(E_ω, 1, :) .* vec_X̂X̂ᵀ)

    stats, E_Tβ, pω, qω
end

function (model::BinaryLogisticRegression)(X::Matrix{T1},
                                           z::Vector{T2},
                                           Nₖ::Union{Vector{T2}, Nothing} = nothing) where {T1 <: AbstractFloat,
                                                                                T2 <: Real}
    stats, E_Tβ, pω, qω = _precompute(model, X, z, Nₖ)

    # KL divegerence betwen qω and pω for each sample.
    q_η, p_η = naturalparam(qω), naturalparam(pω)
    KL = lognorm(pω, perdim = true) - lognorm(qω, perdim = true) - (p_η .- q_η) .* gradlognorm(qω)

    stats' * E_Tβ .- KL .- log(2)
end

function stats_β(model::BinaryLogisticRegression,
                       X::Matrix{T1},
                       z::Vector{T2},
                       Nₖ::Union{Vector{T2}, Nothing} = nothing) where {T1 <: AbstractFloat,
                                                                        T2 <: Real}
    stats, E_Tβ, _, _ = _precompute(model, X, z, Nₖ)
    stats
end

function predict(model::BinaryLogisticRegression,
                 X::Matrix{T}; a::T = 0.368) where T <: AbstractFloat
    # Following https://arxiv.org/pdf/1703.00091.pdf, we use the
    # approximation:
    # ⟨σ(ψ)⟩ ≈ σ(μ / √(1 + a * σ²))
    # where a = 0.368
    # In Bishop 4.5.2 a very similar approximation uses a = π / 8 ≈ 0.392...

    μ, Σ = stdparam(model.β.posterior)
    X̂ = regressors(model, X)

    σ²s = [sum((Σ .* x̂') .* x̂) for x̂ in eachcol(X̂)]
    ψs = (X̂' * μ) ./ sqrt.(1 .+ a * σ²s)
    1 ./ (1 .+ exp.(-ψs))
end

#######################################################################
# K-class Logistic Regression (LR) based on the stick-breaking process
#
# References
# ----------
# [1] Dependent Multinomial Modes Made Easy: Stick-Breaing with the
#     Polya-Gamma Augmentation (https://arxiv.org/pdf/1506.05843.pdf)

struct LogisticRegression{K} <: Model
    # The LR is defined by K - 1 binary LR models (see eq. (7) in [1])
    stickbreaking::Array{BinaryLogisticRegression}

    function LogisticRegression(sb::Array{BinaryLogisticRegression})
        model = new{length(sb) + 1}(sb)

        # We override the stats
        for (k, binarylr) in enumerate(sb)
            binarylr.β.stats = data -> stats_β(model, k, data...)
        end

        return model
    end
end

function LogisticRegression(μ₀::Vector{T}, Σ₀::Matrix{T}, μ::Vector{T},
                            Σ::Matrix{T}, hasbias::Bool; nclasses::Integer) where T <: AbstractFloat
    LogisticRegression([BinaryLogisticRegression(μ₀, Σ₀, deepcopy(μ), deepcopy(Σ), hasbias)
                        for i in 1:nclasses-1])
end

function LogisticRegression(T; inputdim::Integer, nclasses::Integer,
                            initσ::Real = 0.0, pseudocounts::Real = 1.0,
                            hasbias::Bool = true)
    dim = hasbias ? inputdim + 1 : inputdim

    # Prior parameters
    μ₀ = zeros(T, dim)
    Σ₀ = T(1/pseudocounts) * Matrix{T}(I, dim, dim)

    # Initial paremeterization of the posterior
    μ = μ₀ .+ T(initσ) .* randn(T)
    Σ = deepcopy(Σ₀)

    LogisticRegression(μ₀, Σ₀, μ, Σ, hasbias, nclasses = nclasses)
end

function LogisticRegression(;inputdim::Integer, nclasses::Integer,
                            initσ::Real = 0.0, pseudocounts::Real = 1.0,
                            hasbias::Bool = true)
    LogisticRegression(Float64, inputdim = inputdim, nclasses = nclasses,
                       initσ = initσ, pseudocounts = pseudocounts,
                       hasbias = hasbias)
end

getconjugateparams(model::LogisticRegression) = [m.β for m in model.stickbreaking]

function _encode(::LogisticRegression{K}, z::Vector{T}) where {K, T <: Integer}
    onehot = zeros(T, K - 1, length(z))
    for i = 1:K-1
        onehot[i, z .== i] .= 1
    end

    # See eq. (6) in [1]. In our case, N = 1
    Nₖ = zeros(T, K - 1, length(z))
    for i = 1:K-1
        Nₖ[i, z .>= i] .= 1
    end

   onehot, Nₖ
end

function (model::LogisticRegression)(X::Matrix{T1}, z::Vector{T2},) where {T1 <: AbstractFloat,
                                                                           T2 <: Integer}
    onehot, Nₖ = _encode(model, z)
    retval = zeros(T1, size(X, 2))
    for (k, m) in enumerate(model.stickbreaking)
        retval[z .>= k] .+= m(X[:, z .>= k], onehot[k, z .>= k], Nₖ[k, z .>= k])
    end
    retval
end

function stats_β(model::LogisticRegression, k::Integer,
                         X::Matrix{T1},
                         z::Vector{T2}) where {T1 <: AbstractFloat, T2 <: Real}
    onehot, Nₖ = _encode(model, z)
    stats_β(model.stickbreaking[k], X[:, z .>= k], onehot[k, z .>= k],
                  Nₖ[k, z .>= k])
end

function predict(model::LogisticRegression{K}, X::Matrix{T}) where {K, T <: AbstractFloat}
    retval = zeros(T, K, size(X, 2))
    for k in 1:K
        residual = dropdims(sum(retval[1:k, :], dims = 1), dims = 1)
        if k < K
            retval[k, :] = predict(model.stickbreaking[k], X) .* (1 .- residual)
        else
            retval[k, :] = (1 .- residual)
        end
    end
    retval
end


###
#   IntervalUnion Arithmetic
###

# Special case for /. Probably better with promotion rules and convert
for op in (:+, :-, :/, :*, :min, :max, :^, :log, :<, :>)

    if op != :/; @eval ($op)( x::IntervalUnion, y::IntervalUnion) = intervalUnion([$op(xv, yv) for xv in x.v, yv in y.v][:]); end

    @eval ($op)( x::IntervalUnion, y::Interval) = intervalUnion(broadcast($op, x.v, y))
    if op != :/; @eval ($op)( x::Interval, y::IntervalUnion) = intervalUnion(broadcast($op, x, y.v)); end

    @eval ($op)( x::IntervalUnion, n::Real) = intervalUnion(broadcast($op, x.v, n))
    @eval ($op)( n::Real, x::IntervalUnion) = intervalUnion(broadcast($op, n, x.v))

end

for op in (:-, :sin, :cos, :tan, :exp, :log)
    @eval ($op)( x::IntervalUnion) = intervalUnion(broadcast($op, x.v))
end

# Does x/y for y\{0} if 0 ∈ y
function /(x :: IntervalUnion, y :: IntervalUnion)
    yNew = bisect(y, 0)
    return intervalUnion([xv / yv for xv in x.v, yv in yNew.v][:])
end

function /(x :: Interval, y :: IntervalUnion)
    yNew = bisect(y, 0)
    return intervalUnion(broadcast(/, x, yNew.v))
end

function sqrt( x:: IntervalUnion)
    us = sqrt.(x.v)
    return intervalUnion([us; -us])
end


using Distributions
using CausalInference
using LightGraphs

quantile(Normal(), 0.95)

if !isfile(joinpath("data","NCI-60.csv"))
    cd("data") do
        run(`wget http://nugget.unisa.edu.au/ParallelPC/data/real/NCI-60.csv`)
    end
end

data = readdlm("data/NCI-60.csv")
d = size(data,2)
X = data[:,1:d]'
d, n = size(X)
C = Symmetric(cor(X, 2))
p = 0.01
@time h, s = skeleton(d, gausscitest, (C,n), quantile(Normal(), 1-p/2))
println("inferred edges ", ne(h))
writedlm("data/adjac.csv", adjacency_matrix(h))


using ApproxFun, SpecialFunctions, LinearAlgebra, Test
    import DomainSets
    import DomainSets: UnionDomain
    import ApproxFun: ChebyshevDirichlet, Ultraspherical, PiecewiseSegment, ContinuousSpace, space, SpaceOperator,
                        testspace, testbandedoperator, testraggedbelowoperator, testcalculus, testtransforms,
                        testfunctional

@testset "Spaces" begin
    @testset "ChebyshevDirichlet" begin
        testtransforms(ChebyshevDirichlet{1,1}())

        @test Fun(exp,ChebyshevDirichlet{1,1})(.1) ≈ exp(.1)
        @test Fun(Fun(exp,ChebyshevDirichlet{1,1}),Ultraspherical(1))(.1) ≈ exp(.1)

        @test domain(ChebyshevDirichlet{1,1}()) === ChebyshevInterval()
        @test Fun(Fun(exp,ChebyshevDirichlet{1,1}(Segment(-1,1))),Ultraspherical(1))(.1) ≈ exp(.1)

        d = ChebyshevInterval()
        sp = ChebyshevDirichlet{1,1}(d)
        B=Dirichlet(sp)
        D=Derivative(d)
        L=D^2+I
        @test B[1:2,1:4] ≈ [1 -1 0 0; 1 1 0 0]

        @test csc(2)sin(1 - 0.1)  ≈ ([Dirichlet(d);L]\[[1.,0.],0.])(0.1)
        @test csc(2)sin(1 - 0.1)  ≈ ([B;L]\[[1.,0.],0.])(0.1)

        @test norm(([B;L]\[[1.,0],0])-([Dirichlet(d);L]\[[1.,0],0])) <10eps()
    end

    @testset "PiecewiseSpace" begin
        x = Fun(identity,UnionDomain(-1..0, 0..1))
        sp=space(x)
        testtransforms(sp;minpoints=2)

        D = Derivative(sp)
        testbandedoperator(D)

        B = [Dirichlet(sp); continuity(sp,0:1)]
        u = [B; D^2] \ [[1,0],zeros(2),0];
        u2 = [Dirichlet();Derivative(Chebyshev())^2] \ [[1.,0],0]
        @test u(0.) ≈ u2(0.)

        x = Fun(identity, UnionDomain(-10..0, 0..1, 1..15))
        sp=space(x)
        D=Derivative(sp)
        B=Dirichlet(sp)

        u=[B;
            continuity(sp,0:1);
            D^2-x]\[[airyai(-10.),0],zeros(4),0];

        @test u(0.) ≈ airyai(0.)

        s=Fun(sin,-2..2)|>abs
        c=Fun(cos,-2..2)|>abs
        sc=Fun(x -> abs(sin(x))+abs(cos(x)), ContinuousSpace(PiecewiseSegment([-2,(-π/2),0,π/2,2])))
        @test norm(sc-(c+s))<100eps()
    end

    @testset "max/min creates breakpoints" begin
        x = Fun()
        g = 4*(x-0.2)
        f = max(-1,g)
        f2 = min(f,1)

        @test norm(max(x,x)-x)<100eps()
        @test norm(min(x,x)-x)<100eps()

        f3=Fun(x->x<-0.05 ? -1.0 : (x<0.45 ? 4*(x-.2) : 1), (-1..1) \ [-0.05,0.45])
        @test norm(f2.(range(-1,stop=1,length=10))-f3.(range(-1,stop=1,length=10))) < 4eps()
    end

    @testset "Complex piecewise" begin
        x = Fun(identity, Segment(im,0) ∪ Segment(0,1))
        @test x(0.5) ≈ 0.5
        @test x(0.5im) ≈ 0.5im

        @test Fun(Fun(1.0),space(x))(0.5) == 1.0
        @test Fun(Fun(1.0),space(x))(0.5im) == 1.0

        @test (x+1)(0.5) ≈ 1.5
        @test (x-1)(0.5) ≈ -0.5
        @test (1-x)(0.5) ≈ 0.5

        @test sqrt(1-x)(0.2im) ≈ sqrt(1-0.2im)
        @test sqrt(1-x)(0.2) ≈ sqrt(1-0.2)

        w=2/(sqrt(1-x)*sqrt(1+im*x))
        @test w(0.2im) ≈ 2/(sqrt(1-0.2im)*sqrt(1+im*(0.2im)))
        @test w(0.2) ≈ 2/(sqrt(1-0.2)*sqrt(1+im*(0.2)))
    end

    @testset "Equivalent spaces" begin
        @test norm(Fun(cos,Chebyshev)-Fun(cos,Jacobi(-0.5,-0.5)))<100eps()
        @test norm(Fun(cos,Chebyshev)-Fun(cos,JacobiWeight(0,0)))<100eps()
        @test norm(Fun(cos,Jacobi(-0.5,-0.5))-Fun(cos,JacobiWeight(0,0))) < 100eps()
        @test norm(Fun(cos,Chebyshev)-Fun(cos,JacobiWeight(0,0,Jacobi(-0.5,-0.5))))<100eps()
        @test norm(Fun(cos,Jacobi(-0.5,-0.5))-Fun(cos,JacobiWeight(0,0,Jacobi(-0.5,-0.5))))<100eps()
    end

    @testset "ContinuousSpace" begin
        d = PiecewiseSegment(1.,2.,3.,4.)
        S = ContinuousSpace(d)
        testtransforms(S; minpoints=3, invertibletransform=false)

        D = Derivative(S)
        testbandedoperator(D)

        A = [ldirichlet(S); D-I]
        testraggedbelowoperator(A)
        u = [ldirichlet(S); D-I] \ [exp(1.); 0]

        @test u(1.1) ≈ exp(1.1)
        @test u(3.4) ≈ exp(3.4)
        @test last(u) ≈ exp(4)

        d=PiecewiseSegment(0,1.,1+im,im,0.)
        s=Space(d)

        # narrow down bug
        @test s isa ContinuousSpace
        @test ApproxFun.rangetype(s) == Float64
        cs=ApproxFun.canonicalspace(s)
        @test ApproxFun.rangetype(cs) == Float64

        @test conversion_type(s,cs) == s

        @test Fun(exp,d)(.1) ≈ exp(.1)
    end

    @testset "Triple SumSpace" begin
        x=Fun()
        w=log(1-x)+sqrt(1-x^2)
        f=w+x
        @test f(0.1) ≈ (w(0.1)+0.1)
        @test (w+1)(0.1) ≈ (w(0.1)+1)
        @test (w+x+1)(0.1) ≈ (w(0.1)+1.1)
        @test ((w+x)+1)(0.1) ≈ (w(0.1)+1.1)
    end

    @testset "SumSpace bug" begin
        dsp=JacobiWeight(1.,0.,Jacobi(1.,0.,0..1))⊕JacobiWeight(0.5,0.,Jacobi(0.5,-0.5,0..1))
        rsp=Legendre(0..1)⊕JacobiWeight(0.5,0.,Jacobi(0.5,0.5,0..1))

        C=Conversion(dsp,rsp)

        f=Fun(dsp,[1.,2.,3.,4.,5.])
        @test f(0.1) ≈ (C*f)(0.1)
    end

    @testset "Piecewise + Constant" begin
        Γ=Circle() ∪ Circle(0.0,0.4)
        o=ones(Γ)
        @test o(1.) ≈ 1.0
        @test o(0.4) ≈ 1.0

        G=Fun(z->in(z,component(Γ,2)) ? [1 0; -1/z 1] : [z 0; 0 1/z],Γ)
        @test (G-I)(exp(0.1im)) ≈ (G(exp(0.1im))-I)
    end

    @testset "Previoius segfault" begin
        x=Fun(identity,-1..1)
        f=x+sin(2x)*sqrt(1-x^2)
        @test f(0.1) ≈ 0.1+sin(2*0.1)*sqrt(1-0.1^2)
    end

    @testset "Multiple piecewisespace" begin
        x=Fun(identity,-3 .. -2)+Fun(identity,2..3)
        w=sqrt(9-x^2)
        f=w+Fun()
        @test (f+w)(2.5) ≈ 2w(2.5)
        @test (f+w)(.5) ≈ f(.5)
    end

    @testset "Jacobi recurrence bug" begin
        S=Jacobi(0.5,-0.5)
        f=Fun(exp,S)
        @test f(0.1) ≈ exp(0.1)
    end

    @testset "Jacobi evaluation bug" begin
        S=Jacobi(0.5,-0.5)
        f=Fun(exp,S)
        B = ldirichlet(S)
        testfunctional(B)
        @test B*f ≈ exp(-1)
    end

    @testset "cancellation conversion" begin
        x=Fun(0..1)
        f=exp(x)-1
        Fun(f,JacobiWeight(1.,0.,0..1))
    end

    @testset "Hermite" begin
        f=Fun(x->x+x^2,Hermite())
        @test f(1.) ≈ 2.
        @test values(f) ≈ points(f)+points(f).^2

        D = Derivative(Hermite())
        testbandedoperator(D)

        g = D*f
        @test g(1.) ≈ 3.
    end

    @testset "Arc exp" begin
        z=Fun(identity,Arc(0.,.1,0.,π/2))
        @test exp(z)(0.1exp(0.2im)) ≈ exp(0.1exp(0.2im))
    end

    @testset "Extending function" begin
        Γ = Segment(-im,1.0-im) ∪ Curve(Fun(x->exp(0.8im)*(x+x^2-1+im*(x-4x^3+x^4)/6))) ∪ Circle(2.0,0.2)

        @test isempty(component(Γ,1)\component(Γ,1))
        @test Γ \ component(Γ,1) == component(Γ,2) ∪ component(Γ,3)

        @test norm(Fun(ones(component(Γ,1)),Γ) - Fun(x->x ∈ component(Γ,1) ? 1.0 : 0.0,Γ)) == 0
    end

    @testset "Line" begin
        f=Fun(z->2exp(z^2),PeriodicLine(0.,π/2))
        @test f(1.1im) ≈ 2exp(-1.1^2)

        f=Fun(z->2exp(z^2),Line(0.,π/2))
        @test f(1.1im) ≈ 2exp(-1.1^2)
    end

    @testset "Exp Curve" begin
        a = 1+10*im; b = 2-6*im
        d = Curve(Fun(x->1+a*x+x^2+b*x^3))

        x=Fun(d)

        @test exp(x)(1+a*0.1+0.1^2+b*0.1^3) ≈ exp(1+a*0.1+0.1^2+b*0.1^3)
    end

    @testset "ChebyshevDirichlet multiplication" begin
        S=ChebyshevDirichlet()
        x=Fun()
        @test norm((ApproxFun.Recurrence(S)*Fun(exp,S)-Fun(x->x*exp(x),S)).coefficients) < 100eps()
        @test norm((x*Fun(exp,S)-Fun(x->x*exp(x),S)).coefficients) < 100eps()

        @test Integral(S,1)*Fun(S,[1.,2.,3.]) ≈ integrate(Fun(Fun(S,[1.,2.,3.]),Chebyshev()))
    end

    @testset "QuotientSpace" begin
        for (bcs,ret) in ((Dirichlet(Chebyshev()),[1 -1 0 0 0;1 1 0 0 0]),
                          (Neumann(Chebyshev()),[0 1 -4 0 0;0 1 4 0 0]),
                          ([DefiniteIntegral(Chebyshev());SpaceOperator(BasisFunctional(2),Chebyshev(),ConstantSpace())],[2 0 0 0 0;0 1 0 0 0]),
                          (vcat(bvp(Chebyshev(),4)...),[1 -1 1 -1 0;0 1 -4 9 0;1 1 1 1 0;0 1 4 9 0]))
            QS = QuotientSpace(bcs)
            C = Conversion(QS, QS.space)

            norm((bcs*C)[1:size(bcs, 1),1:5] - ret) < 1000eps()
        end
    end

    @testset "Union of ChebyshevDirichlet" begin
        dom = UnionDomain(0..1, 2..3)
        @test components(union(JacobiWeight.(-0.5,-0.5,ChebyshevDirichlet{1,1}.(components(dom)))...)) ==
            (JacobiWeight.(-0.5,-0.5,ChebyshevDirichlet{1,1}.(components(dom)))...,)
    end

    @testset "Ultraspherical special functions" begin
        x = Fun(Ultraspherical(2,0..1))
        sqrt(x)(0.1) ≈ sqrt(0.1)

        f = Fun(x->x*exp(x),Ultraspherical(1,0..1))
        sqrt(f(0.1)) ≈ sqrt(f)(0.1)
    end

    @testset "Fast Ultraspherical*Chebyshev" begin
        f = Fun(exp, Ultraspherical(1))
        g = Fun(exp)

        @test ApproxFun.hasfasttransform(f)
        @test ApproxFun.pointscompatible(f,g)
        @test ApproxFun.hasfasttransformtimes(f,g)
        @test ApproxFun.transformtimes(f,g)(0.1) ≈ exp(0.2)
    end

    @testset "PointSpace" begin
        δ = Fun(ApproxFun.PointSpace([2.0]),[1.0])
        f = Fun(x->cos(50x)) + δ
        g = Fun(x->cos(50x),Ultraspherical(1)) + δ
        @test domain(f) == domain(g)

        @test (f*g)(0.1) ≈ cos(50*0.1)^2
        @test (f*g)(2.0) ≈ 1

        @test exp(f)(0.1) ≈ exp(cos(50*0.1))
        @test exp(f)(2.0) ≈ exp(1)

        @test iszero(sign(-Fun(zero)))
        @test iszero(sign(Fun(zero)))
        @test iszero(abs(-Fun(zero)))
        @test iszero(abs(Fun(zero)))
        @test angle(-Fun(zero)) ≈ Fun(π)
        @test iszero(angle(Fun(zero)))

        @test sign(f)(0.1) ≈ sign(cos(50*0.1))
        @test sign(f)(2.0) ≈ 1
        @test abs(f)(0.1) ≈ abs(cos(50*0.1))
        @test abs(f)(2.0) ≈ 1
        @test angle(f)(0.1) ≈ angle(cos(50*0.1))
        @test angle(f)(2.0) ≈ 0
    end

    @testset "Jump Locations" begin
        x = Fun(UnionDomain(0..1, 2..3))
        @test length(jumplocations(sign(x))) == 0

        x = Fun(Chebyshev(-1..1))
        @test length(jumplocations(x)) == 0
        @test all(jumplocations(sign(x) + sign(x+0.2)) .≈ [-0.2, 0])
    end

    @testset "one for SumSpace" begin
        S = Jacobi(0,1) ⊕ JacobiWeight(1/3,0,Jacobi(1/3,2/3)) ⊕ JacobiWeight(2/3,0,Jacobi(2/3,1/3))
        o = ones(S)
        @test o(0.5) ≈ 1
    end

    @testset "blockbandwidths for FiniteOperator of pointscompatibleace bug" begin
        S = ApproxFun.PointSpace([1.0,2.0])
        @test ApproxFun.blockbandwidths(FiniteOperator([1 2; 3 4],S,S)) == (0,0)
    end

    @testset "SumSpace Conversion" begin
        H = ApproxFun.HeavisideSpace([-1.0,0.0,1.0])
        C = ApproxFun.ContinuousSpace(ApproxFun.PiecewiseSegment([-1.0,0,1]))
        S = H + C
        P = Ultraspherical(1,Segment(-1.0,0.0)) ∪ Ultraspherical(1,Segment(0.0,1.0))
        f = Fun(S, randn(100))
        @test f(0.1) ≈ Fun(f, P)(0.1)

        @test Conversion(S,P)[1,1] == 1.0
        @test Conversion(S,P)[1,2] == 0.5
        @test rangespace(Conversion(S,P)) == Ultraspherical(1,Segment(-1.0,0.0)) ∪ Ultraspherical(1,Segment(0.0,1.0))
        @test domainspace(Conversion(S,P)) ==
            ApproxFun.SumSpace(ApproxFun.HeavisideSpace([-1.0,0.0,1.0]),
                    ApproxFun.ContinuousSpace(ApproxFun.PiecewiseSegment([-1.0,0,1])))

        D = ApproxFun.DiracSpace([-1.0,0.0,1.0])
        S2 = ApproxFun.SumSpace(D , P)
        f = Fun(S, randn(100))
        @test (Conversion(S,S2) * f)(0.1) ≈ f(0.1)
    end

    @testset "Mix Fourier-Chebyshev (#602)" begin
        s = Chebyshev(-π..π)
        a = Fun(t-> 1+sin(cos(2t)), s)
        L = Derivative() + a
        f = Fun(t->exp(sin(10t)), s)
        B = periodic(s,0)
        uChebyshev = [B;L] \ [0.;f]

        s = Fourier(-π..π)
        a = Fun(t-> 1+sin(cos(2t)), s)
        L = Derivative() + a
        f = Fun(t->exp(sin(10t)), s)
        uFourier = L\f

        @test norm(uFourier-uChebyshev) ≤ 100eps()
    end

    @testset "real line integral" begin
        f = x -> x^2/(x^4+1)

        z₁,z₂,z₃,z₄ = exp(im*π/4), exp(3im*π/4), exp(5im*π/4), exp(7im*π/4)

        res₁ = z₁^2 / ((z₁ - z₂)*(z₁ - z₃)*(z₁ - z₄) )
        res₂ = z₂^2 / ((z₂ - z₁)*(z₂ - z₃)*(z₂ - z₄) )

        @test_skip sum(Fun(f, Line())) ≈ 2π*im*(res₁ + res₂)
    end

    @testset "Piecewise Legendre" begin
        f = Fun(x -> exp(-40(x-0.1)^2), Legendre(0..1), 1000)
        @test f(0.2) ≈ exp(-40(0.1)^2)

        f = Fun(x -> exp(-40(x-0.1)^2), Legendre(0..1))
        @test f(0.2) ≈ exp(-40(0.1)^2)

        sp = Legendre(Segment(0 , -1)) ⊕ Legendre(0 .. 1)
        @time f = Fun(x->sign(x)*exp(-40(x-0.1)^2), sp)
        @test f(0.2) ≈ exp(-40(0.1)^2)
        @test f(-0.2) ≈ -exp(-40(0.3)^2)
    end

    @testset "Ambiguous evaluate" begin
        f = Fun(1.0, Chebyshev(NaN..NaN))
        @test f(0.1) ≈ 1.0
    end

    @testset "x + abs(x) (#637)" begin
        x = Fun()
        g = abs(x) + x
        @test ncoefficients(g) == 3
        @test g(0.1) ≈ 0.2
        @test g(-0.2) ≈ 0.0
    end

    @testset "piecewise sample (#635)" begin
        f = abs(Fun(sin, -5..5))
        @test integrate(f)(-4.0) ≈ -(cos(-4.0) - cos(-5.0))
        @test -(cos(-π) - cos(-5.0)) + cos(-3.0) - cos(-π) ≈ integrate(f)(-3.0)
        r = ApproxFun.sample(f,10)
        @test maximum(r) ≤ 5
        @test minimum(r) ≥ -5
    end
end


# function numdecdigits(x)
#     d = 0
#     while round(x-round(x),15-d) != 0
#         x *= 10
#         d += 1
#     end
#     d
# end

# function numintdigits(x)
#     d = 0
#     while round(Int,x) != 0
#         x /= 10
#         d += 1
#     end
#     d
# end

# returns the exponent if x is expressed in scientific notation (base 10)
# e.g.
# 0.00123 = 1.23e-3
# so exponent is -3
function exponent10(x)
    if x == zero(x)
        return 0
    elseif x > zero(x)
        return floor(Int,log10(x))
    else
        return ceil(Int,log10(-x))
    end
end

# return the part of x without the exponent if expressed in scientific notation
# e.g.
# 0.00123 = 1.23e-3
# mantissa = 1.23
mantissa10(x) = x*exp10(-exponent10(x))

latex_expfmt(e) = string(" × 10^{",e,'}')
standard_expfmt(e) = string('e',e)

function formatter(x::AbstractFloat;
                   n::Int = 2,                      # number of significant figures
                   expfmt = latex_expfmt)           # exponent formatter

    (isnan(x) || isinf(x)) && return (string(x),' ',"")

    e = exponent10(x)                               # exponent
    m = round(x*exp10(-e), n-1)                     # mantissa

    # case 1: use scientific notation if number is extra long
    if e > 4 || e < -4
        s = string(format("{:.$(n-1)f}", m), expfmt(e))
        l,r = split(s, '.')
        return (l,'.',r)
    else
        xf = m*exp10(e)

        if n > e+1
            l,r = split(format("{:.$(n-e-1)f}", xf), '.')
            return (l,'.',r)
        elseif n == e+1
            l = format("{:d}", round(Int,xf))
            return (l,'.',"")
        else
            l = format("{:d}", round(Int,xf))
            return (l,' ',"")
        end
    end
end

formatter(x::Integer; kwargs...) = (string(x),' ',"")

formatter(x::AbstractString; kwargs...) = (x,' ',"")

function writetable(io::IO,
                    table::Matrix;
                    header = nothing,                    # column names
                    notes = nothing,
                    expfmt = latex_expfmt,
                    n = 2)                               # number of significant digits

    s = map(x -> formatter(x,n=n,expfmt=expfmt), table)

    if !is(notes, nothing)
        for (k,v) in notes
            i,j = k
            x = s[i,j]
            s[i,j] = (x[1],x[2],x[3] * " ($v)")
        end
    end

    ll,lr = map(x -> length(x[1]), s), map(x -> length(x[3]), s)                    # lengths of left, right parts of columns
    wl,wr = maximum(ll,1), maximum(lr,1)                                            # left, right column part widths
    w = wl+wr+1                                                                     # column widths

    addcols = cols -> string("| ", join(cols, " | "), " |")                         # combine columns in one row
    addparts = (s,wl,wr) -> string(lpad(s[1],wl), s[2], rpad(s[3],wr))              # combine left, middle, right parts of table entry

    s = [addparts(s[i,j],wl[j],wr[j]) for i in 1:size(s,1), j in 1:size(s,2)]

    if !is(header, nothing)
        wh = map(length, header)                                                    # column header widths
        s = [lpad(s[i,j], wh[j]) for i in 1:size(s,1), j in 1:size(s,2)]            # pad column to header width
        pheader = [rpad(header[i], w[i]) for i in 1:length(header)]                 # padded column header
        hs = addcols(pheader)                                                       # header string
        w = map(max,w,wh)
        hline = addcols([repeat("-", l) for l in w])
        ts = join([addcols(s[i,:]) for i in 1:size(s,1)], "\n")
        return write(io, hs, '\n', hline, '\n', ts)
    else
        ts = join([addcols(s[i,:]) for i in 1:size(s,1)], "\n")
        return write(io, ts)
    end
end


function my_rescale(
    X::Vector{Float64},
    Wk::Vector{Float64},
    observed::Vector{Float64},
)
    X = X .- observed[1]
    X = X ./ observed[2]
    X = vcat(X[1] - (X[2] - X[1]), X)
    X = vcat(X, X[end] + (X[2] - X[1]))
    Wk = vcat(zero(Float64), Wk)
    Wk = vcat(Wk, zero(Float64))
    Wk = Wk ./ sum(Wk)
    return X, Wk
end

function cubic_spline_int(X::Vector{Float64}, NewX::Vector{Float64}, Wk::Vector{Float64})
    bw = abs(NewX[2] - NewX[1])
    K = size(X, 1)
    ok = 0
    mainArealiminf = 1
    mainArealimsup = K + 2
    NewX = NewX[mainArealiminf]:bw:(NewX[mainArealiminf]+(bw*(size(Wk, 1)-1)))
    interp_cubic = Interpolations.extrapolate(
        Interpolations.scale(
            Interpolations.interpolate(
                Wk,
                Interpolations.BSpline(Interpolations.Cubic(Interpolations.Line(Interpolations.OnGrid()))),
            ),
            NewX,
        ),
        Interpolations.Line(),
    )

    Wk_cubic = zeros(K)
    for k = 1:K
        Wk_cubic[k] = interp_cubic(X[k])
    end

    Wk = clamp.(Wk_cubic, 1e-20, one(Float64))
    Wk = Wk ./ sum(Wk)
    return Wk::Vector{Float64}
end


# PROJECT_BOX    Project a point onto a box (hyper-rectangle).
#
#   project_box(v,l,u) is the projection of v onto
#   the set { x | l <= x <= u }.

function project_box(
	v,
	l,
	u
)
    x = max(l, min(v, u))
    return x
end


module RationalApproximations

import GenericSVD: svd!
#import GenericLinearAlgebra: qr, eigen
import LinearAlgebra: norm, Diagonal, qr, eigen
import Optim: optimize

export aaa_v1, aaa_v2, locate_extrema_idx, minimax

"""
Represents a rational function of type `(m,m)` in barycentric form.
"""
struct RationalFunc{T<:AbstractFloat}
    α::Vector{T}
    β::Vector{T}
    supp_pt::Vector{T}
end # module

function (r::RationalFunc{T})(x::T) where T <: AbstractFloat
    numer = zero(T)
    denom = zero(T)
    for k = 1:length(r.α)
        δ = x - r.supp_pt[k]
        if abs(δ) ≤ eps(x)
            numer = r.α[k]
            denom = r.β[k]
            break
        else
            numer += r.α[k] / δ
            denom += r.β[k] / δ
        end
    end
    return numer / denom
end

"""
    aaa_v1(F, S, max_m, tol) -> r, err

Adaptive Antoulas-Anderson algorithm computes rational approximations `r[m]`of
type `(m,m)` to a real function `f(x)` based on sample values `F[i] = f(S[i])`.
Stops if `err[m] < tol * norm(F, Inf)` or when `m` reaches `max_m`.

This is version 1 of the algorithm, for which `r[m]` interpolates `f` at the
support points based on `m` weights `w[m]` with `α[j] = w[j] * F[j]` and
`β[j] = w[j]`.
"""
function aaa_v1(F::Vector{T}, S::AbstractVector{T},
                    max_m::Integer, tol::T) where T <: AbstractFloat
    K = length(F)
    if length(S) ≠ K
        throw(DimensionMismatch("F and S must have the same length"))
    end
    avg_F = sum(F) / K
    R = fill(avg_F, K)
    J = collect(1:K)
    supp_pt = T[]
    f = T[]
    w = T[]
    err = T[]
    r = RationalFunc[]
    for m = 1:max_m
        val, i = findmax(abs.(F[J]-R[J]))
        j = J[i]
        push!(supp_pt, S[j])
        push!(f, F[j])
        J = remove(J, i)
        A = T[ (F[J[i]]-f[j])/(S[J[i]]-supp_pt[j]) for i=1:K-m, j=1:m ]
        Fact = svd!(A)
        w = Fact.V[:,m]
        α = w .* f
        β = w
        push!(r, RationalFunc(α, β, copy(supp_pt)))
        R[J] .= r[m].(S[J])
        push!(err, norm(F[J] - R[J], Inf))
        if err[m] < tol * norm(F, Inf)
            r = r[1:m]
            err = err[1:m]
            break
        end
    end
    return r, err
end

"""
    remove(J, i)

Returns the vector obtained from `J` by removing the `i`th element.
"""
function remove(J::Vector, i::Integer)
    m = length(J)
    if i ≤ 0 || i ≥ m+1
        throw(BoundsError(J, i))
    end
    return [J[1:i-1]; J[i+1:m]]
end

"""
    aaa_v1(F, S, max_m, tol) -> r, err

Adaptive Antoulas-Anderson algorithm computes rational approximations `r[m]`of
type `(m,m)` to a real function `f(x)` based on sample values `F[i] = f(S[i])`.
Stops if `err[m] < tol * norm(F, Inf)` or when `m` reaches `max_m`.

This is version 2 of the algorithm, for which the coefficients `α[j]` and `β[j]`
vary independently.
"""
function aaa_v2(F::Vector{T}, S::AbstractVector{T},
                    max_m::Integer, tol::T) where T <: AbstractFloat
    K = length(F)
    if length(S) ≠ K
        throw(DimensionMismatch("F and S must have the same length"))
    end
    avg_F = sum(F) / K
    R = fill(avg_F, K)
    J = collect(1:K)
    supp_pt = T[]
    supp_pt_idx = fill(-1, K)
    resid = F - R
    f = T[]
    w = T[]
    err = T[]
    r = RationalFunc[]
    buffer = Vector{T}(undef, K*2*max_m)
    for m = 1:max_m
        val, l = findmax(abs.(resid[1:K-m+1]))
        i = 0
        for i1 = 1:l
            i += 1
            if supp_pt_idx[i] ≥ 0
                i += 1
            end
        end
        supp_pt_idx[i] = m
        push!(supp_pt, S[i])
        push!(f, F[i])
        A = reshape(view(buffer, 1:(K-m)*2*m), K-m, 2m)
        for j = 1:m
            i1 = 0
            for i = 1:K
                if supp_pt_idx[i] ≤ 0
                    i1 += 1
                    A[i1,j]   = -1 / ( S[i] - supp_pt[j] )
                    A[i1,m+j] =  F[i] / ( S[i] - supp_pt[j] )
                end
            end
        end
        Fact = svd!(A)
        α = Fact.V[1:m,2m]
        β = Fact.V[m+1:2m,2m]
        push!(r, RationalFunc(α, β, copy(supp_pt)))
        i1 = 0
        for i = 1:K
            if supp_pt_idx[i] ≤ 0
                i1 += 1
                resid[i1] = F[i] - r[m](S[i])
            end
        end
        push!(err, norm(resid[1:K-m], Inf))
        if err[m] ≤ tol * norm(F, Inf)
            break
        end
    end
    return r, err
end

"""
    minimax(f, x, supp_pt, Interval, max_iterations, clamp) -> r, zmin, zmax
"""
function minimax(f::Function, x::Vector{T}, supp_pt::Vector{T},
                 Interval::Tuple{T,T}, max_iterations::Integer,
                 clamp=:no) where T <: AbstractFloat
    N = length(supp_pt)
    if length(x) ≠ 2N
        throw(DimensionMismatch("x must be twice as long as supp_pt"))
    end
    α = Vector{T}(undef, N)
    β = similar(α)
    δ = one(T)
    zmax = Vector{T}(undef, max_iterations)
    zmin = similar(zmax)
    r = Vector{RationalFunc{T}}(undef, max_iterations)
    for i = 1:max_iterations
        α, β, λ = equioscillation!(f, x, supp_pt, clamp)
        r[i] = RationalFunc(α, β, supp_pt)
        zmin[i], zmax[i] = locate_extrema!(x, Interval) do t
            return f(t) - r[i](t)
        end
        δ = zmax[i] - zmin[i]
        if δ ≤ 10*eps(zmax[i])
            zmax = zmax[1:i]
            zmin = zmin[1:i]
            r = r[1:i]
            break
        end
    end
    return r, zmin, zmax
end

"""
    locate_extrema_idx(resid, m)
"""
function locate_extrema_idx(resid::Vector{T},
                            m::Integer) where T <: AbstractFloat
    K = length(resid)
    idx = Int64[]
    push!(idx, 1)
    fwd = resid[2] - resid[1]
    σ = sign(fwd)
    for k = 2:K-1
        fwd = resid[k+1] - resid[k]
        if σ * sign(fwd) == -1.0
            push!(idx, k)
            σ = -σ
        end
    end
    push!(idx, K)
    len_idx = length(idx)
    if len_idx < 2m
        error("Found only $len_idx extrema; need $(2m)")
    elseif len_idx > 2m
        var = abs.(diff(resid[idx]))
        I = sortperm(var, rev=true)
        new_idx = [ idx[I[1:2m-1]]; idx[end] ]
        idx = sort(new_idx)
    end
    return idx
end

"""
    equioscillation!(f, x, supp_pt, clamp) -> α, β, λ
"""
function equioscillation!(f::Function, x::Vector{T},
                         supp_pt::Vector{T},
                         clamp::Symbol) where T <: AbstractFloat
    N = length(supp_pt)
    if length(x) ≠ 2N
        throw(DimensionMismatch("x must be twice as long as supp_pt"))
    end
    M = Matrix{T}(undef, 2N, N)
    for k = 1:N
        for l = 1:2N
            Mlk = one(T)
            for p = 1:k-1
                Mlk *= x[l] - supp_pt[p]
            end
            for p = k+1:N
                Mlk *= x[l] - supp_pt[p]
            end
            M[l,k] = Mlk
        end
    end
    d = Vector{T}(undef, 2N)
    for l = 1:2N
        Dl = one(T)
        for p = 1:l-1
            Dl *= x[l] - x[p]
        end
        for p = l+1:2N
            Dl *= x[l] - x[p]
        end
        d[l] = 1/sqrt(abs(Dl))
    end
    D = Diagonal(d)
    Fact = qr(D*M)
    Q1, R = Matrix(Fact.Q), Fact.R
    sgn = Vector{T}(undef, 2N)
    sgn[1] = -one(T)
    for l = 2:2N
        sgn[l] = -sgn[l-1]
    end
    S = Diagonal(sgn)
    F = Diagonal(f.(x))
    A = Q1' * ( S * F ) * Q1
    if clamp == :no
        EFact = eigen(A)
    else
        I0 = Diagonal(ones(2N))
        if clamp == :left
            I0[1,1] = zero(T)
        elseif clamp == :right
            I0[2N,2N] = zero(T)
        elseif clamp == :both
            I0[1,1] = zero(T)
            I0[2N,2N] = zero(T)
        else
            error("illegal value $clamp for argument clamp")
        end
        B = Q1' * I0 * Q1
        EFact = eigen(A, B)
    end
    y = EFact.vectors
    Q1y = Q1 * y
    q = D \ Q1y
    k_no_root = -1
    check_count = 0
    for k = 1:N
        no_root = true
        for l = 2:2N
            if sign(q[l,k]) != sign(q[1,k])
                no_root = false
            end
        end
        if no_root
            k_no_root = k
            check_count += 1
        end
    end
    if check_count ≠ 1
        error("Failed to find unique eigenpair with q≠0")
    end
    λ = EFact.values[k_no_root]
    if clamp == :no
        α = R \ ( Q1' * F * Q1y[:,k_no_root] )
    else
        α = R \ ( Q1' * (F-λ*I0*S) * Q1y[:,k_no_root] )
    end
    β = R \ y[:,k_no_root]
    return α, β, λ
end

"""
    locate_extrema!(resid, x, Interval) -> zmin, zmax
"""
function locate_extrema!(resid::Function, x::Vector{Float64},
                         Interval::Tuple{Float64,Float64})
    n = length(x)
    if resid(x[1]) < resid(x[2])
        sgn = 1
    else
        sgn = -1
    end
    res = optimize(Interval[1], x[2]) do t
        return sgn * resid(t)
    end
    x[1] = res.minimizer
    zmax = zmin = abs(resid(x[1]))
    for l = 2:n-1
        sgn = -sgn
        res = optimize(x[l-1], x[l+1]) do t
            return sgn * resid(t)
        end
        x[l] = res.minimizer
        zmax = max(zmax, abs(resid(x[l])))
        zmin = min(zmin, abs(resid(x[l])))
    end
    sgn = -sgn
    res = optimize(x[n-1], Interval[2]) do t
        return sgn * resid(t)
    end
    x[n] = res.minimizer
    zmax = max(zmax, abs(resid(x[n])))
    zmin = min(zmin, abs(resid(x[n])))
    return zmin, zmax
end

end # module


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

@userplot HScatter

@recipe function f(hs::HScatter; lags=nothing, tol=1e-1, distance=Euclidean())
  # get inputs
  spatialdata = hs.args[1]
  var₁ = hs.args[2]
  var₂ = length(hs.args) == 3 ? hs.args[3] : var₁

  # lookup valid data
  X₁, z₁ = valid(spatialdata, var₁)
  X₂, z₂ = valid(spatialdata, var₂)

  # compute pairwise distance
  m, n = length(z₁), length(z₂)
  pairs = [(i,j) for j in 1:n for i in j:m]
  ds = [evaluate(distance, view(X₁,:,i), view(X₂,:,j)) for (i,j) in pairs]

  # use quartiles by default
  lags == nothing && (lags = quantile(ds, [.00,.25,.50,.75]))

  xguide --> var₁
  yguide --> var₂
  legend --> false
  aspect_ratio --> :equal
  layout --> (1, length(lags))

  for (i, lag) in enumerate(lags)
    # find indices with given lag
    match = findall(abs.(ds .- lag) .< tol)

    if isempty(match)
      @warn "no points were found with lag = $lag, skipping..."
      continue
    end

    # scatter plot coordinates
    mpairs = view(pairs, match)
    x = z₁[first.(mpairs)]
    y = z₂[last.(mpairs)]

    # plot identity line
    @series begin
      subplot := i
      seriestype := :path
      primary := false
      linestyle := :dash
      color := :black

      xmin, xmax = extrema(x)
      ymin, ymax = extrema(y)
      vmin = min(xmin, ymin)
      vmax = max(xmax, ymax)

      [vmin, vmax], [vmin, vmax]
    end

    # plot h-scatter
    @series begin
      subplot := i
      seriestype := :scatter
      title --> @sprintf "lag = %.1f, corr = %.2f" lag cor(x, y)

      x, y
    end
  end
end


#=

  Traffic lights problem in Julia ConstraintSolver.jl

  CSPLib problem 16
  http://www.cs.st-andrews.ac.uk/~ianm/CSPLib/prob/prob016/index.html
  """
  Specification:
  Consider a four way traffic junction with eight traffic lights. Four of the traffic
  lights are for the vehicles and can be represented by the variables V1 to V4 with domains
  {r,ry,g,y} (for red, red-yellow, green and yellow). The other four traffic lights are
  for the pedestrians and can be represented by the variables P1 to P4 with domains {r,g}.

  The constraints on these variables can be modelled by quaternary constraints on
  (Vi, Pi, Vj, Pj ) for 1<=i<=4, j=(1+i)mod 4 which allow just the tuples
  {(r,r,g,g), (ry,r,y,r), (g,g,r,r), (y,r,ry,r)}.

  It would be interesting to consider other types of junction (e.g. five roads
  intersecting) as well as modelling the evolution over time of the traffic light sequence.
  ...

  Results
  Only 2^2 out of the 2^12 possible assignments are solutions.

  (V1,P1,V2,P2,V3,P3,V4,P4) =
     {(r,r,g,g,r,r,g,g), (ry,r,y,r,ry,r,y,r), (g,g,r,r,g,g,r,r), (y,r,ry,r,y,r,ry,r)}
     [(1,1,3,3,1,1,3,3), ( 2,1,4,1, 2,1,4,1), (3,3,1,1,3,3,1,1), (4,1, 2,1,4,1, 2,1)}

  The problem has relative few constraints, but each is very tight. Local propagation
  appears to be rather ineffective on this problem.
  """

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK
const CS = ConstraintSolver
include("constraints_utils.jl")

function traffic_lights(print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS, # <-
                                                            # "traverse_strategy"=>:DBFS,

                                                            "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf, # <-

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>16,

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                        ))

    n=4
    r,ry,g,y = 1:n
    d = Dict(1=>"r",2=>"ry",3=>"g",4=>"y")
    # allowed light combinations
    allowed = [r r g g;
               ry r y r;
               g g r r;
               y r ry r
              ]

    @variable(model, v[1:n],  CS.Integers([r,ry,g,y]))
    @variable(model, p[1:n],  CS.Integers([r,ry,g,y]))

    for i in 1:n, j in 1:n
        if j == (1+i) % n
            @constraint(model, [v[i], p[i], v[j], p[j]] in CS.TableSet(allowed))
        end
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                vv = convert.(Integer,JuMP.value.(v; result=sol))
                pp = convert.(Integer,JuMP.value.(p; result=sol))
                # println("v:$vv  p:$pp")
                for i in 1:n
                print(d[vv[i]]," ", d[pp[i]], " ")
                end
                println()
            end
            println()
        end
    else
        println("status:$status")
    end

    return status
end

@time traffic_lights()


module Regress
using StatsBase

"""
Impute missing values in a column, using the other columns.

# Arguments
- Y: matrix (with preimputed values)
- M: missingness indictor (1 for missing, 0 for observed)
- k: column to impute
- thresh: In regression, only train model with responses less than this value.
"""
function imputecolumn(Y, M, k; thresh=0)
  idx_mis = findall(M[:, k])
  idx_obs = findall((.!M[:, k]) .& (Y[:, k] .< thresh))
  coef = Y[idx_obs, 1:end .!= k] \ Y[idx_obs, k]
  pred = Y[idx_mis, 1:end .!= k] * coef
  return pred
end

"""
Iteratively impute missing values in matrix, using the other columns.

# Arguments
- Y: matrix (with missing values)
- maxiter: Maximum number of iterations to impute missing values.
- tol: if difference between current and next matrices are less than this,
  declare convergence!
- init: initialize missing values at this.
- verbose: if value > 0, print some information about the algorithm.
- thresh: See `imputecolumn`.
"""
function impute(Y; maxiter=30, tol=1e-3, init=missing, verbose=1, thresh=0)
  @assert ndims(Y) == 2

  X = deepcopy(Y)
  M = isnan.(X)
  K = size(X, 2)

  # Initialize at `init` if provided.  Otherwise, initialize at the mean of the
  # negative observed values.
  X[M] .= ismissing(init) ? mean(X[X .< 0]) : init

  # Difference of X between iterations
  diff_X = Float64[]

  for i in 1:maxiter
    X_old = deepcopy(X)
    if verbose > 0
      print("\r$(i)/$(maxiter)")
    end

    for k in 1:K
      X[M[:, k], k] .= imputecolumn(X, M, k, thresh=thresh)
    end

    append!(diff_X, mean(abs.(X_old[M] - X[M])))

    if diff_X[end] < tol
      if verbose > 0
        println()
        println("Convergence defected! Stopping early.")
      end
      break
    end
  end

  if verbose > 0
    println()
  end

  return X, diff_X
end

end  # module Regress

#= Demo
using CSV
using Statistics
using DataFrames
using PyPlot; const plt = PyPlot.plt
using Seaborn; const sns = Seaborn

Y = coalesce.(CSV.read("../../runs/data/cb_transformed.csv"), NaN)
Y_mat = Matrix(Y[Y[!, :sample_id] .== 1, Not(:sample_id)])

Y_complete, diff_Y = Regress.impute(Y_mat, maxiter=50)

plt.plot(diff_Y, marker=:o)

function plotme(j; magthresh=7, xlim=nothing)
  yj = Y_complete[isnan.(Y_mat[:, j]), j]
  # plt.hist(yj[abs.(yj) .< magthresh, :], bins=50, alpha=.5)
  sns.kdeplot(vec(yj[abs.(yj) .< magthresh, :]), alpha=.5, bw=.2)
  axvline(0, lw=2)
  println("Proportion of missing: $(mean(isnan.(Y_mat[:, j])))")
  if xlim != nothing
    plt.xlim(xlim)
  end
  return
end

for j in 1:32
  plotme(j, xlim=(-7, 3));  # 16 is interesting
end
plt.close();

for j in 1:32
  sns.kdeplot(Y_complete[:, j], alpha=.5)
  sleep(.5)
end
plt.xlim((-8, 6))
plt.axvline(0)
plt.close();
=#


"""
   solving ||Ax-b|| via CGLS, This code is a translated version of SOL's matlab code
   the resNE (normalized relative error is ||A'b-A'Ax||/||A'b|| )
"""
function cgls(A, b::Vector{Tv}; params=Dict(),
              x0=0.0, shift=0.0, tol=1e-6, maxIter=30, print_flag=false) where {Tv<:Number}

    # determine the type of A
    if isa(A, Matrix)
       explicitA = true
    elseif isa(A, Function)
       explicitA = false
    else
       error("A must be a matrix or a function")
    end

    # initial guess of the unknowns
    if explicitA

       (m, n) = size(A)

       # provided initial guess
       if x0 != 0.0
          x = copy(x0)
          r = b - A * x
          s = A' * r - shift * x
       # zero vector as initial guess
       else
          x = zeros(eltype(b), n)
          r = copy(b)
          s = A' * r
       end

    else

       m = length(b)

       # provided initial guess
       if x0 != 0.0
          x = copy(x0)
          r = b - A(x, 1; params...)
          s = A(r, 2; params...) - shift * x
          n = length(s)
       # zero vector as initial guess
       else
          r = copy(b)
          s = A(b, 2; params...)
          n = length(s)
          x = zeros(eltype(b), n)
       end

    end

    data_fitting = dot(r,r)
    constraint   = 0.0
    cost0 = data_fitting + constraint

    # initialize some intermidiate vectors
    p = copy(s)
    norms0= vecnorm(s)
    gamma = norms0^2
    normx = vecnorm(x)
    xmax  = normx     # keep the initial one
    resNE = 1.0

    gamma0= copy(gamma)
    delta = 0.0

    # iteration counter and stop condition
    k = 0
    run_flag = true

    if print_flag
       header = "  k         data_fitting           constraint             normx                resNE"
       println(""); println(header);
       @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
    end

    while (k < maxIter) && run_flag

          k = k + 1

          if explicitA
             q = A * p
          else
             q = A(p, 1; params...)
          end

          delta = (vecnorm(q))^2 + shift * (vecnorm(p))^2
          indefinite = delta <= 0 ? true : false
          delta      = delta == 0.? eps(): delta

          alpha = gamma / delta

          x = x + alpha * p
          r = r - alpha * q

          data_fitting = dot(r,r)
          constraint   = shift * dot(x, x)
          cost = data_fitting + constraint

          if explicitA
             s = A' * r - shift * x
          else
             s = A(r, 2; params...) - shift * x
          end

          norms  = norm(s)
          gamma0 = copy(gamma)
          gamma  = norms^2
          beta   = gamma / gamma0

          p = s + beta * p

          # check the stopping crietia
          normx = norm(x)
          xmax  = normx > xmax ? normx : xmax
          if norms <= norms0 * tol || normx * tol >= 1.0
             run_flag = false
          end

          # print information
          # resNE = norms / norms0
          resNE = cost  / cost0
          if print_flag
             @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
          end

    end

    return x, resNE

end

# ==============================================================================
#                 CGLS based on hard drive
# ==============================================================================
"""
   solving ||Ax-b|| via CGLS
"""
function cgls(A::Function, b::Ts; x0="NULL", path="NULL", params=[],
         shift=0.0, tol=1e-6, maxIter=50, print_flag=false) where {Ts<:Union{String, Vector{String}}}

    if path == "NULL"
       path = pwd()
    end

    # create folder to save files of model size
    path_model = joinpath(path, "model")
    if !isdir(path_model)
       mkdir(path_model)
       if !isdir(path_model)
          error("could not create directory")
       end
    end
    x = join([path_model "/x.bin"])
    s = join([path_model "/s.bin"])
    p = join([path_model "/p.bin"])

    # create folders to save residue
    path_residue   = joinpath(path, "residue")
    if !isdir(path_residue)
       mkdir(path_residue)
       if !isdir(path_residue)
          error("could not create directory")
       end
    end

    # create folders to save synthetic data with current model
    path_forward = joinpath(path, "forward")
    if !isdir(path_forward)
       mkdir(path_forward)
       if !isdir(path_forward)
          error("could not create directory")
       end
    end

    # intermidiate variables
    if typeof(b) == String
       r = join([path_residue "/res.bin"])
       q = join([path_forward "/fwd.bin"])
    elseif typeof(b) == Vector{String}
       n = length(b)
       r = Vector{String}(n)
       q = Vector{String}(n)
       for i = 1 : n
           r[i] = join([path_residue "/res" "_" "$i" ".bin"])
           q[i] = join([path_forward "/fwd" "_" "$i" ".bin"])
       end
    end

    # provided initial guess
    if x0 != "NULL"

       cp(x0, x, remove_destination=true)
       dcal = A(x, 1; params=params)
       x_plus_alpha_q!(r, b, -1.0, dcal)

       if typeof(params) <: Dict{Symbol, Any}
          params[:path_m] = s
       elseif typeof(params) == Vector{Dict}
         for i = 1 : length(params)
             params[i][:path_m] = s
         end
       end
       s = A(r, 2; params=params)
       x_plus_alpha_q!(s, -shift, x)

    # zero vector as initial guess
    else

       cp(b, r, remove_destination=true)
       if typeof(params) <: Dict{Symbol, Any}
          params[:path_m] = s
       elseif typeof(params) == Vector{Dict}
          for i = 1 : length(params)
              params[i][:path_m] = s
          end
       end
       s = A(r, 2; params=params)

       # create x
       (hdr, ds) = read_USdata(s)
       tmp = zeros(ds)
       write_USdata(x, hdr, tmp)
    end

    # compute residue
    data_fitting = (norm(r))^2
    constraint   = 0.0
    cost0 = data_fitting + constraint
    convergence = Float64[]; push!(convergence, 1.0);

    # initialize some intermidiate vectors
    cp(s, p, remove_destination=true)

    norms0= norm(s)
    gamma = norms0^2
    normx = norm(x)
    xmax  = normx     # keep the initial one
    resNE = 1.0

    gamma0= copy(gamma)
    delta = 0.0

    # iteration counter and stop condition
    k = 0
    run_flag = true

    if print_flag
       header = "  k         data_fitting           constraint             normx                resNE"
       println(""); println(header);
       @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
    end

    if typeof(params) <: Dict{Symbol, Any}
       params[:path_fwd] = q
    elseif typeof(params) == Vector{Dict}
       for i = 1 : length(params)
           params[i][:path_fwd] = q[i]
       end
    end

    while (k < maxIter) && run_flag

          k = k + 1

          q = A(p, 1; params=params)

          delta = (norm(q))^2 + shift * (norm(p))^2
          indefinite = delta <= 0 ? true : false
          delta      = delta == 0.? eps(): delta

          alpha = gamma / delta

          x_plus_alpha_q!(x,  alpha, p)
          x_plus_alpha_q!(r, -alpha, q)

          data_fitting = (norm(r))^2
          constraint   = shift * (norm(x))^2
          cost = data_fitting + constraint

          # save the intermidiate result
          path_iter = join([path_model "/iteration" "_" "$k" ".bin"])
          cp(x, path_iter, remove_destination=true)

          s = A(r, 2; params=params)
          x_plus_alpha_q!(s, -shift, x)

          norms  = norm(s)
          gamma0 = gamma
          gamma  = norms^2
          beta   = gamma / gamma0

          (hdr, ds) = read_USdata(s)
          (hdr, dp) = read_USdata(p)
          dp = ds + beta * dp
          write_USdata(p, hdr, dp)

          # check the stopping crietia
          normx = norm(x)
          xmax  = normx > xmax ? normx : xmax
          if norms <= norms0 * tol || normx * tol >= 1.0
             run_flag = false
          end

          # print information
          resNE = cost / cost0
          if print_flag
             @printf("%3.0f %20.10e %20.10e %20.10e %20.10e\n", k, data_fitting, constraint, normx, resNE);
          end
          push!(convergence, resNE);
    end

    return x, convergence

end


using Test

@testset "Hyperslice solver" begin include("TestSolver.jl") end
@testset "Simplex/point intersection" begin include("TestSimplexPointIntersection.jl") end
@testset "Additional focus point sampling" begin include("TestSampling.jl") end
@testset "Utilities" begin include("TestUtil.jl") end
@testset "Top-level hyperslice function" begin include("TestHyperslice.jl") end


immutable Cauchy <: ContinuousUnivariateDistribution
    location::Float64
    scale::Float64
    function Cauchy(l::Real, s::Real)
        s > zero(s) || error("scale must be positive")
	new(float64(l), float64(s))
    end
end

Cauchy(l::Real) = Cauchy(l, 1.0)
Cauchy() = Cauchy(0.0, 1.0)

## Support
@continuous_distr_support Cauchy -Inf Inf

## Properties
mean(d::Cauchy) = NaN
median(d::Cauchy) = d.location
mode(d::Cauchy) = d.location

var(d::Cauchy) = NaN
skewness(d::Cauchy) = NaN
kurtosis(d::Cauchy) = NaN

entropy(d::Cauchy) = log(d.scale) + log(4.0 * pi)

## Functions
pdf(d::Cauchy, x::Real) = 1/(pi*d.scale*(1+((x-d.location)/d.scale)^2))
logpdf(d::Cauchy, x::Real) = -log(pi) - log(d.scale) - log1psq((x-d.location)/d.scale)

cdf(d::Cauchy, x::Real) = atan2(one(x),-(x-d.location)/d.scale)/pi
ccdf(d::Cauchy, x::Real) = atan2(one(x),(x-d.location)/d.scale)/pi
logcdf(d::Cauchy, x::Real) = x <= d.location ? log(cdf(d,x)) : log1p(-ccdf(d,x))
logccdf(d::Cauchy, x::Real) = x <= d.location ? log1p(-cdf(d,x)) : log(ccdf(d,x))

quantile(d::Cauchy, p::Real) = @checkquantile p d.location - d.scale*cospi(p)/sinpi(p)
cquantile(d::Cauchy, p::Real) = @checkquantile p d.location + d.scale*cospi(p)/sinpi(p)
invlogcdf(d::Cauchy, lp::Real) = lp < loghalf ? quantile(d,exp(lp)) : cquantile(d,-expm1(lp))
invlogccdf(d::Cauchy, lp::Real) = lp < loghalf ? cquantile(d,exp(lp)) : quantile(d,-expm1(lp))

mgf(d::Cauchy, t::Real) = t == zero(t) ? 1.0 : NaN
cf(d::Cauchy, t::Real) = exp(im * t * d.location - d.scale * abs(t))

## Fitting
# Note: this is not a Maximum Likelihood estimator
function fit{T <: Real}(::Type{Cauchy}, x::Array{T})
    l, u = iqr(x)
    Cauchy(median(x), (u - l) / 2.0)
end


using BoundaryValueDiffEq
using DiffEqBase, DiffEqDevTools
using Test

# First order test
function func_1!(du,u,p,t)
    du[1] = u[2]
    du[2] = 0
end

# Not able to change the initial condition.
# Hard coded solution.
func_1 = ODEFunction(func_1!, analytic=(u0, p, t) -> [5-t,-1])

function boundary!(residual, u, p, t)
    residual[1] = u[1][1]-5
    residual[2] = u[end][1]
end

function boundary_two_point!(residual, u, p, t)
    ua = u[1]
    ub = u[end]
    residual[1] = ua[1]-5
    residual[2] = ub[1]
end

# Second order linear test
function func_2!(du, u, p, t)
    du[1] = u[2]
    du[2] = -u[1]
end

# Not able to change the initial condition.
# Hard coded solution.
func_2!(::Type{Val{:analytic}}, u0, p, t) = [5*(cos(t) - cot(5)*sin(t)),
                                             5*(-cos(t)*cot(5) - sin(t))]
func_2 = ODEFunction(func_2!, analytic=(u0, p, t) -> [5*(cos(t) - cot(5)*sin(t)),
                                                      5*(-cos(t)*cot(5) - sin(t))])
tspan = (0.,5.)
u0 = [5.,-3.5]
probArr = [BVProblem(func_1, boundary!, u0, tspan),
           BVProblem(func_2, boundary!, u0, tspan),
           TwoPointBVProblem(func_1, boundary_two_point!, u0, tspan),
           TwoPointBVProblem(func_2, boundary_two_point!, u0, tspan)]

testTol = 0.2
affineTol = 1e-2
dts = (1/2) .^ (5:-1:1)
order = 4

println("Collocation method (GeneralMIRK)")
println("Affineness Test")
prob = probArr[1]
@time sol = solve(prob, GeneralMIRK4(), dt=0.2)
@test norm(diff(map(x->x[1], sol.u)) .+ 0.2, Inf) + abs(sol[1][1]-5) < affineTol

println("Convergence Test on Linear")
prob = probArr[2]
@time sim = test_convergence(dts,prob,GeneralMIRK4())
@test abs(sim.𝒪est[:final]-order) < testTol

println("Collocation method (MIRK)")
println("Affineness Test")
prob = probArr[3]
@time sol = solve(prob, MIRK4(), dt=0.2)
@test norm(diff(map(x->x[1], sol.u)) .+ 0.2, Inf) .+ abs(sol[1][1]-5) < affineTol

println("Convergence Test on Linear")
prob = probArr[4]
@time sim = test_convergence(dts,prob,MIRK4())
@test abs(sim.𝒪est[:final]-order) < testTol


using FixedSizeArrays

type Heightfield
    w::Float64
    h::Float64
    y::Matrix{Float64}
end

Heightfield(w::Float64,h::Float64, nx::Integer, nz::Integer) =
    Heightfield(w, h, zeros(Float64, nx, nz))

Heightfield(hf::Heightfield) = Heightfield(hf.w, hf.h, copy(hf.y))

function triangle_list(hf::Heightfield)
    nx,nz = size(hf.y)
    x = linspace(-1,1,nx)*hf.w/2
    z = linspace(-1,1,nz)*hf.h/2
    dx = hf.w/(nx-1)
    dz = hf.h/(nz-1)

    vertices = Vector{Point{3,Float64}}()
    for xi in 1:nx
        for zi in 1:nz
            push!(vertices, Point(x[xi], hf.y[xi,zi], z[zi]))
        end
    end

    Dx = Tridiagonal(-ones(nx-1),zeros(nx),ones(nx-1))
    Dz = Tridiagonal(-ones(nz-1),zeros(nz),ones(nz-1))

    dydx = Dx*hf.y
    dydz = (Dz*hf.y')'

    normals = Vector{Vec{3,Float64}}()
    for xi in 1:nx
        for zi in 1:nz
            a = Vec(2dx, dydx[xi,zi], 0.)
            b = Vec(0., dydz[xi,zi], 2dz)
            push!(normals, -normalize(cross(a,b)))
        end
    end

    faces = Vector{Tuple}()
    for xi in 0:nx-2
        for zi in 1:nz-1
            push!(faces, (zi + xi*nz, zi + (xi+1)*nz, zi + 1 + (xi+1)*nz))
            push!(faces, (zi + xi*nz, zi + 1 + (xi+1)*nz, zi + 1 + xi*nz))
        end
    end

    vertices, normals, faces
end

export Heightfield, view, triangle_list


using Plots, LaTeXStrings, Statistics, ProgressBars, JLD
cd(dirname(@__FILE__))
Plots.default(titlefontsize=12, tickfontsize=10, labelfontsize=10,
    fontfamily="Computer Modern")

function Deposition(; len, tot_time, dep_rate, color_step)
    surface = InitialSurf(len)
    AllLens = []
    for n in 0:tot_time
        randsurf = rand(1:len, dep_rate)
        for i in randsurf
            surface = check_and_depose(surface, i, len, (n%color_step) + 1)
        end
        append!(AllLens, Check_Dlen(surface))
    end
    return AllLens
end

function Check_Dlen(surf)
    Dlen = []
    for row in surf
        yindex = []
        for num in 1:length(row)
            if row[num] != 0
                append!(yindex, num)
            end
        end
        append!(Dlen, max(yindex...) - min(yindex...))
    end
    return max(Dlen...)
end

function check_and_depose(surf, index, len, addval) #checking down and sides and first layer
    i1 , i2 = sides(index, len)
    if surf[end][index] != 0
        newsurf = [0 for j = 1:len]
        newsurf[index] = addval
        push!(surf, newsurf)
    end
    for num in 0:length(surf) -1
        if surf[end - num][index] == 0 && (surf[end - num][i1] != 0 || surf[end - num][i2] != 0)
            surf[end - num][index] = addval
            break
        elseif length(surf) > num + 1 && surf[end - 1 - num][index] != 0
            surf[end - num][index] = addval
            break
        end
    end
    return surf
end

function sides(n, L)
    if n == L
        return n-1 , 1
    elseif n == 1
        return L , n+1
    else
        return n-1, n+1
    end
end

function InitialSurf(len)
    pos1 = [0 for i=1:len]
    pos1[floor(Int, len/2)] = 1
    return [pos1]
end

function Linear_fit(;Time, AvgList, steps)
    A = [hcat(Time) reshape(ones(steps), steps, 1)]
    b = reshape(AvgList, steps, 1)
    line = (A \ b)
    return line
end

Parameters = Dict(:len => 200,
                    :tot_time => 403,
                        :dep_rate => 100,
                            :color_step => 100)

up = 5.35; low = 1; steps = 100; SelctedSpan = 56:95; SelectedSteps = length(SelctedSpan);
Time = round.(Int,exp.(low:(up-low)/(steps-1):up))[SelctedSpan]
RawTime = hcat(low:(up-low)/(steps-1):up)[SelctedSpan]

iternum = 1000
allVal = [[0.0 for i ∈ 1:SelectedSteps] for j ∈ 1:iternum]
meanVal = [0.0 for i ∈ 1:SelectedSteps]
vars = [0.0 for i ∈ 1:SelectedSteps]
for i ∈ ProgressBar(1:iternum)
    VarList = Deposition(;Parameters...)[Time]
    allVal[i] = VarList
    meanVal += VarList
end
meanVal /= iternum
for i ∈ 1:SelectedSteps
    vars[i] = std(hcat(allVal...)[i,:])
end

Paraline = Dict(:Time => log.(Time),
                :AvgList => log.(meanVal),
                :steps => SelectedSteps)
Line = Linear_fit(;Paraline...)
X = RawTime
Y = X .* Line[1] .+ Line[2]

plot(X,Y,c = :black,label = L"y = %$(round(Line[1],digits= 3))x + %$(round(Line[2],digits= 3))")
scatter!(RawTime, log.(meanVal),
    c = :steelblue,
    xlabel= L"\mathrm{ln}(Time)",
    ylabel= L"\mathrm{ln}({L_C}_{(t)})",
    title= L"Log-Log\ Plot\ of\ L_C\ (L = %$(Parameters[:len]),\ %$iternum\ runs, %$(Time[end]*Parameters[:dep_rate])\ Particles)",
    label = L"Data\ point",
    ribbon = log.(vars),
    yerror = log.(vars),
    legend = 150)

savefig("../../Figs/Q3/Q3-L(t).pdf")
save("../../Data/Q3/Q3-TotData.jld","Data", allVal)


immutable Chi <: ContinuousUnivariateDistribution
    df::Float64
    function Chi(df::Real)
        df > zero(df) || error("df must be positive")
        new(float64(df))
    end
end

@continuous_distr_support Chi 0.0 Inf

cdf(d::Chi, x::Real) = cdf(Chisq(d.df),x^2)
ccdf(d::Chi, x::Real) = ccdf(Chisq(d.df),x^2)
logcdf(d::Chi, x::Real) = logcdf(Chisq(d.df),x^2)
logccdf(d::Chi, x::Real) = logccdf(Chisq(d.df),x^2)

quantile(d::Chi,p::Real) = sqrt(quantile(Chisq(d.df),p))
cquantile(d::Chi,p::Real) = sqrt(cquantile(Chisq(d.df),p))
invlogcdf(d::Chi,p::Real) = sqrt(invlogcdf(Chisq(d.df),p))
invlogccdf(d::Chi,p::Real) = sqrt(invlogccdf(Chisq(d.df),p))

mean(d::Chi) = √2 * gamma((d.df + 1.0) / 2.0) / gamma(d.df / 2.0)

function mode(d::Chi)
    d.df >= 1.0 || error("Chi distribution has no mode when df < 1")
    sqrt(d.df - 1)
end

var(d::Chi) = d.df - mean(d)^2

function skewness(d::Chi)
    μ, σ = mean(d), std(d)
    (μ / σ^3) * (1.0 - 2.0 * σ^2)
end

function kurtosis(d::Chi)
    μ, σ, γ = mean(d), std(d), skewness(d)
    (2.0 / σ^2) * (1 - μ * σ * γ - σ^2)
end

function pdf(d::Chi, x::Real)
    k = d.df
    (2.0^(1.0 - k / 2.0) * x^(k - 1.0) * exp(-x^2 / 2.0)) / gamma(k / 2.0)
end

function entropy(d::Chi)
    k = d.df
    lgamma(k / 2.0) - log(sqrt(2.0)) -
        ((k - 1.0) / 2.0) * digamma(k / 2.0) + k / 2.0
end

function gradloglik(d::Chi, x::Float64)
  insupport(Chi, x) ? (d.df - 1.0) / x - x : 0.0
end

rand(d::Chi) = sqrt(rand(Chisq(d.df)))


########################### Submodular Minimization ############################
# greedy algorithm for submodular minimization, or maximization??
# WARNING: before using, check min / max
struct Submodular{T, F, X, D} <: Update{T}
    f::F # set-valued input function
    x::X # set to choose from
    δ::D # bound on improvement of adding a certain element
    # function Submodular(f, x::Union{AbstractVector, AbstractSet})
    #     δ = fill(-Inf, length(x))
    #     new{eltype(x), typeof(f), typeof(x), typeof(δ)}(f, x, δ)
    # end
end
function Submodular(f, x::Union{AbstractVector, AbstractSet}, n::Int)
    δ = fill(-Inf, length(x))
    Submodular(f, x, δ)
end

# TODO: parallelization, especially for first pass
function update!(S::Submodular, x::Union{AbstractVector, AbstractSet})
    # length(x) ≥ S.n ? (return x) : nothing # only add if we haven't hit maximum cardinality
    f, δ, grid = S.f, S.δ, S.x
    ind = sortperm(δ) # sort grid according to increasing δ (decreasing in magnitude)
    permute!(δ, ind)
    permute!(grid, ind)
    fx = f(x)
    min_i = 0. # keeps track of smallest (i.e. greatest magnitude of) δ
    for (i, xs) in enumerate(grid) # parallelizing here does not seem to benefit much
        if min_i > δ[i] # if previously recorded δ is smaller (larger in magnitude) than the current minimum
            push!(x, xs) # already minimized allocation for (x, xs)?
            δ[i] = f(x) - fx
            min_i = min(min_i, δ[i])
            pop!(x)
        end
    end
    push!(x, grid[argmin(δ)]) # choose next point greedily
end

# Old code:
# function submodular_minimization(x::AbstractArray{T}, grid::AbstractArray{T},
#                             objective::Function,
#                             num_next::Int = 1) where {T}
#     num_next < 1 && error("num_next < 1")
#
#     δ = fill(-Inf, length(grid))
#     base_objective = objective(x)
#     #
#     Threads.@threads for i in 1:length(grid)
#         δ[i] = objective(vcat(x, grid[i])) - base_objective
#     end
#     next = grid[argmin(δ)] # choose point which reduces the objective most
#
#     if num_next == 1
#         return [next], δ
#     else
#         recursive_next, δ = submodular_minimization(vcat(x, next), grid,
#                                                 objective, num_next-1, δ)
#         return vcat(next, recursive_next), δ # could instead return, and spawn task to complete the recursive step
#     end
# end
#
# # posterior uncertainty sampling for batch learning of optical data
# function submodular_minimization(x::AbstractArray{T}, grid::AbstractArray{T},
#                                 objective::Function, num_next::Int,
#                                 δ::Array{<:Real}) where {T}
#     length(δ) ≠ length(grid) && error("length(δ) ≠ length(grid)")
#     num_next < 1 && error("num_next < 1")
#     # sort grid according to increasing δ (decreasing in magnitude)
#     ind = sortperm(δ)
#     δ = δ[ind]
#     grid = grid[ind]
#
#     base_objective = objective(x)
#     min_i = 0 # keeps track of smallest (i.e. greatest magnitude of) δ
#     for (i, xs) in enumerate(grid) # parallelizing here does not seem to benefit much
#         if min_i > δ[i] # if previously recorded δ is smaller (larger in magnitude) than the current minimum
#             δ[i] = objective(vcat(x, xs)) - base_objective
#             min_i = min(min_i, δ[i])
#         end
#     end
#     next = grid[argmin(δ)] # choose the point which reduces posterior uncertainty most
#
#     if num_next == 1
#         return [next], δ
#     else
#         recursive_next, δ = submodular_minimization(vcat(x, next), grid,
#                                                     objective, num_next-1, δ)
#         return vcat(next, recursive_next), δ
#     end
# end


 n=rand(1:10)
 @test diag(Matrix(Toeplitz(collect(-n:n)))) == zeros(n+1)