using LinearAlgebra

export fromroots,
       truncate!,
       chop!,
       coeffs,
       degree,
       domain,
       mapdomain,
       order,
       hasnan,
       roots,
       companion,
       vander,
       fit,
       integrate,
       derivative,
       variable,
       isintegral,
       ismonic

"""
    fromroots(::AbstractVector{<:Number}; var=:x)
    fromroots(::Type{<:AbstractPolynomial}, ::AbstractVector{<:Number}; var=:x)

Construct a polynomial of the given type given the roots. If no type is given, defaults to `Polynomial`.

# Examples
```jldoctest common
julia> using Polynomials

julia> r = [3, 2]; # (x - 3)(x - 2)

julia> fromroots(r)
Polynomial(6 - 5*x + x^2)
```
"""
function fromroots(P::Type{<:AbstractPolynomial}, roots::AbstractVector; var::SymbolLike = :x)
    x = variable(P, var)
    p =  prod(x - r for r in roots)
    return truncate!(p)
end
fromroots(r::AbstractVector{<:Number}; var::SymbolLike = :x) =
    fromroots(Polynomial, r, var = var)

"""
    fromroots(::AbstractMatrix{<:Number}; var=:x)
    fromroots(::Type{<:AbstractPolynomial}, ::AbstractMatrix{<:Number}; var=:x)

Construct a polynomial of the given type using the eigenvalues of the given matrix as the roots. If no type is given, defaults to `Polynomial`.

# Examples
```jldoctest common
julia> using Polynomials

julia> A = [1 2; 3 4]; # (x - 5.37228)(x + 0.37228)

julia> fromroots(A)
Polynomial(-1.9999999999999998 - 5.0*x + 1.0*x^2)
```
"""
fromroots(P::Type{<:AbstractPolynomial},
    A::AbstractMatrix{T};
    var::SymbolLike = :x,) where {T <: Number} = fromroots(P, eigvals(A), var = var)
fromroots(A::AbstractMatrix{T}; var::SymbolLike = :x) where {T <: Number} =
    fromroots(Polynomial, eigvals(A), var = var)

"""
    fit(x, y, deg=length(x) - 1; [weights], var=:x)
    fit(::Type{<:AbstractPolynomial}, x, y, deg=length(x)-1; [weights], var=:x)

Fit the given data as a polynomial type with the given degree. Uses
linear least squares to minimize the norm  `||y - V⋅β||^2`, where `V` is
the Vandermonde matrix and `β` are the coefficients of the polynomial
fit.

This will automatically scale your data to the [`domain`](@ref) of the
polynomial type using [`mapdomain`](@ref). The default polynomial type
is [`Polynomial`](@ref).

## Weights

Weights may be assigned to the points by specifying a vector or matrix of weights.

When specified as a vector, `[w₁,…,wₙ]`, the weights should be
non-negative as the minimization problem is `argmin_β Σᵢ wᵢ |yᵢ - Σⱼ
Vᵢⱼ βⱼ|² = argmin_β || √(W)⋅(y - V(x)β)||²`, where, `W` the digonal
matrix formed from `[w₁,…,wₙ]`, is used for the solution, `V` being
the Vandermonde matrix of `x` corresponding to the specified
degree. This parameterization of the weights is different from that of
`numpy.polyfit`, where the weights would be specified through
`[ω₁,ω₂,…,ωₙ] = [√w₁, √w₂,…,√wₙ]`
with the answer solving
`argminᵦ | (ωᵢ⋅yᵢ- ΣⱼVᵢⱼ(ω⋅x) βⱼ) |^2`.

When specified as a matrix, `W`, the solution is through the normal
equations `(VᵀWV)β = (Vᵀy)`, again `V` being the Vandermonde matrix of
`x` corresponding to the specified degree.

(In statistics, the vector case corresponds to weighted least squares,
where weights are typically given by `wᵢ = 1/σᵢ²`, the `σᵢ²` being the
variance of the measurement; the matrix specification follows that of
the generalized least squares estimator with `W = Σ⁻¹`, the inverse of
the variance-covariance matrix.)

## large degree

For fitting with a large degree, the Vandermonde matrix is exponentially ill-conditioned. The [`ArnoldiFit`](@ref) type introduces an Arnoldi orthogonalization that fixes this problem.

"""
function fit(P::Type{<:AbstractPolynomial},
             x::AbstractVector{T},
             y::AbstractVector{T},
             deg::Integer = length(x) - 1;
             weights = nothing,
             var = :x,) where {T}
    _fit(P, x, y, deg; weights=weights, var=var)
end

fit(P::Type{<:AbstractPolynomial},
    x,
    y,
    deg::Integer = length(x) - 1;
    weights = nothing,
    var = :x,) = fit′(P, promote(collect(x), collect(y))..., deg; weights = weights, var = var)

#  avoid issue  214
fit′(P::Type{<:AbstractPolynomial}, x, y, args...;kwargs...) = throw(ArgumentError("x and y do not produce abstract   vectors"))
fit′(P::Type{<:AbstractPolynomial},
     x::AbstractVector{T},
     y::AbstractVector{T},
     args...; kwargs...) where {T} = fit(P,x,y,args...;  kwargs...)

fit(x::AbstractVector,
    y::AbstractVector,
    deg::Integer = length(x) - 1;
    weights = nothing,
    var = :x,) = fit(Polynomial, x, y, deg; weights = weights, var = var)

function _fit(P::Type{<:AbstractPolynomial},
             x::AbstractVector{T},
             y::AbstractVector{T},
             deg::Integer = length(x) - 1;
             weights = nothing,
             var = :x,) where {T}
    x = mapdomain(P, x)
    vand = vander(P, x, deg)
    if weights !== nothing
        coeffs = _wlstsq(vand, y, weights)
    else
        coeffs = qr(vand) \ y
    end
    R = float(T)
    return P(R.(coeffs), var)
end

# Weighted linear least squares
_wlstsq(vand, y, W::Number) = _wlstsq(vand, y, fill!(similar(y), W))
function _wlstsq(vand, y, w::AbstractVector)
    W = Diagonal(sqrt.(w))
    qr(W * vand) \ (W * y)
end
_wlstsq(vand, y, W::AbstractMatrix) = qr(vand' * W * vand) \ (vand' * W * y)

"""
    roots(::AbstractPolynomial; kwargs...)

Returns the roots, or zeros, of the given polynomial.

For non-factored, standard basis polynomials the roots are calculated via the eigenvalues of the companion matrix. The `kwargs` are passed to the `LinearAlgeebra.eigvals` call.

!!! note
    The default `roots` implementation is for polynomials in the
    standard basis. The companion matrix approach is reasonably fast
    and accurate for modest-size polynomials. However, other packages
    in the `Julia` ecosystem may be of interest and are mentioned in the documentation.

"""
function roots(q::AbstractPolynomial{T}; kwargs...) where {T}

    p = convert(Polynomial{T},  q)
    roots(p; kwargs...)

end

"""
    companion(::AbstractPolynomial)

Return the companion matrix for the given polynomial.

# References
[Companion Matrix](https://en.wikipedia.org/wiki/Companion_matrix)
"""
companion(::AbstractPolynomial)

"""
    vander(::Type{AbstractPolynomial}, x::AbstractVector, deg::Integer)

Calculate the pseudo-Vandermonde matrix of the given polynomial type with the given degree.

# References
[Vandermonde Matrix](https://en.wikipedia.org/wiki/Vandermonde_matrix)
"""
vander(::Type{<:AbstractPolynomial}, x::AbstractVector, deg::Integer)

"""
    integrate(p::AbstractPolynomial)

Return an antiderivative for `p`
"""
integrate(P::AbstractPolynomial) = throw(ArgumentError("`integrate` not implemented for polynomials of type $P"))

"""
    integrate(::AbstractPolynomial, C)

Returns the indefinite integral of the polynomial with constant `C` when expressed in the standard basis.
"""
function integrate(p::P, C) where {P <: AbstractPolynomial}
    ∫p = integrate(p)
    isnan(C) && return ⟒(P){eltype(∫p+C), indeterminate(∫p)}([C])
    ∫p + (C - constantterm(∫p))
end

"""
    integrate(::AbstractPolynomial, a, b)

Compute the definite integral of the given polynomial from `a` to `b`. Will throw an error if either `a` or `b` are out of the polynomial's domain.
"""
function integrate(p::AbstractPolynomial, a, b)
    P = integrate(p)
    return P(b) - P(a)
end

"""
    derivative(::AbstractPolynomial, order::Int = 1)

Returns a polynomial that is the `order`th derivative of the given polynomial. `order` must be non-negative.
"""
derivative(::AbstractPolynomial, ::Int)

"""
    truncate!(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0)

In-place version of [`truncate`](@ref)
"""
function truncate!(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
                   atol::Real = 0,) where {T}
    truncate!(p.coeffs, rtol=rtol, atol=atol)
    return chop!(p, rtol = rtol, atol = atol)
end

## truncate! underlying storage type
function truncate!(ps::Vector{T};
                   rtol::Real = Base.rtoldefault(real(T)),
                   atol::Real = 0,) where {T}
    max_coeff = norm(ps, Inf)
    thresh = max_coeff * rtol + atol
    for (i,pᵢ) ∈ pairs(ps)
        if abs(pᵢ) <= thresh
            ps[i] = zero(T)
        end
    end
    nothing
end

function truncate!(ps::Dict{Int,T};
                   rtol::Real = Base.rtoldefault(real(T)),
                   atol::Real = 0,) where {T}

    max_coeff = norm(values(ps), Inf)
    thresh = max_coeff * rtol + atol

    for (k,val) in  ps
        if abs(val) <= thresh
            pop!(ps,k)
        end
    end
    nothing
end

truncate!(ps::NTuple; kwargs...) = throw(ArgumentError("`truncate!` not defined."))
Base.truncate(ps::NTuple{0}; kwargs...) = ps
function Base.truncate(ps::NTuple{N,T};
              rtol::Real = Base.rtoldefault(real(T)),
              atol::Real = 0,) where {N,T}
    thresh = norm(ps, Inf) * rtol + atol
    return NTuple{N,T}(abs(pᵢ) <= thresh ? zero(T) : pᵢ for pᵢ ∈ values(ps))
end

"""
    truncate(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0)

Rounds off coefficients close to zero, as determined by `rtol` and `atol`, and then chops any leading zeros. Returns a new polynomial.
"""
function Base.truncate(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
    atol::Real = 0,) where {T}
    truncate!(deepcopy(p), rtol = rtol, atol = atol)
end

"""
    chop!(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0))

In-place version of [`chop`](@ref)
"""
function chop!(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
               atol::Real = 0,) where {T}
    chop!(p.coeffs, rtol=rtol, atol=atol)
    return p
end

# chop! underlying storage type
function chop!(ps::Vector{T};
               rtol::Real = Base.rtoldefault(real(T)),
               atol::Real = 0,) where {T}

    tol = norm(ps) * rtol + atol
    for i = lastindex(ps):-1:1
        val = ps[i]
        if abs(val) > tol
            resize!(ps, i);
            return nothing
        end
    end
    resize!(ps, 1)
    return nothing
end

function chop!(ps::Dict{Int,T};
    rtol::Real = Base.rtoldefault(real(T)),
    atol::Real = 0,) where {T}

    tol = norm(values(ps)) * rtol + atol

    for k in sort(collect(keys(ps)), by=x->x[1], rev=true)
        if  abs(ps[k]) > tol
            return nothing
        end
        pop!(ps, k)
    end

    return nothing
end

chop!(ps::NTuple; kwargs...) = throw(ArgumentError("chop! not defined"))
Base.chop(ps::NTuple{0}; kwargs...) = ps
function Base.chop(ps::NTuple{N,T};
              rtol::Real = Base.rtoldefault(real(T)),
              atol::Real = 0,) where {N,T}
    thresh = norm(ps, Inf) * rtol + atol
    for i in N:-1:1
        if abs(ps[i]) > thresh
            return ps[1:i]
        end
    end
    return NTuple{0,T}()
end

"""
    chop(::AbstractPolynomial{T};
        rtol::Real = Base.rtoldefault(real(T)), atol::Real = 0))

Removes any leading coefficients that are approximately 0 (using `rtol` and `atol`). Returns a polynomial whose degree will guaranteed to be equal to or less than the given polynomial's.
"""
function Base.chop(p::AbstractPolynomial{T};
    rtol::Real = Base.rtoldefault(real(T)),
    atol::Real = 0,) where {T}
    chop!(deepcopy(p), rtol = rtol, atol = atol)
end

"""
    check_same_variable(p::AbstractPolynomial, q::AbstractPolynomial)

Check if either `p` or `q` is constant or if `p` and `q` share the same variable
"""
check_same_variable(p::AbstractPolynomial, q::AbstractPolynomial) =
    (isconstant(p) || isconstant(q)) || indeterminate(p) ==  indeterminate(q)

function assert_same_variable(p::AbstractPolynomial, q::AbstractPolynomial)
    check_same_variable(p,q) || throw(ArgumentError("Polynomials have different indeterminates"))
end

function assert_same_variable(X::Symbol, Y::Symbol)
    X == Y ||  throw(ArgumentError("Polynomials have different indeterminates"))
end

#=
Linear Algebra =#
"""
    norm(::AbstractPolynomial, p=2)

Calculates the p-norm of the polynomial's coefficients
"""
function LinearAlgebra.norm(q::AbstractPolynomial, p::Real = 2)
    vs = values(q)
    isempty(vs) && return zero(eltype(q))
    return norm(vs, p) # if vs=() must be handled in special type
end

"""
    conj(::AbstractPolynomial)

Returns the complex conjugate of the polynomial
"""
LinearAlgebra.conj(p::P) where {P <: AbstractPolynomial} = map(conj, p)
LinearAlgebra.adjoint(p::P) where {P <: AbstractPolynomial} = map(adjoint, p)
LinearAlgebra.adjoint(A::VecOrMat{<:AbstractPolynomial}) = adjoint.(permutedims(A)) # default has indeterminate indeterminate
LinearAlgebra.transpose(p::AbstractPolynomial) = p
LinearAlgebra.transpose!(p::AbstractPolynomial) = p

#=
Conversions =#
Base.convert(::Type{P}, p::P) where {P <: AbstractPolynomial} = p
Base.convert(P::Type{<:AbstractPolynomial}, x) = P(x)
function Base.convert(::Type{S}, p::P) where {S <: Number,T, P<:AbstractPolynomial{T}}
    isconstant(p) && return convert(S, constantterm(p))
    throw(ArgumentError("Can't convert a nonconstant polynomial to type $S"))
end
function Base.convert(::Type{T}, p::P) where {T, P<:AbstractPolynomial{T}}
    isconstant(p) && return constantterm(p)
    throw(ArgumentError("Can't convert a nonconstant polynomial to type $S"))
end

# Methods to ensure that matrices of polynomials behave as desired
Base.promote_rule(::Type{<:AbstractPolynomial{T}},
                  ::Type{<:AbstractPolynomial{S}},
                  ) where {T,S} = Polynomial{promote_type(T, S)}
