"""
    ngcd(p,q, [k]; kwargs...)

Find numerical GCD of polynomials `p` and `q`. Refer to [`NGCD.ngcd`](@ref) for details.

In the case `degree(p) ≫ degree(q)`,  a heuristic is employed to first call one step of the Euclidean gcd approach, and then call `ngcd` with relaxed tolerances.

"""
function ngcd(p::P, q::Q, args...;kwargs...) where {T, S, P<:StandardBasisPolynomial{T}, Q <: StandardBasisPolynomial{S}}

    degree(p) < 0  && return (u=q,      v=p, w=one(q),  θ=NaN, κ=NaN)
    degree(p) == 0 && return (u=one(q), v=p, w=zero(q), θ=NaN, κ=NaN)
    degree(q) < 0  && return (u=one(q), v=p, w=zero(q), θ=NaN, κ=NaN)
    degree(q) == 0 && return (u=one(p), v=p, w=q,       θ=NaN, κ=NaN)
    assert_same_variable(p,q)

    p′,q′ = promote(p,q)

    if degree(p′) > 5*degree(q′)
        return  ngcd′(p′, q′, args...; kwargs...)
    end

    ps, qs = [p′[i] for i in eachindex(p′)], [q′[i] for i in eachindex(q′)] # need vectors, want copy
    u,v,w,Θ,κ = NGCD.ngcd(ps, qs, args...; kwargs...)
    P′ = ⟒(typeof(p′))
    (u=P′(u, indeterminate(p)), v=P′(v, indeterminate(p)), w = P′(w, indeterminate(p)), θ=Θ, κ=κ)

end

"""
    ngcd′(p,q)

When degree(p) ≫ degree(q), this uses a early call to `divrem` to bring about commensurate degrees
before calling `ngcd`.
"""
function ngcd′(p::StandardBasisPolynomial{T}, q::StandardBasisPolynomial{T};
               atol = eps(real(float(T))),
               rtol = atol,
               satol= atol,
               srtol= rtol,
               kwargs...
               ) where {T}

    a, b = divrem(p,q)

    # check if a=u (p,q) ≈ (aq,q)
    if isapprox(p, a*q, atol=atol, rtol=rtol)
        return (u=a, v=p, w=q, θ=NaN, κ=NaN)
    else
        ngcd(q, b; atol=100atol, rtol=100rtol, kwargs...)
    end
end

module NGCD
using LinearAlgebra

"""
    ngcd(ps::Vector{T}, qs::Vector{T}, [k::Int]; scale::Bool=true, atol=eps(T), rtol=eps(T), satol=atol, srtol=rtol)

Return `u, v, w, Θ, κ` where ``u⋅v ≈ ps`` and ``u⋅w ≈ qs`` (polynomial multiplication); `Θ` (`\\Theta[tab]`) is the residual error (``‖ [u⋅v,u⋅w] - [ps,qs] ‖``); and `κ` (`\\kappa[tab]`) is the numerical gcd condition number estimate. When `scale=true` (the default), ``u⋅v ≈ ps/‖ps‖`` and ``u⋅w ≈ qs/‖qs‖``

The numerical GCD problem is defined in [1] (5.4). Let ``(p,q)`` be a
polynomial pair with degree m,n. Let Ρmn be set of all such pairs. Any
given pair of polynomials has an exact greatest common divisor, ``u``, of
degree ``k``, defined up to constant factors. Let ``Ρᵏmn`` be the manifold of
all such ``(p,q)`` pairs with exact gcd of degree ``k``. A given pair ``(p,q)`` with exact gcd of degree ``j`` will
have some distance ``Θᵏ`` from ``Pᵏ``.  For a given threshold ``ϵ > 0`` a numerical GCD
of ``(p,q)`` within ``ϵ`` is an exact GCD of a pair ``(p̂,q̂)`` in ``Ρᵏ`` with

``‖ (p,q) - (p̂,q̂) ‖ <= Θᵏ``, where ``k`` is the largest value for which ``Θᵏ < ϵ``.

(In the ``ϵ → 0`` limit, this would be the exact GCD.)

Suppose ``(p,q)`` is an ``ϵ`` pertubation from ``(p̂,q̂)`` where ``(p̂,q̂)`` has an exact gcd of degree ``k``, then ``degree(gcdₑ(p,q)) = k``; as ``ϵ → 0``, ``gcdₑ(p,q) → gcd(p̂, q̂)``, and

``limsup_{(p,q)→(p̂,q̂)} inf{ ‖ (u,v,w) - (û,v̂,ŵ) ‖} / ‖ (p,q) - (p̂,q̂) ‖ < κₑ(p,q)``.

``κ`` is called the numerical GCD condition number.

The Zeng algorithm proposes a degree for ``u`` and *if* a triple ``(u,v,w)`` with ``u`` of degree ``k`` and ``(u⋅v, u⋅w)`` in ``Ρᵏmn`` can be found satisfying ``‖ (u⋅v, u⋅w) - (p,q) ‖ < ϵ`` then ``(u,v,w)`` is returned; otherwise the proposed degree is reduced and the process repeats. If not terminated, at degree ``0`` a constant gcd is returned.

The initial proposed degree is the first ``j``,  ``j=min(m,n):-1:1``, where ``Sⱼ`` is believed to have a singular value of ``0`` (``Sⱼ`` being related to the Sylvester matrix of `ps` and `qs`). The verification of the proposed degree is done using a Gauss-Newton iteration scheme holding the degree of ``u`` constant.

## Scaling:

If `scale=true` (the default when a polynomial norm is large), the gcd of ``p/‖p‖`` and ``q/‖q‖`` is identified. Scaling can reduce the condition numbers significantly.

## Tolerances:

There are two places where tolerances are utilized:

* in the identification of the rank of ``Sⱼ`` a value ``σ₁ = ‖Rx‖`` is identified. To test if this is zero a tolerance of `max(satol, ‖R‖ₒₒ ⋅ srtol)` is used.

* to test if ``(u ⋅ v, u ⋅ w) ≈ (p,q)`` a tolerance of `max(atol, λ⋅rtol)` is used with `λ` chosen to be  ``(‖(p,q)‖⋅n⋅m)⋅κ′⋅‖A‖ₒₒ`` to reflect the scale of ``p`` and ``q`` and an estimate for the condition number of ``A`` (a Jacobian involved in the solution).

This seems to work well for a reasaonable range of polynomials, however there can be issues: when the degree of ``p`` is much larger than the degree of ``q``, these choices can fail; when a higher rank is proposed, then too large a tolerance for `rtol` or `atol` can lead to a false verification; when a tolerance for `atol` or `rtol` is too strict, then a degree may not be verified.

!!! note:
    These tolerances are adjusted from those proposed in [1].

## Specified degree:

When `k` is specified, a value for ``(u,v,w)`` is identified with ``degree(u)=k``. No tolerances are utilized in computing ``Θᵏ``.

Output:

The function outputs a named tuple with names (`u`, `v`, `w`, `Θ`, `κ`). The components `u`,`v`,`w` estimate the gcd and give the divisors. The value `Θ` estimates ``Θᵏ`` and `κ` estimates the numerical condition number.

Example:

```
using Polynomials
x = variable(Polynomial{Float64})
p = (x+10)*(x^9 + x^8/3 + 1)
q = (x+10)*(x^9 + x^8/7 - 6/7)
gcd(p,q) # u a constant
gcd(p,q, method=:numerical)  # u a degree 1 polynomial
Polynomials.NGCD.ngcd(coeffs(p), coeffs(q), verbose=true) # to see some computations
```

Reference:

[1] The Numerical Greatest Common Divisor of Univariate Polynomials
by Zhonggang Zeng;
[url](http://homepages.neiu.edu/~zzeng/uvgcd.pdf);
[doi](https://doi.org/10.1090/conm/556/11014)

Note: Based on work by Andreas Varga; Requires `VERSION >= v"1.2"`.

"""
function ngcd(ps::Vector{T},
              qs::Vector{T};
              scale::Bool=false, #norm(ps) > 1e6, #false,
              atol = eps(real(T)),
              rtol = Base.rtoldefault(real(T)),
              satol = eps(real(T))^(5/6),
              srtol = eps(real(T)),
              verbose=false,
              minⱼ = -1
              ) where {T <: AbstractFloat}

    m, n = _degree.((ps, qs))

    if m < n
        out = ngcd(qs, ps; atol=atol, rtol=rtol, satol=satol, srtol=srtol, verbose=verbose)
        return (u=out.u, v=out.w, w=out.v, Θ=out.Θ, κ=out.κ)
    end

    if scale
        ps = ps./norm(ps)
        qs = qs./norm(qs)
    end

    # storage
    A0 = zeros(T, m+1, 2)
    A0[:,1] = ps
    A0[end-n:end,2] = qs

    # pre-allocate storage for Sylvester Matrices, S₁, S₂...
    Q = zeros(T, m + n, m + n)
    R = zeros(T, m + n, m + n)
    Sₓ = hcat(convmtx(ps,1),  convmtx(qs, m-n+1))
    local x::Vector{T}

    j = n  # We count down Sn, S_{n-1}, ..., S₂, S₁

    F = qr(Sₓ)
    nr, nc = size(Sₓ) # m+1, m-n+2
    Q[1:nr, 1:nr] = F.Q
    R[1:nc, 1:nc] = F.R

    while true

        V = view(R, 1:nc, 1:nc)
        flag, σ, x = smallest_singular_value(V, satol *  sqrt(1 + m - j), srtol)
        verbose && println("------ degree $j ----- σ₁: $σ  --- $flag")

        if (flag == :iszero || flag == :ispossible)

            u, v, w = initial_uvw(Val(flag), j, ps, qs, x)
            flag, ρ₁, σ₂, ρ = refine_uvw!(u,v,w, ps, qs, atol, rtol)

            verbose && println("   --- Θᵏ: $ρ₁ --- $flag (ρ=$(ρ))")

            if flag == :convergence
                return (u=u, v=v, w=w, Θ=ρ₁, κ=σ₂) # (u,v,w) verified
            end
        end

        # reduce possible degree of u and try again with Sⱼ₋₁
        # unless we hit specified minimum, in which case return it
        if j == minⱼ
            u, v, w = initial_uvw(Val(:ispossible), j, ps, qs, x)
            flag, ρ₁, σ₂, ρ = refine_uvw!(u,v,w, ps, qs, atol, rtol)
            return (u=u, v=v, w=w, Θ=ρ₁, κ=σ₂)
        end

        j -= 1
        nr += 1
        nc += 2
        nc > nr && break

        extend_QR!(Q,R, nr, nc, A0) # before Q⋅R = Sⱼ, now Q⋅R = Sⱼ₋₁

    end

    # u is a constant
    verbose && println("------ GCD is constant ------")
    u, v, w = initial_uvw(Val(:constant), j, ps, qs, x)
    flag, ρ₁, κ, ρ = refine_uvw!(u,v,w, ps, qs, atol, rtol)
    return (u=u, v=v, w=w, Θ=ρ₁, κ=κ)

end

function ngcd(ps::Vector{T}, qs::Vector{S}; kwargs...) where {T, S}
    ps′,qs′ = promote(float.(ps), float(qs))
    ngcd(ps′, qs′; kwargs...)
end

# fix the degree, k
function ngcd(ps::Vector{T},
              qs::Vector{T},
              k::Int;
              kwargs...
              ) where {T <: AbstractFloat}

    m, n = _degree.((ps,qs))

    if m < n
        out = ngcd(qs, ps, k, atol=Inf, rtol=Inf)
        return (u=out.u, v=out.w, w=out.v, Θ=out.Θ, κ=out.κ)
    end

    #    u,v,w = initial_uvw(Val(:iszero), k, ps, qs, nothing)
    Sⱼ = [convmtx(ps, n-k+1) convmtx(qs, m-k+1)]
    F = qr(Sⱼ)
    flag, σ, x = smallest_singular_value(F.R, eps(T) *  sqrt(1 + m - k), eps(T))
    if flag != :iszero
        w, v = x[1:n-k+1], -x[n-k+2:end]
        u  = [convmtx(v,k+1); convmtx(w, k+1)] \ [ps; qs]
    else
        u,v,w = initial_uvw(Val(flag), k, ps, qs, nothing)
    end
    flag, ρ₁, κ, ρ = refine_uvw!(u,v,w, ps, qs, Inf, Inf)
    return (u=u, v=v, w=w, Θ=ρ₁, κ=κ)

end

function ngcd(ps::Vector{T}, qs::Vector{S}, k::Int; kwargs...) where {T, S}
    ps′,qs′ = promote(float.(ps), float(qs))
    ngcd(ps′, qs′, k; kwargs...)
end

## -----

# return guess at smallest singular value and right sinuglar value, x
# for an upper triangular matrix, V
function smallest_singular_value(V::AbstractArray{T,2},
                                 atol=eps(T),
                                 rtol=zero(T)) where {T}

    R = UpperTriangular(V)
    k = size(R)[1]/2
    if iszero(det(R))
        return (:iszero, zero(T), T[])
    end

    m,n = size(R)

    # we are testing if ‖Ax‖ ≈ 0
    # If x is a perfect 0, but x is approximated by x' due to round off
    # then ‖A(x-x')‖ <= ‖A‖⋅‖x - x'‖ so we use ‖A‖ as scale factor
    δ = max(atol,  norm(R,Inf) * rtol)

    x = ones(T, n)
    y = zeros(T, m)
    σ₀ = σ₁ = Inf*one(real(T))
    steps, minᵢ = 1, 5

    while true
        y .= R' \ x # use iteration to converge on singular value
        x .= R  \ y
        x ./= norm(x,2)
        σ₁ = norm(R * x, 2)

        if (steps <= 50) && (steps <= minᵢ || σ₁ < 0.05*σ₀) # decreasing, keep going
            σ₀ = σ₁
        else
            break
        end

        steps += 1
    end

    if σ₁ < δ
        return (:ispossible, σ₁, x)
    else
        return (:constant, σ₁, x)
    end

end

