import Base.∈

export constrained_dimensions,
       tosimplehrep,
       remove_redundant_constraints,
       remove_redundant_constraints!,
       linear_map

"""
    ∈(x::AbstractVector{N}, P::AbstractPolyhedron{N})::Bool where {N<:Real}

Check whether a given point is contained in a polyhedron.

### Input

- `x` -- point/vector
- `P` -- polyhedron

### Output

`true` iff ``x ∈ P``.

### Algorithm

This implementation checks if the point lies inside each defining half-space.
"""
function ∈(x::AbstractVector{N}, P::AbstractPolyhedron{N})::Bool where {N<:Real}
    @assert length(x) == dim(P) "a $(length(x))-dimensional point cannot be " *
        "an element of a $(dim(P))-dimensional set"

    for c in constraints_list(P)
        if dot(c.a, x) > c.b
            return false
        end
    end
    return true
end

"""
    constrained_dimensions(P::AbstractPolyhedron)::Vector{Int} where {N<:Real}

Return the indices in which a polyhedron is constrained.

### Input

- `P` -- polyhedron

### Output

A vector of ascending indices `i` such that the polyhedron is constrained in
dimension `i`.

### Examples

A 2D polyhedron with constraint ``x1 ≥ 0`` is constrained in dimension 1 only.
"""
function constrained_dimensions(P::AbstractPolyhedron)::Vector{Int}
    zero_indices = zeros(Int, dim(P))
    for constraint in constraints_list(P)
        for i in constrained_dimensions(constraint)
            zero_indices[i] = i
        end
    end
    return filter(x -> x != 0, zero_indices)
end

"""
    tosimplehrep(constraints::AbstractVector{LC})
        where {N<:Real, LC<:LinearConstraint{N}}

Return the simple H-representation ``Ax ≤ b`` from a list of linear constraints.

### Input

- `constraints` -- a list of linear constraints

### Output

The tuple `(A, b)` where `A` is the matrix of normal directions and `b` is the
vector of offsets.
"""
function tosimplehrep(constraints::AbstractVector{LC}
                     ) where {N<:Real, LC<:LinearConstraint{N}}
    n = length(constraints)
    if n == 0
        A = Matrix{N}(undef, 0, 0)
        b = Vector{N}(undef, 0)
        return (A, b)
    end
    A = zeros(N, n, dim(first(constraints)))
    b = zeros(N, n)
    @inbounds begin
        for (i, Pi) in enumerate(constraints)
            A[i, :] = Pi.a
            b[i] = Pi.b
        end
    end
    return (A, b)
end

"""
     remove_redundant_constraints!(constraints::AbstractVector{LC};
         [backend]=GLPKSolverLP())::Bool where {N<:Real,
                                                LC<:LinearConstraint{N}}

Remove the redundant constraints of a given list of linear constraints; the list
is updated in-place.

### Input

- `constraints` -- list of constraints
- `backend`     -- (optional, default: `GLPKSolverLP`) numeric LP solver backend

### Output

`true` if the function was successful and the list of constraints `constraints`
is modified by removing the redundant constraints, and `false` only if the
constraints are infeasible.

### Notes

Note that the result may be `true` even if the constraints are infeasible.
For example, ``x ≤ 0 && x ≥ 1`` will return `true` without removing any
constraint.
To check if the constraints are infeasible, use
`isempty(HPolyhedron(constraints)`.

### Algorithm

If there are `m` constraints in `n` dimensions, this function checks one by one
if each of the `m` constraints is implied by the remaining ones.

To check if the `k`-th constraint is redundant, an LP is formulated using the
constraints that have not yet being removed.
If, at an intermediate step, it is detected that a subgroup of the constraints
is infeasible, this function returns `false`.
If the calculation finished successfully, this function returns `true`.

For details, see [Fukuda's Polyhedra
FAQ](https://www.cs.mcgill.ca/~fukuda/soft/polyfaq/node24.html).
"""
function remove_redundant_constraints!(constraints::AbstractVector{LC};
                                       backend=GLPKSolverLP()
                                      )::Bool where {N<:Real,
                                                     LC<:LinearConstraint{N}}

    A, b = tosimplehrep(constraints)
    m, n = size(A)
    non_redundant_indices = 1:m

    i = 1 # counter over reduced constraints

    for j in 1:m    # loop over original constraints
        α = A[j, :]
        Ar = A[non_redundant_indices, :]
        br = b[non_redundant_indices]
        br[i] = b[j] + one(N)
        lp = linprog(-α, Ar, '<', br, -Inf, Inf, backend)
        if lp.status == :Infeasible
            # the polyhedron is empty
            return false
        elseif lp.status == :Optimal
            objval = -lp.objval
            if objval <= b[j]
                # the constraint is redundant
                non_redundant_indices = setdiff(non_redundant_indices, j)
            else
                # the constraint is not redundant
                i = i+1
            end
        else
            error("LP is not optimal; the status of the LP is $(lp.status)")
        end
    end

    deleteat!(constraints, setdiff(1:m, non_redundant_indices))
    return true
end

"""
    remove_redundant_constraints(constraints::AbstractVector{LC};
        backend=GLPKSolverLP()) where {N<:Real, LC<:LinearConstraint{N}}

Remove the redundant constraints of a given list of linear constraints.

### Input

- `constraints` -- list of constraints
- `backend`     -- (optional, default: `GLPKSolverLP`) numeric LP solver backend

### Output

The list of constraints with the redundant ones removed, or an empty set if the
constraints are infeasible.

### Algorithm

See
[`remove_redundant_constraints!(::AbstractVector{<:LinearConstraint{<:Real}})`](@ref)
for details.
"""
function remove_redundant_constraints(constraints::AbstractVector{LC};
                                      backend=GLPKSolverLP()
                                     ) where {N<:Real, LC<:LinearConstraint{N}}
    constraints_copy = copy(constraints)
    if remove_redundant_constraints!(constraints_copy, backend=backend)
        return constraints_copy
    else  # the constraints are infeasible
        return EmptySet{N}()
    end
end

"""
    linear_map(M::AbstractMatrix{N},
               P::AbstractPolyhedron{N};
               check_invertibility::Bool=true,
               cond_tol::Number=DEFAULT_COND_TOL,
               use_inv::Bool=!issparse(M)
               ) where {N<:Real}

Concrete linear map of a polyhedron in constraint representation.

### Input

- `M` -- matrix
- `P` -- abstract polyhedron
- `check_invertibility` -- (optional, deault: `true`) check if the linear map is
                           invertible, in which case this function uses the matrix
                           inverse; if the invertibility check fails, or if
                           this flag is set to `false`, use the vertex representation
                           to compute the linear map (see below for details)
- `cond_tol` -- (optional) tolerance of matrix condition (used to check whether
                the matrix is invertible)
- `use_inv`  -- (optional, default: `false` if `M` is sparse and `true`
                otherwise) whether to compute the full left division through
                `inv(M)`, or to use the left division for each vector; see below

### Output

The type of the result is "as close as possible" to the the type of `P`.
Let `(m, n)` be the size of `M`, where `m ≠ n` is allowed for rectangular maps.

To fix the type of the output to something different than the default value,
consider post-processing the result of this function with a call to a suitable
`convert` method.

In particular, the output depends on the type of `P`, on `m`, and the algorithm
that was used:

- If the vertex-based approach was used:

    - If `P` is a `VPolygon` and `m = 2` then the output is a `VPolygon`.
    - If `P` is a `VPolytope` then the output is a `VPolytope`.
    - Otherwise, the output is an `Interval` if `m = 1`, a `VPolygon` if `m = 2`
      and a `VPolytope` in other cases.

- If the invertibility criterion was used:

    - The types of `HalfSpace`, `Hyperplane`, `Line` and `AbstractHPolygon` are
      preserved.
    - If `P` is an `AbstractPolytope`, then the output is an `Interval` if `m = 1`,
      an `HPolygon` if `m = 2` and an `HPolytope` in other cases.
    - Otherwise, the output is an `HPolyhedron`.

### Algorithm

This function implements two algorithms for the linear map:

- If the matrix ``M`` is invertible (which we check with a sufficient condition),
  then ``y = M x`` implies ``x = \\text{inv}(M) y`` and we transform the
  constraint system ``A x ≤ b`` to ``A \\text{inv}(M) y ≤ b``.
- Otherwise, we transform the polyhedron to vertex representation and apply the map
  to each vertex, returning a polyhedron in vertex representation.

Note that the vertex representation (second approach) is only available if the
polyhedron is bounded. Hence we check boundedness first.

To switch off the check for invertibility, set the option
`check_invertibility=false`. If `M` is not invertible and the polyhedron
is unbounded, this function returns an exception.

The option `use_inv` lets the user control - in case `M` is invertible - if
the full matrix inverse is computed, or only the left division on the normal
vectors. Note that this helps as a workaround when `M` is a sparse matrix, since
the `inv` function is not available for sparse matrices. In this case, either
use the option `use_inv=false` or convert the type of `M` as in
`linear_map(Matrix(M), P)`.

Internally, this function operates at the level of the `AbstractPolyhedron`
interface, but the actual algorithm uses dispatch on the concrete type of `P`,
depending on the algorithm that is used:

-  `_linear_map_vrep(M, P)` if the vertex approach is used
-  `_linear_map_hrep(M, P, use_inv)` if the invertibility criterion is used

New subtypes of the interface should write their own `_linear_map_vrep`
(resp. `_linear_map_hrep`) for special handling of the linear map; otherwise
the fallback implementation for `AbstractPolyhedron` is used (see below).
"""
function linear_map(M::AbstractMatrix{N},
                    P::AbstractPolyhedron{N};
                    check_invertibility::Bool=true,
                    cond_tol::Number=DEFAULT_COND_TOL,
                    use_inv::Bool=!issparse(M)
                   ) where {N<:Real}
    @assert dim(P) == size(M, 2) "a linear map of size $(size(M)) cannot be " *
        "applied to a set of dimension $(dim(P))"

    if !check_invertibility || !isinvertible(M; cond_tol=cond_tol)
        # vertex representation is enforced or the matrix is not invertible => use vertex approach
        return _linear_map_vrep(M, P)
    else
        # matrix M is invertible => use H-rep approach
        # the normal vectors are vec(c.a' * inv(M))
        return _linear_map_hrep(M, P, use_inv)
    end
end

# handle different numeric types
function linear_map(M::AbstractMatrix{NM},
                    P::AbstractPolyhedron{NP};
                    kwargs...) where {NM<:Real, NP<:Real}
    N = promote_type(NM, NP)
    if N != NP
        error("conversion between numeric types of polyhedra not implemented " *
            "yet (see #1181)")
    else
        return linear_map(N.(M), P; kwargs...)
    end
end

function _linear_map_vrep(M::AbstractMatrix{N}, P::AbstractPolyhedron{N}) where {N<:Real}
    if !isbounded(P)
        throw(ArgumentError("the linear map in vertex representation for an " *
            "unbounded set is not defined"))
    end
    require(:Polyhedra; fun_name="linear_map",
            explanation="of a $(typeof(P)) by a non-invertible matrix")
    # since P is bounded, we pass an HPolytope and then convert it to vertex representation
    P = tovrep(HPolytope(constraints_list(P), check_boundedness=false))
    return _linear_map_vrep(M, P)
end

function _linear_map_hrep(M::AbstractMatrix{N}, P::AbstractPolyhedron{N},
                          use_inv::Bool) where {N<:Real}
    constraints = _linear_map_hrep_helper(M, P, use_inv)
    return HPolyhedron(constraints)
end

function _linear_map_hrep_helper(M::AbstractMatrix{N}, P::AbstractPolyhedron{N},
                                 use_inv::Bool) where {N<:Real}
    constraints_P = constraints_list(P)
    constraints_MP = similar(constraints_P)
    if use_inv
        invM = inv(M)
        @inbounds for (i, c) in enumerate(constraints_P)
            constraints_MP[i] = LinearConstraint(vec(_At_mul_B(c.a, invM)), c.b)
        end
    else
        # take left division for each constraint c, transpose(M) \ c.a
        @inbounds for (i, c) in enumerate(constraints_P)
            constraints_MP[i] = LinearConstraint(_At_ldiv_B(M, c.a), c.b)
        end
    end
    return constraints_MP
end

"""
    plot_recipe(P::AbstractPolyhedron{N}, [ε]::N=zero(N)) where {N<:Real}

Convert a (bounded) polyhedron to a pair `(x, y)` of points for plotting.

### Input

- `P` -- bounded polyhedron
- `ε` -- (optional, default: `0`) ignored, used for dispatch

### Output

A pair `(x, y)` of points that can be plotted.

### Algorithm

We first assert that `P` is bounded (i.e., that `P` is a polytope).

One-dimensional polytopes are converted to an `Interval`.
Three-dimensional or higher-dimensional polytopes are not supported.

For two-dimensional polytopes (i.e., polygons) we compute their set of vertices
using `vertices_list` and then plot the convex hull of these vertices.
"""
function plot_recipe(P::AbstractPolyhedron{N}, ε::N=zero(N)) where {N<:Real}
    @assert dim(P) <= 2 "cannot plot a $(dim(P))-dimensional $(typeof(P))"
    @assert isbounded(P) "cannot plot an unbounded $(typeof(P))"

    if dim(P) == 1
        return plot_recipe(convert(Interval, P), ε)
    else
        vlist = transpose(hcat(convex_hull(vertices_list(P))...))
        if isempty(vlist)
            @warn "received a polyhedron with no vertices during plotting"
            return plot_recipe(EmptySet{N}(), ε)
        end
        x, y = vlist[:, 1], vlist[:, 2]

        if length(x) > 1
            # add first vertex to "close" the polygon
            push!(x, x[1])
            push!(y, y[1])
        end

        return x, y
    end
end


### Combination methods for p-values ###