Base.promote_rule(::Type{P},::Type{Q}) where {T,X, P<:AbstractPolynomial{T,X},
                                              S,   Q<:AbstractPolynomial{S,X}} =
                                                   Polynomial{promote_type(T, S),X}
Base.promote_rule(::Type{P},::Type{Q}) where {T,X, P<:AbstractPolynomial{T,X},
                                              S,Y, Q<:AbstractPolynomial{S,Y}} =
                                                  assert_same_variable(X,Y)

#=
Inspection =#
"""
    length(::AbstractPolynomial)

The length of the polynomial.
"""
Base.length(p::AbstractPolynomial) = length(coeffs(p))
"""
    size(::AbstractPolynomial, [i])

Returns the size of the polynomials coefficients, along axis `i` if provided.
"""
Base.size(p::AbstractPolynomial) = size(coeffs(p))
Base.size(p::AbstractPolynomial, i::Integer) = size(coeffs(p), i)
Base.eltype(p::AbstractPolynomial{T}) where {T} = T
# in  analogy  with  polynomial as a Vector{T} with different operations defined.
Base.eltype(::Type{<:AbstractPolynomial}) = Float64
Base.eltype(::Type{<:AbstractPolynomial{T}}) where {T} = T
_eltype(::Type{<:AbstractPolynomial}) = nothing
_eltype(::Type{<:AbstractPolynomial{T}}) where {T} = T
function _eltype(P::Type{<:AbstractPolynomial}, p::AbstractPolynomial)
    T′ = _eltype(P)
    T = T′ == nothing ? eltype(p) : T′
    T
end
Base.iszero(p::AbstractPolynomial) = all(iszero, p)

# See discussions in https://github.com/JuliaMath/Polynomials.jl/issues/258
"""
    all(pred, poly::AbstractPolynomial)

Test whether all coefficients of an `AbstractPolynomial` satisfy predicate `pred`.

You can implement `isreal`, etc., to a `Polynomial` by using `all`.
"""
Base.all(pred, p::AbstractPolynomial) = all(pred, values(p))
"""
    any(pred, poly::AbstractPolynomial)

Test whether any coefficient of an `AbstractPolynomial` satisfies predicate `pred`.
"""
Base.any(pred, p::AbstractPolynomial) = any(pred, values(p))

"""
    map(fn, p::AbstractPolynomial, args...)

Transform coefficients of `p` by applying a function (or other callables) `fn` to each of them.

You can implement `real`, etc., to a `Polynomial` by using `map`.
"""
Base.map(fn, p::P, args...)  where {P<:AbstractPolynomial} = _convert(p, map(fn, coeffs(p), args...))

"""
    isreal(p::AbstractPolynomial)

Determine whether a polynomial is a real polynomial, i.e., having only real numbers as coefficients.

See also: [`real`](@ref)
"""
Base.isreal(p::AbstractPolynomial) = all(isreal, p)
"""
    real(p::AbstractPolynomial)

Construct a real polynomial from the real parts of the coefficients of `p`.

See also: [`isreal`](@ref)

!!! note
    This could cause losing terms in `p`. This method is usually called on polynomials like `p = Polynomial([1, 2 + 0im, 3.0, 4.0 + 0.0im])` where you want to chop the imaginary parts of the coefficients of `p`.
"""
Base.real(p::AbstractPolynomial) = map(real, p)

"""
    isintegral(p::AbstractPolynomial)

Determine whether a polynomial is an integer polynomial, i.e., having only integers as coefficients.
"""
isintegral(p::AbstractPolynomial) = all(isinteger, p)

"""
    ismonic(p::AbstractPolynomial)

Determine whether a polynomial is a monic polynomial, i.e., its leading coefficient is one.
"""
ismonic(p::AbstractPolynomial) = isone(convert(Polynomial, p)[end])

"`hasnan(p::AbstractPolynomial)` are any coefficients `NaN`"
hasnan(p::AbstractPolynomial) = any(hasnan, p)
hasnan(p::AbstractArray) = any(hasnan.(p))
hasnan(x) = isnan(x)

"""
    isconstant(::AbstractPolynomial)

Is the polynomial  `p` a constant.
"""
isconstant(p::AbstractPolynomial) = degree(p) <= 0

"""
    coeffs(::AbstractPolynomial)

Return the coefficient vector. For a standard basis polynomial these are `[a_0, a_1, ..., a_n]`.
"""
coeffs(p::AbstractPolynomial) = p.coeffs

# hook in for offset coefficients of Laurent Polynomials
_coeffs(p::AbstractPolynomial) = coeffs(p)

# specialize this to p[0] when basis vector is 1
"""
    constantterm(p::AbstractPolynomial)

return `p(0)`, the constant term in the standard basis
"""
constantterm(p::AbstractPolynomial{T}) where {T} = p(zero(T))

"""
    degree(::AbstractPolynomial)

Return the degree of the polynomial, i.e. the highest exponent in the polynomial that
has a nonzero coefficient. The degree of the zero polynomial is defined to be -1. The default method assumes the basis polynomial, `βₖ` has degree `k`.
"""
degree(p::AbstractPolynomial) = iszero(p) ? -1 : lastindex(p)

"""
    domain(::Type{<:AbstractPolynomial})

Returns the domain of the polynomial.
"""
domain(::Type{<:AbstractPolynomial})
domain(::P) where {P <: AbstractPolynomial} = domain(P)

"""
    mapdomain(::Type{<:AbstractPolynomial}, x::AbstractArray)
    mapdomain(::AbstractPolynomial, x::AbstractArray)

Given values of x that are assumed to be unbounded (-∞, ∞), return values rescaled to the domain of the given polynomial.

# Examples
```jldoctest  common
julia> using Polynomials

julia> x = -10:10
-10:10

julia> extrema(mapdomain(ChebyshevT, x))
(-1.0, 1.0)

```
"""
function mapdomain(P::Type{<:AbstractPolynomial}, x::AbstractArray)
    d = domain(P)
    x = collect(x)
    x_zerod = x .- minimum(x)
    x_scaled = x_zerod .* (last(d) - first(d)) ./ maximum(x_zerod)
    x_scaled .+= first(d)
    return x_scaled
end
mapdomain(::P, x::AbstractArray) where {P <: AbstractPolynomial} = mapdomain(P, x)

#=
indexing =#
Base.firstindex(p::AbstractPolynomial) = 0
Base.lastindex(p::AbstractPolynomial) = length(p) - 1
Base.eachindex(p::AbstractPolynomial) = 0:length(p) - 1
Base.broadcastable(p::AbstractPolynomial) = Ref(p)

# getindex
function Base.getindex(p::AbstractPolynomial{T}, idx::Int) where {T}
    m,M = firstindex(p), lastindex(p)
    idx < m && throw(BoundsError(p, idx))
    idx > M && return zero(T)
    p.coeffs[idx - m + 1]
end
Base.getindex(p::AbstractPolynomial, idx::Number) = getindex(p, convert(Int, idx))
Base.getindex(p::AbstractPolynomial, indices) = [p[i] for i in indices]
Base.getindex(p::AbstractPolynomial, ::Colon) = coeffs(p)

# setindex
function Base.setindex!(p::AbstractPolynomial, value, idx::Int)
    n = length(coeffs(p))
    if n ≤ idx
        resize!(p.coeffs, idx + 1)
        p.coeffs[n + 1:idx] .= 0
    end
    p.coeffs[idx + 1] = value
    return p
end

Base.setindex!(p::AbstractPolynomial, value, idx::Number) =
    setindex!(p, value, convert(Int, idx))
#Base.setindex!(p::AbstractPolynomial, value::Number, indices) =
#    [setindex!(p, value, i) for i in indices]
Base.setindex!(p::AbstractPolynomial, values, indices) =
    [setindex!(p, v, i) for (v, i) in tuple.(values, indices)]
#    [setindex!(p, v, i) for (v, i) in zip(values, indices)]
#Base.setindex!(p::AbstractPolynomial, value, ::Colon) =
#    setindex!(p, value, eachindex(p))
Base.setindex!(p::AbstractPolynomial, values, ::Colon) =
#        [setindex!(p, v, i) for (v, i) in zip(values, eachindex(p))]
    [setindex!(p, v, i) for (v, i) in tuple.(values, eachindex(p))]

#=
Iteration =#
## XXX breaking change in v2.0.0
#
# we want to keep iteration close to iteration over the coefficients as a vector:
# `iterate` iterates over coefficients, includes 0s
# `collect(T, p)` yields coefficients of `p` converted to type `T`
# `keys(p)` an iterator spanning `firstindex` to `lastindex` which *may* skip 0 terms (SparsePolynomial)
#    and *may* not be in order (SparsePolynomial)
# `values(p)` `pᵢ` for each `i` in `keys(p)`
# `pairs(p)`: `i => pᵢ` possibly skipping over values of `i` with `pᵢ == 0` (SparsePolynomial)
#    and possibly non ordered (SparsePolynomial)
# `monomials(p)`: iterates over pᵢ ⋅ basis(p, i) i  ∈ keys(p)
function Base.iterate(p::AbstractPolynomial, state=nothing)
    i = firstindex(p)
    if state == nothing
        return (p[i], i)
    else
        j = lastindex(p)
        if i <= state < j
            return (p[state+1], state+1)
        end
        return nothing
    end
end

# pairs map i -> aᵢ *possibly* skipping over ai == 0
# cf. abstractdict.jl
struct PolynomialKeys{P}
    p::P
end
struct PolynomialValues{P}
    p::P
end
Base.keys(p::AbstractPolynomial) =  PolynomialKeys(p)
Base.values(p::AbstractPolynomial) =  PolynomialValues(p)
Base.length(p::PolynomialValues) = length(p.p.coeffs)
Base.length(p::PolynomialKeys) = length(p.p.coeffs)
Base.size(p::Union{PolynomialValues, PolynomialKeys}) = (length(p),)
function Base.iterate(v::PolynomialKeys, state=nothing)
    i = firstindex(v.p)
    state==nothing && return (i, i)
    j = lastindex(v.p)
    i <= state < j && return (state+1, state+1)
    return nothing
end

function Base.iterate(v::PolynomialValues, state=nothing)
    i = firstindex(v.p)
    state==nothing && return (v.p[i], i)
    j = lastindex(v.p)
    i <= state < j && return (v.p[state+1], state+1)
    return nothing
end

# iterate over monomials of the polynomial
struct Monomials{P}
    p::P
end
"""
    monomials(p::AbstractPolynomial)

Returns an iterator over the terms, `pᵢ⋅basis(p,i)`, of the polynomial for each `i` in `keys(p)`.
"""
monomials(p) = Monomials(p)
function Base.iterate(v::Monomials, state...)
    y = iterate(pairs(v.p), state...)
    y == nothing && return nothing
    kv, s = y
    return (kv[2]*basis(v.p, kv[1]), s)
end
Base.length(v::Monomials) = length(keys(v.p))

#=
identity =#
Base.copy(p::P) where {P <: AbstractPolynomial} = _convert(p, copy(coeffs(p)))
Base.hash(p::AbstractPolynomial, h::UInt) = hash(indeterminate(p), hash(coeffs(p), h))

# get symbol of polynomial. (e.g. `:x` from 1x^2 + 2x^3...
#_indeterminate(::Type{P}) where {T, X, P <: AbstractPolynomial{T, X}} = X
_indeterminate(::Type{P}) where {P <: AbstractPolynomial} = nothing
_indeterminate(::Type{P}) where {T, X, P <: AbstractPolynomial{T,X}} = X
function indeterminate(::Type{P}) where {P <: AbstractPolynomial}
    X = _indeterminate(P)
    X == nothing ? :x : X
end
indeterminate(p::P) where {P <: AbstractPolynomial} = _indeterminate(P)
function indeterminate(PP::Type{P}, p::AbstractPolynomial{T,Y}) where {P <: AbstractPolynomial, T,Y}
    X = _indeterminate(PP)
    X == nothing && return Y
    assert_same_variable(X,Y)
    return X
    #X = _indeterminate(PP) == nothing ? indeterminate(p) :  _indeterminate(PP)
end
function indeterminate(PP::Type{P}, x::Symbol) where {P <: AbstractPolynomial}
    X = _indeterminate(PP) == nothing ? x :  _indeterminate(PP)
end

#=
zero, one, variable, basis =#

"""
    zero(::Type{<:AbstractPolynomial})
    zero(::AbstractPolynomial)

Returns a representation of 0 as the given polynomial.
"""
function Base.zero(::Type{P}) where {P<:AbstractPolynomial}
    T,X = eltype(P), indeterminate(P)
    ⟒(P){T,X}(zeros(T,1))
end
Base.zero(::Type{P}, var::SymbolLike) where {P <: AbstractPolynomial} = zero(⟒(P){eltype(P),Symbol(var)}) #default 0⋅b₀
Base.zero(p::P, var=indeterminate(p)) where {P <: AbstractPolynomial} = zero(P, var)
"""
    one(::Type{<:AbstractPolynomial})
    one(::AbstractPolynomial)

Returns a representation of 1 as the given polynomial.
"""
Base.one(::Type{P}) where {P<:AbstractPolynomial} = throw(ArgumentError("No default method defined")) # no default method
Base.one(::Type{P}, var::SymbolLike) where {P <: AbstractPolynomial} = one(⟒(P){eltype(P), Symbol(var == nothing ? :x : var)})
Base.one(p::P, var=indeterminate(p)) where {P <: AbstractPolynomial} = one(P, var)

Base.oneunit(::Type{P}, args...) where {P <: AbstractPolynomial} = one(P, args...)
Base.oneunit(p::P, args...) where {P <: AbstractPolynomial} = one(p, args...)

"""
    variable(var=:x)
    variable(::Type{<:AbstractPolynomial}, var=:x)
    variable(p::AbstractPolynomial, var=indeterminate(p))

Return the monomial `x` in the indicated polynomial basis.  If no type is give, will default to [`Polynomial`](@ref). Equivalent  to  `P(var)`.

# Examples
```jldoctest  common
julia> using Polynomials

julia> x = variable()
Polynomial(x)

julia> p = 100 + 24x - 3x^2
Polynomial(100 + 24*x - 3*x^2)

julia> roots((x - 3) * (x + 2))
2-element Vector{Float64}:
 -2.0
  3.0

```
"""
variable(::Type{P}) where {P <: AbstractPolynomial} = throw(ArgumentError("No default method defined")) # no default
variable(::Type{P}, var::SymbolLike) where {P <: AbstractPolynomial} = variable(⟒(P){eltype(P),Symbol(var)})
variable(p::AbstractPolynomial, var = indeterminate(p)) = variable(typeof(p), var)
variable(var::SymbolLike = :x) = variable(Polynomial{Int}, var)

# basis
# var is a positional argument, not a keyword; can't deprecate so we do `_var; var=_var`
# return the kth basis polynomial for the given polynomial type, e.g. x^k for Polynomial{T}
function basis(::Type{P}, k::Int) where {P<:AbstractPolynomial}
    T,X = eltype(P), indeterminate(P)
    zs = zeros(T, k+1)
    zs[end] = one(T)
    ⟒(P){eltype(P), indeterminate(P)}(zs)