# extend QR to next size
# Q gets a 1 in nc,nc, 0s should be elswhere
function extend_QR!(Q,R, nr, nc, A0)

    #old Q is m x m, old R is n x n; we add to these
    n = nc-2
    m = nr - 1
    k,l = size(A0)

    # add two columns to R
    # need to apply Q to top part of new columns
    R[nr-k+1:nr, (nc-1):nc] = A0
    R[1:nr-1, (nc-1):nc] = (view(Q, 1:nr-1, 1:nr-1))' *  R[1:nr-1, (nc-1):nc]

    # extend Q with row and column with identity
    Q[nr,nr] = 1

    # Make R upper triangular using Givens rotations
    for j in nr-1:-1:nc-1
        gj,_ = givens(R[j,nc-1], R[j+1,nc-1], j, j+1)
        rmul!(Q, gj')
        lmul!(gj, R)
    end

    for j in nr-1:-1:nc
        gj,_ = givens(R[j,nc], R[j+1,nc], j, j+1)
        rmul!(Q, gj')
        lmul!(gj, R)
    end

    return nothing

end

## --------------------------------------------------
## Refine u,v,w

## Find u₀,v₀,w₀ from right singular vector
function initial_uvw(::Val{:ispossible}, j, ps, qs, x)

    # Sk*[w;-v] = 0, so pick out v,w after applying permuation
    m,n = _degree.((ps, qs))
    vᵢ = vcat(2:m-n+2, m-n+4:2:length(x))
    wᵢ = vcat(1, (m-n+3):2:length(x))
    v = -x[vᵢ]
    w = x[wᵢ]
    # p194 3.9 C_k(v) u = p or Ck(w) u = q; this uses 10.2
    u = [convmtx(v,j+1); convmtx(w, j+1)] \ [ps; qs]

    return u, v, w

end

function initial_uvw(::Val{:iszero}, j, ps::Vector{T}, qs, x) where {T}
    m,n = _degree.((ps,qs))
    S = [convmtx(ps, n-j+1) convmtx(qs, m-j+1)]

    F = qr(S)
    R = UpperTriangular(F.R)

    if iszero(det(R))
        x = eigvals(R)[:,1]
    else
        x = ones(T, size(R,2))
        x .= R \ (R' \ (x/norm(x)))
        x ./= norm(x)
    end

    w = x[1:n-j+1]
    v = -x[(n-j+2):end]

    u = [convmtx(v, j+1); convmtx(w, j+1)] \ [ps;qs]

    (u,v,w)
end

function initial_uvw(::Val{:constant}, j, ps::Vector{T}, qs, x) where {T}
    u = ones(T,1)
    w = qs
    v = ps
    u,v,w
end

# find estimate for σ₂, used in a condition number (κ = 1/σ)
function σ₂(J)
    F = qr(J)
    flag, σ, x = smallest_singular_value(F.R)
    σ
end

## attempt to refine u,v,w
## check that [u ⊗ v; u ⊗ w] ≈ [p; q]
function refine_uvw!(u::Vector{T}, v, w, p, q, atol, rtol) where {T}

    m, n, l =  _degree.((u, v, w))

    uv = u ⊗ v; uw = u ⊗ w
    ρ₀, ρ₁ = one(T), residual_error(p,q,uv,uw)

    # storage
    A = zeros(T, JF_size(u, v, w)...)
    b = zeros(T, 1 + length(p) + length(q))
    Δf = zeros(T, m + n + l + 3)
    steps = 0

    h, β =  u, norm(u)^2
    minᵢ, Maxᵢ = 5, 20
    κ = NaN

    JF!(A, h, u, v, w)
    Fmp!(b,  h'*u - β, p, q, uv, uw)

    while ρ₁ > 0.0

        # update A, b, then solve A\b
        qrsolve!(Δf, A, b)

        # m + n = degree(p)
        # m + l = degree(q)
        # b has length degree(p)+degree(q) + 3
        Δv, Δw, Δu = Δf[1:(n+1)], Δf[(n+1) .+ (1:l+1)], Δf[n+1+l+1 + 1:end]

        uv, uw =  (u-Δu) ⊗ (v-Δv),  (u-Δu) ⊗ (w-Δw)
        ρ₀, ρ′ = ρ₁, residual_error(p, q, uv, uw)

        # don't worry about first few, but aftewards each step must be productive
        # though we can have really bad first steps, which we cap
        if  (steps <= Maxᵢ) && (steps <= minᵢ || ρ′ < 0.95 * ρ₀) && (  ρ′ < 100*ρ₁ )
            ρ₁ = ρ′
            u .-= Δu; v .-= Δv; w .-= Δw
            steps += 1
        else
            break
        end

        # update A,b for next iteration
        JF!(A, h, u, v, w)
        Fmp!(b,  h'*u - β, p, q, uv, uw)

    end

    # this is a heuristic
    # sensitivity is Δu / Δp <= ‖ A+ ‖ = κ
    # we use an estimate for ‖(p,q)‖ error times ‖A⁺‖⋅‖A‖ₒₒ
    κ = 1/σ₂(A) # ≈ ‖A⁺‖
    λ = norm((norm(p), norm(q))) * (m * n) * min(1, κ) * norm(A, Inf)
    ρ = max(atol, rtol * λ)

    if ρ₁ <= ρ
        return :convergence, ρ₁, κ, ρ
    else
        return :no_convergence, ρ₁, κ, ρ
    end

end

function qrsolve!(y::Vector{T}, A, b) where {T}
    y .= qr(A) \ b
end

# # Fast least-squares solver for full column rank Hessenberg-like matrices
# # By Andreas Varga
function qrsolve!(y::Vector{T}, A, b) where {T <: Float64}
    Base.require_one_based_indexing(A)
    m, n = size(A)
    m < n && error("Column dimension exceeds row dimension")
    _, τ = LinearAlgebra.LAPACK.geqrf!(A)
    T <: Complex ? tran = 'C' : tran = 'T'
    LinearAlgebra.LAPACK.ormqr!('L', tran, A, τ, view(b,:,1:1))
    y .= UpperTriangular(triu(A[1:n,:]))\b[1:n]
end

## Jacobian F(u,v,w) = [p,p'] is J(u,v,w)
function JF_size(u, v, w)

    m, k, j = _degree(u), _degree(v), _degree(w)
    n, l = m + k, m + j

    ai, aj = convmtx_size(v, m + 1)
    bi, bj = convmtx_size(u, k + 1)
    di, dj = convmtx_size(w, m + 1)
    fi, fj = convmtx_size(u, j + 1)
    ci, cj = ai, fj
    ei, ej = di, bj

    (1 + ai + di, aj + bj + cj)
end

# Jacobian matrix
function JF(u::Vector{U}, v::Vector{V}, w::Vector{W}) where {U,V,W}
    R = promote_type(U,V, W)
    M = zeros(R, JF_size(u, v, w)...)
    JF!(M, u, v, w)
    M
end

function JF!(M, h,  u::Vector{T}, v, w) where {T}

    du, dv, dw = _degree(u), _degree(v), _degree(w)
    m, n = du + dv, du + dw

    # JF_size should return these
    r11,c11 = convmtx_size(u, dv+1)
    r13,c13 = convmtx_size(v, du+1)
    r22,c22 = convmtx_size(u, dw+1)
    r23,c23 = convmtx_size(w, du+1)

    J11 = view(M, 1:r11, 1:c11)
    J13 = view(M, 1:r13, c11 + c22 .+ (1:c23))
    J22 = view(M, r11 .+ (1:r22), c11 .+ (1:c22))
    J23 = view(M, r13 .+ (1:r23), (c11 + c22) .+ (1:c23))

    convmtx!(J11, u, dv+1)
    convmtx!(J13, v, du+1)
    convmtx!(J22, u, dw+1)
    convmtx!(J23, w, du+1)
    M[end, end-du:end] = h'

    return nothing
end

## compute F(u,v,w) - [p, p'] = [u*v, u*w] - [p, p']
function Fmp!(b, γ, p, q, uv, uw)
    b[end] = γ
    for i in 1:1+length(p)-1
        j = i
        b[i] = uv[j] - p[j]
    end
    for i in 1+length(p):length(b)-1
        j = i - length(p)
        b[i] = uw[j] - q[j]
    end
    return nothing
end

# norm( [u ⊗ v; u ⊗ w] .- [p; q], 2)
function residual_error(p::Vector{T}, q, uv, uw) where {T}
    tot = zero(T)
    for (pᵢ, uvᵢ) in zip(p,uv)
        tot += (pᵢ-uvᵢ)^2
    end
    for (qᵢ, uwᵢ) in zip(q, uw)
        tot += (qᵢ-uwᵢ)^2
    end
    sqrt(tot)
end

## --------------------------------------------------
## utils

_degree(p::Vector) = length(p) - 1

"""
    conv(p,q)

Convolution of two vectors (or product of two polynomials).
Also `⊗`
"""
function conv(p::Vector{T}, q::Vector{T}) where {T}
    R = promote_type(T,eltype(q))
    n = length(p)-1
    m = length(q)-1
    a = zeros(R, m+n+1)

    @inbounds for i = 1:n+1
        for j = 1:m+1
            a[i+j-1] += p[i] * q[j]
        end
    end
    a
end
⊗(p::Vector{T}, q::Vector{T}) where {T} = conv(p,q)

"""
    convmtx(v, n::Int)
    convmtx!(M,v,n)
    convmtx_size(v,n)

Convolution matrix.
C = convmtx(v,n) returns the convolution matrix C for a vector v.
If q is a column vector of length n, then C*q is the same as conv(v,q).

"""

function convmtx!(C, v::AbstractVector{T}, n::Int) where T

    #   Form C as the Toeplitz matrix
    #   C = Toeplitz([v; zeros(n-1)],[v[1]; zeros(n-1));  put Toeplitz code inline

    nv = length(v)-1

    @inbounds for j = 1:n
        C[j:j+nv,j] = v
    end

    nothing

end

convmtx_size(v::AbstractVector, n) = (n + length(v) - 1, n)
function convmtx(v::AbstractVector{T}, n::Int) where {T}
    C = zeros(T, convmtx_size(v, n)...)
    convmtx!(C, v, n)
    C
end

end


include("../src/cel_bulirsch.jl")
include("../src/transit_structure.jl")
include("../src/transit_poly_struct.jl")
include("../src/Mn_compute.jl")
using QuadGK

function Mn_num(k2::T,m::T) where {T <: Real}
# Numerically integrates M_n(k^2)/(4br)^m (note: m can be a half-integer).
# See 11/10/2018 notes.
f(x) = (k2-sin(x)^2)^m
g(x) = (k2-sin(x)^2)^m*sin(x)^2
if k2 < 1.0
  kap2 = convert(T,asin(sqrt(big(k2))))
  Mn,error = quadgk(f,-kap2,kap2,rtol=1e-15)
  Nn,error = quadgk(g,-kap2,kap2,rtol=1e-15)
else
  pi2 = 0.5*pi
  Mn,error = quadgk(f,-pi2,pi2,rtol=1e-15)
  Nn,error = quadgk(g,-pi2,pi2,rtol=1e-15)
end
return Mn,Nn
end

n = 30
t= transit_init(0.1,0.5,ones(n)/n,false)
nb = 7
rgrid = [0.01,0.1,1.0,10.0,100.0]
for r in rgrid
Mnr = zeros(t.n_max+1)
Nnr = zeros(t.n_max+1)
Mnl = zeros(t.n_max+1)
Nnl = zeros(t.n_max+1)
for i=1:nb
#r = 10.^(rand()*3-2); ; b=rand()*((1+r)-b0)+b0; fourbr = 4*b*r
  b0 = maximum([1e-8,r-1])
  b = b0+i/nb*(1+r-b0)
  t.r = r; t.b = b; t.fourbr = 4*b*r; t.sqbr = sqrt(b*r)
  t.sqarea = sqarea_triangle(1.0,r,b)
  t.onembmr2 = (1.0-(r-b)^2); t.onembpr2 = 1.0-(b+r)^2
  t.sqonembmr2 = sqrt(t.onembmr2)
  t.k2 = t.onembmr2/t.fourbr;
  t.onemr2mb2 = 1-r^2-b^2
  if t.k2 < 1.0
    println("k2: ",t.k2)
    t.k2inv=inv(t.k2); t.kc = sqrt(1.0-t.k2); t.k = sqrt(t.k2); t.kap0 = 2*asin(t.k)
    t.Eofk = cel_bulirsch(t.k2,t.kc,1.0,1.0,1.0-t.k2)
    t.Em1mKdm = cel_bulirsch(t.k2,t.kc,1.0,1.0,0.0)
    t.kite_area2 = sqrt(abs(t.sqarea))
  else
    println("k2: ",t.k2)
    t.k2inv=inv(t.k2); t.kc = sqrt(1.0-t.k2inv); t.k = sqrt(t.k2); t.kap0 = pi
    t.Eofk = cel_bulirsch(t.k2inv,t.kc,1.0,1.0,1.0-t.k2inv)
    t.Em1mKdm = cel_bulirsch(t.k2inv,t.kc,1.0,1.0,0.0)
  end
  Mn_raise!(t)
  Mnr .= t.Mn
  Nn_raise!(t)
  Nnr .= t.Nn
  Mn_lower!(t)
  Nn_lower!(t)
  Mnl .= t.Mn
  Nnl .= t.Nn
  Mnn = zeros(t.n_max+1)
  # Now, compute the N_n functions: N_n = (4br)^{n/2} \int_{-\kappa/2}^{\kappa/2} (k^2-\sin^2\phi)^{n/2} \sin^2\phi d\phi
  for m=0:t.n_max
    Mn_n,Nn_num = Mn_num(t.k2,float(m/2))
    Mnn[m+1] = Mn_n*sqrt(t.fourbr)^m
    Nn_num *= sqrt(t.fourbr)^m
    println(Mnr[m+1]," ",Mnl[m+1]," ",Mnn[m+1]," ",Mnr[m+1]-Mnn[m+1]," ",Mnl[m+1]-Mnn[m+1]," ",Mnl[m+1]-Mnr[m+1])
    # Now use function which doesn't involve division by b:
    if t.k2 < 0.5
      println("n: ",m," N_n: ",Nnl[m+1]," N_n num: ",Nn_num)
    else
      println("n: ",m," N_n: ",Nnr[m+1]," N_n num: ",Nn_num)
    end
  end
  println("k2: ",t.k2," r: ",r," b: ",b," max(M_n-M_{n,num}): ",maximum(abs,Mnr-Mnn)," ",maximum(abs,Mnl-Mnn))
  read(STDIN,Char)
end
end


@diffusion_process LorenzAux{K,R} begin
    :dimensions
    process --> 3
    wiener --> 3

    :parameters
    _ --> (3, K)
    σ --> K

    :auxiliary_info
    t0 --> Float64
    T --> Float64
    vT --> R

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

DiffusionDefinition.B(t, P::LorenzAux) = @SMatrix [
    -P.p1         P.p1        0.0;
     P.p2-P.vT[3]  -1.0    -P.vT[1];
     P.vT[2]      P.vT[1]    -P.p3
]

DiffusionDefinition.β(t, P::LorenzAux) = @SVector[
    0.0,
    P.vT[1]*P.vT[3],
    -P.vT[1]*P.vT[3]
]

DiffusionDefinition.σ(t, P::LorenzAux) = SDiagonal(P.σ, P.σ, P.σ)


######################################################################
# eid.jl
# Generates the expected information density (EID).
######################################################################

"""
`fisher(m::SearchDomain)`

Creates the Fisher information matrix. Vehicle positions are discretized to same level as possible jammer locations.
"""
function fisher2(m::SearchDomain)
	# TODO: multiply by measurement noise
	# TODO: what do we do about rad2deg business? Multiply by constant?
	# TODO: what do we do about xr = yr = 0??
	n = m.num_cells
	F = zeros(n, n, n, n, 2, 2)

	for vx = 1:n
		for vy = 1:n
			for theta_x = 1:n
				for theta_y = 1:n
					# Create the 2x2 matrix
					xr = theta_x - vx
					yr = theta_y - vy
					if xr == 0 && yr == 0
						xr = 1e-6
						yr = 1e-6
					end
					den = (xr^2 + yr^2)^2
					F[vx,vy,theta_x,theta_y,1,1] = -2*yr*xr / den
					F[vx,vy,theta_x,theta_y,2,2] = 2*yr*xr / den
					F[vx,vy,theta_x,theta_y,1,2] = (xr^2 - yr^2) / den
					F[vx,vy,theta_x,theta_y,2,1] = (xr^2 - yr^2) / den
				end
			end
		end
	end
	return F
end

function fisher(m::SearchDomain)
	# TODO: multiply by measurement noise
	# TODO: what do we do about rad2deg business? Multiply by constant?
	# TODO: what do we do about xr = yr = 0??
	n = m.num_cells
	F = zeros(n, n, n, n, 2, 2)

	for vx = 1:n
		for vy = 1:n
			for theta_x = 1:n
				for theta_y = 1:n
					# compute phi
					# Create the 2x2 matrix
					xr = theta_x - vx
					yr = theta_y - vy
					if xr == 0 && yr == 0
						xr = 1e-6
						yr = 1e-6
					end

					alpha = atan2(xr, yr)
					den = (xr^2 + yr^2)

					F[vx,vy,theta_x,theta_y,1,1] = cos(alpha)^2 / den
					F[vx,vy,theta_x,theta_y,2,2] = sin(alpha)^2 / den
					F[vx,vy,theta_x,theta_y,1,2] = -sin(2*alpha) / (2*den)
					F[vx,vy,theta_x,theta_y,2,1] = -sin(2*alpha) / (2*den)
				end
			end
		end
	end
	return F
end

"""
`EID(m::SearchDomain, b::Belief)`

Arguments:

* `m` is the search domain
* `b` is the belief

Returns the expected information density.
"""
function EID(m::SearchDomain, b::Belief)
	n = m.num_cells
	F = m.F
	ei = zeros(n, n, 2, 2)
	eid = zeros(n, n)
	for vx = 1:n
		for vy = 1:n
			for theta_x = 1:n
				for theta_y = 1:n
					p = b[theta_x, theta_y]
					ei[vx, vy, 1, 1] += F[vx, vy, theta_x, theta_y,1,1] * p
					ei[vx, vy, 2, 1] += F[vx, vy, theta_x, theta_y,2,1] * p
					ei[vx, vy, 1, 2] += F[vx, vy, theta_x, theta_y,1,2] * p
					ei[vx, vy, 2, 2] += F[vx, vy, theta_x, theta_y,2,2] * p
				end
			end
			ma = ei[vx, vy, 1, 1]
			mb = ei[vx, vy, 2, 1]
			mc = ei[vx, vy, 1, 2]
			md = ei[vx, vy, 2, 2]
			eid[vx, vy] = ma*md - mb*mc
		end
	end
	return eid
end

function EID(m::SearchDomain)
	return EID(m, m.b)
end

"""
`EID(m::SearchDomain, theta_x::Int64, theta_y::Int64)`

Arguments:

* `m` is the search domain
* `(theta_x, theta_y)` is the jammer location (0-indexed)

Returns the expected information density.
"""
function EID(m::SearchDomain, theta_x::Int, theta_y::Int)
	n = m.num_cells
	F = m.F
	eid = zeros(n, n)
	for vx = 1:n
		for vy = 1:n
			ma = F[vx, vy, theta_x+1, theta_y+1, 1, 1]
			mb = F[vx, vy, theta_x+1, theta_y+1, 1, 2]
			mc = F[vx, vy, theta_x+1, theta_y+1, 2, 1]
			md = F[vx, vy, theta_x+1, theta_y+1, 2, 2]

			eid[vx, vy] = ma*md - mb*mc
		end
	end
	return eid
end


# Probably this should not be used in the periodic system
function init_pot_Hps_HGH_G( Gvec::GvectorsT, center )

    Ns = Gvec.Ns
    Ngvec = prod(Ns)
    Npoints = Ngvec
    Ω = Gvec.Ω
    G2 = Gvec.G2

    Vpot = zeros( Float64,Npoints )

    const zval = 1
    const rloc = 0.2
    const c1 = -4.0663326
    const c2 = 0.6678322

    Vg = zeros(Complex128,Npoints)
    pre1 = -4*pi*zval/Ω
    pre2 = sqrt(8*pi^3)*rloc^3/Ω
    #
    for ig=2:Ngvec
        Gr = sqrt(G2[ig])*rloc
        expGr2 = exp(-0.5*Gr^2)
        Vg[ig] = pre1/G2[ig]*expGr2 + pre2*expGr2 * (c1 + c2*(3-Gr^2))
    end
    # limiting value, with minus sign ?
    Vg[1] = 2*pi*zval*rloc^2 + (2*pi)^1.5 * rloc^3 * (c1 + 3.0*c2)

    V_ionic = real( G_to_R(Ns,Vg) ) * Npoints

    return V_ionic
end

function init_pot_Hps( LF, center )
    Npoints = LF.Nx * LF.Ny * LF.Nz
    Vpot = zeros( Float64,Npoints )

    const rc1 = 0.25
    const rc2 = 0.284
    const aa  = -1.9287
    const bb  = 0.3374

    # TODO Add journal reference for this pseudopotential
    for ip=1:Npoints
        r = norm( LF.lingrid[:,ip] - center[:] )
        Vpot[ip] = -1.0/r * erf( r/rc1 ) + (aa + bb*r^2)*exp(-(r/rc2)^2)
    end
    return Vpot
end

# Probably this should not be used in the periodic system
function init_pot_Hps_HGH( LF, center )
    Npoints = LF.Nx * LF.Ny * LF.Nz
    Vpot = zeros( Float64,Npoints )

    const Zion = 1
    const rloc = 0.2
    const C1 = -4.0663326
    const C2 = 0.6678322

    # TODO Add journal reference
    for ip = 1:Npoints
        r = norm( LF.lingrid[:,ip] - center[:] )
        if r < 1e-7
            println("WARNING: small r:", r);
            println("Using limiting value:")
            Vpot[ip] = -2*Zion/(sqrt(2*pi)*rloc) + C1
        else
            rrloc = r/rloc
            Vpot[ip] = -Zion/r * erf( r/(sqrt(2.0)*rloc) ) +
                     (C1 + C2*rrloc^2)*exp(-0.5*(rrloc)^2)
        end
    end
    return Vpot
end


#=

Implementation of commutator-free exponential time propagators, as
described in

  Alvermann, A., Fehske, H., & Littlewood, P. B. (2012). Numerical
  time propagation of quantum systems in radiation fields. New
  J. Phys., 14(10),
  105008. http://dx.doi.org/10.1088/1367-2630/14/10/105008

=#
struct CFET4BfCPropagator{E<:Exponentiator} <: MagnusPropagator
    B::LinearMap
    f::Function
    C::LinearMap
    a::Number
    Exp::E
end
CFET4BfCPropagator(B::LinearMap,
                   f::Function,
                   C::LinearMap,
                   Exp::E) where E<:Exponentiator =
                       CFET4BfCPropagator(B, f, C, 1, Exp)

function (p::CFET4BfCPropagator{E})(t::T, τ::T,
                                    v::AbstractVector, w::AbstractVector) where {E<:Exponentiator,
                                                                                T<:AbstractFloat}
    h1 = T(37/66 - 400/957*sqrt(5/3))
    h2 = T(-4/33)
    h3 = T(37/66 + 400/957*sqrt(5/3))
    h4 = T(-11/162)
    h5 = T(92/81)

    ff1 = p.f(t + T(1/2-sqrt(3/20))*τ)
    ff2 = p.f(t + τ/2)
    ff3 = p.f(t + T(1/2+sqrt(3/20))*τ)

    f1 = h1*ff1 + h2*ff2 + h3*ff3
    f2 = h4*ff1 + h5*ff2 + h4*ff3
    f3 = h3*ff1 + h2*ff2 + h1*ff3

    τ1 = 11τ/40
    τ2 = 9τ/20

    p.Exp(p.B + f3*p.C, p.a*τ1, v, w)
    p.Exp(p.B + f2*p.C, p.a*τ2, w, w)
    p.Exp(p.B + f1*p.C, p.a*τ1, w, w)
end

export CFET4BfCPropagator


# small scrip to compute the solution of Lipman Schinwer equation
# We test the types introduced in FastConvolution.jl
# we test that the application is fast and that the construction
# is performed fast.

using PyPlot
using IterativeSolvers
using SpecialFunctions

@everywhere include("../src/FastConvolution.jl")
@everywhere include("../src/Preconditioner.jl")

FFTW.set_num_threads(4);
BLAS.set_num_threads(4)

#Defining Omega
h = 0.001
k = 1/h

# size of box
a  = 1
x = -a/2:h:a/2
y = -a/2:h:a/2
(n,m) = length(x), length(y)
N = n*m
X = repmat(x, 1, m)[:]
Y = repmat(y', n,1)[:]
# we solve \triangle u + k^2(1 + nu(x))u = 0

# We use the modified quadrature in Ruan and Rohklin
(ppw,D) = referenceValsTrapRule();
D0 = D[1];

spline(y) = (y.<0) +  (y.>=0).*(y.<1).*(2*y.^3 - 3*y.^2 + 1)
window(a1,b1,b2,a2,y) = (y.>=a1).*( (y.<b1).*spline(1/(abs(b1-a1)).*(-y+b1)) + (y.>=b1).*(y.<b2) + (y.>=b2).*spline(1/(abs(b2-a2)).*(y-b2)))

# Defining the smooth perturbation of the slowness
nu(x,y) = 0.05*sin(2*pi*x).*window(-0.48,-0.4,0.4,0.48,y).*window(-0.48,-0.05,0.05,0.48,x);

figure(1); clf();
imshow(real(reshape( (1-nu(X,Y)),n,m)), extent=[x[1], x[end], y[1], y[end]]);cb =  colorbar();

# Sampling the Green's function for the Toeplitz form
Ge = buildGConv(x,y,h,n,m,D0,k);
GFFT = fft(Ge);

fastconv = FastM(GFFT,nu(X,Y),2*n-1,2*m-1,n, m, k);

# assembling the sparsifiying preconditioner
As = buildSparseA(k,X,Y,D0, n ,m);

# assembling As*( I + k^2G*nu)
Mapproxsp = As + k^2*(buildSparseAG(k,X,Y,D0, n ,m)*spdiagm(nu(X,Y)));

# defining the preconditioner
precond = SparsifyingPreconditioner(Mapproxsp, As)

# building the RHS from the incident field
u_plane = exp(k*im*sqrt(1-nu(0*Y,X)).*X);

u_inc = exp(k*im*X);

rhs = -(fastconv*u_inc - u_inc);

# allocating the solution
u = zeros(Complex128,N);

# solving the system using GMRES
@time info =  gmres!(u, fastconv, rhs, Pl = precond)
# println(info[2].residuals[:])

# plotting the solution
figure(2)
clf()
imshow(real(reshape(u+u_inc,n,m)))

figure(3); clf()
imshow(real(reshape((u+u_inc)./u_inc, n,m))); colorbar()
title("\$u(x)/e^{ikx}\$ ")

intx = cumsum(sqrt(1-nu(x,0*x)))*h;
u_gem = exp(k*im*intx);

UUgem = repmat(u_gem, 1,n)

figure(4); clf()
imshow(real(reshape((u+u_inc)./UUgem[:], n,m))); colorbar()
title("\$u(x)/e^{ik \\phi (x)}\$ where \$ \\phi \$ comes from \$ \\int_{0}^{x} \\sqrt{1 + n(y)} dy \$ ")

refFact = 10;
xRef = -a/2:h/refFact:a/2

intxRef = cumsum(sqrt(1-nu(xRef,0*xRef)))*h/refFact;
u_gem = exp(k*im*intxRef)[1:refFact:end];

UUgem = repmat(u_gem, 1,n)

figure(14); clf()
imshow(real(reshape(-(u+u_inc)./UUgem[:], n,m))); colorbar()
title("\$u(x)/e^{ik \\phi (x)}\$ where \$ \\phi \$ comes from \$ \\int_{0}^{x} \\sqrt{1 + n(y)} dy \$ ")

UU = reshape(u+u_inc,n,m)
ucomp = UU[:,round(Integer,floor(end/2))]

# there is a two that had to disappear... I need to find the bug... :/
phi_1 = -1im*rhs./u_inc

figure(5); clf()
imshow(real(reshape(phi_1, n,m))); colorbar()
title("phase function \$phi_1(x)\$ ")

u_rytov = u_inc.*exp(1im*phi_1)

figure(6); clf()
imshow(real(reshape(u_rytov, n,m))); colorbar()
title("\$u_{rytov}\$" )

figure(7); clf()
imshow(real(reshape((u+u_inc)./ u_rytov, n,m))); colorbar()
title("\$u(x)/u_{rytov}(x)\$ ")

#Defining the smooth perturbation of the slowness
nu1(x,y) = nu(x,y).*window(-0.28,-0.10,0.10,0.28,y);

figure(24); clf();
imshow(real(reshape( (1-nu1(X,Y)),n,m)), extent=[x[1], x[end], y[1], y[end]]);cb =  colorbar();

fastconvWindow = FastM(GFFT,nu1(X,Y),2*n-1,2*m-1,n, m, k);

rhsWindow = -(fastconvWindow*u_inc - u_inc);

# there is a two that had to disappear... I need to find the bug... :/
phi_Window = -1im*rhsWindow./u_inc

figure(15); clf()
imshow(real(reshape(phi_1, n,m))); colorbar()
title("phase function phi_1(x) ")

u_rytovWindow = u_inc.*exp(1im*phi_Window)

figure(16); clf()
imshow(real(reshape(u_rytovWindow, n,m))); colorbar()
title("u_rytov ")

figure(17); clf()
imshow(real(reshape((u+u_inc)./ u_rytovWindow, n,m))); colorbar()
title("u/u_rytov ")

UU = reshape((u+u_inc)./u_rytov,n,m)
ucomp = UU[:,round(Integer,floor(end/2))]

figure(20); clf();
plot(real(ucomp[:]))

# u_inc = exp(k*im*X);

# rhs1 = -(fastconv1*u_inc - u_inc);

# # allocating the solution
# u1 = zeros(Complex128,N);

# # solving the system using GMRES
# @time info =  gmres!(u1, fastconv1, rhs1, precond1)
# println(info[2].residuals[:])

# # plotting the solution
# figure(5)
# clf()
# imshow(real(reshape(u1+u_inc,n,m)))

# figure(6); clf()
# imshow(real(reshape((u1+u_inc)./u_inc, n,m))); colorbar()


module BasisFunctionExpansions
using Clustering, LinearAlgebra, Statistics
export BasisFunctionExpansion, UniformRBFE, MultiUniformRBFE, MultiDiagonalRBFE, MultiRBFE, BasisFunctionApproximation
export get_centers, get_centers_multi, get_centers_automatic, quadform, γ2σ, σ2γ
export toeplitz, getARregressor, getARXregressor, LPVSS, predict, output_variance

## Types
abstract type BasisFunctionExpansion end

# function Base.show{N}(io::IO,b::BasisFunctionExpansion{N})
#     s = string(typeof(b),"\n")
#     for fn in fieldnames(b)
#         s *= string(fn) * ": " * string(getfield(b,fn)) * "\n"
#     end
#     print(io,s)
# end

struct BasisFunctionApproximation{BFE <: BasisFunctionExpansion}
    bfe::BFE
    linear_combination::Vector{Float64}
end

"""
    BasisFunctionApproximation(y::Vector, v, bfe::BasisFunctionExpansion, λ = 0)

Perform parameter identification to identify the Function `y = ϕ(v)`, where `ϕ` is a Basis Function Expansion of type `bfe`.
`λ` is an optional regularization parameter (L² regularization).
"""
function BasisFunctionApproximation(y::AbstractVector,v,bfe::BasisFunctionExpansion, λ = 0)
    A = bfe(v)
    p = size(A,2)
    if λ == 0
        x = A\y
    else
        x = [A; λ*I]\[y;fill(0,p)]
    end
    BasisFunctionApproximation(bfe,x)
end

function (bfa::BasisFunctionApproximation)(v)
    A = bfa.bfe(v)
    return A*bfa.linear_combination
end

### UniformRBFE ================================================================
"""
A Uniform RBFE has the same variance for all basis functions
"""
struct UniformRBFE{F,VT} <: BasisFunctionExpansion
    activation::F
    μ::VT
    σ::Float64
end

"""
    UniformRBFE(μ::Vector, σ::Float, activation)

Supply all parameters. OBS! `σ` can not be an integer, must be some kind of AbstractFloat
"""
function UniformRBFE(μ::AbstractVector, σ::AbstractFloat, activation)
    UniformRBFE(v->activation(v,μ,σ2γ(σ)),μ,σ)
end

"""
    UniformRBFE(v::Vector, Nv::Int; normalize=false, coulomb=false)

Supply scheduling signal and number of basis functions For automatic selection of centers and widths

The keyword `normalize` determines weather or not basis function activations are normalized to sum to one for each datapoint, normalized networks tend to extrapolate better ["The normalized radial basis function neural network" DOI: 10.1109/ICSMC.1998.728118](http://ieeexplore.ieee.org/document/728118/)
"""
function UniformRBFE(v::AbstractVector, Nv::Int; normalize=false, coulomb=false)
    activation, μ, γ = basis_activation_func_automatic(v,Nv,normalize,coulomb)
    UniformRBFE(activation,μ,γ2σ(γ))
end

### RBFE =======================================================================
"""
A `MultiUniformRBFE` has the same diagonal covariance matrix for all basis functions
See also `MultiDiagonalRBFE`, which has different covariance matrices for all basis functions
"""
struct MultiUniformRBFE{F,MT<:AbstractMatrix,VT<:AbstractVector} <: BasisFunctionExpansion
    activation::F
    μ::MT
    Σ::VT
end

"""
    MultiUniformRBFE(μ::Matrix, Σ::Vector, activation)

Supply all parameters. Σ is the diagonal of the covariance matrix
"""
function MultiUniformRBFE(μ::AbstractMatrix, Σ::AbstractVector, activation)
    MultiUniformRBFE(v->activation(v,μ,σ2γ(Σ)),μ,Σ)
end

"""
    MultiUniformRBFE(v::AbstractVector, Nv::Vector{Int}; normalize=false, coulomb=false)

Supply scheduling signal and number of basis functions For automatic selection of centers and widths

The keyword `normalize` determines weather or not basis function activations are normalized to sum to one for each datapoint, normalized networks tend to extrapolate better ["The normalized radial basis function neural network" DOI: 10.1109/ICSMC.1998.728118](http://ieeexplore.ieee.org/document/728118/)
"""
function MultiUniformRBFE(v::AbstractMatrix, Nv::AbstractVector{Int}; normalize=false, coulomb=false)
    @assert !coulomb "Coulomb not yet supported for multi-dimensional BFEs"
    @assert length(Nv) == size(v,2)
    activation, μ, γ = basis_activation_func_automatic(v,Nv,normalize,coulomb)
    MultiUniformRBFE(activation,μ,γ2σ.(γ))
end

## MultiDiagonalRBFE =======================================================================
"""
A `MultiDiagonalRBFE` has different diagonal covariance matrices for all basis functions
See also `MultiUniformRBFE`, which has the same covariance matrix for all basis functions
"""
struct MultiDiagonalRBFE{F,MT<:AbstractMatrix,VT<:AbstractVector} <: BasisFunctionExpansion
    activation::F
    μ::MT
    Σ::Vector{VT}
end

"""
    MultiDiagonalRBFE(μ::Matrix, Σ::Vector{Vector{Float64}}, activation)

Supply all parameters. Σ is the diagonals of the covariance matrices
"""
function MultiDiagonalRBFE(μ::AbstractMatrix, Σ::AbstractVector{T}, activation) where T <: AbstractVector
    MultiDiagonalRBFE(v->activation(v,μ,σ2γ(Σ)),μ,Σ)
end

"""
    MultiDiagonalRBFE(v::AbstractVector, nc; normalize=false, coulomb=false)

Supply scheduling signal `v` and numer of centers `nc` For automatic selection of covariance matrices and centers using K-means.

The keyword `normalize` determines weather or not basis function activations are normalized to sum to one for each datapoint, normalized networks tend to extrapolate better ["The normalized radial basis function neural network" DOI: 10.1109/ICSMC.1998.728118](http://ieeexplore.ieee.org/document/728118/)
"""
function MultiDiagonalRBFE(v::AbstractMatrix, nc; normalize=false, coulomb=false)
    @assert !coulomb "Coulomb not yet supported for multi-dimensional BFEs"
    μ, Σ = get_centers_Kmeans(v, nc)
    μ = hcat(μ...)
    Σ = [2diag(Σi) for Σi in Σ] # Heuristically inflate covariance by 2 "On the Kernel Widths in Radial-Basis Function Networks" NABIL BENOUDJIT and MICHEL VERLEYSEN
    activation = basis_activation_func(μ,σ2γ(Σ),normalize,coulomb)
    MultiDiagonalRBFE(activation,μ,Σ)
end

## MultiRBFE =======================================================================
"""
A `MultiRBFE` has different diagonal covariance matrices for all basis functions
See also `MultiUniformRBFE`, which has the same covariance matrix for all basis functions
"""
struct MultiRBFE{F,MT<:AbstractMatrix,MTs} <: BasisFunctionExpansion
    activation::F
    μ::MT
    Σ::Vector{MTs}
end

"""
    MultiRBFE(μ::Matrix, Σ::Vector{Vector{Float64}}, activation)

Supply all parameters. Σ is the diagonals of the covariance matrices
"""
function MultiRBFE(μ::AbstractMatrix, Σ::AbstractVector{T}, activation) where T <: AbstractVector
    MultiRBFE(v->activation(v,μ,σ2γ(Σ)),μ,Σ)
end

"""
    MultiRBFE(v::AbstractVector, nc; normalize=false, coulomb=false)

Supply scheduling signal `v` and numer of centers `nc` For automatic selection of covariance matrices and centers using K-means.

The keyword `normalize` determines weather or not basis function activations are normalized to sum to one for each datapoint, normalized networks tend to extrapolate better ["The normalized radial basis function neural network" DOI: 10.1109/ICSMC.1998.728118](http://ieeexplore.ieee.org/document/728118/)
"""
function MultiRBFE(v::AbstractMatrix, nc; normalize=false, coulomb=false)
    @assert !coulomb "Coulomb not yet supported for multi-dimensional BFEs"
    μ, Σ = get_centers_Kmeans(v, nc)
    μ = hcat(μ...)
    Σ .*= 3 # Heuristically inflate covariance by 3 "On the Kernel Widths in Radial-Basis Function Networks" NABIL BENOUDJIT and MICHEL VERLEYSEN
    activation = basis_activation_func(μ,σ2γ(Σ),normalize,coulomb)
    MultiRBFE(activation,μ,Σ)
end

## Squared exponential functions

squared_exponential(v::Real,vc,gamma) = exp.(-gamma*(v.-vc).^2)
squared_exponential(v::AbstractVector,vc,gamma::Number) = exp.(-gamma*(v.-vc').^2)
squared_exponential(v::AbstractVector,vc,gamma::AbstractVector) = exp.(-((vc'.-v').^2)*gamma)
function squared_exponential(v::AbstractMatrix{T},vc,gamma::AbstractVector) where T
    a = Matrix{T}(undef,size(v,1),size(vc,2))
    for i = 1:size(v,1)
        a[i,:] = exp.(-sum(gamma.*(v[i,:].-vc).^2, dims=1))
    end
    a
end

function squared_exponential(v::AbstractMatrix,vc,gamma::AbstractVector{T}) where T <: AbstractVector
    a = zeros(size(v,1),size(vc,2))
    for j = 1:size(vc,2)
        for i = 1:size(v,1)
            a[i,j] += exp(-gamma[j]⋅(v[i,:].-vc[:,j]).^2)
        end
    end
    a
end

function squared_exponential(v::AbstractMatrix,vc,gamma::AbstractVector{T}) where T <: AbstractMatrix
    a = zeros(size(v,1),size(vc,2))
    for j = 1:size(vc,2)
        for i = 1:size(v,1)
            d = v[i,:].-vc[:,j]
            a[i,j] += exp(-quadform(d,gamma[j]))
        end
    end
    a
end

function normalized_squared_exponential(v,vc,gamma::Number)
    r = squared_exponential(v,vc,gamma)
    r ./= (sum(r,dims=2) .+ 1e-8)
end

function normalized_squared_exponential(v::AbstractVector,vc,gamma::AbstractVector)
    r = squared_exponential(v,vc,gamma)
    r ./= (sum(r) + 1e-8)
end

function normalized_squared_exponential(v::AbstractMatrix,vc,gamma::AbstractVector)
    r = squared_exponential(v,vc,gamma)
    r ./= (sum(r,dims=2) .+ 1e-8)
end

squared_exponential_coulomb(v,vc,gamma) = squared_exponential(v,vc,gamma).*(sign(v) .== sign(vc))

function normalized_squared_exponential_coulomb(v,vc,gamma)
    r = squared_exponential_coulomb(v,vc,gamma)
    r ./= (sum(r,dims=2) .+ 1e-8)
end

function squared_exponential(v::AbstractMatrix,vc, sigma, velocity::Int=0)
    error("This function has not yet been revised")
    @assert size(v,2) == length(sigma)
    N_basis = size(vc,2)
    y       = fill(0,N_basis)
    iSIGMA  = sigma.^-2

    if velocity > 0
        y = [sign(vc[velocity,i]) == sign(v[velocity]) ? mnorm_pdf(v,vc[:,i],iSIGMA) : 0 for i = 1:N_basis]
    else
        y = [mnorm_pdf(v,vc[:,i],iSIGMA) for i = 1:N_basis]
    end
end

"""
    basis_activation_func_automatic(v,Nv,normalize,coulomb)

Returns a func v->ϕ(v) ∈ ℜ(Nv) that calculates the activation of `Nv` basis functions spread out to cover v nicely. If coulomb is true, then we get twice the number of basis functions, `2Nv`, with a hard split at `v=0` (useful to model Coulomb friction). coulomb is not yet fully supported for all expansion types.

The keyword `normalize` determines weather or not basis function activations are normalized to sum to one for each datapoint, normalized networks tend to extrapolate better ["The normalized radial basis function neural network" DOI: 10.1109/ICSMC.1998.728118](http://ieeexplore.ieee.org/document/728118/)
"""
function basis_activation_func_automatic(v,Nv,normalize,coulomb=false)
    vc,gamma = get_centers_automatic(v,Nv,coulomb)
    K = basis_activation_func(vc,gamma,normalize,coulomb)
    K,vc,gamma
end

function basis_activation_func(vc,gamma,normalize,coulomb=false)
    if coulomb
        K = normalize ? v -> normalized_squared_exponential_coulomb(v,vc,gamma) : v -> squared_exponential_coulomb(v,vc,gamma) # Use coulomb basis function instead
    else
        K = normalize ? v -> normalized_squared_exponential(v,vc,gamma) : v -> squared_exponential(v,vc,gamma)
    end
end

"""
    vc,γ = get_centers_automatic(v::AbstractVector,Nv::Int,coulomb = false)
"""
function get_centers_automatic(v::AbstractVector,Nv::Int,coulomb = false)
    if coulomb # If Coulomb setting is activated, double the number of basis functions and clip the activation at zero velocity (useful for data that exhibits a discontinuity at v=0, like coulomb friction)
        vc    = range(0, stop=maximum(abs.(v)), length=Nv+2)
        vc    = vc[2:end-1]
        vc    = [-vc[end:-1:1]; vc]
        Nv    = 2Nv
        gamma = (Nv/(abs(vc[1]-vc[end])))^2
    else
        vc    = range(minimum(v), stop=maximum(v), length=Nv)
        gamma = (Nv/(abs(vc[1]-vc[end])))^2
    end
    vc,gamma
end

"""
    vc,γ = get_centers_automatic(v::AbstractMatrix, Nv::AbstractVector{Int}, coulomb=false, coulombdims=0)
"""
function get_centers_automatic(v::AbstractMatrix, Nv::AbstractVector{Int}, coulomb=false, coulombdims=0)
    @assert !coulomb "Coulomb not yet supported for multi-dimensional BFEs"
    @assert size(v,2) == length(Nv) "size(v,2) != length(Nv)"
    dims   = size(v,2)
    minq   = minimum(v,dims=1)[:]
    maxq   = maximum(v,dims=1)[:]
    bounds = [minq maxq]
    get_centers(bounds, Nv)
end

"""
    vc,γ = get_centers(bounds, Nv, coulomb=false, coulombdims=0)
"""
function get_centers(bounds, Nv, coulomb=false, coulombdims=0)
    # TODO: split centers on velocity dim
    @assert !coulomb "Coulomb not yet supported for multi-dimensional BFEs"
    dims     = length(Nv)
    interval = [(bounds[n,2]-bounds[n,1])/Nv[n] for n = 1:dims]
    C = [range(bounds[n,1]+interval[n]/2, stop=bounds[n,2]-interval[n]/2, length=(Nv)[n]) for n = 1:dims]

    Nbasis  = prod(Nv)
    centers = zeros(dims, Nbasis)
    v       = Nbasis
    h       = 1
    for i = 1:dims
        v = v ÷ Nv[i]
        centers[i,:] = vec(repeat(C[i]',v,h))'
        h *= Nv[i]
    end
    centers, (1 ./interval).^2
end

function get_centers_Kmeans(v, nc::Int; verbose=false)
    iters = 21
    n_state = size(v,2)
    errorvec = fill(0.,iters)
    methods = [:rand;:kmpp]
    Σ = [fill(0.,n_state,n_state) for i = 1:nc, j = 1:iters]
    μ = [fill(0.,n_state) for i = 1:nc, j = 1:iters]

    for iter = 1:iters
        clusterresult = Clustering.kmeans(copy(v'), nc; maxiter=200, display=:none, init=iter<iters ? methods[iter%2+1] : :kmcen)
        for i = 1:nc
            si = 1+(i-1)n_state*2
            μ[i,iter] .= clusterresult.centers[:,i]
            Σ[i,iter] .= cov(v[clusterresult.assignments .== i,:])
        end
        errorvec[iter] = clusterresult.totalcost
    end
    verbose && println("Std in errors among initial centers: ", round(std(errorvec), digits=6))
    ind = argmin(errorvec)
    return μ[:,ind], Σ[:,ind]
end

## Utility functions

function quadform(x::AbstractVector,A::AbstractMatrix)
    x⋅(A*x)
end

function quadform(x::AbstractVector,A::AbstractVector)
    s = zero(eltype(x))
    for i in 1:length(x)
        s += A[i]*x[i]^2
    end
    s
end

γ2σ(γ) = √(1 ./(2γ))
σ2γ(σ) = 1 ./(2σ.^2)

function γ2σ(γ::Vector{T}) where T <: AbstractVector
    [1 ./(2γ) for γi in γ]
end

function σ2γ(σ::Vector{T}) where T <: AbstractVector
    [1 ./(2σi) for σi in σ]
end

function γ2σ(γ::Vector{T}) where T <: AbstractMatrix
    [inv(2γ) for γi in γ]
end

function σ2γ(σ::Vector{T}) where T <: AbstractMatrix
    [inv(2σi) for σi in σ]
end

function meshgrid(a::AbstractVector,b::AbstractVector)
    grid_a = [i for i in a, j in b]
    grid_b = [j for i in a, j in b]
    grid_a, grid_b
end

function meshgrid(a::AbstractVector,b::AbstractVector,c::AbstractVector)
    grid_a = [i for i in a, j in b, k in c]
    grid_b = [j for i in a, j in b, k in c]
    grid_c = [k for i in a, j in b, k in c]
    grid_a, grid_b, grid_c
end
# This code must be run after all type definitions!

# for T in subtypes(BasisFunctionExpansion)
#     (b::T)(x) = b.activation(x)
# end

(b::UniformRBFE)(x)       = b.activation(x)
(b::MultiUniformRBFE)(x)  = b.activation(x)
(b::MultiDiagonalRBFE)(x) = b.activation(x)
(b::MultiRBFE)(x)         = b.activation(x)

## Plot tools ==================================================================
include("plotting.jl")
include("dynamics.jl")
# plot

"""
    x,xm,u,n,m = testdata(T,r=1)

Generate `T` time steps of state-space data where the A-matrix changes from
`A = [0.95 0.1; 0 0.95]` to `A = [0.5 0.05; 0 0.5]` at time `t=T÷2`
`x,xm,u,n,m` = (state,noisy state, input, statesize, inputsize)
`r` is the seed to the random number generator.
"""
function testdata(T_,r=1)
    Random.seed!(r)
    n,m      = 2,1
    At_      = [0.95 0.1; 0 0.95]
    Bt_      = reshape([0.2; 1],2,1)
    u        = randn(1,T_)
    x        = fill(0,n,T_)
    for t = 1:T_-1
        if t == T_÷2
            At_ = [0.5 0.05; 0 0.5]
        end
        x[:,t+1] = At_*x[:,t] + Bt_*u[:,t] #+ 0.2randn(n)
    end
    xm = x + 0.2randn(size(x));
    x',xm',u',n,m
end

end # module


using FFTW

using Oceananigans.BoundaryConditions: PBC, ZFBC

#=
These functions return the transforms required to solve Poisson's equation with
periodic boundary conditions or staggered Neumann boundary.

Fast Fourier transforms (FFTs) are used in the periodic dimensions and
real-to-real discrete cosine transforms are used in the wall-bounded dimensions.
Note that the DCT-II is used for the DCT and the DCT-III for the IDCT
which correspond to REDFT10 and REDFT01 in FFTW.

They operatore on an array with the shape of `A`, which is needed to plan
efficient transforms. `A` will be mutated.
=#

plan_forward_transform(A::Array, ::PBC, dims, planner_flag=FFTW.PATIENT) =
    FFTW.plan_fft!(A, dims, flags=planner_flag)

plan_forward_transform(A::Array, ::ZFBC, dims, planner_flag=FFTW.PATIENT) =
    FFTW.plan_r2r!(A, FFTW.REDFT10, dims, flags=planner_flag)

plan_backward_transform(A::Array, ::PBC, dims, planner_flag=FFTW.PATIENT) =
    FFTW.plan_ifft!(A, dims, flags=planner_flag)

plan_backward_transform(A::Array, ::ZFBC, dims, planner_flag=FFTW.PATIENT) =
    FFTW.plan_r2r!(A, FFTW.REDFT01, dims, flags=planner_flag)

@hascuda begin
     plan_forward_transform(A::CuArray, ::Union{PBC, ZFBC}, dims) = plan_fft!(A, dims)
    plan_backward_transform(A::CuArray, ::Union{PBC, ZFBC}, dims) = plan_ifft!(A, dims)
end


function k(α,β,γ)

			x̂ = cartesian(α)
			ξ = cartesian(β)[1]
			η = cartesian(γ)[1]

			ŷ = [ξ*(1-η), η]

			#jacobian of y--->ŷ = (1-y[2])

			chartI = simplex(p1,x̂,p0)
			chartII = simplex(p2,x̂,p1)
			chartIII = simplex(p0,x̂,p2)

			n1 = neighborhood(chartI, ŷ)
			n2 = neighborhood(chartII, ŷ)
			n3 = neighborhood(chartIII, ŷ)

			yI = cartesian(n1)
			yII = cartesian(n2)
			yIII = cartesian(n3)

			return(Kernel(x̂,yI)*jacobian(n1)*(1-η) +
					Kernel(x̂,yII)*jacobian(n2)*(1-η) +
					Kernel(x̂,yIII)*jacobian(n3)*(1-η))
end

function verifintegral1(sourcechart::CompScienceMeshes.Simplex{3,2,1,3,Float64},
	 testchart::CompScienceMeshes.Simplex{3,2,1,3,Float64}, integrand, accuracy::Int64)

	global Kernel, p0, p1, p2

	Kernel = integrand

	p0 = Sourcechart.vertices[1]
	p1 = Sourcechart.vertices[2]
	p2 = Sourcechart.vertices[3]

	qps1 = quadpoints(Sourcechart, accuracy)

	path = simplex(point(0), point(1))
	qps2 = quadpoints(path, accuracy)

	result = sum(w*w2*w1*k(α,β,γ) for (β,w1) in qps2, (γ,w2) in qps2, (α,w) in qps1)

	return (result)

end

function verifintegral2(sourcechart::CompScienceMeshes.Simplex{3,2,1,3,Float64},
	 testchart::CompScienceMeshes.Simplex{3,2,1,3,Float64}, integrand, accuracy::Int64)

	qps1 = quadpoints(sourcechart, accuracy)
	qps2 = quadpoints(testchart, accuracy)

	result = sum(w2*w1*integrand(cartesian(x),cartesian(y))
				for (x,w2) in qps2, (y,w1) in qps1)

	return (result)

end


"""
    BigleafConstants(;...)

Constants used troughout the Bigleaf.jl Package

Default values can be overridden by the named arguments of the constructor:

- cp           : Specific heat of air for constant pressure (J K-1 kg-1)
- Rgas         : Universal gas constant (J mol-1 K-1)
- Rv           : Gas constant of water vapor (J kg-1 K-1) (Stull 1988 p.641)
- Rd           : Gas constant of dry air (J kg-1 K-1) (Foken p. 245)
- Md           : Molar mass of dry air (kg mol-1)
- Mw           : Molar mass of water vapor (kg mol-1)
- eps          : Ratio of the molecular weight of water vapor to dry air (=Mw/Md)
- g            : Gravitational acceleration (m s-2)
- solar_constant: Solar constant (W m-2)
- pressure0    : Reference atmospheric pressure at sea level (Pa)
- Tair0        : Reference air temperature (K)
- k            : von Karman constant
- Cmol         : Molar mass of carbon (kg mol-1)
- Omol         : Molar mass of oxygen (kg mol-1)
- H2Omol       : Molar mass of water (kg mol-1)
- sigma        : Stefan-Boltzmann constant (W m-2 K-4)
- Pr           : Prandtl number
- Sc_CO2       : Schmidt number for CO2
- Kelvin       : Conversion degree Celsius to Kelvin
- DwDc         : Ratio of the molecular diffusivities for water vapor and CO2
- days2seconds : Seconds per day
- kPa2Pa       : Conversion kilopascal (kPa) to pascal (Pa)
- Pa2kPa       : Conversion pascal (Pa) to kilopascal (kPa)
- umol2mol     : Conversion micromole (umol) to mole (mol)
- mol2umol     : Conversion mole (mol) to micromole (umol)
- kg2g         : Conversion kilogram (kg) to gram (g)
- g2kg         : Conversion gram (g) to kilogram (kg)
- kJ2J         : Conversion kilojoule (kJ) to joule (J)
- J2kJ         : Conversion joule (J) to kilojoule (kJ)
- se_median    : Conversion standard error (SE) of the mean to SE of the median
  (http://influentialpoints.com/Training/standard_error_of_median.htm)
- frac2percent : Conversion between fraction and percent

## Examples
```jldoctest; output=false
BigleafConstants().g == 9.81
# on the moon change gravity constant to 1/6 that of the earth
BigleafConstants(g = BigleafConstants().g/6).g == 9.81/6
# output
true
```
"""
@with_kw struct BigleafConstants{FT,IT,RT}
  cp::FT         = 1004.834        # specific heat of air for constant pressure (J K-1 kg-1)
  Rgas::FT       = 8.31451         # universal gas constant (J mol-1 K-1)
  Rv::FT         = 461.5           # gas constant of water vapor (J kg-1 K-1) (Stull 1988 p.641)
  Rd::FT         = 287.0586        # gas constant of dry air (J kg-1 K-1) (Foken 2008 p. 245)
  Md::FT         = 0.0289645       # molar mass of dry air (kg mol-1)
  Mw::FT         = 0.0180153       # molar mass of water vapor (kg mol-1)
  eps::FT        = 0.622           # ratio of the molecular weight of water vapor to dry air
    # (=Mw/Md)
  g::FT          = 9.81            # gravitational acceleration (m s-2)
  solar_constant::FT = 1366.1      # solar constant, i.e. solar radation at earth distance from
  # the sun (W m-2)
  pressure0::FT  = 101325.0        # reference atmospheric pressure at sea level (Pa)
  Tair0::FT      = 273.15          # reference air temperature (K)
  k::FT          = 0.41            # von Karman constant
  Cmol::FT       = 0.012011        # molar mass of carbon (kg mol-1)
  Omol::FT       = 0.0159994       # molar mass of oxygen (kg mol-1)
  H2Omol::FT     = 0.01801528      # molar mass of water (kg mol-1)
  sigma::FT      = 5.670367e-08    # Stefan-Boltzmann constant (W m-2 K-4)
  Pr::FT         = 0.71            # Prandtl number
  Sc_CO2::FT     = 1.07            # Schmidt number for CO2 (Hicks et al. 1987)
  ## Conversion constants
  Kelvin::FT       = 273.15        # conversion degree Celsius to Kelvin
  DwDc::FT         = 1.6           # Ratio of the molecular diffusivities for water vapor and CO2
  days2seconds::IT = 86400         # seconds per day
  kPa2Pa::IT       = 1000          # conversion kilopascal (kPa) to pascal (Pa)
  Pa2kPa::RT       = 1/1000        # conversion pascal (Pa) to kilopascal (kPa)
  umol2mol::RT     = 1/1e6         # conversion micromole (umol) to mole (mol)
  mol2umol::IT     = 1_000_000     # conversion mole (mol) to micromole (umol)
  kg2g::IT         = 1000          # conversion kilogram (kg) to gram (g)
  g2kg::RT         = 1/1000        # conversion gram (g) to kilogram (kg)
  kJ2J::IT         = 1000          # conversion kilojoule (kJ) to joule (J)
  J2kJ::RT         = 1/1000        # conversion joule (J) to kilojoule (kJ)
  se_median::FT    = 1.253         # conversion standard error (SE) of the mean to SE of the
  # median (http://influentialpoints.com/Training/standard_error_of_median.htm)
  frac2percent::IT = 100           # conversion between fraction and percent
end


## Packages and settings.
include("../src/StencilBuilder.jl")

using BenchmarkTools
using LoopVectorization
using .StencilBuilder

float_type = Float32

## CPU dynamics kernel.
function dynamics_w_kernel!(
    wt, u, v, w, s,
    visc,
    dxi, dyi, dzi, dzhi,
    is, ie, js, je, ks, ke)

    alpha = 9.81/300;

    @fast3d begin
        @fd (wt, u, v, w, s) wt += (
            - gradx(interpz(u) * interpx(w)) + visc * (gradx(gradx(w)))
            - grady(interpz(v) * interpy(w)) + visc * (grady(grady(w)))
            - gradz(interpz(w) * interpz(w)) + visc * (gradz(gradz(w)))
            + alpha*interpz(s) )
    end
end

## User input
itot = 384; jtot = 384; ktot = 384;
igc = 1; jgc = 1; kgc = 1;

icells = itot + 2igc; jcells = jtot + 2jgc; kcells = ktot + 2kgc
is = igc+1; ie = igc+itot; js = jgc+1; je = jgc+jtot; ks = kgc+1; ke = kgc+ktot

wt = rand(float_type, icells, jcells, kcells)
u = rand(float_type, icells, jcells, kcells)
v = rand(float_type, icells, jcells, kcells)
w = rand(float_type, icells, jcells, kcells)
s = rand(float_type, icells, jcells, kcells)
dxi = rand(float_type)
dyi = rand(float_type)
dzi = rand(float_type, kcells)
dzhi = rand(float_type, kcells)
visc = convert(float_type, 1)

## Run CPU kernel
@btime dynamics_w_kernel!(
    $wt, $u, $v, $w, $s,
    $visc,
    $dxi, $dyi, $dzi, $dzhi,
    $is, $ie, $js, $je, $ks, $ke)


@testset "Hessian of Lagrangian" begin
    MOI = DTO.MOI
    # horizon
    T = 3

    # acrobot
    num_state = 4
    num_action = 1
    num_parameter = 0
    parameter_dimensions = [num_parameter for t = 1:T]

    function acrobot(x, u, w)
        # dimensions
        n = 4
        m = 1
        d = 0

        # link 1
        mass1 = 1.0
        inertia1 = 0.33
        length1 = 1.0
        lengthcom1 = 0.5

        # link 2
        mass2 = 1.0
        inertia2 = 0.33
        length2 = 1.0
        lengthcom2 = 0.5

        gravity = 9.81
        friction1 = 0.1
        friction2 = 0.1

        # mass matrix
        function M(x, w)
            a = (inertia1 + inertia2 + mass2 * length1 * length1
                + 2.0 * mass2 * length1 * lengthcom2 * cos(x[2]))

            b = inertia2 + mass2 * length1 * lengthcom2 * cos(x[2])

            c = inertia2

            return [a b; b c]
        end

        function Minv(x, w)
            a = (inertia1 + inertia2 + mass2 * length1 * length1
                + 2.0 * mass2 * length1 * lengthcom2 * cos(x[2]))

            b = inertia2 + mass2 * length1 * lengthcom2 * cos(x[2])

            c = inertia2

            return 1.0 / (a * c - b * b) * [c -b; -b a]
        end

        # dynamics bias
        function τ(x, w)
            a = (-1.0 * mass1 * gravity * lengthcom1 * sin(x[1])
                - mass2 * gravity * (length1 * sin(x[1])
                + lengthcom2 * sin(x[1] + x[2])))

            b = -1.0 * mass2 * gravity * lengthcom2 * sin(x[1] + x[2])

            return [a; b]
        end

        function C(x, w)
            a = -2.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
            b = -1.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
            c = mass2 * length1 * lengthcom2 * sin(x[2]) * x[3]
            d = 0.0

            return [a b; c d]
        end

        # input Jacobian
        function B(x, w)
            [0.0; 1.0]
        end

        # dynamics
        q = view(x, 1:2)
        v = view(x, 3:4)

        qdd = Minv(q, w) * (-1.0 * C(x, w) * v
                + τ(q, w) + B(q, w) * u[1] - [friction1; friction2] .* v)

        return [x[3]; x[4]; qdd[1]; qdd[2]]
    end

    function midpoint_implicit(y, x, u, w)
        h = 0.05 # timestep
        y - (x + h * acrobot(0.5 * (x + y), u, w))
    end

    dt = Dynamics(midpoint_implicit, num_state, num_state, num_action, num_parameter=num_parameter, evaluate_hessian=true)
    dynamics = [dt for t = 1:T-1]

    # initial state
    x1 = [0.0; 0.0; 0.0; 0.0]

    # goal state
    xT = [0.0; π; 0.0; 0.0]

    # objective
    ot = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4]) + 0.1 * dot(u, u)
    oT = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4])
    objt = Cost(ot, num_state, num_action, num_parameter=num_parameter, evaluate_hessian=true)
    objT = Cost(oT, num_state, 0, num_parameter=num_parameter, evaluate_hessian=true)
    objective = [[objt for t = 1:T-1]..., objT]

    # constraints
    bnd1 = Bound(num_state, num_action, state_lower=x1, state_upper=x1)
    bndt = Bound(num_state, num_action)
    bndT = Bound(num_state, 0, state_lower=xT, state_upper=xT)
    bounds = [bnd1, [bndt for t = 2:T-1]..., bndT]

    ct = (x, u, w) -> [-5.0 * ones(num_action) - cos.(u) .* sum(x.^2); cos.(x) .* tan.(u) - 5.0 * ones(num_state)]
    cT = (x, u, w) -> sin.(x.^3.0)
    cont = Constraint(ct, num_state, num_action, num_parameter=num_parameter, indices_inequality=collect(1:(num_action + num_state)), evaluate_hessian=true)
    conT = Constraint(cT, num_state, 0, num_parameter=num_parameter, evaluate_hessian=true)
    constraints = [[cont for t = 1:T-1]..., conT]

    # data
    # trajopt = DTO.TrajectoryOptimizationData(objective, dynamics, constraints, bounds)
    # nlp = DTO.NLPData(trajopt, evaluate_hessian=true)
    solver = Solver(dynamics, objective, constraints, bounds, evaluate_hessian=true)

    # Lagrangian
    function lagrangian(z)
        x1 = z[1:num_state]
        u1 = z[num_state .+ (1:num_action)]
        x2 = z[num_state + num_action .+ (1:num_state)]
        u2 = z[num_state + num_action + num_state .+ (1:num_action)]
        x3 = z[num_state + num_action + num_state + num_action .+ (1:num_state)]
        λ1_dyn = z[num_state + num_action + num_state + num_action + num_state .+ (1:num_state)]
        λ2_dyn = z[num_state + num_action + num_state + num_action + num_state + num_state .+ (1:num_state)]

        λ1_stage = z[num_state + num_action + num_state + num_action + num_state + num_state + num_state .+ (1:(num_action + num_state))]
        λ2_stage = z[num_state + num_action + num_state + num_action + num_state + num_state + num_state + num_action + num_state .+ (1:(num_action + num_state))]
        λ3_stage = z[num_state + num_action + num_state + num_action + num_state + num_state + num_state + num_action + num_state + num_action + num_state .+ (1:num_state)]

        L = 0.0
        L += ot(x1, u1, zeros(num_parameter))
        L += ot(x2, u2, zeros(num_parameter))
        L += oT(x3, zeros(0), zeros(num_parameter))
        L += dot(λ1_dyn, midpoint_implicit(x2, x1, u1, zeros(num_parameter)))
        L += dot(λ2_dyn, midpoint_implicit(x3, x2, u2, zeros(num_parameter)))
        L += dot(λ1_stage, ct(x1, u1, zeros(num_parameter)))
        L += dot(λ2_stage, ct(x2, u2, zeros(num_parameter)))
        L += dot(λ3_stage, cT(x3, zeros(0), zeros(num_parameter)))
        return L
    end

    nz = num_state + num_action + num_state + num_action + num_state + num_state + num_state + num_action + num_state + num_action + num_state + num_state
    np = num_state + num_action + num_state + num_action + num_state
    nd = num_state + num_state + num_action + num_state + num_action + num_state + num_state
    @variables z[1:nz]
    L = lagrangian(z)
    Lxx = Symbolics.hessian(L, z[1:np])
    Lxx_sp = Symbolics.sparsehessian(L, z[1:np])
    spar = [findnz(Lxx_sp)[1:2]...]
    Lxx_func = eval(Symbolics.build_function(Lxx, z)[1])
    Lxx_sp_func = eval(Symbolics.build_function(Lxx_sp.nzval, z)[1])

    z0 = rand(nz)
    nh = length(solver.nlp.hessian_lagrangian_sparsity)
    h0 = zeros(nh)

    σ = 1.0
    fill!(h0, 0.0)
    DTO.trajectory!(solver.nlp.trajopt.states, solver.nlp.trajopt.actions, z0[1:np],
        solver.nlp.indices.states, solver.nlp.indices.actions)
    DTO.duals!(solver.nlp.trajopt.duals_dynamics, solver.nlp.trajopt.duals_constraints, solver.nlp.duals_general, z0[np .+ (1:nd)], solver.nlp.indices.dynamics_constraints, solver.nlp.indices.stage_constraints, solver.nlp.indices.general_constraint)
    DTO.hessian!(h0, solver.nlp.indices.objective_hessians, solver.nlp.trajopt.objective, solver.nlp.trajopt.states, solver.nlp.trajopt.actions, solver.nlp.trajopt.parameters, σ)
    DTO.hessian_lagrangian!(h0, solver.nlp.indices.dynamics_hessians, solver.nlp.trajopt.dynamics, solver.nlp.trajopt.states, solver.nlp.trajopt.actions, solver.nlp.trajopt.parameters, solver.nlp.trajopt.duals_dynamics)
    DTO.hessian_lagrangian!(h0, solver.nlp.indices.stage_hessians, solver.nlp.trajopt.constraints, solver.nlp.trajopt.states, solver.nlp.trajopt.actions, solver.nlp.trajopt.parameters, solver.nlp.trajopt.duals_constraints)

    sparsity_objective_hessians = DTO.sparsity_hessian(objective, solver.nlp.trajopt.state_dimensions, solver.nlp.trajopt.action_dimensions)
    sparsity_dynamics_hessians = DTO.sparsity_hessian(dynamics, solver.nlp.trajopt.state_dimensions, solver.nlp.trajopt.action_dimensions)
    sparsity_constraint_hessian = DTO.sparsity_hessian(constraints, solver.nlp.trajopt.state_dimensions, solver.nlp.trajopt.action_dimensions)
    hessian_sparsity = collect([sparsity_objective_hessians..., sparsity_dynamics_hessians..., sparsity_constraint_hessian...])
    sp_key = sort(unique(hessian_sparsity))

    idx_objective_hessians = DTO.hessian_indices(objective, sp_key, solver.nlp.trajopt.state_dimensions, solver.nlp.trajopt.action_dimensions)
    idx_dynamics_hessians = DTO.hessian_indices(dynamics, sp_key, solver.nlp.trajopt.state_dimensions, solver.nlp.trajopt.action_dimensions)
    idx_con_hess = DTO.hessian_indices(constraints, sp_key, solver.nlp.trajopt.state_dimensions, solver.nlp.trajopt.action_dimensions)

    # indices
    @test sp_key[vcat(idx_objective_hessians...)] == sparsity_objective_hessians
    @test sp_key[vcat(idx_dynamics_hessians...)] == sparsity_dynamics_hessians
    @test sp_key[vcat(idx_con_hess...)] == sparsity_constraint_hessian

    # Hessian
    h0_full = zeros(np, np)
    for (i, h) in enumerate(h0)
        h0_full[sp_key[i][1], sp_key[i][2]] = h
    end
    @test norm(h0_full - Lxx_func(z0)) < 1.0e-8
    @test norm(norm(h0 - Lxx_sp_func(z0))) < 1.0e-8

    h0 = zeros(nh)
    MOI.eval_hessian_lagrangian(solver.nlp, h0, z0[1:np], σ, z0[np .+ (1:nd)])
    @test norm(norm(h0 - Lxx_sp_func(z0))) < 1.0e-8

    # a = z0[1:np]
    # b = z0[np .+ (1:nd)]
    # info = @benchmark MOI.evaluate_hessianian_lagrangian($solver, $h0, $a, $σ, $b)
end


using Flux
using Flux: @epochs
using Flux: params
using Flux: throttle
using Plots
using Distributions
using DataFrames
using CSV
using StaticArrays

@inline phi(x) = max(x^3,0)
@inline f(x) = 1

#build a block
block_layer = SkipConnection(Chain(Dense(10,10, phi),Dense(10,10, phi)), +)

#construct network skeleton
M = Chain(Dense(2,10),
block_layer,
block_layer,
block_layer,
block_layer,
Dense(10,1))

#check number of params
params(M)

#inner loss function
function inner_loss(N, space)
    a = rand(Uniform(-1,1),2,N)
    loss = 0
    wt = 4*ones(N)/N
    l11 = [1, 0]
    l21 = [2, 0]
    l12 = [0, 1]
    l22 = [0, 2]

    @inbounds loss += ((0.5*((-1/12*M(a .+ space.*l21) .+ 2/3*M(a .+ space.*l11) .- 2/3*M(a .- space.*l11) .+ 1/12*M(a .- space.*l21))./space).^2) *wt)[1]
    @inbounds loss += ((0.5*((-1/12*M(a .+ space.*l22) .+ 2/3*M(a .+ space.*l12) .- 2/3*M(a .- space.*l12) .+ 1/12*M(a .- space.*l22))./space).^2) *wt)[1]
    @inbounds loss += (-1 .*M(a)*wt)[1]

    return loss
end

#boundary condition loss function
#add [0,1) as a boundary
function bc_loss(N, β)
    dist = Uniform(-1,1)
    a = hcat(vcat(rand(dist, 1,floor(Int32, N/5)), ones(1,floor(Int32, N/5))),vcat(rand(dist, 1,floor(Int32, N/5)), ones(1,floor(Int32, N/5)).*-1),
    vcat(ones(1,floor(Int32, N/5)), rand(dist, 1,floor(Int32, N/5))), vcat(ones(1,floor(Int32, N/5)).*-1, rand(dist, 1,floor(Int32, N/5))),
    vcat(zeros(1,floor(Int32, N/5)), rand(1,floor(Int32, N/5))))
    out_original = M(a)
    loss = (β/N)*sum(abs2, out_original)
    return loss
end

#add losses
function total_loss(N, space = 0.01, β = 750)
    return inner_loss(N, space) + bc_loss(floor(Int32, N/2), β)
end

#train model
function train_model(N = 100, train_time = 20)
    df = DataFrame(time = Int[], in_loss = Float64[],bound_loss = Float64[])
    c = 1
    n = floor(Int32, N/2)
    evalcb = function()
        @show(total_loss(N))
        in_loss = inner_loss(N, 0.01)
        bound_loss = bc_loss(n, 750)
        push!(df, [c, in_loss, bound_loss])
        c += 1
        if c == train_time
            CSV.write("$(c)_2d_cutsq.csv",df)
        end
    end
    opt = ADAM()
    dataset = [(N) for i=1:train_time]
    Flux.train!(total_loss, params(M), zip(dataset), opt, cb=throttle(evalcb, 10))
end

@epochs 1000 train_model(4000)

#plot the solution over a mesh in the domain
function plot_surface()
    data = zeros(201,201)
    grid_space = [i for i in -100:100]
    for i in grid_space
        for j in grid_space
            point = [i/100 j/100]
            @inbounds data[i+101,j+101] = M(point')[1]
        end
    end

    gr()
    heatmap(1:size(data,1),
        1:size(data,2), data,
        c = cgrad([:blue, :white,:red, :yellow]),
        xlabel = "X", ylabel = "Y",
        title = "2D Cut Square Fitted Solution")

    savefig("2d_cutsq_fitted.png")
end

plot_surface()


module clcircularplatecl_examples
using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtools.MeshExportModule
using Statistics: mean

# Clamped square plate with concentrated force
#  Data listed in the Simo 1990 paper 'A class of... '
#  Analytical solution for the vertical deflection under the load
#    Zienkiewicz, Taylor, The finite element method, fifth edition, volume 2,
#     analyt_sol = 3*(1-nu^2)*Magnitude*R^2/(4*pi*E*thickness^3);
E = 1e7;
nu = 0.3;
Magnitude = 10;
R = 100.0;
thickness = 2.5;
tolerance = 0.0001*thickness;
@show analyt_sol = 3*(1-nu^2)*Magnitude*R^2/(4*pi*E*thickness^3);

function clcircularplatecl_h8_full()
	nt = 1
	for nperradius in [2, 4, 8]
		nt=nt+1;
    	fens, fes  = Q4circlen(R, nperradius)
    	fens, fes = H8extrudeQ4(fens, fes, nt, (x, k) -> [x[1],x[2],k*thickness/nt])

    	MR = DeforModelRed3D
    	material = MatDeforElastIso(MR, E, nu)
    	femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3, 2)), material)

    	boundaryfes = meshboundary(fes);
    	topl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf thickness thickness], inflate =   tolerance);
    	botl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf 0.0 0.0], inflate =   tolerance);
    	x0l = selectelem(fens,boundaryfes, box =  [0.0 0.0 -Inf Inf 0.0 thickness], inflate =   tolerance);
    	y0l = selectelem(fens,boundaryfes, box =  [-Inf Inf 0.0 0.0 0.0 thickness], inflate =   tolerance);
    	cyll = setdiff(1:count(boundaryfes),topl,botl,x0l,y0l);

    	geom = NodalField(fens.xyz)
    	u = NodalField(zeros(size(fens.xyz,1), 3)) # displacement field

    	x0nl = connectednodes(subset(boundaryfes, x0l))
    	setebc!(u,x0nl,true,1,0.0)

    	y0nl = connectednodes(subset(boundaryfes, y0l))
    	setebc!(u,y0nl,true,2,0.0)

    	cylnl = connectednodes(subset(boundaryfes, cyll))
    	setebc!(u,cylnl,true,1,0.0)
    	setebc!(u,cylnl,true,2,0.0)
    	setebc!(u,cylnl,true,3,0.0)

    	applyebc!(u)
    	numberdofs!(u)

    	cnl = selectnode(fens; box = [0 0 0 0 0 thickness], inflate  =  tolerance)

    	nfemm = FEMMBase(IntegDomain(FESetP1(reshape(cnl, length(cnl), 1)), PointRule()))
  		F = distribloads(nfemm, geom, u, ForceIntensity([0; 0; Magnitude/4/length(cnl)]), 3)

  		associategeometry!(femm, geom)

  		K = stiffness(femm, geom, u)

  		scattersysvec!(u, K\F)

  		u0z = mean(u.values[cnl, 3]);
  		println("Deflection under the load: $(round((u0z / analyt_sol)* 100000)/100000*100) %")

    	# File =  "clcircularplatecl_1_$(nperradius).vtk"
    	# vtkexportmesh(File, fens, fes; vectors = [("u", u.values)])
    	# @async run(`"paraview.exe" $File`)

    end

    true
end # clcircularplatecl_h8_full

function clcircularplatecl_h8_uri()
	nt = 1
	for nperradius in [2, 4, 8]
		nt=nt+1;
		fens, fes  = Q4circlen(R, nperradius)
		fens, fes = H8extrudeQ4(fens, fes, nt, (x, k) -> [x[1],x[2],k*thickness/nt])

    	MR = DeforModelRed3D
    	material = MatDeforElastIso(MR, E, nu)
    	femm = FEMMDeforLinear(MR, IntegDomain(fes, GaussRule(3, 1)), material)

    	boundaryfes = meshboundary(fes);
    	topl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf thickness thickness], inflate =   tolerance);
    	botl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf 0.0 0.0], inflate =   tolerance);
    	x0l = selectelem(fens,boundaryfes, box =  [0.0 0.0 -Inf Inf 0.0 thickness], inflate =   tolerance);
    	y0l = selectelem(fens,boundaryfes, box =  [-Inf Inf 0.0 0.0 0.0 thickness], inflate =   tolerance);
    	cyll = setdiff(1:count(boundaryfes),topl,botl,x0l,y0l);

    	geom = NodalField(fens.xyz)
    	u = NodalField(zeros(size(fens.xyz,1), 3)) # displacement field

    	x0nl = connectednodes(subset(boundaryfes, x0l))
    	setebc!(u,x0nl,true,1,0.0)

    	y0nl = connectednodes(subset(boundaryfes, y0l))
    	setebc!(u,y0nl,true,2,0.0)

    	cylnl = connectednodes(subset(boundaryfes, cyll))
    	setebc!(u,cylnl,true,1,0.0)
    	setebc!(u,cylnl,true,2,0.0)
    	setebc!(u,cylnl,true,3,0.0)

    	applyebc!(u)
    	numberdofs!(u)

    	cnl = selectnode(fens; box = [0 0 0 0 0 thickness], inflate  =  tolerance)

    	nfemm = FEMMBase(IntegDomain(FESetP1(reshape(cnl, length(cnl), 1)), PointRule()))
  		F = distribloads(nfemm, geom, u, ForceIntensity([0; 0; Magnitude/4/length(cnl)]), 3)

  		associategeometry!(femm, geom)

  		K = stiffness(femm, geom, u)

  		scattersysvec!(u, K\F)

  		u0z = mean(u.values[cnl, 3]);
  		println("Deflection under the load: $(round((u0z / analyt_sol)* 100000)/100000*100) %")

    	# File =  "clcircularplatecl_1_$(nperradius).vtk"
    	# vtkexportmesh(File, fens, fes; vectors = [("u", u.values)])
    	# @async run(`"paraview.exe" $File`)

    end

    true