"""
    combine(PValues, <:PValueCombination)

Combine p-values

# Examples

```jldoctest
julia> pvals = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pvals, Fisher())
0.007616871850449092

julia> combine(pvals, Stouffer())
0.00709832618126593
```

# See also

`PValueCombination`s:

[`Fisher`](@ref)
[`Logit`](@ref)
[`Stouffer`](@ref)
[`Tippett`](@ref)
[`Simes`](@ref)
[`Wilkinson`](@ref)
[`Minimum`](@ref)

"""
function combine end

function combine(pValues::AbstractVector{T}, method::M)::T where {T <: AbstractFloat,M <: PValueCombination}
    return combine(PValues(pValues), method)
end

## Fisher combination ##

"""
Fisher's p-value combination

# Examples

```jldoctest
julia> pvals = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pvals, Fisher())
0.007616871850449092
```

# References

Fisher, R.A. (1925).
Statistical methods for research workers
(Genesis Publishing Pvt Ltd).

"""
struct Fisher <: PValueCombination
end

function combine(pValues::PValues{T}, method::Fisher)::T where T <: AbstractFloat
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    if minimum(pValues) == 0
        return NaN
    end
    x = -2 * sum(log.(pValues))
    p = ccdf(Chisq(2n), Float64(x))
    return p
end

## Logit combination ##

"""
Logit p-value combination

# Examples

```jldoctest
julia> pvals = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pvals, Logit())
0.006434494635148462
```

# References

Mudholkar, G.S., and George, E.O. (1977).
The Logit Statistic for Combining Probabilities - An Overview
(Rochester University NY, Dept of Statistics).

"""
struct Logit <: PValueCombination
end

function combine(pValues::PValues{T}, method::Logit)::T where T <: AbstractFloat
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    if minimum(pValues) == 0 || maximum(pValues) == 1
        return NaN
    end
    c = sqrt((5n + 2) * n * pi^2 / ((5n + 4) * 3))
    x = -sum(log.(pValues ./ (1 .- pValues))) / c
    p = ccdf(TDist(5n + 4), x)
    return p
end

## Stouffer combination ##

"""
Stouffer's p-value combination

# Examples

```jldoctest
julia> pvals = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pvals, Stouffer())
0.00709832618126593

julia> weights = [1.0, 2.0, 0.4, 1.5];

julia> combine(pvals, weights, Stouffer())
0.007331653763696742
```

# References

Stouffer, S.A. (1949).
The American soldier. Vol. 1: Adjustment during army life
(Princeton University Press).

Liptak, T. (1958).
On the combination of independent tests.
Magyar Tud Akad Mat Kutato Int Kozl 3, 171–197.

"""
struct Stouffer <: PValueCombination
end

function combine(pValues::PValues{T}, method::Stouffer)::T where T <: AbstractFloat
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    if minimum(pValues) == 0 || maximum(pValues) == 1
        return NaN
    end
    z = cquantile.(Ref(Normal()), pValues)
    z = sum(z) ./ sqrt(n)
    p = ccdf(Normal(), z)
    return p
end

function combine(pValues::PValues{T}, weights::AbstractVector{R}, method::Stouffer)::T where {T <: AbstractFloat,R <: Real}
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    if minimum(pValues) == 0 || maximum(pValues) == 1
        return NaN
    end
    z = cquantile.(Ref(Normal()), pValues) .* weights
    z = sum(z) ./ sqrt(sum(abs2, weights))
    p = ccdf(Normal(), z)
    return p
end

function combine(pValues::AbstractVector{T}, weights::AbstractVector{R}, method::Stouffer)::T where {T <: AbstractFloat,R <: Real}
    return combine(PValues(pValues), weights, method)
end

## Tippett combination ##

"""
Tippett's p-value combination

# Examples

```jldoctest
julia> pvals = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pvals, Tippett())
0.039403990000000055
```

# References

Tippett, L.H.C. (1931). The Methods of Statistics. An introduction mainly for
workers in the biological sciences.

"""
struct Tippett <: PValueCombination
end

function combine(pValues::PValues{T}, method::Tippett)::T where T <: AbstractFloat
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    p = 1 - (1 - minimum(pValues))^n
    return p
end

## Simes combination ##

"""
Simes's p-value combination

# Examples

```jldoctest
julia> pvals = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pvals, Simes())
0.04
```

# References

Simes, R.J. (1986). An improved Bonferroni procedure for multiple tests of
significance. Biometrika 73, 751–754.

"""
struct Simes <: PValueCombination
end

function combine(pValues::PValues{T}, method::Simes)::T where T <: AbstractFloat
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    pValues = sort(pValues)
    p = n * minimum(pValues ./ (1:n))
    return p
end

## Wilkinson combination ##

"""
Wilkinson's p-value combination

# Examples

```jldoctest
julia> pv = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pv, Wilkinson(1))  # combination with rank 1
0.03940399000000003

julia> combine(pv, Wilkinson(4))  # combination with rank 4
0.0625
```

# References

Wilkinson, B. (1951). A statistical consideration in psychological research.
Psychological Bulletin 48, 156.

"""
struct Wilkinson <: PValueCombination
    rank::Integer

    function Wilkinson(rank)
        if rank < 1
            throw(ArgumentError("Rank must be positive."))
        end
        return new(rank)
    end
end

function combine(pValues::PValues{T}, method::Wilkinson)::T where T <: AbstractFloat
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    rank = method.rank
    if rank < 1 || rank > n
        throw(ArgumentError("Rank must be in 1,..,$(n)"))
    end
    p_rank = partialsort(pValues, rank)
    p = cdf(Beta(rank, n - rank + 1), Float64(p_rank))
    return p
end

## Generalised minimum combination ##

"""
Minimum of adjusted p-value combination

# Examples

```jldoctest
julia> pv = PValues([0.01, 0.02, 0.3, 0.5]);

julia> combine(pv, Minimum(BenjaminiHochberg()))
0.04

julia> combine(pv, Minimum(ForwardStop()))
0.01005033585350145
```

"""
struct Minimum <: PValueCombination
    adjustment::PValueAdjustment
end

function combine(pValues::PValues{T}, method::Minimum)::T where T <: AbstractFloat
    n = length(pValues)
    if n == 1
        return pValues[1]
    end
    padj = adjust(pValues, method.adjustment)
    p = minimum(padj)
    return p
end


struct Nested <: InferenceType end
const NESTED = Nested()

nested_results(traces, log_weights) = NonparametricSamplingResults{Nested}(NESTED, log_weights, Array{Any, 1}(), traces)

function sample(r :: SamplingResults{Nested}, k, n :: Int)
    v = r[k]
    log_Z = logsumexp(r.log_weights)
    weights = StatsBase.Weights(exp.(r.log_weights .- log_Z))
    if n == 1
        StatsBase.sample(v, weights)
    else
        StatsBase.sample(v, weights, n)
    end
end

@doc raw"""
    function nested(f, replace_fn; params = (), num_points :: Int64 = 1)

Generic implementation of nested sampling (Skilling, Nested sampling for general Bayesian computation, Bayesian Analysis, 2006).
The number of sampling iterations is a function of `num_points` aka ``N``, and the empirical entropy of the sampling distribution, given at the
``n``-th iteration by ``H_n \approx \sum_k \hat p_k \log \hat p_k^{-1}``, where ``\hat p_k = L_k w_k / Z_k``, ``L_k`` is the likelihood
value, ``w_k`` is the difference in prior mass, and ``Z_k`` is the current estimate of the partition function. The number of sampling iterations
is equal to ``\min_n \{n > 0: n > NH_n\}``.

Args:
+ `f`: stochastic function. Must have signature `f(t :: Trace, params...)`
+ `replace_fn`: function that returns a tuple `(new_trace :: Trace, new_log_likelihood :: Float64)`. The input signature of this function
    must be `replace_fn(f :: F, log_likelihood :: Float64, params...) where F <: Function`. It must guarantee that `new_log_likelihood > log_likelihood`.
+ `params`: any parameters to pass to `f`
+ `num_points`: the number of likelihood points to keep track of
"""
function nested(f, replace_fn; params = (), num_points :: Int64 = 1)
    traces = [trace() for _ in 1:num_points]
    for t in traces
        f(t, params...)
    end
    loglikelihoods = [loglikelihood(t) for t in traces]
    weights = Float64[]
    pts = Array{Trace, 1}()
    x_last = 1.0
    Z = 0.0  # partition function estimate
    n = 1  # number of iterations

    while true
        min_l_ix = argmin(loglikelihoods)
        log_l_n = loglikelihoods[min_l_ix]

        x_n = exp(-n / num_points)
        w_n = x_last - x_n
        x_last = x_n

        weight_n = log_l_n + log(w_n)
        Z += exp(weight_n)
        push!(weights, weight_n)
        push!(pts, traces[min_l_ix])

        new_t, this_log_l = replace_fn(f, log_l_n, params...)
        traces[min_l_ix] = new_t
        loglikelihoods[min_l_ix] = this_log_l

        # estimate entropy
        normed_weights = weights .- log(Z)
        if n > 1
            entropy = -1.0 * sum(exp.(normed_weights) .* normed_weights)
            breakpoint = Int(num_points * ceil(entropy))
            if n > breakpoint
                break
            end
        end
        n += 1
    end
    # approximation for the remainder of the sampled traces
    append!(pts, traces)
    x_last = exp(-(n - 1) / num_points)
    x_n = exp(-n / num_points)
    w_n = x_last - x_n
    for t in traces
        push!(weights, loglikelihood(t) + log(w_n))
    end
    nested_results(pts, weights)
end

@doc raw"""
    nested(f; params = (), num_points :: Int64 = 1)

Run nested sampling using internal rejection method.
"""
nested(f; params = (), num_points :: Int64 = 1) = nested(f, rejection; params = params, num_points = num_points)

export nested, sample


#Came from GLAbstraction/GLUtils.jl
#question: Used?
function close_to_square(n::Real)
    # a cannot be greater than the square root of n
    # b cannot be smaller than the square root of n
    # we get the maximum allowed value of a
    amax = floor(Int, sqrt(n))
    if 0 == rem(n, amax)
        # special case where n is a square number
        return (amax, div(n, amax))
    end
    # Get its prime factors of n
    primeFactors = factor(n)
    # Start with a factor 1 in the list of candidates for a
    candidates = Int[1]
    for (f, _) in primeFactors
        # Add new candidates which are obtained by multiplying
        # existing candidates with the new prime factor f
        # Set union ensures that duplicate candidates are removed
        candidates = union(candidates, f .* candidates)
        # throw out candidates which are larger than amax
        filter!(x -> x <= amax, candidates)
    end
    # Take the largest factor in the list d
    return (candidates[end], div(n, candidates[end]))
end

#Came from GLAbstraction/GLCamera.jl
w_component(::Point{N,T}) where {N,T} = T(1)
w_component(::Vec{N,T}) where {N,T} = T(0)

# @component struct AABB
#     origin  ::Vec3f0
#     diagonal::Vec3f0
# end

# function aabb_ray_intersect(aabb::AABB, entity_pos::Point3f0, ray_origin::Point3f0, ray_direction::Vec3f0)
#     dirfrac = 1.0f0 ./ ray_direction
#     right   = aabb.origin + aabb.diagonal

#     real_origin = aabb.origin + entity_pos
#     real_right  = right + entity_pos
#     t1 = (real_origin - ray_origin) .* dirfrac
#     t2 = (real_right - ray_origin)  .* dirfrac
#     tsmaller = min.(t1, t2)
#     tbigger  = max.(t1,t2)

#     tmin = maximum(tsmaller)
#     tmax = minimum(tbigger)

#     return tmin < tmax
# end


function test_tf_idf()
    println("-- Testing TF-IDF content-based recommender")

    m = [1  0  1  0  1  1  0  0  0  1
         0  1  1  1  0  0  0  1  0  0
         0  0  0  1  1  1  0  0  0  0
         0  0  1  1  0  0  1  1  0  0
         0  1  0  0  0  0  0  0  1  1
         1  0  0  1  0  0  0  0  0  0
         0  0  0  0  0  0  0  1  0  1
         0  0  1  1  0  0  1  0  0  1
         0  0  0  0  0  1  0  0  1  0
         0  1  0  0  1  0  1  0  0  0
         0  0  1  0  1  0  0  0  1  0
         1  0  0  0  0  1  1  0  0  0
         0  0  1  1  1  0  0  1  0  0
         0  1  1  1  0  0  0  0  1  0
         0  0  0  1  0  1  1  1  0  0
         1  0  0  0  0  1  0  0  1  0
         0  1  1  1  0  0  0  1  0  0
         0  0  0  1  0  0  0  0  1  0
         0  1  1  0  1  0  1  0  0  1
         0  0  1  1  0  0  1  0  1  0]

    tf = sum(m, dims=1)
    idf = 1 ./ tf

    users = [ 1  -1
             -1   1
              0   0
              0   1
              0   0
              1   0
              0   0
              0   0
              0   0
              0   0
              0   0
              0  -1
              0   0
              0   0
              0   0
              1   0
              0   1
              0   0
             -1   0
              0   0]
    uv1 = sparse(users[:, 1])
    uv2 = sparse(users[:, 2])

    n_doc = size(m, 1)
    docs = [i for i in 1:20]

    # Case: basic
    data = DataAccessor(m)
    set_user_attribute(data, 1, uv1)
    set_user_attribute(data, 2, uv2)

    recommender = TFIDF(data, tf, ones(size(tf))) # do not use IDF
    rec = recommend(recommender, 1, n_doc, docs)
    @test first(rec[2]) == 1
    @test first(rec[3]) == 12
    @test last(rec[2]) == 4 && last(rec[3]) == 4

    # Case: normalized matrix
    m_normalized = m ./ sqrt.(sum(m.^2, dims=2))
    data = DataAccessor(m_normalized)
    set_user_attribute(data, 1, uv1)
    set_user_attribute(data, 2, uv2)

    recommender = TFIDF(data, tf, ones(size(tf))) # do not use IDF
    rec = recommend(recommender, 1, n_doc, docs)
    @test first(rec[5]) == 1
    @test isapprox(last(rec[5]), 1.0090, atol=1e-4)
    rec = recommend(recommender, 2, n_doc, docs)
    @test first(rec[10]) == 7
    @test isapprox(last(rec[10]), 0.7444, atol=1e-4)
    @test first(rec[11]) == 19
    @test isapprox(last(rec[11]), 0.4834, atol=1e-4)

    # Case: normalized matrix & using IDF
    recommender = TFIDF(data, tf, idf)
    rec = recommend(recommender, 1, n_doc, docs)
    @test first(rec[4]) == 1
    @test isapprox(last(rec[4]), 0.2476, atol=1e-4)