end
function basis(::Type{P}, k::Int, _var::SymbolLike; var=_var) where {P <: AbstractPolynomial}
    T,X = eltype(P), Symbol(var)
    basis(⟒(P){T,X}, k)
end
basis(p::P, k::Int, _var=indeterminate(p); var=_var) where {P<:AbstractPolynomial} = basis(P, k, var)

#=
arithmetic =#
Base.:-(p::P) where {P <: AbstractPolynomial} = _convert(p, -coeffs(p))

Base.:*(p::AbstractPolynomial, c::Number) = scalar_mult(p, c)
Base.:*(c::Number, p::AbstractPolynomial) = scalar_mult(c, p)
Base.:*(c::T, p::P) where {T, P <: AbstractPolynomial{T}} = scalar_mult(c, p)
Base.:*(p::P, c::T) where {T, P <: AbstractPolynomial{T}} = scalar_mult(p, c)

# implicitly identify c::Number with a constant polynomials
Base.:+(c::Number, p::AbstractPolynomial) = +(p, c)
Base.:-(p::AbstractPolynomial, c::Number) = +(p, -c)
Base.:-(c::Number, p::AbstractPolynomial) = +(-p, c)

# scalar operations
# no generic p+c, as polynomial addition falls back to scalar ops
#function Base.:+(p::P, n::Number) where {P <: AbstractPolynomial}
#    p1, p2 = promote(p, n)
#    return p1 + p2
#end

Base.:-(p1::AbstractPolynomial, p2::AbstractPolynomial) = +(p1, -p2)

## addition
## Fall back addition is possible as vector addition with padding by 0s
## Subtypes will likely want to implement both:
## +(p::P,c::Number) and +(p::P, q::Q) where {T,S,X,P<:SubtypePolynomial{T,X},Q<:SubtypePolynomial{S,X}}
## though the default for poly+poly isn't terrible

# polynomial + scalar; implicit identification of c with c*one(P)
Base.:+(p::P, c::T) where {T,X, P<:AbstractPolynomial{T,X}} = p + c * one(P)

function Base.:+(p::P, c::S) where {T,X, P<:AbstractPolynomial{T,X}, S}
    R = promote_type(T,S)
    q = convert(⟒(P){R,X}, p)
    q + R(c)
end

# polynomial + polynomial when different types
function Base.:+(p::P, q::Q) where {T,X,P <: AbstractPolynomial{T,X}, S,Y,Q <: AbstractPolynomial{S,Y}}
    isconstant(p) && return constantterm(p) + q
    isconstant(q) && return p + constantterm(q)
    assert_same_variable(X,Y)
    sum(promote(p,q))

end

# Works when p,q of same type.
# For Immutable, must remove N,M bit;
# for others, can widen to Type{T,X}, Type{S,X} to avoid a promotion
function Base.:+(p::P, q::P) where {T,X,P<:AbstractPolynomial{T,X}}
    cs = degree(p) >= degree(q)  ? ⊕(P, p.coeffs, q.coeffs) : ⊕(P, q.coeffs, p.coeffs)
    return P(cs)
end

# addition of polynomials is just vector space addition, so can be done regardless
# of basis, as long as the same. These ⊕ methods try to find a performant means to add
# to sets of coefficients based on the storage type. These assume n1 >= n2
function ⊕(P::Type{<:AbstractPolynomial}, p1::Vector{T}, p2::Vector{S}) where {T,S}

    n1, n2 = length(p1), length(p2)
    R = promote_type(T,S)

    cs = collect(R,p1)
    for i in 1:n2
        cs[i] += p2[i]
    end

    return cs
end

# Padded vector sum of two tuples assuming N ≥ M
@generated function ⊕(P::Type{<:AbstractPolynomial}, p1::NTuple{N,T}, p2::NTuple{M,S}) where {T,N,S,M}

    exprs = Any[nothing for i = 1:N]
    for i in  1:M
        exprs[i] = :(p1[$i] + p2[$i])
    end
    for i in M+1:N
        exprs[i] =:(p1[$i])
    end

    return quote
        Base.@_inline_meta
        tuple($(exprs...))
    end

end

# addition when a dictionary is used for storage
function ⊕(P::Type{<:AbstractPolynomial}, p1::Dict{Int,T}, p2::Dict{Int,S}) where {T,S}

    R = promote_type(T,S)
    p = Dict{Int, R}()

    # this allocates in the union
#    for i in union(eachindex(p1), eachindex(p2))
#        p[i] = p1[i] + p2[i]
#    end

    for (i,pi) ∈ pairs(p1)
        @inbounds p[i] = pi + get(p2, i, zero(R))
    end
    for (i,pi) ∈ pairs(p2)
        if iszero(get(p,i,zero(R)))
            @inbounds p[i] = get(p1, i, zero(R)) + pi
        end
    end

    return  p

end

## -- multiplication

# this fall back not necessarily efficient (e.g., sparse)
function scalar_mult(p::P, c::S) where {S, T, X, P<:AbstractPolynomial{T,X}}
    R = Base.promote_op(*, T, S) # typeof(one(T)*one(S))?
    𝐏 = ⟒(P){R,X}
    𝐏([pᵢ * c for pᵢ ∈ coeffs(p)])
end

function scalar_mult(c::S, p::P) where {S, T, X, P<:AbstractPolynomial{T, X}}
    R = Base.promote_op(*, T, S)
    𝐏 = ⟒(P){R,X}
    𝐏([c * pᵢ for pᵢ ∈ coeffs(p)])
end

function Base.:/(p::P, c::S) where {P <: AbstractPolynomial,S}
    _convert(p, coeffs(p) ./ c)
end

## polynomial p*q
## Polynomial multiplication formula depend on the particular basis used. The subtype must implement
function Base.:*(p1::P, p2::Q) where {T,X,P <: AbstractPolynomial{T,X},S,Y,Q <: AbstractPolynomial{S,Y}}
    isconstant(p1) && return constantterm(p1) * p2
    isconstant(p2) && return p1 * constantterm(p2)
    assert_same_variable(X, Y)
    p1, p2 = promote(p1, p2)
    return p1 * p2
end

Base.:^(p::AbstractPolynomial, n::Integer) = Base.power_by_squaring(p, n)

function Base.divrem(num::P, den::O) where {P <: AbstractPolynomial,O <: AbstractPolynomial}
    n, d = promote(num, den)
    return divrem(n, d)
end

"""
    gcd(a::AbstractPolynomial, b::AbstractPolynomial; atol::Real=0, rtol::Real=Base.rtoldefault)

Find the greatest common denominator of two polynomials recursively using
[Euclid's algorithm](http://en.wikipedia.org/wiki/Polynomial_greatest_common_divisor#Euclid.27s_algorithm).

# Examples

```jldoctest common
julia> using Polynomials

julia> gcd(fromroots([1, 1, 2]), fromroots([1, 2, 3]))
Polynomial(4.0 - 6.0*x + 2.0*x^2)

```
"""
function Base.gcd(p1::AbstractPolynomial{T}, p2::AbstractPolynomial{S}; kwargs...) where {T,S}
    gcd(promote(p1, p2)...; kwargs...)
end

function Base.gcd(p1::AbstractPolynomial{T}, p2::AbstractPolynomial{T};
                  atol::Real=zero(real(T)),
                  rtol::Real=Base.rtoldefault(real(T))
                  ) where {T}

    r₀, r₁ = p1, p2
    iter = 1
    itermax = length(r₁)

    while !iszero(r₁) && iter ≤ itermax
        _, rtemp = divrem(r₀, r₁)
        r₀ = r₁
        r₁ = truncate(rtemp; atol=atol, rtol=rtol)
        iter += 1
    end
    return r₀
end

"""
    uvw(p,q; kwargs...)

return `u` the gcd of `p` and `q`, and `v` and `w`, where `u*v = p` and `u*w = q`.
"""
uvw(p::AbstractPolynomial, q::AbstractPolynomial; kwargs...) = uvw(promote(p,q)...; kwargs...)
uvw(p1::P, p2::P; kwargs...) where {P <:AbstractPolynomial} = throw(ArgumentError("uvw not defind"))

"""
    div(::AbstractPolynomial, ::AbstractPolynomial)
"""
Base.div(n::AbstractPolynomial, d::AbstractPolynomial) = divrem(n, d)[1]

"""
    rem(::AbstractPolynomial, ::AbstractPolynomial)
"""
Base.rem(n::AbstractPolynomial, d::AbstractPolynomial) = divrem(n, d)[2]

#=
Comparisons =#
Base.isequal(p1::P, p2::P) where {P <: AbstractPolynomial} = hash(p1) == hash(p2)
Base.:(==)(p1::AbstractPolynomial, p2::AbstractPolynomial) =
    check_same_variable(p1,p2) && (coeffs(p1) == coeffs(p2))
Base.:(==)(p::AbstractPolynomial, n::Number) = degree(p) <= 0 && constantterm(p) == n
Base.:(==)(n::Number, p::AbstractPolynomial) = p == n

function Base.isapprox(p1::AbstractPolynomial{T,X},
                       p2::AbstractPolynomial{S,Y};
                       rtol::Real = (Base.rtoldefault(T, S, 0)),
                       atol::Real = 0,) where {T,X,S,Y}
    assert_same_variable(p1, p2)
    (hasnan(p1) || hasnan(p2)) && return false  # NaN poisons comparisons
    # copy over from abstractarray.jl
    Δ  = norm(p1-p2)
    if isfinite(Δ)
        return Δ <= max(atol, rtol*max(norm(p1), norm(p2)))
    else
        for i in 0:max(degree(p1), degree(p2))
            isapprox(p1[i], p2[i]; rtol=rtol, atol=atol) || return false
        end
        return true
    end
end

function Base.isapprox(p1::AbstractPolynomial{T},
                       n::S;
                       rtol::Real = (Base.rtoldefault(T, S, 0)),
                       atol::Real = 0,) where {T,S}
    return isapprox(p1, _convert(p1, [n]))
end

Base.isapprox(n::S,
    p1::AbstractPolynomial{T};
    rtol::Real = (Base.rtoldefault(T, S, 0)),
    atol::Real = 0,) where {T,S} = isapprox(p1, n, rtol = rtol, atol = atol)


#=

    This is a port of the R2 model SurveyResultUnbias2.cs

    With the datafiles
    - survey_result_unbias2_population.csv
    - survey_result_unbias2_gender.csv
    - survey_result_unbias2_person_gender.csv
    - survey_result_unbias2_answer.csv

    Note: there is only 5 observations in the answer file, but reading the files makes it look
          more general...

    Result from the R2 model:
    ```
    <0>[0] Mean: 0.590981              bias[0]
    <0>[0] Variance: 0.0427666

    <0>[1] Mean: 0.472561              bias[1]
    <0>[1] Variance: 0.0494528

    <1>[0] Mean: 591.075               votes[0]
    <1>[0] Variance: 43109

    <1>[1] Mean: 944.156               votes[1]
    <1>[1] Variance: 197880
    ```

    This model:

Summary Statistics
  parameters        mean        std   naive_se      mcse         ess      rhat   ess_per_sec
      Symbol     Float64    Float64    Float64   Float64     Float64   Float64       Float64

     bias[1]      0.5994     0.2008     0.0020    0.0025   4269.9355    0.9999      385.8259
     bias[2]      0.5022     0.2250     0.0023    0.0035   3757.7959    1.0002      339.5496
    votes[1]    599.4513   201.1412     2.0114    2.5556   4243.2179    0.9999      383.4118
    votes[2]   1004.9705   450.3128     4.5031    6.9752   3788.1458    1.0002      342.2920

Quantiles
  parameters       2.5%      25.0%       50.0%       75.0%       97.5%
      Symbol    Float64    Float64     Float64     Float64     Float64

     bias[1]     0.1927     0.4562      0.6150      0.7582      0.9293
     bias[2]     0.1006     0.3207      0.5015      0.6832      0.9061
    votes[1]   191.9750   456.0000    614.0000    754.2500    931.0000
    votes[2]   201.0000   643.0000   1002.0000   1369.0000   1812.0000

=#

using Turing, StatsPlots, Distributions, StatsBase
using CSV
include("jl_utils.jl")

@model function survey_result_unbias2(population, person_gender, answer)
    pop_len = length(population)
    answer_len = length(answer)

    bias ~ filldist(Beta(1,1),pop_len)
    votes = tzeros(pop_len)
    for i in 1:pop_len
        votes[i] ~ Binomial(population[i],bias[i])
    end

    ansBias =  Vector{Real}(undef, answer_len) # tzeros(person_gender_len)
    for i in 1:answer_len
        ansBias[i] = bias[person_gender[i]+1]
        answer[i] ~ Bernoulli(ansBias[i])
    end

end

# 1000 2000
population = parse.(Int64,split(readline("survey_result_unbias2_population.csv"),","))
# 0 1
# gender = parse.(Int64,split(readline("survey_result_unbias2_gender.csv"),","))
# 0 1 0 1 0
person_gender = parse.(Int64,split(readline("survey_result_unbias2_person_gender.csv"),","))
# 1 0 1 1 0
answer = parse.(Int64,split(readline("survey_result_unbias2_answer.csv"),","))

model = survey_result_unbias2(population, person_gender, answer)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 1_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 1_000)
# chns = sample(model, IS(), 1_000)

# chns = sample(model, NUTS(), 1_000)
# chns = sample(model, HMC(0.1,6), 1_000)

display(chns)


## Adapted from  doc/libsundials-serial-dev/examples/ida/serial/idaRoberts_dns.c and
##               sundialsTB/ida/examples_ser/midasRoberts_dns.m

## /*
##  * -----------------------------------------------------------------
##  * $Revision: 1.2 $
##  * $Date: 2009/01/21 21:46:40 $
##  * -----------------------------------------------------------------
##  * Programmer(s): Allan Taylor, Alan Hindmarsh and
##  *                Radu Serban @ LLNL
##  * -----------------------------------------------------------------
##  * This simple example problem for IDA, due to Robertson,
##  * is from chemical kinetics, and consists of the following three
##  * equations:
##  *
##  *      dy1/dt = -.04*y1 + 1.e4*y2*y3
##  *      dy2/dt = .04*y1 - 1.e4*y2*y3 - 3.e7*y2**2
##  *         0   = y1 + y2 + y3 - 1
##  *
##  * on the interval from t = 0.0 to t = 4.e10, with initial
##  * conditions: y1 = 1, y2 = y3 = 0.
##  *
##  * While integrating the system, we also use the rootfinding
##  * feature to find the points at which y1 = 1e-4 or at which
##  * y3 = 0.01.
##  *
##  * The problem is solved with IDA using IDADENSE for the linear
##  * solver, with a user-supplied Jacobian. Output is printed at
##  * t = .4, 4, 40, ..., 4e10.
##  * -----------------------------------------------------------------
##  */

using Sundials