end # clcircularplatecl_h8_uri

function clcircularplatecl_h8_ms()
    nt = 1
    for nperradius in [2, 4, 8]
    	nt=nt+1;
    	fens, fes  = Q4circlen(R, nperradius)
    	fens, fes = H8extrudeQ4(fens, fes, nt, (x, k) -> [x[1],x[2],k*thickness/nt])

    	MR = DeforModelRed3D
    	material = MatDeforElastIso(MR, E, nu)
    	femm = FEMMDeforLinearMSH8(MR, IntegDomain(fes, GaussRule(3, 2)), material)

    	boundaryfes = meshboundary(fes);
    	topl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf thickness thickness], inflate =   tolerance);
    	botl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf 0.0 0.0], inflate =   tolerance);
    	x0l = selectelem(fens,boundaryfes, box =  [0.0 0.0 -Inf Inf 0.0 thickness], inflate =   tolerance);
    	y0l = selectelem(fens,boundaryfes, box =  [-Inf Inf 0.0 0.0 0.0 thickness], inflate =   tolerance);
    	cyll = setdiff(1:count(boundaryfes),topl,botl,x0l,y0l);

    	geom = NodalField(fens.xyz)
    	u = NodalField(zeros(size(fens.xyz,1), 3)) # displacement field

    	x0nl = connectednodes(subset(boundaryfes, x0l))
    	setebc!(u,x0nl,true,1,0.0)

    	y0nl = connectednodes(subset(boundaryfes, y0l))
    	setebc!(u,y0nl,true,2,0.0)

    	cylnl = connectednodes(subset(boundaryfes, cyll))
    	setebc!(u,cylnl,true,1,0.0)
    	setebc!(u,cylnl,true,2,0.0)
    	setebc!(u,cylnl,true,3,0.0)

    	applyebc!(u)
    	numberdofs!(u)

    	cnl = selectnode(fens; box = [0 0 0 0 0 thickness], inflate  =  tolerance)

    	nfemm = FEMMBase(IntegDomain(FESetP1(reshape(cnl, length(cnl), 1)), PointRule()))
  		F = distribloads(nfemm, geom, u, ForceIntensity([0; 0; Magnitude/4/length(cnl)]), 3)

  		associategeometry!(femm, geom)
  		@show minimum(femm.phis), maximum(femm.phis)

  		K = stiffness(femm, geom, u)

  		scattersysvec!(u, K\F)

  		u0z = mean(u.values[cnl, 3]);
  		println("Deflection under the load: $(round((u0z / analyt_sol)* 100000)/100000*100) %")

    	# File =  "clcircularplatecl_1_$(nperradius).vtk"
    	# vtkexportmesh(File, fens, fes; vectors = [("u", u.values)])
    	# @async run(`"paraview.exe" $File`)

    end

    true
