using Pumas, LinearAlgebra, Plots, Query

iv = @model begin
  @param   begin
    tvcl ∈ RealDomain(lower=0)
    tvv ∈ RealDomain(lower=0)
    Ω ∈ PDiagDomain(2)
    σ_prop ∈ RealDomain(lower=0)
  end

  @random begin
    η ~ MvNormal(Ω)
  end

  @pre begin
    CL = tvcl * (Wt/70)^0.75 * exp(η[1])
    V  = tvv * (Wt/70) * exp(η[2])
  end

  @covariates Wt

  @dynamics ImmediateAbsorptionModel

  @derived begin
      cp := @. (Central / V)
      dv ~ @. Normal(cp, sqrt(cp^2*σ_prop))
    end
end

ev = DosageRegimen(2000, time=0)
choose_covariates() = (wt = rand(55:80))
pop =  Population(map(i -> Subject(id=i, evs=ev, cvs = (Wt = choose_covariates(),)),1:24))
param = (
  tvcl = 11.5,
  tvv  = 50,
  Ω = Diagonal([0.04,0.04]),
  σ_prop = 0.001
  )

sd_obstimes = [0, 0.25, 0.5, 0.75, 1, 2, 4, 8,
                    12, 16, 20, 21, 24]
obs = simobs(iv, pop, param, obstimes=sd_obstimes)
iv_sim = DataFrame(obs)|>
          @filter(!(_.time == 0 && _.evid == 0)) |>
          @mutate(cmt = 1) |>
          @mutate(route = "iv") |>
          DataFrame

timeu = u"hr"
concu = u"mg/L"
amtu  = u"mg"

iv_nca = read_nca(iv_sim, id=:id, time=:time, conc=:dv, amt=:amt,
    route=:route,timeu=timeu, concu=concu, amtu=amtu,llq=0.4concu)

NCA.auc(iv_nca)
NCA.auc_extrap_percent(iv_nca)
report = NCAReport(iv_nca)

iv_sim_7percBLQ = iv_sim |>
                  @filter()
sd_obstimes_7percBLQ = [0, 0.25, 0.5, 0.75, 1, 2, 4, 8, 12, 16, 20]
sd_obstimes_18percBLQ = [0, 0.25, 0.5, 0.75, 1, 2, 4, 8, 12, 16, 20, 21, 24]

plot(obs,  yscale = :log10)


module StochasticVolatility

using ArgCheck
using Distributions
using Parameters
using DynamicHMC
using StatsBase
using Base.Test
using ContinuousTransformations

export StochasticVolatility

"""
    simulate_stochastic(ρ, σ_v, ϵs, νs)

Take in the parameter values (ρ, σ) for the latent volatility process, the errors ϵs used for the measurement equation and the errors νs used for the transition equation.

The discrete-time version of the Ornstein-Ulenbeck Stochastic - volatility model:

    y_t = x_t + ϵ_t where ϵ_t ∼ χ^2(1)
    x_t = ρ * x_(t-1) + σ * ν_t  where ν_t ∼ N(0, 1)

"""
function simulate_stochastic(ρ, σ, ϵs, νs)
    N = length(ϵs)
    @argcheck N == length(νs)
    x₀ = νs[1]*σ*(1 - ρ^2)^(-0.5)
    xs = Vector{typeof(x₀)}(N)
    for i in 1:N
        xs[i] = (i == 1) ? x₀ : (ρ*xs[i-1] + σ*νs[i])
    end
    xs + log.(ϵs) + 1.27
end

simulate_stochastic(ρ, σ, N) = simulate_stochastic(ρ, σ, rand(Chisq(1), N), randn(N))

struct StochasticVolatility{T, Prior_ρ, Prior_σ, Ttrans}
    "observed data"
    ys::Vector{T}
    "prior for ρ (persistence)"
    prior_ρ::Prior_ρ
    "prior for σ_v (volatility of volatility)"
    prior_σ::Prior_σ
    "χ^2 draws for simulation"
    ϵ::Vector{T}
    "Normal(0,1) draws for simulation"
    ν::Vector{T}
    "Transformations cached"
    transformation::Ttrans
end

## In this form, we use an AR(2) process of the first differences with an intercept as the auxiliary model.

"""
    lag(xs, n, K)

Lag-`n` operator on vector `xs` (maximum `K` lags).
"""
lag(xs, n, K) = xs[((K+1):end)-n]

"""
    lag_matrix(xs, ns, K = maximum(ns))

Matrix with differently lagged xs.
"""
function lag_matrix(xs, ns, K = maximum(ns))
    M = Matrix{eltype(xs)}(length(xs)-K, maximum(ns))
    for i ∈ ns
        M[:, i] = lag(xs, i, K)
    end
    M
end

"first auxiliary regression y, X, meant to capture first differences"
function yX1(zs, K)
    Δs = diff(zs)
    lag(Δs, 0, K), hcat(lag_matrix(Δs, 1:K, K), ones(eltype(zs), length(Δs)-K), lag(zs, 1, K+1))
end

"second auxiliary regression y, X, meant to capture levels"
function yX2(zs, K)
    lag(zs, 0, K), hcat(ones(eltype(zs), length(zs)-K), lag_matrix(zs, 1:K, K))
end

"Constructor which calculates cached values and buffers. Use this unless you know what you are doing."
ℝto(dist::Uniform) = transformation_to(Segment(minimum(dist), maximum(dist)), LOGISTIC)
ℝto(::InverseGamma) = transformation_to(ℝ⁺)

"""
    StochasticVolatility(ys, prior_ρ, prior_σ, M)

Convenience constructor for StochasticVolatility.
Take in the observed data, the priors, and number of simulations (M).
"""

function StochasticVolatility(ys, prior_ρ, prior_σ, M)
    transformation = TransformationTuple((ℝto(prior_ρ), (ℝto(prior_σ))))
    StochasticVolatility(ys, prior_ρ, prior_σ, rand(Chisq(1), M), randn(M), transformation)
end

function (pp::StochasticVolatility)(θ)
    @unpack ys, prior_ρ, prior_σ, ν, ϵ, transformation = pp
    ρ, σ = transformation(θ)
    logprior = logpdf(prior_ρ, ρ) + logpdf(prior_σ, σ)
    N = length(ϵ)

    # Generating xs, which is the latent volatility process

    xs = simulate_stochastic(ρ, σ, ϵ, ν)
    Y_1, X_1 = yX1(xs, 2)
    β₁ = qrfact(X_1, Val{true}) \ Y_1
    v₁ = mean(abs2,  Y_1 - X_1*β₁)
    Y_2, X_2 = yX2(xs, 2)
    β₂ = qrfact(X_2, Val{true}) \ Y_2
    v₂ = mean(abs2,  Y_2 - X_2*β₂)
    # We work with first differences
    y₁, X₁ = yX1(ys, 2)
    log_likelihood1 = sum(logpdf.(Normal(0, √v₁), y₁ - X₁ * β₁))
    y₂, X₂ = yX2(ys, 2)
    log_likelihood2 = sum(logpdf.(Normal(0, √v₂), y₂ - X₂ * β₂))
    logprior + log_likelihood1 + log_likelihood2
end

"""

    path(parts...)

`parts...` appended to the directory containing `src/` etc. Useful for saving graphs and data.

"""

path(parts...) = normpath(joinpath(splitdir(Base.find_in_path("StochasticVolatility"))[1], "..", parts...))

end


using LevelSetShapes
using AbbreviatedTypes
using LinearAlgebra
using Random: MersenneTwister
using Statistics: mean
using Test

include("simple.jl")


@doc """
Symbolic Cholesky factorization.  The fill-reducing permutation
is calculated externally, for example using Metis.NodeND
"""->
type CholSymb{T<:Integer}    # symbolic Cholesky
    pinv::Vector{T}          # fill reducing perm for Chol
    cp::Vector{T}            # column pointers
    parent::Vector{T}        # elimination tree
end

@doc """
Find the nonzero pattern of Cholesky L(k,1:k-1) using etree and triu(A(:,k))
"""->
function ereach{Tv,Ti}(A::Symmetric{Tv,SparseMatrixCSC{Tv,Ti}}, k, parent)
    if A.uplo != 'U'
        error("ereach requires the upper triangle of A to be stored, i.e. A.uplo == 'U'")
    end
    k = convert(Ti,k)
    Ad = A.data
    Ai = rowvals(Ad)
    s = IntSet([k])                     # non-zeros in row k of L
    v = IntSet([k])                     # visited nodes
    for i in sub(Ai,nzrange(Ad,k))
        if i > k
            continue                    # use only upper triangle
        end
        while !in(i,v)                  # traverse up etree
            push!(s,i)
            push!(v,i)
            i = parent[i]
        end
    end
    collect(s)
end

# find nonzero pattern of Cholesky L[k,1:k-1] using etree and triu(A[:,k])
# based on cs_ereach p. 43, "Direct Methods for Sparse Linear Systems"
function ereach{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}, k::Integer, parent::Vector{Ti})
    m,n = size(A); Ap = A.colptr; Ai = A.rowval
    s = Ti[]; sizehint!(s, n)            # to be used as a stack
    visited = falses(n)
    visited[k] = true
    for p in Ap[k]:(Ap[k+1] - 1)
        i = Ai[p]                # A[i,k] is nonzero
        if i > k continue end    # only use upper triangular part of A
        while !visited[i]        # traverse up etree
            push!(s,i)           # L[k,i] is nonzero
            visited[i] = true
            i = parent[i]
        end
    end
    s
end

@doc """
Determine the pattern of the sparse Cholesky factor of C

Returns the elimination tree and a lower triangular sparse matrix filled with ones.
"""->
function cholpattern{Tv,Ti}(C::Symmetric{Tv,SparseMatrixCSC{Tv,Ti}})
    C.uplo == 'U' || error("cholpattern requires the upper triangle of C to be stored")
    n = size(C,2)
    parent = etree(C)
    Cd = C.data
    Ci = rowvals(Cd)
    cp = ones(Ti,n+1)
    rv = Ti[]
    s = sizehint!(IntSet([]),n) # nonzeros on column j of L' (row j of L)
    v = sizehint!(IntSet([]),n) # visited nodes
    for j in 1:n
        empty!(s)
        empty!(v)
        push!(s,j)
        push!(v,j)
        for i in sub(Ci,nzrange(Cd,j))
            i > j && break              # upper triangle only
            while !in(i,v)
                push!(s,i)
                push!(v,i)
                i = parent[i]
            end
        end
        append!(rv,collect(s))
        cp[j+1] += convert(Ti,length(rv))
    end
    parent,LowerTriangular(transpose(SparseMatrixCSC(n,n,cp,rv,ones(cp[n+1]))))
end

## write a function that copies the lower triangle of the symmetric matrix to L

function LLcholfact{Tv,Ti}(C::Symmetric{Tv,SparseMatrixCSC{Tv,Ti}})
    parent,L = cholpattern(C)
    n = size(C,2)
    Ct = ltri!(transpose(C.data)) # for now evaluate the lower triangle
    Cti = rowvals(Ct)
    Ctx = nonzeros(Ct)                 # these values will be modified
    Ld = L.data
    Li = rowvals(Ld)
    Lx = fill!(0.,nonzeros(Ld))         # initialize contents of L to zero
    for k in 1:n
        ## spread values of k'th column of A into k'th column of L
        Akr = nzrange(Ct,k)
        Lkr = nzrange(Ld,k)
        Lki = sub(Li,Lkr)
        Lkx = sub(Lx,Lkr)
        if length(Akr) == length(Lkr)
            copy!(Lkx,sub(Ax,Akr))
        else
            Aki = sub(Cti,Akr)
            Akx = sub(Ctx,Akr)
            for i in eachindex(Aki)
                Lkx[searchsortedfirst(Lki,Aki[i])] = Akx[i]
            end
        end
    end
end

function cholsymb{Tv,Ti}(S::Symmetric{Tv,SparseMatrixCSC{Tv,Ti}})
    _,pinv = Metis.nodeND(S)
    C = symperm(S,pinv)
    parent,post = etree(C,true)
    cp = colcounts(C,parent,post)
    CholSymb(pinv,convert(Vector{Ti},cp),parent)
end

function cholfact{Tv,Ti}(A::Symmetric{Tv,SparseMatrixCSC{Tv,Ti}}, S::CholSymb{Ti})
    n = size(A,2)
    cp = S.cp; pinv = S.pinv; parent = S.parent
## maybe change the natural case (i.e. no permutation) to pinv = Array(Ti, 0)
    C  = (pinv == 1:n) ? A : symperm(A, pinv)
    Cp = C.colptr; Ci = C.rowval; Cx = C.nzval
    Lp = copy(cp); Li = Array(Ti, cp[n+1]); Lx = Array(Tv, cp[n+1])
    c  = copy(cp[1:n])
    s  = Array(Ti, n)                   # Ti workspace
    x  = Array(Tv, n)                   # Tv workspace
    for k in 1:n                        # compute L[k,:] for L*L' = C
        ## Nonzero pattern of L[k,:]
        top = cs_ereach(C, k, parent, s, c) # find pattern of L(k,:)
        x[k] = 0                            # x[1:k] is now zero
        for p in Cp[k]:(Cp[k+1]-1)          # x = full(triu(C(:,k)))
            if (Ci[p] <= k) x[Ci[p]] = Cx[p] end
        end
        d = x[k]                        # d = C[k,k]
        x[k] = zero(Tv)                 # clear x for k+1st iteration
        ## Triangular solve
        while top <= n             # solve L[0:k-1,0:k-1] * x = C[:,k]
            i = s[top]             # s[top..n-1] is pattern of L[k,:]
            lki = x[i]/Lx[Lp[i]]   # L[k,i] = x[i]/L[i,i]
            x[i] = zero(Tv)        # clear x for k+1st iteration
            for p in (Lp[i]+1):(c[i]-1)
                x[Li[p]] -= Lx[p] * lki
            end
            d -= lki * lki ;            # d = d - L[k,i]*L[k,i]
            p = c[i]
            c[i] += 1
            Li[p] = k                   # store L[k,i] in column i
            Lx[p] = lki
        end
        ## Compute L[k,k]
        if (d <= 0) error("Matrix is not positive definite, detected at row $k") end
        p = c[k]
        c[k] += 1
        Li[p] = k                 # store L[k,k] = sqrt(d) in column k
        Lx[p] = sqrt(d)
    end
    SparseMatrixCSC(n, n, Lp, Li, Lx)