## Define the system residual function.
function resrob(tres, yy, yp, rr, user_data)
    yval = Sundials.asarray(yy)
    ypval = Sundials.asarray(yp)
    rval = Sundials.asarray(rr)
    rval[1]  = -0.04*yval[1] + 1.0e4*yval[2]*yval[3]
    rval[2]  = -rval[1] - 3.0e7*yval[2]*yval[2] - ypval[2]
    rval[1] -=  ypval[1]
    rval[3]  =  yval[1] + yval[2] + yval[3] - 1.0
    return int32(0)   # indicates normal return
end

## Root function routine. Compute functions g_i(t,y) for i = 0,1.
function grob(t, yy, yp, gout, user_data)
    yval = Sundials.asarray(yy)
    gval = Sundials.asarray(gout, (2,))
    gval[1] = yval[1] - 0.0001
    gval[2] = yval[3] - 0.01
    return int32(0)   # indicates normal return
end

## Define the Jacobian function. BROKEN - JJ is wrong
function jacrob(Neq, tt, cj, yy, yp, resvec,
                JJ, user_data, tempv1, tempv2, tempv3)
    yval = Sundials.asarray(yy)
    JJ = pointer_to_array(convert(Ptr{Float64}, JJ), (3,3))
    JJ[1,1] = -0.04 - cj
    JJ[2,1] = 0.04
    JJ[3,1] = 1.0
    JJ[1,2] = 1.0e4*yval[3]
    JJ[2,2] = -1.0e4*yval[3] - 6.0e7*yval[2] - cj
    JJ[3,2] = 1.0
    JJ[1,3] = 1.0e4*yval[2]
    JJ[2,3] = -1.0e4*yval[2]
    JJ[3,3] = 1.0
    return int32(0)   # indicates normal return
end

neq = int32(3)
nout = 12
t0 = 0.0
yy = [1.0,0.0,0.0]
yp = [-0.04,0.04,0.0]
rtol = 1e-4
avtol = [1e-8, 1e-14, 1e-6]
tout1 = 0.4

mem = Sundials.IDACreate()
retval = Sundials.IDAInit(mem, resrob, t0, yy, yp)
retval = Sundials.IDASVtolerances(mem, rtol, avtol)

## Call IDARootInit to specify the root function grob with 2 components
retval = Sundials.IDARootInit(mem, int32(2), grob)

## Call IDADense and set up the linear solver.
retval = Sundials.IDADense(mem, neq)
## retval = Sundials.IDADlsSetDenseJacFn(mem, jacrob)

iout = 0
tout = tout1
tret = [1.0]
rootsfound = int32([0, 0])

while true
    retval = Sundials.IDASolve(mem, tout, tret, yy, yp, Sundials.IDA_NORMAL)
    println("T = ", tout, ", Y = ", yy)
    if retval == Sundials.IDA_ROOT_RETURN
        retvalr = Sundials.IDAGetRootInfo(mem, rootsfound)
        println("roots = ", rootsfound)
    end
    if retval == Sundials.IDA_SUCCESS
        iout += 1
        tout *= 10.0
    end
    if iout == nout break end
end


function compute_kinetic_limit(parameter_dictionary::Dict{String,Any})::Float64

    try

        # grab the parameter values from the dictionary -
        E1 = parameter_dictionary["PFK_concentration"]
        ATP = parameter_dictionary["ATP_concentration"]
        F6P = parameter_dictionary["F6P_concentration"]
        kcat = parameter_dictionary["kcat"]
        K_ATP = parameter_dictionary["K_ATP"]
        K_F6P = parameter_dictionary["K_F6P"]

        # compute the kinetic limit -
        kinetic_limit = kcat * E1 * (ATP / (K_ATP + ATP)) * (F6P / (K_F6P + F6P))

        # return -
        return kinetic_limit
    catch error
        throw(error)
    end
end

function compute_activity_variable(effector::Float64, parameter_dictionary::Dict{String,Any})::Float64

    try

        # grab the parameter values from the dictionary -
        W1 = parameter_dictionary["W1"]
        W2 = parameter_dictionary["W2"]
        KD = parameter_dictionary["KD"]
        n = parameter_dictionary["n"]

        # compute the f -
        f = ((effector / KD)^n) / (1 + (effector / KD)^n)

        # compute the v-variable -
        v_variable = (W1 + W2 * f) / (1 + W1 + W2 * f)

        # return -
        return v_variable
    catch error
        throw(error)
    end
end

function compute_activity_variable(effector_array::Array{Float64,1}, parameter_dictionary::Dict{String,Any})::Array{Float64,1}

    # initialize -
    v_variable_array = Array{Float64,1}()

    try

        # process each effector -
        for effector_value in effector_array

            # compute -
            value = compute_activity_variable(effector_value, parameter_dictionary)

            # cache -
            push!(v_variable_array, value)
        end

        # return -
        return v_variable_array
    catch error
        throw(error)
    end
end


# Source: https://www.johndcook.com/blog/2020/02/08/arenstorf-orbit/

using Revise
using LinearAlgebra
using Plots

includet("src/InternalArray.jl")
includet("src/FixedStepMethods.jl")
includet("src/MixedStepMethods.jl")
IArray = InternalArray

mmoon = 0.012277471
mearth = 1-mmoon
p0 = [0.994, 0]
v0 = [0, -2.001585106]
target_err = 1e-4

function derivative(t, x)
    if ndims(x) == 2
        r = zeros(size(x))
        for i in 1:size(x, 2)
            r[:,i] = derivative(t, x[:,i])
        end
        return r
    end
    D1 = ((x[1]+mmoon)*(x[1]+mmoon) + x[2]*x[2])
    D1 *= sqrt(D1)
    D2 = ((x[1]-mearth)*(x[1]-mearth) + x[2]*x[2])
    D2 *= sqrt(D2)
    return [
        x[3],
        x[4],
        x[1] + 2*x[4] - mearth*(x[1]+mmoon)/D1 - mmoon*(x[1]-mearth)/D2,
        x[2] - 2*x[3] - mearth*x[2]/D1 - mmoon*x[2]/D2
    ]
end

t, x = mixed_step_ode_solver(
    derivative,
    0,
    [p0; v0],
    (t, x) ->
        last(t) > 5 && norm(p0 .- last(x)[1:2]) < target_err*last(t) || last(t) > 50,
    min_step = 1e-7,
    max_step = 0.1,
    target_err = target_err,
    method = mod_euler_step,
);

plot(
    IArray(x, 1),
    IArray(x, 2),
    label = "satellite",
    title = "Arenstorf orbit",
    aspect_ratio = :equal,
    framestyle = :box,
    xlims = (-1.4, 1.2),
    legend = :bottomright
)
plot!(cos, sin, 0, 2π, label = "moon", color = :black, style = :dot)
scatter!((0, 0), label = "", color = :black, markersize = 1.5)
# savefig("media/arenstorf_orbit.png")

using IJulia

px = [x[i][1] for i in 1:length(x)]
py = [x[i][2] for i in 1:length(x)]
anim = Animation()
i = 1
for j = 0:0.1:last(t)
    while i < length(t) && t[i] < j
        i += 1
    end
    plot(
        px[1:i],
        py[1:i],
        label = "satellite",
        title = "Arenstorf orbit",
        aspect_ratio = :equal,
        framestyle = :box,
        xlims = (-1.4, 1.2),
        ylims = (-1.3, 1.4),
        legend = :bottomright,
    )
    inds = range(0, 2*π*j/last(t), length = 150)
    plot!(cos.(inds), sin.(inds), label = "moon", color = :black, style = :dot)
    scatter!((0, 0), label = "", color = :black, markersize = 1.5)
    frame(anim)
end
gif(anim, "media/arenstorf_orbit.gif", fps = 15)


# Some constant parameters for the Jaynes-Cummings model

# Convenience symbol
const ⊗ = kron

# Physical quantities
const ħ = Complex128(1.05457173E-34) # in Joules*seconds

# Matrices
const σ_z      = Array{Complex128}([1 0 ; 0 -1])
const σ_plus   = Array{Complex128}([0 0 ; 1 0])
const σ_minus  = Array{Complex128}([0 1 ; 0 0])

# Wigner function space
const MAXDISP = 2.5 # maximum displacement operator value, real and imaginary


## Test Laplace transform pairs

"""
    TransformPair{T,V}

A pair of functions for analyzing an inverse Laplace transform method.
Field `ft` is the real-space function. Field `fs` is the Laplace-space function.
"""
struct TransformPair{T,V}
    ft::T
    fs::V
end

"""
    ILtPair(ilt::AbstractILt, ft::Function)

return an object of type `ILtPair` that associates `ilt` the inverse Laplace transform of
a function with it "exact" numerical inverse `ft`. Calling `abserr(p, t)` returns the
absolute error between the inverse transform and the exact value.

# Example

This example compares the inversion using the Weeks algorithm of the Laplace transform
of `cos(t)` to its exact value at `t=1.0`.
```
julia> p = ILtPair(Weeks(s -> s/(1+s^2)), cos);
julia> abserr(p, 1.0)

0.0
```
"""
struct ILtPair{T,V}
    ilt::T
    ft::V
end

# type ILtPair{T,V} <: AbstractILt
"""
    iltpair_power(n)

Return a `TransformPair` for the power function `x^n`. This can be used
to construct an `ILTPair`.

```julia-repl
julia> p  = Talbot(iltpair_power(3));

julia> Float64(abserr(p, 1)) # test Talbot method for Laplace transform of x^3, evaluated at 1
2.0820366247539812e-26
```
"""
function iltpair_power(n)
    fs = s -> s^(-n - 1) * SpecialFunctions.gamma(1 + n)
    ft = t -> t^n
    return TransformPair(ft, fs)
end

"""
    abserr(p::ILtPair, t)

Compute the absolute error between the estimated inverse Laplace transform and
"exact" numerical solution contained in `p` at the point `t`. See `ILtPair`.
"""
abserr(p::ILtPair, t) = abs(p.ilt(t) - p.ft(t))

## FIXME: Consider forwarding from MacroTools
# Create ILtPair with ilt type
for ilttype in (:Talbot, :GWR, :Weeks)
    @eval $(ilttype)(p::TransformPair, args...) = ILtPair($(ilttype)(p.fs, args...), p.ft)
end

for f in (:optimize,)
    @eval $(f)(p::ILtPair, args...) = $(f)(p.ilt, args...)
end


"""
NsolPDE(n; sham=1, resdec=.5, rtol=1.e-7, atol=1.e-10)

Solve the Elliptic PDE using nsol.jl on an n x n grid. You give me
n and (optionally) sham and resdec and I return the output of nsol.
"""
function NsolPDE(n; sham = 1, resdec = 0.5, rtol = 1.e-7, atol = 1.e-10)
    # Get some room for the residual
    u0 = zeros(n * n)
    FV = copy(u0)
    # Get the precomputed data from pdeinit
    pdata = pdeinit(n)
    # Storage for the Jacobian
    J = copy(pdata.D2)
    # Call the solver
    hout = nsol(
        pdeF!,
        u0,
        FV,
        J,
        pdeJ!;
        rtol = rtol,
        atol = atol,
        pdata = pdata,
        sham = sham,
        resdec = resdec,
    )
    return hout
end


function power_posterior(cm::ConditionalModel)
    m = cm.model

    pr = prior(m, observed(cm)...)
    lik = likelihood(m, observed(cm)...)

    avs = argvals(cm)
    obs = observations(cm)
    function f(t, x)
        logdensity(pr(avs), x) + t * logdensity(lik(merge(avs, x)) | obs)
    end

    return f
end

# EXAMPLE
# m = @model begin
#     x ~ Normal()
#     y ~ Normal(x,1) |> iid(3)
#     return y
# end;

# prior(m, :y)

# likelihood(m, :y)

# y = rand(m())

# f = power_posterior(m() | (;y))

# f(1.0, (x=0.2,))
# f(0.0, (x=0.2,))


using Distributed
using Images
using FileIO
using Mmap

# calculate a distributed fractal
width = 40000
height = 20000

rmin = -2.5
rmax = 1.5

imin = -1.25
imax = 1.25

iter = 500
epsilon = 0.25

mapfile = "zset-$(width)-$(height)"
s = open(mapfile)
n = read(s, Int)
println("size of zset is $n")
if n != (width * height)
	@warn("zset doesn't match dims")
	exit()
end

zset = Mmap.mmap(s, Array{ComplexF64,1}, n)
close(s)

num_procs =  0
hosts = ["robert@mbp2018.local"]

for i in 1:num_procs
	println("started proc $(i)")
	sshflags=`-i /Users/robert/src/julia/juliaSnippets/distributed/ms-keys`

	addprocs(hosts, tunnel=false, sshflags=sshflags, max_parallel=10, dir="/Users/robert/tmp", topology=:all_to_all)

end

# define mandel everywhere
include("mandel.jl")

println("starting pmap on $(workers())")
#
#generate image from results
# time measurements
print("starting...\n")
tStart=time()
#zset = Array{ComplexF64,1}(undef,0)
m_set = mandelbrot_set(zset, rmin, rmax, imin, imax, width, height, iter, epsilon)
tStop = time()

# write the image-file
img = colorview(RGB, m_set)
save("mandel.jpg", img)

print("done. took ", tStop-tStart, " seconds\n");


using DReal
reset_ctx!()

a = Var(Float64,"aexpop",0.,1.)
b2 = (a < 0.1) | (a > 0.9)
push_ctx!()
add!(a>0.0)
add!(a<1.0)
add!(b2)
is_satisfiable()
pop_ctx!()
push_ctx!()
add!(a>0.0)
add!(a<1.0)
add!(!b2)
is_satisfiable()


using ApproxFun, SingularIntegralEquations, Test
    import SingularIntegralEquations: ⁻, ⁺