end

test_tf_idf()


using SimpleSudoku
using Test

@testset "SimpleSudoku.jl" begin

    # New York Times 2022 02 14 Hard
    grid_str = """
            300007400
            000010083
            500020700
            980000006
            000005001
            010000020
            620070000
            009600000
            050400000"""

    grid = mapreduce(vcat, split(grid_str, "\n")) do str
        reshape(parse.(Int8, split(str, "")), 1, 9)
    end

    solved_grid, _ = solve_sudoku(grid)

    solved_grid
    #  3  6  1  9  8  7  4  5  2
    #  2  9  7  5  1  4  6  8  3
    #  5  4  8  3  2  6  7  1  9
    #  9  8  5  2  4  1  3  7  6
    #  7  3  2  8  6  5  9  4  1
    #  4  1  6  7  3  9  5  2  8
    #  6  2  4  1  7  3  8  9  5
    #  1  7  9  6  5  8  2  3  4
    #  8  5  3  4  9  2  1  6  7

end


@testset "Homology" begin
    cplx = SimplicialComplex(Simplex(1,2,3),
                        Simplex(2,4),
                        Simplex(3,4),
                        Simplex(5,4),
                        Simplex(6))

    h = homology(cplx)
    @test eltype(h) == Tuple{Int64,Int64,Int64}
    @test grouptype(supertype(typeof(h))) == Nothing
    @test grouptype(typeof(h)) == Int
    @test length(h) == 3

    b, t, _ = group(h,0);
    @test b == 2
    @test t == 0

    itm, st = iterate(h)
    @test itm[1] == 0
    @test itm[2] == 2
    @test st[1] == 1

    itm, st = iterate(h, st)
    @test itm[1] == 1
    @test itm[2] == 1
    @test st[1] == 2

    itm, st = iterate(h, st)
    @test itm[1] == 2
    @test itm[2] == 0
    @test st[1] == 3

    @test iterate(h, st) === nothing

    @test ComputationalHomology.betti(h) == [2,1,0]
    @test ComputationalHomology.euler(h) == 1

    g = withgenerators(h)
    @test eltype(g) == Tuple{Int64,Int64,Int64,Dict{Chain,Int64}}
    @test length(g) == 3

    itm, st = iterate(g)
    @test itm[1] == 0
    @test itm[2] == 2
    @testset for (ch, d) in itm[4]
        @test d == 0
        @test ch[1][1] == 1
        @test ch[1][2] in [6,5]
    end

    itm, st = iterate(g, st)
    @test itm[1] == 1
    @test itm[2] == 1
    @test first(itm[4])[2] == 0
    @testset for (a,b) in zip( simplify(first(itm[4])[1]), Chain(1, Int) + (-1, 4) + (1,3) + (1,5))
        @test a == b
    end

    itm, st = iterate(g, st)
    @test itm[1] == 2
    @test itm[2] == 0
    @test length(itm[4]) == 0

    @test iterate(g, st) === nothing

    @test length(generators(g)) == 3
end


import Rotations
import StaticArrays: StaticVector, @SVector
import PoseComposition: Pose
import LinearAlgebra: I
import Base: convert

function Pose(orientation::Rotations.Rotation{3})::Pose
    return Pose(zeros(3), orientation)
end

const IDENTITY_ORIENTATION = Rotations.UnitQuaternion(1.0,0.0,0.0,0.0)

function Pose(pos::AbstractVector{<:Real})::Pose
    return Pose(pos, IDENTITY_ORIENTATION)
end

function Pose(x::Real, y::Real, z::Real)::Pose
    return Pose([x,y,z], IDENTITY_ORIENTATION)
end

const IDENTITY_POSE = Pose(zeros(3),IDENTITY_ORIENTATION)

rotation_pose(pose::Pose) = Pose(zeros(3), pose.orientation)
translation_pose(pose::Pose) = Pose(pose.pos, IDENTITY_ORIENTATION)

"""
    c_relative_to_a::Pose = get_c_relative_to_a(b_relative_to_a::Pose, c_relative_to_b::Pose)

Given the pose of b relative to a, and the pose of c relative to b, return the pose of c relative to a.
"""
function get_c_relative_to_a(b_relative_to_a::Pose, c_relative_to_b::Pose)
    pos = b_relative_to_a.pos + (b_relative_to_a.orientation * c_relative_to_b.pos)
    orientation = b_relative_to_a.orientation * c_relative_to_b.orientation
    return Pose(pos, orientation)
end

function pose_pow(pose::Pose, t::Int)
    if t == 0
        return IDENTITY_POSE
    end
    p = pose
    for i in 1:(t-1)
        p = get_c_relative_to_a(p, pose)
    end
    p
end

"""
    b_relative_to_a::Pose = inverse_pose(a_relative_to_b::Pose)

Given the pose of b relative to a, return the pose of a relative to b.
"""
function inverse_pose(a_relative_to_b::Pose)
    pos = a_relative_to_b.orientation' * (-a_relative_to_b.pos)
    orientation = inv(a_relative_to_b.orientation)
    b_relative_to_a = Pose(pos, orientation)
    return b_relative_to_a
end

"""
    points = get_points_in_frame_b(points_in_frame_a::Matrix{<:Real}, b_relative_to_a::Pose)

Given coordinates of points in one coordinate frame, get the coordinates of the points in another coordinate frame.

The points that are taken as input and returned as output are both 3 x n matrices.
"""
function get_points_in_frame_b(points_in_frame_a::Matrix{<:Real}, b_relative_to_a::Pose)
    if size(points_in_frame_a)[1] != 3
        error("expected an 3 x n matrix")
    end
    n = size(points_in_frame_a[2])
    return b_relative_to_a.orientation' * (points_in_frame_a .- b_relative_to_a.pos)
end

function move_points_to_frame_b(points_in_frame_a::Matrix{<:Real}, b_relative_to_a::Pose)
    if size(points_in_frame_a)[1] != 3
        error("expected an 3 x n matrix")
    end
    n = size(points_in_frame_a[2])
    return b_relative_to_a.orientation * points_in_frame_a .+ b_relative_to_a.pos
end

function convert(::Type{Matrix{Float32}}, pose::Pose)
    R = Matrix{Float32}(pose.orientation)
    mat = Matrix{Float32}(I, 4, 4)
    mat[1:3,1:3] = R
    mat[1:3,4] = pose.pos
    return mat
end

export Pose, close, isapprox, IDENTITY_ORIENTATION, IDENTITY_POSE, inverse_pose, get_points_in_frame_b,
        get_c_relative_to_a, move_points_to_frame_b, pose_pow, rotation_pose, translation_pose


using PyCall
#pygui(:qt5) # good, default
# pygui(:wx) # crashes
# pygui(:gtk3)
using Random
import PyPlot as plt
#const plt = PyPlot

# https://github.com/JuliaPy/PyPlot.jl
function learn_basic_usage()
    # use x = linspace(0,2*pi,1000) in Julia 0.6
    x = range(0; stop=2*pi, length=1000)
    y = sin.(3 * x + 4 * cos.(2 * x))
    plt.plot(x, y, color="red", linewidth=2.0, linestyle="--")
    plt.title("A sinusoidally modulated sinusoid")
end

# https://matplotlib.org/stable/tutorials/introductory/pyplot.html
function learn_intro_to_pyplot_1()
    plt.clf() # clear figure, cla(), close()
    # x data are [0, 1, 2, 3]
    plt.plot([1, 2, 3, 4])
    plt.ylabel("Some numbers")
    plt.show()

    # https://stackoverflow.com/questions/6916978/how-do-i-tell-matplotlib-to-create-a-second-new-plot-then-later-plot-on-the-o
end
function learn_intro_to_pyplot_2()
    plt.clf() # clear figure, cla(), close()
    plt.plot([1, 2, 3, 4], [1, 4, 9, 16]) # "b-"
    plt.show()
end
function learn_intro_to_pyplot_3()
    plt.clf() # clear figure, cla(), close()
    plt.plot([1, 2, 3, 4], [1, 4, 9, 16], "ro")
    plt.axis([0, 6, 0, 20])
    plt.show()
end

function learn_intro_to_pyplot_4()
    plt.clf() # clear figure, cla(), close()
    t = range(0, step=0.2, stop=5)
    @show t
    plt.plot(
        t, t, "r--",
        t, t.^2, "bs",
        t, t.^3, "g^"
    )
    plt.show()
end

function learn_intro_to_pyplot_5()
    plt.clf() # clear figure, cla(), close()
    data = Dict(
        "a" => range(0,length=50),
        "c" => rand(0:50, 50),
        "d" => randn(50),
    )
    data["b"] = data["a"] + 10 * randn(50)
    data["d"] = abs.(data["d"]) .* 100

    plt.scatter("a", "b", c="c", s="d", data=data)
    plt.xlabel("entry a")
    plt.ylabel("entry b")
    plt.show()
end

#learn_basic_usage()
#learn_intro_to_pyplot_1()
#learn_intro_to_pyplot_2()
#learn_intro_to_pyplot_3()
#learn_intro_to_pyplot_4()
learn_intro_to_pyplot_5()

nothing


using DiffEqFlux, StaticArrays, Test

fd = FastDense(2,25,tanh)
pd = initial_params(fd)
fd(ones(2),pd)

f1 = FastDense(2,25,tanh)
f2 = FastDense(25,2,tanh)
p1 = initial_params(f1)
p2 = initial_params(f2)
@test FastChain(f1,f2)(ones(2),[p1;p2]) == f2(f1(ones(2),p1),p2)

f = FastChain(FastDense(2,25,tanh),FastDense(25,2,tanh))
p = initial_params(f)
@test f(ones(2),p) == f2(f1(ones(2),p[1:length(p1)]),p[length(p1)+1:end])

fd1 = FastDense(2,25,tanh,bias=false)
pd1 = initial_params(fd1)
fd1(ones(2),pd1)

f3 = FastDense(2,25,tanh,bias=false)
f4 = FastDense(25,2,tanh,bias=false)
p3 = initial_params(f3)
p4 = initial_params(f4)

@test FastChain(f3,f4)(ones(2),[p3;p4]) == f4(f3(ones(2),p3),p4)

f5 = FastChain(FastDense(2,25,tanh,bias=false),FastDense(25,2,tanh,bias=false))
p5 = initial_params(f5)
@test f5(ones(2),p5) == f4(f3(ones(2),p5[1:length(p3)]),p5[length(p3)+1:end])

fs = StaticDense(2,25,tanh)
x = rand(2)

@test fs(x,pd) ≈ fd(x,pd)
fdgrad = Flux.Zygote.gradient((x,p)->sum(fd(x,p)),x,pd)
fsgrad = Flux.Zygote.gradient((x,p)->sum(fs(x,p)),x,pd)
@test fdgrad[1] ≈ fsgrad[1]
@test fdgrad[2] ≈ fsgrad[2] rtol=1e-5

# Now test vs Zygote
struct TestDense{F,F2} <: DiffEqFlux.FastLayer
  out::Int
  in::Int
  σ::F
  initial_params::F2
  function TestDense(in::Integer, out::Integer, σ = identity;
                 initW = Flux.glorot_uniform, initb = Flux.zeros)
    initial_params() = vcat(vec(initW(out, in)),initb(out))
    new{typeof(σ),typeof(initial_params)}(out,in,σ,initial_params)
  end
end
(f::TestDense)(x,p) = f.σ.(reshape(p[1:(f.out*f.in)],f.out,f.in)*x .+ p[(f.out*f.in+1):end])
ts = StaticDense(2,25,tanh)
testgrad = Flux.Zygote.gradient((x,p)->sum(ts(x,p)),x,pd)
@test fdgrad[1] ≈ testgrad[1]
@test fdgrad[2] ≈ testgrad[2] rtol=1e-5

fsgrad = Flux.Zygote.gradient((x,p)->sum(fs(x,p)),SVector{2}(x),SVector{75}(pd))
@test fsgrad[1] isa SArray
@test fsgrad[2] isa SArray

layer = FastDense(3, 4, bias=false)
p = initial_params(layer)

rand_loss(p) = begin
    y = sum(abs2, layer(randn(3, 5), p))
end

y, back = Zygote.pullback(rand_loss, p)
grad = back(1.0)[1]
@test length(grad) == length(p)

layer = StaticDense(3, 4, bias=false)
p = initial_params(layer)

rand_loss(p) = begin
    y = sum(abs2, layer(randn(3), p))
end

y, back = Zygote.pullback(rand_loss, p)
grad = back(1.0)[1]
@test length(grad) == length(p)


using BifurcationInference
using ForwardDiff

using LaTeXStrings,Plots
using Plots.Measures

####################################################################################################
####################################################################################################
######################################################################################## saddle-node
include("minimal/saddle-node.jl"); X.roots .= [[[-1.0]]]
plot(F,[5/2,-1],X)
savefig("docs/figures/saddle-node.pdf")

####################################################################################################
####################################################################################################
########################################################################################## pitchfork
include("minimal/pitchfork.jl")
plot(F,[1/2,-1],X)
savefig("docs/figures/pitchfork.pdf")

####################################################################################################
####################################################################################################
################################################################################ bifurcation measure
measure(x) = 1/(1+abs(f(x)/ForwardDiff.derivative(f,x)))

layout = @layout [a;b{1.0w,0.5h}]
default(); default(grid=false,label="",margin=0mm)
plot(layout = layout, link = :x, size=(400,500) )

hline!([0],subplot=2,linewidth=1,ylim=(-5,5),color=:black, xlabel=L"\mathrm{arclength,}s", xmirror=true, topmargin=-5mm,
	ylabel=L"\mathrm{determinant}\,\quad\left|\!\!\!\!\frac{\partial F_{\theta}}{\partial u}\right|")

