export  convert2affinePCE,
        calculateAffinePCE,
        assign2multi,
        evaluatePCE,
        sampleMeasure,
        samplePCE,
        mean,
        var,
        std

# auxiliary functions
function _createMethodVector(n::Int, word::String="adaptiverejection")
    [word for i in 1:n]
end

function _createMethodVector(m::ProductMeasure, word::String="adaptiverejection")
    _createMethodVector(length(m.measures), word)
end

function _createMethodVector(mop::MultiOrthoPoly, word::String="adaptiverejection")
    _createMethodVector(mop.measure, word)
end

function _checkStandardDevation(σ::Real)
    σ < 0 && throw(DomainError(σ, "σ has to be non-negative"))
    σ < 1e-4 && @warn "σ is close to zero (σ = $σ)"
end

function _checkKind(kind::String)
    lowercase(kind) ∉ ["lbub", "μσ"] && throw(DomainError(kind, "this kind is not supported"))
    lowercase(kind)
end

function _checkBounds(lb::Real, ub::Real)
    lb >= ub && throw(DomainError((lb, ub), "inconsistent bounds"))
end

function _checkNumberOfSamples(n::Int)
    n < 1 && throw(DomainError(n, "invalid number of samples"))
end
"""
    calculateAffinePCE(α::AbstractVector{<:Real})
Computes the affine PCE coefficients ``x_0`` and ``x_1`` from recurrence coefficients ``\alpha``.
"""
function calculateAffinePCE(α::AbstractVector{<:Real})
    length(α) < 1 && throw(DomainError(length(α), "not enough recursion coefficients"))
    return [α[1], 1.]
end

calculateAffinePCE(op::AbstractOrthoPoly) = calculateAffinePCE(op.α)

function calculateAffinePCE(i::Int, mop::MultiOrthoPoly)
    p = length(mop.name)
    i > p && throw(InconsistencyError("mop is $p-variate, but the PC for the $i-variate entry was requested"))
    calculateAffinePCE(mop.uni[i])
end

"""
```
convert2affinePCE(mu::Real, sigma::Real, op::AbstractCanonicalOrthoPoly; kind::String)
```
Computes the affine PCE coefficients ``x_0`` and ``x_1`` from
```math
X = a_1 + a_2 \\Xi = x_0 + x_1 \\phi_1(\\Xi),
```
where ``\\phi_1(t) = t-\\alpha_0`` is the first-order monic basis polynomial.

Works for subtypes of AbstractCanonicalOrthoPoly. The keyword `kind in ["lbub", "μσ"]`
specifies whether `p1` and `p2` have the meaning of lower/upper bounds or mean/standard deviation.
"""
function convert2affinePCE(a1::Real,a2::Real,α0::Real)
    [ a1 + α0 * a2; a2 ]
end

function convert2affinePCE(mu::Real, sigma::Real, op::GaussOrthoPoly)
    _checkStandardDevation(sigma)
    convert2affinePCE(mu, sigma, first(op.α))
end

function convert2affinePCE(par1::Real, par2::Real, op::Uniform01OrthoPoly; kind::String="lbub")
    kind = _checkKind(kind)
    a1, a2 =
        if kind == "lbub"
            _checkBounds(par1, par2)
            par1, par2 - par1
        elseif kind == "μσ"
            _checkStandardDevation(par2)
            par1 - sqrt(3)*par2 , 2*sqrt(3)*par2
        end
    convert2affinePCE(a1, a2, first(op.α))
end

function convert2affinePCE(par1::Real, par2::Real, op::Uniform_11OrthoPoly; kind::String="lbub")
    kind = _checkKind(kind)
    a1, a2 =
        if kind == "lbub"
            _checkBounds(par1, par2)
            0.5 * (par1 + par2), 0.5 * (par2 - par1)
        elseif kind == "μσ"
            _checkStandardDevation(par2)
            par1, sqrt(3) * par2
        end
    convert2affinePCE(a1, a2, first(op.α))
end

function convert2affinePCE(p1::Real, p2::Real, op::Beta01OrthoPoly; kind::String="lbub")
    kind = _checkKind(kind)
    α, β = op.measure.ashapeParameter, op.measure.bshapeParameter
    a1, a2 =
        if kind == "lbub"
            _checkBounds(p1, p2)
            a1, a2 = p1, p2-p1
        elseif kind == "μσ"
            _checkStandardDevation(p2)
            a1, a2 = p1 - sqrt(α / β) * sqrt(1 + α + β) * p2, (α + β) * sqrt((α + β + 1) / (α * β)) * p2
        end
    convert2affinePCE(a1, a2, first(op.α))
end

function convert2affinePCE(p1::Real, p2::Real, op::GammaOrthoPoly)
    throw(error("convert2affine not yet implemented for $(typeof(op))"))
end

function convert2affinePCE(p1::Real, p2::Real, op::LogisticOrthoPoly)
    _checkStandardDevation(p2)
    convert2affinePCE(p1, p2, first(op.α))
end

function assign2multi(x::AbstractVector{<:Real},i::Int,ind::AbstractMatrix{<:Int})
    l, p = size(ind)
    nx, deg = length(x), ind[end,end]
    nx > deg + 1 && throw(DomainError(nx, "inconsistent number of coefficients ($nx vs $(deg+1))"))
    i > p && throw(DomainError((i,p), "basis is $p-variate, you requested $i-variate"))
    myind = findUnivariateIndices(i,ind)[1:nx]
    y = spzeros(Float64,l)
    y[myind] = x
    return y
end

"""
__Univariate__
```
sampleMeasure(n::Int,name::String,w::Function,dom::Tuple{<:Real,<:Real},symm::Bool,d::Dict;method::String="adaptiverejection")
sampleMeasure(n::Int,m::Measure;method::String="adaptiverejection")
sampleMeasure(n::Int,op::AbstractOrthoPoly;method::String="adaptiverejection")
```
Draw `n` samples from the measure `m` described by its
- `name`
- weight function `w`,
- domain `dom`,
- symmetry property `symm`,
- and, if applicable, parameters stored in the dictionary `d`.
By default an adaptive rejection sampling method is used (from [AdaptiveRejectionSampling.jl](https://github.com/mauriciogtec/AdaptiveRejectionSampling.jl)),
unless it is a common random variable for which [Distributions.jl](https://github.com/JuliaStats/Distributions.jl) is used.

The function is dispatched to accept `AbstractOrthoPoly`.

__Multivariate__
```
sampleMeasure(n::Int,m::ProductMeasure;method::Vector{String}=["adaptiverejection" for i=1:length(m.name)])
sampleMeasure(n::Int,mop::MultiOrthoPoly;method::Vector{String}=["adaptiverejection" for i=1:length(mop.meas.name)])
```
Multivariate extension which provides array of samples with `n` rows and
as many columns as the multimeasure has univariate measures.
"""
function sampleMeasure(n::Int, w::Function, dom::Tuple{<:Real,<:Real}; method::String="adaptiverejection")
    _checkNumberOfSamples(n)
    method = lowercase(method)

    if method == "adaptiverejection"
        # works only if w is log-concave
        sampler = RejectionSampler(w, dom)
        return run_sampler!(sampler, n)
    elseif method == "rejection"
        # all purpose method but needs a solid envelope PDF
        throw(error("method $method not yet implemented"))
    elseif method == "inversecdf"
        throw(error("method $method not yet implemented"))
        # requires CDF and reliable root-finding method
    else
        throw(error("method $method not implemented"))
    end
end

function sampleMeasure(n::Int, meas::AbstractMeasure; method::String="adaptiverejection")
    sampleMeasure(n, meas.w, meas.dom; method=method)
end

sampleMeasure(n::Int, op::AbstractOrthoPoly; method::String="adaptiverejection") = sampleMeasure(n,op.measure; method=method)
sampleMeasure(n::Int, op::AbstractCanonicalOrthoPoly) = sampleMeasure(n, op.measure::AbstractCanonicalMeasure)

function sampleMeasure(n::Int, dist::Distribution{Univariate,Continuous})
    _checkNumberOfSamples(n)
    rand(dist, n)
end

sampleMeasure(n::Int, meas::GaussMeasure) = sampleMeasure(n, Normal())
sampleMeasure(n::Int, meas::Uniform01Measure) = sampleMeasure(n, Uniform())
sampleMeasure(n::Int, meas::Beta01Measure) = sampleMeasure(n, Beta(meas.ashapeParameter, meas.bshapeParameter))
sampleMeasure(n::Int, meas::GammaMeasure) = sampleMeasure(n, Gamma(meas.shapeParameter, 1/meas.rateParameter))
sampleMeasure(n::Int, meas::LogisticMeasure) = sampleMeasure(n, Logistic())

function sampleMeasure(n::Int, meas::AbstractCanonicalMeasure; method::String="adaptiverejection")
    @warn "ignoring keyword method; sampling from Distributions.jl instead"
    sampleMeasure(n, meas)
end

function sampleMeasure(n::Int, measure::ProductMeasure; method::Vector{String}=_createMethodVector(measure))
    samples = Matrix{Float64}(undef, n, 0)
    for (k, unimeasure) in enumerate(measure.measures)
        samples = hcat(samples, sampleMeasure(n, unimeasure; method = method[k]))
    end
    samples
end

sampleMeasure(n::Int, mop::MultiOrthoPoly; method::Vector{String}=_createMethodVector(mop)) = sampleMeasure(n,mop.measure; method=method)

"""
    evaluatePCE(x::AbstractVector{<:Real},ξ::AbstractVector{<:Real},α::AbstractVector{<:Real},β::AbstractVector{<:Real})
Evaluation of polynomial chaos expansion
```math
\\mathsf{x} = \\sum_{i=0}^{L} x_i \\phi_i{\\xi_j},
```
where `L+1 = length(x)` and ``x_j`` is the ``j``th sample where ``j=1,\\dots,m``
with `m = length(ξ)`.
"""
function evaluatePCE(x::AbstractVector{<:Real},ξ::AbstractVector{<:Real},α::AbstractVector{<:Real},β::AbstractVector{<:Real})
    length(α) != length(β) && throw(InconsistencyError("inconsistent number of recurrence coefficients"))
    Nsmpl = length(ξ)
    _checkNumberOfSamples(Nsmpl)
    Nx, Nrec = length(x), length(α)
    Nx > Nrec && throw(InconsistencyError("not enough recursion coefficients"))
    if Nrec > Nx α,β = α[1:Nx], β[1:Nx] end
    ϕ = zeros(Float64, Nsmpl, Nx)
    for n in 1:Nx
        ϕ[:,n] = evaluate(n-1,ξ,α,β)
    end
    ϕ*x
end
evaluatePCE(x::AbstractVector{<:Real},ξ::Real,α::AbstractVector{<:Real},β::AbstractVector{<:Real}) = evaluatePCE(x,[ξ],α,β)
evaluatePCE(x::AbstractVector{<:Real},ξ::AbstractVector{<:Real},op::AbstractOrthoPoly) = evaluatePCE(x,ξ,op.α,op.β)

function evaluatePCE(x::AbstractVector{<:Real},ξ::Matrix{<:Real},α::AbstractVector{<:AbstractVector{<:Real}},β::AbstractVector{<:AbstractVector{<:Real}},ind::AbstractMatrix{Int})
    Nsmpl = size(ξ,1)
    _checkNumberOfSamples(Nsmpl)
    !(length(α) == length(β) == size(ξ,2) == size(ind,2)) && throw(InconsistencyError("inconsistent number of coefficients"))
    Nx = length(x)
    Nx > size(ind,1) && throw(InconsistencyError("too few pc coefficients (resp: too small basis)"))
    ϕ = zeros(Float64,Nsmpl,Nx)
    for n in 1:Nx
        ϕ[:,n] = evaluate(ind[n,:],ξ,α,β)
    end
    ϕ*x
end
function evaluatePCE(x::AbstractVector{<:Real},ξ::Matrix{<:Real},mOP::MultiOrthoPoly)
    a,b = coeffs(mOP)
    evaluatePCE(x,ξ,a,b,mOP.ind)
end

"""
__Univariate__
```
samplePCE(n::Int,x::AbstractVector{<:Real},op::AbstractOrthoPoly;method::String="adaptiverejection")
```
Combines [`sampleMeasure`](@ref) and [`evaluatePCE`](@ref), i.e. it first draws `n` samples
from the measure, then evaluates the PCE for those samples.

__Multivariate__
```
samplePCE(n::Int,x::AbstractVector{<:Real},mop::MultiOrthoPoly;method::Vector{String}=["adaptiverejection" for i=1:length(mop.meas.name)])
```
"""
function samplePCE(n::Int,x::AbstractVector{<:Real},op::AbstractOrthoPoly;method::String="adaptiverejection")
    ξ = sampleMeasure(n,op;method=method)
    evaluatePCE(x,ξ,op)
end

function samplePCE(n::Int,x::AbstractVector{<:Real},op::AbstractCanonicalOrthoPoly;method::String="adaptiverejection")
    ξ = sampleMeasure(n,op)
    evaluatePCE(x,ξ,op)
end

function samplePCE(n::Int,x::AbstractVector{<:Real},mop::MultiOrthoPoly;method::Vector{String} = _createMethodVector(mop))
    ξ = sampleMeasure(n, mop; method=method)
    evaluatePCE(x, ξ, mop)
end

"""
__Univariate__
```
mean(x::AbstractVector,op::AbstractOrthoPoly)
```
__Multivariate__
```
mean(x::AbstractVector,mop::MultiOrthoPoly)
```
compute mean of random variable with PCE `x`
"""
mean(x::AbstractVector,op::AbstractOrthoPoly) = x[1]*computeSP2(0,op.β)