@testset "Wiener–Hopf" begin
    # Scalar case

    G=Fun(z->2+cos(z+1/z),Circle()) # the symbol of the Toeplitz operator
    T=ToeplitzOperator(G)

    C  = Cauchy(-1)
    V  = (I+(1-G)*C)\(G-1)

    Φmi = 1+C*V
    Φp = V+Φmi

    L  = ToeplitzOperator(1/Φmi)
    U  = ToeplitzOperator(Φp)

    @test norm((T-U*L)[1:10,1:10]) < 10eps()  # check the accuracy

    # Matrix case

    G=Fun(z->[-1 -3; -3 -1]/z +
             [ 2  2;  1 -3] +
             [ 2 -1;  1  2]*z,Circle())

    H = Hilbert(space(G),1)
    @test domainspace(H) == rangespace(H) == space(G)
    C  = Cauchy(-1)

    A=(I+(I-G)*C)

    F=(G-I)[:,1]
    F̃ = vec(F)

    C2=promotedomainspace(C.ops[2],space(F))

    testblockbandedoperator(Hilbert(space(F)[1]))
    testblockbandedoperator(Hilbert(space(F)))
    testblockbandedoperator(C2)

    @test norm((C*F - [C*F[1];C*F[2]]).coefficients) == 0
    @test norm((C*G - [C*G[1] C*G[3];C*G[2] C*G[4]]).coefficients) == 0
    @test cauchy(F,exp(0.1im)⁻) ≈ (C*F)(exp(0.1im))
    @test cauchy(G,exp(0.1im)⁻) ≈ (C*G)(exp(0.1im))

    V1 = A\F
    Ṽ1 = A\F̃

    A1=choosespaces(A,(G-I)[:,1])
    A2=choosespaces(A,Fun((G-I)[:,1]))

    @test A1\Fun((G-I)[:,1])  == V1
    @test A1\(G-I)[:,1]  == V1

    QR=qrfact(A1)

    @test QR\Fun((G-I)[:,1]) == V1
    @test QR\(G-I)[:,1] == V1

    @test norm((V1-Ṽ1).coefficients) == 0
    @test norm((A*V1-F).coefficients) < 100eps()

    @test norm((F-Fun((G-I)[:,1])).coefficients) == 0
    @test Fun(V1) == V1

    Ṽ = QR\(G-I)
    V  = (I+(I-G)*C)\(G-I)

    @test map(f->f(exp(0.1im)),Array(G-I)) ≈ (G-I)(exp(0.1im))

    @test (G-I)[:,1]==Fun((G-I)[:,1],rangespace(QR))
    @test (G-I)[:,1]==Fun(vec((G-I)[:,1]),rangespace(QR))

    @test norm((V-Ṽ).coefficients) == 0

    @test norm((V1-V[:,1]).coefficients) == 0

    V2  = A\(G-I)[:,2]
    @test norm((V2-V[:,2]).coefficients) == 0

    @test norm((A*V[:,1]-(G[:,1]-[1,0])).coefficients) < 100eps()

    z=exp(0.1im)
    @test V(z)+(I-G(z))*cauchy(V,(z)⁻) ≈ G(z)-I

    @test cauchy(V[1,1],exp(0.1im)⁻) ≈ (C*V[1,1])(exp(0.1im))
    @test cauchy(V[2,1],exp(0.1im)⁻) ≈ (C*V[2,1])(exp(0.1im))
    @test cauchy(V[:,1],exp(0.1im)⁻) ≈ (C*V[:,1])(exp(0.1im))
    @test cauchy(V,exp(0.1im)⁻) ≈ (C*V)(exp(0.1im))

    Φmi = I+C*V
    Φp = V+Φmi

    @test Φmi(z) ≈ (I+cauchy(V,(z)⁻))
    @test Φp(z) ≈ (I+cauchy(V,(z)⁺))
    @test Φp(z)*inv(Φmi(z)) ≈ G(z)

    Φm=inv.(Φmi)
    @test Φm(z) ≈ inv(Φmi(z))

    T=ToeplitzOperator(G)

    L  = ToeplitzOperator(Φm)
    U  = ToeplitzOperator(Φp)

    @test norm((T-U*L)[1:10,1:10]) < 100eps()  # check the accuracy
end


############################################################################################
#Markov Kernel
markov_MV_data = [randn(_rng, 2) for _ in 1:N]
param_HMM_MV = (;
    μ = Param([[-2., -2.], [0., 0.], [2., 2.]],
        [MvNormal( [i, i], [1., 1.] ) for i in -1.0:1:1.]),
    σ = Param([[ 2.2 0.2 ; 0.2 2.2], [ 1.1 0.1 ; 0.1 1.1], [ 2.2 0.2 ; 0.2 2.2] ],
        [InverseWishart(10, [0.8 0.5 ; 0.5 .8] ) for i in 1:3]),
    p = Param([[.6, .2, .2],  [.1, .1, .8], [.1, .1, .8]], [Dirichlet(3,3) for i in 1:3]),
    latent = Param(markov_latent, Fixed()),
)
struct HMM_MV <: ModelName end
hmm_MV = ModelWrapper(HMM_MV(), param_HMM_MV)
markov_MV_objective = Objective(hmm_MV, markov_MV_data, :latent)

################################################################################
function get_dynamics(model::ModelWrapper{<:HMM_MV}, θ)
    @unpack μ, σ, p = θ
    dynamicsᵉ = [ MvNormal(μ[iter], σ[iter]) for iter in eachindex(μ) ]
    dynamicsˢ = [ Categorical( p[iter] ) for iter in eachindex(μ) ]
    return dynamicsᵉ, dynamicsˢ
end

function BaytesFilters.dynamics(objective::Objective{<:ModelWrapper{<:HMM_MV}})
    @unpack model, data = objective
    dynamicsᵉ, dynamicsˢ = get_dynamics(model, model.val)

    initialˢ = Categorical(fill(1 / length(dynamicsᵉ), length(dynamicsᵉ)))
    transition(particles, iter) = dynamicsˢ[particles[iter - 1]]
    evidence(particles, iter) = dynamicsᵉ[particles[iter]]

    return Markov(initialˢ, transition, evidence)
end
#dynamics(markov_MV_objective)

############################################################################################
"Forward Filter HMM"
function filter_forward(objective::Objective{<:ModelWrapper{HMM_MV}})
    @unpack model, data = objective
## Map Parameter to observation and state probabilities
    @unpack p = model.val
    dynamicsᵉ, _ = get_dynamics(model, model.val)
#!NOTE: Working straight with parameter instead of distributions here as easier to implement
    dynamicsˢ           = transpose( reduce(hcat, p ) ) #[ Categorical(p[iter]) for iter in eachindex(μ) ]
    initialˢ            = get_stationary( dynamicsˢ )
    structure = BaytesCore.ByRows()
## Assign Log likelihood
    ℓℒᵈᵃᵗᵃ = zeros(Float64, size(data,1), size(dynamicsˢ, 1) )
    Base.Threads.@threads for state in Base.OneTo( size(ℓℒᵈᵃᵗᵃ, 2) )
    for iter in Base.OneTo( size(ℓℒᵈᵃᵗᵃ, 1) )
            ℓℒᵈᵃᵗᵃ[iter, state] += logpdf( dynamicsᵉ[state], grab(data, iter, structure) )
        end
    end
## Initialize
    Nstates = size(dynamicsˢ, 1)
    α = zeros( size(ℓℒᵈᵃᵗᵃ) ) # α represents P( sₜ | e₁:ₜ ) for each t, which is numerically more stable than classical forward probabilities p(sₜ, e₁:ₜ)
    c = vec_maximum( @view(ℓℒᵈᵃᵗᵃ[1,:]) ) # Used for stable logsum() calculations for incremental log likelihood addition, i.e.: log( exp(x) + exp(y) ) == x + log(1 + exp(y -x ) ) for y >= x.
    ℓℒ = 0.0 #Log likelihood container
## Calculate initial probabilities
    for stateₜ in Base.OneTo(Nstates)
         α[1,stateₜ] += initialˢ[stateₜ] * exp(ℓℒᵈᵃᵗᵃ[1,stateₜ]-c) #Calculate initial p(s₁, e₁) ∝ P(s₁) * p(e₁ | s₁)
    end
## Normalize P( s₁, e₁ ) and return normalizing constant P(e₁), which can be used to calculate ℓℒ = P(e₁:ₜ) = p(e₁) ∏_k p(eₖ | e₁:ₖ₋₁)
    norm = sum( @view(α[1,:]) )
    α[1,:] ./= norm
    ℓℒ += log(norm) + c# log(norm) = p(eₜ | e₁:ₜ₋₁), c is constant that is added back after removing on top for numerical stability
## Loop through sequence
    for t = 2:size(α, 1)
        c = vec_maximum( @view(ℓℒᵈᵃᵗᵃ[t,:]) ) # Used for stable logsum() calculations, i.e.: log( exp(x) + exp(y) ) == x + log(1 + exp(y -x ) ) for y >= x.
        ## Calculate ∑_sₜ₋₁ P( sₜ | sₜ₋₁) * P(sₜ₋₁ | e₁:ₜ₋₁) - we sum over sₜ₋₁ - states are per row
        for stateₜ in Base.OneTo(Nstates)
            for stateₜ₋₁ in Base.OneTo(Nstates)
                α[t,stateₜ] += dynamicsˢ[stateₜ₋₁, stateₜ] * α[t-1,stateₜ₋₁] # * for both log and non-log version as inside log( sum(probability(...)))
            end
            ## Then multiply with ℒᵈᵃᵗᵃ corrector P( eₜ | sₜ )
            α[t,stateₜ] *= exp(ℓℒᵈᵃᵗᵃ[t,stateₜ]-c) # - c for higher numerical stability, will be added back to likelihood increment below
        end
        ## Normalize α and obtain P( eₜ | e₁:ₜ₋₁)
        norm = sum( @view(α[t,:]) )
        α[t,:] ./= norm
        ## Add normalizing constant p(eₜ | e₁:ₜ₋₁) to likelihood term
        ℓℒ += log(norm)+c
    end
    return (α, ℓℒ)
end
filter_forward(markov_MV_objective)

############################################################################################
function ModelWrappers.simulate(rng::Random.AbstractRNG, model::ModelWrapper{HMM_MV}; Nsamples = 1000)
    dynamicsᵉ, dynamicsˢ = get_dynamics(model, model.val)
    latentⁿᵉʷ = Vector{eltype(model.val.latent)}(undef, Nsamples)
    observedⁿᵉʷ = [zeros(2) for _ in 1:Nsamples]

    stateₜ = rand(rng, dynamicsˢ[ rand(1:length(dynamicsˢ) ) ] )
    latentⁿᵉʷ[1] = stateₜ
    observedⁿᵉʷ[1] = rand(rng, dynamicsᵉ[stateₜ])

    for iter in 2:size(observedⁿᵉʷ,1)
            stateₜ = rand(rng, dynamicsˢ[ stateₜ ]) #stateₜ for t-1 overwritten
            latentⁿᵉʷ[iter] = stateₜ
            observedⁿᵉʷ[iter] = rand(rng, dynamicsᵉ[stateₜ])
    end
    return observedⁿᵉʷ, latentⁿᵉʷ
end


# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

include("common_io.jl")
using LinearAlgebra

struct Wishart
  gamma::Float64
  m::Int
end

function logsumexp(x)
  mx = maximum(x)
  log.(sum(exp.(x .- mx))) .+ mx
end

sumsq(v) = sum(abs2, v)

function ltri_unpack(D, LT)
  d=length(D)
  make_row(r::Int, L) = hcat(reshape([ L[i] for i=1:r-1 ],1,r-1), D[r], zeros(1,d-r))
  row_start(r::Int) = div((r-1)*(r-2),2)
  inds(r) = row_start(r) .+ (1:r-1)
  vcat([ make_row(r, LT[inds(r)]) for r=1:d ]...)
end

function get_Q(d,icf)
  ltri_unpack(exp.(icf[1:d]),icf[d+1:end])
end

# Gradient helpers
function pack(alphas,means,icf)
  [alphas[:];means[:];icf[:]]
end

function unpack(d,k,packed)
  alphas = reshape(packed[1:k],1,k)
  off = k
  means = reshape(packed[(1:d*k) .+ off],d,k)
  icf_sz = div(d*(d + 1),2)
  off += d*k
  icf = reshape(packed[off+1:end],icf_sz,k)
  (alphas,means,icf)
end

function log_gamma_distrib(a, p)
  out = 0.25 * p * (p - 1) * 1.1447298858494002 #convert(Float64, log(pi))
	for j in 1:p
    out += (logabsgamma(a + 0.5*(1 - j)))[1]
  end
	out
end

function log_wishart_prior(wishart::Wishart, sum_qs, Qs, icf)
  p = size(Qs[1],1)
  n = p + wishart.m + 1
  C = n*p*(log(wishart.gamma) - 0.5*log(2)) - log_gamma_distrib(0.5*n, p)

  frobenius = 0.
  for Q in Qs
    frobenius += sum(abs2,diag(Q))
  end
  frobenius += sum(abs2,icf[d+1:end,:])
	0.5*wishart.gamma^2 * frobenius - wishart.m*sum(sum_qs) - k*C
end

# IO
function read_gmm_instance(fn,replicate_point)
  fid = open(fn)
  lines = readlines(fid)
  close(fid)
  line=split(lines[1]," ")
  d = parse(Int,line[1])
  k = parse(Int,line[2])
  n = parse(Int,line[3])
  icf_sz = div(d*(d + 1),2)
  off = 1

  alphas = zeros(Float64,1,k)
  for i in 1:k
    alphas[i] = parse(Float64,lines[i+off])
  end
  off += k

  means = zeros(Float64,d,k)
  for ik in 1:k
    line=split(lines[ik+off]," ")
    for id in 1:d
      means[id,ik] = parse(Float64,line[id])
    end
  end
  off += k

  icf = zeros(Float64,icf_sz,k)
  for ik in 1:k
    line=split(lines[ik+off]," ")
    for i in 1:icf_sz
      icf[i,ik] = parse(Float64,line[i])
    end
  end
  off += k

  if replicate_point
    x_ = zeros(Float64,d,1)
    line=split(lines[1+off]," ")
    for id in 1:d
      x_[id] = parse(Float64,line[id])
    end
    x = repeat(x_,1,n)
    off += 1
  else
    x = zeros(Float64,d,n)
    for ix in 1:n
      line=split(lines[ix+off]," ")
      for id in 1:d
        x[id,ix] = parse(Float64,line[id])
      end
    end
    off += n
  end
  line = split(lines[1+off]," ")
  wishart = Wishart(parse(Float64,line[1]),parse(Int,line[2]))
  (alphas,means,icf,x,wishart)
end


using Revise # reduce recompilation time
using Plots
# using Documenter
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D

using SetupDG

push!(LOAD_PATH, "./examples/EntropyStableEuler.jl/src")
using EntropyStableEuler
using EntropyStableEuler.Fluxes1D

function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,(rhou,),E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end
unorm(U) = sum(map((x->x.^2),U))
function pfun_nd(rho, rhoU, E)
    rhoUnorm2 = unorm(rhoU)./rho
    return @. (γ-1)*(E - .5*rhoUnorm2)
end

function primitive_to_conservative_hardcode(rho,U,p)
    rhoU = rho.*U
    Unorm = unorm(U)
    E = @. p/(γ-1) + .5*rho*Unorm
    return (rho,rhoU,E)
end

function euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
    rhoR,uR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))

    unorm = (@. uL*uR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. f4aux*uavg)

    return (FxS1,FxS2,FxS3)
end

function euler_fluxes(rhoL,uL,betaL,rhoR,uR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
                           rhoR,uR,betaR,rhologR,betalogR)
end

const TOL = 1e-16
"Approximation parameters"
N = 2 # The order of approximation
K = 50#800
T = 1.0#3.0

# Becker viscous shocktube
const γ = 1.4
const M_0 = 20.0
const mu = 0.001
const lambda = 2/3*mu
const Pr = 3/4
const cp = γ/(γ-1)
const cv = 1/(γ-1)
const kappa = mu*cp/Pr

const v_inf = 0.2
const rho_0 = 1.0
const v_0 = 1.0
const m_0 = rho_0*v_0
const v_1 = (γ-1+2/M_0^2)/(γ+1)
const v_01 = sqrt(v_0*v_1)