plot!([-1.52,-1.52],[0,1],subplot=1,linewidth=2,color=:gray)
plot!([-1.52,-1.52],[0,5],subplot=2,linewidth=2,color=:gray)

plot!([1.33,1.33],[0,1],subplot=1,linewidth=2,color=:gray)
plot!([1.33,1.33],[0,5],subplot=2,linewidth=2,color=:gray)

plot!([1.67,1.67],[0,1],subplot=1,linewidth=2,color=:gray)
plot!([1.67,1.67],[0,5],subplot=2,linewidth=2,color=:gray)

p = -2:0.01:2
f(p) = - 2 - (p-1/2)^3 + 3.1(p-1/2)
plot!( p, measure, subplot=1, linewidth=2, color=:red, ylim=(0,1) )
plot!( p, f, subplot=2, linewidth=2, color=:red )

f(p) = - 2 - p^3 + p
plot!( p, measure, subplot=1, linewidth=2, color=:gold, ylabel=L"\mathrm{Measure}\,\,\varphi_{\theta}(s)", ylim=(0,1) )
plot!( p, f, subplot=2, linewidth=2, color=:gold )

yticks!([0,1],subplot=1)
yticks!([-2,0,2],subplot=2)
xticks!([NaN],subplot=1)

scatter!([1.67],[0],subplot=2,marker=:star,color=:black,markersize=7)
scatter!([1.33],[0],subplot=2,marker=:star,color=:black,markersize=7)
scatter!([-1.52],[0],subplot=2,marker=:circle,color=:black,markersize=4)

savefig("docs/figures/bifurcation-measure.pdf")

####################################################################################################
####################################################################################################
############################################################################################ scaling
include("scaling/scaling.jl")
using StatsBase: mean,std
using LaTeXStrings

continuation_mean = Float64[]
continuation_std = Float64[]

forward_mean = Float64[]
forward_std = Float64[]

backward_mean = Float64[]
backward_std = Float64[]

for n ∈	41:100

	ts = [scaling_continuation(n,3) for _ ∈ 1:10]
	push!(continuation_mean,mean(ts))
	push!(continuation_std,std(ts))

	ts = [scaling_forward(n,3) for _ ∈ 1:10]
	push!(forward_mean,mean(ts))
	push!(forward_std,std(ts))

	# ts = [scaling_backward(n,3) for _ ∈ 1:10]
	# push!(backward_mean,mean(ts))
	# push!(backward_std,std(ts))
end

default();default(grid=false,label="", yscale=:log10, xscale=:log10,ms=3, ylim=(8e-2,1e1), msc=:auto, lc=:black,
	xlabel=L"\mathrm{States}\,\,N",ylabel=L"\mathrm{Execution\,\,Time}\,\,/\,\,\mathrm{sec}", size=(400,400),
	linewidth=2, legend=:topleft)

plot( 10:40, x->x^2/180, label=L"N^2", lc=:red )
scatter!( 1:40, backward_mean,yerror=backward_std, color=:black, label="Cost Gradient" )

# plot!( 1:100, x->x^3/205800, lc=:red )
# scatter!( 1:83, forward_mean, yerror=forward_std, color=:gray, label="Cost Function" )

# plot!( 20:60, x->x^2/1580, lc=:red )
# scatter!( 1:83, continuation_mean, yerror=continuation_std, color=:gold, label="Continuation" )

xticks!([1,3,5,10,20,40],["1","3","5","10","20","40"])
yticks!([1e-1,1e0,1e1],["0.1","1","10"])
savefig("docs/figures/scaling.pdf")

####################################################################################################
####################################################################################################
################################################################################ two-state-optima
using BifurcationInference,Plots,StaticArrays
using Clustering,JLD
using StatsBase: median

trajectories, losses, bifurcations, K = Vector{Float64}[], Float64[], Float64[], 0
for file ∈ readdir(join=true,"trajectories")
	data = load(file)

	mask = @. (data["bifurcations"]==2)&(data["losses"]<0.1)
	if ~any(x->any(isnan.(x)),data["trajectory"])

		append!(trajectories,data["trajectory"][mask])
		append!(bifurcations,data["bifurcations"][mask])
		append!(losses,data["losses"][mask])
		K += 1
	end
end

K/length(readdir("trajectories"))
optima = hcat(trajectories...)

@. optima[end,:] = abs(optima[end,:])
M,N = size(optima)
clusters = dbscan(optima,0.5,min_cluster_size=40)

layout = @layout [ a{0.7w} [b{0.5h}; b{0.5h} ] ]
default(); default(msc=:auto,label="",size=(600,450),grid=false)

fig = scatter([0],[0],markersize=0,title=L"\mathrm{Parameter\,\ Estimates}\,\,\theta^{*}", layout=layout, subplot=1,
	ylabel=L"\mathrm{Sensitivity/Degredation\,\,Ratio}\quad k/\mu_1",xlabel=L"\mathrm{Basal\,\,Production\,\,of}\,\,u_1\,,\quad a_1")

color = [:lightblue,:pink]
include("applied/two-state.jl")
for (j,cluster) ∈ enumerate(clusters)
	x,y = 10 .^(getindex(optima,3,cluster.core_indices)), getindex(optima,5,cluster.core_indices) ./ 10 .^getindex(optima,1,cluster.core_indices)
	@assert all( i -> sign(getindex(optima,3,i)) == sign(getindex(optima,4,i)), cluster.core_indices )

	scatter!(x,y,color=color[j],subplot=1,xscale=:log10,xlim=(1e-3,1e3),yscale=:log10,ylim=(1e-1,1e2))
	annotate!(median(x),median(y),j,subplot=1)

	parameters = getindex(optima,:,cluster.core_indices)
	if j == 1 X = StateSpace( 2, 0:0.002:10, [4,5] ) else X = StateSpace( 2, 0:0.01:10, [4,5] ) end

	plot!(F,parameters[:,rand(1:length(cluster.core_indices))],X,
		determinant=false,xlim=(0,10),ylim=(0.01,100),yscale=:log10,
		subplot=j+1, framestyle = :box)
	annotate!(8.8,30,j,subplot=j+1)
end

xlabel!("",subplot=2)
ylabel!("",subplot=2)
ylabel!("",subplot=3)
yticks!([NaN],subplot=2)
yticks!([NaN],subplot=3)

savefig("docs/figures/two-state-optima.pdf")


module MultiFESpacesTests

using Test
using Gridap

model = CartesianDiscreteModel(domain=(0.0,1.0,0.0,1.0), partition=(4,4))

order = 1
diritag = "boundary"
fespace = H1ConformingFESpace(Float64,model,order,diritag)

ufun1(x) = x[1] + x[2]
U1 = TrialFESpace(fespace,ufun1)

ufun2(x) = x[1] + x[2]
U2 = TrialFESpace(fespace,ufun2)

U = [U1,U2]

@test string(U) == "Vector of FESpace objects with 2 fields"

U = MultiFESpace([U1,U2])

@test string(U) == "MultiFESpace object with 2 fields"

@test length(U) == 2
@test U[1] === U1
@test U[2] === U2

V1, state = iterate(U)
@test V1 === U1

V2, state = iterate(U,state)
@test V2 === U2

@test num_free_dofs(U) == num_free_dofs(U1) + num_free_dofs(U2)

a(v,u) = varinner(v,u)

bfun(x) = x[2]
b(v) = varinner(v,CellField(trian,bfun))

trian = Triangulation(model)
quad = CellQuadrature(trian,degree=2)

v1 = CellBasis(V1)
v2 = CellBasis(V2)

u1 = CellBasis(U1)
u2 = CellBasis(U2)

mat11 = integrate(a(v1,u1),trian,quad)
mat12 = integrate(a(v1,u2),trian,quad)
mat22 = integrate(a(v2,u2),trian,quad)

mat = MultiCellMatrix([mat11,mat12,mat22],[(1,1),(1,2),(2,2)])

vec1 = integrate(b(v1),trian,quad)
vec2 = integrate(b(v2),trian,quad)

vec = MultiCellVector([vec1,vec2],[(1,),(2,)])

cellids = IdentityCellNumber(Int,length(vec1))
_vec = apply_constraints(U,vec,cellids)
@test _vec.blocks[1] === vec1
@test _vec.blocks[2] === vec2

_mat = apply_constraints_rows(U,mat,cellids)
@test _mat.blocks[1] === mat11
@test _mat.blocks[2] === mat12
@test _mat.blocks[3] === mat22

_mat = apply_constraints_cols(U,mat,cellids)
@test _mat.blocks[1] === mat11
@test _mat.blocks[2] === mat12
@test _mat.blocks[3] === mat22

dofs = celldofids(U)
@test dofs.blocks[1] === celldofids(U1)
@test dofs.blocks[2] === celldofids(U2)

n = num_free_dofs(U)
n1 = num_free_dofs(U[1])
n2 = num_free_dofs(U[2])
v = rand(n)

f1 = restrict_to_field(U,v,1)
@test f1 == v[1:n1]

f2 = restrict_to_field(U,v,2)
@test f2 == v[1+n1:n]

end # module MultiFESpacesTests


"""
    load_gid_mesh(filename) -> mesh
"""
function load_gid_mesh(fn::AbstractString)
    open(fn, "r") do io
        load_gid_mesh(io)
    end
end

"""
    load_gid_mesh(stream) ->mesh
"""
function load_gid_mesh(file)

    lineCount = countlines(file);
    seekstart(file);
    readline(file); # MESH dimension 3 ElemType Triangle Nnode 3
    readline(file); # Coordinates

    # Read the vertices
    vertices = Vector{Pt{3,Float64}}(undef,lineCount)
    vertexCount = 0;
    while !eof(file)
            line = split(readline(file));
            size(line, 1) == 4 || break; # End Coordinates
            vertexCount += 1;
            #vertices[vertexCount] = Pt{3,Float64}(map(Meta.parse, line[2:4]));
            vertices[vertexCount] = Pt{3,Float64}(parse.(Float64,line[2:4]))
    end

    readline(file); # "\n"
    readline(file); # Elements

    # Read the triangles
    triangles = Vector{Pt{3,Int}}(undef,lineCount-vertexCount);
    triangleCount = 0;
    while !eof(file)
            line = split(readline(file));
            size(line, 1) == 4 || break; # End Elements
            triangleCount += 1;
            triangles[triangleCount] = Pt{3,Int}(map(Meta.parse, line[2:4]));
    end

    Mesh(vertices[1:vertexCount], triangles[1:triangleCount])
end


# Node tests
a = Node("a",Intervals([(0.0,1.0),(1.0,2.0)]))
b = Node("a",Intervals([(0.0,1.0),(1.0,2.0)]))
c = Node("c",Intervals([(0.000001,0.99)]))

@test a==b
@test b==a
@test a==a
@test a!=c
@test b!=c
@test c ⊆ a
@test a ⊆ b
@test a ⊈ c
@test b ⊈ c
@test a ⊆ a
@test (0.5,0.99) ⊆ a
@test (0.6,1.5) ⊆ b
@test (-0.001,0.9) ⊈ c
@test (0.0,0.99) ⊆ c
@test 0.5 ∈ a
@test 0.98 ∈ c
@test 2.5 ∉ a
@test 2.0 ∉ c
@test (a ∩ c).list == [(0.000001,0.99)]
@test (a ∪ c).list == [(0.0,2.0)]

a = Node("a",Intervals([(0.0,1.0),(1.0,2.0)]))
b = Node("b",Intervals([(0.0,1.5),(1.6,2.5)]))
c = Node("c",Intervals([(1.0,1.6)]))
d = Node("d",Intervals([(0.3,1.2)]))
L = Vector([a,b,c])
@test a ∈ L
@test b ∈ L
@test d ∉ L
@test c ∈ L
@test get_idx(a,L) == [1]
@test get_idx(b,L) == [2]
@test get_idx(d,L) == []
@test get_idx("a",L) == [1]
@test get_idx("b",L) == [2]
@test get_idx("d",L) == []
@test get_idx(1.0,L) == [1,2,3]
@test get_idx(1.61,L) == [1,2]
@test get_idx(2.4,L) == [2]
@test get_idx(2.6,L) == []
L2 = Vector([a,b])
@test L2 ⊆ L
@test L ∩ L2 == [a,b]
@test L ∪ L2 == [a,b,c]
L2 = Vector([a,b,d])
@test L2 ⊈ L
@test L ∩ L2 == [a,b]
@test L ∪ L2 == [a,b,c,d]
@test L2 ∪ L == [a,b,d,c]
L2 = Vector([])
@test L2 ⊆ L
@test L ∩ L2 == []
@test L ∪ L2 == L
L2 = [Node("a",Intervals([(0.3,0.9)])),
      Node("x",Intervals([(1.3,1.5)])),
      Node("z",Intervals([(1.2,1.3)]))]
@test all(L2 .⊆ L)
L2 = [Node("a",Intervals([(0.3,0.9)])),
      Node("b",Intervals([(1.3,1.9)])),
      Node("z",Intervals([(1.2,1.3)]))]
@test !all(L2 .⊆ L)

a1 = Node("a",Intervals([(0.5,1.0),(1.5,2.5)]))
a2 = Node("a",Intervals([(0.0,0.45),(0.6,1.0),(1.1,1.4),(2.3,3.0)]))
a3 = merge(a1,a2)
@test a3.name == "a"
@test a3.presence == Intervals([(0.0,0.45),(0.5,1.0),(1.1,1.4),(1.5,3.0)])
merge!(a1,a2)
@test a1.name == "a"
@test a1.presence == Intervals([(0.0,0.45),(0.5,1.0),(1.1,1.4),(1.5,3.0)])


using LinearAlgebra, DiffEqOperators, Random, Test
################################################################################
# Test 2d extension
################################################################################

#Create Array
n = 8
m = 15
A = rand(n,m)