function mean(x::AbstractVector,mop::MultiOrthoPoly)
    nunc = length(mop.uni)
    x[1]*computeSP(zeros(Int64,nunc),mop)
end

"""
__Univariate__
```
var(x::AbstractVector,op::AbstractOrthoPoly)
var(x::AbstractVector,t2::Tensor)
```
__Multivariate__
```
var(x::AbstractVector,mop::MultiOrthoPoly)
var(x::AbstractVector,t2::Tensor)
```
compute variance of random variable with PCE `x`
"""
function var(x::AbstractVector,op::AbstractOrthoPoly)
    t = computeSP2(op)
    # length(t2) > length(x) && throw(InconsistencyError("cannot compute variance; too many PCE coefficients"))
    sum( x[i]^2*t[i] for i in 2:length(x) )
end

function var(x::AbstractVector,t2::AbstractTensor)
    sum( x[i]^2*t2.get([i-1,i-1]) for i in 2:length(x) )
end

function var(x::AbstractVector,mop::MultiOrthoPoly)
    length(x) > size(mop.ind,1) && throw(InconsistencyError("cannot compute variance; too many PCE coefficients"))
    var(x,Tensor(2,mop))
end

"""
__Univariate__
```
std(x::AbstractVector,op::AbstractOrthoPoly)
```
__Multivariate__
```
std(x::AbstractVector,mop::MultiOrthoPoly)
```
compute standard deviation of random variable with PCE `x`
"""
std(x::AbstractVector,op::AbstractOrthoPoly) = sqrt(var(x,op))
std(x::AbstractVector,mop::MultiOrthoPoly) = sqrt(var(x,mop))


######### Stan program example  ###########

using Stan, Mamba

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  binorm = "
  transformed data {
      matrix[2,2] Sigma;
      vector[2] mu;

      mu[1] <- 0.0;
      mu[2] <- 0.0;
      Sigma[1,1] <- 1.0;
      Sigma[2,2] <- 1.0;
      Sigma[1,2] <- 0.10;
      Sigma[2,1] <- 0.10;
  }
  parameters {
      vector[2] y;
  }
  model {
        y ~ multi_normal(mu,Sigma);
  }
  "

  global stanmodel, rc, sim1, sim
  stanmodel = Stanmodel(name="binormal", model=binorm, Sample(save_warmup=true));

  rc, sim1 = stan(stanmodel, CmdStanDir=CMDSTAN_HOME)

  if rc == 0
    ## Subset Sampler Output
    sim = sim1[1:size(sim1, 1), ["lp__", "y.1", "y.2"], 1:size(sim1, 3)]

    describe(sim)
    println()

    ## Brooks, Gelman and Rubin Convergence Diagnostic
    try
      gelmandiag(sim1, mpsrf=true, transform=true) |> display
    catch e
      #println(e)
      gelmandiag(sim, mpsrf=false, transform=true) |> display
    end

    ## Geweke Convergence Diagnostic
    gewekediag(sim) |> display

    ## Highest Posterior Density Intervals
    hpd(sim) |> display

    ## Cross-Correlations
    cor(sim) |> display

    ## Lag-Autocorrelations
    autocor(sim) |> display

    ## Plotting
    p = plot(sim, [:trace, :mean, :density, :autocor], legend=true);
    draw(p, ncol=4, filename="$(stanmodel.name)-summaryplot", fmt=:svg)
    draw(p, ncol=4, filename="$(stanmodel.name)-summaryplot", fmt=:pdf)
  end # rc == 0
end # cd


## Bulge chasing for the RDS and CSS case

## A bulge is introduced by a unitary transformation. In the single shift case this is just U^t A U; for
## the double shift case this is (VU)^T A (VU).
## In either case, we have three steps:
##
## * create the unitary transformations (U or (V,U)) -- create_bulge
## * absorb the left side -- absorb_Ut
## * chase the right side down until it is absorbed -- absorb_U
##
function bulge_step(QF::QFactorization{T, S, Vt}, RF, storage, ctr, m=0) where {T, S, Vt}

    create_bulge(QF, RF, storage, ctr)
    absorb_Ut(QF, RF, storage, ctr) # see CSS or RDS
    absorb_U(QF, RF, storage, ctr)

    return nothing
end

##################################################
## Absorb U
## After U absorbtion we have either
## A * U or (W,A) * V * U
## THis chases U or (V*U) through until the value can be absorbed
##
## This just pushes work to passthrough_triu and passthrough_Q
##
## * passthrough_triu is basically passthrough(RF) with some efficiencies for first
## time through
##
## * passthrough_Q passes through chain and return false, until fusion when
## it returns true
##
function absorb_U(QF::QFactorization{T, S, Vt}, RF, storage, ctr) where {T, S, Vt}

    flag = false
    while !flag

        passthrough_triu(QF, RF, storage, ctr, Val(:right))
        flag = passthrough_Q(QF, RF, storage,  ctr, Val(:right))
        ## implicit unitary operation to move U from left to right

    end

    return nothing

end


immutable Octonion{T<:Real} <: Number
  s::T
  v1::T
  v2::T
  v3::T
  v4::T
  v5::T
  v6::T
  v7::T
  norm::Bool
end

Octonion( s::Real, v1::Real, v2::Real, v3::Real, v4::Real, v5::Real, v6::Real, v7::Real, n::Bool = false) =
  Octonion( promote( s, v1, v2, v3, v4, v5, v6, v7 )..., n )
Octonion( x::Real ) = Octonion( x, zero(x), zero(x), zero(x), zero(x), zero(x), zero(x), zero(x), abs(x) == one(x) )
Octonion( z::Complex ) = Octonion( z.re, z.im, zero(z.re), zero(z.re), zero(z.re), zero(z.re), zero(z.re), zero(z.re), abs(z) == one(z.re) )
Octonion( q::Quaternion ) = Octonion( q.s, q.v1, q.v2, q.v3, zero(q.s), zero(q.s), zero(q.s), zero(q.s), q.norm )
Octonion( s::Real, a::Vector ) = Octonion( s, a[1], a[2], a[3], a[4], a[5], a[6], a[7] )
Octonion( a::Vector ) = Octonion( 0, a[1], a[2], a[3], a[4], a[5], a[6], a[7] )

convert{T}(::Type{Octonion{T}}, x::Real) =
  Octonion(convert(T,x), convert(T,0), convert(T,0), convert(T,0), convert(T,0), convert(T,0), convert(T,0), convert(T,0))
convert{T}(::Type{Octonion{T}}, z::Complex) =
  Octonion(convert(T,real(z)), convert(T,imag(z)), convert(T,0), convert(T,0), convert(T,0), convert(T,0), convert(T,0), convert(T,0))
convert{T}(::Type{Octonion{T}}, q::Quaternion) =
  Octonion(convert(T,real(z)), convert(T,q.v1), convert(T,q.v2), convert(T,q.v3), convert(T,0), convert(T,0), convert(T,0), convert(T,0))
convert{T<:Real}(::Type{Octonion{T}}, o::Octonion{T}) = o
convert{T}(::Type{Octonion{T}}, o::Octonion) =
  Octonion(convert(T,o.s), convert(T,o.v1), convert(T,o.v2), convert(T,o.v3), convert(T,o.v4), convert(T,o.v5), convert(T,o.v6), convert(T,o.v7), o.norm)

promote_rule{T<:Real}(::Type{Octonion{T}}, ::Type{T}) = Octonion{T}
promote_rule{T<:Real}(::Type{Octonion}, ::Type{T}) = Octonion
promote_rule{T<:Real,S<:Real}(::Type{Octonion{T}}, ::Type{S}) = Octonion{promote_type(T,S)}
promote_rule{T<:Real,S<:Real}(::Type{Complex{T}}, ::Type{Octonion{S}}) = Octonion{promote_type(T,S)}
promote_rule{T<:Real,S<:Real}(::Type{Quaternion{T}}, ::Type{Octonion{S}}) = Octonion{promote_type(T,S)}
promote_rule{T<:Real,S<:Real}(::Type{Octonion{T}}, ::Type{Octonion{S}}) = Octonion{promote_type(T,S)}

octo( p, v1, v2, v3, v4, v5, v6, v7, n = false ) = Octonion( p, v1, v2, v3, v4, v5, v6, v7, n )
octo( x ) = Octonion( x )
octo( s, a ) = Octonion( s, a )

function show(io::IO, o::Octonion)
  pm(x) = x < 0 ? " - $(-x)" : " + $x"
  print(io, o.s, pm(o.v1), "im", pm(o.v2), "jm", pm(o.v3), "km", pm(o.v4), "ilm", pm(o.v5), "jlm", pm(o.v6), "klm", pm(o.v7), "lm")
end

real( o::Octonion ) = o.s
imag( o::Octonion ) = [ o.v1, o.v2, o.v3, o.v4, o.v5, o.v6, o.v7 ]

(/)( o::Octonion, x::Real ) = Octonion( o.s/x, o.v1/x, o.v2/x, o.v3/x, o.v4/x, o.v5/x, o.v6/x, o.v7/x )

conj( o::Octonion ) = Octonion( o.s, -o.v1, -o.v2, -o.v3, -o.v4, -o.v5, -o.v6, -o.v7, o.norm )
abs( o::Octonion ) = sqrt(o.s*o.s + o.v1*o.v1 + o.v2*o.v2 + o.v3*o.v3 + o.v4*o.v4 + o.v5*o.v5 + o.v6*o.v6 + o.v7*o.v7 )
abs2( o::Octonion ) = o.s*o.s + o.v1*o.v1 + o.v2*o.v2  + o.v3*o.v3 + o.v4*o.v4 + o.v5*o.v5 + o.v6*o.v6 + o.v7*o.v7
inv( o::Octonion ) = o.norm ? conj(o) : conj(o)/abs2(o)

function normalize( o::Octonion )
  if ( o.norm )
    return o
  end
  o = o / abs( o )
  Octonion( o.s, o.v1, o.v2, o.v3, o.v4, o.v5, o.v6, o.v7, true )
end

function normalizea( o::Octonion )
  if ( o.norm )
    return ( o, one( o.s ) )
  end
  a = abs( o )
  o = o / a
  ( Octonion( o.s, o.v1, o.v2, o.v3, o.v4, o.v5, o.v6, o.v7, true ), a )
end

(-)( o::Octonion) = Octonion( -o.s, -o.v1, -o.v2, -o.v3, -o.v4, -o.v5, -o.v6, -o.v7, o.norm )

(+)( o::Octonion, w::Octonion ) = Octonion( o.s + w.s,
                                            o.v1 + w.v1,
                                            o.v2 + w.v2,
                                            o.v3 + w.v3,
                                            o.v4 + w.v4,
                                            o.v5 + w.v5,
                                            o.v6 + w.v6,
                                            o.v7 + w.v7 )

(-)( o::Octonion, w::Octonion ) = Octonion( o.s - w.s,
                                            o.v1 - w.v1,
                                            o.v2 - w.v2,
                                            o.v3 - w.v3,
                                            o.v4 - w.v4,
                                            o.v5 - w.v5,
                                            o.v6 - w.v6,
                                            o.v7 - w.v7 )

macro ocmul( i, j )
  si = @compat Symbol( "v$i" )
  sj = @compat Symbol( "v$j" )
  esc( :( o.$si * w.$sj - o.$sj * w.$si ) )
end

(*)( o::Octonion, w::Octonion ) =
  Octonion(
            o.s*w.s - o.v1*w.v1 - o.v2*w.v2 - o.v3*w.v3 - o.v4*w.v4 - o.v5*w.v5 - o.v6*w.v6 - o.v7*w.v7,
            o.s*w.v1 + o.v1*w.s + @ocmul( 2, 3 ) + @ocmul( 6, 5 ) + @ocmul( 7, 4 ),
            o.s*w.v2 + o.v2*w.s + @ocmul( 3, 1 ) + @ocmul( 4, 6 ) + @ocmul( 7, 5 ),
            o.s*w.v3 + o.v3*w.s + @ocmul( 1, 2 ) + @ocmul( 5, 4 ) + @ocmul( 7, 6 ),
            o.s*w.v4 + o.v4*w.s + @ocmul( 1, 7 ) + @ocmul( 3, 5 ) + @ocmul( 6, 2 ),
            o.s*w.v5 + o.v5*w.s + @ocmul( 1, 6 ) + @ocmul( 2, 7 ) + @ocmul( 4, 3 ),
            o.s*w.v6 + o.v6*w.s + @ocmul( 2, 4 ) + @ocmul( 3, 7 ) + @ocmul( 5, 1 ),
            o.s*w.v7 + o.v7*w.s + @ocmul( 4, 1 ) + @ocmul( 5, 2 ) + @ocmul( 6, 3 )
          )

(/)( o::Octonion, w::Octonion ) = o * inv(w)

function exp( o::Octonion )
  s = o.s
  se = exp( s )
  scale = se
  th = abs( Octonion( imag( o)  ) )
  if th > 0
    scale *= sin( th ) / th
  end
  Octonion( se * cos( th ),
            scale * o.v1,
            scale * o.v2,
            scale * o.v3,
            scale * o.v4,
            scale * o.v5,
            scale * o.v6,
            scale * o.v7,
            abs( s ) == 0 )
end

function log( o::Octonion )
  o, a = normalizea( o )
  s = o.s
  M = abs( Octonion( imag( o ) ) )
  th = atan2( M, s )
  if M > 0
    M = th / M
    return Octonion( log( a ),
                     o.v1 * M,
                     o.v2 * M,
                     o.v3 * M,
                     o.v4 * M,
                     o.v5 * M,
                     o.v6 * M,
                     o.v7 * M )
  else
    return Octonion( log( a ), th , 0.0, 0.0 )
  end