end # clcircularplatecl_h8_ms

function clcircularplatecl_h8u_ms()
    nt = 1
    for nperradius in [2, 4, 8]
    	nt=nt+1;
    	fens,fes = T4quartercyln(R, thickness, nperradius, nt)
    	fens,fes = T4toH8(fens,fes)

    	MR = DeforModelRed3D
    	material = MatDeforElastIso(MR, E, nu)
    	femm = FEMMDeforLinearMSH8(MR, IntegDomain(fes, GaussRule(3, 2)), material)

    	boundaryfes = meshboundary(fes);
    	topl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf thickness thickness], inflate =   tolerance);
    	botl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf 0.0 0.0], inflate =   tolerance);
    	x0l = selectelem(fens,boundaryfes, box =  [0.0 0.0 -Inf Inf 0.0 thickness], inflate =   tolerance);
    	y0l = selectelem(fens,boundaryfes, box =  [-Inf Inf 0.0 0.0 0.0 thickness], inflate =   tolerance);
    	cyll = setdiff(1:count(boundaryfes),topl,botl,x0l,y0l);

    	geom = NodalField(fens.xyz)
    	u = NodalField(zeros(size(fens.xyz,1), 3)) # displacement field

    	x0nl = connectednodes(subset(boundaryfes, x0l))
    	setebc!(u,x0nl,true,1,0.0)

    	y0nl = connectednodes(subset(boundaryfes, y0l))
    	setebc!(u,y0nl,true,2,0.0)

    	cylnl = connectednodes(subset(boundaryfes, cyll))
    	setebc!(u,cylnl,true,1,0.0)
    	setebc!(u,cylnl,true,2,0.0)
    	setebc!(u,cylnl,true,3,0.0)

    	applyebc!(u)
    	numberdofs!(u)

    	cnl = selectnode(fens; box = [0 0 0 0 0 thickness], inflate  =  tolerance)

    	nfemm = FEMMBase(IntegDomain(FESetP1(reshape(cnl, length(cnl), 1)), PointRule()))
  		F = distribloads(nfemm, geom, u, ForceIntensity([0; 0; Magnitude/4/length(cnl)]), 3)

  		associategeometry!(femm, geom)
  		@show minimum(femm.phis), maximum(femm.phis)

  		K = stiffness(femm, geom, u)

  		scattersysvec!(u, K\F)

  		u0z = mean(u.values[cnl, 3]);
  		println("Deflection under the load: $(round((u0z / analyt_sol)* 100000)/100000*100) %")

    	File =  "clcircularplatecl_h8u_ms_$(nperradius).vtk"
    	vtkexportmesh(File, fens, fes; vectors = [("u", u.values)])
    	@async run(`"paraview.exe" $File`)

    end

    true