#Create atomic BC
q1 = RobinBC((1.0, 2.0, 3.0), (0.0, -1.0, 2.0), 0.1, 4.0)
q2 = PeriodicBC(Float64)

BCx = vcat(fill(q1, div(m,2)), fill(q2, m-div(m,2)))  #The size of BCx has to be all size components *except* for x
BCy = vcat(fill(q1, div(n,2)), fill(q2, n-div(n,2)))

Qx = MultiDimBC{1}(BCx)
Qy = MultiDimBC{2}(BCy)

Ax = Qx*A
Ay = Qy*A

@test size(Ax)[1] == size(A)[1]+2
@test size(Ay)[2] == size(A)[2]+2

for j in 1:m
    @test Ax[:, j]  == Array(BCx[j]*A[:, j])
end
for i in 1:n
    @test Ay[i,:] == Array(BCy[i]*A[i,:])
end

################################################################################
# Test 3d extension
################################################################################

#Create Array
n = 8
m = 11
o = 12
A = rand(n,m, o)

#Create atomic BC
q1 = RobinBC((1.0, 2.0, 3.0), (0.0, -1.0, 2.0), 0.1, 4.0)
q2 = PeriodicBC(Float64)

BCx = vcat(fill(q1, (div(m,2), o)), fill(q2, (m-div(m,2), o)))  #The size of BCx has to be all size components *except* for x
BCy = vcat(fill(q1, (div(n,2), o)), fill(q2, (n-div(n,2), o)))
BCz = fill(Dirichlet0BC(Float64), (n,m))

Qx = MultiDimBC{1}(BCx)
Qy = MultiDimBC{2}(BCy)
Qz = Dirichlet0BC{3}(Float64, size(A)) #Test the other constructor

Q1 = (Qx+Qy+Qz)
Q2 = compose(Qx,Qy,Qz) #test addition combinations
@test_broken Q1 == Q2 #This fails
@test all(Q1.BCs .== Q2.BCs) # but this passes so it does actually work
@test_broken Qtmp = Qx + Qz
@test_skip Qz+Qx+Qy == Qy+Qtmp
Ax = Qx*A
Ay = Qy*A
Az = Qz*A

@test size(Ax)[1] == size(A)[1]+2
@test size(Ay)[2] == size(A)[2]+2
@test size(Az)[3] == size(A)[3]+2
for j in 1:m, k in 1:o
    @test Ax[:, j, k] == Array(BCx[j, k]*A[:, j, k])
end
for i in 1:n, k in 1:o
    @test Ay[i, :, k] == Array(BCy[i, k]*A[i, :, k])
end
for i in 1:n, j in 1:m
    @test Az[i, j, :] == Array(BCz[i, j]*A[i, j, :])
end

#test compositions to higher dimension
for N in 2:7
    sizes = rand(4:7, N)
    A = rand(sizes...)

    Q1_N = Neumann0BC(Float64, Tuple(ones(N)), 3.0, size(A))

    Q = compose(Q1_N...)

    A1_N = Q1_N.*fill(A, N)

    A_extended = Q*A
    @test Array(A_extended) == Array(compose(A1_N...))
end


test_str=".#.
..#
###"

test_list = split(test_str, '\n')

list = readlines("day17_input.txt")

function cycle3d(list, cycles)
    m = length(list)
    width = m+2*cycles+2
    depth = 1+2*cycles+2
    space = fill(false, width, width, depth) # m + 2*cycles (might grow once on both sides at each cycle) + 2 (padding)
    for i in 1:m
        space[1+cycles+i, 2+cycles:m+1+cycles, 2+cycles] = map(==('#'), collect(list[i]))
    end
    previous_space = copy(space)
    for cycle = 1:cycles
#=          for k in 2:depth-1
            for j in 2:width-1
                for i in 2:width-1 =#
          for k in 2+cycles-cycle:2+cycles+cycle
            for j in 2+cycles-cycle:m+1+cycles+cycle
                for i in 2+cycles-cycle:m+1+cycles+cycle
                    active = sum(previous_space[i-1:i+1,j-1:j+1,k-1:k+1])
                    if space[i,j,k] && active ∉ [3,4]
                        space[i,j,k] = false
                    elseif !space[i,j,k] && active == 3
                        space[i,j,k] = true
                    end
                end
            end
        end
        previous_space = deepcopy(space)
    end
    return count(space)
end

@show cycle3d(test_list, 6) == 112
@show cycle3d(list, 6) == 310

function cycle4d(list, cycles)
    m = length(list)
    width = m+2*cycles+2
    depth = 1+2*cycles+2
    space = fill(false, width, width, depth, depth)
    for i in 1:m
        space[1+cycles+i, 2+cycles:m+1+cycles, 2+cycles, 2+cycles] = map(==('#'), collect(list[i]))
    end
    previous_space = copy(space)
    for cycle = 1:cycles
#=         for l in 2:depth-1
            for k in 2:depth-1
                for j in 2:width-1
                    for i in 2:width-1 =#
        for l in 2+cycles-cycle:2+cycles+cycle
            for k in 2+cycles-cycle:2+cycles+cycle
                for j in 2+cycles-cycle:m+1+cycles+cycle
                    for i in 2+cycles-cycle:m+1+cycles+cycle
                        active = sum(previous_space[i-1:i+1,j-1:j+1,k-1:k+1,l-1:l+1])
                        if space[i,j,k,l] && active ∉ [3,4]
                            space[i,j,k,l] = false
                        elseif !space[i,j,k,l] && active == 3
                            space[i,j,k,l] = true
                        end
                    end
                end
            end
        end
        previous_space = deepcopy(space)
    end
    return count(space)
end

@show cycle4d(test_list, 6) == 848
@show cycle4d(list, 6) == 2056

nothing

#=
```julia
julia> @btime cycle3d(list,6)
  1.858 ms (13496 allocations: 1.50 MiB)

julia> @btime cycle4d(list,6)
  39.148 ms (283334 allocations: 31.09 MiB)
```
 =#

function cycle3d_onepass(list, cycles)
    m = length(list)
    width = m+2*cycles+2
    depth = 3+2*cycles
    space = fill(false, width, width, depth) # m + 2*cycles (might grow once on both sides at each cycle) + 2 (padding)
    for i in 1:m
        space[1+cycles+i, 2+cycles:m+1+cycles, 2+cycles] = map(==('#'), collect(list[i]))
    end
    sumxyz = fill(0, width, width, depth)
    for cycle = 1:cycles
        sumxyz[:,:,2:end-1] = space[:,:,3:end] + space[:,:,2:end-1] + space[:,:,1:end-2]
        sumxyz[:,2:end-1,:] = sumxyz[:,3:end,:] + sumxyz[:,2:end-1,:] + sumxyz[:,1:end-2,:]
        sumxyz[2:end-1,:,:] = sumxyz[3:end,:,:] + sumxyz[2:end-1,:,:] + sumxyz[1:end-2,:,:]
        space = map(x -> (x[1] && x[2] ∉ [3,4]) ? false : (!x[1] && x[2] == 3) ? true : x[1], zip(space, sumxyz))
    end
    return count(space)
end

function cycle4d_onepass(list, cycles)
    m = length(list)
    width = m+2*cycles+2
    depth = 3+2*cycles
    space = fill(false, width, width, depth, depth)
    for i in 1:m
        space[1+cycles+i, 2+cycles:m+1+cycles, 2+cycles, 2+cycles] =
            map(==('#'), collect(list[i]))
    end
    sumxyzw = fill(0, width, width, depth, depth)
    for cycle = 1:cycles
        sumxyzw[:,:,:,2:end-1] = space[:,:,:,3:end] + space[:,:,:,2:end-1] + space[:,:,:,1:end-2]
        sumxyzw[:,:,2:end-1,:] = sumxyzw[:,:,3:end,:] + sumxyzw[:,:,2:end-1,:] + sumxyzw[:,:,1:end-2,:]
        sumxyzw[:,2:end-1,:,:] = sumxyzw[:,3:end,:,:] + sumxyzw[:,2:end-1,:,:] + sumxyzw[:,1:end-2,:,:]
        sumxyzw[2:end-1,:,:,:] = sumxyzw[3:end,:,:,:] + sumxyzw[2:end-1,:,:,:] + sumxyzw[1:end-2,:,:,:]
        space = map(x -> (x[1] && x[2] ∉ [3,4]) ? false : (!x[1] && x[2] == 3) ? true : x[1], zip(space, sumxyzw))
    end
    return count(space)
end

@show cycle3d_onepass(list, 6) == 310
@show cycle4d_onepass(list, 6) == 2056

nothing

#=
#
# It turns out the nested loop is still faster because I only consider the points
# in space next to points which are possibly active. I mean, the range of the loops
# grows with the number of cycles, so I minimize the number of neighborhood counts.
# In the one-pass code, I end up counting everything. Of course I can modify the
# code for the range of the sum to grow with the cycles as well, but it is enough
# for the moment.

```julia
julia> @btime cycle3d(list, 6);
  1.802 ms (14266 allocations: 1.57 MiB)

julia> @btime cycle4d(list, 6);
  40.654 ms (288967 allocations: 31.61 MiB)

julia> @btime cycle3d_onepass(list, 6);
  2.482 ms (921 allocations: 3.00 MiB)

julia> @btime cycle4d_onepass(list, 6);
  63.279 ms (5986 allocations: 61.57 MiB)
```
=#


function recover_dynamics(X, t, svname, grpname, p_ = Float32[1.3, 0.9, 0.8, 1.8])

    @info "Started run $svname $(grpname)."

    # Define the network 2->5->5->5->2
    U = FastChain(
        FastDense(2, 5, rbf),
        FastDense(5, 5, rbf),
        FastDense(5, 5, rbf),
        FastDense(5, 2),
    )
    # Get the initial parameters
    p = initial_params(U)

    # Define the hybrid model
    function ude_dynamics!(du, u, p, t, p_true)
        û = U(u, p) # Network prediction
        du[1] = p_true[1] * u[1] + û[1]
        du[2] = -p_true[4] * u[2] + û[2]
    end

    # Closure with the known parameter
    nn_dynamics!(du, u, p, t) = ude_dynamics!(du, u, p, t, p_)
    # Define the problem
    prob_nn = ODEProblem(nn_dynamics!, X[:, 1], (t[1], t[end]), p)

    ## Function to train the network
    # Define a predictor
    function predict(θ, X = X[:, 1], T = t)
        Array(
            solve(
                prob_nn,
                Vern7(),
                u0 = X,
                p = θ,
                tspan = (T[1], T[end]),
                saveat = T,
                abstol = 1e-6,
                reltol = 1e-6,
                sensealg = ForwardDiffSensitivity(),
            ),
        )
    end

    # Simple L2 loss
    function loss(θ)
        X̂ = predict(θ)
        sum(abs2, X .- X̂) + Float32(1e-4)*sum(abs2, θ)/length(θ)
    end

    # Container to track the losses
    losses = Float32[]

    # Callback to show the loss during training
    callback(θ, l) = begin
        push!(losses, l)
        false
    end

    ## Training

    # First train with ADAM for better convergence -> move the parameters into a
    # favourable starting positing for BFGS
    res1 = DiffEqFlux.sciml_train(
        loss,
        p,
        ADAM(0.1f0),
        cb = callback,
        maxiters = 200,
    )
    println("$svname $(grpname) : Training loss after $(length(losses)) iterations: $(losses[end])")
    # Train with BFGS
    res2 = DiffEqFlux.sciml_train(
        loss,
        res1.minimizer,
        BFGS(initial_stepnorm = 0.01f0),
        cb = callback,
        maxiters = 10000,
    )
    println(
        "$svname $(grpname) : Final training loss after $(length(losses)) iterations: $(losses[end])",
    )

    p_trained = res2.minimizer

    X̂ = predict(p_trained)
    Ŷ = U(X̂, p_trained)

    # Create a Basis
    @variables u[1:2]
    # Generate the basis functions, multivariate polynomials up to deg 5
    # and sine
    b = [polynomial_basis(u, 5); sin.(u)]
    basis = Basis(b, u)

    # Create an optimizer for the SINDy problem
    opt = SR3(Float32(1e-2), Float32(0.1))
    # Create the thresholds which should be used in the search process
    λ = Float32.(exp10.(-7:0.1:7))
    # Target function to choose the results from; x = L0 of coefficients and L2-Error of the model
    g(x) = x[1] < 1 ? Inf : norm(x, 2)

    # Test on uode derivative data
    Ψ = SINDy(X̂, Ŷ, basis, λ,  opt, g = g, maxiter = 50000, normalize = true, denoise = true, convergence_error = Float32(1e-10)) # Succeed

    # Extract the parameter
    p̂ = parameters(Ψ)

    # Just the equations
    b = Basis((u, p, t) -> Ψ(u, ones(Float32, length(p̂)), t), u)

    # Retune for better parameters -> we could also use DiffEqFlux or other parameter estimation tools here.
    Ψf = SINDy(
        X̂,
        Ŷ,
        b,
        STRRidge(0.01f0),
        maxiter = 100,
        convergence_error = Float32(1e-18),
    ) # Succeed

    p̂ = parameters(Ψf)

    @info "$svname $(grpname) : Found equations : $(Ψf.equations.eqs)"
    @info "$svname $(grpname) : Parameter estimation : $(p̂)"

    ## Save the results

    jldopen("$(svname)_recovery_loop.jld2", "a+") do file
        mygroup = JLD2.Group(file, grpname)
        mygroup["X"] = X
        mygroup["t"] = t
        mygroup["initial_parameters"] = p
        mygroup["trained_parameters"] = p_trained
        mygroup["losses"] = losses
        mygroup["result"] = Ψf
    end

    @info "Finished run $svname $(grpname)."

    return
end


const DS = DataStructure101

export Heap, MaxHeap, MinHeap
export level, root, parent, left_child, right_child, heapify!, build!

abstract type Heap end

mutable struct MaxHeap <: Heap
    data::Vector{Int64}
    heapified::Bool

    MaxHeap() = new(Int64[], true)
end

mutable struct MinHeap <: Heap
    data::Vector{Int64}
    heapified::Bool

    MinHeap() = new(Int64[], true)