end

(^)(o::Octonion, w::Octonion) = exp( w * log( o ) )

function sqrt( o::Octonion )
  exp( 0.5 * log( o ) )
end

octorand() = octo( randn(), randn(), randn(), randn(), randn(), randn(), randn(), randn() )


#=

  Just forgotten puzzle (Enigma 1517) in Julia ConstraintSolver.jl

  From http://www.f1compiler.com/samples/Enigma 201517.f1.html
  """
  Enigma 1517 Bob Walker, New Scientist magazine, October 25, 2008.

  Joe was furious when he forgot one of his bank account numbers.
  He remembered that it had all the digits 0 to 9 in some order, so he tried
  the following four sets without success:

      9 4 6 2 1 5 7 8 3 0
      8 6 0 4 3 9 1 2 5 7
      1 6 4 0 2 9 7 8 5 3
      6 8 2 4 3 1 9 0 7 5

  When Joe finally remembered his account number, he realised that in each set
  just four of the digits were in their correct position and that, if one knew
  that, it was possible to work out his account number.
  What was it?
  """

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

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

function just_forgotten(print_solutions=true,all_solutions=true)

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

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

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

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

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 100, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

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

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    a = resize_matrix([[9,4,6,2,1,5,7,8,3,0],
                       [8,6,0,4,3,9,1,2,5,7],
                       [1,6,4,0,2,9,7,8,5,3],
                       [6,8,2,4,3,1,9,0,7,5]])
    rows,cols = size(a)
    @variable(model, 0 <= x[1:cols] <= 9, Int)
    @constraint(model, x in CS.AllDifferentSet())

    for row in eachrow(a)
        # These two versions don't work in ConstraintSolver.jl v0.6.0
        # @constraint(model, sum(x .== row) == 4)
        # @constraint(model, sum([x[i] == row[i] for i in 1:cols]) == 4)

        # booleans to the rescue!
        b = @variable(model, [1:cols], Bin)
        for i in 1:cols
            @constraint(model,b[i] := {x[i] == row[i]})
        end
        @constraint(model, sum(b) == 4)
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                println("x:$x_val")

            end
        end
    else
        println("status:$status")
    end

    return status
end

println("Proving unicity")
@time just_forgotten(true,true)
println("\n'First' solution")
@time just_forgotten(true,false)


using Nettle
using Formatting
using Iterators

_regex = r"\w+"
# default in py : r"[\w\u4e00-\u9fcc]+"
# going with the more straightforward get the words

type Simhash
    # this, unfortunately, does not hash on init
    value
    f::Int

    function Simhash(value::BigInt, f::Int)
        new(value, f)
    end
    function Simhash(value::Simhash, f::Int)
        new(value.value, f)
    end
    function Simhash(value::AbstractString, f::Int)
        v = build_by_text(value, f)
        new(v, f)
    end
    function Simhash(value::Dict, f::Int)
        v = build_by_features(value, f)
        new(v, f)
    end
    function Simhash(value::AbstractArray, f::Int)
        v = build_by_features(value, f)
        new(v, f)
    end
    function Simhash(s::Simhash)
        new(s.value, s.f)
    end
end

function _hashfunc(x)
    h = Hasher("md5")
    update!(h, x)
    return parse(BigInt, hexdigest!(h), 16)
end

function distance(one::Simhash, another::Simhash)
    # this is ridiculous, julia
    x =  BigInt((one.value $ another.value) & ((BigInt(1) << BigInt(one.f)) - BigInt(1)))
    ans = 0
    while x > 0
        ans += 1
        x &= x - 1
    end
    return ans
end

function _tokenize(content::String)
    content = lowercase(content)
    content = join(matchall(_regex, content), "")
    return _slide(content)
end

function _slide(content::String, width=4)
    return [content[i:min(i + width-1, length(content))] for i in 1:max(length(content) - width, 1) + 1]
end

function build_by_text(content::String, f)
    features = _tokenize(content)
    fdict = Dict{String, Integer}(g[1] => length(g) for g = groupby(x -> x, sort(features)))
    return build_by_features(fdict, f)
end

function build_by_features(features::Any, f)
    v = zeros(f)
    masks = [BigInt(1) << BigInt(i-1) for i in 1:f]

    if isa(features, Dict)
        # to list of tuples
        features = collect(features)
    end

    for feat in features
        # strings are by default unicode
        if isa(feat, String)
            h = _hashfunc(feat)
            w = 1
        else
            # without assertion for the iterables
            h = _hashfunc(feat[1])
            w = feat[2]
        end

        for i in 1:f
            v[i] += h & masks[i] != 0 ? w : -w
        end
    end

    ans = 0
    for i in 1:f
        if v[i] >= 0
            ans |= masks[i]
        end
    end

    # numeric rep of the string/iterable
    return ans
end


"""
`MESH` - Type defining a graph with nodes in two separate instances.

* src_nodes: Nx2 array of points in global space (original instance)
* dst_nodes: Nx2 array of points in global space (updated instance)
* edges: NxM sparse incidence matrix, each column is zero except at two elements
    with value -1 and 1
    i.e.
    edges    1  2  3  4
    nodes
        A    1  1  0  0
        B   -1  0  1  1
        C    0 -1 -1  0
        D    0  0  0 -1
        E    0  0  0  0

```
            A - B   E
             \ / \
              C   D
```

"""
type Mesh
  src_nodes::Array
  dst_nodes::Array
  edges::SparseMatrixCSC{Int64, Int64}
end

Mesh() = Mesh([], [], spzeros(0,0))

"""
`CREATE_MESH` - Create mesh for image with global offset, given point spacing

  ```mesh = create_mesh(img, offset, dist)```

* img: 2D array representing image
* offset: 2-element array for location of img's [0,0] pixel in global space
* dist: positive number for distance between two neighboring nodes in mesh

Created triangular mesh forms equilateral triangles, so all interior nodes
are equidistant from immediate neighbors.

```
   *-----*
  / \   / \
 /   \ /   \
*-----*-----*
 \   / \   /
  \ /   \ /
   *-----*

All edges are dist
```
"""
function create_mesh(img, offset, dist)
  sz = collect(size(img))
  spacing = [dist*sin(pi/3), dist]
  dims = floor(Int64, sz./spacing + 1)
  margin = sz.%spacing / 2

  n = max(ij_to_index(dims, dims...), ij_to_index(dims, dims[1], dims[2]-1))
  m = 0
  nodes = zeros(Float64, (n,2))
  edges = spzeros(Float64, n, 3*n)

  for i in 1:dims[1], j in 1:dims[2]
    k = ij_to_index(dims, i, j)
    if k == 0
      continue
    end
    nodes[k,:] = ij_to_coordinates(dims, i, j, offset+margin, spacing)
    if (j != 1)
      m += 1; edges[k, m] = -1; edges[ij_to_index(dims, i, j-1), m] = 1;
    end

    if (i != 1)
      if iseven(i) || j != dims[2]
        m += 1; edges[k, m] = -1; edges[ij_to_index(dims, i-1, j), m] = 1;
      end
      if iseven(i) && (j != dims[2])
        m += 1; edges[k, m] = -1; edges[ij_to_index(dims, i-1, j+1), m] = 1;
      end
      if isodd(i) && (j != 1)
        m += 1; edges[k, m] = -1; edges[ij_to_index(dims, i-1, j-1), m] = 1;
      end
      if isodd(i) && ((j == 1) || (j == dims[2]))
        m += 1; edges[k, m] = -1; edges[ij_to_index(dims, i-2, j), m] = 1;
      end
    end
  end

  edges = edges[:, 1:m];
  return Mesh(nodes, nodes, edges)
end

"""
`IJ_TO_INDEX` - Convert mesh row & column to node index

```
ind = ij_to_index(dims, i, j)
```

* dims: tuple of dimension sizes in i & j
* i: row index
* j: column index
* ind: int representing index of the node in the mesh
"""
function ij_to_index(dims, i, j)
  ind = 0
  if iseven(i) && (j == dims[2])
    return ind
  end
  if ((i < 1) || (j < 1) || (i > dims[1]) || (j > dims[2]))
    return ind
  end
  ind += div(i-1, 2) * (dims[2] - 1)  #even rows
  ind += div(i, 2) * dims[2]          #odd rows
  ind += j
  ind = convert(Int64, ind)
  return ind
end

"""
`IJ_TO_COORDINATES` - Convert mesh row & column to coordinate in global space

```
[pi, pj] = ij_to_coordinates(dims, i, j, offset, spacing)
```

* dims: tuple of dimension sizes in i & j
* i: row index
* j: column index
* offset: 2-element array for location of img's [0,0] pixel in global space
* spacing: 2-element array for distance between mesh rows and columns
* [pi, pj]: global coordinates of the mesh node
"""
function ij_to_coordinates(dims, i, j, offset, spacing)
  if iseven(i) && (j == dims[2])
    return [0, 0]
  end
  pi = (i-1)*spacing[1] + offset[1]
  if iseven(i)
    pj = (j-0.5)*spacing[2] + offset[2]
  else
    pj = (j-1)*spacing[2] + offset[2]
  end
  return [pi, pj]
end

"""
`GET_TRIANGLES` - Use mesh edge-node incidence matrix to form triangle array

```
triangles = get_triangles(mesh)
```

See `INCIDENCE_TO_TRIANGLES` documentation for more information.
"""
function get_triangles(mesh::Mesh)
  node_dict = incidence_to_dict(mesh.edges)
  return dict_to_triangles(node_dict)
end

function count_nodes(mesh::Mesh)
  return size(src_nodes, 1)
end


import LinearAlgebra: I, Symmetric, cond, diagm, dot
using RunLengthArrays
import Healpix

################################################################################
# Pointing generation (used for testing)

function build_F(ntod, baseline_lengths)
    F = zeros(ntod, length(baseline_lengths))

    startidx = 1
    for (i, cur_len) in enumerate(baseline_lengths)
        nelems = min(ntod - startidx + 1, cur_len)
        F[startidx:(startidx + nelems - 1), i] .= 1
        startidx += nelems
    end

    F
end

function build_P_and_psi(pixidx, npix)
    ntod = length(pixidx)
    P = zeros(ntod, 3npix)
    psi = Array{Float64}(undef, ntod)

    curangle = 0.0
    for (idx, curpix) in enumerate(pixidx)
        psi[idx] = curangle

        P[idx, 1 + 3(curpix - 1)] = 1
        P[idx, 2 + 3(curpix - 1)] = cos(2curangle)
        P[idx, 3 + 3(curpix - 1)] = sin(2curangle)

        curangle += 2π / ntod
    end

    P, psi
end

generate_pixidx(ntod, npix) = Iterators.take(Iterators.cycle(1:npix), ntod) |> collect

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

# Here we build a number of vectors and matrices that solve the full map-making
# equation by explicitly building all the matrices in full form. The values here
# will be compared with the ones produced by the map-maker.
#
# Unlike the destriping code, we do *not* split baselines/TODs/etc. among MPI
# processes: each MPI process carries the full computation.

NTOD = 512
BASELINES = RunLengthArray{Int, Float64}(
    Int[NTOD/8, NTOD/4, NTOD/4, 3NTOD/8],
    Float64[2, -3, 1, 0],
)

# Map with NSIDE = 1
skymap = Healpix.PolarizedMap{Float64, Healpix.RingOrder}(
    Float64[1,   3,   -4,    6,    5,   -1,   -4,    2,   8,   0,   -1,   -2],
    Float64[1.0, 1.8,  0.7,  6.8, -2.2,  6.0,  0.3, -0.3, 4.7, 7.2,  5.1,  7.2],
    Float64[1.2, 0.4,  0.7, -1.4,  5.4,  4.9,  4.1, -1.0, 4.2, 5.9,  2.4,  3.8],
)