end # clcircularplatecl_h8u_ms

function clcircularplatecl_h8_export()
    nt = 1
    for nperradius in [2, 4, 8]
    	nt=nt+1;
    	fens, fes  = Q4circlen(R, nperradius)
    	fens, fes = H8extrudeQ4(fens, fes, nt, (x, k) -> [x[1],x[2],k*thickness/nt])

    	MR = DeforModelRed3D
    	material = MatDeforElastIso(MR, E, nu)
    	femm = FEMMDeforLinearMSH8(MR, IntegDomain(fes, GaussRule(3, 2)), material)

    	boundaryfes = meshboundary(fes);
    	topl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf thickness thickness], inflate =   tolerance);
    	botl = selectelem(fens,boundaryfes, box =  [-Inf Inf -Inf Inf 0.0 0.0], inflate =   tolerance);
    	x0l = selectelem(fens,boundaryfes, box =  [0.0 0.0 -Inf Inf 0.0 thickness], inflate =   tolerance);
    	y0l = selectelem(fens,boundaryfes, box =  [-Inf Inf 0.0 0.0 0.0 thickness], inflate =   tolerance);
    	cyll = setdiff(1:count(boundaryfes),topl,botl,x0l,y0l);

    	geom = NodalField(fens.xyz)
    	u = NodalField(zeros(size(fens.xyz,1), 3)) # displacement field

    	x0nl = connectednodes(subset(boundaryfes, x0l))
    	setebc!(u,x0nl,true,1,0.0)

    	y0nl = connectednodes(subset(boundaryfes, y0l))
    	setebc!(u,y0nl,true,2,0.0)

    	cylnl = connectednodes(subset(boundaryfes, cyll))
    	setebc!(u,cylnl,true,1,0.0)
    	setebc!(u,cylnl,true,2,0.0)
    	setebc!(u,cylnl,true,3,0.0)

    	applyebc!(u)
    	numberdofs!(u)

    	cnl = selectnode(fens; box = [0 0 0 0 0 thickness], inflate  =  tolerance)

    	AE = AbaqusExporter("clcircularplatecl_h8_export_$(nperradius)");
    	HEADING(AE, "Clamped square plate with concentrated force");
    	PART(AE, "part1");
    	END_PART(AE);
    	ASSEMBLY(AE, "ASSEM1");
    	INSTANCE(AE, "INSTNC1", "PART1");
    	NODE(AE, fens.xyz);
    	ELEMENT(AE, "c3d8rh", "AllElements", 1, connasarray(femm.integdomain.fes))
    	NSET_NSET(AE, "cnl", cnl)
    	NSET_NSET(AE, "x0nl", x0nl)
    	NSET_NSET(AE, "y0nl", y0nl)
    	NSET_NSET(AE, "cylnl", cylnl)
    	ORIENTATION(AE, "GlobalOrientation", vec([1. 0 0]), vec([0 1. 0]));
    	SOLID_SECTION(AE, "elasticity", "GlobalOrientation", "AllElements", "Hourglassctl");
    	END_INSTANCE(AE);
    	END_ASSEMBLY(AE);
    	MATERIAL(AE, "elasticity")
    	ELASTIC(AE, E, nu)
    	SECTION_CONTROLS(AE, "Hourglassctl", "HOURGLASS=ENHANCED")
    	STEP_PERTURBATION_STATIC(AE)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.x0nl", 1)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.y0nl", 2)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.cylnl", 1)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.cylnl", 2)
    	BOUNDARY(AE, "ASSEM1.INSTNC1.cylnl", 3)
    	CLOAD(AE, "ASSEM1.INSTNC1.cnl", 3, Magnitude/4/length(cnl))
    	END_STEP(AE)
    	close(AE)

    	# File =  "clcircularplatecl_h8_export_$(nperradius).vtk"
    	# vtkexportmesh(File, fens, fes; vectors = [("u", u.values)])
    	# @async run(`"paraview.exe" $File`)

    end

    true