end

Base.getindex(heap::Heap, i::Int64) = heap.data[i]

Base.setindex!(heap::Heap, v::Int64, i::Int64) = (heap.data[i] = v)

Base.length(heap::Heap) = length(heap.data)

level(heap::Heap, i::Int64) = floor(Int, log2(i)) + 1

root(heap::Heap) = 1

parent(heap::Heap, i::Int64) = (i > length(heap) || i < 1) ? -1 : floor(Int, i/2)

left_child(heap::Heap, i::Int64) = 2 * i

right_child(heap::Heap, i::Int64) = 2 * i + 1

heapified(heap::Heap) = heap.heapified

function tree_repr(heap::Heap, i::Int64, n::Int64; tree_str::String="", level::Int64=0)
    right_index = right_child(heap, i)
    left_index = left_child(heap, i)
    length = Base.length(heap)

    (right_index <= n) && (tree_str = tree_repr(
        heap,
        right_index,
        length,
        tree_str=tree_str,
        level=level+1
    ))

    tree_str = "$(tree_str)\n$("\t"^level)HeapNode($(heap[i]))"

    (left_index <= n) && (tree_str = tree_repr(
        heap,
        left_index,
        length,
        tree_str=tree_str,
        level=level+1
    ))

    return tree_str
end

function Base.show(io::IO, heap::Heap)
    print(io, "Heap(\n$(tree_repr(heap, root(heap), length(heap)))\n)")
end

function heapify!(heap::MaxHeap, i::Int64, n::Int64)
    largest_index = i
    right_index = right_child(heap, i)
    left_index = left_child(heap, i)

    (right_index <= n && heap[i] < heap[right_index]) &&
        (largest_index = right_index)
    (left_index <= n && heap[largest_index] < heap[left_index]) &&
        (largest_index = left_index)

    if largest_index != i
        heap[i], heap[largest_index] = heap[largest_index], heap[i]
        heapify!(heap, largest_index, n)
    end
end

function heapify!(heap::MinHeap, i::Int64, n::Int64)
    smallest_index = i
    right_index = right_child(heap, i)
    left_index = left_child(heap, i)

    (right_index <= n && heap[i] > heap[right_index]) &&
        (smallest_index = right_index)
    (left_index <= n && heap[smallest_index] > heap[left_index]) &&
        (smallest_index = left_index)

    if smallest_index != i
        heap[i], heap[smallest_index] = heap[smallest_index], heap[i]
        heapify!(heap, smallest_index, n)
    end
end

function build!(heap::Heap)
    heapified(heap) && return

    len = length(heap)
    p = parent(heap, len)
    for i = p:-1:1
        heapify!(heap, i, len)
    end

    heap.heapified = true
end

function Base.push!(heap::Heap, v::Int64)
    push!(heap.data, v)
    heap.heapified = false
end

function Base.pop!(heap::Heap)
    !heapified(heap) && throw("Not Heap")

    len = length(heap)
    heap[1], heap[len] = heap[len], heap[1]
    value = pop!(heap.data)
    heapify!(heap, 1, len-1)

    return value
end

function Base.sort!(heap::Heap)
    !heapified(heap) && throw("Not Heap")

    sorted = Int64[]
    len = length(heap)
    for i=1:len
        push!(sorted, pop!(heap))
    end

    return sorted
end


function gen_newton_al_functions(solver::Solver)
    n,m,N = get_sizes(solver)
    p,pI,pE = get_num_constraints(solver)
    p_N,pI_N,pE_N = get_num_terminal_constraints(solver)

    c_function!, c_jacobian!, c_labels, cI!, cE! = generate_constraint_functions(solver.obj)
    f = solver.fd
    s_x = reshape(1:(N*n),n,N)
    s_u = reshape(N*n .+ (1:(m*(N-1))),m,N-1)
    s_λ = (N*n + (N-1)*m) .+ (1:((N-1)*pE + pE_N))
    s_μ = (N*n + (N-1)*m + length(s_λ)) .+ (1:((N-1)*pI + pI_N))
    s_ν = (N*n + (N-1)*m + length(s_λ) + length(s_μ)) .+ (1:N*n)
    s_s = (N*n + (N-1)*m + length(s_λ) + length(s_μ) + N*n) .+ (1:length(s_μ))
    sze = (X=s_x, U=s_u, λ=s_λ, μ=s_μ, ν=s_ν, s=s_s)

    function newton_cost(V::Vector,ρ=1)
        X,U,λ,μ,ν,s = unpackV(V)
        e = cE(V)
        i = cI(V) + 0.5s.^2
        dd = d(V)
        J = cost(solver,X,U) + λ'e + μ'i + ν'dd + 0.5*ρ*(e'e + i'i + dd'dd)
        return J
    end

    function cI(V)
        X,U,λ,μ,ν,s = unpackV(V)
        C = zeros(eltype(V),pI,N-1)
        for k = 1:N-1
            cI!(view(C,1:pI,k),X[:,k],U[:,k])
        end
        CN = zeros(eltype(V),pI_N)
        cI!(CN,X[:,N])
        CI = [vec(C); CN]
        return CI
    end

    function cE(V)
        X,U,λ,μ,ν,s = unpackV(V)
        C = zeros(eltype(V),pE,N-1)
        for k = 1:N-1
            cE!(view(C,1:pE,k),X[:,k],U[:,k])
        end
        CN = zeros(eltype(V),pE_N)
        cE!(CN,X[:,N])
        CE = [vec(C); CN]
        return CE
    end

    function d(V)
        X,U,λ,μ,ν,s = unpackV(V)
        D = zeros(eltype(V),n,N)
        D[:,1] = X[:,1] - solver.obj.x0
        for k = 2:N
            f(view(D,1:n,k),X[:,k-1],U[:,k-1])
            D[:,k] -= X[:,k]
        end
        return vec(D)
    end

    function active_set(V,eps=1e-2)
        X,U,λ,μ,ν,s = unpackV(V)
        a = trues(length(V))
        ci = cI(V)
        c_inds = ci .>= -eps
        a[sze.μ] = c_inds
        a[sze.s] = c_inds
        return a
    end

    function packV(X,U,λ,μ,ν,s)
        V = [vec(X);vec(U);λ;μ;ν;s]
    end

    function unpackV(V)
        X = V[sze.X]
        U = V[sze.U]
        λ = V[sze.λ]
        μ = V[sze.μ]
        ν = V[sze.ν]
        s = V[sze.s]
        return X,U,λ,μ,ν,s
    end
    return newton_cost, packV, unpackV, cI, cE, d, active_set, sze
end

function init_newton_results(results::SolverIterResults,Solver::Solver)
    n,m,N = get_sizes(solver)
    p,pI,pE = get_num_constraints(solver)
    pN,pI_N,pE_N = get_num_terminal_constraints(solver)

    X = to_array(res.X)
    U = to_array(res.U)
    μ = vcat([k < N ? res.λ[k][1:pI] : res.λ[k][1:pI_N] for k = 1:N]...)
    λ = vcat([k < N ? res.λ[k][pI .+ (1:pE)] : res.λ[k][pI_N .+ (1:pE_N)] for k = 1:N]...)
    ν = vec(to_array(res.s))
    s = vcat([k < N ? res.C[k][1:pI] : res.C[k][1:pI_N] for k = 1:N]...)
    s = sqrt.(2*max.(0,-s))
    packV(X,U,λ,μ,ν,s)
end

function line_search(V,H,g,a=trues(length(V)))
    J0 = mycost(V)
    α = 1.
    V_ = copy(V)5.17, 6.2 & A5.1, 6.9 & A5.2, A3.18, A3.20, A11.4, and A13.2.
    delta = (H[a,a]\g[a])
    V_[a] = V[a] - α*delta
    J = mycost(V_)
    iter = 1
    while J > J0
        α /= 2
        V_[a] = V[a] - α*delta
        J = mycost(V_)
        iter += 1
        @show J, α
        iter > 10 ? break : nothing
    end
    @show J0-J
    return V_
end

# Block move
model, obj = TrajectoryOptimization.Dynamics.double_integrator
u_min = -0.2
u_max = 0.2
x_bnd = 2
obj_con = TrajectoryOptimization.ConstrainedObjective(obj, tf=5.0,use_xf_equality_constraint=true,u_min=u_min, u_max=u_max, x_max=x_bnd, x_min=-x_bnd)#, x_max=x_max)

solver = TrajectoryOptimization.Solver(model,obj_con,integration=:rk4,N=5,opts=opts)
Random.seed!(1)
U = rand(solver.model.m, solver.N)

res, stats = TrajectoryOptimization.solve(solver,U)
@assert max_violation(res) < opts.constraint_tolerance

newton_cost, packV, unpackV, cI, cE, d, active_set = gen_newton_al_functions(solver)
max_c(V) = max(maximum(cI(V)),norm(cE(V),Inf),norm(d(V),Inf))
V = init_newton_results(res,solver)

cost(solver,res)
newton_cost(V)
max_violation(res)
max_c(V)
g = ForwardDiff.gradient(newton_cost,V)
H = ForwardDiff.hessian(newton_cost,V)
a = active_set(V,1.0)
rank(H[a,a])
size(H[a,a])
V[a] = V[a] - (H[a,a] + 0I)\g[a]
# V = V - H\g
newton_cost(V)
max_c(V)

# dubinscar
model,obj = Dynamics.dubinscar_parallelpark
solver = Solver(model,obj,N=21)
solver.opts.cost_tolerance = 1e-4
solver.opts.constraint_tolerance = 1e-3
solver.opts.penalty_max = 1e4
n,m,N = get_sizes(solver)
Random.seed!(1)
U0 = rand(m,N-1)
res,stats = solve(solver,U0)
cost(solver,res)
λ_update_default!(res,solver);
update_constraints!(res,solver)
@assert max_violation(res) < solver.opts.constraint_tolerance
ρ = maximum(maximum.(res.μ))
maximum(vcat(res.C...) .* vcat(res.active_set...))

newton_cost, packV, unpackV, cI, cE, d, active_set = gen_newton_al_functions(solver)
mycost(V) = newton_cost(V,1e4)
max_c(V) = max(maximum(cI(V)),norm(cE(V),Inf),norm(d(V),Inf))
# c_all = [cI(V); abs.(cE(V)); abs.(d(V))]
# argmax(c_all)
V = init_newton_results(res,solver)

cost(solver,res)
mycost(V)
max_violation(res)
max_c(V)
g = ForwardDiff.gradient(mycost,V)
H = ForwardDiff.hessian(mycost,V)
a = active_set(V,1e-3)
rank(H)
rank(H[a,a])
size(H[a,a])
minimum(abs.(eigvals(H)))
# V[a] - H[a,a]\g[a]
V1 = line_search(V,H+1000I*0,g,a)
# V = V - H\g
mycost(V1)
max_c(V1)
ForwardDiff.gradient(mycost,V1)
copyto!(V,V1)
max_c(V1)

J = Float64[]
c_max = Float64[]

## quadrotor
model,obj = Dynamics.quadrotor
n = 13
m = 4

# Unconstrained
Q = (1e-3)*Matrix(I,n,n)
R = (1e-2)*Matrix(I,m,m)
# R = (1e-1)*Matrix(I,m,m)
Qf = (1000.0)*Matrix(I,n,n)
tf = 1.0
dt = 0.1

# -initial state
x0 = zeros(n)
x0[1:3] = [0.; 0.; 0.]
q0 = [1.;0.;0.;0.]
x0[4:7] = q0

# -final state
xf = copy(x0)
xf[1:3] = [0.;10.;0.] # xyz position
xf[4:7] = q0
u_max = 10.0
u_min = 0.0

obj_uncon = LQRObjective(Q, R, Qf, tf, x0, xf)
obj_con = TrajectoryOptimization.ConstrainedObjective(obj_uncon,u_min=u_min,u_max=u_max)#,u_min=u_min)
solver = Solver(model,obj_con,N=10)

solver.obj
solver.opts.cost_tolerance = 1e-5
solver.opts.constraint_tolerance = 1e-4
solver.opts.penalty_max = 1e5
n,m,N = get_sizes(solver)
Random.seed!(1)
U = 0.5*9.81/4.0*ones(solver.model.m, solver.N-1)
res,stats = solve(solver,U)
plot(res.U)
plot(res.X)
cost(solver,res)
λ_update_default!(res,solver);
update_constraints!(res,solver)
max_violation(res)
@assert max_violation(res) < solver.opts.constraint_tolerance
ρ = maximum(maximum.(res.μ))
maximum(vcat(res.C...) .* vcat(res.active_set...))

newton_cost, packV, unpackV, cI, cE, d, active_set = gen_newton_al_functions(solver)
mycost(V) = newton_cost(V,ρ)
max_c(V) = max(maximum(cI(V)),norm(cE(V),Inf),norm(d(V),Inf))
# max_c(V) = max(maximum(cE(V)),norm(d(V),Inf))

# c_all = [cI(V); abs.(cE(V)); abs.(d(V))]
# argmax(c_all)
V = init_newton_results(res,solver)

cost(solver,res)
mycost(V)
max_violation(res)
max_c(V)
g = ForwardDiff.gradient(mycost,V)
H = ForwardDiff.hessian(mycost,V)
a = active_set(V,1e-3)
X,U,λ,μ,ν,s = unpackV(V)
13*10 + 4*9 + length(λ) + length(ν)# + 36*2
sum(a)
rank(H)
size(H)
rank(H[a,a])
size(H[a,a])
rank(H[1:309,1:309])
# V[a] - H[a,a]\g[a]
V1 = line_search(V,H,g)
# V = V - H\g
mycost(V1)
max_c(V1)
ForwardDiff.gradient(mycost,V1)
copyto!(V,V1)
max_c(V1)
minimum(abs.(eigvals(H)))
J = Float64[]
c_max = Float64[]
for i = 1:10
    g = ForwardDiff.gradient(mycost,V)
    H = ForwardDiff.hessian(mycost,V)
    a = active_set(V,1e-3)
    V1 = line_search(V,H+0I,g,a)

    push!(J,mycost(V1))
    push!(c_max,max_c(V1))
    push!(grad, norm(ForwardDiff.gradient(mycost,V1)))
    copyto!(V,V1)