# The order is I1, Q1, U1, I2, Q2, U2…
m = reshape([skymap.i skymap.q skymap.u]', 3length(skymap.i))

NPIX = length(skymap.i)
PIXIDX = generate_pixidx(NTOD, NPIX)

diagCw = [1.1 for _ in 1:NTOD]
Cw = diagm(0 => diagCw)
invCw = inv(Cw)

P, PSI = build_P_and_psi(PIXIDX, NPIX)
F = build_F(NTOD, runs(BASELINES))
M = P' * invCw * P
invM = inv(M)
Z = I - P * invM * P' * invCw
D = F' * invCw * Z * F

y = P * m + F * values(BASELINES)

# These are the terms in the destriping equation "Ax = b"
A = F' * invCw * Z * F
b = F' * invCw * Z * y

baseline_guess = inv(F' * invCw * Z * F) * F' * invCw * Z * y
# Remove the average
baseline_guess .-= sum(baseline_guess) / length(baseline_guess)
# This is the maximum-likelihood map
map_guess = inv(P' * invCw * P) * P' * invCw * (y - F * baseline_guess)

################################################################################
# True map-maker

# The code that follows structures the data depending whether we are
# using MPI or not:
#
# 1. When we are *not* using MPI, there are two observations
# 2. When we are using MPI, we force the usage of two processes, and each
#    process holds only one observation.

obs_range1 = 1:sum(runs(BASELINES)[1:2])
obs_range2 = (sum(runs(BASELINES)[1:2]) + 1):NTOD

mpicomm = if Harlequin.use_mpi()
    # If use_mpi returns true, this means that this file
    # ("test_mapmaking.jl") is being run within "test_mpi.jl", which
    # has already imported the MPI library. So we can freely use
    # functions in the MPI package

    MPI.Init()
    MPI.COMM_WORLD
else
    nothing
end

# Range of validity of the array of baselines
baseline_range = if Harlequin.use_mpi()
    (MPI.Comm_rank(mpicomm) == 0) ? (1:2) : (3:4)
else
    1:4
end

rank_baselines = if Harlequin.use_mpi()
    @assert MPI.Comm_size(mpicomm) == 2

    RunLengthArray{Int, Float64}(
        runs(BASELINES)[baseline_range],
        values(BASELINES)[baseline_range],
    )
else
    BASELINES
end

rank_runs = if Harlequin.use_mpi()
    [runs(rank_baselines)]
else
    [runs(rank_baselines)[1:2], runs(rank_baselines)[3:4]]
end

rank_baselines = if Harlequin.use_mpi()
    [collect(rank_baselines)]
else
    [collect(BASELINES)[obs_range1], collect(BASELINES)[obs_range2]]
end

time_samp = 0.01 # Integration time for one sample; it is not really used

# This holds the same information as matrix M = P' ⋅ Cw^-1 ⋅ P
nobs_matrix = [NobsMatrixElement{Float64}() for i in 1:NPIX]

obs_list = [
    Observation{Float64}(
        pixidx = PIXIDX[obs_range1],
        psi_angle = PSI[obs_range1],
        tod = y[obs_range1],
        sigma_squared = diagCw[obs_range1],
        name = "Det1",
    )
    Observation{Float64}(
        pixidx = PIXIDX[obs_range2],
        psi_angle = PSI[obs_range2],
        tod = y[obs_range2],
        sigma_squared = diagCw[obs_range2],
        name = "Det2",
    )
]

# When using MPI, we only use two processes and assign the two
# observations to each of them
obs_list = Harlequin.use_mpi() ? [obs_list[MPI.Comm_rank(mpicomm) + 1]] : obs_list

let io = IOBuffer()
    show(io, obs_list[1])
    show(io, "text/plain", obs_list[1])
end

compute_nobs_matrix!(nobs_matrix, obs_list, comm = mpicomm)

# Test that NobsMatrixElement objects can be showed
let io = IOBuffer()
    show(io, nobs_matrix[1])
    show(io, "text/plain", nobs_matrix[1])
end

# Verify that nobs_matrix is really a representation of matrix M
for i in 1:NPIX
    startidx = 1 + 3(i - 1)
    subm = @view M[startidx:(startidx + 2), startidx:(startidx + 2)]
    @test nobs_matrix[i].invm ≈ inv(subm)
    @test Symmetric(nobs_matrix[i].m) * nobs_matrix[i].invm ≈ I
end

# Test compute_z_and_group!

@test length(runs(BASELINES)) == 4

dest_baselines = if use_mpi()
    [
        Float64[0.0, 0.0],
    ]
else
    [
        Float64[0.0, 0.0],
        Float64[0.0, 0.0],
    ]
end

destriping_data = DestripingData{Float64, Healpix.RingOrder}(
    1,
    obs_list,
    rank_runs,
    max_iterations = 15,
    threshold = 1e-14,
    use_preconditioner = true,
    comm = mpicomm,
)

# Check that the "show" method works in all its forms

let io = IOBuffer()
    show(io, destriping_data)
    show(io, "text/plain", destriping_data)
end

compute_z_and_group!(
    dest_baselines,
    rank_runs,
    destriping_data,
    obs_list,
    comm = mpicomm,
)

@test collect(Iterators.flatten(dest_baselines)) ≈ b[baseline_range]

left_side = if use_mpi()
    [
        Float64[0.0, 0.0],
    ]
else
    [
        Float64[0.0, 0.0],
        Float64[0.0, 0.0],
    ]
end

reset_maps!(destriping_data)

compute_z_and_group!(
    left_side,
    rank_baselines,
    rank_runs,
    destriping_data,
    obs_list,
    comm = mpicomm,
)

@test collect(Iterators.flatten(left_side)) ≈ (A * values(BASELINES))[baseline_range]

reset_maps!(destriping_data)

destripe!(
    obs_list,
    destriping_data,
    comm = mpicomm,
)

estimated_baselines = collect(Iterators.flatten([x.values for x in destriping_data.baselines]))
@test estimated_baselines ≈ values(BASELINES)[baseline_range]

@test (abs.(destriping_data.skymap.i .- skymap.i) |> maximum) < 0.3
@test (abs.(destriping_data.skymap.q .- skymap.q) |> maximum) < 0.3
@test (abs.(destriping_data.skymap.u .- skymap.u) |> maximum) < 0.3


export LogNormal, momentMatching, logMomentMatching

"""
Description:

    A log-normal node with location-scale parameterization:

    f(out,m,s) = logN(out|m, s) = 1/out (2π s)^{-1/2} exp(-1/(2s) (log(out) - m)^2))

Interfaces:

    1. out
    2. m (location)
    3. s (squared scale)

Construction:

    LogNormal(out, m, s, id=:some_id)
"""
mutable struct LogNormal <: SoftFactor
    id::Symbol
    interfaces::Vector{Interface}
    i::Dict{Symbol,Interface}

    function LogNormal(out, m, s; id=generateId(LogNormal))
        @ensureVariables(out, m, s)
        self = new(id, Array{Interface}(undef, 3), Dict{Symbol,Interface}())
        addNode!(currentGraph(), self)
        self.i[:out] = self.interfaces[1] = associate!(Interface(self), out)
        self.i[:m] = self.interfaces[2] = associate!(Interface(self), m)
        self.i[:s] = self.interfaces[3] = associate!(Interface(self), s)

        return self
    end
end

slug(::Type{LogNormal}) = "log𝒩"

format(dist::ProbabilityDistribution{Univariate, LogNormal}) = "$(slug(LogNormal))(m=$(format(dist.params[:m])), s=$(format(dist.params[:s])))"

ProbabilityDistribution(::Type{Univariate}, ::Type{LogNormal}; m::Float64=1.0, s::Float64=1.0) = ProbabilityDistribution{Univariate, LogNormal}(Dict(:m=>m, :s=>s))
ProbabilityDistribution(::Type{LogNormal}; m::Float64=1.0, s::Float64=1.0) = ProbabilityDistribution{Univariate, LogNormal}(Dict(:m=>m, :s=>s))

dims(dist::ProbabilityDistribution{Univariate, LogNormal}) = 1

vague(::Type{LogNormal}) = ProbabilityDistribution(Univariate, LogNormal, m=1.0, s=huge)

unsafeMean(dist::ProbabilityDistribution{Univariate, LogNormal}) = exp(dist.params[:m] + 0.5*dist.params[:s])
unsafeLogMean(dist::ProbabilityDistribution{Univariate, LogNormal}) = dist.params[:m]

unsafeVar(dist::ProbabilityDistribution{Univariate, LogNormal}) = (exp(dist.params[:s]) - 1.0)*exp(2.0*dist.params[:m] + dist.params[:s])
unsafeLogVar(dist::ProbabilityDistribution{Univariate, LogNormal}) = dist.params[:s]

unsafeCov(dist::ProbabilityDistribution{Univariate, LogNormal}) = unsafeVar(dist)
unsafeLogCov(dist::ProbabilityDistribution{Univariate, LogNormal}) = dist.params[:s]

logPdf(dist::ProbabilityDistribution{Univariate, LogNormal},x) = -0.5*(log(2pi)+log(dist.params[:s])) -log(x) -0.5*(log(x)-dist.params[:m])^2/dist.params[:s]
isProper(dist::ProbabilityDistribution{Univariate, LogNormal}) = (dist.params[:s] > 0.0)

function sample(dist::ProbabilityDistribution{Univariate, LogNormal})
    return exp(dist.params[:m]+sqrt(dist.params[:s])*randn())
end

"""
Gamma approximation to the log-normal distribution using Laplace's method
"""
laplace(::Type{Gamma}, dist::ProbabilityDistribution{Univariate, LogNormal}) = ProbabilityDistribution(Univariate, Gamma, a=1/dist.params[:s], b=1/dist.params[:s]*exp(-dist.params[:m]))

@symmetrical function prod!(x::ProbabilityDistribution{Univariate, LogNormal},
                            y::ProbabilityDistribution{Univariate, Gamma},
                            z::ProbabilityDistribution{Univariate, Gamma}=ProbabilityDistribution(Univariate, Gamma, a=1.0, b=1.0))

    x_approx = laplace(Gamma, x)
    z.params[:a] = x_approx.params[:a] + y.params[:a] - 1.0
    z.params[:b] = x_approx.params[:b] + y.params[:b]

    return z
end

@symmetrical function prod!(x::ProbabilityDistribution{Univariate, LogNormal},
                            y::ProbabilityDistribution{Univariate, PointMass},
                            z::ProbabilityDistribution{Univariate, PointMass}=ProbabilityDistribution(Univariate, PointMass, m=0.0))

    (y.params[:m] > 0.0) || error("PointMass location $(y.params[:m]) should be positive")
    z.params[:m] = y.params[:m]

    return z
end

# Entropy functional
function differentialEntropy(dist::ProbabilityDistribution{Univariate, LogNormal})
    0.5*log(dist.params[:s]) +
    dist.params[:m] + 0.5 +
    0.5*log(2*pi)
end

# Average energy functional
function averageEnergy(::Type{LogNormal}, marg_out::ProbabilityDistribution{Univariate}, marg_m::ProbabilityDistribution{Univariate}, marg_s::ProbabilityDistribution{Univariate})
    unsafeLogMean(marg_out) +
    0.5*log(2*pi) +
    0.5*unsafeLogMean(marg_s) +
    0.5*unsafeInverseMean(marg_s)*( unsafeCov(marg_m) + unsafeLogCov(marg_out) + (unsafeLogMean(marg_out) - unsafeMean(marg_m))^2 )
end


using BatchedBLAS, Test, CUDA, LinearAlgebra, SymmetricFormats

L=16; N=4
A = reshape(1.0:L*L*N, L,L,N);
APU = hcat([SymmetricPacked(x, :U).tri for x in eachslice(A, dims=3)]...);
APL = hcat([SymmetricPacked(x, :L).tri for x in eachslice(A, dims=3)]...);
x = reshape(1.0:L*N, L,N);
y = reshape(L*N:-1:1.0, L,N);
o = collect(1.0:N);
alpha_scalar = beta_scalar = 0.5
alpha_vector = range(0,1, length=N)
beta_vector = range(1,0, length=N)
cualpha_vector = CuArray(alpha_vector)
cubeta_vector = CuArray(beta_vector)

test_equality(::Type{<:Integer}, cpu, gpu) = @test maximum(abs.(cpu-Array(gpu))) < 1
test_equality(::Type{<:AbstractFloat}, cpu, gpu) = @test isapprox(cpu, Array(gpu))

test_types = (Float64, Int32)
@testset "A=$TAo, x=$Tx, y=$Ty, α=$Talpha, β=$Tbeta" for
         TAo in test_types, Tx in test_types, Ty in test_types,
         Talpha in (Float64, Vector{Float64}), Tbeta in (Float64, Vector{Float64})

    cuA = CuArray{TAo}(A)
    cuAPU = CuArray{TAo}(APU)
    cuAPL = CuArray{TAo}(APL)
    cux = CuArray{Tx}(x)
    cuy = CuArray{Ty}(y)
    cuo = CuArray{TAo}(o)

    @testset "dot!" begin
        o_cpu=copy(o);
        @views for i=1:N
            o_cpu[i] = x[:,i]' * y[:,i]
        end
        o_gpu=copy(cuo);
        batched_dot!(o_gpu, cux, cuy)
        test_equality(TAo, o_cpu, o_gpu)
    end

    @testset "gemv!" for trans in ['N', 'T', 'C']
        y_cpu=copy(y);
        for i=1:N
            thisalpha = Talpha<:Vector ? alpha_vector[i] : alpha_scalar
            thisbeta = Tbeta<:Vector ? beta_vector[i] : beta_scalar
            BLAS.gemv!(trans, thisalpha, A[:,:,i], x[:,i], thisbeta, @view y_cpu[:,i])
        end
        y_gpu=copy(cuy);
        thisalpha = Talpha<:Vector ? cualpha_vector : alpha_scalar
        thisbeta = Tbeta<:Vector ? cubeta_vector : beta_scalar
        batched_gemv!(trans, thisalpha, cuA, cux, thisbeta, y_gpu)
        test_equality(Ty, y_cpu, Array(y_gpu))
    end

    @testset "symv! & spmv!" for uplo in ['U', 'L']
        y_cpu=copy(y);
        for i=1:N
            thisalpha = Talpha<:Vector ? alpha_vector[i] : alpha_scalar
            thisbeta = Tbeta<:Vector ? beta_vector[i] : beta_scalar
            BLAS.symv!(uplo, thisalpha, A[:,:,i], x[:,i], thisbeta, @view y_cpu[:,i])
        end
        y_gpu=copy(cuy);
        thisalpha = Talpha<:Vector ? cualpha_vector : alpha_scalar
        thisbeta = Tbeta<:Vector ? cubeta_vector : beta_scalar
        batched_symv!(uplo, thisalpha, cuA, cux, thisbeta, y_gpu)
        test_equality(Ty, y_cpu, Array(y_gpu))

        y_cpu=copy(y);
        for i=1:N
            thisalpha = Talpha<:Vector ? alpha_vector[i] : alpha_scalar
            thisbeta = Tbeta<:Vector ? beta_vector[i] : beta_scalar
            BLAS.spmv!(uplo, thisalpha, APU[:,i], x[:,i], thisbeta, @view y_cpu[:,i])
        end
        y_gpu=copy(cuy);
        thisalpha = Talpha<:Vector ? cualpha_vector : alpha_scalar
        thisbeta = Tbeta<:Vector ? cubeta_vector : beta_scalar
        batched_spmv!(uplo, thisalpha, cuAPU, cux, thisbeta, y_gpu)
        test_equality(Ty, y_cpu, Array(y_gpu))
    end

    @testset "ger!" begin
        A_cpu=copy(A);
        for i=1:N
            thisalpha = Talpha<:Vector ? alpha_vector[i] : alpha_scalar
            BLAS.ger!(thisalpha, x[:,i], y[:,i], @view A_cpu[:,:,i])
        end
        A_gpu=copy(cuA);
        thisalpha = Talpha<:Vector ? cualpha_vector : alpha_scalar
        batched_ger!(thisalpha, cux, cuy, A_gpu)
        test_equality(TAo, A_cpu, Array(A_gpu))
    end

    @testset "syr! & spr!" for uplo in ['U', 'L']
        A_cpu=copy(A);
        for i=1:N
            thisalpha = Talpha<:Vector ? alpha_vector[i] : alpha_scalar
            BLAS.syr!(uplo, thisalpha, x[:,i], @view A_cpu[:,:,i])
        end
        A_gpu=copy(cuA);
        thisalpha = Talpha<:Vector ? cualpha_vector : alpha_scalar
        batched_syr!(uplo, thisalpha, cux, A_gpu)
        test_equality(TAo, A_cpu, Array(A_gpu))

        APU_cpu=copy(APU);
        for i=1:N
            thisalpha = Talpha<:Vector ? alpha_vector[i] : alpha_scalar
            spr!(uplo, thisalpha, x[:,i], @view APU_cpu[:,i])
        end
        APU_gpu=copy(cuAPU);
        thisalpha = Talpha<:Vector ? cualpha_vector : alpha_scalar
        batched_spr!(uplo, thisalpha, cux, APU_gpu)
        test_equality(TAo, APU_cpu, Array(APU_gpu))
    end

end


module MinimumDivergence

##############################################################################
##
## Dependencies
##
##############################################################################

using Calculus
using Divergences
using Ipopt
using ArrayViews
using PDMats
using StatsBase
using MathProgBase
using ForwardDiff

##############################################################################
##
## Extend methods
##
##############################################################################

import Base.show
import Calculus: gradient, hessian
import Divergences: Divergence
import StatsBase: coef, coeftable, confint, loglikelihood, nobs, stderr, vcov
import MathProgBase.MathProgSolverInterface

##############################################################################
##
## Exported methods
##
##############################################################################

##############################################################################
##
## Load file
##
##############################################################################

include("types.jl")
include("smoothing.jl")
include("momfun.jl")
include("interface.jl")
include("api.jl")
include("methods.jl")
include("utils.jl")

global __λ
global __p

# include("MathProgBase.jl")
# include("smoothing.jl")
# include("momfun.jl")
# include("md.jl")
# include("methods.jl")
# include("utils.jl")

export MomentFunction,
       MinimumDivergenceProblem,
       nobs, nmom, npar

end # module


module TestSpatialDiscretization

using Modia

# Desired:
#   using ModiaMath: plot
#
# In order that these packages need not to be defined in the user environment, they are included via Modia:
using Modia.ModiaMath: plot

const n = 5
@model SpatialDiscretization begin
    u = Var(start=zeros(n))
    c = 1
    @equations begin
        # der(u[2:n]) = ( u[2:n] - u[1:n-1] ) / c
        der(u) = [1; ( u[2:n] - u[1:n - 1] ) / c]
    end
end

result = simulate(SpatialDiscretization, 1)
plot(result, "u", figure=1)

splice(uodd, ueven) = [if isodd(i); uodd[div(i,2)+1] else ueven[div(i,2)] end for i in 1:length(ueven)+length(uodd)]

@model SpatialDiscretization2 begin
    u = Var(size=(5,))
    ueven = Var(start=[1,2])
    uodd = Var(size=(3,))
    @equations begin
        u = splice(uodd, ueven)
        der(ueven) = - u[[2,4]] + u[[1,3]]
        uodd = ones(3)*time
    end
end

result = simulate(SpatialDiscretization2, 1, logTranslation=true, storeEliminated=false)
plot(result, "ueven", figure=2)

n = 100
@model SpatialDiscretization3 begin
    u = Var(size=(2*n+1,))
    ueven = Var(start=[i for i in 1:n])
    uodd = Var(size=(n+1,))
    @equations begin
        u = splice(uodd, ueven)
        der(ueven) = - u[2*(1:n)] + u[2(1:n)-1]
        uodd = ones(n+1)*time
    end
end

@static if VERSION < v"0.7.0-DEV.2005"
    @time result = simulate(SpatialDiscretization3, 1, logTranslation=true, storeEliminated=false)
    plot(result, "ueven", figure=3)
end

@model SpatialDiscretization4 begin
    u = Var(size=(2*n+1,))
    ueven = Var(start=[i for i in 1:n])
    uodd = Var(size=(n+1,))
    @equations begin
        u = splice(uodd, ueven)
        der(ueven) = ueven
        uodd = zeros(n+1)
    end
end

@time result = simulate(SpatialDiscretization4, 1, logTranslation=true, storeEliminated=false)
plot(result, "ueven", figure=4)

maskEven(n) = [if isodd(i); 0 else 1 end for i in 1:n]

@model SpatialDiscretization5 begin
    u = Float(start=[if iseven(i); div(i,2) else 0 end for i in 1:2*n+1], state=false)
    @equations begin
        broadcast(*, maskEven(2*n+1), der(u)) = u
    end
end

# @time result = simulate(SpatialDiscretization5, 1, logTranslation=true, storeEliminated=false, removeSingularities=false)
# plot(result, "u", figure=5)

end


geneLen  = Dict{String, Int64}()

f = open( "hg19_gencode_basic.gtf", "r" )
n = countlines( f )  # 计算数据行数
seekstart( f )       # 返回文件起始位置

for i = 1:n
    s = split( readline(f), "\t")
    a_start = parse(Int64, s[4])
    a_end = parse(Int64, s[5])

    gene_Info = s[9]
    geneID_pre = split(gene_Info, ";")[1]
    geneID = replace(geneID_pre, "gene_id " => "")
    geneID_a = strip(geneID, '"')

    if !haskey(geneLen, geneID_a)
        geneLen[geneID_a] = a_end - a_start
    else
        geneLen[geneID_a] += a_end - a_start
    end

end

for (k,v) in geneLen

    println("$k length is $v")
end


using FastGaussQuadrature
using FixedSizeArrays
using FastAnonymous
using PyPlot
using FEFunctions
using EvolvingCurves

clogspace(start::Real, stop::Real,  n::Integer = 50) = return start*(stop/start).^linspace(0,1,n)

iround(x::Real) = convert(Int,round(x))
iround(x::AbstractVector) = [iround(x) for x in x]

# Set up simulation parameters

function parameters(nP)
    assemble_kwargs(;kwargs...) = kwargs

    P = linspace(0,2*π,nP+1)
    x = HermiteFEFunction(Vec{2,Float64}, p -> [cos(p), sin(p)], P)
    c = CubicFEFunction(Float64, p -> 0.5+20*exp(-(2*(p<π?p:p-2π))^2), P)
    f = @anon c -> c^2

    return x,c,f,assemble_kwargs(α = 1e-4, kb = 0.2)#, kb = 0.2, kc = 0.05, ks = 10
end

# Gather convergence data

using JLD
function convergence_run(;reuse = false)
    if !isdir("convergence") mkdir("convergence"); end

    function solve(nP,nT)
        T = linspace(0,2e-3,nT+1)
        x,c,f,kwargs = parameters(nP)
        return evolve(x,c,f,T; kwargs...)
    end

    nPmax = 100
    nTmax = 10000

    nPref = iround(nPmax*10.0^(1/4))
    nTref = nTmax*10
    nP = unique(iround(clogspace(2,nPmax,40)))
    nT = unique(iround(clogspace(10,nTmax,40)))

    println("Computing reference solution")
    if !reuse
        @time xref,cref = solve(nPref, nTref)
        save("convergence/refsol.jld",
            "nPref",nPref,
            "nTref",nTref,
            "xref", xref,
            "cref", cref
        )
    end
    data = load("convergence/refsol.jld")
    xref = data["xref"]
    cref = data["cref"]

    xref = sample(xref,5)
    cref = sample(cref,5)

    xP = Array{Float64}(nP)
    cP = Array{Float64}(nP)
    for i = 1:length(nP)
        println("nP-study: computing errors for nP = ", nP[i])
        @time x,c = solve(nP[i], nTref)
        @time xP[i], cP[i] = linf_errors(x,c, xref,cref)
    end
    println()

    xT = Array{Float64}(nT)
    cT = Array{Float64}(nT)
    for i = 1:length(nT)
        println("nT-study: computing errors for nT = ", nT[i])
        @time x,c = solve(nPref, nT[i])
        @time xT[i], cT[i] = linf_errors(x,c, xref,cref)
    end
    println()

    save("convergence/data.jld",
        "nP",nP,
        "xP",xP,
        "cP",cP,
        "nT",nT,
        "xT",xT,
        "cT",cT
    )
end

# Plot convergence data (run convergence_run() first)

function convergence_plot()
    data = load("convergence/data.jld")
    nP = data["nP"]
    xP = data["xP"]
    cP = data["cP"]
    nT = data["nT"]
    xT = data["xT"]
    cT = data["cT"]

    fig = figure(figsize = (4,3))
    idx = nP .> 20
    loglog(nP[idx],2*cP[end]*nP[end]^4./nP[idx].^4, "k-", lw = 1.5)
    loglog(nP[idx],2*xP[end]*nP[end]^4./nP[idx].^4, "k-", lw = 1.5)
    loglog(nP,xP, "b-o", mec = "b", label="x", ms = 2.5, mfc="w", mew = 1.5, lw = 1.5)
    loglog(nP,cP, "g-D", mec = "g", label="c", ms = 2, mfc="w", mew = 1.5, lw = 1.5)
    xlabel("# mesh points")
    ylabel("Error")
    ylim([1e-6,1e1])
    tight_layout()
    savefig("convergence/convergence_nP.png")
    close(fig)

    fig = figure(figsize = (4,3))
    loglog(nT,2*cT[end]*nT[end]./nT, "k-", lw = 1.5)
    loglog(nT,2*xT[end]*nT[end]./nT, "k-", lw = 1.5)
    loglog(nT,xT, "b-o", mec = "b", label="x", ms = 2.5, mfc="w", mew = 1.5, lw = 1.5)
    loglog(nT,cT, "g-D", mec = "g", label="c", ms = 2, mfc="w", mew = 1.5, lw = 1.5)
    xlabel("# time steps")
    ylabel("Error")
    ylim([1e-6,1e1])
    legend(loc="upper right", fontsize=12)
    tight_layout()
    savefig("convergence/convergence_nT.png")
    close(fig)
end

# Animate evolution

function animate()
    if !isdir("animate") mkdir("animate"); end
    for file in readdir("animate") rm("animate/$file"); end

    nt = 50
    T = unique([linspace(0,1e-2,101); clogspace(1e-2,8,80)])
    x,c,f,kwargs = parameters(16)

    fig = figure(figsize=(4,3))
    xlim([-1.2,2.8])
    ylim([-1.2,1.2])
    try
        plotobj = plot(x,c)
        colorbar()

        function saveplot(it,x,c)
            update(plotobj, x,c)
            draw()

            savefig("animate/it$it.png")
        end

        saveplot(1,x,c)
        for it in 2:length(T)
            TT = linspace(T[it-1],T[it], nt)
            x,c = evolve(x,c,f,TT; kwargs...)
            saveplot(it,x,c)
        end
    catch
        rethrow()
    finally
        close(fig)
    end
end

# Demonstrate mesh degeneration

function mesh_degeneration()
    if !isdir("mesh_degeneration") mkdir("mesh_degeneration"); end
    for file in readdir("mesh_degeneration") rm("mesh_degeneration/$file"); end

    for (iα,α) = enumerate((1e-4,1e4))
        nt = 50
        T = unique([linspace(0,1e-3,11); clogspace(1e-3,8,91)])
        x,c,f,kwargs = parameters(20)

        fig = figure(figsize=(4,3))
        xlim([-1.2,2.8])
        ylim([-1.2,1.2])
        axis("off")
        gca()[:get_xaxis]()[:set_visible](false)
        gca()[:get_yaxis]()[:set_visible](false)
        try
            plotobj = plot(x,c, major_size=40, minor_size=2)

            function saveplot(it,x,c)
                update(plotobj, x,c)
                draw()

                savefig("mesh_degeneration/alpha$(iα)_it$it.png")
            end

            saveplot(1,x,c)

            for it in 2:length(T)
                TT = linspace(T[it-1],T[it],nt)
                x,c = evolve(x,c,f,TT; filter(t -> t[1] != :α, kwargs)..., α = α)
                saveplot(it,x,c)
            end
        catch
            rethrow()
        finally
            close(fig)
        end
    end
end


export HJBOneDim, HJBTwoDim

abstract AbstractHJB

type HJBOneDim{T1<:Real} <: AbstractHJB
    b::Function # Drift term
    σ::Function # Volatility term
    f::Function # "Gain" function
    g::Function # Terminal value
    T::T1     # Terminal time
    amin::T1  # Control minimum value
    amax::T1  # Control maximum value
    xmin::T1  # Domain minimum value
    xmax::T1  # Domain maximum value
    # bcond = (true, false) means Dirichlet condition on x=x_{min},
    #         but no condition on x=x_{max}
    bcond::Tuple{Bool,Bool}
    Dfun::Tuple{Function,Function} # Dirichlet condition on boundaries
end

type HJBTwoDim{T1<:Real} <: AbstractHJB
    b::Function      # Drift term
    σ::Function      # Volatility term
    f::Function      # "Gain" function
    g::Function      # Terminal value
    T::T1            # Terminal time
    amin::Vector{T1} # Control minimum value
    amax::Vector{T1} # Control maximum value
    xmin::Vector{T1} # Domain minimum value
    xmax::Vector{T1} # Domain maximum value
    Dbound::Function # Dirichlet condition on x-boundary
    # TODO: assert T>0, amin < amax, xmin < xmax
    # TODO: assert that the vectors are of size 2
    # TODO: make sure b and sigma return a length-2 vector
end


# Gauss-Kronrod quadrature rule, via the QuadGK package, since Genz-Malik
# rule does not handle the 1d case.  We will just use a fixed-order (7) G-K rule.

struct GaussKronrod{T<:AbstractFloat}
    x::Vector{T}
    w::Vector{T}
    wg::Vector{T}
end

# cache the Gauss-Kronrod rules so that we don't
# call QuadGK.kronrod every time.
const gkcache = Dict{Type, GaussKronrod}()

function GaussKronrod(::Type{T}) where {T<:AbstractFloat}
    haskey(gkcache, T) && return gkcache[T]::GaussKronrod{T}
    gkcache[T] = g = GaussKronrod{T}(QuadGK.kronrod(T,7)...)
    return g
end

# further speed up the common case of double precision (25% faster for a trivial integrand)
const gk_float64 = GaussKronrod{Float64}(QuadGK.kronrod(Float64,7)...)
GaussKronrod(::Type{Float64}) = gk_float64

countevals(g::GaussKronrod) = 17

function (g::GaussKronrod{T})(f, a_::SVector{1}, b_::SVector{1}, norm=vecnorm) where {T}
    a = a_[1]
    b = b_[1]
    c = (a+b)*T(0.5)
    Δ = (b-a)*T(0.5)

    fx⁰ = f(SVector(c))                # f(0)
    I = fx⁰ * g.w[end]
    I′ = fx⁰ * g.wg[end]
    @inbounds for i = 1:length(g.x)-1
        Δx = Δ * g.x[i]
        fx = f(SVector(c + Δx)) + f(SVector(c - Δx))
        I += fx * g.w[i]
        if iseven(i)
            I′ += fx * g.wg[i>>1]
        end
    end
    V = abs(Δ)
    I *= V
    I′ *= V
    return I, norm(I - I′), 1
end


function constrain_sdp(model, s::ConstrainedContinuousIdentitySystem, Qi, Ui, Zi, nQ, nZ)
    Ei = s.X[1]; Di = s.X[2];
    @constraint(model, Qi - Ei'*Ui*Ei - Di'*Zi*Di - Array(I(nQ)) in PSDCone())
    @constraint(model, Zi - Array(I(nZ)) in PSDCone())
end

function constrain_sdp(model, s::ConstrainedLinearMap, Qi, Qj, Uij, Zij, nZ, γ)
    Ei = s.X[1]; Di = s.X[2]; Ai = s.A
    @constraint(model, γ^2*Qi - Ai'*Qj*Ai - Ei'*Uij*Ei - Di'*Zij*Di in PSDCone())
    @constraint(model, Zij - Array(I(nZ)) in PSDCone())
end

function sdplyap_feasible_prog(γ, s::StateDepDiscreteSwitchedLinearSystem, optimizer; verbose=0)
    modes=1:nstates(s)
    model = Model(optimizer)
    nA = size(s.resetmaps[1].A, 1)
    nE = size(s.resetmaps[1].X[1], 1)
    nD = size(s.resetmaps[1].X[2], 1)
    if verbose ≤ 0
        set_silent(model)
    end
    Q_vrefs = Dict(state => @variable(model, [i=1:nA, j=1:nA]) for state in modes)
    Ui_vrefs = Dict(state => @variable(model, [i=1:nE, j=1:nE], lower_bound=0) for state in modes)
    Zi_vrefs = Dict(state => @variable(model, [i=1:nD, j=1:nD]) for state in modes)

    for i in modes
        Qi = Q_vrefs[i]
        constrain_sdp(model, HybridSystems.mode(s, i), Qi, Ui_vrefs[i], Zi_vrefs[i], nA, nD)
        for t in out_transitions(s, i)
            j = target(s, t)
            Qj = Q_vrefs[j]
            Uij_ref = @variable(model, [i=1:nE, j=1:nE], lower_bound=0)
            Zij_ref = @variable(model, [i=1:nD, j=1:nD])
            constrain_sdp(model, resetmap(s, t),  Qi, Qj, Uij_ref, Zij_ref, nD, γ)
        end
    end

    if verbose ≥ 2
        print(model)
    end

    optimize!(model)

    if verbose ≥ 3
        print(model)
    end
    if verbose ≥ 2
        try
            @show JuMP.solution_summary(model, verbose=true)
        catch err
            println("solution_summary requires JuMP version 0.21.7 or later")
        end
    elseif verbose ≥ 1
        try
            @show JuMP.solution_summary(model)
        catch err
            println("solution_summary requires JuMP version 0.21.7 or later")
        end
    end

    status = (JuMP.termination_status(model), JuMP.primal_status(model), JuMP.dual_status(model), JuMP.raw_status(model))
end


module ApproxManifoldProducts

using Reexport
@reexport using KernelDensityEstimate
@reexport using TransformUtils

using DocStringExtensions

using NLsolve
using Optim
using CoordinateTransformations
using Requires

import Base: *
# import KernelDensityEstimate: kde!

const KDE = KernelDensityEstimate
const TUs = TransformUtils
const CTs = CoordinateTransformations

# TODO temporary for initial version of on-manifold products
KDE.setForceEvalDirect!(true)

export
  get2DLambda,
  get2DMu,
  get2DMuMin,
  resid2DLinear,
  ensurePeriodicDomains!,
  solveresid2DLinear!,
  solveresid2DLinear,
  ManifoldBelief,
  MB,
  *,

  # APi and util functions
  buildHybridManifoldCallbacks,
  getKDEManifoldBandwidths,
  manifoldProduct,
  manikde!,

  # Supported manifolds
  Manifold,
  EuclideanManifold,
  Euclid2,
  Circular

include("Interface.jl")
include("CommonUtils.jl")
include("Euclidean.jl")
include("CircularUtils.jl")
include("Circular.jl")
include("KernelHilbertEmbeddings.jl")

include("Legacy.jl")
include("API.jl")

function __init__()
  @require Gadfly="c91e804a-d5a3-530f-b6f0-dfbca275c004" begin
    @require Colors="5ae59095-9a9b-59fe-a467-6f913c188581" include("plotting/CircularPlotting.jl")
  end
end

end


"""
    MultivariateUniform(lower, upper)

Uniform distribution in `d` dimensions, where `d == length(lower) == length(upper)`.
# Example
```julia-repl
julia> d = MultivariateUniform([-2., -3], [4, 1])
MultivariateUniform{Float64} in 2 dimensions

julia> logpdf(d, [-3, 0])
-Inf

julia> logpdf(d, [-1, 0])
-3.1780538303479453
```
"""
struct MultivariateUniform{S} <: ContinuousMultivariateDistribution
    lower::Vector{S}
    upper::Vector{S}
    range::Vector{S}
    function MultivariateUniform(lower, upper)
        length(lower) != length(upper) && error("Lower and upper bounds must have the same length.")
        (|)((lower .> upper)...) && error("Upper bounds must not be smaller than lower bounds")
        new{eltype(lower)}(lower, upper, upper .- lower)
    end
end
Base.length(d::MultivariateUniform) = length(d.lower)
lower(d::MultivariateUniform) = d.lower
upper(d::MultivariateUniform) = d.upper
function Base.show(io::IO, ::MIME"text/plain", d::MultivariateUniform{S}) where S
    println(io, "MultivariateUniform{$S} in $(length(d.lower)) dimensions")
end
function Distributions._rand!(rng::Random.AbstractRNG,
                              d::MultivariateUniform,
                              x::AbstractVector{<:Real})
#     length(x) != length(d) && error("length(x) = $(length(x)), but length(distribution) = $(length(d))")
    for i in eachindex(d.lower)
        x[i] = d.lower[i] + rand(rng) * d.range[i]
    end
    x
end
function Distributions._logpdf(d::MultivariateUniform, x::AbstractVector{<:Real})
    for i in eachindex(x)
        (d.lower[i] > x[i] || d.upper[i] < x[i]) && return -Inf
    end
    -sum(log.(d.range))
end
Distributions.logpdf(d::Nothing, ::Any) = 0.

"""
    TruncatedMultivariateNormal(mvnormal, lower, upper)

Multivariate normal distribution `mvnormal` truncated to a box
given by `lower` and `upper` bounds.
Simple rejection sampling is implemented.
IMPORTANT: `logpdf` is not properly normalized.
# Example
```julia-repl
julia> using Distributions

julia> d = TruncatedMultivariateNormal(MultivariateNormal([2., 3.], .3*I), [0, 0], [4, 5]);

julia> logpdf(d, [2, 3])
-0.6339042620834094

julia> logpdf(d, [2, 3]) == logpdf(d.mvnormal, [2, 3])
true

julia> logpdf(d, [-1, 1])
-Inf
```
"""
struct TruncatedMultivariateNormal{B, T} <: ContinuousMultivariateDistribution
    mvnormal::B
    lower::Vector{T}
    upper::Vector{T}
end
"""
    TruncatedMultivariateNormal(m, cov; lower, upper)

Constructs a `TruncatedMultivariateNormal` with mean `m` covariance matrix `cov`
and bounds `lower` and `upper`.
# Example
```julia-repl
julia> d = TruncatedMultivariateNormal([0, 0], 3I;
                                       lower = [-1, -4], upper = [4, 5])
```
"""
function TruncatedMultivariateNormal(m, sig; lower, upper)
    TruncatedMultivariateNormal(MultivariateNormal(m, sig), lower, upper)
end
Base.length(d::TruncatedMultivariateNormal) = length(d.lower)
lower(d::TruncatedMultivariateNormal) = d.lower
upper(d::TruncatedMultivariateNormal) = d.upper
function Distributions._rand!(rng::Random.AbstractRNG,
                              d::TruncatedMultivariateNormal,
                              x::AbstractVector{<:Real})
    # To do this more efficiently
    # see e.g. https://github.com/BrianNaughton/TruncatedMVN.jl/blob/master/src/multivariate.jl
    while true
        Distributions._rand!(rng, d.mvnormal, x)
        for i in eachindex(x)
            (d.lower[i] > x[i] || d.upper[i] < x[i]) && break
            if i == length(x)
                return x
            end
        end
    end
end
function Distributions._logpdf(d::TruncatedMultivariateNormal,
                               x::AbstractVector{<:Real})
    # This is wrong; but it doesn't matter here.
    # A constant is missing due to truncation.
    for i in eachindex(x)
        (d.lower[i] > x[i] || d.upper[i] < x[i]) && return -Inf
    end
    Distributions._logpdf(d.mvnormal, x)
end

"""
    MultiParticleNormal(particles, weights;
                        lower = nothing, upper = nothing,
                        diagonal = false,
                        scaling = 2, regularization = 0.)

Mixture of `k` multivariate Gaussians with means given by an array of `particles`
(`length(particles) == k`), weighted by `weights` and with covariance matrix
given by `scaling * cov(particles) + regularization * I` for each Gaussian,
if `diagonal == false`; otherwise `scaling * std(particles) + regularization`.
Sampling truncated Gaussians with diagonal covariance matrix can be more efficient.
The mean of the distribution is the weighted average of the particles.

If `lower` and `upper` are not `nothing` the Gaussians are truncated to the box
confined by the `lower` and `upper` bounds.
Note that the logpdf of the truncated version is not properly normalized
(see also [`TruncatedMultivariateNormal`](@ref)).
# Example
```julia-repl
julia> d = MultiParticleNormal([[0, 0], [2, 3], [4, 0]], [1/3, 1/3, 1/3]);

julia> mean(d)
2-element Array{Float64,1}:
 2.0
 1.0
```
"""
mutable struct MultiParticleNormal{P, S, B} <: ContinuousMultivariateDistribution
    particles::Vector{P}
    weights::Vector{Float64}
    cov::S
    scaling::Float64
    regularization::Float64
    lower::B
    upper::B
end
"""
    MultiParticleNormal(; kwargs...)

Returns a anonymous function that accepts `particles` as input and creates a
`MultiParticleNormal` distributions with uniform weights. It accepts the same
keyword arguments `kwargs...` as the normal `MultiParticleNormal` constructor.
# Example
```julia-repl
julia> constructor = MultiParticleNormal(scaling = 3, lower = [-1, 0], upper = [5, 5]);

julia> d = constructor([[0, 0], [2, 3], [4, 0], [5, 1]])
MultiParticleNormal(4 particles in 2 dimensions, scaling = 3.0, regularization = 0.0, with bounds)

julia> d.weights
4-element Array{Float64,1}:
 0.25
 0.25
 0.25
 0.25
```
"""
function MultiParticleNormal(; kwargs...)
    particles -> MultiParticleNormal(particles,
                                     fill(1/length(particles), length(particles));
                                     kwargs...)
end
function cov_matrix(::Type{<:PDMats.AbstractPDMat}, particles, weights, scaling, regularization)
    m = scaling .* cov(particles[weights .> 0]) + regularization * I
    for _ in 1:10
        try
            return PDMats.PDMat(m)
        catch err
            if typeof(err)!=LinearAlgebra.PosDefException
                throw(err)
            end
            n = size(m, 1)
            ϵ = 1e-6 * tr(m)/n
            @inbounds for i in 1:n
                m[i, i] += ϵ
            end
        end
    end
end
function cov_matrix(::Type{<:AbstractVector}, particles, weights, scaling, regularization)
    scaling .* std(particles[weights .> 0]) .+ regularization
end
function MultiParticleNormal(particles, weights;
                             lower = nothing, upper = nothing,
                             diagonal = false,
                             scaling = 2, regularization = 0.)
    sum(weights) ≈ 1 || error("Particle weights must be normalized.")
    MultiParticleNormal(particles, weights,
                        cov_matrix(diagonal ? Vector : PDMats.PDMat,
                                   particles, weights, scaling, regularization),
                        float(scaling), float(regularization),
                        lower, upper)
end
Base.length(d::MultiParticleNormal) = size(d.cov, 1)
lower(d::MultiParticleNormal{<:Any, <:Any, <:AbstractVector}) = d.lower
upper(d::MultiParticleNormal{<:Any, <:Any, <:AbstractVector}) = d.upper
function Base.show(io::IO, ::MIME"text/plain", d::MultiParticleNormal{P,S,B}) where {P, S, B}
    println(io, "MultiParticleNormal($(length(d.particles)) particles in $(size(d.cov, 1)) dimensions, scaling = $(round(d.scaling, sigdigits = 3)), regularization = $(round(d.regularization, sigdigits = 3)), $(B <: Nothing ? "without" : "with") bounds)")
end
function Distributions._rand!(rng::Random.AbstractRNG,
                              d::MultiParticleNormal,
                              x::AbstractVector{<:Real})
    idx = wsample(d.weights)
    Distributions._rand!(rng, d, x, idx)
end
function Distributions._rand!(rng::Random.AbstractRNG,
                              d::MultiParticleNormal{<:Any, <:Any, <:Nothing},
                              x::AbstractVector{<:Real},
                              idx::Int)
    Distributions.add!(unwhiten!(d.cov, randn!(rng, x)), d.particles[idx])
end
function Distributions._rand!(rng::Random.AbstractRNG,
                              d::MultiParticleNormal,
                              x::AbstractVector{<:Real},
                              idx::Int)
    while true
        Distributions.add!(unwhiten!(d.cov, randn!(rng, x)), d.particles[idx])
        for i in eachindex(x)
            (d.lower[i] > x[i] || d.upper[i] < x[i]) && break
            if i == length(x)
                return x
            end
        end
    end
        # To do this more efficiently
        # see e.g. https://github.com/BrianNaughton/TruncatedMVN.jl/blob/master/src/multivariate.jl
        # return gibbsTMVN(d.particles[idx], d.cov, d.lower, d.upper, diagm(ones(length(d.upper))), 1)
end
function Distributions._rand!(rng::Random.AbstractRNG,
                              d::MultiParticleNormal{<:Any, <:AbstractVector},
                              x::AbstractVector{<:Real},
                              idx::Int)
    particle = d.particles[idx]
    for i in eachindex(x)
        a = d.lower[i]
        b = d.upper[i]
        r = b - a
        m = particle[i]
        sig = d.cov[i]
        if r/sig <= 2.5066282746310002
            denom = 1/(2*sig^2)
            while true
                tmp = rand() * r + a
                if exp(-(tmp-m)^2*denom) > rand()
                    x[i] = tmp
                    break
                end
            end
        else
            while true
                tmp = m + randn() * sig
                if a <= tmp && b >= tmp
                    x[i] = tmp
                    break
                end
            end
        end
    end
    x
end
function _lpdf(d, x)
    log(sum(d.weights[i] * exp(logpdf(MultivariateNormal(d.particles[i], d.cov), x))
            for i in eachindex(d.weights)))
end
function Distributions._logpdf(d::MultiParticleNormal{<:Any, <:Any, <:Nothing},
                               x::AbstractVector{<:Real})
    _lpdf(d, x)
end
function Distributions._logpdf(d::MultiParticleNormal, x::AbstractVector{<:Real})
    # This is wrong if upper and lower are finite; but it doesn't matter here.
    # A constant is missing due to truncation.
    for i in eachindex(x)
        (d.lower[i] > x[i] || d.upper[i] < x[i]) && return -Inf
    end
    _lpdf(d, x)
end
function Distributions.mean(d::MultiParticleNormal)
    sum(d.weights[i] .* d.particles[i] for i in eachindex(d.weights))
end
function update!(d::MultiParticleNormal{<:Any, S}, particles, weights; log_weights = true) where S
    log_weights ? weights!(d.weights, weights) : d.weights .= weights
    for i in eachindex(particles)
        d.particles[i] .= particles[i]
    end
    d.cov = cov_matrix(S, d.particles, d.weights, d.scaling, d.regularization)
    d
end

lower(d) = fill(-Inf, length(d))
upper(d) = fill(Inf, length(d))


using Test

using MarkovChains
@testset "inftime_solve" begin
@testset "acyclic1" begin
    chain = ContMarkovChain()
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    n3 = add_state!(chain)
    add_transition!(chain, n1, n2, 2.0)
    add_transition!(chain, n2, n3, 4.0)
    init_prob = fill(0.0, state_count(chain))
    init_prob[1] = 1.0
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), n1:n3))
    @test cumtime[n1] ≈ 0.5
    @test cumtime[n2] ≈ 0.25
    @test cumtime[n3] ≈ Inf
    prob = collect(map(state -> state_prob(res, state), n1:n3))
    @test prob[n1] ≈ 0
    @test prob[n2] ≈ 0
    @test prob[n3] ≈ 1

    @test mean_time_to_absorption(chain, init_prob) == 0.75
end
@testset "irreducible1" begin
    chain = ContMarkovChain()
    n0 = add_state!(chain)
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    n3 = add_state!(chain)
    add_transition!(chain, n0, n1, 1.0)
    add_transition!(chain, n1, n2, 1.0)
    add_transition!(chain, n2, n3, 1.0)
    add_transition!(chain, n3, n2, 3.0)
    add_transition!(chain, n2, n1, 2.0)
    add_transition!(chain, n1, n0, 1.0)
    Q = trans_rate_matrix(chain)
    init_prob = fill(0.0, state_count(chain))
    init_prob[1] = 0.1
    init_prob[2] = 0.9
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), n0:n3))
    for t in cumtime
        @test t == Inf
    end
    prob = collect(map(state -> state_prob(res, state), n0:n3))
    @test prob ≈ [0.375, 0.375, 0.1875, 0.0625]