end


# @everywhere using RegularizedLeastSquares, LinearOperators, TensorDecompositions, LinearAlgebra
#@everywhere using RegularizedLeastSquares, LinearOperators, LinearAlgebra
export smRecovery

###################################################################
# System Matrix recovery using DCT-sparsity and low-rank constraing
###################################################################
function smRecovery(y::Matrix{T}, samplingIdx::Array{Int64}, params::Dict) where T
  # constrain number of threads for distributed computation
  BLAS.set_num_threads(1)

  shape = params[:shape]

  # sampling operator
  P = SamplingOp(samplingIdx, shape)

  # sparsifying transformation and regularization
  if !haskey(params, :sparseTrafo)
    params[:sparseTrafo] = DCTOp(ComplexF64,params[:shape],2)
  end
  reg1Name = get(params, :reg1, "L1")
  reg1 = Regularization(reg1Name, params[:λ1]; params...)

  # low rank regularization
  lrProx = get(params, :reg2, "Nothing")
  λ_lr = get(params, :λ_lr, 1.0)
  reg2 = lrReg(lrProx, λ_lr, params)

  # precalculate inverse for split Bregman iteration
  precon = invPrecon(samplingIdx, prod(shape); params...)
  params[:precon] = precon
  params[:iterationsCG] = 1

  # normalized measurement
  y2, y_norm = getNormalizedMeasurement(y)

  # # reconstruction
  # sf = map(x->splitBregman(P,x,[reg1,reg2]; params...),y2)
  #
  # sfMat = zeros(ComplexF64,prod(shape),size(y,2))
  # for i=1:size(y,2)
  #   # undo normalization
  #   sfMat[:,i] .= y_norm[i]*sf[i]
  # end

  # reconstruction
  sfMat = zeros(ComplexF64,prod(shape),size(y,2))
  solver = SplitBregman(P;reg=[reg1,reg2], params...)
  for k=1:size(y,2)
    sfMat[:,k] = solve(solver, y2[k])
    # undo normalization
    sfMat[:,k] *= y_norm[k]
  end

  return sfMat
end

#########################
# low rank regularization
#########################
function lrReg(lrProx::AbstractString, λ::AbstractFloat, params::Dict{Symbol,Any})
  if lrProx == "LR"
    proxMap = (x,y;kargs...)->proxLR!(x,y,params[:shape];kargs...)
    reg = Regularization(prox! = proxMap, λ=λ, params=params)
  elseif lrProx == "FR"
    proxMap = (x,y;kargs...)->proxFR!(x,y,params[:shape];kargs...)
    reg = Regularization(prox! = proxMap, λ=λ, params=params)
  elseif lrProx == "Nothing"
    reg = Regularization(prox! = (x,y;kargs...)->x, λ=λ, params=params)
  else
    error("proximal map $(lrProx) is not supported")
  end
  return reg
end

# low rank regularization for 3d (using HOSVD)
function proxLR!(x, λ::Real, shape::NTuple{3,Int64}; kargs...)
  guvw_r = hosvd(reshape(real.(x), shape),shape)
  guvw_i = hosvd(reshape(imag.(x), shape),shape)
  proxL1!(guvw_r.core,λ)
  proxL1!(guvw_i.core,λ)
  x[:] .= vec(compose(guvw_r)+1im*compose(guvw_i))
end