end

plot(-diff(J),yscale=:log10)
plot(c_max,yscale=:log10)
plot(grad,yscale=:log10)


struct RevisedAffineForm{N,T<:AbstractFloat}
    c::T
    ϵ::SVector{N,T}
    Δ::T
end

function RevisedAffineForm(
    c::T, ϵ::SVector{N,S}, Δ::R
) where {T<:AbstractFloat,N,S<:AbstractFloat,R<:AbstractFloat}
    Q = promote_type(T, S, R)
    return RevisedAffineForm(convert(Q, c), convert(SVector{N,Q}, ϵ), convert(Q, Δ))
end

function RevisedAffineForm(
    c::Interval{T}, ϵ::SVector{N,Interval{T}}, Δ::T
) where {N,T<:AbstractFloat}
    Δ = Δ ⊕₊ radius(c)
    @inbounds for ϵᵢ in ϵ
        Δ = Δ ⊕₊ radius(ϵᵢ)
    end
    return RevisedAffineForm(mid(c), mid.(ϵ), Δ)
end

RevisedAffineForm(c, ϵ, Δ) = RevisedAffineForm(float(c), SVector(float(ϵ)...), float(Δ))

function RevisedAffineForm(x::T, n::Int) where {T<:Real}
    TS = float(T)
    return RevisedAffineForm(convert(TS, x), SVector{n,TS}(zeros(TS, n)), zero(TS))
end

function RevisedAffineForm(x::Interval{T}, n::Int) where {T<:AbstractFloat}
    c, r = midpoint_radius(x)
    return RevisedAffineForm(c, SVector{n,T}(zeros(T, n)), r)
end

function ±(x::RevisedAffineForm, Δ::Real)
    Δ < 0 && throw(ArgumentError("cumulated error must be non-negative"))
    return RevisedAffineForm(x.c, x.ϵ, x.Δ ⊕₊ float(Δ))
end

############
# PRINTING #
############

function _lower(n::Int)
    idx = ('₀', '₁', '₂', '₃', '₄', '₅', '₆', '₇', '₈', '₉')
    s = ""
    ds = reverse!(digits(n))
    for d in ds
        s *= idx[d + 1]
    end
    return s
end

function Base.show(io::IO, x::RevisedAffineForm{N}) where {N}
    print(io, x.c)
    names = isempty(config[:names]) ? ['ϵ' * _lower(i) for i in 1:N] : config[:names]
    @inbounds for (e, n) in zip(x.ϵ, names)
        iszero(e) && continue
        print(io, e > 0 ? " + " : " - ")
        print(io, abs(e), n)
    end
    x.Δ > 0 && print(io, " ± ", x.Δ)
    return nothing
end

############
# EQUALITY #
############
function ==(x::RevisedAffineForm{N}, y::RevisedAffineForm{N}) where {N}
    return x.c == y.c && x.ϵ == y.ϵ && x.Δ == y.Δ
end


"""
```
pseudo_transient_relaxation(stategrid::StateGrid, value_functions::NTuple{N, T}, Gs::NTuple{N, T},
    μ::AbstractArray{S}, Σ²::AbstractArray{S}, Δ::S;
    uniform::Bool = false, bc::Vector{Tuple{S, S}} = Vector{Tuple{T, T}}(undef, 0),
    Q = nothing, banded::Bool = false) where {S <: Real, T <: AbstractArray{<: Real}, N}

pseudo_transient_relaxation!(stategrid::StateGrid, new_value_functions::NTuple{N, T},
    value_functions::NTuple{N, T}, Gs::NTuple{N, T},
    μ::AbstractArray{S}, Σ²::AbstractArray{S}, Δ::S;
    uniform::Bool = false, bc::Vector{Tuple{S, S}} = Vector{Tuple{S, S}}(undef, 0),
    L₁ = nothing, L₂ = nothing, Q = nothing,
    banded::Bool = false) where {S <: Real, T <: AbstractArray{<: Real}, N}
```

The homogeneous diffusion must take the form

```
dXₜ / Xₜ = μₜ dt + Σₜ dW_t,
```

where `Σₜ` is a matrix and every component of the n-dimensional Brownian motion `Wₜ`
is assumed to be independent of the other components.

For a broad class of problems, the HJB can be represented in the form

```
0 = μᵥ - G(v)
```

where `G(v)` collects all the terms unrelated to `μᵥ` and thus contains the HJB's nonlinearities.
By Ito's lemma, this equation can be rewritten as

```
v G(v) = ∇v⋅(μₜXₜ) + ½ Tr[(ΣₜXₜ)ᵀ * (H∘v) * (ΣₜXₜ)] + ∂v / ∂t.
```

Because this PDE is derived from an HJB,
it is solved backward in time. The pseudo-transient relaxation method solves this problem
progressively by treating the nonlinear term `G(v)` as constant and solving the linear part
via an implicit Euler step.

Let `A` be the matrix approximation of the infinitessimal generator of an Ito diffusion. Then
the implicit Euler step is
```
-∂v / ∂t ≈ (vₙ₊₁ - vₙ) / dt + G(vₙ) vₙ₊₁ = A vₙ₊₁
(I / dt + diag(G(vₙ)) - A) vₙ₊₁   = vₙ / dt
(I + dt * (diag(G(vₙ)) - A)) vₙ₊₁ = vₙ
```
In general, `A` may be affine such that `A v = L v + b`, hence

```
(I + dt * (diag(G(vₙ)) - A)) vₙ₊₁ = (vₙ + b).
```

It is also convenient to transform `dt` from `(0, ∞)` to `(0, 1)` by
`Δ = dt / (1 + dt) ⇒ dt = Δ / (1 - Δ)`, hence

```
(I + Δ / (1 - Δ) * (diag(G(vₙ)) - A)) vₙ₊₁  = (vₙ + b)
(diag(1 - Δ)  + Δ * (diag(G(vₙ)) - A)) vₙ₊₁ = (1 - Δ) (vₙ + b)
```
"""
# TODO: Fix Σ² to take in a matrix and to not assume it has been squared already
function pseudo_transient_relaxation(stategrid::StateGrid,
                                     value_functions::NTuple{N, T}, Gs::NTuple{N, T},
                                     μ::AbstractArray{S}, Σ²::AbstractArray{S}, Δ::S;
                                     uniform::Bool = false, bc::Vector{Tuple{S, S}} = Vector{Tuple{T, T}}(undef, 0),
                                     Q = nothing, banded::Bool = false) where {S <: Real, T <: AbstractArray{<: Real}, N}
    new_value_functions = map(x -> similar(x), value_functions)
    _, _, err = pseudo_transient_relaxation!(stategrid, new_value_functions, value_functions, Gs, μ, Σ², Δ,
                                             uniform = uniform, bc = bc, Q = Q, banded = banded)
    return new_value_functions, err
end
function pseudo_transient_relaxation!(stategrid::StateGrid, new_value_functions::NTuple{N, T},
                                      value_functions::NTuple{N, T}, Gs::NTuple{N, T},
                                      μ::AbstractArray{S}, Σ²::AbstractArray{S}, Δ::S;
                                      uniform::Bool = false, bc::Vector{Tuple{S, S}} = Vector{Tuple{S, S}}(undef, 0),
                                      L₁ = nothing, L₂ = nothing, Q = nothing,
                                      banded::Bool = false) where {S <: Real, T <: AbstractArray{<: Real}, N}
    @assert (ndims(μ) == ndims(Σ²) && length(μ) == length(Σ²))  "The dimensions of `drift` and `volatility` must match."
    @assert length(value_functions) == length(Gs) "For each value function, we need the associated nonlinear component."
    if ndims(μ) == 1
        return _ptr_1D!(stategrid, new_value_functions, value_functions, Gs, μ, Σ², Δ;
                        uniform = uniform, bc = bc, L₁ = L₁, L₂ = L₂, Q = Q, banded = banded)
    else
        error("not implemented yet")
    end
end

function _ptr_1D!(stategrid::StateGrid, new_value_functions::NTuple{N, T},
                  value_functions::NTuple{N, T}, Gs::NTuple{N, T},
                  μ::AbstractVector{S}, Σ²::AbstractVector{S}, Δ::S; uniform::Bool = false,
                  bc::Vector{Tuple{S, S}} = Vector{Tuple{S, S}}(undef, 0), L₁ = nothing, L₂ = nothing,
                  Q = nothing, banded::Bool = false) where {S <: Real, T <: AbstractVector{<: Real}, N}

    if banded
        @warn "The BandedMatrix concretization for DiffEqOperators does not work for GhostDerivativeOperator types yet. The output is an Array."
    end

    # Set up finite differences and construct the infinitessimal generator
    x = values(stategrid.x)[1]
    n = length(stategrid)
    dx = uniform ?  x[2] - x[1] : nonuniform_grid_spacing(stategrid, bc)

    # If no DiffEqOperator objects are passed, then we assume the user wants a BandedMatrix w/first-order upwind first derivatives
    # and second order central difference second derivatives. Reflecting boundaries are assumed.
    if isnothing(Q) && isnothing(L₁) && isnothing(L₂)
        # Set up the implicit time step
        for (nvf, vf, G) in zip(new_value_functions, value_functions, Gs)
            _default_ptr_1D!(nvf, vf, G, x, dx, n, μ, Σ², Δ)
        end
    else
        if banded
            # Calculate number of diagonals for BandedMatrix representation
            L₁_n_offdiags = if isnothing(L₁)
                # first order upwinded by default
                1
            elseif typeof(L₁.parameters)[3] || L₁.approximation_order > 2
                # If upwinded, need diagonals for stencil length (minus 0th diagonal) on both sides of 0th diagonal
                # If CenteredDifference, approximation order > 2, need whole stencil length for boundaries,
                # and subtract 1 to not double count 0th diagonal
                L₁.stencil_length - 1
            elseif L₁.approximation_order == 2 # CenteredDifference, approximation order = 2
                1 # just need 1st upper and lower diagonals, plus main diagonal b/c handling boundary values
            end
            L₂_n_offdiags = if isnothing(L₂) # see CenteredDifference cases of L₁_ndiags
                1
            elseif L₂.approximation_order == 2
                1
            else
                L₂.stencil_length - 1
            end
        end

        if isnothing(Q) # Assume reflecting boundaries,
            # typical case for models where the volatility vanishes at boundaries and drift point "inward"
            Q = RobinBC((0., 1., 0.), (0., 1., 0.), dx)
        end
        if isnothing(L₁)
            # Filling coefficients inside this function can cause undefineds to be used in the concretization, reason unknown
            L₁ = UpwindDifference(1, 1, dx, n, ones(n)) # So just use 1s first
        end
        L₁.coefficients .= μ .* x
        L₁ = sparse(L₁ * Q)
        L₂ = sparse((Σ² .* x.^2 ./ 2.) * (isnothing(L₂) ? CenteredDifference(2, 2, dx, n) : L₂) * Q)

        # Set up the implicit time step
        if banded
            # Currently, BandedMatrix concretization doesn't work for GhostDerivativeOperator,
            # so form the BandedMatrix from the sparse concretization. The hope is that time is
            # saved during the left divide step.
            min_n_offdiags = min(L₁_n_offdiags, L₂_n_offdiags)
            diags_dict = Dict(i => diag(L₁[1], i) + diag(L₂[1], i) for i in 1:min_n_offdiags)
            for i in -min_n_offdiags:-1
                diags_dict[i] = -Δ * (diag(L₁[1], i) + diag(L₂[1], i)) # multiplied by -Δ here to obtain further speed up
            end
            orig_0th_diag = -Δ .* (Array(diag(L₁[1], 0)) + Array(diag(L₂[1], 0))) # Handle 0th diagonal separately b/c will add terms
            diags_dict[0] = similar(orig_0th_diag)                                # Also convert them to Array rather than SparseVector

            for (nvf, vf, G) in zip(new_value_functions, value_functions, Gs)
                diags_dict[0] .= (1 - Δ) .+ Δ .* G + orig_0th_diag # add nonlinear terms to main diagonal
                A = BandedMatrix(diags_dict...)

                # nvf .= A \ (1. - Δ) .* (vf + b)
                ldiv!(nvf, qr(A), (1. - Δ) .* (vf + b)) # Typically faster b/c fewer allocations, qr also faster than factorize/cholesky
            end
        else
            A = L₁[1] + L₂[1]
            b = L₁[2] + L₂[2]

            for (nvf, vf, G) in zip(new_value_functions, value_functions, Gs)
                nvf .= (Diagonal((1 - Δ) .+ Δ .* G) - (Δ .* A)) \ ((1 - Δ) .* (vf + b))
            end
        end
    end

    return new_value_functions, value_functions
end

function _default_ptr_1D!(nvf::T, vf::T, G::T, x::AbstractVector{S},
                          dx::AbstractVector{S}, n::Int, μ::AbstractVector{S}, Σ²::AbstractVector{S},
                          Δ::S) where {N <: Int, T <: AbstractVector{<: Real}, S <: Real}

    # Process Σ² matrix
    Σ²in = similar(Σ²)
    Σ²in[2:n - 1] .= Σ²[2:n - 1] .* x[2:n - 1].^2 ./ 2. # In non-uniform case, we utilize Fornberg weights, so just calculate
    Σ²in[1] = 0.                                        # the raw coeficients.
    Σ²in[n] = 0.

    # Populate diagonals and construct BandedMatrix
    DD, D0, DU = _centered_difference_reflecting_bc_weights(2, x, -Δ .* Σ²in) # Initialize diagonals for centered difference operator
    μx         = μ .* x
    DU       .-= Δ .* (max.( μx[1:n - 1], 0.) ./ dx[1:n - 1]) # up diagonal
    DD       .-= Δ .* (max.(-μx[2:n],     0.) ./ dx[2:n])     # down diagonal, note FD scheme makes DD negative, multiplied by negative drift ⇒ positive, then subtracted ⇒ negative

    D0          .+= (1 - Δ) .+ Δ .* G
    D0[1:n - 1] .+= Δ .* (max.( μx[1:n - 1], 0.) ./ dx[1:n - 1]) # Also have to add to the main diagonal
    D0[2:n]     .+= Δ .* (max.(-μx[2:n],     0.) ./ dx[2:n])     # for the first-order FD operator

    A             = BandedMatrix(0 => D0, 1 => DU, -1 => DD)

    # Solve linear system
    # nvf .= A \ ((1. - Δ) .* vf)
    ldiv!(nvf, qr(A), (1. - Δ) .* vf) # Typically faster b/c fewer allocations, qr also faster than factorize/cholesky