end
@testset "multiple_comps" begin
    chain = ContMarkovChain()
    t0 = add_state!(chain)
    t1 = add_state!(chain)
    t2 = add_state!(chain)
    m1 = add_state!(chain)
    m2 = add_state!(chain)
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    add_transition!(chain, t0, t1, 1.0)
    add_transition!(chain, t1, t2, 1.0)
    add_transition!(chain, t2, t1, 2.0)
    add_transition!(chain, t1, m1, 1.0)
    add_transition!(chain, t2, n1, 3.0)
    add_transition!(chain, m1, m2, 2.0)
    add_transition!(chain, m2, m1, 1.0)
    add_transition!(chain, n1, n2, 2.0)
    add_transition!(chain, n2, n1, 1.0)
    init_prob = fill(0.0, state_count(chain))
    init_prob[1] = 1.0
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), t0:n2))
    for t in cumtime[[m1, m2, n1, n2]]
        @test t == Inf
    end
    for t in cumtime[[t0, t1, t2]]
        @test t > 0
    end
    prob = collect(map(state -> state_prob(res, state), t0:n2))
    for p in prob[[t0, t1, t2]]
        @test p == 0.0
    end
    for p in prob[[m1, m2, n1, n2]]
        @test p > 0
    end

    @test isinf(mean_time_to_absorption(chain, init_prob))