const uL = v_0+v_inf
const uR = v_1+v_inf
const rhoL = m_0/v_0
const rhoR = m_0/v_1
const eL = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_0^2)
const eR = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-v_1^2)
const pL = (γ-1)*rhoL*eL
const pR = (γ-1)*rhoR*eR
const EL = pL/(γ-1)+0.5*rhoL*uL^2
const ER = pR/(γ-1)+0.5*rhoR*uR^2

const Bl = -1.0
const Br = 1.5
gr(size=(300,300),ylims=(0,6.0),legend=false,markerstrokewidth=1,markersize=2)
plot()

"Mesh related variables"
VX = LinRange(Bl,Br,K+1)
EToV = transpose(reshape(sort([1:K; 2:K+1]),2,K))

"Initialize reference element"
r,_ = gauss_lobatto_quad(0,0,N)   # Reference nodes
VDM = vandermonde_1D(N,r)         # modal to nodal
Dr = grad_vandermonde_1D(N,r)/VDM # nodal differentiation
V1 = vandermonde_1D(1,r)/vandermonde_1D(1,[-1;1]) # nodal linear interpolation

Nq = N
rq,wq = gauss_lobatto_quad(0,0,Nq)
Vq = vandermonde_1D(N,rq)/VDM
M = Vq'*diagm(wq)*Vq
Mlump = zeros(size(M))
Mlump_inv = zeros(size(M))
for i = 1:Nq+1
    Mlump[i,i] = sum(M[i,:])
    Mlump_inv[i,i] = 1.0/Mlump[i,i]
end
rf = [-1.0;1.0]
nrJ = [-1.0;1.0]
Vf = vandermonde_1D(N,rf)/VDM
LIFT = Mlump_inv*Vf'

# operators
Qr = M*Dr
B = zeros(N+1,N+1)
B[1,1] = -1
B[end,end] = 1
L = Array(spdiagm(0=>-2*ones(N+1), 1=>ones(N), -1=>ones(N)))
L[1,1] = -1
L[end,end] = -1
psi = pinv(L)*-1/2*B*ones(N+1)
S0 = zeros(N+1,N+1)
for i = 1:N+1
    for j = 1:N+1
        if L[i,j] != 0
            S0[i,j] = psi[j] - psi[i]
        end
    end
end
Qr0 = S0+1/2*B