end # clcircularplatecl_h8_export

#     true
# end # twisted_beam_export

function allrun()
    println("#####################################################")
    println("# clcircularplatecl_h8_full ")
    clcircularplatecl_h8_full()
    println("#####################################################")
    println("# clcircularplatecl_h8_uri ")
    clcircularplatecl_h8_uri()
    println("#####################################################")
    println("# clcircularplatecl_h8_ms ")
    clcircularplatecl_h8_ms()
    println("#####################################################")
    println("# clcircularplatecl_h8_export ")
    clcircularplatecl_h8_export()
    return true
end # function allrun

end # module clcircularplatecl_examples


using Distributed
@everywhere begin
	import DPFEHM
	import GaussianRandomFields
	import Optim
	import PyPlot
	import Random
	import Zygote

	function inversemodel(sigma, lambda, mu, deltah, sqrtnumobs, alpha, beta, gamma, iterations)
		#sigma is the standard deviation of logK
		#lambda is the correlation length
		#mu is the mean of logK
		#deltah is the pressure drop from the left boundary to the right boundary
		#sqrtnumobs is the square root of the number of observation wells (which are distributed on a regular sqrtnumobs-by-sqrtnumobs grid)
		#alpha is the weight associated with head measurements
		#beta is the weight associated with the logK measurements
		#gamma is the weight associated with the darcy velocity measurements
		#iterations is the maximum number of iterations that the inverse analysis will perform

		function getobsnodes(coords, obslocs)
			obsnodes = Array{Int}(undef, length(obslocs))
			for i = 1:length(obslocs)
				obsnodes[i] = findmin(map(j->sum((obslocs[i] .- coords[:, j]) .^ 2), 1:size(coords, 2)))[2]
			end
			return obsnodes
		end

		mins = [0.0, 0.0]#meters
		maxs = [100.0, 100.0]#meters
		ns = [101, 101]
		dx = (maxs[1] - mins[1]) / (ns[1] - 1)
		dy = (maxs[2] - mins[2]) / (ns[2] - 1)
		num_eigenvectors = 200
		x0 = zeros(num_eigenvectors)

		eigenvalue_error = fill(NaN, num_eigenvectors) #RMS of eigenvalues
		logK_error = fill(NaN, iterations + 1) #RMS of logK
		obj_func_values = fill(NaN, iterations + 1) #RMS of Loss

		#set up the eigenvector parameterization of the geostatistical field
		cov = GaussianRandomFields.CovarianceFunction(2, GaussianRandomFields.Matern(lambda, 1; σ=sigma))
		x_pts = range(mins[1], maxs[1]; length=ns[1])
		y_pts = range(mins[2], maxs[2]; length=ns[2])
		grf = GaussianRandomFields.GaussianRandomField(cov, GaussianRandomFields.KarhunenLoeve(num_eigenvectors), x_pts, y_pts)
		logKs = GaussianRandomFields.sample(grf)
		parameterization = copy(grf.data.eigenfunc)
		sigmas = copy(grf.data.eigenval)

		#set up the mesh  and boundary conditions
		coords, neighbors, areasoverlengths, _ = DPFEHM.regulargrid2d(mins, maxs, ns, 1.0)
		Qs = zeros(size(coords, 2))
		boundaryhead(x, y) = deltah * (x - maxs[1]) / (mins[1] - maxs[1])
		dirichletnodes = Int[]
		dirichleths = zeros(size(coords, 2))
		for i = 1:size(coords, 2)
			if coords[1, i] == mins[1] || coords[1, i] == maxs[1] || coords[2, i] == mins[2] || coords[2, i] == maxs[2]
				push!(dirichletnodes, i)
				dirichleths[i] = boundaryhead(coords[1:2, i]...)
			end
		end

		logKs2Ks_neighbors(Ks) = exp.(0.5 * (Ks[map(p->p[1], neighbors)] .+ Ks[map(p->p[2], neighbors)]))#Zygote differentiates this efficiently but the definitions above are ineffecient with Zygote
		function solveforh(logKs, dirichleths)
			@assert length(logKs) == length(Qs)
			if maximum(logKs) - minimum(logKs) > 25
				return fill(NaN, length(Qs))#this is needed to prevent the solver from blowing up if the line search takes us somewhere crazy
			else
				Ks_neighbors = logKs2Ks_neighbors(logKs)
				return DPFEHM.groundwater_steadystate(Ks_neighbors, neighbors, areasoverlengths, dirichletnodes, dirichleths, Qs)
			end
		end
		solveforheigs(x) = solveforh(x2logKs(x), dirichleths)
		x2logKs(x) = reshape(parameterization * (sigmas .* x), ns...)

		#set up the observations
		obslocs_x = range(mins[1], maxs[1]; length=sqrtnumobs + 2)[2:end - 1]
		obslocs_y = range(mins[2], maxs[2]; length=sqrtnumobs + 2)[2:end - 1]
		obslocs = collect(Iterators.product(obslocs_x, obslocs_y))[:]
		x_true = randn(num_eigenvectors)
		logKs_true = x2logKs(x_true)
		h_true = solveforheigs(x_true)
		obsnodes = getobsnodes(coords, obslocs)
		h_obser = h_true[obsnodes]
		logK_obser = logKs_true[obsnodes]
		hmat = reshape(h_true, reverse(ns)...)
		Ks = exp.(logKs)
		vy = Ks .* vcat(fill(0.0, ns[1])', hmat[3:end, :] - hmat[1:end - 2, :], fill(0.0, ns[1])') / (2 * dy)
		vx = Ks .* hcat(fill(0.0, ns[2]), hmat[:, 3:end] - hmat[:, 1:end - 2], fill(0.0, ns[2])) / (2 * dx)
		vel_obser = map(i->sqrt(vx[i] ^ 2 + vy[i] ^ 2), obsnodes)

		#set up the loss function
		function f(logKs)
			h_obssigma = 1e-3
			logK_obssigma = 1e-1
			vel_obssigma = 1e-1
			h = solveforh(logKs, dirichleths)
			h_error = sum((h[obsnodes] - h_obser) .^ 2 ./ h_obssigma ^ 2)
			K_error = sum((logKs[obsnodes] - logK_obser) .^ 2 ./ logK_obssigma ^ 2)
			Ks = exp.(logKs)
			hmat = reshape(h, reverse(ns)...)
			vy = Ks .* vcat(fill(0.0, ns[1])', hmat[3:end, :] - hmat[1:end - 2, :], fill(0.0, ns[1])') / (2 * dy)
			vx = Ks .* hcat(fill(0.0, ns[2]), hmat[:, 3:end] - hmat[:, 1:end - 2], fill(0.0, ns[2])) / (2 * dx)
			vel_pred = map(i->sqrt(vx[i] ^ 2 + vy[i] ^ 2), obsnodes)
			vel_error = sum((vel_pred - vel_obser) .^ 2 ./ vel_obssigma ^ 2)
			error_sys = alpha * h_error + beta * K_error + gamma * vel_error
			return error_sys
		end
		feigs(x) = f(x2logKs(x)) + sum(x .^ 2)

		#solve the inverse problem
		options = Optim.Options(iterations=iterations, store_trace=true, show_trace=false, extended_trace=true, x_tol=1e-6, time_limit=60 * 60 * 2)
		opt = Optim.optimize(feigs, x->Zygote.gradient(feigs, x)[1], x0, Optim.LBFGS(), options; inplace=false)
		logK_error[1:length(opt.trace)] = map(j->(sqrt.((sum((x2logKs(opt.trace[j].metadata["x"])-x2logKs(x_true)).^2))/num_eigenvectors)),1:length(opt.trace))
		obj_func_values[1:length(opt.trace)] = map(j->(sqrt.(opt.trace[j].value)), 1:length(opt.trace))
		x_est = opt.minimizer
		eigenvalue_error .= x_est - x_true
		return eigenvalue_error, logK_error, obj_func_values
	end