end
using SparseArrays
@testset "multiple_comps_init_bottom" begin
    chain = ContMarkovChain()
    t0 = add_state!(chain)
    t1 = add_state!(chain)
    t2 = add_state!(chain)
    m1 = add_state!(chain)
    m2 = add_state!(chain)
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    add_transition!(chain, t0, t1, 1.0)
    add_transition!(chain, t1, t2, 1.0)
    add_transition!(chain, t2, t1, 2.0)
    add_transition!(chain, t1, m1, 1.0)
    add_transition!(chain, t2, n1, 3.0)
    add_transition!(chain, m1, m2, 2.0)
    add_transition!(chain, m2, m1, 1.0)
    add_transition!(chain, n1, n2, 2.0)
    add_transition!(chain, n2, n1, 1.0)
    init_prob = sparsevec([n2, n1, m1], [0.7, 0.2, 0.1])
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), t0:n2))
    for t in cumtime[[m1, m2, n1, n2]]
        @test t == Inf
    end
    for t in cumtime[[t0, t1, t2]]
        @test t == 0.0
    end
    prob = collect(map(state -> state_prob(res, state), t0:n2))
    for p in prob[[t0, t1, t2]]
        @test p == 0.0
    end
    @test prob[n1] == 1.0 / 3.0 * 0.9
    @test prob[n2] == 2.0 / 3.0 * 0.9
    @test prob[m1] == 1.0 / 3.0 * 0.1
    @test prob[m2] == 2.0 / 3.0 * 0.1