# low rank regularization for 2d (using SVD)
function proxLR!(x, λ::Real, shape::NTuple{2,Int64}; kargs...)
  U,S,V = svd(reshape(x,shape))
  proxL1!(S,λ)
  x[:] .= vec(U*Matrix(Diagonal(S))*V')
end

# fixed rank constraing for 3d (using HOSVD)
function proxFR!(x, λ::Real, shape::NTuple{3,Int64}; rank::NTuple{3,Int64}=(1,1,1), kargs...)
  guvw_r = hosvd(reshape(real.(x), shape),rank)
  guvw_i = hosvd(reshape(imag.(x), shape),rank)
  x[:] = vec(compose(guvw_r)+1im*compose(guvw_i))
end

# fixed rank constraing for 3d (using SVD)
function proxFR!(x, λ::Real, shape::NTuple{2,Int64}; rank::Int64=1, kargs...)
  U,S,V = svd(reshape(x, shape))
  S[rank+1:end] .= 0.0
  x[:] .= vec(U*Matrix(Diagonal(S))*V')
end

##########################
# inverse of normal matrix
##########################
mutable struct diagPrecon{T}
  diag::Vector{T}
end

function ldiv!(y::Vector{T},P::diagPrecon{T},x::Vector{T}) where T
  y[:] .= P.diag .* x
end

function ldiv!(P::diagPrecon{T},x::Vector{T}) where T
  x[:] .= P.diag .* x
end

\(P::diagPrecon{T}, x::Vector{T}) where T = P.diag .* x

function invPrecon(samplingIdx, ncol; μ::Float64 = 1.e-2, λ::Float64 = 1.e2, ρ::Float64=1.0, kargs...)
  normalP = normalPDiag(samplingIdx, ncol)
  diag = [μ*normalP[i]+λ+ρ for i=1:length(normalP)]
  diag = diag.^-1
  return diagPrecon(diag)
end

function normalPDiag(samplingIdx, ncol; T=ComplexF64)
  normalP = zeros(T,ncol)
  for i=1:length(samplingIdx)
    normalP[samplingIdx[i]] = 1.0
  end
  return normalP
end

function getNormalizedMeasurement(y::Matrix{T}) where T
  y_norm = [norm(y[:,k]) for k=1:size(y,2)]
  y2 = @DArray [ComplexF64.(y[:,k]/y_norm[k]) for k=1:size(y,2)]
  return y2, y_norm
end


# forward-mode evaluation of an expression tree
# tree is represented as Vector{NodeData} and adjacency matrix from adjmat()
# assumes values of expressions have already been computed
# partials_storage[k] is the partial derivative of nd[k].parent with respect to the value
# of node k. It's efficient to compute this at the same time as the value of the parent.
# Since we use it in reverse mode and in dual forward mode.
# Note that partials_storage makes a subtle assumption that we have a tree instead of
# a general DAG. If we have a DAG, then need to associate storage with each edge of the DAG.
# user_input_buffer and user_output_buffer are used as temporary storage
# when handling user-defined functions
function forward_eval{T}(storage::Vector{T},partials_storage::Vector{T},nd::Vector{NodeData},adj,const_values,parameter_values,x_values::Vector{T},subexpression_values,user_input_buffer=[],user_output_buffer=[])

    @assert length(storage) >= length(nd)
    @assert length(partials_storage) >= length(nd)

    # nd is already in order such that parents always appear before children
    # so a backwards pass through nd is a forward pass through the tree

    children_arr = rowvals(adj)

    for k in length(nd):-1:1
        # compute the value of node k
        @inbounds nod = nd[k]
        partials_storage[k] = zero(T)
        if nod.nodetype == VARIABLE
            @inbounds storage[k] = x_values[nod.index]
        elseif nod.nodetype == VALUE
            @inbounds storage[k] = const_values[nod.index]
        elseif nod.nodetype == SUBEXPRESSION
            @inbounds storage[k] = subexpression_values[nod.index]
        elseif nod.nodetype == PARAMETER
            @inbounds storage[k] = parameter_values[nod.index]
        elseif nod.nodetype == CALL
            op = nod.index
            @inbounds children_idx = nzrange(adj,k)
            #@show children_idx
            n_children = length(children_idx)
            if op == 1 # :+
                tmp_sum = zero(T)
                # sum over children
                for c_idx in children_idx
                    @inbounds ix = children_arr[c_idx]
                    @inbounds partials_storage[ix] = one(T)
                    @inbounds tmp_sum += storage[ix]
                end
                storage[k] = tmp_sum
            elseif op == 2 # :-
                child1 = first(children_idx)
                @assert n_children == 2
                @inbounds ix1 = children_arr[child1]
                @inbounds ix2 = children_arr[child1+1]
                @inbounds tmp_sub = storage[ix1]
                @inbounds tmp_sub -= storage[ix2]
                @inbounds partials_storage[ix1] = one(T)
                @inbounds partials_storage[ix2] = -one(T)
                storage[k] = tmp_sub
            elseif op == 3 # :*
                tmp_prod = one(T)
                for c_idx in children_idx
                    @inbounds tmp_prod *= storage[children_arr[c_idx]]
                end
                if tmp_prod == zero(T) # inefficient
                    for c_idx in children_idx
                        prod_others = one(T)
                        for c_idx2 in children_idx
                            (c_idx == c_idx2) && continue
                            prod_others *= storage[children_arr[c_idx2]]
                        end
                        partials_storage[children_arr[c_idx]] = prod_others
                    end
                else
                    for c_idx in children_idx
                        ix = children_arr[c_idx]
                        partials_storage[ix] = tmp_prod/storage[ix]
                    end
                end
                @inbounds storage[k] = tmp_prod
            elseif op == 4 # :^
                @assert n_children == 2
                idx1 = first(children_idx)
                idx2 = last(children_idx)
                @inbounds ix1 = children_arr[idx1]
                @inbounds ix2 = children_arr[idx2]
                @inbounds base = storage[ix1]
                @inbounds exponent = storage[ix2]
                if exponent == 2
                    @inbounds storage[k] = base*base
                    @inbounds partials_storage[ix1] = 2*base
                else
                    storage[k] = pow(base,exponent)
                    partials_storage[ix1] = exponent*pow(base,exponent-1)
                end
                partials_storage[ix2] = storage[k]*log(base)
            elseif op == 5 # :/
                @assert n_children == 2
                idx1 = first(children_idx)
                idx2 = last(children_idx)
                @inbounds ix1 = children_arr[idx1]
                @inbounds ix2 = children_arr[idx2]
                @inbounds numerator = storage[ix1]
                @inbounds denominator = storage[ix2]
                recip_denominator = 1/denominator
                @inbounds partials_storage[ix1] = recip_denominator
                partials_storage[ix2] = -numerator*recip_denominator*recip_denominator
                storage[k] = numerator*recip_denominator
            elseif op == 6 # ifelse
                @assert n_children == 3
                idx1 = first(children_idx)
                @inbounds condition = storage[children_arr[idx1]]
                @inbounds lhs = storage[children_arr[idx1+1]]
                @inbounds rhs = storage[children_arr[idx1+2]]
                @inbounds partials_storage[children_arr[idx1+1]] = condition == 1
                @inbounds partials_storage[children_arr[idx1+2]] = !(condition == 1)
                storage[k] = ifelse(condition == 1, lhs, rhs)
            elseif op >= USER_OPERATOR_ID_START
                evaluator = user_operator_map[op]
                f_input = Compat.view(user_input_buffer, 1:n_children)
                grad_output = Compat.view(user_output_buffer, 1:n_children)
                r = 1
                for c_idx in children_idx
                    ix = children_arr[c_idx]
                    f_input[r] = storage[ix]
                    grad_output[r] = 0.0
                    r += 1
                end
                fval = MathProgBase.eval_f(evaluator, f_input)::T
                MathProgBase.eval_grad_f(evaluator, grad_output, f_input)
                storage[k] = fval
                r = 1
                for c_idx in children_idx
                    ix = children_arr[c_idx]
                    partials_storage[ix] = grad_output[r]
                    r += 1
                end
            else
                error("Unsupported operation $(operators[op])")
            end
        elseif nod.nodetype == CALLUNIVAR # univariate function
            op = nod.index
            @inbounds child_idx = children_arr[adj.colptr[k]]
            #@assert child_idx == children_arr[first(nzrange(adj,k))]
            child_val = storage[child_idx]
            if op >= USER_UNIVAR_OPERATOR_ID_START
                f = user_univariate_operator_f[op]
                fprime = user_univariate_operator_fprime[op]
                fval = f(child_val)::T
                fprimeval = fprime(child_val)::T
            else
                fval, fprimeval = eval_univariate(op, child_val)
            end
            @inbounds partials_storage[child_idx] = fprimeval
            @inbounds storage[k] = fval
        elseif nod.nodetype == COMPARISON
            op = nod.index

            @inbounds children_idx = nzrange(adj,k)
            n_children = length(children_idx)
            result = true
            for r in 1:n_children-1
                partials_storage[children_arr[children_idx[r]]] = zero(T)
                cval_lhs = storage[children_arr[children_idx[r]]]
                cval_rhs = storage[children_arr[children_idx[r+1]]]
                if op == 1
                    result &= cval_lhs <= cval_rhs
                elseif op == 2
                    result &= cval_lhs == cval_rhs
                elseif op == 3
                    result &= cval_lhs >= cval_rhs
                elseif op == 4
                    result &= cval_lhs < cval_rhs
                elseif op == 5
                    result &= cval_lhs > cval_rhs
                end
            end
            partials_storage[children_arr[children_idx[n_children]]] = zero(T)
            storage[k] = result
        elseif nod.nodetype == LOGIC
            op = nod.index
            @inbounds children_idx = nzrange(adj,k)
            # boolean values are stored as floats
            partials_storage[children_arr[first(children_idx)]] = zero(T)
            partials_storage[children_arr[last(children_idx)]] = zero(T)
            cval_lhs = (storage[children_arr[first(children_idx)]] == 1)
            cval_rhs = (storage[children_arr[last(children_idx)]] == 1)
            if op == 1
                storage[k] = cval_lhs && cval_rhs
            elseif op == 2
                storage[k] = cval_lhs || cval_rhs
            end
        end

    end
    #@show storage

    return storage[1]

end

export forward_eval

# Evaluate directional derivatives of an expression tree.
# This is equivalent to evaluating the expression tree using DualNumbers or ForwardDiff,
# but instead we keep the storage for the epsilon components separate so that we don't
# need to recompute the real components.
# Computes partials_storage_ϵ as well
# We assume that forward_eval has already been called.
function forward_eval_ϵ{N,T}(storage::Vector{T},storage_ϵ::DenseVector{ForwardDiff.Partials{N,T}},partials_storage::Vector{T},partials_storage_ϵ::DenseVector{ForwardDiff.Partials{N,T}},nd::Vector{NodeData},adj,x_values_ϵ,subexpression_values_ϵ)

    @assert length(storage_ϵ) >= length(nd)
    @assert length(partials_storage_ϵ) >= length(nd)

    zero_ϵ = zero(ForwardDiff.Partials{N,T})

    # nd is already in order such that parents always appear before children
    # so a backwards pass through nd is a forward pass through the tree

    children_arr = rowvals(adj)

    for k in length(nd):-1:1
        # compute the value of node k
        @inbounds nod = nd[k]
        partials_storage_ϵ[k] = zero_ϵ
        if nod.nodetype == VARIABLE
            @inbounds storage_ϵ[k] = x_values_ϵ[nod.index]
        elseif nod.nodetype == VALUE
            @inbounds storage_ϵ[k] = zero_ϵ
        elseif nod.nodetype == SUBEXPRESSION
            @inbounds storage_ϵ[k] = subexpression_values_ϵ[nod.index]
        elseif nod.nodetype == PARAMETER
            @inbounds storage_ϵ[k] = zero_ϵ
        else
            ϵtmp = zero_ϵ
            @inbounds children_idx = nzrange(adj,k)
            for c_idx in children_idx
                ix = children_arr[c_idx]
                @inbounds partial = partials_storage[ix]
                if !isfinite(partials_storage[ix]) && storage_ϵ[ix] == zero_ϵ
                    continue
                end
                ϵtmp += storage_ϵ[ix]*partials_storage[ix]
            end
            storage_ϵ[k] = ϵtmp

            if nod.nodetype == CALL
                op = nod.index
                n_children = length(children_idx)
                if op == 3 # :*
                    # Lazy approach for now
                    anyzero = false
                    tmp_prod = one(ForwardDiff.Dual{N,T})
                    for c_idx in children_idx
                        ix = children_arr[c_idx]
                        sval = storage[ix]
                        gnum = ForwardDiff.Dual(sval,storage_ϵ[ix])
                        tmp_prod *= gnum
                        anyzero = ifelse(sval*sval == zero(T), true, anyzero)
                    end
                    # By a quirk of floating-point numbers, we can have
                    # anyzero == true && ForwardDiff.value(tmp_prod) != zero(T)
                    if anyzero # inefficient
                        for c_idx in children_idx
                            prod_others = one(ForwardDiff.Dual{N,T})
                            for c_idx2 in children_idx
                                (c_idx == c_idx2) && continue
                                ix = children_arr[c_idx2]
                                gnum = ForwardDiff.Dual(storage[ix],storage_ϵ[ix])
                                prod_others *= gnum
                            end
                            partials_storage_ϵ[children_arr[c_idx]] = ForwardDiff.partials(prod_others)
                        end
                    else
                        for c_idx in children_idx
                            ix = children_arr[c_idx]
                            prod_others = tmp_prod/ForwardDiff.Dual(storage[ix],storage_ϵ[ix])
                            partials_storage_ϵ[ix] = ForwardDiff.partials(prod_others)
                        end
                    end
                elseif op == 4 # :^
                    @assert n_children == 2
                    idx1 = first(children_idx)
                    idx2 = last(children_idx)
                    @inbounds ix1 = children_arr[idx1]
                    @inbounds ix2 = children_arr[idx2]
                    @inbounds base = storage[ix1]
                    @inbounds base_ϵ = storage_ϵ[ix1]
                    @inbounds exponent = storage[ix2]
                    @inbounds exponent_ϵ = storage_ϵ[ix2]
                    base_gnum = ForwardDiff.Dual(base,base_ϵ)
                    exponent_gnum = ForwardDiff.Dual(exponent,exponent_ϵ)
                    if exponent == 2
                        partials_storage_ϵ[ix1] = 2*base_ϵ
                    else
                        partials_storage_ϵ[ix1] = ForwardDiff.partials(exponent_gnum*pow(base_gnum,exponent_gnum-1))
                    end
                    result_gnum = ForwardDiff.Dual(storage[k],storage_ϵ[k])
                    partials_storage_ϵ[ix2] = ForwardDiff.partials(result_gnum*log(base_gnum))
                elseif op == 5 # :/
                    @assert n_children == 2
                    idx1 = first(children_idx)
                    idx2 = last(children_idx)
                    @inbounds ix1 = children_arr[idx1]
                    @inbounds ix2 = children_arr[idx2]
                    @inbounds numerator = storage[ix1]
                    @inbounds numerator_ϵ = storage_ϵ[ix1]
                    @inbounds denominator = storage[ix2]
                    @inbounds denominator_ϵ = storage_ϵ[ix2]
                    recip_denominator = 1/ForwardDiff.Dual(denominator,denominator_ϵ)
                    partials_storage_ϵ[ix1] = ForwardDiff.partials(recip_denominator)
                    partials_storage_ϵ[ix2] = ForwardDiff.partials(-ForwardDiff.Dual(numerator,numerator_ϵ)*recip_denominator*recip_denominator)
                elseif op >= USER_OPERATOR_ID_START
                    error("User-defined operators not supported for hessian computations")
                end
            elseif nod.nodetype == CALLUNIVAR # univariate function
                op = nod.index
                @inbounds child_idx = children_arr[adj.colptr[k]]
                child_val = storage[child_idx]
                if op >= USER_UNIVAR_OPERATOR_ID_START
                    fprimeprime = user_univariate_operator_fprimeprime[op](child_val)::T
                else
                    fprimeprime = eval_univariate_2nd_deriv(op, child_val,storage[k])
                end
                partials_storage_ϵ[child_idx] = fprimeprime*storage_ϵ[child_idx]
            end
        end

    end

    return storage_ϵ[1]

end

export forward_eval_ϵ

switchblock = Expr(:block)
for i = 1:length(univariate_operators)
    op = univariate_operators[i]
    deriv_expr = univariate_operator_deriv[i]
	ex = :(return $op(x), $deriv_expr::T)
    push!(switchblock.args,i,ex)
end
switchexpr = Expr(:macrocall, Expr(:.,:Lazy,quot(Symbol("@switch"))), :operator_id,switchblock)

@eval @inline function eval_univariate{T}(operator_id,x::T)
    $switchexpr
end

# TODO: optimize sin/cos/exp
switchblock = Expr(:block)
for i = 1:length(univariate_operators)
    op = univariate_operators[i]
    if op == :asec || op == :acsc || op == :asecd || op == :acscd || op == :acsch || op == :trigamma
        # there's an abs in the derivative that Calculus can't symbolically differentiate
        continue
    end
    if i in 1:3 # :+, :-, :abs
        deriv_expr = :(zero(T))
    elseif op == :sin || op == :cos
        deriv_expr = :(-fval)
    elseif op == :exp
        deriv_expr = :(fval)
    else
        deriv_expr = Calculus.differentiate(univariate_operator_deriv[i],:x)
    end
	ex = :(return $deriv_expr::T)
    push!(switchblock.args,i,ex)
end
switchexpr = Expr(:macrocall, Expr(:.,:Lazy,quot(Symbol("@switch"))), :operator_id,switchblock)

@eval @inline function eval_univariate_2nd_deriv{T}(operator_id,x::T,fval::T)
    $switchexpr
end


"""
Abstract type for models using a Gaussian/normal-inverse-gamma conjugate prior.
"""
abstract type AbstractBayesianLM end

"""
    BayesLM(d::Int; <keyword arguments>)

Construct a Bayesian linear model.

# Arguments

- `λ::Float64=1.0`: Prior scaling.
- `shape0::Float64=1e-3`: Inverse-gamma prior shape hyperparameter.
- `scale0::Float64=1e-3`: Inverse-gamma prior scale hyperparameter.
"""
mutable struct BayesLM <: AbstractBayesianLM
    shape0::Float64
    scale0::Float64
    β0::Matrix{Float64}
    Σ0::Hermitian{Float64,Matrix{Float64}}
    Λ0::Hermitian{Float64,Matrix{Float64}}
    shape::Float64
    scale::Float64
    β::Matrix{Float64}
    Σ::Hermitian{Float64,Matrix{Float64}}
    Λ::Hermitian{Float64,Matrix{Float64}}

    function BayesLM(d::Int64; λ::Float64=1.0, shape0::Float64=1e-3, scale0::Float64=1e-3)
        if d <= 0 || shape0 <= 0 || scale0 <= 0
            msg = "d, shape0 and scale0 must be strictly positive"
            throw(ArgumentError(msg))
        end
        β0, Σ0, Λ0 = zeros(d, 1),
        Hermitian(diagm(ones(d)) * λ),
        Hermitian(diagm(ones(d)) / λ)
        β, Σ, Λ = deepcopy(β0), deepcopy(Σ0), deepcopy(Λ0)
        return new(shape0, scale0, β0, Σ0, Λ0, shape0, scale0, β, Σ, Λ)
    end
end

function sherman_morrison_inv(Ainv::AbstractMatrix, u::AbstractMatrix, v::AbstractMatrix)
    if size(u) != size(v)
        throw(DimensionMismatch("u and v have mismatched dimensions"))
    elseif size(u, 2) != 1
        throw(ArgumentError("u and v should have exactly one column"))
    elseif size(Ainv) != (size(u, 1), size(u, 1))
        throw(DimensionMismatch("size of A does not match u/v"))
    end
    numer = Ainv * u * v' * Ainv
    denom = 1.0 + (v' * Ainv * u)[1, 1]
    return Ainv - numer / denom
end

function check_regression_data(X::AbstractMatrix, y::AbstractMatrix)
    if size(y, 1) != 1
        throw(ArgumentError("y should have exactly 1 row"))
    elseif size(X, 2) != size(y, 2)
        throw(DimensionMismatch("X and y don't match in dimension 2"))
    end
end

function fit!(lm::BayesLM, X::AbstractMatrix, y::AbstractMatrix)
    check_regression_data(X, y)
    if size(lm.β, 1) != size(X, 1)
        msg = "X does not have the same number of predictors as the model"
        throw(DimensionMismatch(msg))
    end
    G = Hermitian(X * X')
    Λ = G + lm.Λ
    Σ = size(X, 2) == 1 ? Hermitian(sherman_morrison_inv(lm.Σ, X, X)) : inv(Λ)
    β = Σ * (X * y' + lm.Λ * lm.β)
    shape = lm.shape + size(y, 2) / 2
    scale = lm.scale + 0.5 * (y * y' - β' * Λ * β + lm.β' * lm.Λ * lm.β)[1, 1]
    return lm.shape, lm.scale, lm.β, lm.Σ, lm.Λ = shape, scale, β, Σ, Λ
end

function (lm::BayesLM)(X::AbstractMatrix)
    if size(X, 1) != size(lm.β, 1)
        msg = "X does not have the same number of predictors as the model"
        throw(DimensionMismatch(msg))
    end
    return lm.β' * X
end

function shape_scale(lm::BayesLM)
    return lm.shape, lm.scale
end

"""
    std(model::AbstractBayesianLM)

Return the posterior mean of the inverse-gamma distribution.
"""
function std(model::AbstractBayesianLM)
    shape, scale = shape_scale(model)
    if shape <= 1
        throw(ArgumentError("normal-inverse-gamma has undefined mean for shape ≤ 1"))
    end
    return sqrt(scale / (shape - 1))
end


using DirichletMixtureModels

importall DirichletMixtureModels

import Distributions:
        UnivariateDistribution,
        MultivariateDistribution,
        Exponential,
        Normal,
        NormalCanon,
        MvNormal,
        NormalStats,
        rand,
        suffstats,
        pdf

import ConjugatePriors:
    NormalGamma,
    NormalWishart,
    posterior_canon,
    rand

function example_pdf(y::Float64,μ,λ,μ0,n0,α0,β0)
  pdf(NormalCanon(λ*μ, λ), y)
end
function example_pdf(y::Array{Float64},μ,λ,μ0,n0,α0,β0)
  pdf.(NormalCanon(λ*μ, λ), y)
end
function example_post(Y::Array{Float64,1}, μ0,n0,α0,β0)
  p=posterior_canon(NormalGamma(μ0,n0,α0,β0),suffstats(Normal,Y))
  rand(p)
end
function example_post(y::Float64, μ0,n0,α0,β0)
  p=posterior_canon(NormalGamma(μ0,n0,α0,β0),suffstats(Normal,[y]))
  rand(p)
end
function example_pri(μ0,n0,α0,β0)
  rand(NormalGamma(μ0,n0,α0,β0))
end
function example_marg(y::Float64,μ0,n0,α0,β0)
  gamma(α0+1./2)/gamma(α0) * sqrt(n0/(n0+1)) * 1/sqrt(2*π) * β0^α0 /
    (β0+n0/2/(n0+1)*(y-μ0)^2)^(α0+1./2)
end


###########################################################################
# Scaled Least-squares
###########################################################################
function bearing_sls(angles, gains, ref_angles, ref_gains)
	@assert length(angles) == length(gains)
	num_gains = length(gains)
	best_err = Inf
	best_shift = 0
	best_scale = 0
	temp_gains = copy(gains)
	for scale_factor = 0:100
		temp_gains += 1
		for shift = 0:359
			err = 0.0
			for j = 1:num_gains
				if validgain(gains[j])
					ref_gain = interp_linear(angles[j], shift, ref_angles, ref_gains)
					err += (temp_gains[j] - ref_gain)^2
				end
			end
			if err < best_err
				best_err = err
				best_shift = shift
				best_scale = scale_factor
			end
		end
	end
	return 360 - best_shift
end

function bearing_sls(dir_file::AbstractString, ref_file::AbstractString; samples::Int64=0, sense::Int64=1)
	test_gp = GainPattern(dir_file)
	ref_gp = GainPattern(ref_file)

	num_samples = length(test_gp.angles)
	num_exp = length(test_gp.samples[1])

	# Allowing user to specify count
	if (samples > 0) && (samples < num_samples)
		num_samples = samples
	end

	aoa_array = zeros(num_exp)
	angles = zeros(num_samples)
	gains = zeros(num_samples)
	for i = 1:num_exp
		if sense > 0
			for j = 1:num_samples
				# Loop through all samples in this experiment...
				angles[j] = test_gp.angles[j]
				gains[j] = test_gp.samples[j][i]
			end
		elseif sense == -1
			# backwards loop, starting from last measurement
			for j = 1:num_samples
				angles[j] = test_gp.angles[end - j + 1]
				gains[j] = test_gp.samples[end - j + 1][i]
			end
		elseif sense == -2
			# backwards loop, starting from origin
			for j = 2:num_samples
				angles[j] = test_gp.angles[end - j + 2]
				gains[j] = test_gp.samples[end - j + 2][i]
			end
			angles[1] = test_gp.angles[1]
			gains[1] = test_gp.samples[1][i]
		end
		aoa_array[i] = bearing_sls(angles, gains, ref_gp.angles, ref_gp.meangains)
	end
	return aoa_array
end


naolaplacian00(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = ud2/(2*sqrt(pi)*r)
naolaplacian1m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3/pi)*sin(θ)*sin(ϕ)*(-2*u + r^2*ud2))/(2*r^3)
naolaplacian10(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3/pi)*cos(θ)*(-2*u + r^2*ud2))/(2*r^3)
naolaplacian11(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3/pi)*cos(ϕ)*sin(θ)*(-2*u + r^2*ud2))/ (2*r^3)
naolaplacian2m2(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(15/pi)*sin(θ)^2*sin(2*ϕ)*(-6*u + r^2*ud2))/(4*r^3))
naolaplacian2m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/pi)*sin(2*θ)*sin(ϕ)*(-6*u + r^2*ud2))/(4*r^3)
naolaplacian20(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(5/pi)*(1 + 3*cos(2*θ))*(-6*u + r^2*ud2))/(8*r^3)
naolaplacian21(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/pi)*cos(ϕ)*sin(2*θ)*(-6*u + r^2*ud2))/(4*r^3)
naolaplacian22(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/pi)*cos(2*ϕ)*sin(θ)^2*(-6*u + r^2*ud2))/(4*r^3)
naolaplacian3m3(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(35/(2*pi))*sin(θ)^3*sin(3*ϕ)*(-12*u + r^2*ud2))/(4*r^3)
naolaplacian3m2(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(105/pi)*cos(θ)*sin(θ)^2*sin(2*ϕ)*(-12*u + r^2*ud2))/(4*r^3))
naolaplacian3m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(21/(2*pi))*(3 + 5*cos(2*θ))*sin(θ)*sin(ϕ)*(-12*u + r^2*ud2))/(8*r^3)
naolaplacian30(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(7/pi)*(3*cos(θ) + 5*cos(3*θ))*(-12*u + r^2*ud2))/(16*r^3)
naolaplacian31(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(21/(2*pi))*(3 + 5*cos(2*θ))*cos(ϕ)*sin(θ)*(-12*u + r^2*ud2))/(8*r^3)
naolaplacian32(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(105/pi)*cos(θ)*cos(2*ϕ)*sin(θ)^2*(-12*u + r^2*ud2))/(4*r^3)
naolaplacian33(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(35/(2*pi))*cos(3*ϕ)*sin(θ)^3*(-12*u + r^2*ud2))/(4*r^3)
naolaplacian4m4(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(35/pi)*sin(θ)^4*sin(4*ϕ)*(-20*u + r^2*ud2))/(16*r^3))
naolaplacian4m3(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/(2*pi))*cos(θ)*sin(θ)^3*sin(3*ϕ)*(-20*u + r^2*ud2))/(4*r^3)
naolaplacian4m2(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(5/pi)*(5 + 7*cos(2*θ))*sin(θ)^2*sin(2*ϕ)*(-20*u + r^2*ud2))/(16*r^3))
naolaplacian4m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5/(2*pi))*(2*sin(2*θ) + 7*sin(4*θ))*sin(ϕ)*(-20*u + r^2*ud2))/(32*r^3)
naolaplacian40(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*(9 + 20*cos(2*θ) + 35*cos(4*θ))*(-20*u + r^2*ud2))/(128*sqrt(pi)*r^3)
naolaplacian41(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5/(2*pi))*cos(ϕ)*(2*sin(2*θ) + 7*sin(4*θ))*(-20*u + r^2*ud2))/(32*r^3)
naolaplacian42(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5/pi)*(5 + 7*cos(2*θ))*cos(2*ϕ)*sin(θ)^2*(-20*u + r^2*ud2))/(16*r^3)
naolaplacian43(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/(2*pi))*cos(θ)*cos(3*ϕ)*sin(θ)^3*(-20*u + r^2*ud2))/(4*r^3)
naolaplacian44(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/pi)*cos(4*ϕ)*sin(θ)^4*(-20*u + r^2*ud2))/(16*r^3)
naolaplacian5m5(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(77/(2*pi))*sin(θ)^5*sin(5*ϕ)*(-30*u + r^2*ud2))/(16*r^3)
naolaplacian5m4(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(385/pi)*cos(θ)*sin(θ)^4*sin(4*ϕ)*(-30*u + r^2*ud2))/(16*r^3))
naolaplacian5m3(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(385/(2*pi))*(7 + 9*cos(2*θ))*sin(θ)^3*sin(3*ϕ)*(-30*u + r^2*ud2))/(32*r^3)
naolaplacian5m2(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(1155/pi)*(5*cos(θ) + 3*cos(3*θ))*sin(θ)^2*sin(2*ϕ)*(-30*u + r^2*ud2))/(32*r^3))
naolaplacian5m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(165/pi)*(15 + 28*cos(2*θ) + 21*cos(4*θ))*sin(θ)*sin(ϕ)*(-30*u + r^2*ud2))/(128*r^3)
naolaplacian50(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(11/pi)*(30*cos(θ) + 35*cos(3*θ) + 63*cos(5*θ))*(-30*u + r^2*ud2))/(256*r^3)
naolaplacian51(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(165/pi)*(15 + 28*cos(2*θ) + 21*cos(4*θ))*cos(ϕ)*sin(θ)*(-30*u + r^2*ud2))/(128*r^3)
naolaplacian52(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1155/pi)*(5*cos(θ) + 3*cos(3*θ))*cos(2*ϕ)*sin(θ)^2*(-30*u + r^2*ud2))/(32*r^3)
naolaplacian53(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(385/(2*pi))*(7 + 9*cos(2*θ))*cos(3*ϕ)*sin(θ)^3*(-30*u + r^2*ud2))/(32*r^3)
naolaplacian54(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/pi)*cos(θ)*cos(4*ϕ)*sin(θ)^4*(-30*u + r^2*ud2))/(16*r^3)
naolaplacian55(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(77/(2*pi))*cos(5*ϕ)*sin(θ)^5*(-30*u + r^2*ud2))/(16*r^3)
naolaplacian6m6(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(3003/(2*pi))*sin(θ)^6*sin(6*ϕ)*(-42*u + r^2*ud2))/(32*r^3))
naolaplacian6m5(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(1001/(2*pi))*cos(θ)*sin(θ)^5*sin(5*ϕ)*(-42*u + r^2*ud2))/(16*r^3)
naolaplacian6m4(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(91/pi)*(9 + 11*cos(2*θ))*sin(θ)^4*sin(4*ϕ)*(-42*u + r^2*ud2))/(64*r^3))
naolaplacian6m3(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1365/(2*pi))*(21*cos(θ) + 11*cos(3*θ))*sin(θ)^3*sin(3*ϕ)*(-42*u + r^2*ud2))/(64*r^3)
naolaplacian6m2(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(1365/(2*pi))*(35 + 60*cos(2*θ) + 33*cos(4*θ))*sin(θ)^2*sin(2*ϕ)*(-42*u + r^2*ud2))/(256*r^3))
naolaplacian6m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(273/pi)*(50*cos(θ) + 45*cos(3*θ) + 33*cos(5*θ))*sin(θ)*sin(ϕ)*(-42*u + r^2*ud2))/(256*r^3)
naolaplacian60(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(13/pi)*(50 + 105*cos(2*θ) + 126*cos(4*θ) + 231*cos(6*θ))*(-42*u + r^2*ud2))/(1024*r^3)
naolaplacian61(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(273/pi)*(50*cos(θ) + 45*cos(3*θ) + 33*cos(5*θ))*cos(ϕ)*sin(θ)*(-42*u + r^2*ud2))/(256*r^3)
naolaplacian62(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1365/(2*pi))*(35 + 60*cos(2*θ) + 33*cos(4*θ))*cos(2*ϕ)*sin(θ)^2*(-42*u + r^2*ud2))/(256*r^3)
naolaplacian63(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(1365/(2*pi))*(21*cos(θ) + 11*cos(3*θ))*cos(3*ϕ)*sin(θ)^3*(-42*u + r^2*ud2))/(64*r^3)
naolaplacian64(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(91/pi)*(9 + 11*cos(2*θ))*cos(4*ϕ)*sin(θ)^4*(-42*u + r^2*ud2))/(64*r^3)
naolaplacian65(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(1001/(2*pi))*cos(θ)*cos(5*ϕ)*sin(θ)^5*(-42*u + r^2*ud2))/(16*r^3)
naolaplacian66(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(3003/(2*pi))*cos(6*ϕ)*sin(θ)^6*(-42*u + r^2*ud2))/(32*r^3)
naolaplacian7m7(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(715/pi)*sin(θ)^7*sin(7*ϕ)*(-56*u + r^2*ud2))/(64*r^3)
naolaplacian7m6(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(5005/(2*pi))*cos(θ)*sin(θ)^6*sin(6*ϕ)*(-56*u + r^2*ud2))/(32*r^3))
naolaplacian7m5(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/pi)*(11 + 13*cos(2*θ))*sin(θ)^5*sin(5*ϕ)*(-56*u + r^2*ud2))/(128*r^3)
naolaplacian7m4(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(385/pi)*(27*cos(θ) + 13*cos(3*θ))*sin(θ)^4*sin(4*ϕ)*(-56*u + r^2*ud2))/(128*r^3))
naolaplacian7m3(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/pi)*(189 + 308*cos(2*θ) + 143*cos(4*θ))*sin(θ)^3*sin(3*ϕ)*(-56*u + r^2*ud2))/(512*r^3)
naolaplacian7m2(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(35/(2*pi))*(350*cos(θ) + 275*cos(3*θ) + 143*cos(5*θ))*sin(θ)^2*sin(2*ϕ)*(-56*u + r^2*ud2))/(512*r^3))
naolaplacian7m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(105/pi)*(350 + 675*cos(2*θ) + 594*cos(4*θ) + 429*cos(6*θ))*sin(θ)*sin(ϕ)*(-56*u + r^2*ud2))/(2048*r^3)
naolaplacian70(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(15/pi)*(175*cos(θ) + 189*cos(3*θ) + 231*cos(5*θ) + 429*cos(7*θ))*(-56*u + r^2*ud2))/(2048*r^3)
naolaplacian71(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(105/pi)*(350 + 675*cos(2*θ) + 594*cos(4*θ) + 429*cos(6*θ))*cos(ϕ)*sin(θ)*(-56*u + r^2*ud2))/(2048*r^3)
naolaplacian72(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/(2*pi))*(350*cos(θ) + 275*cos(3*θ) + 143*cos(5*θ))*cos(2*ϕ)*sin(θ)^2*(-56*u + r^2*ud2))/(512*r^3)
naolaplacian73(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(35/pi)*(189 + 308*cos(2*θ) + 143*cos(4*θ))*cos(3*ϕ)*sin(θ)^3*(-56*u + r^2*ud2))/(512*r^3)
naolaplacian74(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/pi)*(27*cos(θ) + 13*cos(3*θ))*cos(4*ϕ)*sin(θ)^4*(-56*u + r^2*ud2))/(128*r^3)
naolaplacian75(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(385/pi)*(11 + 13*cos(2*θ))*cos(5*ϕ)*sin(θ)^5*(-56*u + r^2*ud2))/(128*r^3)
naolaplacian76(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(5005/(2*pi))*cos(θ)*cos(6*ϕ)*sin(θ)^6*(-56*u + r^2*ud2))/(32*r^3)
naolaplacian77(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(715/pi)*cos(7*ϕ)*sin(θ)^7*(-56*u + r^2*ud2))/(64*r^3)
naolaplacian8m8(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(12155/pi)*sin(θ)^8*sin(8*ϕ)*(-72*u + r^2*ud2))/(256*r^3))
naolaplacian8m7(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(12155/pi)*cos(θ)*sin(θ)^7*sin(7*ϕ)*(-72*u + r^2*ud2))/(64*r^3)
naolaplacian8m6(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((sqrt(7293/(2*pi))*(13 + 15*cos(2*θ))*sin(θ)^6*sin(6*ϕ)*(-72*u + r^2*ud2))/(128*r^3))
naolaplacian8m5(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17017/pi)*(11*cos(θ) + 5*cos(3*θ))*sin(θ)^5*sin(5*ϕ)*(-72*u + r^2*ud2))/(256*r^3)
naolaplacian8m4(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(1309/pi)*(99 + 156*cos(2*θ) + 65*cos(4*θ))*sin(θ)^4*sin(4*ϕ)*(-72*u + r^2*ud2))/(1024*r^3))
naolaplacian8m3(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(19635/pi)*(126*cos(θ) + 91*cos(3*θ) + 39*cos(5*θ))*sin(θ)^3*sin(3*ϕ)*(-72*u + r^2*ud2))/(1024*r^3)
naolaplacian8m2(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = -((3*sqrt(595/(2*pi))*(210 + 385*cos(2*θ) + 286*cos(4*θ) + 143*cos(6*θ))*sin(θ)^2*sin(2*ϕ)*(-72*u + r^2*ud2))/(2048*r^3))
naolaplacian8m1(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17/pi)*(70*sin(2*θ) + 154*sin(4*θ) + 286*sin(6*θ) + 715*sin(8*θ))*sin(ϕ)*(-72*u + r^2*ud2))/(8192*r^3)
naolaplacian80(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(17/pi)*(1225 + 2520*cos(2*θ) + 2772*cos(4*θ) + 3432*cos(6*θ) + 6435*cos(8*θ))*(-72*u + r^2*ud2))/(32768*r^3)
naolaplacian81(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17/pi)*cos(ϕ)*(70*sin(2*θ) + 154*sin(4*θ) + 286*sin(6*θ) + 715*sin(8*θ))*(-72*u + r^2*ud2))/(8192*r^3)
naolaplacian82(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(595/(2*pi))*(210 + 385*cos(2*θ) + 286*cos(4*θ) + 143*cos(6*θ))*cos(2*ϕ)*sin(θ)^2*(-72*u + r^2*ud2))/(2048*r^3)
naolaplacian83(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(19635/pi)*(126*cos(θ) + 91*cos(3*θ) + 39*cos(5*θ))*cos(3*ϕ)*sin(θ)^3*(-72*u + r^2*ud2))/(1024*r^3)
naolaplacian84(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(1309/pi)*(99 + 156*cos(2*θ) + 65*cos(4*θ))*cos(4*ϕ)*sin(θ)^4*(-72*u + r^2*ud2))/(1024*r^3)
naolaplacian85(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(17017/pi)*(11*cos(θ) + 5*cos(3*θ))*cos(5*ϕ)*sin(θ)^5*(-72*u + r^2*ud2))/(256*r^3)
naolaplacian86(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (sqrt(7293/(2*pi))*(13 + 15*cos(2*θ))*cos(6*ϕ)*sin(θ)^6*(-72*u + r^2*ud2))/(128*r^3)
naolaplacian87(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(12155/pi)*cos(θ)*cos(7*ϕ)*sin(θ)^7*(-72*u + r^2*ud2))/(64*r^3)
naolaplacian88(u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = (3*sqrt(12155/pi)*cos(8*ϕ)*sin(θ)^8*(-72*u + r^2*ud2))/(256*r^3)

naolaplacian(l::Int, m::Int, u::Float64, ud2::Float64, r::Float64, θ::Float64, ϕ::Float64) = naolaplaciantab[l,m](u,ud2,r,θ,ϕ)

naolaplaciantab = Dict(
    (0,0) => naolaplacian00,
    (1,-1) => naolaplacian1m1,
    (1,0) => naolaplacian10,
    (1,1) => naolaplacian11,
    (2,-2) => naolaplacian2m2,
    (2,-1) => naolaplacian2m1,
    (2,0) => naolaplacian20,
    (2,1) => naolaplacian21,
    (2,2) => naolaplacian22,
    (3,-3) => naolaplacian3m3,
    (3,-2) => naolaplacian3m2,
    (3,-1) => naolaplacian3m1,
    (3,0) => naolaplacian30,
    (3,1) => naolaplacian31,
    (3,2) => naolaplacian32,
    (3,3) => naolaplacian33,
    (4,-4) => naolaplacian4m4,
    (4,-3) => naolaplacian4m3,
    (4,-2) => naolaplacian4m2,
    (4,-1) => naolaplacian4m1,
    (4,0) => naolaplacian40,
    (4,1) => naolaplacian41,
    (4,2) => naolaplacian42,
    (4,3) => naolaplacian43,
    (4,4) => naolaplacian44,
    (5,-5) => naolaplacian5m5,
    (5,-4) => naolaplacian5m4,
    (5,-3) => naolaplacian5m3,
    (5,-2) => naolaplacian5m2,
    (5,-1) => naolaplacian5m1,
    (5,0) => naolaplacian50,
    (5,1) => naolaplacian51,
    (5,2) => naolaplacian52,
    (5,3) => naolaplacian53,
    (5,4) => naolaplacian54,
    (5,5) => naolaplacian55,
    (6,-6) => naolaplacian6m6,
    (6,-5) => naolaplacian6m5,
    (6,-4) => naolaplacian6m4,
    (6,-3) => naolaplacian6m3,
    (6,-2) => naolaplacian6m2,
    (6,-1) => naolaplacian6m1,
    (6,0) => naolaplacian60,
    (6,1) => naolaplacian61,
    (6,2) => naolaplacian62,
    (6,3) => naolaplacian63,
    (6,4) => naolaplacian64,
    (6,5) => naolaplacian65,
    (6,6) => naolaplacian66,
    (7,-7) => naolaplacian7m7,
    (7,-6) => naolaplacian7m6,
    (7,-5) => naolaplacian7m5,
    (7,-4) => naolaplacian7m4,
    (7,-3) => naolaplacian7m3,
    (7,-2) => naolaplacian7m2,
    (7,-1) => naolaplacian7m1,
    (7,0) => naolaplacian70,
    (7,1) => naolaplacian71,
    (7,2) => naolaplacian72,
    (7,3) => naolaplacian73,
    (7,4) => naolaplacian74,
    (7,5) => naolaplacian75,
    (7,6) => naolaplacian76,
    (7,7) => naolaplacian77,
    (8,-8) => naolaplacian8m8,
    (8,-7) => naolaplacian8m7,
    (8,-6) => naolaplacian8m6,
    (8,-5) => naolaplacian8m5,
    (8,-4) => naolaplacian8m4,
    (8,-3) => naolaplacian8m3,
    (8,-2) => naolaplacian8m2,
    (8,-1) => naolaplacian8m1,
    (8,0) => naolaplacian80,
    (8,1) => naolaplacian81,
    (8,2) => naolaplacian82,
    (8,3) => naolaplacian83,
    (8,4) => naolaplacian84,
    (8,5) => naolaplacian85,
    (8,6) => naolaplacian86,
    (8,7) => naolaplacian87,
    (8,8) => naolaplacian88
)


using ModelingToolkit, DifferentialEquations
@variables t
D = Differential(t)
sts = @variables x(t) RHS(t)  # RHS is observed
ps = @parameters τ=0.3      # parameters
@named m = ODESystem([ RHS  ~ (1 - x)/τ, D(x) ~ RHS ], t, sts, ps)
ms = structural_simplify(m)
prob = ODEProblem(ms, [x => 1.1], (0.0,1.0), [])
sol = solve(prob, Tsit5())

sol(0.3, idxs=[RHS])


########
# Set up parameters
########
parameters = (
    ρₒ = 1.0,  # reference density
    cₛ = 1e-2, # sound speed
    ℓᵐ = 10,   # jet thickness, (larger is thinner)
    ℓ = 20,   # perturbation thickness, (larger is thinner)
    m = 2,    # number of sign changes on equator for perturbation
    ϕᵖ = π / 2 * 0.05, # of centerdness of perturbation
    ϵ = 0.3,  # perturbation amplitude
    vˢ = 5e-4, # velocity scale
    α = 2e-4,
    Ω = 1e-3,
    λ_coupler = (500 / 60 / 86400), #(L_airsea / τ_airsea)
)

########
# Set up inital condition
########
uᵐ(p, λ, ϕ, r) = p.ℓᵐ * sech(p.ℓᵐ * ϕ)^2
vᵐ(p, λ, ϕ, r) = 0.0
hᵐ(p, λ, ϕ, r) = 0.0

u1(p, λ, ϕ, r) = p.ℓ * 2 * (ϕ - p.ϕᵖ) * exp(-p.ℓ * (ϕ - p.ϕᵖ)^2) * cos(ϕ) * cos(2 * (ϕ - p.ϕᵖ)) * sin(p.m * λ)
u2(p, λ, ϕ, r) = exp(-p.ℓ * (ϕ - p.ϕᵖ)^2) * sin(ϕ) * cos(2 * (ϕ - p.ϕᵖ)) * sin(p.m * λ)
u3(p, λ, ϕ, r) = 2 * exp(-p.ℓ * (ϕ - p.ϕᵖ)^2) * cos(ϕ) * sin(2 * (ϕ - p.ϕᵖ)) * sin(p.m * λ)
uᵖ(p, λ, ϕ, r) = u1(p, λ, ϕ, r) + u2(p, λ, ϕ, r) + u3(p, λ, ϕ, r)
vᵖ(p, λ, ϕ, r) = p.m * exp(-p.ℓ * (ϕ - p.ϕᵖ)^2) * cos(2 * (ϕ - p.ϕᵖ)) * cos(p.m * λ)
hᵖ(p, λ, ϕ, r) = 0.0

ρ₀(p, λ, ϕ, r) = p.ρₒ
ρuʳᵃᵈ(p, λ, ϕ, r) = 0
ρuˡᵃᵗ(p, λ, ϕ, r) = p.vˢ * ρ₀(p, λ, ϕ, r) * (p.ϵ * vᵖ(p, λ, ϕ, r))
ρuˡᵒⁿ(p, λ, ϕ, r) = p.vˢ * ρ₀(p, λ, ϕ, r) * (uᵐ(p, λ, ϕ, r) + p.ϵ * uᵖ(p, λ, ϕ, r))
ρθ₀B(p, λ, ϕ, r) = ρ₀(p, λ, ϕ, r) * tanh(p.ℓᵐ * ϕ)
ρθ₀A(p, λ, ϕ, r) = 0

# Cartesian Representation for initialization (boiler plate really)
ρ₀ᶜᵃʳᵗ(p, x...) = ρ₀(p, lon(x...), lat(x...), rad(x...))
ρu⃗₀ᶜᵃʳᵗ(p, x...) = (
    ρuʳᵃᵈ(p, lon(x...), lat(x...), rad(x...)) * r̂(x...) +
    ρuˡᵃᵗ(p, lon(x...), lat(x...), rad(x...)) * ϕ̂(x...) +
    ρuˡᵒⁿ(p, lon(x...), lat(x...), rad(x...)) * λ̂(x...)
)
ρθ₀ᶜᵃʳᵗA(p, x...) = ρθ₀A(p, lon(x...), lat(x...), rad(x...))
ρθ₀ᶜᵃʳᵗB(p, x...) = ρθ₀B(p, lon(x...), lat(x...), rad(x...))


# Source code from this PR
# https://github.com/JuliaStats/HypothesisTests.jl/pull/124
using HypothesisTests
using LinearAlgebra
using StatsFuns: norminvcdf, normccdf
export ShapiroWilkTest, SWCoeffs

#=
From:
PATRICK ROYSTON
Approximating the Shapiro-Wilk W-test for non-normality
*Statistics and Computing* (1992) **2**, 117-119
DOI: [10.1007/BF01891203](https://doi.org/10.1007/BF01891203)
=#

# TODO: Rerun simulation and polynomial fitting

const ROYSTON_COEFFS = Dict{String, Vector{Float64}}(
"C1" => [0.0E0, 0.221157E0, -0.147981E0, -0.207119E1, 0.4434685E1, -0.2706056E1],
"C2" => [0.0E0, 0.42981E-1, -0.293762E0, -0.1752461E1, 0.5682633E1, -0.3582633E1],
"C3" => [0.5440E0, -0.39978E0, 0.25054E-1, -0.6714E-3],
"C4" => [0.13822E1, -0.77857E0, 0.62767E-1, -0.20322E-2],
"C5" => [-0.15861E1, -0.31082E0, -0.83751E-1, 0.38915E-2],
"C6" => [-0.4803E0, -0.82676E-1, 0.30302E-2],
"C7" => [0.164E0, 0.533E0],
"C8" => [0.1736E0, 0.315E0],
"C9" => [0.256E0, -0.635E-2],
"G"  => [-0.2273E1, 0.459E0]
)

for (s,c) in ROYSTON_COEFFS
    @eval $(Symbol("_"*s))(x) = Base.Math.@horner(x, $(c...))
end

#=
The following hardcoded constants has been replaced by more precise values:

SQRTH = sqrt(2.0)/2.0 # 0.70711E0
TH = 3/8 # 0.375E0
SMALL = eps(1.0) # 1E-19
PI6 = π/6 # 0.1909859E1
STQR = asin(sqrt(0.75)) # 0.1047198E1
=#

struct SWCoeffs <: AbstractVector{Float64}
    N::Int
    A::Vector{Float64}
end

Base.size(SWc::SWCoeffs) = (SWc.N,)
Base.IndexStyle(::Type{SWCoeffs}) = IndexLinear()

function Base.getindex(SWc::SWCoeffs, i::Int)
    if i <= lastindex(SWc.A)
        return SWc.A[i]
    elseif i <= length(SWc)
        if isodd(SWc.N) && i == div(SWc.N, 2) + 1
            return 0.0
        else
            return -SWc.A[SWc.N+1-i]
        end
    else
        throw(BoundsError(SWc, i))
    end
end

function SWCoeffs(N::Int)
    if N < 3
        throw(ArgumentError("N must be greater than or equal to 3: $N"))
    elseif N == 3 # exact
        return SWCoeffs(N, [sqrt(2.0)/2.0])
    else
        # Weisberg&Bingham 1975 statistic; store only positive half of m:
        # it is (anti-)symmetric; hence '2' factor below
        m = [-norminvcdf((i - 3/8)/(N + 1/4)) for i in 1:div(N,2)]
        mᵀm = 2sum(abs2, m)

        x = 1/sqrt(N)

        a₁ = m[1]/sqrt(mᵀm) + _C1(x) # aₙ = cₙ + (...)

        if N ≤ 5
            ϕ = (mᵀm - 2m[1]^2)/(1 - 2a₁^2)
            m = m/sqrt(ϕ) # A, but reusing m to save allocs
            m[1] = a₁
        else
            a₂ = m[2]/sqrt(mᵀm) + _C2(x) # aₙ₋₁ = cₙ₋₁ + (...)
            ϕ = (mᵀm - 2m[1]^2 - 2m[2]^2)/(1 - 2a₁^2 - 2a₂^2)
            m = m/sqrt(ϕ) # A, but reusing m to save allocs
            m[1], m[2] = a₁, a₂
        end

        return SWCoeffs(N, m)
    end
end

function swstat(X::AbstractArray{T}, A::SWCoeffs) where T<:Real

    if X[end] - X[1] < lastindex(X)*eps(eltype(X))
        throw("Data seems to be constant!")
    end

    AX = dot(A,X)
    S² = sum(abs2, X.-mean(X))

    return AX^2/S²
end

function pvalue(W::T, A::SWCoeffs, N1=A.N) where T<:Real

    if A.N == 3 # exact by Shapiro&Wilk 1965
        return π/6 * (asin(sqrt(W)) - asin(sqrt(0.75)))
    elseif A.N ≤ 11
        γ = _G(A.N)
        if γ ≤ log(1 - W)
            return zero(W)
        end
        w = -log(γ - log(1 - W))
        μ = _C3(A.N)
        σ = exp(_C4(A.N))
    else
        w = log(1 - W)
        μ = _C5(log(A.N))
        σ = exp(_C6(log(A.N)))
    end

    if (A.N - N1) > 0
        throw("Not implemented yet!")
    end

    return normccdf((w - μ)/σ)
end

struct ShapiroWilkTest <: HypothesisTests.HypothesisTest
    SWc::SWCoeffs            # Expectation of order statistics for Shapiro-Wilk test
    W::Float64            # test statistic
    N1::Int               #
end

HypothesisTests.testname(::ShapiroWilkTest) = "Shapiro-Wilk normality test"
HypothesisTests.population_param_of_interest(t::ShapiroWilkTest) =
("Squared correlation of data and SWCoeffes (W)", 1.0, t.W)

HypothesisTests.default_tail(::ShapiroWilkTest) = :left

function HypothesisTests.show_params(io::IO, t::ShapiroWilkTest, ident)
    println(io, ident, "number of observations:         ", t.SWc.N)
    println(io, ident, "censored ratio:                 ", (t.SWc.N - t.N1)/t.SWc.N)
    println(io, ident, "W-statistic:                    ", t.W)
end

HypothesisTests.pvalue(t::ShapiroWilkTest) = pvalue(t.W, t.SWc, t.N1)

function ShapiroWilkTest(
    X::AbstractArray{T}, SWc::SWCoeffs=SWCoeffs(length(X)), N1=length(X)) where {T<:Real}

    N = length(X)
    #fatal errors
    if N < 3
        throw("Need at least 3 samples.")
    elseif N1 > N
        throw("N1 must be less than or equal to length(X)")
    elseif length(SWc) ≠ length(X)
        throw("Length of the sample differs from Shapiro-Wilk coefficients!")
    end

    #non-fatal errors
    if N > 5000
        @warn("p-value may be unreliable for samples larger than 5000 points")
    elseif (N1 < N) && (N < 20)
        @warn("Number of samples is < 20. Censoring may produce unreliable p-value.")
    elseif (N - N1)/N > 0.8
        @warn("(N - N1)/N > 0.8. Censoring too much data may produce unreliable p-value.")
    end

    if !issorted(view(X, 1:N1))
        @warn("Shapiro-Wilk requires sorted data")
        W = swstat(sort(X[1:N1]), SWc)
    else
        W = swstat(view(X, 1:N1), SWc)
    end

    return ShapiroWilkTest(SWc, W, N1)
end

#=
# To compare with the standard ALGORITHM AS R94 fortran subroutine
#  * grab scipys (swilk.f)[https://github.com/scipy/scipy/blob/master/scipy/stats/statlib/swilk.f];
#  * compile
#  ```
#  gfortran -shared -fPIC -o swilk.so swilk.f
#  gfortran -fdefault-integer-8 -fdefault-real-8 -shared -fPIC swilk.f -o swilk64.so
#  ```

for (lib, I, F) in (("./swilk64.so", Int64, Float64),
                    ("./swilk.so"  , Int32, Float32))
    @eval begin
        function swilkfort!(X::AbstractVector{$F}, A::AbstractVector{$F}, computeA=true)

            w, pval = Ref{$F}(0.0), Ref{$F}(0.0)
            ifault = Ref{$I}(0)

            ccall((:swilk_, $lib),
                Void,
                (
                Ref{Bool},  # INIT if false compute SWCoeffs in A, else use A
                Ref{$F},    # X    sample
                Ref{$I},    # N    samples length
                Ref{$I},    # N1   (upper) uncensored data length
                Ref{$I},    # N2   length of A
                Ref{$F},    # A    A
                Ref{$F},    # W    W-statistic
                Ref{$F},    # PW   p-value
                Ref{$I},    # IFAULT error code (see swilk.f for meaning)
                ),
                !computeA, X, length(X), length(X), length(A), A, w, pval, ifault)
            return (w[], pval[], ifault[], A)
        end
        swilkfort(X::Vector{$F}) = swilkfort!(X, zeros($F, div(length(X),2)))
    end
end
=#


#

julia> using Checkers
julia> f(x) = x^2
julia> @test_formany -10<x<10, f(x) >= 0
#Test Passed
  Expression: (:((-10 < x < 10, f(x) >= 0)), :(mode = test_formany))

julia> @test_forall x in -10:10, f(x) >= 0
#Test Passed
  Expression: (x in -10:10, f(x) >= 0)
    @test_forall x in 0:10, f(x) >= 0

julia> @test_forall x in -10:10, f(x) > 0 #Should fail b/c f(0) = 0
#Test Failed
  Expression: (x in -10:10, f(x) > 0)

#tests the conditional property of Julia
 julia> @test_formany ntests = 100 maxtests = 100_000 0 < x < 10,  0 < y < 10,  x < y --> x^2 < y^2
#100 test will be run
#Passes if 100 tests meet the crierea of holding true. Will fail becuase generating 100 pairs of x,y, some will bound to fail condition

#another test the conditional property of Julia
false.
#100,000 tests will performed until 100 passing test
julia> @test_formany ntests = 100 maxtests = 100_000 0 < x < 10,  0 < y < 10,  x < y --> x^2 < y^2
#Same as top but will pass because 100000test will be performed.


@doc raw"""
`SDE`: Stratonovich Stochastic Differential Equation

Defines a stochastic differential initial value problem
```math
\begin{align*}
\dq (t) &= v(t, q(t)) \, dt + B(t, q(t)) \circ dW , & q(t_{0}) &= q_{0} ,
\end{align*}
```
with drift vector field ``v``, diffusion matrix ``B``,
initial conditions ``q_{0}``, the dynamical variable ``q``
taking values in ``\mathbb{R}^{d}``, and the m-dimensional Wiener process W

### Fields

* `d`:  dimension of dynamical variable ``q`` and the vector field ``v``
* `m`:  dimension of the Wiener process
* `ni`: number of initial conditions
* `ns`: number of sample paths
* `v`:  function computing the deterministic vector field
* `B`:  function computing the d x m diffusion matrix
* `t₀`: initial time
* `q₀`: initial condition for dynamical variable ``q`` (may be a random variable itself)

### Parameters

* `N`: dimension of nitial condition array: N=1 - single, N=2 - multiple

The functions `v` and `B`, providing the drift vector field and diffusion matrix,
`v(t, q, v)` and `B(t, q, B, col=0)`, where `t` is the current time, `q` is the
current solution vector, and `v` and `B` are the variables which hold the result
of evaluating the vector field ``v`` and the matrix ``B`` on `t` and `q` (if col==0),
or the column col of the matrix B (if col>0).

### Example

```julia
    function v(λ, t, q, v)
        v[1] = λ*q[1]
        v[2] = λ*q[2]
    end

    function B(μ, t, q, B; col=0)
        if col==0 #whole matrix
            B[1,1] = μ*q[1]
            B[2,1] = μ*q[2]
        elseif col==1
            #just first column
        end
    end

    t₀ = 0.
    q₀ = [1., 1.]
    λ  = 2.
    μ  = 1.

    v_sde = (t, q, v) -> v(λ, t, q, v)
    B_sde = (t, q, B) -> B(μ, t, q, B)

    sde = SDE(v_sde, B_sde, t₀, q₀)
```
"""
struct SDE{dType <: Number, tType <: Number, vType <: Function, BType <: Function,
           pType <: Union{Tuple,Nothing}, N} <: AbstractEquationSDE{dType, tType}

    d::Int
    m::Int
    ni::Int
    ns::Int
    v::vType
    B::BType
    t₀::tType
    q₀::Array{dType, N}
    parameters::pType
    periodicity::Vector{dType}

    function SDE(m, ns, v::vType, B::BType, t₀::tType, q₀::DenseArray{dType};
                 parameters=nothing, periodicity=zeros(dType,size(q₀,1))) where {
                        dType <: Number, tType <: Number, vType <: Function, BType <: Function}

        N  = ndims(q₀)
        d  = size(q₀,1)
        ni = size(q₀,2)

        @assert N ∈ (1,2)
        @assert ni ≥ 1
        @assert ns ≥ 1
        @assert ni == 1 || ns == 1
        # either multiple deterministic initial conditions and one sample path
        # or one deterministic initial condition and multiple sample paths

        new{dType,tType,vType,BType,typeof(parameters),N}(d, m, ni, ns, v, B, t₀, q₀, parameters, periodicity)
    end
end

function SDE(m::Int, ns::Int, v::Function, B::Function, q₀::DenseArray{DT}; kwargs...) where {DT}
    SDE(m, ns, v, B, zero(DT), q₀; kwargs...)
end

Base.hash(sde::SDE, h::UInt) = hash(sde.d, hash(sde.m, hash(sde.ni, hash(sde.ns,
                               hash(sde.v, hash(sde.B, hash(sde.t₀, hash(sde.q₀,
                               hash(sde.parameters, hash(sde.periodicity, h))))))))))

Base.:(==)(sde1::SDE, sde2::SDE) = (
                                sde1.d == sde2.d
                             && sde1.m == sde2.m
                             && sde1.ni == sde2.ni
                             && sde1.ns == sde2.ns
                             && sde1.v == sde2.v
                             && sde1.B == sde2.B
                             && sde1.t₀ == sde2.t₀
                             && sde1.q₀ == sde2.q₀
                             && sde1.parameters == sde2.parameters
                             && sde1.periodicity == sde2.periodicity)

function Base.similar(sde::SDE, q₀::AbstractArray)
    similar(sde, sde.t₀, q₀)
end

function Base.similar(sde::SDE, q₀::AbstractArray, ns::Int)
    similar(sde, sde.t₀, q₀, ns)
end

function Base.similar(sde::SDE, t₀::TT, q₀::AbstractArray{DT,N}, ns::Int=(N > 1 ? 1 : sde.ns)) where {DT <: Number, TT <: Number, N}
    @assert sde.d == size(q₀,1)
    SDE(sde.m, ns, sde.v, sde.B, t₀, q₀; parameters=sde.parameters, periodicity=sde.periodicity)
end

@inline Base.ndims(sde::SDE) = sde.d

@inline periodicity(equation::SDE) = equation.periodicity


using Plots, Random, Statistics

N = 1000000
σ_η = 0.01  # variance
σ_ε = 0.05  # variance

# IID process
y = sqrt(σ_ε).*randn(N)
# Check data generation
abs(var(y) - σ_ε) < 0.001

# Plot discretizaiton vs simulated data
tauc = tauchen(19,0.,sqrt(σ_ε))
histogram(y, bins = range(tauc.state_values[1].-step(tauc.state_values)/2, step=step(tauc.state_values),
                    length=length(tauc.state_values)+1), normalize = :probability)
scatter!(tauc.state_values,tauc.p[1,:])

rouw = rouwenhorst(19,0.,sqrt(σ_ε))
histogram!(y, bins = range(rouw.state_values[1].-step(rouw.state_values)/2, step=step(rouw.state_values),
                    length=length(rouw.state_values)+1), normalize = :probability, alpha = 0.5)
compare = scatter!(rouw.state_values,rouw.p[1,:],alpha=0.5)
plot(compare)

# Sum of IID processes
y = sqrt(σ_ε).*randn(N) + sqrt(σ_η).*randn(N)
# Check data generation
abs(var(y) - (σ_ε + σ_η)) < 0.001

# Plot discretizaiton vs simulated data
tauc_ε = tauchen(19,0.,sqrt(σ_ε))
tauc_η = tauchen(19,0.,sqrt(σ_η))
plot(tauc_ε.state_values, tauc_η.state_values)
histogram(y, bins = range(tauc.state_values[1].-step(tauc.state_values)/2, step=step(tauc.state_values),
                    length=length(tauc.state_values)+1), normalize = :probability)
scatter!(tauc.state_values,tauc.p[1,:])

rouw = rouwenhorst(19,0.,sqrt(σ_ε))
histogram!(y, bins = range(rouw.state_values[1].-step(rouw.state_values)/2, step=step(rouw.state_values),
                    length=length(rouw.state_values)+1), normalize = :probability, alpha = 0.5)
compare = scatter!(rouw.state_values,rouw.p[1,:],alpha=0.5)
plot(compare)

# With persistent process
κ = zeros(Float64,N,T_ret-1) # *** Get from data
σ_η = 0.01
σ_z0 = 0.15
σ_ε = 0.05

function net_labor_income(N, T_ret, σ_ε, σ_η, σ_z0, κ)
    Y = zeros(Float64, N, T_ret-1)
    z_t1 = sqrt(σ_z0) .* randn(N)

    for t in 1:(T_ret-1)
        z_t = z_t1 + sqrt(σ_η) .* randn(N)
        Y[:,t] = exp.( κ[t] .+ z_t + sqrt(σ_ε) .* randn(N) )
        z_t1 = copy(z_t)
    end
    return Y
end

N = 100000
T_ret = 2
Y_l = net_labor_income(N, T_ret, σ_ε, σ_η, 0., κ)
histogram(Y_l, alpha = 0.5)
tauc_ε = tauchen(19,0.,σ_ε)
tauc_η1 = tauchen(39,0.,var(Y_l[:,1]))

# unconditional variance of z_t
T_ret = 40
ρ = 1.
σ

plot(σ)

using QuantEcon
rouwenhorst(19,0.,sqrt(0.2)).p == rouwenhorst(19,0.,sqrt(0.1)).p

import QuantEcon.rouwenhorst
function rouwenhorst(N::Integer, ρ::Real, σ::Real, μ::Real=0.0; stationary::Bool = true)
    if stationary
        σ_y = σ / sqrt(1-ρ^2)
        m = μ / (1 - ρ)
    else
        σ_y = σ
        m = 0.0  #μ not used if non-stationary
    end

    p  = (1+ρ)/2
    Θ = [p 1-p; 1-p p]
    ψ = sqrt(N-1) * σ_y
    state_values, p = _rouwenhorst(p, p, m, ψ, N)
    MarkovChain(p, state_values)
end

function _rouwenhorst(p::Real, q::Real, m::Real, Δ::Real, n::Integer)
    if n == 2
        return [m-Δ, m+Δ],  [p 1-p; 1-q q]
    else
        _, θ_nm1 = _rouwenhorst(p, q, m, Δ, n-1)
        θN = p    *[θ_nm1 zeros(n-1, 1); zeros(1, n)] +
             (1-p)*[zeros(n-1, 1) θ_nm1; zeros(1, n)] +
             q    *[zeros(1, n); zeros(n-1, 1) θ_nm1] +
             (1-q)*[zeros(1, n); θ_nm1 zeros(n-1, 1)]

        θN[2:end-1, :] ./= 2

        return range(m-Δ, stop=m+Δ, length=n), θN
    end
end

function rouwenhorst(N::Integer, ρ::Real, σ::Real, σ_0::Real, T::Int)
    σs = zeros(eltype(σ_0), T, 1)
    σs[1] = σ_0
    [ σs[t] = ρ^2*σs[t-1] + σ  for t in 2:T ] # from Fella et al. 2019
    θNs = rouwenhorst.(N, ρ, σs, stationary = false)
    return θNs
    # unnecessary to run this multiple times because transition matrix is same
end

T_ret = 40
σ_η, σ_z0
test = rouwenhorst(19, 1., σ_η, σ_z0, T_ret)
test[1].state_values
test[2].state_values
test[1].p == test[30].p

stop

plot(rouwenhorst(19,0.,sqrt(σ_ε),2.).state_values)
y = μ + ρ

#old

#
# σ_η = 0.01
# σ_ε = 0.05
#
# # Set up check
# N = 1000000
# y = sqrt(σ_ε).*randn(N)
# abs(var(y) - σ_ε) < 0.001
# tauc = tauchen(19,0.,sqrt(σ_ε))
# rouw = rouwenhorst(19,0.,sqrt(σ_ε))
# histogram(y, bins = range(tauc.state_values[1].-step(tauc.state_values)/2, step=step(tauc.state_values),
#                     length=length(tauc.state_values)+1), normalize = :probability)
# scatter!(tauc.state_values,tauc.p[1,:])
# histogram(y, bins = range(rouw.state_values[1].-step(rouw.state_values)/2, step=step(rouw.state_values),
#                     length=length(rouw.state_values)+1), normalize = :probability, alpha = 0.5)
# compare = scatter!(rouw.state_values,rouw.p[1,:],alpha=0.5)
# plot(compare)
#
#
#
#
#
# π_εr = rouwenhorst(19,0.,sqrt(σ_ε)).p
# distt = π_εt[10,:]
# distr = π_εr[10,:]
# plot()
# # histogram!(y, bins = range(-3*sqrt(σ_ε)-3*sqrt(σ_ε)/19, stop = 3*sqrt(σ_ε)+3*sqrt(σ_ε)/19, length = 20), normalize = :probability)
# histogram!(y, bins = range(-3*sqrt(var(y))-3*sqrt(var(y))/19, stop = 3*sqrt(var(y))+3*sqrt(var(y))/19, length = 20), normalize = :probability)
# scatter!(range(-3*sqrt(var(y)), stop = 3*sqrt(var(y)), length = 19),distt)
# sqrt(var(y))*sqrt(18)
# scatter!(range(-3*sqrt(var(y)), stop = 3*sqrt(var(y)), length = 19),distr)
# ylims!(0.,0.2)
# count( x -> (x < sqrt(var(y))/6 && x > -sqrt(var(y))/6), y)/N
# π_ε[10,10]
#
#
#
# # test sum of random vars
# σ_η = 0.2
# # σ_ε
# N = 1000000
# y = sqrt(σ_ε).*randn(N) .+ sqrt(σ_η).*randn(N)
# abs(var(y) - (σ_ε + σ_η)) < 0.001
# π_ε = tauchen(19,0.,sqrt(σ_ε)).p
# π_η = tauchen(19,0.,sqrt(σ_η),5).p
# dist = (π_ε * π_η)[10,:]
# # count( x -> (x < sqrt(σ_ε)/6 && x > -sqrt(σ_ε)/6), y)/N
# # count( x -> (x < sqrt(var(y))/6 && x > -sqrt(var(y))/6), y)/N
# (π_ε * π_η)[10,10]
# plot()
# histogram!(y, bins = range(-3*sqrt(var(y))-3*sqrt(var(y))/19, stop = 3*sqrt(var(y))+3*sqrt(var(y))/19, length = 20), normalize = :probability)
# scatter!(range(-3*sqrt(var(y)), stop = 3*sqrt(var(y)), length = 19),dist)
# ylims!(0.,0.2)
#
#
#
#
#
#
#
# y = sqrt(σ_ε).*randn(N) .+ sqrt(σ_η).*randn(N)
# histogram(y, bins = 100, normalize = true)
# count( x -> (x < 0 && x > -0.05), y)/N
# π_ε = rouwenhorst(19,0.,σ_ε)
# π_η = rouwenhorst(19,0.,σ_η)
#
#
# plot(π_η[1])
# π_η.p[1,:] .== π_η.p[2]
# π_η.p[10,10]


using OrdinaryDiffEq, DiffEqBase, Base.Test
srand(100)

linear = (t,u) -> (1.01*u)
analytic_linear = (t,u0) -> u0*exp(1.01*t)
prob = ODETestProblem(linear,1/2,analytic_linear,(0.0,1.0))

sol =solve(prob,Tsit5(),dt=1//2^(6),tstops=[1/2])

@test 1//2 ∈ sol.t

sol =solve(prob,RK4(),dt=1//3,tstops=[1/2])

@test sol.t == [0,1/3,1/2,1/3+1/2,1]

sol =solve(prob,RK4(),tstops=[1/5,1/4,1/3,1/2,3/4])

@test sol.t == [0,1/5,1/4,1/3,1/2,3/4,1]

sol =solve(prob,RK4(),tstops=[0,1/5,1/4,1/3,1/2,3/4,1])

@test sol.t == [0,1/5,1/4,1/3,1/2,3/4,1]

sol = solve(prob,RK4(),tstops=0:1//16:1)

@test sol.t == collect(0:1//16:1)

sol = solve(prob,RK4(),tstops=linspace(0,1,100))

@test sol.t == collect(linspace(0,1,100))


### Types

immutable SparseVector{Tv,Ti<:Integer} <: AbstractSparseVector{Tv,Ti}
    n::Int              # the number of elements
    nzind::Vector{Ti}   # the indices of nonzeros
    nzval::Vector{Tv}   # the values of nonzeros

    function SparseVector(n::Integer, nzind::Vector{Ti}, nzval::Vector{Tv})
        n >= 0 || throw(ArgumentError("The number of elements must be non-negative."))
        length(nzind) == length(nzval) ||
            throw(DimensionMismatch("The lengths of nzind and nzval are inconsistent."))
        new(convert(Int, n), nzind, nzval)
    end
end

SparseVector{Tv,Ti}(n::Integer, nzind::Vector{Ti}, nzval::Vector{Tv}) =
    SparseVector{Tv,Ti}(n, nzind, nzval)

### Basic properties

length(x::SparseVector) = x.n
size(x::SparseVector) = (x.n,)
nnz(x::SparseVector) = length(x.nzval)
countnz(x::SparseVector) = countnz(x.nzval)
nonzeros(x::SparseVector) = x.nzval
nonzeroinds(x::SparseVector) = x.nzind

### Construct empty sparse vector

sparsevector{T}(::Type{T}, len::Integer) = SparseVector(len, Int[], T[])

### Construction from lists of indices and values

function _sparsevector!{Ti<:Integer}(I::Vector{Ti}, V::Vector, len::Integer)
    # pre-condition: no duplicate indexes in I
    if !isempty(I)
        p = sortperm(I)
        permute!(I, p)
        permute!(V, p)
    end
    SparseVector(len, I, V)
end

function _sparsevector!{Tv,Ti<:Integer}(I::Vector{Ti}, V::Vector{Tv}, len::Integer, combine::BinaryOp)
    if !isempty(I)
        p = sortperm(I)
        permute!(I, p)
        permute!(V, p)
        m = length(I)

        # advances to the first non-zero element
        r = 1     # index of last processed entry
        while r <= m
            if V[r] == zero(Tv)
                r += 1
            else
                break
            end
        end
        r > m && SparseVector(len, Ti[], Tv[])

        # move r-th to l-th
        l = 1       # length of processed part
        i = I[r]    # row-index of current element
        if r > 1
            I[l] = i; V[l] = V[r]
        end

        # main loop
        while r < m
            r += 1
            i2 = I[r]
            if i2 == i  # accumulate r-th to the l-th entry
                V[l] = call(combine, V[l], V[r])
            else  # advance l, and move r-th to l-th
                pv = V[l]
                if pv != zero(Tv)
                    l += 1
                end
                i = i2
                if l < r
                    I[l] = i; V[l] = V[r]
                end
            end
        end
        if V[l] == zero(Tv)
            l -= 1
        end
        if l < m
            resize!(I, l)
            resize!(V, l)
        end
    end
    SparseVector(len, I, V)
end

function sparsevector{Tv,Ti<:Integer}(I::AbstractVector{Ti}, V::AbstractVector{Tv}, combine::BinaryOp)
    length(I) == length(V) ||
        throw(DimensionMismatch("The lengths of I and V are inconsistent."))
    len = 0
    for i in I
        i >= 1 || error("Index must be positive.")
        if i > len
            len = i
        end
    end
    _sparsevector!(_copy_convert(Ti, I), _copy_convert(Tv, V), len, combine)
end

function sparsevector{Tv,Ti<:Integer}(I::AbstractVector{Ti}, V::AbstractVector{Tv}, len::Integer, combine::BinaryOp)
    length(I) == length(V) ||
        throw(DimensionMismatch("The lengths of I and V are inconsistent."))
    maxi = convert(Ti, len)
    for i in I
        1 <= i <= maxi || error("An index is out of bound.")
    end
    _sparsevector!(_copy_convert(Ti, I), _copy_convert(Tv, V), len, combine)
end

sparsevector{Ti<:Integer}(I::AbstractVector{Ti}, V::AbstractVector) =
    sparsevector(I, V, AddFun())

sparsevector{Ti<:Integer}(I::AbstractVector{Ti}, V::AbstractVector, len::Integer) =
    sparsevector(I, V, len, AddFun())

sparsevector{Ti<:Integer}(I::AbstractVector{Ti}, v::Number, combine::BinaryOp) =
    sparsevector(I, fill(v, length(I)), combine)

sparsevector{Ti<:Integer}(I::AbstractVector{Ti}, v::Number, len::Integer, combine::BinaryOp) =
    sparsevector(I, fill(v, length(I)), len, combine)

sparsevector{Ti<:Integer}(I::AbstractVector{Ti}, v::Number) =
    sparsevector(I, v, AddFun())

sparsevector{Ti<:Integer}(I::AbstractVector{Ti}, v::Number, len::Integer) =
    sparsevector(I, v, len, AddFun())

### Construction from dictionary

function sparsevector{Tv,Ti<:Integer}(dict::Associative{Ti,Tv})
    m = length(dict)
    nzind = Array(Ti, m)
    nzval = Array(Tv, m)

    cnt = 0
    len = zero(Ti)
    for (k, v) in dict
        k >= 1 || error("Index must be positive.")
        if k > len
            len = k
        end
        if v != zero(v)
            cnt += 1
            @inbounds nzind[cnt] = k
            @inbounds nzval[cnt] = v
        end
    end
    resize!(nzind, cnt)
    resize!(nzval, cnt)
    _sparsevector!(nzind, nzval, len)
end

function sparsevector{Tv,Ti<:Integer}(dict::Associative{Ti,Tv}, len::Integer)
    m = length(dict)
    nzind = Array(Ti, m)
    nzval = Array(Tv, m)

    cnt = 0
    maxk = convert(Ti, len)
    for (k, v) in dict
        1 <= k <= maxk || error("An index (key) is out of bound.")
        if v != zero(v)
            cnt += 1
            @inbounds nzind[cnt] = k
            @inbounds nzval[cnt] = v
        end
    end
    resize!(nzind, cnt)
    resize!(nzval, cnt)
    _sparsevector!(nzind, nzval, len)
end

### Element access

function setindex!{Tv,Ti<:Integer}(x::SparseVector{Tv,Ti}, v::Tv, i::Ti)
    nzind = nonzeroinds(x)
    nzval = nonzeros(x)

    m = length(nzind)
    k = searchsortedfirst(nzind, i)
    if 1 <= k <= m && nzind[k] == i  # i found
        if v == zero(v)
            deleteat!(nzind, k)
            deleteat!(nzval, k)
        else
            nzval[k] = v
        end
    else  # i not found
        if v != zero(v)
            insert!(nzind, k, i)
            insert!(nzval, k, v)
        end
    end
    x
end

setindex!{Tv, Ti<:Integer}(x::SparseVector{Tv,Ti}, v, i::Integer) =
    setindex!(x, convert(Tv, v), convert(Ti, i))

### Conversion

# convert SparseMatrixCSC to SparseVector
function convert{Tv,Ti<:Integer}(::Type{SparseVector{Tv,Ti}}, s::SparseMatrixCSC{Tv,Ti})
    size(s, 2) == 1 || throw(ArgumentError("The input argument must have a single-column."))
    SparseVector(s.m, s.rowval, s.nzval)
end

convert{Tv,Ti}(::Type{SparseVector{Tv}}, s::SparseMatrixCSC{Tv,Ti}) =
    convert(SparseVector{Tv,Ti}, s)

convert{Tv,Ti}(::Type{SparseVector}, s::SparseMatrixCSC{Tv,Ti}) =
    convert(SparseVector{Tv,Ti}, s)

# convert Vector to SparseVector

function _dense2sparsevec{Tv}(s::Vector{Tv}, initcap::Int)
    # pre-condition: initcap > 0
    n = length(s)
    cap = initcap
    nzind = Array(Int, cap)
    nzval = Array(Tv, cap)
    c = 0
    @inbounds for i = 1:n
        v = s[i]
        if v != zero(v)
            if c >= cap
                cap *= 2
                resize!(nzind, cap)
                resize!(nzval, cap)
            end
            c += 1
            nzind[c] = i
            nzval[c] = v
        end
    end
    if c < cap
        resize!(nzind, c)
        resize!(nzval, c)
    end
    SparseVector(n, nzind, nzval)
end

convert{Tv}(::Type{SparseVector{Tv,Int}}, s::Vector{Tv}) =
    _dense2sparsevec(s, max(8, div(length(s), 8)))

convert{Tv}(::Type{SparseVector{Tv}}, s::Vector{Tv}) =
    convert(SparseVector{Tv,Int}, s)

convert{Tv}(::Type{SparseVector}, s::Vector{Tv}) =
    convert(SparseVector{Tv,Int}, s)

# convert between different types of SparseVector
convert{Tv,Ti,TvS,TiS}(::Type{SparseVector{Tv,Ti}}, s::SparseVector{TvS,TiS}) =
    SparseVector{Tv,Ti}(s.n, convert(Vector{Ti}, s.nzind), convert(Vector{Tv}, s.nzval))

convert{Tv,TvS,TiS}(::Type{SparseVector{Tv}}, s::SparseVector{TvS,TiS}) =
    SparseVector{Tv,TiS}(s.n, s.nzind, convert(Vector{Tv}, s.nzval))

### Rand Construction

@compat function sprand{T}(n::Integer, p::AbstractFloat, rfn::Function, ::Type{T})
    I = randsubseq(1:convert(Int, n), p)
    V = rfn(T, length(I))
    SparseVector(n, I, V)
end

@compat function sprand(n::Integer, p::AbstractFloat, rfn::Function)
    I = randsubseq(1:convert(Int, n), p)
    V = rfn(length(I))
    SparseVector(n, I, V)
end

@compat sprand{T}(n::Integer, p::AbstractFloat, ::Type{T}) = sprand(n, p, rand, T)
@compat sprand(n::Integer, p::AbstractFloat) = sprand(n, p, rand)
@compat sprandn(n::Integer, p::AbstractFloat) = sprand(n, p, randn)


# ---
# title: 315. Count of Smaller Numbers After Self
# id: problem315
# author: zhwang
# date: 2022-03-26
# difficulty: Hard
# categories: Binary Search, Divide and Conquer, Sort, Binary Indexed Tree, Segment Tree
# link: <https://leetcode.com/problems/count-of-smaller-numbers-after-self/description/>
# hidden: true
# ---
#
# You are given an integer array _nums_ and you have to return a new _counts_
# array. The _counts_ array has the property where `counts[i]` is the number of
# smaller elements to the right of `nums[i]`.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [5,2,6,1]
#     Output: [2,1,1,0]
#     Explanation:
#     To the right of 5 there are **2** smaller elements (2 and 1).
#     To the right of 2 there is only **1** smaller element (1).
#     To the right of 6 there is **1** smaller element (1).
#     To the right of 1 there is **0** smaller element.
#
#
#
#
# **Constraints:**
#
#   * `0 <= nums.length <= 10^5`
#   * `-10^4 <= nums[i] <= 10^4`
#
#
## @lc code=start
using LeetCode

## method 1: using merge sort
function count_smaller_method1(arr::Vector{Int})::Vector{Int}
    res = fill(0, length(arr))
    function merge_sort(arr::AbstractVector)::AbstractVector
        ## arr = [(idx1, num1), ...]
        (n = length(arr)) == 1 && return arr
        mid = n >> 1
        lpart = merge_sort(@view(arr[1:mid]))
        rpart = merge_sort(@view(arr[mid + 1:end]))
        return merge_sorted(lpart, rpart)
    end
    function merge_sorted(lpart::AbstractVector, rpart::AbstractVector)::Vector{Tuple{Int,Int}}
        l1, l2 = length(lpart), length(rpart)
        combined = Vector{Tuple{Int,Int}}(undef, l1 + l2)
        p1 = p2 = 1
        for pos in eachindex(combined)
            if p2 > l2 || p1 <= l1 && lpart[p1][2] <= rpart[p2][2]
                combined[pos] = lpart[p1]
                res[lpart[p1][1]] += (p2 - 1)
                p1 += 1
            else
                combined[pos] = rpart[p2]
                p2 += 1
            end
        end
        combined
    end

    merge_sort(collect(enumerate(arr)))
    return res
end

## method 2: using binary search
function search_left_border(left::Int, right::Int, key::Function)
    ## similar to findfirst(key, left:right) but using bs method
    while left <= right
        mid = left + (right - left) >> 1
        if key(mid) ## lies in the right part
            right = mid - 1
        else
            left = mid + 1
        end
    end
    left
end
function count_smaller_method2(arr::Vector{Int})::Vector{Int}
    queue, res = Int[], Array{Int}(undef, length(reverse!(arr)))
    for (i, num) in enumerate(arr)
        loc = search_left_border(1, length(queue), x -> queue[x] >= num)
        res[i] = loc - 1
        insert!(queue, loc, num)
    end
    reverse!(arr)
    return reverse!(res)
end
## @lc code=end

## @lc test=start
@testset "315.count-of-smaller-numbers-after-self.jl" begin
    f1 = count_smaller_method1
    f2 = count_smaller_method2
    @test f1([5,2,6,1]) == f2([5,2,6,1]) == [2, 1, 1, 0]
    @test f1([2,0,1]) == f2([2,0,1]) == [2, 0, 0]
    @test f1([-1]) == f2([-1]) == [0]
    @test f1([-1, -1]) == f2([-1, -1]) == [0, 0]
    @test f1([7,6,5,4,3,2,1]) == f2([7,6,5,4,3,2,1]) == [6, 5, 4, 3, 2, 1, 0]
end
## @lc test=end

## @lc info=start
# link: [solution to 315]()
## @lc info=end


#---------------------------------------------------------------------
#	Biconnected components
#---------------------------------------------------------------------

"""
	verts2verts(EV::Lar.Cells)::Lar.Cells

Adjacency lists of vertices of a cellular 1-complex.

# Example

```julia
julia> V,(VV,EV,FV) = Lar.cuboidGrid([3,3],true);

julia> verts2verts(EV::Lar.Cells)::Lar.Cells
16-element Array{Array{Int64,1},1}:
 [2, 5]
 [1, 3, 6]
 [2, 4, 7]
 [3, 8]
 [1, 6, 9]
 [2, 5, 7, 10]
 [3, 6, 8, 11]
 [4, 7, 12]
 [5, 10, 13]
 [6, 9, 11, 14]
 [7, 10, 12, 15]
 [8, 11, 16]
 [9, 14]
 [10, 13, 15]
 [11, 14, 16]
 [12, 15]
```
"""
function verts2verts(EV::Lar.Cells)::Lar.Cells
    cscEV = Lar.characteristicMatrix(EV)
    cscVE = convert(Lar.ChainOp, cscEV')
    cscVV = cscVE * cscEV
    rows,cols,data = findnz(cscVV)
    VV = [ findnz(cscVV[k,:])[1] for k=1:size(cscVV,2) ]
    return [setdiff(vs,[k]) for (k,vs) in enumerate(VV)] # remove reflexive pairs
end

"""
	DFV_visit( VV::cells, out::Array, count::Int, visited::Array, parent::Array, d::Array, low::Array, stack::Array, u::Int )::Array

Hopcroft-Tarjan algorithm. Depth-First-Visit recursive algorithm.
"""
function DFV_visit( VV::Lar.Cells, out::Array, count::Int, visited::Array, parent::Array, d::Array, low::Array, stack::Array, u::Int )::Array

    visited[u] = true
    count += 1
    d[u] = count
    low[u] = d[u]
    for v in VV[u]
        if ! visited[v]
            push!(stack, [(u,v)])
            parent[v] = u
            DFV_visit( VV,out,count,visited,parent,d,low,stack, v )
            if low[v] >= d[u]
                push!(out, [outputComp(stack,u,v)])
            end
            low[u] = min( low[u], low[v] )
        else
            if ! (parent[u]==v) && (d[v] < d[u])
                push!(stack, [(u,v)])
            end
            low[u] = min( low[u], d[v] )
        end
    end
    out
end

visited=[]; count=0; out=[]; parent=[]; d=[];low=[]; stack=[]; u=1

"""
	outputComp(stack::Array, u::Int, v::Int)::Array

Internal output of biconnected components. Utility function for DFV_visit graph algorithm (Hopcrof, Tarjan (1973)).
"""
function outputComp(stack::Array, u::Int, v::Int)::Array
    out = []
    while true
        e = pop!(stack)[1]
        push!(out,e)
        if e == (u,v)
        	break
        end
    end
    return [out]
end

"""
	biconnectedComponent(model)

Main procedure for computation of biconnected components of a graph
represented by a LAR unsigned pair (Points, Cells).

# Example

```julia
julia> V =  [3.0  6.0  6.0  6.0  6.0  0.0  3.0  10.0  10.0  10.0  0.0  3.0  0.0;
 			 2.0  2.0  0.0  4.0  8.0  4.0  4.0   4.0   0.0   8.0  8.0  0.0  0.0]

julia> EV = [[1, 2], [2, 3], [2, 4], [5, 10], [3, 12], [5, 11], [6, 7], [8, 9],
			 [6, 13], [12, 13], [1, 7], [1, 12], [4, 7], [4, 8], [4, 5], [8, 10]]

julia> VV = [[k] for k=1:size(V,2)]

julia> Plasm.view( Plasm.numbering(3)((V,[VV, EV])) )
```

```julia
model = (V,EV)
V,EVs = biconnectedComponent(model)
EW = convert(Lar.Cells, cat(EVs))
Plasm.view( Plasm.numbering(3)((V,[VV,EW])) )
```

Another example:

```julia
V = [0. 1  0  0 -1;
	 0  0  1 -1  0]
EV = [[1,2],[2,3],[1,4],[4,5],[5,1],[1,3]  ,[2,4]]
model = (V,EV)
Plasm.view( Plasm.numbering(1)((V,[[[k] for k=1:size(V,2)],EV])) )
V,EVs = Lar.biconnectedComponent(model)

cscEV = Lar.characteristicMatrix(EV)
biconcomp = Lar.Arrangement.biconnected_components(cscEV)

Matrix(Lar.characteristicMatrix(EV))

V,EVs = biconnectedComponent(model)
EW = convert(Lar.Cells, cat(EVs))
Plasm.view( Plasm.numbering(.3)((V,[VV,EW])) )
```

"""
function biconnectedComponent(model)
    W,EV = model
    V = collect(1:size(W,2))
    count = 0
    stack,out = [],[]
    visited = [false for v in V]
    parent = Union{Int, Array{Any,1}}[[] for v in V]
    d = Any[0 for v in V]
    low = Any[0 for v in V]
    VV = Lar.verts2verts(EV)
    out = Any[]
    for u in V
        if ! visited[u]
            out = DFV_visit( VV,out,count,visited,parent,d,low,stack, u )
        end
    end
    out = [component for component in out if length(component) >= 1]
    EVs = [[map(sort∘collect,edges) for edges in cat(comp) if length(edges)>1]
    		for comp in out]
    EVs = cat(filter(x->!isempty(x), EVs))
    EVs = sort(EVs, lt=(x,y)->length(x)>length(y))
    return W, EVs
end

"""
	depth_first_search(EV::Lar.Cells [, start::Int=1])::Tuple{Lar.Cells, Lar.Cells}

*Depth-First-Search algorithm* (Tarjan, 1972) to compute a *spanning tree* of an undirected graph. The graph is entered as an array of edges, given as an array `EV` of 1-cells. Complexity linear in the number of nodes and edges: O(V, E).

Return a partition of `EV` into two arrays `spanningtree` and `fronds` which define a ``palm tree``.  The algorithm may `start` from any vertex index.

# Example

```
julia> EV = [[1,2],[2,3],[3,4],[4,5],[5,6],[1,6],
       [6,7],[4,7],[2,7],[5,8],[1,8],[3,8]];

julia> depth_first_search(EV,8)
(Array{Int64,1}[[8, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]], Array{Int64,1}[[6, 1], [7, 2], [7, 4], [5, 8], [3, 8]])

julia> depth_first_search(EV)
(Array{Int64,1}[[1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7], [5, 8]], Array{Int64,1}[[6, 1], [7, 2], [7, 4], [8, 1], [8, 3]])
```
"""
function depth_first_search(EV::Lar.Cells, start::Int=1)::
		Tuple{Lar.Cells, Lar.Cells}
	# initialize the data structures and start DFS
    VV = Lar.verts2verts(EV)
    spanningtree = Array{Int,1}[];
    fronds = Array{Int,1}[]
	number = zeros(Int, length(VV))

	function DFS(v::Int, u::Int)
	# vertex u is the father of
	# vertex v in the spanning tree being constructed
		number[v] = i; i += 1
		for w in VV[v]
			if number[w] == 0
				# w is not yet numbered
				push!(spanningtree, [v,w])
				DFS(w, v)
			elseif (number[w] < number[v]) && (w != u)
				push!(fronds, [v,w])
			end
		end
	end

	# initialization of DFS algorithm
	i = 1
	DFS(start, 1)
	return spanningtree::Lar.Cells, fronds::Lar.Cells
end

"""
	edge_biconnect(EV::Lar.Cells)

Compute maximal 2-edge-connected components of an undirected graph.

# Examples

```julia
V = hcat([[1.,2],[1,3],[0,3],[1,4],[2,3],[2,2],[1,1],[1,0],[0,1]]...)
EV = [[3,4],[2,4],[2,3],[2,5],[1,2],[5,6],[1,6],[1,7],[7,9],[8,9],[7,8]]
using Plasm
Plasm.view( Plasm.numbering(1.3)((V,[[[k] for k=1:size(V,2)], EV])) )

julia> edge_biconnect(EV::Lar.Cells)
4-element Array{Array{Array{Int64,1},1},1}:
 [[4, 2], [3, 4], [2, 3]]
 [[6, 1], [5, 6], [2, 5], [1, 2]]
 [[9, 7], [8, 9], [7, 8]]
 [[1, 7]]
```

```julia
julia> V = hcat([[0,4],[4,4],[4,0],[2,0],[0,0],[2,2],[1,2],[3,2],[1,3],[3,3]]...)
julia> EV = [[1,2],[9,10],[1,5],[7,9],[8,10],[2,3],[6,7],[6,8],[4,6],[4,5],[3,4]]
julia> Plasm.view( Plasm.numbering(1.3)((V,[[[k] for k=1:size(V,2)], EV])) )

julia> EVs = edge_biconnect(EV::Lar.Cells)
3-element Array{Array{Array{Int64,1},1},1}:
 [[8, 6], [10, 8], [9, 10], [7, 9], [6, 7]]
 [[4, 6]]
 [[5, 1], [4, 5], [3, 4], [2, 3], [1, 2]]

julia> map(Lar.depth_first_search, EVs)
```
"""
function edge_biconnect(EV::Lar.Cells)
    VV = Lar.verts2verts(EV)
	number = zeros(Int, length(VV))
	lowpt = zeros(Int, length(VV))
	components = Array{Array{Int,1},1}[]

	function biconnect(v, u)
		number[v] = i; i +=  1
		lowpt[v] = number[v]
		for w in VV[v]
			if number[w] == 0  # w is not yet numbered
				push!(stack, [v, w]) # add [v, w] to edges stack

				biconnect(w, v)

				lowpt[v] = min(lowpt[v], lowpt[w])
				if lowpt[w] ≥ number[v]
					# start new biconnected component
					component = Array{Int,1}[]
					while number[stack[end][1]] ≥ number[w]
						u1,u2 = pop!(stack)
						push!(component, [u1,u2])
					end
					v,w = pop!(stack)
					push!(component, [v,w])
					push!(components, component)
				end
			else
				if number[w] < number[v] && w ≠ u
					push!(stack, [v,w]) # add (v, w) to edge stack;
					lowpt[v] = min(lowpt[v], number[w]);
				end
			end
		end
	end

	i = 1
	stack = Array{Int64,1}[]
	for w=1:length(VV)
		if number[w] == 0 # w is not yet numbered
			biconnect(w, 0)
		end
	end
	return components
end