end

function _default_ptr_1D!(nvf::T, vf::T, G::T, x::AbstractVector{S},
                          dx::S, n::Int, μ::AbstractVector{S}, Σ²::AbstractVector{S},
                          Δ::S) where {N <: Int, T <: AbstractVector{<: Real}, S <: Real}

    # Process Σ² matrix
    Σ²in = similar(Σ²)                                           # In non-uniform case, for 2:n - 1, want to divide by the
    Σ²in[2:n - 1] .= Σ²[2:n - 1] .* x[2:n - 1].^2 ./ (2 .* dx^2) # forward and backward difference, but for the boundaries,
    Σ²in[1] = 0.                                                 # we impose reflecting boundaries.
    Σ²in[n] = 0.                                                 # Note that dx is length (n + 1).

    # Populate diagonals and construct BandedMatrix
    μx = μ .* x
    DU = -Δ .* (max.( μx[1:n - 1], 0.) ./ dx + Σ²in[1:n - 1]) # up diagonal
    DD = -Δ .* (max.(-μx[2:n],     0.) ./ dx + Σ²in[2:n])     # down diagonal, note FD scheme makes DD negative, multiplied by negative drift ⇒ positive, then subtracted ⇒ negative

    D0            = (1 - Δ) .+ Δ .* G
    D0[1:n - 1] .-= DU
    D0[2:n]     .-= DD

    A = BandedMatrix(0 => D0, 1 => DU, -1 => DD)

    # Solve linear system
    # nvf .= A \ ((1. - Δ) .* vf)
    ldiv!(nvf, qr(A), (1. - Δ) .* vf) # Typically faster b/c fewer allocations, qr also faster than factorize/cholesky
end


import Polynomials.hasneg
import Polynomials.isneg

export AbstractLagrangePolynomial
export LagrangePoly
export LGRPoly
export derivmatrix
export lgr_points
export lgr_weights

abstract type AbstractLagrangePolynomial{T<:Number}<:AbstractPolynomial{T} end

coeffs(p::AbstractLagrangePolynomial) = p.y # temp fix, make this work with methods that require coefficients

convert(P::Type{<:AbstractLagrangePolynomial}, p::AbstractLagrangePolynomial) where {T} = P(p.x, p.y, domain(p), p.var)

Base.eachindex(p::AbstractLagrangePolynomial) = 1:length(p)

Base.getindex(p::AbstractLagrangePolynomial, i::Int) = (p.x[i],p.y[i])
Base.setindex!(p::AbstractLagrangePolynomial, y::Number, i::Int) = p.y[i] = y # add case for updating x and y
Base.firstindex(p::AbstractLagrangePolynomial) = 1
Base.lastindex(p::AbstractLagrangePolynomial) = length(p.x)

Base.length(p::AbstractLagrangePolynomial) = length(p.x)

function (lp::AbstractLagrangePolynomial{T})(x::S) where {T,S<:Number}
    if x ∉ domain(lp)
        pDomain = domain(lp)
        throw(DomainError(x,"outside of domain $pDomain"))
    end
    return lagrange_eval_weights(lp,x,lp.y)
end

function showterm(io::IO, ::Type{<:AbstractLagrangePolynomial{T}}, pj::Any, var, j, first::Bool, mimetype) where {N, T}
    !first &&  print(io, " ")
    print(io, "$(pj)")
    return true
end

mutable struct LagrangePoly{T<:Number} <: AbstractLagrangePolynomial{T} # consider making this immutable for performance
    x::Vector{T}
    y::Vector{T}
    weights::Vector{T}
    domain::Interval{T}
    var::Symbol
    function LagrangePoly{T}(x::Vector{T}, y::Vector{T}, domain::Interval{T}, var::Symbol) where {T <: Number}
        return new{T}(x,y,lagrange_bary_weights(x),domain, var)
    end
end

@register LagrangePoly

function LagrangePoly(x::Vector{T}, y::Vector{T}, lower=min(x...), upper=max(x...); var::Symbol=:x, containlower::Bool=true, containupper::Bool=true) where {T}
    lower = convert(T, lower)
    upper = convert(T, upper)
    return LagrangePoly{T}(x,y,Interval(lower, upper, containlower, containupper),var)
end

function lagrange_bary_weights(x::AbstractVector{T}) where {T}
    numPoints = length(x)
    w = ones(T,numPoints)
    for j in 2:numPoints
        @inbounds xj = x[j]
        for k in 1:j-1
            @inbounds w[k] *= x[k] - xj
        end
        for k in 1:j-1
            @inbounds w[j] *= xj - x[k]
        end
    end
    return 1 ./ w
end

function lagrange_eval_weights(p::AbstractLagrangePolynomial{T}, xeval, y) where {T}
    x = p.x
    w = p.weights
    num = zero(T)
    denom = zero(T)
    for j = 1:length(x)
        @inbounds wx = w[j] / (xeval - x[j])
        isinf(wx) && return y[j]
        @inbounds num += wx*y[j]
        denom +=wx
    end
    return num/denom
end

domain(p::AbstractLagrangePolynomial) = p.domain

function fit(P::Type{<:AbstractLagrangePolynomial}, x::AbstractVector{T}, y::AbstractVector{T}; kwargs...) where {T} # make this fit the Polynomials.jl version of this command
    return LagrangePoly(x,y; kwargs...)
end

function fit(P::Type{<:AbstractLagrangePolynomial}, x::AbstractVector, yFunc::Function; kwargs...)
    return LagrangePoly(x, map(x -> yFunc(x), x); kwargs...)
end

function derivmatrix(p::AbstractLagrangePolynomial{T}) where {T}
    w = p.weights
    x = p.x
    numPoints = length(x)
    D = zeros(T, numPoints, numPoints)
    for i in 1:numPoints
        diag = zero(T)
        xi = x[i]
        wi = w[i]
        for j = 1:i-1
            @inbounds Dij = w[j] / ( wi * (xi - x[j]) )
            diag -= Dij
            @inbounds D[i,j] = Dij
        end
        for j = i+1:numPoints
            @inbounds Dij = w[j] / ( wi * (xi - x[j]) )
            diag -= Dij
            @inbounds D[i,j] = Dij
        end
        @inbounds D[i,i] = diag
    end
    return D
end

function derivative(p::AbstractLagrangePolynomial{T}, order::Integer = 1) where {T}
    pder = deepcopy(p)
    pder.y = derivmatrix(p) * p.y
    return order == 1 ? pder : derivative(pder, order - 1) # make this better by numerically calculating higher derivatives, use the Klein and Berrut method in that book you found
end

# Legendre-Gauss-Radau

# todo precompile first few (hudred)? weights (probs doesn't matter)

mutable struct LGRPoly{T<:Number} <: AbstractLagrangePolynomial{T}
    x::Vector{T}
    y::Vector{T}
    weights::Vector{T}
    var::Symbol
    function LGRPoly{T}(y::Vector{T}, var::Symbol) where {T <: Number}
        lgrpoints = lgr_points(length(y))
        return new{T}(lgrpoints,y,lagrange_bary_weights(lgrpoints), var)
    end
end
@register LGRPoly
LGRPoly(y::AbstractVector{T}; var::SymbolLike = :x) where {T} = LGRPoly{T}(y, Symbol(var)) # add function where flipped = true

function LGRPoly(yFunc::Function, numPoints::Int; kwargs...)
    y = map(yFunc, lgr_points(numPoints))
    LGRPoly(y; kwargs...)
end

function lgr_points(numPoints::Int) # find when this becomes slower than fast gaussradau and use that instead
    numPoints<2 && error("number of points must be greater than one")
    N = 1:numPoints-2
    a = zeros(numPoints-1)
    b = zeros(numPoints-2)
    a[1] = 1 / 3
    a[2:end] = map((N)-> 1 / (4*N^2 + 8N +3),N)
    b = map((N)-> (N^2 + N)^0.5 / (2N+1),N)
    J = SymTridiagonal(a,b)
    return pushfirst!(eigvals(J),-1)
end

domain(P::Type{<:LGRPoly}) = Interval(-1,1)
domain(p::LGRPoly{T}) where {T} = Interval{T}(-1,1)

fit(P::Type{<:AbstractLagrangePolynomial}, y::AbstractVector; kwargs...) = LGRPoly(y;kwargs...)
fit(P::Type{<:AbstractLagrangePolynomial}, yFunc::Function, numPoints::Int; kwargs...) = LGRPoly(yFunc, numPoints;kwargs...)

Base.convert(P::Type{<:LGRPoly}, p::LGRPoly) where {T} = P(p.y, p.var)

function lgr_weights(p::LGRPoly{T}) where {T} # definitely precompute these
    xlen = length(p)
    return gaussradau(xlen)[2]
end

function lgr_weights(numPoints::Int) # definitely precompute these
    return gaussradau(numPoints)[2]
end

function integrate(p::LGRPoly)
    return sum(lgr_weights(p).*p.y)
end


F, z = CyclotomicField(2)
@test Hecke.force_coerce_cyclo(F, z) == z


logbesseli(ν, x) = log(besselix(ν * one(real(x)), x)) + abs(real(x))

# hypergeometric functions with matrix arguments
function pFq end

pFq(::Tuple{}, ::Tuple{}, x) = exp(tr(x)) # ₀F₀
# TODO: Implement ₀F₁
# see P Koev, A Edelman. The efficient evaluation of the hypergeometric function of a
# matrix argument. Math. Comp. 75 (2006), 833-846
# pFq(a::NTuple{p}, b::NTuple{q}, x)

logpFq(a, b, z) = log(pFq(a, b, z))


"""
    randomcenters(X, k) -> Vector{Float}

    Return a vector of `k` cluster centers chosen from `X` data matrix.
"""
function randomcenters(X::AbstractMatrix{<:Real}, k::Integer;)
    d, n = size(X)
    i = sample(1:size(X, 2), k)
    return X[:, i]
end

function kmppcenters(
    X::AbstractMatrix{<:Real},
    k::Integer;
    metric::PreMetric = SqEuclidean(),
)
    n = size(X, 2)

    D = eltype(X)
    iseeds = Vector{Int}(undef, k)
    # randomly pick the first center
    p = rand(1:n)
    iseeds[1] = p

    if k > 1
        mincosts = colwise(metric, X, view(X, :, p))
        mincosts[p] = 0

        # pick remaining (with a chance proportional to mincosts)
        tmpcosts = zeros(n)
        for j = 2:k
            p = wsample(1:n, mincosts)
            iseeds[j] = p

            # update mincosts
            c = view(X, :, p)
            colwise!(tmpcosts, metric, X, view(X, :, p))
            updatemin!(mincosts, tmpcosts)
            mincosts[p] = 0
        end
    end

    return X[:, iseeds]
end

"""
    isodata_s(X, k) -> Vector{Float}

    Return a vector of `k` cluster centers calculated from `X` data matrix.
    Starting with one cluster center, this function assigns the data to the clusters,
    then splits the clusters until there are `k` clusters.
"""
function isodata_s(X, k)
    d, n = size(X)
    assignments = Vector{Int}(undef, n)
    centers = randomcenters(X, 1)

    dist = pairwise(Euclidean(), centers, X, dims = 2)
    D = eltype(dist)

    costs = Vector{D}(undef, n)
    counts = Vector{Int}(undef, k)

    while true
        # compute number of clusters
        NROWS = size(centers, 2)

        # assign cluster members
        counts = Vector{Int}(undef, NROWS)
        dist = pairwise(Euclidean(), centers, X, dims = 2)

        update_assignments!(dist, assignments, costs, counts)
        update_centers!(X, assignments, centers, counts)

        # check number of centroids
        if NROWS == k
            return centers
        elseif NROWS > k
            while NROWS > k
                # discard small clusters (fewest assignments), ties broken arbitrarily
                to_delete = findfirst(x -> x ∈ sort(counts)[1:NROWS-k], counts)
                centers = discard([to_delete], centers)
                NROWS = size(centers, 2)

                counts = Vector{Int}(undef, NROWS)
                dist = pairwise(Euclidean(), centers, X, dims = 2)
                update_assignments!(dist, assignments, costs, counts)
                update_centers!(X, assignments, centers, counts)
            end
            return centers
        end

        # split clusters
        sd = Array{D}(undef, d, NROWS)
        update_sd!(X, assignments, centers, counts, sd)

        centers = split!(collect(1:NROWS), centers, sd)
    end

    return centers
end

function updatemin!(r::AbstractArray, x::AbstractArray)
    n = length(r)
    length(x) == n || throw(DimensionMismatch("Inconsistent array lengths."))
    @inbounds for i = 1:n
        xi = x[i]
        if xi < r[i]
            r[i] = xi
        end
    end
    return r
end


@diffusion_process FitzHughDiffusionConjugAuxSimple begin
    :dimensions
    process --> 2
    wiener --> 1

    :parameters
    (ϵ, s, γ, β, σ, t, T) --> Float64
    (u, v) --> ℝ{2}

    :additional
    constdiff --> true
    linear --> true
end

B(t, P::FitzHughDiffusionConjugAuxSimple) = @SMatrix [0.0  1.0; 0.0 0.0]

β(t, P::FitzHughDiffusionConjugAuxSimple) = ℝ{2}(0.0, 0.0)

σ(t, P::FitzHughDiffusionAux) = ℝ{2}(0.0, P.σ)