end
@testset "one_state" begin
    chain = ContMarkovChain()
    t0 = add_state!(chain)
    init_prob = sparsevec([t0], [1.0])
    res = inftime_solve(chain, init_prob)
    @test state_cumtime(res, t0) == Inf
    @test state_prob(res, t0) == 1.0
    @test mean_time_to_absorption(chain, init_prob) == 0.0
end
@testset "zero_state" begin
    chain = ContMarkovChain()
    init_prob = sparsevec([], [])
    res = inftime_solve(chain, init_prob)
end

end


using Revise
using AdFem
using PyPlot
using LinearAlgebra
using Statistics
using MAT

mmesh = Mesh(joinpath(PDATA, "twoholes_large.stl"))
mmesh = Mesh(mmesh.nodes * 10, mmesh.elems, -1, 2)

ν = Variable(0.01)
K = ν*constant(compute_fem_laplace_matrix(mmesh))
B = constant(compute_interaction_matrix(mmesh))
Z = [K -B'
    -B spzero(size(B,1))]

bd1 = bcnode((x,y)->y>0.23-1e-5, mmesh)
bd2 = bcnode((x,y)->y<=0.23-1e-5, mmesh)
bd = [bd1;bd2]
bd = [bd; bd .+ mmesh.ndof; 2mmesh.ndof + 1]

rhs = zeros(2mmesh.ndof + mmesh.nelem)
bdval = zeros(length(bd))
bdval[1:length(bd1)] .= 1.0
Z, rhs = impose_Dirichlet_boundary_conditions(Z, rhs, bd, bdval)
sol = Z\rhs
U = matread("fenics/fwd1.mat")["u"]

loss = sum((sol[2mmesh.ndof+1:end] - U[2mmesh.ndof+1:end])^2)
sess = Session(); init(sess)

BFGS!(sess, loss)


# ==================================
# Beam
#
# system type: continuous LTI system
# state dimension: 348
# input dimension: 1
# ==================================
using ReachabilityBenchmarks, MathematicalSystems, LazySets, MAT