# Drop zeros
Qr = Matrix(droptol!(sparse(Qr),TOL))
Qr0 = Matrix(droptol!(sparse(Qr0),TOL))
B = Matrix(droptol!(sparse(B),TOL))
L = Matrix(droptol!(sparse(L),TOL))
S0 = Matrix(droptol!(sparse(S0),TOL))
S = Matrix(droptol!(sparse((Qr-Qr')/2),TOL))

"""High order mesh"""
x = V1*VX[transpose(EToV)]
xf = Vf*x
mapM = reshape(1:2*K,2,K)
mapP = copy(mapM)
mapP[1,2:end] .= mapM[2,1:end-1]
mapP[2,1:end-1] .= mapM[1,2:end]

# """Periodic"""
# mapP[1] = mapM[end]
# mapP[end] = mapP[1]

"""Geometric factors"""
J = repeat(transpose(diff(VX)/2),N+1,1)
nxJ = repeat([-1;1],1,K)
rxJ = 1.0

"""Geometric factors"""
J = repeat(transpose(diff(VX)/2),N+1,1)
nxJ = repeat([-1;1],1,K)
rxJ = 1.0

"""Initial condition"""
function bisection_solve_velocity(x,max_iter,tol)
    v_L = v_1
    v_R = v_0
    num_iter = 0

    L_k = kappa/m_0/cv
    f(v) = -x+2*L_k/(γ+1)*(v_0/(v_0-v_1)*log((v_0-v)/(v_0-v_01))-v_1/(v_0-v_1)*log((v-v_1)/(v_01-v_1)))

    v_new = (v_L+v_R)/2
    while num_iter < max_iter
        v_new = (v_L+v_R)/2

        if abs(f(v_new)) < tol
            return v_new
        elseif sign(f(v_L)) == sign(f(v_new))
            v_L = v_new
        else
            v_R = v_new
        end
        num_iter += 1
    end

    return v_new
end

const max_iter = 10000
const tol = 1e-14

function exact_sol_viscous_shocktube(x,t)
    u   = bisection_solve_velocity(x-v_inf*t,max_iter,tol)
    rho = m_0/u
    e   = 1/(2*γ)*((γ+1)/(γ-1)*v_01^2-u^2)
    return rho, rho*(v_inf+u), rho*(e+1/2*(v_inf+u)^2)
end

U = exact_sol_viscous_shocktube.(x,0.0)
U = ([x[1] for x in U], [x[2] for x in U], [x[3] for x in U])

function zhang_wavespd(rhoL,rhouL,EL,sigma2L,sigma3L,pL,rhoR,rhouR,ER,sigma2R,sigma3R,pR)
    uL   = rhouL/rhoL
    eL   = (EL - .5*rhoL*uL^2)/rhoL
    tauL = sigma2L
    qL   = uL*tauL-sigma3L
    uR   = rhouR/rhoR
    eR   = (ER - .5*rhoR*uR^2)/rhoR
    tauR = sigma2R
    qR   = uR*tauR-sigma3R
    wavespdL = abs(uL)+1/(2*rhoL^2*eL)*(sqrt(rhoL^2*qL^2+2*rhoL^2*eL*abs(tauL-pL)^2)+rhoL*abs(qL))
    wavespdR = abs(uR)+1/(2*rhoR^2*eR)*(sqrt(rhoR^2*qR^2+2*rhoR^2*eR*abs(tauR-pR)^2)+rhoR*abs(qR))
    return max(wavespdL,wavespdR)
end

function zhang_wavespd(rhoL,rhouL,EL,sigma2L,sigma3L,pL)
    uL   = rhouL/rhoL
    eL   = (EL - .5*rhoL*uL^2)/rhoL
    tauL = sigma2L
    qL   = uL*tauL-sigma3L
    wavespdL = abs(uL)+1/(2*rhoL^2*eL)*(sqrt(rhoL^2*qL^2+2*rhoL^2*eL*abs(tauL-pL)^2)+rhoL*abs(qL))
    return wavespdL
end

function limiting_param(U_low, P_ij)
    l = 1.0
    # Limit density
    if U_low[1] + P_ij[1] < -TOL
        l = min(abs(U_low[1])/(abs(P_ij[1])+1e-14), 1.0)
    end

    # limiting internal energy (via quadratic function)
    a = P_ij[1]*P_ij[3]-1.0/2.0*P_ij[2]^2
    b = U_low[3]*P_ij[1]+U_low[1]*P_ij[3]-U_low[2]*P_ij[2]
    c = U_low[3]*U_low[1]-1.0/2.0*U_low[2]^2

    l_eps_ij = 1.0
    if b^2-4*a*c >= 0
        r1 = (-b+sqrt(b^2-4*a*c))/(2*a)
        r2 = (-b-sqrt(b^2-4*a*c))/(2*a)
        if r1 > TOL && r2 > TOL
            l_eps_ij = min(r1,r2)
        elseif r1 > TOL && r2 < -TOL
            l_eps_ij = r1
        elseif r2 > TOL && r1 < -TOL
            l_eps_ij = r2
        end
    end

    l = min(l,l_eps_ij)
    return l
end

function flux_lowIDP(U_i,U_j,f_i,f_j,c_ij,wavespd)
    return c_ij*(f_i+f_j)-abs(c_ij)*wavespd*(U_j-U_i)
end

function flux_ES(rho_i,u_i,beta_i,rho_j,u_j,beta_j,c_ij)
    return 2*c_ij.*euler_fluxes(rho_i,u_i,beta_i,rho_j,u_j,beta_j)
end

function flux_viscous(sigma_i,sigma_j,c_ij)
    return c_ij*(sigma_i+sigma_j)
end

function rhs_viscous(U,K,N,Mlump_inv,S)
    J = (Br-Bl)/K/2 # assume uniform interval
    Nc = 3
    rhsU  = [zeros(N+1,K) for i = 1:Nc]
    theta = [zeros(N+1,K) for i = 1:Nc]
    sigma = [zeros(N+1,K) for i = 1:Nc]

    p = pfun_nd.(U[1],U[2],U[3])

    VU = v_ufun(U...)
    VUf = (x->Vf*x).(VU)
    VUP = (x->x[mapP]).(VUf)
    VL = v_ufun(rhoL,rhoL*uL,EL)
    VUP[1][1] = VL[1]
    VUP[2][1] = VL[2]
    VUP[3][1] = VL[3]
    VUP[1][end] = VUf[1][end]
    VUP[2][end] = VUf[2][end]
    VUP[3][end] = VUf[3][end]

    surfx(uP,uf) = LIFT*(@. .5*(uP-uf)*nxJ)
    # Strong form, dv/dx
    VUx = (x->Dr*x).(VU)
    VUx = VUx .+ surfx.(VUP,VUf)
    VUx = VUx./J

    # σ = K dv/dx
    for k = 1:K
        for i = 1:N+1
            Kx = zeros(3,3)
            v1 = VU[1][i,k]
            v2 = VU[2][i,k]
            v4 = VU[3][i,k]
            Kx[2,2] = -(2*mu-lambda)/v4
            Kx[2,3] = (2*mu-lambda)*v2/v4^2
            Kx[3,2] = Kx[2,3]
            Kx[3,3] = -(2*mu-lambda)*v2^2/v4^3+kappa/cv/v4^2

            sigma[2][i,k] += Kx[2,2]*VUx[2][i,k] + Kx[2,3]*VUx[3][i,k]
            sigma[3][i,k] += Kx[3,2]*VUx[2][i,k] + Kx[3,3]*VUx[3][i,k]
        end
    end

    sxf = (x->Vf*x).(sigma)
    sxP = (x->x[mapP]).(sxf)
    sxP[1][1] = sxf[1][1]
    sxP[2][1] = sxf[2][1]
    sxP[3][1] = sxf[3][1]
    sxP[1][end] = sxf[1][end]
    sxP[2][end] = sxf[2][end]
    sxP[3][end] = sxf[3][end]
    # strong form, dσ/dx
    penalization(uP,uf) = LIFT*(@. -.5*(uP-uf))
    sigmax = (x->Dr*x).(sigma)
    sigmax = sigmax .+ surfx.(sxP,sxf) #.+ penalization.(VUP,VUf)
    sigmax = sigmax./J

    return sigmax, sigma
end

function rhs_IDP(U,K,N,Mlump_inv,S0,S,wq)
    p = pfun_nd.(U[1],U[2],U[3])
    flux = zero.(U)
    @. flux[1] = U[2]
    @. flux[2] = U[2]^2/U[1]+p
    @. flux[3] = U[3]*U[2]/U[1]+p*U[2]/U[1]

    Ub = zero.(U)
    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[1]/(2*p)

    J = (Br-Bl)/K/2 # assume uniform interval

    # Low order and high order algebraic fluxes
    F_low       = [zeros(N+1,N+1),zeros(N+1,N+1),zeros(N+1,N+1)]
    F_high      = [zeros(N+1,N+1),zeros(N+1,N+1),zeros(N+1,N+1)]
    F_low_P     = [zeros(2),zeros(2),zeros(2)] # 1: left boundary, 2: right boundary
    F_high_P    = [zeros(2),zeros(2),zeros(2)]
    L           = zeros(N+1,N+1) # Array of limiting params
    rhsU        = [zeros(N+1,K),zeros(N+1,K),zeros(N+1,K)]

    # TODO: redundant!
    _,sigma = rhs_viscous(U,K,N,Mlump_inv,S)
    wavespd_arr = zeros(N+1,K)
    beta_arr = zeros(N+1,K)
    for k = 1:K
        for i = 1:N+1
            wavespd_arr[i,k] = wavespeed_1D(U[1][i,k],U[2][i,k],U[3][i,k])
            beta_arr[i,k] = zhang_wavespd(U[1][i,k],U[2][i,k],U[3][i,k],sigma[2][i,k],sigma[3][i,k],p[i,k])
        end
    end

    L_plot = zeros(N+1,K)

    #dt = Inf
    # TODO: hardcoded
    dt = 1e-4
    d_ii_arr = zeros(N+1,K)
    for k = 1:K
        L = zeros(N+1,N+1)
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    wavespd = max(wavespd_arr[i,k],wavespd_arr[j,k],beta_arr[i,k],beta_arr[j,k])
                    fluxS = flux_ES(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],S[i,j])
                    for c = 1:3
                        F_low[c][i,j]  = -flux_viscous(sigma[c][i,k],sigma[c][j,k],S[i,j])+flux_lowIDP(U[c][i,k],U[c][j,k],flux[c][i,k],flux[c][j,k],S0[i,j],wavespd)
                        F_high[c][i,j] = -flux_viscous(sigma[c][i,k],sigma[c][j,k],S[i,j])+fluxS[c]
                    end
                    d_ii_arr[i,k] += wavespd
                end
            end
        end

        EL = pL/(γ-1)+0.5*rhoL*uL^2
        ER = pR/(γ-1)+0.5*rhoR*uR^2
        U_left   = (k == 1) ? [rhoL; rhoL*uL; EL]                 : [U[1][end,k-1]; U[2][end,k-1]; U[3][end,k-1]]
        Ub_left  = (k == 1) ? [rhoL; uL; rhoL/(2*pL)]             : [Ub[1][end,k-1]; Ub[2][end,k-1]; Ub[3][end,k-1]]
        f_left   = (k == 1) ? [rhoL*uL; rhoL*uL^2+pL; uL*(pL+EL)] : [flux[1][end,k-1]; flux[2][end,k-1]; flux[3][end,k-1]]
        U_right  = (k == K) ? [rhoR; rhoR*uR; ER]                 : [U[1][1,k+1]; U[2][1,k+1]; U[3][1,k+1]]
        Ub_right = (k == K) ? [rhoR; uR; rhoR/(2*pR)]             : [Ub[1][1,k+1]; Ub[2][1,k+1]; Ub[3][1,k+1]]
        f_right  = (k == K) ? [rhoR*uR; rhoR*uR^2+pR; uR*(pR+ER)] : [flux[1][1,k+1]; flux[2][1,k+1]; flux[3][1,k+1]]
        p_left   = (k == 1) ? pL : p[end,k-1]
        p_right  = (k == K) ? pR : p[1,k+1]
        # Assume velocity doesn't change at boundary, isothermal
        TL = pL/rhoL/(γ-1)/cv
        TR = pR/rhoR/(γ-1)/cv
        sigma_left  = (k == 1) ? [sigma[1][1,1];sigma[2][1,1];sigma[3][1,1]]       : [sigma[1][end,k-1]; sigma[2][end,k-1]; sigma[3][end,k-1]]
        sigma_right = (k == K) ? [sigma[1][end,k];sigma[2][end,k];sigma[3][end,k]] : [sigma[1][1,k+1]; sigma[2][1,k+1]; sigma[3][1,k+1]]
        wavespd_l = max(wavespd_arr[1,k],wavespeed_1D(U_left[1],U_left[2],U_left[3]),beta_arr[1,k],zhang_wavespd(U_left[1],U_left[2],U_left[3],sigma_left[2],sigma_left[3],p_left))
        wavespd_r = max(wavespd_arr[end,k],wavespeed_1D(U_right[1],U_right[2],U_right[3]),beta_arr[end,k],zhang_wavespd(U_right[1],U_right[2],U_right[3],sigma_right[2],sigma_right[3],p_right))

        d_ii_arr[1,k] += wavespd_l
        d_ii_arr[end,k] += wavespd_r

        for c = 1:3
            F_low_P[c][1] = -flux_viscous(sigma[c][1,k],sigma_left[c],-1/2)+flux_lowIDP(U[c][1,k],U_left[c],flux[c][1,k],f_left[c],-0.5,wavespd_l)
            F_low_P[c][2] = -flux_viscous(sigma[c][end,k],sigma_right[c],1/2)+flux_lowIDP(U[c][end,k],U_right[c],flux[c][end,k],f_right[c],0.5,wavespd_r)

            if k == 1
                F_low_P[c][1] = -flux_viscous(sigma[c][1,k],sigma_left[c],-1/2)-0.5*(flux[c][1,k]+f_left[c])
            end

            if k == K
                F_low_P[c][2] = -flux_viscous(sigma[c][end,k],sigma_right[c],1/2)+0.5*(flux[c][end,k]+f_right[c])
            end

            F_high_P[c][1] = F_low_P[c][1]
            F_high_P[c][2] = F_low_P[c][2]
        end

        P_ij = zeros(3,1)
        # TODO: redundant
        U_low =  [zeros(N+1,1),zeros(N+1,1),zeros(N+1,1)]
        for c = 1:3
            U_low[c] .= sum(-F_low[c],dims=2)
            U_low[c][1] -= F_low_P[c][1]
            U_low[c][end] -= F_low_P[c][2]
            U_low[c] .= U[c][:,k]+dt/J*Mlump_inv*U_low[c]
        end

        for i = 1:N+1
            lambda_j = 1/N
            m_i = J*wq[i]
            for j = 1:N+1
                if i != j
                    for c = 1:3
                        P_ij[c] = dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                    end
                    L[i,j] = limiting_param([U_low[1][i]; U_low[2][i]; U_low[3][i]],P_ij)
                end
            end
        end

        # Symmetrize limiting parameters
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    l_ij = min(L[i,j],L[j,i])
                    L[i,j] = l_ij
                    L[j,i] = l_ij
                end
            end
        end

        # elementwise limiting
        l = 1.0
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    l = min(l,L[i,j])
                end
            end
        end
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    L[i,j] = l
                    L[j,i] = l
                end
            end
        end

        L = ones(size(L))

        for c = 1:3
            # With limiting
            rhsU[c][:,k] = sum((L.-1).*F_low[c] - L.*F_high[c],dims=2)

            rhsU[c][1,k] += -F_high_P[c][1]
            rhsU[c][N+1,k] += -F_high_P[c][2]

            rhsU[c][:,k] .= 1/J*Mlump_inv*rhsU[c][:,k]
        end

        for i = 1:N+1
            L_plot[i,k] = l
        end
    end

    dt = min(dt,minimum(J*wq./d_ii_arr./2.0))

    return rhsU,dt,L_plot
end

function rhs_IDP(U,K,N,Mlump_inv,S0,S,wq,dt)
    p = pfun_nd.(U[1],U[2],U[3])
    flux = zero.(U)
    @. flux[1] = U[2]
    @. flux[2] = U[2]^2/U[1]+p
    @. flux[3] = U[3]*U[2]/U[1]+p*U[2]/U[1]

    Ub = zero.(U)
    @. Ub[1] = U[1]
    @. Ub[2] = U[2]/U[1]
    @. Ub[3] = U[1]/(2*p)

    J = (Br-Bl)/K/2 # assume uniform interval

    # Low order and high order algebraic fluxes
    F_low       = [zeros(N+1,N+1),zeros(N+1,N+1),zeros(N+1,N+1)]
    F_high      = [zeros(N+1,N+1),zeros(N+1,N+1),zeros(N+1,N+1)]
    F_low_P     = [zeros(2),zeros(2),zeros(2)] # 1: left boundary, 2: right boundary
    F_high_P    = [zeros(2),zeros(2),zeros(2)]
    L           = zeros(N+1,N+1) # Array of limiting params
    rhsU        = [zeros(N+1,K),zeros(N+1,K),zeros(N+1,K)]

    # TODO: redundant!
    _,sigma = rhs_viscous(U,K,N,Mlump_inv,S)
    wavespd_arr = zeros(N+1,K)
    beta_arr = zeros(N+1,K)
    for k = 1:K
        for i = 1:N+1
            wavespd_arr[i,k] = wavespeed_1D(U[1][i,k],U[2][i,k],U[3][i,k])
            beta_arr[i,k] = zhang_wavespd(U[1][i,k],U[2][i,k],U[3][i,k],sigma[2][i,k],sigma[3][i,k],p[i,k])
        end
    end

    L_plot = zeros(N+1,K)

    for k = 1:K
        L = zeros(N+1,N+1)
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    wavespd = max(wavespd_arr[i,k],wavespd_arr[j,k],beta_arr[i,k],beta_arr[j,k])
                    fluxS = flux_ES(Ub[1][i,k],Ub[2][i,k],Ub[3][i,k],Ub[1][j,k],Ub[2][j,k],Ub[3][j,k],S[i,j])
                    for c = 1:3
                        F_low[c][i,j]  = -flux_viscous(sigma[c][i,k],sigma[c][j,k],S[i,j])+flux_lowIDP(U[c][i,k],U[c][j,k],flux[c][i,k],flux[c][j,k],S0[i,j],wavespd)
                        F_high[c][i,j] = -flux_viscous(sigma[c][i,k],sigma[c][j,k],S[i,j])+fluxS[c]
                    end
                end
            end
        end

        U_left   = (k == 1) ? [rhoL; rhoL*uL; EL]                 : [U[1][end,k-1]; U[2][end,k-1]; U[3][end,k-1]]
        Ub_left  = (k == 1) ? [rhoL; uL; rhoL/(2*pL)]             : [Ub[1][end,k-1]; Ub[2][end,k-1]; Ub[3][end,k-1]]
        f_left   = (k == 1) ? [rhoL*uL; rhoL*uL^2+pL; uL*(pL+EL)] : [flux[1][end,k-1]; flux[2][end,k-1]; flux[3][end,k-1]]
        U_right  = (k == K) ? [rhoR; rhoR*uR; ER]                 : [U[1][1,k+1]; U[2][1,k+1]; U[3][1,k+1]]
        Ub_right = (k == K) ? [rhoR; uR; rhoR/(2*pR)]             : [Ub[1][1,k+1]; Ub[2][1,k+1]; Ub[3][1,k+1]]
        f_right  = (k == K) ? [rhoR*uR; rhoR*uR^2+pR; uR*(pR+ER)] : [flux[1][1,k+1]; flux[2][1,k+1]; flux[3][1,k+1]]
        p_left   = (k == 1) ? pL : p[end,k-1]
        p_right  = (k == K) ? pR : p[1,k+1]
        # Assume velocity doesn't change at boundary, isothermal
        TL = pL/rhoL/(γ-1)/cv
        TR = pR/rhoR/(γ-1)/cv
        sigma_left  = (k == 1) ? [sigma[1][1,1];sigma[2][1,1];sigma[3][1,1]]       : [sigma[1][end,k-1]; sigma[2][end,k-1]; sigma[3][end,k-1]]
        sigma_right = (k == K) ? [sigma[1][end,k];sigma[2][end,k];sigma[3][end,k]] : [sigma[1][1,k+1]; sigma[2][1,k+1]; sigma[3][1,k+1]]
        wavespd_l = max(wavespd_arr[1,k],wavespeed_1D(U_left[1],U_left[2],U_left[3]),beta_arr[1,k],zhang_wavespd(U_left[1],U_left[2],U_left[3],sigma_left[2],sigma_left[3],p_left))
        wavespd_r = max(wavespd_arr[end,k],wavespeed_1D(U_right[1],U_right[2],U_right[3]),beta_arr[end,k],zhang_wavespd(U_right[1],U_right[2],U_right[3],sigma_right[2],sigma_right[3],p_right))

        for c = 1:3
            F_low_P[c][1] = -flux_viscous(sigma[c][1,k],sigma_left[c],-1/2)+flux_lowIDP(U[c][1,k],U_left[c],flux[c][1,k],f_left[c],-0.5,wavespd_l)
            F_low_P[c][2] = -flux_viscous(sigma[c][end,k],sigma_right[c],1/2)+flux_lowIDP(U[c][end,k],U_right[c],flux[c][end,k],f_right[c],0.5,wavespd_r)

            F_high_P[c][1] = F_low_P[c][1]
            F_high_P[c][2] = F_low_P[c][2]
        end

        P_ij = zeros(3,1)
        # TODO: redundant
        U_low =  [zeros(N+1,1),zeros(N+1,1),zeros(N+1,1)]
        for c = 1:3
            U_low[c] .= sum(-F_low[c],dims=2)
            U_low[c][1] -= F_low_P[c][1]
            U_low[c][end] -= F_low_P[c][2]
            U_low[c] .= U[c][:,k]+dt/J*Mlump_inv*U_low[c]
        end

        for i = 1:N+1
            lambda_j = 1/N
            m_i = J*wq[i]
            for j = 1:N+1
                if i != j
                    for c = 1:3
                        P_ij[c] = dt/(m_i*lambda_j)*(F_low[c][i,j]-F_high[c][i,j])
                    end
                    L[i,j] = limiting_param([U_low[1][i]; U_low[2][i]; U_low[3][i]],P_ij)
                end
            end
        end

        # Symmetrize limiting parameters
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    l_ij = min(L[i,j],L[j,i])
                    L[i,j] = l_ij
                    L[j,i] = l_ij
                end
            end
        end

        # elementwise limiting
        l = 1.0
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    l = min(l,L[i,j])
                end
            end
        end
        for i = 1:N+1
            for j = 1:N+1
                if i != j
                    L[i,j] = l
                    L[j,i] = l
                end
            end
        end

        # L = ones(size(L))

        for c = 1:3
            # With limiting
            rhsU[c][:,k] = sum((L.-1).*F_low[c] - L.*F_high[c],dims=2)

            rhsU[c][1,k] += -F_high_P[c][1]
            rhsU[c][N+1,k] += -F_high_P[c][2]

            rhsU[c][:,k] .= 1/J*Mlump_inv*rhsU[c][:,k]
        end

        for i = 1:N+1
            L_plot[i,k] = l
        end
    end

    return rhsU,L_plot
end

# Time stepping
"Time integration"
t = 0.0
U = collect(U)
resU = [zeros(size(x)),zeros(size(x)),zeros(size(x))]
resW = [zeros(size(x)),zeros(size(x)),zeros(size(x))]
resZ = [zeros(size(x)),zeros(size(x)),zeros(size(x))]

Vp = vandermonde_1D(N,LinRange(-1,1,10))/VDM

anim = Animation()
const ptL = Bl+(Br-Bl)/K/(N+1)/2
const ptR = Br-(Br-Bl)/K/(N+1)/2
const hplot = (Br-Bl)/K/(N+1)
i = 1

while t < T

    # SSPRK(3,3)
    # rhsU,dt,_ = rhs_IDP(U,K,N,Mlump_inv,S0,S,wq)
    # dt = min(dt,T-t)
    dt = min(1e-5,T)
    rhsU,_ = rhs_IDP(U,K,N,Mlump_inv,S0,S,wq,dt)
    @. resW = U + dt*rhsU
    rhsU,_ = rhs_IDP(resW,K,N,Mlump_inv,S0,S,wq,dt)
    @. resZ = resW+dt*rhsU
    @. resW = 3/4*U+1/4*resZ
    rhsU,L_plot = rhs_IDP(resW,K,N,Mlump_inv,S0,S,wq,dt)
    @. resZ = resW+dt*rhsU
    @. U = 1/3*U+2/3*resZ

    global t = t + dt
    global i = i + 1
    println("Current time $t with time step size $dt, and final time $T")
    if mod(i,1000) == 1
        plot(x,U[1])
        # for k = 1:K
        #     plot!(ptL+(k-1)*hplot*(N+1):hplot:ptL+k*hplot*(N+1), 1 .-L_plot[:,k],st=:bar,alpha=0.2)
        # end
        frame(anim)
    end
end

plot(x,U[1])
gif(anim,"~/Desktop/tmp.gif",fps=15)

exact_U = @. exact_sol_viscous_shocktube.(x,T)
exact_rho = [x[1] for x in exact_U]
exact_u = [x[2] for x in exact_U]
exact_u = exact_u./exact_rho
exact_E = [x[3] for x in exact_U]

rho = U[1]
u = U[2]./U[1]
E = U[3]
p = pfun_nd.(U[1],U[2],U[3])
J = (Br-Bl)/K/2

Linferr = maximum(abs.(exact_rho-rho))/maximum(abs.(exact_rho)) +
          maximum(abs.(exact_u-u))/maximum(abs.(exact_u)) +
          maximum(abs.(exact_E-E))/maximum(abs.(exact_E))

L1err = sum(J*wq.*abs.(exact_rho-rho))/sum(J*wq.*abs.(rho)) +
        sum(J*wq.*abs.(exact_u-u))/sum(J*wq.*abs.(u)) +
        sum(J*wq.*abs.(exact_E-E))/sum(J*wq.*abs.(E))
println("N = $N, K = $K")
println("L1 error is $L1err")
println("Linf error is $Linferr")


##################################################
# Uncertain populations
################################################
x = UncertainValue([1, 2, 3], rand(3))
y = UncertainValue([1, UncertainValue(Normal, 0, 1), 3], rand(3))

@test x isa UncertainScalarPopulation{T1, T2} where {T1 <: Number, T2}
@test y isa UncertainScalarPopulation{T1, T2} where {T1 <: AbstractUncertainValue, T2}

##################################################
# Uncertain theoretical distributions
################################################

@test UncertainValue(Normal(1, 2)) isa UncertainScalarNormallyDistributed
@test UncertainValue(Normal(1, 2)) isa UncertainScalarNormallyDistributed
@test UncertainValue(Truncated(Normal(1, 2), 0.2, 0.8)) isa ConstrainedUncertainScalarValueTwoParameter

# Uncertain normally distributed values
@test UncertainValue(Normal, 1, 0.2, nσ = 2, trunc_lower = -5) isa UncertainScalarNormallyDistributed
@test UncertainValue(Normal, -3, 0.2) isa UncertainScalarNormallyDistributed
@test UncertainValue(Normal, 0, 2) isa UncertainScalarNormallyDistributed
@test UncertainValue(Normal, -1, 0.1) isa UncertainScalarNormallyDistributed
@test UncertainValue(Normal, -1, 0.1, nσ = 2, trunc_upper = 4) isa UncertainScalarNormallyDistributed
@test UncertainValue(Normal, 5.0, 0.2) isa UncertainScalarNormallyDistributed

# Uncertain uniformly distributed values
@test UncertainValue(Uniform, 1, 7) isa UncertainScalarUniformlyDistributed
@test UncertainValue(Uniform, -1, 7) isa UncertainScalarUniformlyDistributed
@test UncertainValue(Uniform, -6, -2) isa UncertainScalarUniformlyDistributed

# Uncertain beta-distributed values
@test UncertainValue(Beta, 1, 7) isa UncertainScalarBetaDistributed
@test_broken UncertainValue(Beta, -1, 7) isa UncertainScalarBetaDistributed
@test_broken UncertainValue(Beta, -1, -2) isa UncertainScalarBetaDistributed
@test_broken UncertainValue(Beta, 4, -2) isa UncertainScalarBetaDistributed

# Uncertain beta-distributed values
@test UncertainValue(BetaBinomial, 10, 1, 7) isa UncertainScalarBetaBinomialDistributed
@test_broken UncertainValue(BetaBinomial, 10, -1, 7) isa UncertainScalarBetaBinomialDistributed
@test_broken UncertainValue(BetaBinomial, 10, -1, -2) isa UncertainScalarBetaBinomialDistributed
@test_broken UncertainValue(BetaBinomial, 1, 2) isa UncertainScalarBetaBinomialDistributed

# Uncertain beta prime-distributed values
@test UncertainValue(BetaPrime, 1, 7) isa UncertainScalarBetaPrimeDistributed
@test_broken UncertainValue(BetaPrime, -1, 7) isa UncertainScalarBetaPrimeDistributed
@test_broken UncertainValue(BetaPrime, -1, -2) isa UncertainScalarBetaPrimeDistributed
@test_broken UncertainValue(BetaPrime, 2, -2) isa UncertainScalarBetaPrimeDistributed

# Uncertain gamma-distributed values
@test UncertainValue(Gamma, 1, 7) isa UncertainScalarGammaDistributed
@test UncertainValue(Gamma, 1, 7, trunc_upper = 4, trunc_lower = 1) isa UncertainScalarGammaDistributed
@test_broken UncertainValue(Gamma, -1, 7) isa UncertainScalarGammaDistributed
@test_broken UncertainValue(Gamma, -1, -2) isa UncertainScalarGammaDistributed
@test_broken UncertainValue(Gamma, -1, -2) isa UncertainScalarGammaDistributed

# Uncertain Fréchet-distributed values
@test UncertainValue(Frechet, 1, 7) isa UncertainScalarFrechetDistributed

# Uncertain Binomial-distributed values
@test UncertainValue(Binomial, 50, 0.4) isa UncertainScalarBinomialDistributed

################################################
# Uncertain values from empirical distributions
################################################
empirical_uniform = rand(Uniform(), 100)
empirical_normal = rand(Normal(), 100)
empirical_beta = rand(Beta(), 100)

@test UncertainValue(Uniform, empirical_uniform) isa UncertainScalarTheoreticalFit
@test UncertainValue(Normal, empirical_normal) isa UncertainScalarTheoreticalFit
@test UncertainValue(Beta, empirical_beta) isa UncertainScalarTheoreticalFit

##################################################
# Uncertain values from kernel density estimates
################################################

# Implicit constructor
@test UncertainValue(empirical_uniform) isa UncertainScalarKDE
@test UncertainValue(empirical_normal) isa UncertainScalarKDE
@test UncertainValue(empirical_beta) isa UncertainScalarKDE

# Explicit constructor
@test UncertainValue(UnivariateKDE, empirical_uniform) isa UncertainScalarKDE
@test UncertainValue(UnivariateKDE, empirical_normal) isa UncertainScalarKDE
@test UncertainValue(UnivariateKDE, empirical_beta) isa UncertainScalarKDE

# Empirical cumulative distribution function
d = Normal()
some_sample = rand(d, 1000)
uv = UncertainValue(some_sample)
uv_ecdf = UncertainData.UncertainValues.ecdf(uv)

tol = 1e-7
@test all(uv_ecdf .>= 0.0 - tol)
@test all(uv_ecdf .<= 1.0 + tol)

# Quantiles (empirical and true qantiles should be close for large samples)
large_sample = rand(d, Int(1e6))
uv = UncertainValue(UnivariateKDE, large_sample)
@test abs(quantile(uv, 0.8) - quantile(d, 0.8)) < 1e-2


export StraightLineHomotopy

"""
    StraightLineHomotopy(G::System, F::System; gamma = 1.0)
    StraightLineHomotopy(G::AbstractSystem, F::AbstractSystem; gamma = 1.0)

Constructs the straight line homotopy ``H(x, t) = γ t G(x) + (1-t) F(x)`` where
``γ`` is `gamma`.
"""
struct StraightLineHomotopy{S<:AbstractSystem,T<:AbstractSystem} <: AbstractHomotopy
    start::S
    target::T
    γ::ComplexF64

    u::Vector{ComplexF64}
    ū::Vector{ComplexDF64}
    U::Matrix{ComplexF64}

    dv_start::LA.Transpose{ComplexF64,Matrix{ComplexF64}}
    dv_target::LA.Transpose{ComplexF64,Matrix{ComplexF64}}
end

function StraightLineHomotopy(start::ModelKit.System, target::ModelKit.System; kwargs...)
    StraightLineHomotopy(ModelKitSystem(start), ModelKitSystem(target); kwargs...)
end
function StraightLineHomotopy(
    start::AbstractSystem,
    target::AbstractSystem;
    γ = 1.0,
    gamma = γ,
)
    size(start) == size(target) ||
        throw(ArgumentError("Start and target do not have the same size, got $(size(start)) and $(size(target))"))

    m, n = size(start)
    u = zeros(ComplexF64, m)
    ū = zeros(ComplexDF64, m)
    U = zeros(ComplexF64, m, n)

    dv_start = LA.transpose(zeros(ComplexF64, 5, m))
    dv_target = LA.transpose(zeros(ComplexF64, 5, m))

    StraightLineHomotopy(start, target, ComplexF64(gamma), u, ū, U, dv_start, dv_target)
end

Base.size(H::StraightLineHomotopy) = size(H.start)

function Base.show(io::IO, mime::MIME"text/plain", H::StraightLineHomotopy)
    println(io, typeof(H), ":")
    println(io, "γ: ", H.γ)
    println(io, "\nG: ")
    show(io, mime, H.start)
    println(io, "\n\nF:")
    show(io, mime, H.target)
end

function evaluate!(u, H::StraightLineHomotopy, x::AbstractVector{T}, t) where {T}
    evaluate!(u, H.start, x)

    if T isa ComplexDF64 || T isa DoubleF64
        evaluate!(H.ū, H.target, x)
        ts, tt = H.γ .* t, 1 - t
        for i = 1:size(H, 1)
            @inbounds u[i] = ts * u[i] + tt * H.ū[i]
        end
    else
        evaluate!(H.u, H.target, x)
        ts, tt = H.γ .* t, 1 - t
        for i = 1:size(H, 1)
            @inbounds u[i] = ts * u[i] + tt * H.u[i]
        end
    end
    u
end

function evaluate_and_jacobian!(
    u,
    U,
    H::StraightLineHomotopy,
    x::AbstractVector{T},
    t,
) where {T}
    evaluate_and_jacobian!(u, U, H.start, x)
    evaluate_and_jacobian!(H.u, H.U, H.target, x)
    ts, tt = H.γ .* t, 1 - t
    for i = 1:size(H, 1)
        @inbounds u[i] = ts * u[i] + tt * H.u[i]
    end
    for j = 1:size(H, 2), i = 1:size(H, 1)
        @inbounds U[i, j] = ts * U[i, j] + tt * H.U[i, j]
    end
    nothing
end
function taylor!(u, ::Val{1}, H::StraightLineHomotopy, x, t)
    evaluate!(u, H.start, x)
    evaluate!(H.u, H.target, x)
    for i = 1:size(H, 1)
        @inbounds u[i] = H.γ * u[i] - H.u[i]
    end
    u
end
function taylor!(u, v::Val{K}, H::StraightLineHomotopy, tx::TaylorVector, t) where {K}
    taylor!(H.dv_start, v, H.start, tx)
    taylor!(H.dv_target, v, H.target, tx)

    for i = 1:size(H, 1)
        start = H.γ * (H.dv_start[i, K] + t * H.dv_start[i, K+1])
        target = (1 - t) * H.dv_target[i, K+1] - H.dv_target[i, K]
        u[i] = start + target
    end
    u
end


function brent(f::Function, x0::Number, x1::Number, args::Tuple=();
               xtol::AbstractFloat=1e-7, ytol=2eps(Float64),
               maxiter::Integer=50)
    EPS = eps(Float64)
    y0 = f(x0,args...)
    y1 = f(x1,args...)
    if abs(y0) < abs(y1)
        # Swap lower and upper bounds.
        x0, x1 = x1, x0
        y0, y1 = y1, y0
    end
    x2 = x0
    y2 = y0
    x3 = x2
    bisection = true
    for q in 1:maxiter
        # x-tolerance.
        if abs(x1-x0) < xtol
            return x1
        end

        # Use inverse quadratic interpolation if f(x0)!=f(x1)!=f(x2)
        # and linear interpolation (secant method) otherwise.
        if abs(y0-y2) > ytol && abs(y1-y2) > ytol
            x = x0*y1*y2/((y0-y1)*(y0-y2)) +
                x1*y0*y2/((y1-y0)*(y1-y2)) +
                x2*y0*y1/((y2-y0)*(y2-y1))
        else
            x = x1 - y1 * (x1-x0)/(y1-y0)
        end

        # Use bisection method if satisfies the conditions.
        delta = abs(2EPS*abs(x1))
        min1 = abs(x-x1)
        min2 = abs(x1-x2)
        min3 = abs(x2-x3)
        if (x < (3x0+x1)/4 && x > x1) ||
           (bisection && min1 >= min2/2) ||
           (!bisection && min1 >= min3/2) ||
           (bisection && min2 < delta) ||
           (!bisection && min3 < delta)
            x = (x0+x1)/2
            bisection = true
        else
            bisection = false
        end

        y = f(x,args...)
        # y-tolerance.
        if abs(y) < ytol
            return x
        end
        x3 = x2
        x2 = x1
        if sign(y0) != sign(y)
            x1 = x
            y1 = y
        else
            x0 = x
            y0 = y
        end
        if abs(y0) < abs(y1)
            # Swap lower and upper bounds.
            x0, x1 = x1, x0
            y0, y1 = y1, y0
        end
    end
    error("Max iteration exceeded")
end


using BivMatFun;
using BenchmarkTools;
using Printf;
using LinearAlgebra;
using DelimitedFiles;
using MAT;

include("experiments_common.jl")

function run_test()

    success = true;

    matrices = [ "jordbloc1", "grcar", "smoke", "kahan2", "lesp", "sampling", "grcar-randn" ];

    m = 64;
    n = m;
    a = .5;

    fncts = [
    Dict("name" => "sqrt", "f" => (x,y,i,j) -> sqrt(x+y)),
    Dict("name" => "invsqrt", "f" => (x,y,i,j)  -> 1 ./ sqrt(x+y)),
    Dict("name" => "phi1", "f" => (x,y,i,j) -> expm1(x+y) ./ (x+y)),
    Dict("name" => "expa", "f" => (x,y,i,j) -> exp((x+y).^a))
    ]

    shift = 1.00;

    lesp = n -> convert(Matrix{ComplexF64}, diagm(-1 => 1 ./ (2:n), 0 => -(5 : 2 : 2*n+3), 1 => 2 : n));
    smoke = n -> diagm(0 => exp.(2im * pi * (1 : n) ./ n), 1 => complex(ones(n-1)), -n+1 => [ 1.0 + 0im ]);

    # err_fun2m, err_diag, err_diagm, time_fun2m, time_diag, time_diagm, nblocksA, nblocksB, d_fun2m, d_fundiag, errest
    data = zeros(length(fncts), length(matrices), 11);

    for i = 1 : length(fncts)
        for j = 1 : length(matrices);
            matrix_name = matrices[j];

            fname = fncts[i]["name"];
            f = fncts[i]["f"];

            vars = matread("$(data_directory)/experiment2_$(matrix_name)_$(fname).mat")
            A = vars["A"]; B = vars["B"]; C = vars["C"]; Y = vars["Y"]; errest = vars["errest"]

            Y2, infom = fun2m(f, A, B, C, method = BivMatFun.Diag, bs = max(m,n));

            if ! is_ci_test()
                tdiagm = @benchmark fun2m($f, $A, $B, $C, method = $BivMatFun.Diag, bs = $max($m, $n))
                tdiagm = median(tdiagm.times) / 1e9;
            else
                tdiagm = 0.0
            end

            XD = BivMatFun.diag_fun(f, A, B, C);
            if ! is_ci_test()
                tdiag = @benchmark BivMatFun.diag_fun($f, $A, $B, $C)
                tdiag = median(tdiag.times) / 1e9;
            else
                tdiag = 0.0
            end

            X, info = fun2m(f, A, B, C, method = BivMatFun.Diag);
            if ! is_ci_test()
                tsp = @benchmark fun2m($f, $A, $B, $C, method = BivMatFun.Diag);
                tsp = median(tsp.times) / 1e9;
            else
                tsp = 0.0
            end

            #% err_fun2m, err_diag, time_fun2m, time_diag, time_diagm, nblocksA, nblocksB, d_fun2m, d_fundiagm, errest
            data[i, j, 1] = norm(X - Y) / norm(Y);
            data[i, j, 2] = norm(XD - Y) / norm(Y);
            data[i, j, 3] = norm(Y2 - Y) / norm(Y);
            data[i, j, 4] = tsp;
            data[i, j, 5] = tdiag;
            data[i, j, 6] = tdiagm;
            data[i, j, 7] = info.nblocksA;
            data[i, j, 8] = info.nblocksB;
            data[i, j, 9] = info.digits;
            data[i, j, 10] = infom.digits;
            data[i, j, 11] = errest;

            # In these tests, it is expected that fun2m with Diag should work in all cases, and that
            # the method using Taylor may fail. Hence, we only check if the residue is sufficiently
            # for the case with Diag; we allow some tolerance, because the error estimate is rather
            # rough.
            success &= data[i, j, 1] < errest * 1e2;

            # writedlm("experiment2.dat", data, '\t');

            @printf("%s: %s: |X - Y|/|Y|: %e (nblocks = %d, %d, time = %fs, digits = %d); errest = %e\n",
            fname, matrix_name, norm(X - Y) / norm(Y), info.nblocksA, info.nblocksB, tsp, info.digits, errest);
            @printf("%s: %s: |Y2 - Y| / |Y|: %e, time = %fs\n", fname, matrix_name, norm(Y2 - Y) / norm(Y), tdiagm)
            @printf("%s: %s: |XD - Y|/|Y|: %e, time = %fs\n", fname, matrix_name, norm(XD - Y) / norm(Y), tdiag);
        end

        # We format the data in a way that is easier to handle and ready to include
        # in the paper TeX files.
        writedlm("exp2-$i.dat", data[i,:,:], '\t')
    end

    return success;

end

run_test();