end#end the distributed portion


"""Experimental package that evaluates pseudoinverses such as `inv(A'A)A'b` lazily, as if `A\b`.
Note that many normal matrix operations will fail, because only a few sample methods have been
defined.
"""
module LazyPseudoinverse

using LinearAlgebra, PDMats
import Base: *, \, promote_op, inv, collect, Matrix, size # overloads
import LinearAlgebra: matprod
import PDMats.PDMat

export MatrixInnerProd, adjointmul, InvMatrixInnerProd, LazyLeftPseudoinverse

abstract type AbstractMatrixInnerProd{T} <: AbstractMatrix{T} end
abstract type AbstractInvMatrixInnerProd{T} <: AbstractMatrix{T} end

"""
`MatrixInnerProd(A::AbstractMatrix)`

Construct a matrix inner product `A'A` from `A`. This stores `A` as a field and
can act like a positive definite matrix. For example, `inv` and `\` are defined,
although `inv` constructs a lazy InvMatrixInnerProd.
Use `collect` to materialize an actual matrix. It can also be converted to
`PDMat(MatrixInnerProd(A))` if `using PDMats`.
"""
struct MatrixInnerProd{T,S<:AbstractMatrix{T}} <: AbstractMatrixInnerProd{T}
    A::S
end

Base.collect(AtA::MatrixInnerProd) = adjointmul(AtA.A',AtA.A)
Base.Matrix(AtA::MatrixInnerProd) = collect(AtA)
Base.size(AtA::MatrixInnerProd) = (size(AtA.A,2),size(AtA.A,2))
# When user shows MatrixInnerProd, print the result (A'A) as if normal, even though it's really lazy
Base.print_array(io::IO, A::MatrixInnerProd) = Base.print_array(io, adjointmul(A.A',A.A))
\(AtA::MatrixInnerProd, B::AbstractVector) = cholesky(collect(AtA)) \ B
\(AtA::MatrixInnerProd, B::AbstractMatrix) = cholesky(collect(AtA)) \ B
*(AtA::MatrixInnerProd, B::AbstractVector) = collect(AtA)*B
*(AtA::MatrixInnerProd, B::AbstractMatrix) = collect(AtA)*B
# Only a few methods defined here, so many matrix operations will fail, e.g. A'A+B

"""
`adjointmul(A', B) performs inner product of `A` and `B`, `A'B`. This is originally `*`
from `matmul.jl` (which has been type-pirated for LazyPseudoinverse.jl). To actually
complete `A'*B`, use `adjointmul(A', B)`
"""
function adjointmul(A::Adjoint,B::AbstractMatrix) # A'*B
    TS = promote_op(matprod, eltype(A), eltype(B)) # this is what * normally does in matmul.jl:151
    mul!(similar(B, TS, (size(A,1), size(B,2))), A, B)
end

"""
`A'*B` is overloaded to perform a normal matrix multiplication, except when fed `A'A` which
returns a `MatrixInnerProd` which stores `A` and acts like positive definite matrix.
Warning: This is type piracy of `Base.*`, be careful!
"""
function *(A::Adjoint,B::AbstractMatrix) # A'*B
    if A.parent === B # might be a bit looser to do ==
        MatrixInnerProd(A.parent)
    else
        adjointmul(A, B) # do a regular matmul
    end
end
PDMats.PDMat(A::MatrixInnerProd) = PDMat(adjointmul(A.A',A.A)) # sample conversion to positive definite matrix

"""
`InvMatrixInnerProd(A::AbstractMatrix)`

Construct an inverse matrix inner product `inv(A'A`) from `A`. This stores `A` as a field and
can act like an inverted positive definite matrix.
Use `collect` to materialize an actual matrix. It can also be converted to
`PDMat(MatrixInnerProd(A))` when `using PDMats`.
"""
struct InvMatrixInnerProd{T,S<:AbstractMatrix{T}} <: AbstractInvMatrixInnerProd{T}
    A::S
end
Base.collect(invAtA::InvMatrixInnerProd) = inv(cholesky(adjointmul(invAtA.A',invAtA.A)))
Base.Matrix(invAtA::InvMatrixInnerProd) = collect(invAtA)
Base.size(invAtA::InvMatrixInnerProd) = (size(invAtA.A,2),size(invAtA.A,2))
# When user shows InvMatrixInnerProd, print the result inv(A'A) as computed regularly:
Base.print_array(io::IO, A::InvMatrixInnerProd) = Base.print_array(io, inv(adjointmul(A.A',A.A)))
PDMats.PDMat(A::InvMatrixInnerProd) = PDMat(inv(adjointmul(A.A',A.A))) # Cholesky decomposition

inv(AtA::MatrixInnerProd) = InvMatrixInnerProd(AtA.A) # lazy inverse

function *(invAtA::InvMatrixInnerProd,B::Adjoint) # inv(A'A)*B'
    if invAtA.A === B.parent  # B = A, so inv(A'A)*A'
        LazyLeftPseudoinverse(invAtA.A) # return the pseudo-inverse, ready to be right-multiplied
    else
        cholesky(adjointmul(invAtA.A',invAtA.A))\B # use linear solve instead of explicit inverse
    end
end

# Ensure that stuff like inv(A'A)B, or B*inv(A'A) still works as normal, although to get
# useful functionality many other methods need to be defined
*(invAtA::InvMatrixInnerProd,B::AbstractVector) = cholesky(adjointmul(invAtA.A',invAtA.A))\B
*(invAtA::InvMatrixInnerProd,B::AbstractMatrix) = cholesky(adjointmul(invAtA.A',invAtA.A))\B
*(B::AbstractMatrix,invAtA::InvMatrixInnerProd) = B*collect(invAtA)

"""
`LazyLeftPseudoinverse(A::AbstractMatrix)`

Construct equivalent of `inv(A'A)A'` from `A`. This stores `A` as a field and
can act (and displays) like a pseudo-inverse `pinv(A)`. But when actually multiplied
against `b`, it computes `A\\b` using QR.
Use `collect` to materialize an actual matrix.
"""
struct LazyLeftPseudoinverse{T,S<:AbstractMatrix{T}} <: AbstractMatrix{T}
    A::S
end

Base.size(pending::LazyLeftPseudoinverse) = reverse(size(pending.A))
Base.print_array(io::IO, pending::LazyLeftPseudoinverse) = Base.print_array(io, pinv(pending.A))
*(pending::LazyLeftPseudoinverse, B::AbstractVector) = pending.A \ B
*(pending::LazyLeftPseudoinverse, B::AbstractMatrix) = pending.A \ B
*(B::AbstractMatrix,pending::LazyLeftPseudoinverse) = B*pinv(pending.A)
Base.collect(lazyLeftpinv::LazyLeftPseudoinverse) = pinv(lazyLeftpinv.A)
Base.Matrix(lazyLeftpinv::LazyLeftPseudoinverse) = collect(lazyLeftpinv)

end # module


using Lazy

curry(f, x) = (xs...) -> f(x, xs...)

fibs = @lazy 0:1:(fibs + tail(fibs))

under = curry(takewhile, (x -> x <= 4e6))
even = curry(filter, iseven)
solve = (sum ∘ even ∘ under)

println(solve(fibs))


struct Simulation{DT <: Number, FT <: Function}
    Δt::DT
    f::FT
    x::OffsetArray{DT,3,Array{DT,3}}

    function Simulation(f::FT, x₀::AbstractArray{DT,2}, Δt::DT, nt::Int) where {DT, FT}
        x = OffsetArray(zeros(DT, size(x₀)..., nt+1), axes(x₀)..., 0:nt)
        x[:,:,0] .= x₀
        new{DT,FT}(Δt, f, x)
    end
end

ntimesteps(sim::Simulation) = lastindex(sim.x,3)
nsamples(sim::Simulation) = length(axes(sim.x,2))
ndims(sim::Simulation) = length(axes(sim.x,1))

eachtimestep(sim::Simulation) = axes(sim.x,3)[1:end]
eachsample(sim::Simulation) = axes(sim.x,2)

function run!(sim::Simulation{DT}) where {DT}
    ẋ = zeros(DT, ndims(sim))
    for n in eachtimestep(sim)
        for i in eachsample(sim)
            sim.f(ẋ, sim.x[:,i,n-1])
            sim.x[:,i,n] .= sim.x[:,i,n-1] .+ sim.Δt .* ẋ
        end
    end
    return sim.x
end