function beam_model()
    file = matopen(@relpath "beam.mat")

    # system matrix
    A = read(file, "A")

    # input matrix
    B = read(file, "B")

    # state domain
    X = Universe(348)

    # input domain
    U = BallInf([0.5], 0.3)

    # continuous LTI system
    S = @system(x' = Ax + Bu, x ∈ X, u ∈ U)

    return S
end


    # This routine gets the coefficients for the non-linea2
# limb-darkening law fit to a function which goes to
# zero at mu=0 and mu=1:

include("/Users/ericagol/Computer/Julia/regress.jl")

function get_coeff!(mu::Array{T,1},fnon::Array{T,1},i::Int64,nmu::Int64) where {T <: Real}
# i is the order of the fit; nmu is number of mu points which are fit.
fn = zeros(T,i-1,nmu)
for j=2:i
  fn[j-1,:] .= (j+2)*mu.^j-j*mu.^(j-2)-2mu
end
ci,cov = regress(fn,fnon,ones(BigFloat,nmu))
fmod = zeros(BigFloat,nmu)
for j=1:i-1
  fmod += ci[j]*fn[j,:]
end
return ci,fmod
end


"""
```julia
    SoftMaxLikelihood()
```
Multiclass likelihood with [Softmax transformation](https://en.wikipedia.org/wiki/Softmax_function):

```math
p(y=i|{fₖ}) = exp(fᵢ)/ ∑ₖexp(fₖ)
```

There is no possible augmentation for this likelihood
"""
struct SoftMaxLikelihood{T<:Real} <: MultiClassLikelihood{T}
    Y::AbstractVector{BitVector} #Mapping from instances to classes
    class_mapping::AbstractVector{Any} # Classes labels mapping
    ind_mapping::Dict{Any,Int} # Mapping from label to index
    y_class::AbstractVector{Int64} #GP Index for each sample
    θ::AbstractVector{AbstractVector{T}} # Variational parameter of Polya-Gamma distribution
    function SoftMaxLikelihood{T}() where {T<:Real}
        new{T}()
    end
    function SoftMaxLikelihood{T}(Y::AbstractVector{<:BitVector},
    class_mapping::AbstractVector, ind_mapping::Dict{<:Any,<:Int},y_class::AbstractVector{<:Int}) where {T<:Real}
        new{T}(Y,class_mapping,ind_mapping,y_class)
    end
end

function SoftMaxLikelihood()
    SoftMaxLikelihood{Float64}()
end

function pdf(l::SoftMaxLikelihood,f::AbstractVector)
    StatsFuns.softmax(f)
end

function pdf(l::SoftMaxLikelihood,y::Int,f::AbstractVector{<:Real})
    StatsFuns.softmax(f)[y]
end

function Base.show(io::IO,model::SoftMaxLikelihood{T}) where T
    print(io,"Softmax likelihood")
end

function init_likelihood(likelihood::SoftMaxLikelihood{T},nLatent::Integer,nSamplesUsed::Integer) where T
    if inference isa GibbsSampling
        θ = [abs.(rand(T,nSamplesUsed))*2 for i in 1:nLatent]
        LogisticSoftMaxLikelihood{T}(likelihood.Y,likelihood.class_mapping,likelihood.ind_mapping,likelihood.y_class,θ)
    else
        return likelihood
    end
end

function sample_local!(model::VGP{T,<:SoftMaxLikelihood,<:GibbsSampling}) where {T}
    pg = PolyaGammaDist()
    model.likelihood.θ .= broadcast((y::BitVector,γ::AbstractVector{<:Real},μ::AbstractVector{<:Real},i::Int64)->draw.([pg],1.0,μ-logsumexp()),model.likelihood.Y,model.likelihood.γ,model.μ,1:model.nLatent)
    return nothing #TODO FINISH AT SOME POINT
end

function grad_samples(model::AbstractGP{T,<:SoftMaxLikelihood},samples::AbstractMatrix{T},index::Integer) where {T}
    class = model.likelihood.y_class[index]
    grad_μ = zeros(T,model.nLatent)
    grad_Σ = zeros(T,model.nLatent)
    nSamples = size(samples,1)
    samples .= mapslices(StatsFuns.softmax,samples,dims=2)
    t = 0.0
    @inbounds for i in 1:nSamples
        s = samples[i,class]::T
        @views g_μ = grad_softmax(samples[i,:],class)/s
        grad_μ += g_μ
        @views h = diaghessian_softmax(samples[i,:],class)/s
        grad_Σ += h - abs2.(g_μ)
    end
    for k in 1:model.nLatent
        model.inference.∇μE[k][index] = grad_μ[k]/nSamples
        model.inference.∇ΣE[k][index] = 0.5.*grad_Σ[k]/nSamples
    end
end

function log_like_samples(model::AbstractGP{T,<:SoftMaxLikelihood},samples::AbstractMatrix,index::Integer) where {T}
    class = model.likelihood.y_class[index]
    nSamples = size(samples,1)
    loglike = mapslices(logsumexp,samples,dims=2)/nSamples
end

function grad_softmax(s::AbstractVector{<:Real},i::Integer)
    return (δ.(i,eachindex(s))-s)*s[i]
end

function diaghessian_softmax(s::AbstractVector{<:Real},i::Integer)
    return s[i]*(abs2.(δ.(i,eachindex(s))-s) - s.*(1.0.-s))
end

function hessian_softmax(s::AbstractVector{<:Real},i::Integer)
    m = length(s)
    hessian = zeros(m,m)
    for j in 1:m
        for k in 1:m
            hessian[j,k] = s[i]*((δ(i,k)-s[k])*(δ(i,j)-s[j])-s[j]*(δ(j,k)-s[k]))
        end
    end
    return hessian
end


function jacobian(m::KrusellSmith)

    # Load in endogenous state and eq cond indices
    endo = augment_model_states(m.endogenous_states_unnormalized,
                                n_model_states_unnormalized(m))
    eq   = m.equilibrium_conditions

    # Load in grid settings
    nw::Int = get_setting(m, :nw)
    ns::Int = get_setting(m, :ns)

    zhi::Float64 = get_setting(m, :zhi)
    zlo::Float64 = get_setting(m, :zlo)
    smoother::Float64 = get_setting(m, :smoother)

    wgrid = get_grid(m, :wgrid)
    sgrid = get_grid(m, :sgrid)

    # Make the Jacobian
    JJ = zeros(4*nw+2,8*nw+4)

    # Create auxiliary variables
    g::Array{Float64,1} = trunc_lognpdf(sgrid, m[:μ_s].value, m[:σ_s].value)                   # density g (of skill, s) evaluated on sgrid

    m[:Wstar] = MPL(1.0, m[:Lstar].value, m[:Kstar].value, m[:α].value)            # Steady state wages
    m[:Rstar] = MPK(1.0, m[:Lstar].value, m[:Kstar].value, m[:α].value, m[:δ].value)     # Steady state net return on capital

    # KF equation
    KFmollificand::Array{Float64, 2} = (repeat(wgrid.points, 1, nw*ns) - m[:Rstar]*repeat(wgrid.points'-m[:cstar].value', nw, ns))/m[:Wstar] - kron(sgrid.points',ones(nw, nw))
    μRHS::Array{Float64,1} = mollifier.(KFmollificand, zhi, zlo, smoother)*kron(sgrid.weights.*(g/m[:Wstar]), wgrid.weights.*m[:μstar].value)

    # These matrices (from the Euler equation) correspond to the matrices in the PDF documentation on Dropbox
    # (note that Γ is Λ removing the Σ_i^nx, and the ω^x_i terms from the summation.
    ξ = Array{Float64, 2}(undef, nw, nw) #zeros(nw, nw)
    Γ = Array{Float64, 2}(undef, nw, nw) #zeros(nw, nw)

    # These matrices (from the KF equation) corresp. to the lyx document on dropbox, script A and ζ
    AA = Array{Float64,2}(undef, nw, nw)
    ζ = Array{Float64, 2}(undef, nw, nw)

    # i_w is the ith entry in the nw discretization of x
    # i_wp is i_w' the grid of potential future x's
    for i_w in 1:nw
        for i_wp in 1:nw
            # The following hold the sums over i=1̣̣̣̣, ..., ns
            sum_ns_ξ::Float64 = 0.0
            sum_ns_Γ::Float64 = 0.0
            sum_ns_AA::Float64 = 0.0
            sum_ns_ζ::Float64 = 0.0
            # The first part of ξ and Γ which only relies on i_w' so we calculate outside of 1:ns loop
            front = ((m[:cstar].value[i_wp].^(-m[:γ]))/m[:Wstar])
            first_part_mollificand = (wgrid.points[i_wp] - m[:Rstar] * (wgrid.points[i_w] -
                                                                        m[:cstar].value[i_w]))/m[:Wstar]
            for iss in 1:ns
                mollificand::Float64 = first_part_mollificand - sgrid.points[iss]
                # dm and mm are the mollifiers which correspond to q(mollificand) in the paper
                dm = dmollifier(mollificand, zhi, zlo, smoother)
                mm = mollifier(mollificand, zhi, zlo, smoother)
                skill_distrib_weights = g[iss] * sgrid.weights[iss]

                sum_ns_ξ += front * dm * skill_distrib_weights
                sum_ns_Γ += front * mm * skill_distrib_weights
                sum_ns_AA+= (1.0/m[:Wstar]) * mm * skill_distrib_weights
                sum_ns_ζ += (1.0/m[:Wstar]) * dm * skill_distrib_weights
            end
            ξ[i_w,i_wp] = sum_ns_ξ
            Γ[i_w,i_wp] = sum_ns_Γ
            AA[i_w,i_wp] = sum_ns_AA
            ζ[i_w,i_wp] = sum_ns_ζ
        end
    end

    dRdK::Float64 = -(1-m[:α])*((m[:Rstar]+m[:δ]-1)/m[:Kstar])
    dWdK::Float64 = (m[:α]*m[:Wstar]/m[:Kstar])
    dRdZ::Float64 = m[:Rstar]+m[:δ]-1.0
    dWdZ::Float64 = m[:Wstar].value
    lRHS = m[:β]*m[:Rstar]*Γ*wgrid.weights

    # Fill in Jacobian

    # Euler equation (EE)
    term1_EE::Array{Float64, 1} = m[:β]*Γ*wgrid.weights - m[:β]*(m[:Rstar]/m[:Wstar]) * ((wgrid.points-m[:cstar].value) .* (ξ*wgrid.weights))
    term2_EE::Array{Float64, 1} = -(1.0/m[:Wstar]) * (lRHS+((m[:β]*m[:Rstar])/m[:Wstar]) * ξ.*(repeat(wgrid.points', nw, 1) - m[:Rstar]*repeat(wgrid.points-m[:cstar].value,1,nw))*wgrid.weights)

    JJ[eq[:eq_euler], endo[:K′_t]]  = term1_EE*dRdK + term2_EE*dWdK
    JJ[eq[:eq_euler], endo[:z′_t]] = term1_EE*dRdZ + term2_EE*dWdZ

    JJ[eq[:eq_euler], endo[:l′_t]] = m[:β]*m[:Rstar]*Γ*Matrix(Diagonal(((wgrid.weights.*(m[:lstar].value.^(-(1.0+m[:γ])/m[:γ]))
                                                                         .*(m[:lstar].value.^(-1.0/m[:γ]) .<= wgrid.points))./m[:cstar].value)))

    JJ[eq[:eq_euler], endo[:l_t]]  = -(eye(nw) + Matrix(Diagonal((m[:β]/m[:γ])*(m[:Rstar]*m[:Rstar]/m[:Wstar])*(ξ*wgrid.weights)
                                                                 .*(m[:lstar].value.^(-(1.0+m[:γ])/m[:γ])).*(m[:lstar].value.^(-1.0/m[:γ]) .<= wgrid.points))))

    # KF Equation
    JJ[eq[:eq_kolmogorov_fwd], endo[:μ_t1]]   = AA'*Matrix(Diagonal(wgrid.weights))

    JJ[eq[:eq_kolmogorov_fwd], endo[:l_t1]] = (-(m[:Rstar]/m[:Wstar])*(1.0/m[:γ])*ζ') * Matrix(Diagonal( m[:μstar].value.*wgrid.weights.*(m[:lstar].value.^(-(1.0+m[:γ])/m[:γ])) .*(m[:lstar].value.^(-1.0/m[:γ]) .<= wgrid.points)))

    term1_KF::Array{Float64,1} = -(1.0/m[:Wstar])*(ζ'.*repeat((wgrid.points-m[:cstar].value)',nw,1))*(m[:μstar].value.*wgrid.weights)
    term2_KF::Array{Float64,1} = -(μRHS/m[:Wstar] + (1.0/(m[:Wstar]*m[:Wstar]))*(ζ'.*(repeat(wgrid.points,1,nw) - m[:Rstar]*repeat( (wgrid.points-m[:cstar].value)',nw,1)))*(m[:μstar].value.*wgrid.weights))

    JJ[eq[:eq_kolmogorov_fwd], endo[:K_t]]  = term1_KF*dRdK + term2_KF*dWdK
    JJ[eq[:eq_kolmogorov_fwd], endo[:z_t]]   = term1_KF*dRdZ + term2_KF*dWdZ

    JJ[eq[:eq_kolmogorov_fwd], endo[:μ_t]]    = -eye(nw)

    # DEFN of LM(t+1) = M(t)
    JJ[eq[:eq_μ], endo[:μ′_t1]]  = eye(nw)

    JJ[eq[:eq_μ], endo[:μ_t]]    = -eye(nw)

    # DEFN of LELL(t+1) = ELL(t)
    JJ[eq[:eq_l], endo[:l′_t1]]= eye(nw)

    JJ[eq[:eq_l], endo[:l_t]]  = -eye(nw)

    # LOM K
    JJ[eq[:eq_K_law_of_motion], endo[:l_t]]  = -(1.0/m[:γ])*m[:μstar].value'*(Matrix(Diagonal(wgrid.weights.*(m[:lstar].value.^(-(1.0+m[:γ])/m[:γ]))
                                                                                .*(m[:lstar].value.^(-1.0/m[:γ]) .<= wgrid.points))))

    JJ[eq[:eq_K_law_of_motion], endo[:μ_t]]  .= -(wgrid.points-m[:cstar].value)'*Matrix(Diagonal(wgrid.weights))

    JJ[eq[:eq_K_law_of_motion], endo[:K′_t]]  .= 1.0

    # TFP
    JJ[eq[:eq_TFP], endo[:z′_t]]   .= 1.0

    JJ[eq[:eq_TFP], endo[:z_t]]    .= -m[:ρ_z]

    if !m.testing && get_setting(m, :normalize_distr_variables)
        JJ = normalize(m, JJ)
    end

    return JJ
end

function normalize(m::KrusellSmith, JJ::Matrix{Float64})
    nw = get_setting(m, :nw)

    P = eye(nw)
    P[:,1] = ones(nw)

    Q = Matrix{Float64}(undef, nw, nw)
    Q::Matrix{Float64}, _ = qr(P)

    S = Matrix{Float64}(undef, nw, nw-1)
    S::Matrix{Float64} = Q[:, 2:end]'

    Qf = zeros(4*nw, 4*nw+2)
    Qf[1:nw, 1:nw] = eye(nw)
    Qf[nw+1:2*nw-1, nw+1:2*nw] = S
    Qf[2*nw:3*nw-2, 2*nw+1:3*nw] = S
    Qf[3*nw-1:4*nw, 3*nw+1:4*nw+2] = eye(nw+2)

    # Qx and Qy are for normalizing any variables that represent distributions
    # the S component is for normalizing a distribution, the other identity portions
    # are for non-distributional variables
    # The distribution in the states, x, is the lagged μ
    # The distribution in the jumps, y, is the current μ
    Qx = cat(S, eye(nw), [1], [1], dims = [1 2])   # 161 x 162
    Qy = cat(S, eye(nw), dims = [1 2])             # 159 x 160

    m <= Setting(:n_predetermined_variables, size(Qx, 1))

    # whenever outputs and/or inputs are densities,
    # pre/postmultiply the jacobians from step 2 by an appropriate
    # matrix so things integrate to 1

    Qright = cat(Qx',Qy',Qx',Qy', dims = [1 2])
    Jac1 = Qf*JJ*Qright

    return Jac1
end