text:
module GenericSchur
using LinearAlgebra

# This is the main public interface of the package.
# We extend these functions without further note.
import LinearAlgebra: schur!, eigvals!, eigvecs, eigen!

# Wrappers like `schur` and `eigvals` should just work.

# Some other stdlib functions (e.g., lmul!) are extended, but either non-piratically
# or with explicit qualification.

# These are introduced here.
export triangularize, eigvalscond, subspacesep, balance!

const STypes = Union{AbstractFloat, Complex{T} where T<:AbstractFloat}

if VERSION < v"1.2-"
    eigsortby = nothing
else
    using LinearAlgebra: eigsortby, sorteig!
end
using LinearAlgebra: checksquare

schur!(A::StridedMatrix{T}; kwargs...) where {T <: STypes} = gschur!(A; kwargs...)

function eigvals!(A::StridedMatrix{T};
                  sortby::Union{Function,Nothing}=eigsortby, kwargs...
                  ) where {T <: STypes}
    S = gschur!(A; wantZ=false, kwargs...)
    if sortby !== nothing
        return sorteig!(S.values, sortby)
    else
        return S.values
    end
end

# This is probably the best we can do unless LinearAlgebra coöperates
"""
    eigvecs(S::Schur{Complex{<:AbstractFloat}}; left=false) -> Matrix

Compute right or left eigenvectors from a Schur decomposition.
Eigenvectors are returned as columns of a matrix, ordered to match `S.values`.
The returned eigenvectors have unit Euclidean norm, and the largest
elements are real.
"""
function eigvecs(S::Schur{Complex{T}}; left::Bool=false) where {T <: AbstractFloat}
    if left
        v = _gleigvecs!(S.T,S.Z)
    else
        v = _geigvecs!(S.T,S.Z)
    end
    _enormalize!(v)
    v
end

function eigen!(A::StridedMatrix{T}; permute::Bool=true, scale::Bool=true,
                sortby::Union{Function,Nothing}=eigsortby, kwargs...
                ) where {T <: STypes}
    n = checksquare(A)
    if permute || scale
        A, B = balance!(A, scale=scale, permute=permute)
    end
    if T <: Real
        S = triangularize(schur(A))
        VT = Complex{T}
    else
        S = schur(A)
        VT = T
    end
    if isempty(kwargs)
        v = _geigvecs!(S.T,S.Z)
        if permute || scale
            lmul!(B, v)
        end
        _enormalize!(v)
        if sortby !== nothing
            return LinearAlgebra.Eigen(sorteig!(S.values, v, sortby)...)
        else
            return LinearAlgebra.Eigen(S.values,v)
        end
    end

    # recip. cond. nr. needs both sets of eigvecs
    # but compute it for BALANCED forms for consistency
    do_vr = get(kwargs, :jvr, true) || get(kwargs, :jce, false)
    if do_vr
        v = _geigvecs!(S.T,S.Z)
    else
        v = zeros(VT,0,0)
    end
    do_vl = get(kwargs, :jvl, false) || get(kwargs, :jce, false)
    if do_vl
        vl = _gleigvecs!(S.T,S.Z)
    else
        vl = zeros(VT,0,0)
    end

    do_econd = get(kwargs, :jce, false)
    rconde = zeros(real(T), do_econd ? n : 0)
    if do_econd
        for j in 1:n
            vlj = view(vl,:,j)
            vrj = view(v,:,j)
            rconde[j] = abs(dot(vlj,vrj)) / norm(vlj) / norm(vrj)
        end
        if !get(kwargs, :jvr, true)
            v = zeros(VT,0,0)
        end
        if !get(kwargs, :jvl, true)
            vl = zeros(VT,0,0)
        end
    end

    if get(kwargs, :jvr, true)
        if permute || scale
            lmul!(B, v)
        end
        _enormalize!(v)
    end
    if get(kwargs, :jvl, true)
        if permute || scale
            ldiv!(B, vl)
        end
        _enormalize!(vl)
    end

    do_vcond = get(kwargs, :jcv, false)
    rcondv = zeros(real(T), do_vcond ? n : 0)
    if do_vcond
        # warning: this is very expensive
        Ttmp = similar(S.T)
        sel = falses(n)
        for j in 1:n
            copyto!(Ttmp, S.T)
            Stmp = Schur(Ttmp,similar(S.Z,0,0),copy(S.values))
            sel[j] = true
            ordschur!(Stmp,sel)
            sel[j] = false
            if do_econd
                rconde[j] = eigvalscond(Stmp,1)
            end
            if do_vcond
                rcondv[j] = subspacesep(Stmp,1)
            end
        end
    end
    if sortby !== nothing
        return LinearAlgebra.Eigen(sorteig!(S.values, v, sortby, vl, rconde, rcondv)...)
    else
        return LinearAlgebra.Eigen(S.values,v,vl,false,rconde,rcondv)
    end
end

using LinearAlgebra: RealHermSymComplexHerm, Algorithm, QRIteration

function eigen!(A::RealHermSymComplexHerm{<:STypes, <:StridedMatrix},
                              alg::Algorithm = QRIteration();
                              kwargs...)
    geigen!(A, alg; kwargs...)
end

function eigvals!(A::RealHermSymComplexHerm{<:STypes, <:StridedMatrix},
                                alg::Algorithm = QRIteration();
                                kwargs...)
    geigvals!(A, alg; kwargs...)
end

function eigen!(A::SymTridiagonal{<:AbstractFloat},
                              alg::Algorithm = QRIteration();
                              kwargs...)
    geigen!(A, alg; kwargs...)
end

function eigvals!(A::SymTridiagonal{<:AbstractFloat},
                              alg::Algorithm = QRIteration();
                              kwargs...)
    geigvals!(A, alg; kwargs...)
end

# The variants should be defined here for logical coherence, but not until
# we've registered nontrivial implementations.
############################################################################
# Internal implementations follow

using LinearAlgebra: Givens, Rotation
using Printf

if VERSION >= v"1.2"
    using Base: require_one_based_indexing
else
    function require_one_based_indexing(A::AbstractArray)
        if Base.has_offset_axes(A)
            throw(ArgumentError("offset axes are not supported"))
        end
    end
end

function geigen!(A::RealHermSymComplexHerm{<:STypes, <:StridedMatrix},
                 alg::Algorithm = QRIteration();
                 sortby::Union{Function,Nothing}=eigsortby
                 )
    H = hessenberg!(A)
    V = _materializeQ(H)
    S = gschur!(H, V, alg)
    λ = S.values
    LinearAlgebra.Eigen(sorteig!(λ, V, eigsortby)...)
end

function geigvals!(A::RealHermSymComplexHerm{<:STypes, <:StridedMatrix},
                 alg::Algorithm = QRIteration();
                 sortby::Union{Function,Nothing}=eigsortby
                 )
    H = hessenberg!(A)
    S = gschur!(H, nothing, alg)
    λ = S.values
    sorteig!(λ, sortby)
end

# fallback methods; only some algorithms have selection capability

function geigen!(A::RealHermSymComplexHerm{<:STypes, <:StridedMatrix},
                 irange::UnitRange,
                 alg::Algorithm = QRIteration; kwargs...
                 )
    throw(ArgumentError("eigenvalue selection is not implemented for $alg"))
end
function geigen!(A::RealHermSymComplexHerm{<:STypes, <:StridedMatrix},
                 vl::Real, vu::Real,
                 alg::Algorithm = QRIteration; kwargs...
                 )
    throw(ArgumentError("eigenvalue selection is not implemented for $alg"))
end

# debugging stuff
macro mydebug(expr); nothing; end

function _fmt_nr(z::Complex)
    s = @sprintf("%10.3e%+10.3eim", reim(z)...)
    return s
end
function _fmt_nr(x::Real)
    s = @sprintf("%10.3e", x)
    return s
end

include("util.jl")
include("hessenberg.jl")

if VERSION < v"1.3"
    _getdata(H::HessenbergFactorization) = H.data
    const HessenbergArg = HessenbergFactorization
else
    _getdata(H::Hessenberg) = H.H.data
    const HessenbergArg = Hessenberg
end

include("householder.jl")
include("balance.jl")

#
# portions translated from LAPACK::zlahqr
# LAPACK Copyright:
# Univ. of Tennessee
# Univ. of California Berkeley
# Univ. of Colorado Denver
# NAG Ltd.
function gschur!(H::HessenbergArg{Complex{RT}}, Z=nothing;
                  maxiter = 100*size(H, 1), maxinner = 30*size(H, 1),
                  checksd=true, kwargs...
                 ) where {RT <: AbstractFloat}
    T = Complex{RT}
    n = size(H, 1)
    istart = 1
    iend = n
    w = Vector{T}(undef, n)
    HH = _getdata(H)

    if checksd
        for j in 1:n-1
            isreal(HH[j+1,j]) || throw(ArgumentError("algorithm assumes real subdiagonal"))
        end
    end

    ulp = eps(RT)
    smallnum = safemin(RT) * (n / ulp)
    rzero = zero(RT)
    half = 1 / RT(2)
    threeq = 3 / RT(4)

    for j=1:n-1
        HH[j+2:n,j] .= zero(T)
    end

    # iteration count
    it = 0

    @inbounds while iend >= 1
        # key for comparison to zlahqr:
        # I => iend
        # L => istart
        # K => _istart + 1
        istart = 1
        for its=0:maxinner
            it += 1
            if it > maxiter
                throw(ArgumentError("iteration limit $maxiter reached"))
            end

            # Determine if the matrix splits.
            # Find lowest positioned subdiagonal "zero" if any; reset istart if found.
            for _istart in iend-1:-1:istart
                if abs1(HH[_istart+1, _istart]) <= smallnum
                    istart = _istart + 1
                    @mydebug println("Split1! at istart = $istart, subdiag is ",
                                     _fmt_nr(HH[istart, istart-1]))
                    break
                end
                # deflation criterion from Ahues & Tisseur (LAWN 122, 1997)
                tst = abs1(HH[_istart,_istart]) + abs1(HH[_istart+1,_istart+1])
                if tst == 0
                    if (_istart-1 >= 1) tst += abs(real(HH[_istart,_istart-1])) end
                    if (_istart+2 <= n) tst += abs(real(HH[_istart+2,_istart+1])) end
                end
                if abs(real(HH[_istart+1,_istart])) <= ulp*tst
                    ab = max(abs1(HH[_istart+1,_istart]),
                             abs1(HH[_istart,_istart+1]))
                    ba = min(abs1(HH[_istart+1,_istart]),
                             abs1(HH[_istart,_istart+1]))
                    aa = max(abs1(HH[_istart+1,_istart+1]),
                             abs1(HH[_istart,_istart]-HH[_istart+1,_istart+1]))
                    bb = min(abs1(HH[_istart+1,_istart+1]),
                             abs1(HH[_istart,_istart]-HH[_istart+1,_istart+1]))
                    s = aa + ab
                    if ba * (ab / s) <= max(smallnum, ulp * (bb * (aa / s)))
                        istart = _istart + 1
                        @mydebug println("Split2! at istart = $istart, subdiag is ",
                                         _fmt_nr(HH[istart, istart-1]))
                        break
                    end
                end
            end # check for split

            if istart > 1
                # clean up
                HH[istart, istart-1] = zero(T)
            end

            # if block size is one we deflate
            if istart >= iend
                iend -= 1
                @mydebug println("Bottom deflation! Block size is one. New iend is $iend")
                break # from "inner" loop
            end

            # select shift
            # exceptional shift logic adapted from LAPACK zlahqr
            if its % 30 == 10
                s = threeq * abs(real(HH[istart+1,istart]))
                t = s + HH[istart,istart]
            elseif its % 30 == 20
                s = threeq * abs(real(HH[iend,iend-1]))
                t = s + HH[iend,iend]
            else
                # Wilkinson's shift
                t = HH[iend,iend]
                u = sqrt(HH[iend-1,iend]) * sqrt(HH[iend,iend-1])
                s = abs1(u)
                if s ≠ rzero
                    x = half * (HH[iend-1,iend-1] - t)
                    sx = abs1(x)
                    s = max(s, abs1(x))
                    y = s * sqrt( (x/s)^2 + (u/s)^2)
                    if sx > rzero
                        if real(x / sx) * real(y) + imag(x / sx) * imag(y) < rzero
                            y = -y
                        end
                    end
                    t -= u * (u / (x+y))
                end
            end # shift selection

            # run a QR iteration
            # following zlahqr, only use single-shift
            singleShiftQR!(HH, Z, t, istart, iend)

        end # inner loop
    end # outer loop
    w = diag(HH)
    return Schur(triu(HH), Z === nothing ? Matrix{T}(undef,0,0) : Z, w)
end

"""
gschur(A::StridedMatrix) -> F::Schur

Computes the Schur factorization of matrix `A` using a generic implementation.
See `LinearAlgebra.schur` for usage.
"""
gschur(A::StridedMatrix{T}; kwargs...) where {T} = gschur!(Matrix(A); kwargs...)

"""
gschur!(A::StridedMatrix) -> F::Schur

Destructive version of `gschur` (q.v.).
"""
function gschur!(A::StridedMatrix{Complex{T}}; wantZ::Bool=true, scale::Bool=true,
                 kwargs...) where T <: AbstractFloat
    n = checksquare(A)
    if scale
        scaleA, cscale, anrm = _scale!(A)
    else
        scaleA = false
    end
    H = _hessenberg!(A)
    if wantZ
        Z = _materializeQ(H)
        S = gschur!(H, Z; checksd=false, kwargs...)
    else
        S = gschur!(H; checksd=false, kwargs...)
    end
    if scaleA
        safescale!(S.T, cscale, anrm)
        S.values .= diag(S.T, 0)
    end
    S
end

function singleShiftQR!(HH::StridedMatrix{T}, Z, shift::Number, istart::Integer, iend::Integer) where {T <: Complex}
    n = size(HH, 1)
    ulp = eps(real(eltype(HH)))

    @mydebug Hsave = Z*HH*Z'
    @mydebug function dcheck(str)
        de = norm(Hsave - Z*HH*Z')
        if de > 1e-6
            println()
            @warn "$str decomp err $de"
        else
            print(str," decomp err ", de)
        end
    end

    # key:
    # istart => L
    # iend => I
    # istart1 => M

    # look for two consecutive small subdiagonals
    istart1 = -1
    h11s = zero(eltype(HH))
    h21 = zero(real(eltype(HH)))
    flag = false
    v = zeros(T,2)
    @inbounds for mm = iend-1:-1:istart+1
        # determine the effect of starting the single-shift Francis
        # iteration at row mm: see if this would make HH[mm,mm-1] tiny.
        h11 = HH[mm,mm]
        h22 = HH[mm+1,mm+1]
        h11s = h11 - shift
        h21 = real(HH[mm+1,mm]) # for reflector
        s = abs1(h11s) + abs(h21)
        h11s /= s
        h21 /= s
        v .= (h11s, T(h21))
        h10 = real(HH[mm,mm-1])
        if abs(h10) * abs(h21) <= ulp * (abs1(h11s) * (abs1(h11) + abs1(h22)))
            istart1 = mm
            flag = true
            break
        end
    end
    @inbounds begin
      if !flag
        istart1 = istart
        h11 = HH[istart,istart]
        h22 = HH[istart+1,istart+1]
        h11s = h11 - shift
        h21 = real(HH[istart+1,istart])
        s = abs1(h11s) + abs(h21)
        h11s /= s
        h21 /= s
        v .= (h11s,T(h21))
      end
    end
    @mydebug println("QR sweep $(istart1):$(iend) shift ", _fmt_nr(shift),
                     " subdiag ", _fmt_nr(HH[iend,iend-1]))

    @inbounds for k = istart1:iend-1
        if k > istart1
            # prepare to chase bulge
            v .= (HH[k,k-1], HH[k+1,k-1])
        # else
        #    use v from above to create a bulge
        end
        τ1 = _reflector!(v)
        # τ1 = LinearAlgebra.reflector!(v)
        if k > istart1
            HH[k,k-1] = v[1]
            HH[k+1,k-1] = zero(T)
        end
        v2 = v[2]
        τ2 = real(τ1*v2)
        # lmul!(R,view(HH,:,k:n))
        for j = k:n
            ss = τ1' * HH[k,j] + τ2 * HH[k+1,j]
            HH[k,j] -= ss
            HH[k+1,j] -= ss * v2
        end
        # rmul!(view(HH,1:min(iend,k+2),:),R)
        for j = 1:min(k+2,iend)
            ss = τ1 * HH[j,k] + τ2 * HH[j,k+1]
            HH[j,k] -= ss
            HH[j,k+1] -= ss * v2'
        end
        if !(Z === nothing)
            # rmul!(Z,R)
            for j = 1:n
                ss = τ1 * Z[j,k] + τ2 * Z[j,k+1]
                Z[j,k] -= ss
                Z[j,k+1] -= ss * v2'
            end
        end
        if (k == istart1) && (istart1 > istart)
            # If the QR step started at a row below istart because two consecutive
            # small subdiagonals were found, extra scaling must be performed
            # to ensure reality of HH[istart1,istart1-1]
            t = one(T) - τ1
            t /= abs(t)
            HH[istart1+1,istart1] *= t'
            if istart1 + 2 <= iend
                HH[istart1+2, istart1+1] *= t
            end
            for j = istart1:iend
                if j != istart1 + 1
                    if n > j
                        HH[j,j+1:n] *= t
                    end
                    HH[1:j-1,j] *= t'
                    if !(Z === nothing)
                        Z[1:n, j] *= t'
                    end
                end
            end
        end
        @mydebug dcheck(" QR k=$k")
    end
    # ensure reality of tail
    @inbounds begin
      t = HH[iend,iend-1]
      if imag(t) != 0
        rt = abs(t)
        HH[iend, iend-1] = rt
        t /= rt
        if n > iend
            HH[iend, iend+1:n] *= t'
        end
        HH[1:iend-1, iend] *= t
        if !(Z === nothing)
            Z[1:n, iend] *= t
        end
      end
    end
#    @mydebug dcheck(" QR post")
    @mydebug println()
    return HH
end

# based on LAPACK dlahqr
"""
gschur!(H::Hessenberg, Z) -> F::Schur

Compute the Schur decomposition of a Hessenberg matrix.  Subdiagonals of `H` must be real.
If `Z` is provided, it is updated with the unitary transformations of the decomposition.
"""
function gschur!(H::HessenbergArg{T}, Z::Union{Nothing, AbstractMatrix}=nothing;
                  tol = eps(real(T)),
                  maxiter = 100*size(H, 1), standardize = nothing,
                  kwargs...) where {T <: AbstractFloat}
    if !(standardize === nothing)
        @warn "obsolete keyword `standardize` in gschur!" maxlog=1
    end
    n = size(H, 1)
    if !(Z === nothing)
        size(Z, 2) == n || throw(DimensionMismatch("second dimension of Z must match H"))
    end
    istart = 1
    iend = n
    HH = _getdata(H)
    triu!(HH,-1)
    τ = Rotation(Givens{T}[])
    w = Vector{Complex{T}}(undef, n)
    iwcur = n
    function putw!(w,z)
        w[iwcur] = z
        iwcur -= 1
    end

    smallnum = floatmin(T) * (n / eps(T))
    threeq = 3 / T(4)
    m7_16 = -7 / T(16)

    # iteration count
    iter = 0

    # Main loop.
    @inbounds while iend >= 1
        # eigenvalues in diagonal entries iend+1:n have converged.
        istart = 1
        iterqr = 0
        deflate = false
        while true
            iter += 1
            if iter > maxiter
                throw(ArgumentError("iteration limit $maxiter reached"))
            end

            # istart => L
            # iend => I

            # Determine if the matrix splits.
            split = false
            for k in iend:-1:istart+1
                if abs(HH[k, k-1]) < smallnum
                    split = true
                else
                    Hkk = HH[k, k]
                    Hkm1km1 = HH[k-1, k-1]
                    t = abs(Hkm1km1) + abs(Hkk)
                    if t == 0
                        if k > 2
                            t += abs(HH[k-1, k-2])
                        end
                        if k+1 <= n
                            t += abs(HH[k+1, k])
                        end
                    end
                    # conservative deflation criterion from Ahues & Tisseur
                    aHkkm1 = abs(HH[k, k-1])
                    if aHkkm1 <= t * eps(T)
                        aHkm1k = abs(HH[k-1, k])
                        ab = max(aHkkm1, aHkm1k)
                        ba = min(aHkkm1, aHkm1k)
                        aa = max(abs(Hkk), abs(Hkm1km1 - Hkk))
                        bb = min(abs(Hkk), abs(Hkm1km1 - Hkk))
                        s = aa + bb
                        if ba * (ab/s) <= max(smallnum, eps(T) * (bb * (aa / s)))
                            split = true
                        end
                    end
                end
                if split
                    istart = k
                    @mydebug println("Split at $istart, subdiagonal element is ",
                                     _fmt_nr(HH[istart, istart-1]))
                    break
                end
            end
            if !split
                istart = 1
            end
            if istart > 1
                HH[istart, istart-1] = zero(T)
            end
            if istart >= iend-1
                deflate = true
                break
            end

            # run a QR iteration
            iterqr += 1
            if iterqr == 10
                # exceptional shift from top
                s = abs(HH[istart+1,istart]) + abs(HH[istart+2,istart+1])
                H11 = threeq * s + HH[istart,istart]
                H12 = m7_16 * s
                H21 = s
                H22 = H11
            elseif iterqr == 20
                # exceptional shift from bottom
                s = abs(HH[iend,iend-1]) + abs(HH[iend-1,iend-2])
                H11 = threeq * s + HH[iend,iend]
                H12 = m7_16 * s
                H21 = s
                H22 = H11
            else
                # Francis' double shift (i.e. 2nd degree Rayleigh quotient)
                H11 = HH[iend-1,iend-1]
                H21 = HH[iend,iend-1]
                H12 = HH[iend-1,iend]
                H22 = HH[iend,iend]
            end
            s = abs(H11) + abs(H12) + abs(H21) + abs(H22)
            if s == 0
                r1r,r2r,r1i,r2i = zero(T),zero(T),zero(T),zero(T)
            else
                H11 /= s
                H12 /= s
                H21 /= s
                H22 /= s
                tr = (H11+H22) / T(2)
                d = (H11-tr) * (H22-tr) - H12 * H21
                rtd = sqrt(abs(d))
                if d >= 0
                    # conjugate pair
                    r1r = tr*s
                    r2r = r1r
                    r1i = rtd * s
                    r2i = -r1i
                else
                    # duplicate real shifts
                    r1r = tr + rtd
                    r2r = tr - rtd
                    if abs(r1r - H22) <= abs(r2r - H22)
                        r1r *= s
                        r2r = r1r
                    else
                        r2r *= s
                        r1r = r2r
                    end
                    r1i, r2i = zero(T), zero(T)
                end
            end
            doubleShiftQR!(HH, Z, complex(r1r, r1i), complex(r2r, r2i), istart, iend)
        end
        if deflate && (istart >= iend)
            # if block size is one we deflate
            thisw = HH[iend,iend]
            putw!(w, thisw)
            iend = istart-1
            @mydebug println("Deflate one. New iend is $iend. w=",_fmt_nr(thisw))
        elseif deflate && (istart + 1 == iend)
            # and the same for a 2x2 block
            H2 = HH[iend-1:iend,iend-1:iend]
            G2,w1,w2 = _gs2x2!(H2,iend)
            putw!(w,w2)
            putw!(w,w1)
            lmul!(G2,view(HH,:,istart:n))
            rmul!(view(HH,1:iend,:),G2')
            HH[iend-1:iend,iend-1:iend] .= H2 # clean
            if iend > 2
                HH[iend-1,iend-2] = 0
            end
            Z === nothing || rmul!(Z,G2')
        end
        iend = istart-1
        @mydebug begin
            println("Deflate two. New iend is $iend w1=",_fmt_nr(w1),
                    " w2=",_fmt_nr(w2))
        end
    end

    TT = triu(HH,-1)
    return Schur{T,typeof(TT)}(TT, Z === nothing ? similar(TT,0,0) : Z, w)
end

# compute Schur decomposition of real 2x2 in standard form
# return corresponding Givens and eigenvalues
# Translated from LAPACK::dlanv2
# Copyright:
# Univ. of Tennessee
# Univ. of California Berkeley
# Univ. of Colorado Denver
# NAG Ltd.
function _gs2x2!(H2::StridedMatrix{T},jj) where {T <: Real}
    a,b,c,d = H2[1,1], H2[1,2], H2[2,1], H2[2,2]
    sgn(x) = (x < 0) ? -one(T) : one(T) # fortran sign differs from Julia
    half = one(T) / 2
    small = 4eps(T) # how big discriminant must be for easy reality check
    if c==0
        cs = one(T)
        sn = zero(T)
    elseif b==0
        # swap rows/cols
        cs = zero(T)
        sn = one(T)
        a,b,c,d = d,-c,zero(T),a
    elseif ((a-d) == 0) && (b*c < 0)
        # nothing to do
        cs = one(T)
        sn = zero(T)
    else
        asubd = a-d
        p = half*asubd
        bcmax = max(abs(b),abs(c))
        bcmis = min(abs(b),abs(c)) * sgn(b) * sgn(c)
        scale = max(abs(p), bcmax)
        z = (p / scale) * p + (bcmax / scale) * bcmis
        # if z is of order machine accuracy: postpone decision
        if z >= small
            # real eigenvalues
            z = p + sqrt(scale) * sqrt(z) * sgn(p)
            a = d + z
            d -= (bcmax / z) * bcmis
            τ = hypot(c,z)
            cs = z / τ
            sn = c / τ
            b -= c
            c = zero(T)
        else
            # complex or almost equal real eigenvalues
            σ = b + c
            τ = hypot(σ, asubd)
            cs = sqrt(half * (one(T) + abs(σ) / τ))
            sn = -(p / (τ * cs)) * sgn(σ)
            # apply rotations
            aa = a*cs + b*sn
            bb = -a*sn + b*cs
            cc = c*cs + d*sn
            dd = -c*sn + d*cs
            a = aa*cs + cc*sn
            b = bb*cs + dd*sn
            c = -aa*sn + cc*cs
            d = -bb*sn + dd*cs
            midad = half * (a+d)
            a = midad
            d = a
            if (c != 0)
                if (b != 0)
                    if b*c >= 0
                        # real eigenvalues
                        sab = sqrt(abs(b))
                        sac = sqrt(abs(c))
                        p = sab*sac*sgn(c)
                        τ = one(T) / sqrt(abs(b+c))
                        a = midad + p
                        d = midad - p
                        b -= c
                        c = 0
                        cs1 = sab*τ
                        sn1 = sac*τ
                        cs, sn = cs*cs1 - sn*sn1, cs*sn1 + sn*cs1
                    end
                else
                    b,c = -c,zero(T)
                    cs,sn = -sn,cs
                end
            end
        end
    end

    if c==0
        w1,w2 = a,d
    else
        rti = sqrt(abs(b))*sqrt(abs(c))
        w1 = a + rti*im
        w2 = d - rti*im
    end
    H2[1,1], H2[1,2], H2[2,1], H2[2,2] = a,b,c,d
    G = Givens(jj-1,jj,cs,sn)
    return G,w1,w2
end

function gschur!(A::StridedMatrix{T}; wantZ::Bool=true, scale::Bool=true,
                 kwargs...) where {T <: AbstractFloat}
    n = checksquare(A)
    if scale
        scaleA, cscale, anrm = _scale!(A)
    else
        scaleA = false
    end
    H = _hessenberg!(A)
    if wantZ
        Z = _materializeQ(H)
        S = gschur!(H, Z; kwargs...)
    else
        S = gschur!(H; kwargs...)
    end
    if scaleA
        safescale!(S.T, cscale, anrm)
        safescale!(S.values, cscale, anrm)
    end
    S
end

function doubleShiftQR!(H::StridedMatrix{T}, Z, shift1, shift2,
                        istart::Integer, iend::Integer) where {T <: Real}
    n = size(H,1)
    @mydebug Hsave = Z*H*Z'
    @mydebug function dcheck(str)
        print(str," decomp err ", norm(Hsave - Z*H*Z'))
    end
    i1 = 1
    i2 = n
    r1r, r1i = reim(shift1)
    r2r, r2i = reim(shift2)
    v = zeros(T,3)
    # look for two consecutive small subdiagonals
    mx = istart
    for m=iend-2:-1:istart
        # Determine whether starting double-shift QR iteration
        # at row m would make H[m,m-1] negligible.
        # Use scaling to avoid over/underflow.
        H21s = H[m+1,m]
        s = abs(H[m,m] - r2r) + abs(r2i) + abs(H21s)
        H21s /= s
        v[1] = H21s * H[m,m+1] +
            (H[m,m] - r1r) * ((H[m,m]-r2r) / s) - r1i*(r2i / s)
        v[2] = H21s * (H[m,m] + H[m+1,m+1] - r1r - r2r)
        v[3] = H21s * H[m+2,m+1]
        s = sum(abs.(v))
        v ./= s
        if (m > istart) && (abs(H[m,m-1]) * (abs(v[2]) + abs(v[3])) <=
                            eps(T) * abs(v[1]) *
                            (abs(H[m-1,m-1]) + abs(H[m,m]) + abs(H[m+1,m+1])))
            mx = m
            break
        end
    end
    @mydebug println("QR sweep $mx:$iend, shifts ",_fmt_nr(shift1)," ",_fmt_nr(shift2),
                     " subdiags ", _fmt_nr(H[iend-1,iend-2])," ",_fmt_nr(H[iend,iend-1]))
    for k=mx:iend-1
        # first iteration creates a bulge using reflection based on v
        # subsequent iterations use reflections to restore Hessenberg form
        # in column k-1
        nr =  min(3, iend-k+1) # order of G
        if k > mx
            v[1:nr] .= H[k:k+nr-1,k-1]
        end
        @mydebug println(" nr=$nr v=$v ")
        τ1 = _reflector!(view(v,1:nr))
        if k > mx
            H[k,k-1] = v[1]
            H[k+1,k-1] = zero(T)
            if k < iend-1
                H[k+2,k-1] = zero(T)
            end
        elseif mx > istart
            # avoid problems when v[2] and v[3] underflow
            H[k,k-1] *= (one(T) - τ1)
        end
        v2 = v[2]
        τ2 = τ1 * v2
        if nr == 3
            v3 = v[3]
            τ3 = τ1 * v3
            # apply G from left; transform H[:,k:i2]
            for j=k:i2
                ss = H[k,j] + v2 * H[k+1,j] + v3 * H[k+2,j]
                H[k,j] -= ss * τ1
                H[k+1,j] -= ss * τ2
                H[k+2,j] -= ss * τ3
            end
            # apply G from right; transform H[i1:min(k+3,iend),:]
            for j=i1:min(k+3,iend)
                ss = H[j,k] + v2 * H[j,k+1] + v3 * H[j,k+2]
                H[j,k] -= ss * τ1
                H[j,k+1] -= ss * τ2
                H[j,k+2] -= ss * τ3
            end
            if !(Z === nothing)
                for j in 1:size(Z,1)
                    ss = Z[j,k] + v2 * Z[j,k+1] + v3 * Z[j,k+2]
                    Z[j,k] -= ss * τ1
                    Z[j,k+1] -= ss * τ2
                    Z[j,k+2] -= ss * τ3
                end
            end
        elseif nr == 2
            # apply G from left; transform H[:,k:i2]
            for j=k:i2
                ss = H[k,j] + v2 * H[k+1,j]
                H[k,j] -= ss * τ1
                H[k+1,j] -= ss * τ2
            end
            # apply G from right; transform H[i1:iend,:]
            for j=i1:iend
                ss = H[j,k] + v2 * H[j,k+1]
                H[j,k] -= ss * τ1
                H[j,k+1] -= ss * τ2
            end
            if !(Z === nothing)
                for j in 1:size(Z,1)
                    ss = Z[j,k] + v2 * Z[j,k+1]
                    Z[j,k] -= ss * τ1
                    Z[j,k+1] -= ss * τ2
                end
            end
        end
        @mydebug dcheck(" QR k=$k ")
    end # k loop
    @mydebug println()
end

include("vectors.jl")
include("triang.jl")

include("symtridiag.jl")

include("norm1est.jl")
include("sylvester.jl")

include("generalized.jl")
include("gsylvester.jl")

include("rordschur.jl")
include("ordschur.jl")
include("gcondition.jl")

end # module


text:
include("rates.jl")

"""
    zero_radiation_populations(line::HydrogenicLine,
                               atmos::Atmosphere)

For a given atom density, calculate the populations according to zero-radiation,
on a regular grid.
"""
function zero_radiation_populations(line::HydrogenicLine,
                                    atmos::Atmosphere)

    nz, nx, ny = size(atmos.temperature)
    nλ = length(line.λ)

    J_zero = zeros(Float64, (nλ, nz, nx, ny))u"kW*m^-2*nm^-1"

    LTE_pops = LTE_populations(line, atmos)

    γ = γ_constant(line,
                   atmos.temperature,
                   (LTE_pops[:, :, :, 1] .+ LTE_pops[:, :, :, 2]),
                   atmos.electron_density)

    damping_λ = Array{Float64, 4}(undef, size(J_zero))
    for l in eachindex(line.λ)
       damping_λ[l, :, :, :] = damping.(γ, line.λ[l], line.ΔD)
    end

    R_zero = calculate_R(atmos, line, J_zero, damping_λ, LTE_pops)
    C_zero = calculate_C(atmos, LTE_pops)

    populations = get_revised_populations(R_zero, C_zero, atmos.hydrogen_populations)

    return populations
end

"""
    zero_radiation_populations(line::HydrogenicLine,
                               sites::VoronoiSites)

For a given atom density, calculate the populations according to zero-radiation,
on an irregular grid.
"""
function zero_radiation_populations(line::HydrogenicLine,
                                    sites::VoronoiSites)

    nλ = length(line.λ)

    J_zero = zeros(Float64, (nλ, sites.n))u"kW*m^-2*nm^-1"

    LTE_pops = LTE_populations(line, sites)

    γ = γ_constant(line,
                   sites.temperature,
                   (LTE_pops[:, 1] .+ LTE_pops[:, 2]),
                   sites.electron_density)

    damping_λ = Matrix{Float64}(undef, size(J_zero))
    for l in eachindex(line.λ)
       damping_λ[l, :] = damping.(γ, line.λ[l], line.ΔD)
    end

    R_zero = calculate_R(sites, line, J_zero, damping_λ, LTE_pops)
    C_zero = calculate_C(sites, LTE_pops)

    populations = get_revised_populations(R_zero, C_zero, sites.hydrogen_populations)

    return populations
end

"""
    LTE_populations(line::HydrogenicLine,
                    atmos::Atmosphere)

Given the atom density, calculate the atom populations according to LTE, on a
regular grid.
"""
function LTE_populations(line::HydrogenicLine,
                         atmos::Atmosphere)

    χ = [line.χi, line.χj, line.χ∞]
    # Ionised hydrogen -> g = 1
    g = [line.gi, line.gj, 1]
    atom_density = atmos.hydrogen_populations*1.0
    nz, nx, ny = size(atom_density)

    n_levels = 3
    n_relative = ones(Float64, nz, nx, ny, n_levels)

    saha_const = (k_B / h) * (2π * m_e) / h
    saha_factor = 2 * ((saha_const * atmos.temperature*1.0).^(3/2) ./ atmos.electron_density*1.0) .|> u"m/m"

    for i=2:n_levels
        ΔE = χ[i] - χ[1]
        n_relative[:,:,:,i] = g[i] / g[1] * exp.(-ΔE ./ (k_B * atmos.temperature))
    end

    # Last level is ionised stage (H II)
    n_relative[:,:,:,n_levels] .*= saha_factor
    n_relative[:,:,:,1] = 1 ./ sum(n_relative, dims=4)[:,:,:,1]
    n_relative[:,:,:,2:end] .*= n_relative[:,:,:,1]

    return n_relative .* atom_density
end

"""
    LTE_populations(line::HydrogenicLine,
                    sites::VoronoiSites)

Given the atom density, calculate the atom populations according to LTE, on an
irregular grid.
"""
function LTE_populations(line::HydrogenicLine,
                         sites::VoronoiSites)

    χ = [line.χi, line.χj, line.χ∞]
    # Ionised hydrogen -> g = 1
    g = [line.gi, line.gj, 1]
    atom_density = sites.hydrogen_populations
    n = length(atom_density)

    n_levels = 3
    n_relative = ones(Float64, n, n_levels)

    saha_const = (k_B / h) * (2π * m_e) / h
    saha_factor = 2 * ((saha_const * sites.temperature).^(3/2) ./ sites.electron_density) .|> u"m/m"

    for i=2:n_levels
        ΔE = χ[i] - χ[1]
        n_relative[:, i] = g[i] / g[1] * exp.(-ΔE ./ (k_B * sites.temperature))
    end

    # Last level is ionised stage (H II)
    n_relative[:,n_levels] .*= saha_factor
    n_relative[:,1] = 1 ./ sum(n_relative, dims=2)[:,1]
    n_relative[:,2:end] .*= n_relative[:,1]

    return n_relative .* atom_density
end

"""
    get_revised_populations(R::Array{<:Unitful.Frequency, 5},
                            C::Array{<:Unitful.Frequency, 5},
                            atom_density::Array{<:NumberDensity,3})

Update populations on a regular grid with Statistical Equilibrium
"""
function get_revised_populations(R::Array{<:Unitful.Frequency, 5},
                                 C::Array{<:Unitful.Frequency, 5},
                                 atom_density::Array{<:NumberDensity,3})

    # Total rates
    P = R .+ C

    n_levels = size(P)[1] - 1
    nz, nx, ny = size(atom_density)

    A = Array{Float64, 5}(undef, (n_levels, n_levels, nz, nx, ny))u"s^-1"
    b = Array{Float64, 4}(undef, (n_levels, nz, nx, ny))u"s^-1*m^-3"
    populations = Array{Float64, 4}(undef, (nz, nx, ny, n_levels + 1))u"m^-3"

    for r = 1:n_levels
        A[r, r, :, :, :] = P[1, r+1, :, :, :] .+ P[r+1, 1, :, :, :]
        for c = setdiff(1:n_levels, r)
            A[r, c, :, :, :] = P[1, r+1, :, :, :] .- P[c+1, r+1, :, :, :]
            A[r, r, :, :, :] .+= P[r+1, c+1, :, :, :]
        end

        b[r, :, :, :] = atom_density .* P[1, r+1, :, :, :]
    end

    for k = 1:nz
        for i = 1:nx
            for j = 1:ny
                populations[k, i, j, 2:end] .= inv(A[:, :, k, i, j])*b[:, k, i, j]
            end
        end
    end

    populations[: ,: ,: , 1] = atom_density .- sum(populations[:, :, :, 2:end], dims=4)[:, :, :, 1]

    return populations
end

"""
    get_revised_populations(R::Array{<:Unitful.Frequency, 3},
                            C::Array{<:Unitful.Frequency, 3},
                            atom_density::Vector{<:NumberDensity})

Update populations on an irregular grid with Statistical Equilibrium
"""
function get_revised_populations(R::Array{<:Unitful.Frequency, 3},
                                 C::Array{<:Unitful.Frequency, 3},
                                 atom_density::Vector{<:NumberDensity})

    P = R .+ C

    n_levels = size(P)[1] - 1
    n = length(atom_density)

    A = Array{Float64, 3}(undef, (n_levels, n_levels, n))u"s^-1"
    b = Matrix{Float64}(undef, (n_levels, n))u"s^-1*m^-3"
    populations = Matrix{Float64}(undef, (n, n_levels + 1))u"m^-3"

    for r=1:n_levels
        A[r,r,:] = P[1,r+1,:] .+ P[r+1,1,:]
        for c=setdiff(1:n_levels, r)
            A[r,c,:] = P[1,r+1,:] .- P[c+1,r+1,:]
            A[r,r,:] .+= P[r+1,c+1,:]
        end

        b[r,:,:,:] = atom_density .* P[1,r+1,:,:,:]
    end

    for ii=1:n
        populations[ii, 2:end] .= inv(A[:, :, ii])*b[:, ii]
    end

    populations[:, 1] = atom_density .- sum(populations[:, 2:end], dims=2)[:,1]

    return populations
end

ext: jl

max_stars_repo_name: meudnaes/VoronoiRT

max_issues_repo_name: meudnaes/VoronoiRT

max_forks_repo_name: meudnaes/VoronoiRT


text:
logistic(x) = 1 ./ (1 + exp.(-x))
@scalardiff logistic(x::Number) 1 (logistic(x) .* (1 .- logistic(x)))

function autoencoder_cost(We1, We2, Wd, b1, b2, x)
    firstLayer = logistic(We1 * x .+ b1)
    encodedInput = logistic(We2 * firstLayer .+ b2)
    reconstructedInput = logistic(Wd * encodedInput)
    cost = sum((reconstructedInput .- x) .^ 2.0)
    return cost
end

function mlp1(w1, w2, w3, x1)
    xx2 = w1 * x1
    x2 = log.(1. + exp.(xx2))
    xx3 = w2 * x2
    x3 = log.(1. + exp.(xx3))
    x4 = w3 * x3
    sum(1. ./ (1. + exp.(-x4)))
end

function mlp2(w1, w2, w3, b1, b2,  b3, x1)
    x2 = logistic(w1 * x1 .+ b1)
    x3 = logistic(w2 * x2 .+ b2)
    x4 = logistic(w3 * x3 .+ b3)
    sum(x4)
end

function rnn(Wxh, Whh, Wy, hprev, x, y)
    h = tanh.(Whh * hprev + Wxh * x)
    yhat = Why * h
    cost = sum((yhat .- y) .^ 2.0)
    return h, cost
end


text:
# This function is very important for the computation of Floquet multipliers and checks that the eigensolvers compute the eigenvalues with largest modulus instead of their default behaviour which is with largest real part. If this option is not properly set, bifurcations of periodic orbits will be missed.
function checkFloquetOptions(eigls::AbstractEigenSolver)
	if eigls isa DefaultEig
		return @set eigls.which = abs
	end

	if eigls isa EigArpack
		return setproperties(eigls; which = :LM, by = abs)
	end

	if eigls isa EigArnoldiMethod
		return setproperties(eigls; which = ArnoldiMethod.LM(), by = abs)
	end

	if eigls isa EigKrylovKit
		return @set eigls.which = :LM
	end
end
####################################################################################################
# Computation of Floquet Coefficients for periodic orbits problems based on Finite Differences

"""
	floquet = FloquetQaDTrap(eigsolver::AbstractEigenSolver)

This composite type implements the computation of the eigenvalues of the monodromy matrix in the case of periodic orbits problems based on Finite Differences (Trapeze method), also called the Floquet multipliers. The method, dubbed Quick and Dirty (QaD), is not numerically very precise for large / small Floquet exponents. It allows, nevertheless, to detect bifurcations. The arguments are as follows:
- `eigsolver::AbstractEigenSolver` solver used to compute the eigenvalues.

If `eigsolver isa DefaultEig`, then the monodromy matrix is formed and its eigenvalues are computed. Otherwise, a Matrix-Free version of the monodromy is used.

!!! danger "Floquet multipliers computation"
    The computation of Floquet multipliers is necessary for the detection of bifurcations of periodic orbits (which is done by analyzing the Floquet exponents obtained from the Floquet multipliers). Hence, the eigensolver `eigsolver` needs to compute the eigenvalues with largest modulus (and not with largest real part which is their default behavior). This can be done by changing the option `which = :LM` of `eigsolver`. Nevertheless, note that for most implemented eigensolvers in the current Package, the proper option is set.
"""
struct FloquetQaDTrap{E <: AbstractEigenSolver } <: AbstractFloquetSolver
	eigsolver::E
	function FloquetQaDTrap(eigls::AbstractEigenSolver)
		eigls2 = checkFloquetOptions(eigls)
		return new{typeof(eigls2)}(eigls2)
	end
end

"""
Matrix-Free expression expression of the Monodromy matrix for the periodic problem computed at the space-time guess: `u0`
"""
function MonodromyQaDFD(poPb::PeriodicOrbitTrapProblem, u0::AbstractVector, par, du::AbstractVector)
	# extraction of various constants
	M = poPb.M
	N = poPb.N

	# period of the cycle
	T = extractPeriodFDTrap(u0)

	# time step
	h =  T * getTimeStep(poPb, 1)
	Typeh = typeof(h)

	out = copy(du)

	u0c = extractTimeSlices(u0, N, M)

	@views out .= out .+ h/2 .* apply(poPb.J(u0c[:, M-1], par), out)
	# res = (I - h/2 * poPb.J(u0c[:, 1])) \ out
	@views res, _ = poPb.linsolver(poPb.J(u0c[:, 1], par), out; a₀ = convert(Typeh, 1), a₁ = -h/2)
	out .= res

	for ii in 2:M-1
		h =  T * getTimeStep(poPb, ii)
		@views out .= out .+ h/2 .* apply(poPb.J(u0c[:, ii-1], par), out)
		# res = (I - h/2 * poPb.J(u0c[:, ii])) \ out
		@views res, _ = poPb.linsolver(poPb.J(u0c[:, ii], par), out; a₀ = convert(Typeh, 1), a₁ = -h/2)
		out .= res
	end

	return out
end

function MonodromyQaDFD(::Val{:ExtractEigenVector}, poPb::PeriodicOrbitTrapProblem, u0::AbstractVector, par, du::AbstractVector)
	# extraction of various constants
	M = poPb.M
	N = poPb.N

	# period of the cycle
	T = extractPeriodFDTrap(u0)

	# time step
	h =  T * getTimeStep(poPb, 1)
	Typeh = typeof(h)

	out = copy(du)

	u0c = extractTimeSlices(u0, N, M)

	@views out .= out .+ h/2 .* apply(poPb.J(u0c[:, M-1], par), out)
	# res = (I - h/2 * poPb.J(u0c[:, 1])) \ out
	@views res, _ = poPb.linsolver(poPb.J(u0c[:, 1], par), out; a₀ = convert(Typeh, 1), a₁ = -h/2)
	out .= res
	out_a = [copy(out)]
	# push!(out_a, copy(out))

	for ii in 2:M-1
		h =  T * getTimeStep(poPb, ii)
		@views out .= out .+ h/2 .* apply(poPb.J(u0c[:, ii-1], par), out)
		# res = (I - h/2 * poPb.J(u0c[:, ii])) \ out
		@views res, _ = poPb.linsolver(poPb.J(u0c[:, ii], par), out; a₀ = convert(Typeh, 1), a₁ = -h/2)
		out .= res
		push!(out_a, copy(out))
	end
	push!(out_a, copy(du))

	return out_a
end

# Compute the monodromy matrix at `u0` explicitely, not suitable for large systems
function MonodromyQaDFD(poPb::PeriodicOrbitTrapProblem, u0::AbstractVector, par)
	# extraction of various constants
	M = poPb.M
	N = poPb.N

	# period of the cycle
	T = extractPeriodFDTrap(u0)

	# time step
	h =  T * getTimeStep(poPb, 1)

	u0c = extractTimeSlices(u0, N, M)

	@views mono = Array(I - h/2 * (poPb.J(u0c[:, 1], par))) \ Array(I + h/2 * poPb.J(u0c[:, M-1], par))
	temp = similar(mono)

	for ii in 2:M-1
		# for some reason, the next line is faster than doing (I - h/2 * (poPb.J(u0c[:, ii]))) \ ...
		# also I - h/2 .* J seems to hurt (a little) performances
		h =  T * getTimeStep(poPb, ii)
		@views temp = Array(I - h/2 * (poPb.J(u0c[:, ii], par))) \ Array(I + h/2 * poPb.J(u0c[:, ii-1], par))
		mono .= temp * mono
	end
	return mono
end

function (fl::FloquetQaDTrap)(J, nev)
	if fl.eigsolver isa DefaultEig
		# we build the monodromy matrix and compute the spectrum
		monodromy = MonodromyQaDFD(J.pb, J.orbitguess0, J.par)
	else
		# we use a Matrix Free version
		monodromy = dx -> MonodromyQaDFD(J.pb, J.orbitguess0, J.par, dx)
	end
	vals, vecs, cv, info = fl.eigsolver(monodromy, nev)
	# the `vals` should be sorted by largest modulus, but we need the log of them sorted this way
	logvals = log.(complex.(vals))
	I = sortperm(logvals, by = x-> real(x), rev = true)
	# Base.display(logvals)
	return logvals[I], geteigenvector(fl.eigsolver, vecs, I), cv, info
end

####################################################################################################
# Computation of Floquet Coefficients for periodic orbit problems based on Shooting

"""
	floquet = FloquetQaDShooting(eigsolver::AbstractEigenSolver)

This composite type implements the computation of the eigenvalues of the monodromy matrix in the case of periodic orbits problems based on the Shooting method, also called the Floquet multipliers. The method, dubbed Quick and Dirty (QaD), is not numerically very precise for large / small Floquet exponents. It allows, nevertheless, to detect bifurcations. The arguments are as follows:
- `eigsolver::AbstractEigenSolver` solver used to compute the eigenvalues.

If `eigsolver == DefaultEig()`, then the monodromy matrix is formed and its eigenvalues are computed. Otherwise, a Matrix-Free version of the monodromy is used.

!!! danger "Floquet multipliers computation"
    The computation of Floquet multipliers is necessary for the detection of bifurcations of periodic orbits (which is done by analyzing the Floquet exponents obtained from the Floquet multipliers). Hence, the eigensolver `eigsolver` needs to compute the eigenvalues with largest modulus (and not with largest real part which is their default behavior). This can be done by changing the option `which = :LM` of `eigsolver`. Nevertheless, note that for most implemented eigensolvers in the current Package, the proper option is set.
"""
struct FloquetQaDShooting{E <: AbstractEigenSolver } <: AbstractFloquetSolver
	eigsolver::E
	function FloquetQaDShooting(eigls::AbstractEigenSolver)
		eigls2 = checkFloquetOptions(eigls)
		return new{typeof(eigls2)}(eigls2)
	end
end

function (fl::FloquetQaDShooting)(J, nev)
	if fl.eigsolver isa DefaultEig
		@warn "Not implemented yet in a fast way! Need to form the full monodromy matrix, not practical for large scale problems"
		# we build the monodromy matrix and compute the spectrum
		monodromy = MonodromyQaDShooting(J.pb, J.x, J.par)
	else
		# we use a Matrix Free version
		monodromy = dx -> MonodromyQaDShooting(J.pb, J.x, J.par, dx)
	end
	vals, vecs, cv, info = fl.eigsolver(monodromy, nev)

	# the `vals` should be sorted by largest modulus, but we need the log of them sorted this way
	logvals = log.(complex.(vals))
	# Base.display(logvals)
	I = sortperm(logvals, by = x-> real(x), rev = true)
	return logvals[I], geteigenvector(fl.eigsolver, vecs, I), cv, info
end

"""
Matrix-Free expression expression of the Monodromy matrix for the periodic problem based on Shooting computed at the space-time guess: `x`. The dimension of `u0` is N * M + 1 and the one of `du` is N.
"""
function MonodromyQaDShooting(sh::ShootingProblem, x::AbstractVector, p, du::AbstractVector)
	# period of the cycle
	T = extractPeriodShooting(x)

	# extract parameters
	M = length(sh.ds)
	N = div(length(x) - 1, M)

	# extract the time slices
	xv = @view x[1:end-1]
	xc = reshape(xv, N, M)

	out = copy(du)

	for ii in 1:M
		# call the jacobian of the flow
		@views out .= sh.flow(Val(:SerialdFlow), xc[:, ii], p, out, sh.ds[ii] * T).du
	end
	return out
end

# Compute the monodromy matrix at `x` explicitely, not suitable for large systems
function MonodromyQaDShooting(sh::ShootingProblem, x::AbstractVector, p)
	# period of the cycle
	T = extractPeriodShooting(x)

	# extract parameters
	M = length(sh.ds)
	M > 1 && @error "This is not yet a practical approach for multiple shooting"

	N = div(length(x) - 1, M)

	Mono = zeros(N, N)

	# extract the time slices
	xv = @view x[1:end-1]
	xc = reshape(xv, N, M)
	du = zeros(N)

	for ii in 1:N
		du[ii] = 1.0
		# call jacobian of the flow
		@views Mono[:, ii] .= sh.flow(xc[:, 1], p, du, T).du
		du[ii] = 0.0
	end

	return Mono
end

"""
	MonodromyQaDShooting(hpsh::PoincareShootingProblem, x_bar, p, dx_bar::AbstractVector)

Matrix-Free expression of the Monodromy matrix for the periodic problem based on Poincaré Shooting computed at the space-time guess: `x`. The dimension of `x` is N * M and the one of `du` is N. If we denote by
"""
function MonodromyQaDShooting(psh::PoincareShootingProblem, x_bar, p, dx_bar::AbstractVector)
	M = getM(psh)
	Nm1 = div(length(x_bar), M)

	# reshape the period orbit guess into a Matrix
	x_barc = reshape(x_bar, Nm1, M)
	@assert length(dx_bar) == Nm1 "Please provide the right dimension to your matrix-free eigensolver, it must be $Nm1."

	xc = similar(x_bar, Nm1 + 1)
	outbar = copy(dx_bar)
	outc = similar(dx_bar, Nm1 + 1)

	for ii in 1:M
		E!(psh.section,  xc,  view(x_barc, :, ii), ii)
		dE!(psh.section, outc, outbar, ii)
		outc .= diffPoincareMap(psh, xc, p, outc, ii)
		# check to <outc, normals[ii]> = 0
		# println("--> ii=$ii, <out, normali> = ", dot(outc, sh.section.normals[ii]))
		dR!(psh.section, outbar, outc, ii)
	end
	# @assert 1==0
	return outbar

end

function MonodromyQaDShooting(psh::PoincareShootingProblem, x)
	@assert 1==0 "WIP, no done yet!"
end


text:
module multi
include("moo_functions.jl")
include("plot_utils.jl")
using Plots

const POP_SIZE = 100
const DIM = 10
const MAX_GEN = 50
const CROSS_PROB = 0.2
const MUT_PROB = 0.8
const MUT_STEP = 0.05
const REPEATS = 10
const NELITES = 0
const HYP_REF = (11.,11.)

const EXPERIMENT_NAME = "multi"
const COMPARE_TO = "none"
const GENS_PER_FRAME = 5
const LOG_SIZE = div(MAX_GEN, GENS_PER_FRAME)

mutable struct Individual
    x::Vector{Float64}
    f::Tuple{Float64, Float64}
    front::Int64
    ssc::Float64 #secondary sorting criterion
    Individual() = begin
        x = rand(DIM)
        f = (0,0)
        front = 0
        ssc = 0.0
        return new(x,f,front,ssc)
    end
end

Base.isless(ind1::Individual, ind2::Individual) = begin
    return (ind1.front, ind1.f) < (ind2.front, ind2.f)
end

const Population = Vector{Individual}

function isdominated(ind1::Individual, ind2::Individual)::Bool
    return ind1.f[1] >= ind2.f[1] && ind1.f[2] >= ind2.f[2] && ind1.f != ind2.f
end

function isdominated(ind::Individual, pop::Population; first = 1)::Bool
    for k in first:POP_SIZE
        if isdominated(ind, pop[k])
            return true
        end
    end
    return false
end

#find fronts and reorder pop accordingly
function findfronts!(pop::Population)
    for k in 1:POP_SIZE
        pop[k].front = 0
    end
    front = 0
    removed = 0
    while removed < POP_SIZE
        sorted = 0
        for l in (removed+1):POP_SIZE
            if !isdominated(pop[l], pop, first = (removed+1))
                pop[l].front = front
                sorted += 1
                #swap
                k = removed + sorted
                temp = pop[k]
                pop[k] = pop[l]
                pop[l] = temp
            end
        end
        front += 1
        removed += sorted
    end
end

#find ssc using crowding dist
function findssc!(pop::Population)
    for k in 1:POP_SIZE
        ind = pop[k]
        if k == 1 || k == POP_SIZE
            ind.ssc = Inf
            continue
        end
        ind0 = pop[k-1]
        ind1 = pop[k+1]
        if ind0.front < ind.front || ind.front < ind1.front
            ind.ssc = Inf
            continue
        end
        ind.ssc = ind1.f[1] - ind0.f[1] + ind0.f[2] - ind1.f[2]
    end
end

function hypervolume(pop::Population)::Float64
    vol = (HYP_REF[1] - pop[1].f[1])*(HYP_REF[2] - pop[1].f[2])
    for k in 2:POP_SIZE
        if pop[k].front == 0
            break
        end
        vol += (HYP_REF[1] - pop[k].f[1])*(pop[k-1].f[2] - pop[k].f[2])
    end
    return vol
end

function eval!(ind::Individual, fun::Function)
    ind.f = fun(ind.x)
end

function breed!(parents::Population, childs::Population)
    npairs = div(POP_SIZE, 2)
    Threads.@threads for n in 1:npairs
        parent1 = tournament_selection(parents)
        parent2 = tournament_selection(parents)
        child1 = childs[n]
        child2 = childs[n + npairs]
        copy!(child1, parent1)
        copy!(child2, parent2)
        cross!(parent1, parent2, child1, child2)
    end
end

function tournament_selection(pop::Population)::Individual
    ind1 = rand(pop)
    ind2 = rand(pop)
    if ind1.front < ind2.front
        return ind1
    elseif ind1.front == ind2.front && ind1.ssc > ind2.ssc
        return ind1
    end
    return ind2
end

function copy!(dest::Individual, src::Individual)
    dest.x .= src.x
    #dest.mut_var = src.mut_var
end

function cross!(
        parent1::Individual, parent2::Individual,
        child1::Individual, child2::Individual
        )
    if rand() < CROSS_PROB
        cross_point = rand(1:(DIM-1))
        for i in 1:cross_point
            child1.x[i] = parent2.x[i]
            child2.x[i] = parent1.x[i]
        end
    end
end

function mutate!(ind::Individual)
    if rand() < MUT_PROB
        for k in 1:DIM
            ind.x[k] += MUT_STEP*randn()
            ind.x[k] = clamp(ind.x[k], 0., 1.)
        end
    end
end

function eva(fun::Function, opt_hypvol::Float64)
    pop = [Individual() for _ in 1:POP_SIZE]
    buffer = [Individual() for _ in 1:POP_SIZE]
    log_hypvol = Float64[]
    for k in 1:POP_SIZE
        eval!(pop[k], fun)
        eval!(buffer[k], fun)
    end
    findfronts!(pop)
    sort!(pop)
    findssc!(pop)
    #then iterate:
    for gen in 1:MAX_GEN
        #generate offspring
        breed!(pop, buffer)
        #kill the parents, except elites
        Threads.@threads for n in NELITES+1:POP_SIZE
            temp = pop[n]
            pop[n] = buffer[n]
            buffer[n] = temp
        end
        #mutate and calculate fitness
        Threads.@threads for ind in buffer
            mutate!(ind)
            eval!(ind, fun)
        end
        findfronts!(pop)
        sort!(pop)
        findssc!(pop)
        #save results
        if (gen % GENS_PER_FRAME == 0)
            @show(gen)
            obj = opt_hypvol - hypervolume(pop)
            @show(obj)
            push!(log_hypvol, obj)
        end
    end
    #extract nondom front
    po_front = filter(ind -> ind.front == 0, pop)
    return (po_front, log_hypvol)
end

function main()
    funs = moo_functions.zdt()
    opt_hypvols = moo_functions.opt_hvs()
    for n in [1]
        println()
        println("function #", n)
        println("============================================")
        log = zeros(REPEATS, LOG_SIZE)
        for experiment in 1:REPEATS
            println("experiment #", experiment)
            (result, log_hypvol) = eva(funs[n], opt_hypvols[n])
            for gen in 1:LOG_SIZE
                log[experiment, gen] = log_hypvol[gen]
            end
            #save result as a plot
            if experiment == 1
                f1 = [ind.f[1] for ind in result]
                f2 = [ind.f[2] for ind in result]
                fig = scatter(f1, f2, xlabel = "f1", ylabel = "f2")
                savefig(fig, "julia-multi-results/"*EXPERIMENT_NAME*string(n)*"_graph.pdf")
            end
        end
        t = LinRange(0, MAX_GEN*POP_SIZE, LOG_SIZE)
        new_path = "julia-multi-results/"*EXPERIMENT_NAME*string(n)*".log"
        old_path = "julia-multi-results/"*COMPARE_TO*string(n)*".log"
        fig_path = "julia-multi-results/"*EXPERIMENT_NAME*string(n)*".pdf"
        plot_utils.save_log(t, log, new_path)
        if ispath(old_path)
            p = plot_utils.plot_log(old_path, new_path, symbols = ["old","new"])
            savefig(p, fig_path)
        else
            p = plot_utils.plot_log(new_path)
            savefig(p, fig_path)
        end
    end
end

end

max_stars_repo_name: OndrejKincl/evaTeaching-python

max_issues_repo_name: OndrejKincl/evaTeaching-python

max_forks_repo_name: OndrejKincl/evaTeaching-python


text:
"""
    Module Basis1D
General math tools for 1D bases used by DG
"""

module Basis1D

export gauss_lobatto_quad, gauss_quad
export jacobiP, grad_jacobiP
export vandermonde1D, grad_vandermonde1D
export chebyshev
export dPhi1D

using LinearAlgebra
using SpecialFunctions

"""
    gauss_lobatto_quad(degree, α=0, β=0)
Initialize Gauss-Legendre-Lobatto (GLL) quadrature nodes and weights (x,w)
to be able to exactly integrate (α,β) Jacobi polynomial of given degree
"""
function gauss_lobatto_quad(degree, α=0, β=0)
    # Need n points to integrate polynomials of degree 2n-3
    n = ceil(Int, (degree+3)/2)
    p = n-1
    if (α!=0) && (β!=0)
        error("currently only allows alpha == beta == 0")
    end
    x = zeros(p+1)
    w = zeros(p+1)
    if p == 0
        x[1] = 0
        w[1] = 2
    elseif p == 1
        x[1] = -1.0
        x[2] = 1.0
        w[1] = 1.0
        w[2] = 1.0
    else
        xint, w = gauss_quad(2*p-3, α+1, β+1)
        x = [-1 transpose(xint) 1]

        V = vandermonde1D(p,x)
        w = vec(sum(inv(V*V'),dims=2))
    end
    return x[:],w[:]
end

"""
    gauss_quad(degree, α, β)
Initialize Gaussian quadrature nodes and weights (x,w) to be able to exactly
integrate Jacobi Polynomial (α,β) of given degree
"""
function gauss_quad(degree, α=0, β=0)
    # Need n points to integrate polynomials of degree 2n-1
    n = ceil(Int, (degree+1)/2)
    p = n-1
    if p == 0
        x = [-(α-β)/(α+β+2)]
        w = [2]
        return x, w
    end

    J = zeros(p+1, p+1)
    h₁ = @. 2*(0:p)+α+β
    J = diagm(0 => (@. -1/2*(α^2-β^2)/(h₁+2)/h₁), 1 => (@. 2/(h₁[1:p]+2)*sqrt((1:p)*((1:p)+α+β)*((1:p)+α)*((1:p)+β)/(h₁[1:p]+1)/(h₁[1:p]+3))))
    if α+β<10*eps()
        J[1,1] = 0.0
    end
    J = J + transpose(J) # Finalize symmetric tridiagonal Jacobi matrix

    x, V = eigen(J)
    w = @. transpose(V[1,:])^2*2^(α+β+1)/(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+1)
    return x[:], w[:]
end

"""
    grad_jacobiP(p, r, α, β)
Evaluate the derivative of Jacobi Polynomial (α, β) of order p at nodes r
"""
function grad_jacobiP(p, r, α, β)
    dP = zeros(length(r))
    if p != 0
        dP = sqrt(p*(p+α+β+1))*jacobiP(p-1,r,α+1,β+1)
    end
    return dP
end

"""
    jacobiP(p, x, α, β)
Evaluate Jacobi Polynomial (α, β) of order p at points x
"""
function jacobiP(p, x, α, β)
    xp = x
    if size(xp, 2) == 1
        xp = transpose(xp)
    end

    PL = zeros(p+1,length(xp))
    γ₀ = 2^(α+β+1)/(α+β+1)*gamma(α+1)*gamma(β+1)/gamma(α+β+1)
    PL[1,:] .= 1.0/sqrt(γ₀)
    if p == 0
        P = transpose(PL)
        return P
    end
    γ₁ = (α+1)*(β+1)/(α+β+3)*γ₀
    PL[2,:] = ((α+β+2).*xp/2 .+ (α-β)/2)/sqrt(γ₁)
    if p == 1
        P = PL[p+1,:]
        return P
    end

    aold = 2/(2+α+β)*sqrt((α+1)*(β+1)/(α+β+3))

    for i = 1:p-1
        h₁ = 2i+α+β
        anew = 2/(h₁+2)*sqrt((i+1)*(i+1+α+β)*(i+1+α)*(i+1+β)/(h₁+1)/(h₁+3))
        bnew = -(α^2-β^2)/h₁/(h₁+2)
        PL[i+2,:] = 1/anew*(-aold*transpose(PL[i,:]).+(xp.-bnew).*transpose(PL[i+1,:]))
        aold = anew
    end

    P = PL[p+1,:]
    return P;
end

"""
    vandermonde1D(p, r)
Initialize the 1D Vandermonde matrix of order p Legendre polynomials at nodes r
"""
function vandermonde1D(p, r)
    V1D = zeros(length(r), p+1)
    for j = 1:p+1
        V1D[:,j] = jacobiP(j-1, r[:], 0, 0)
    end
    return V1D
end

"""
    grad_vandermonde1D(p, r)
Initialize the 1D gradient Vandermonde matrix of order p Legendre polynomials at nodes r
"""
function grad_vandermonde1D(p, r)
    V1D = zeros(length(r), p+1)
    for j = 1:p+1
        V1D[:,j] = grad_jacobiP(j-1, r[:], 0, 0)
    end
    return V1D
end

"""
    chebyshev(p)
Return the 1D Chebyshev nodes of order p on [-1,1]
"""
function chebyshev(p)
    cheby = cos.((p:-1:0)*π/p)
end

"""
    interpolation_matrix1D(x_from, x_to)
Compute an interpolation matrix from a set of 1D points to another set of 1D points.
Assumes that points interpolating from provide enough accuracy (aka - they are
well spaced out and of high enough order), and define an interval. Points
interpolating onto can be of any size, but must be defined on this same interval.
Interpolation matrix ∈ ℜ^(size of x_to x size of x_from)
"""
function interpolation_matrix1D(x_from, x_to)
    # Create nodal representation of reference bases
    order = size(x_from,1) - 1 # Assumes order = (size of x_from) - 1
    n_from = size(x_from,1)
    l_from = vandermonde1D(order, x_from)

    eye = diagm(0=>ones(n_from))
    V = reshape(l_from, n_from,n_from)
    coeffs_phi = V \ eye

    # Compute reference bases on the output points
    l_to = vandermonde1D(order, x_to)
    n_to = size(l_to,1)
    V_to = reshape(l_to, n_to,n_from)

    # Construct interpolation matrix
    Interp1D = V_to*coeffs_phi
end

"""
    dPhi1D(x_from, x_to)
Compute the gradient of the basis functions defined by 1D points on another
set of 1D points.

Assumes that points interpolating from provide enough accuracy (aka - they are
well spaced out and of high enough order), and define an interval. Points
interpolating onto can be of any size, but must be defined on this same interval.
dPhi ∈ ℜ^(size of x_to × size of x_from × 1)
"""
function dPhi1D(x_from, x_to)
    dim = 1
    # Create nodal representation of reference bases
    order = size(x_from,1) - 1 # Assumes order = (size of x_from) - 1
    n_from = size(x_from,1)
    l_from = vandermonde1D(order, x_from)

    eye = diagm(0=>ones(n_from))
    V = reshape(l_from, n_from,n_from)
    coeffs_phi = V \ eye

    # Compute derivative of reference bases on the output points
    dl_to = grad_vandermonde1D(order, x_to)
    n_to = size(dl_to,1)
    dV = reshape(dl_to, n_to,n_from,dim)
    # Construct gradient of phi = dV*coeffs_phi
    dPhi_to = Array{Float64,3}(undef, n_to,n_from,dim)
    for l = 1:dim
        dPhi_to[:,:,l] = dV[:,:,l]*coeffs_phi
    end
    return dPhi_to
end

end

max_stars_repo_name: NoseKnowsAll/DGToolkit

max_issues_repo_name: NoseKnowsAll/DGToolkit

max_forks_repo_name: NoseKnowsAll/DGToolkit


text:
import Base: (^), (*), (~), (/), (|), (!)

(*)(x::Any, y::AbstractToken) = sSequence(parser(x),y)
(*)(x::AbstractToken, y::Any) = sSequence(x,parser(y))
"""
    (*)(x::Any, y::AbstractToken)
    (*)(x::AbstractToken, y::Any)
    (*)(x::AbstractToken, y::AbstractToken)

Chain parsers in [`sSequence`](@ref).
See also [`@seq`](@ref).
"""
(*)(x::AbstractToken, y::AbstractToken) = sSequence(x,y)

## todo: cuts

"""
    (/)(x::ParserTypes, y::ParserTypes)

`Sequence(PositiveLookbehind(x),y)`

```jldoctest
julia> match("is "/"match", "no match is match").start
13
```

!!! note
    This syntax is reviewed and I your appreciate your comments!
"""
(/)(x::ParserTypes, y::ParserTypes) =
    Sequence(PositiveLookbehind(x),y)

"""
    Base.broadcasted(::typeof((&)), x::CharNotIn, y::CharNotIn)

Character matchers `m` like `Union{CharIn,CharNotIn,T}`, or any
type `T` providing a `ismatch(m::T,c::Char)::Bool` method represent a
lazy bitarray for all characters.
"""
Base.broadcasted(::typeof((&)), x::CharNotIn, y::CharNotIn) =
    CharNotIn(x.pcre*y.pcre, x.sets,y.sets)

Base.broadcasted(::typeof((&)), x::CharIn, y::CharNotIn) =
    CharIn(setdiff(x.sets, y.sets))

Base.broadcasted(::typeof((&)), x::Union{CharIn,CharNotIn}, ::AnyChar) =
    x

Base.broadcasted(::typeof((&)), x::NamedParser, y) =
    NamedParser(x.name,x.parser .& y; doc=x.doc)

Base.broadcasted(::typeof((&)), x::JoinSubstring, y) =
    JoinSubstring(x.parser .& y)

Base.broadcasted(::typeof((&)), x::Transformation, y) =
    Transformation(x.transform, x.parser .& y)

Base.broadcasted(::typeof((&)), x::Repeat, y) =
    Repeat(x.range, x.parser .& y)

Base.broadcasted(::typeof((|)), x::CharIn, y::CharIn) =
    CharIn(x.pcre*y.pcre, x.sets,y.sets)

"""
    (!)(x::AbstractToken)

Parser Transformation getting the matched SubString.

```jldoctest
julia> parse(Repeat(CharIn(:L)),"abc123")
3-element Array{Char,1}:
 'a': ASCII/Unicode U+0061 (category Ll: Letter, lowercase)
 'b': ASCII/Unicode U+0062 (category Ll: Letter, lowercase)
 'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase)

julia> parse(!Repeat(CharIn(:L)),"abc123")
"abc"

```

"""
(!)(x::AbstractToken) = JoinSubstring(x)
using InternedStrings
import InternedStrings: intern
"""
    (!)(x::JoinSubstring)

Parser transformating result `v -> InternedStrings.intern(v)`.
"""
(!)(x::AbstractToken{<:SubString}) =
    instance(String, x)
(!)(x::Transformation{<:SubString}) =
    map(InternedStrings.intern, x.parser)

"""
    (!)(x::NamedParser)

Parser transformating result `v -> v=>x.name`.
"""
(!)(x::NamedParser) = map(v -> v => x.name, x)

(|)(x, y::ParserTypes) = sEither(parser(x),y)
(|)(x::ParserTypes, y) = sEither(x,parser(y))
"""
    (|)(x::AbstractToken, y)
    (|)(x, y::AbstractToken)
    (|)(x::AbstractToken, y::AbstractToken)

Operator syntax for `sEither(x, y)`.

```jldoctest
julia> 'a' | CharIn("AB") | "bc"
|🗄... Either
├─ a
├─ [AB] CharIn
└─ bc
::Union{Char, SubString}

```

"""
(|)(x::ParserTypes, y::ParserTypes) = sEither(x,y)

"""
    (|)(x::AbstractToken{T}, default::Union{T,Missing})

Operator syntax for `Optional(x, default=default)`.

```jldoctest
julia> parser("abc") | "nothing"
|🗄... Either
├─ abc
└─ nothing
::SubString

julia> parser("abc") | missing
abc? |missing
::Union{Missing, SubString}

```

"""
function (|)(x::AbstractToken{T}, default::Union{T,Missing}) where { T }
    Optional(x,default=default)
end
function (|)(x::Char, y::Char)
    CharIn(tuple(x,y))
end
function (|)(x::CharIn, y::Char)
    CharIn(tuple(x.sets...,y))
end

"""
    `(|)(x::Either, T::Type)`

Return new Either with `T` added to result_type(x).
todo: Note that the options array is kept. As a consequence `push!`on result will also push to `x`.
"""
(|)(x::Either, T::Type) =
    Either{Union{result_type(x),T}}(x.options)


text:
module Experiments
using Particles, Distributions, DataFrames, Underscore, StatsBase

export Experiment,
    Result,
    run,
    run_gibbs,
    flatten,
    dictofarrays,
    experiments,
    results_summary

struct Experiment
    params::Dict{Symbol,Any}
    seed::UInt
end

Experiment(params::Dict) = Experiment(params, rand(UInt))

"""
    Result

Holds the result of an experiment:

* `K_modal` the modal number of clusters (highest posterior probability)
* `K_mean` expected number of clusters
* `K_map` number of clusters in particle with highest overall posterior
* `p_of_K` the posterior probability of each number of clusters ``p(K | x^(i), w^(i))``

"""
struct Result
    K_modal::Int
    K_mean::Float64
    K_map::Int
    p_of_K::Vector{Float64}
end

function Result(ps::ParticleFilter)
    K_dist = ncomponents_dist(ps)
    p_of_K = probs(K_dist)
    K_modal = findmax(p_of_K)[2]
    K_mean = mean(K_dist)

    ps = particles(ps)
    best = findmax(marginal_log_posterior(p) for p in ps)

    K_map = ncomponents(ps[best[2]])
    return Result(K_modal, K_mean, K_map, p_of_K)
end

# quick and dirty: just get the entropy and the Result struct
result_entropy_assignsim(ps::ParticleFilter) =
    Result(ps), state_entropy(ps), assignment_similarity(ps)

result_entropy(ps::ParticleFilter) =
    Result(ps), state_entropy(ps)

function data!(ex::Experiment)
    if ex.params[:data] isa Function
        ex.params[:data_f] = ex.params[:data]
        ex.params[:data] = ex.params[:data_f](ex)
    end
end

import Base.run
function run(ex::Experiment, transform=Result; progress=false)
    println(join(["$k=>$v" for (k,v) in ex.params], ", "))
    data!(ex)

    params = ex.params

    prior = NormalInverseChisq(params[:μ_0], params[:σ2_0], params[:κ_0], params[:ν_0])
    ps = ChenLiuParticles(params[:num_particle], prior, ChineseRestaurantProcess(params[:α]))

    results = []
    srand(ex.seed)
    first = 1
    for k in sort(params[:num_obs])
        filter!(ps, view(ex.params[:data], first:k), progress)
        first = k+1
        push!(results, transform(ps))
    end

    return [results...]
end

function Result(gs::GibbsCRPSamples)
    K_dist = ncomponents_dist(gs)
    p_of_K = probs(K_dist)
    K_modal = findmax(p_of_K)[2]
    K_mean = mean(K_dist)

    MAP_components, i =
        findmax(marginal_log_posterior(cs, gs.g.logα) for cs in gs.components)

    K_map = length(MAP_components)
    Result(K_modal, K_mean, K_map, p_of_K)
end

function run_gibbs(ex::Experiment)
    params = ex.params

    prior = NormalInverseChisq(params[:μ_0], params[:σ2_0], params[:κ_0], params[:ν_0])
    samples = fit!(GibbsCRP(prior, params[:α], params[:data]), false,
                   samples=params[:num_particle], burnin=500)

    return Result(samples)
end

using Base.Iterators: product

# each experiment has multiple :num_obs, so we need to flatten...
"""
    flatten(d::Dict, k::Symbol)
    flatten(d::Dict)

Turn a dict of vectors (generally iterables) into an array of dicts.  Every
iterable value becomes a dimension in the resulting array
"""

flatten(d::Dict{K,V}, k::K) where {K,V} = (setindex!(copy(d), n, k) for n in d[k])
function flatten(d::Dict)
    (Dict(kv...) for kv in product([ (k=>s for s in v) for (k,v) in d ]...))
end

flatten(ex::Experiment, k::Symbol) = Experiment.(collect(flatten(ex.params, k)), ex.seed)

"""
    dictofarrays(ds::Array{Dict{K,V},N}) where {K,V,N}

The opposite of flatten: take an array of dicts and turn it into a dict of
arrays.

"""
function dictofarrays(ds::Array{Dict{K,V},N}) where {K,V,N}
    d = Dict{K,Any}()
    for k in keys(first(ds))
        d[k] = getindex.(ds, k)
    end
    return d
end

"""
    experiments(params::Dict)
    experiments(; kw...)

Construct an array of experiments from parameter values supplied as a Dict or
keyword arguments.  Each parameter value that's iterable will expand to a
"dimension" of the experiments array.

# Example

    experiments(μ = [-1, 0, 1],
                σ2 = [.1, 1.0, 10.0],
                x = ([1, 2, 3, 4], ))

Produces 9 `Experiment`s, one of each combination of μ and σ2.  Each of these
has the same value of `x`, since it was passed as a vector wrapped in a tuple.

"""
experiments(params::Dict) = [Experiment(d) for d in flatten(params)]
experiments(; kw...) = experiments(Dict(kw))

Dict(r::Result) = Dict([k=>getfield(r,k) for k in fieldnames(r)])
Dict(e::Experiment, r::Result) = merge(Dict(r), Dict(:seed=>e.seed), e.params)

flatten(es::Array{Experiment,N}, rs::Array{Vector{Result},N}, k::Symbol) where N =
    reduce(append!, [], (zip(flatten(e, k), r) for (e,r) in zip(es, rs)))

function results_summary(ers)
    @_ [Dict(er...) for er in ers] |>
    dictofarrays |>
    DataFrame |>
    by(_, [:num_particle, :num_obs, :α]) do d
        DataFrame(low = mean(d[:K_map] .< 2),
                  success = mean(d[:K_map] .== 2),
                  high = mean(d[:K_map] .> 2),
                  avg_K_map = mean(d[:K_map]),
                  avg_K_mean = mean(d[:K_mean]),
                  avg_p_K2 = mean(get.(d[:p_of_K], 2, 0)))
    end
end

end # module Experiments


text:
using RCall,Blink
w = Window()
sleep(2)

specs = ["CH4","CH3O2NO2","H2","N2O5","H2O2","CH3OH","HONO","CO","HO2NO2","CH3O","CH3NO3","HNO3","CH3OOH","HCHO","CH3O2","NO3","OH","NO2","NO","O3","HO2"]

for i in 1:length(specs)
species = specs[i]
@rput species
rfile = pwd()*"/points.R"
@rput rfile
R"source(rfile)"
sleep(2)

    loadurl(w, "file://"*pwd()*"/index.html")
    sleep(1)

try

    @js w species=$species
    R"names = points($(species),maxno,minno)"
    sleep(2)
    @rget names
    @js w cl=$(names)
    @js w eval("window.coords=cl.coords[0].map((d,i)=>[d,cl.coords[1][i]])")
    sleep(1)
    @js w df=$(names[Symbol("df")])
    @js w sum =  $(normalize(names[Symbol("sum")]))
    @js w cols =  $(names[Symbol("class")])
    @js w points()

    sleep(3)
    @js w savesvg()

  catch e
            println("caught an error $e")

end

end
#

max_stars_repo_name: wolfiex/VisACC

max_issues_repo_name: wolfiex/VisACC

max_forks_repo_name: wolfiex/VisACC


text:
using Revise

using PathTracing

using ImageView, Gtk, GtkReactive

# Create a canvas `c`. There are other approaches, like stealing one from a previous call
# to `imshow`, or using GtkReactive directly.
guidict = imshow_gui((512, 384), name="Path Tracer")
frame = guidict["frame"]
canvas = guidict["canvas"]
# To see anything you have to call `showall` on the window (once)
Gtk.showall(guidict["window"])

# This dance with the frame and zoom region signal `zrsig` is to preserve the
# aspect ratio. Why is this *soo* annoying!? :-( :-(
img = PathTracing.main()
zrsig = Signal(ZoomRegion(img))
imgsig = Signal(img)
imshow(frame, canvas, imgsig, zrsig)

@async Revise.entr([], all=true, postpone=true) do
    global img = PathTracing.main()
    push!(imgsig, img)
    push!(zrsig, ZoomRegion(img))
end


text:
function f(x::AbstractIntVar,alpha::Float64)
    return 1/(assignedValue(x)+1)^alpha
end

"""
    struct CPReward2 <: AbstractReward end
CPReward2 is a variant of CPReward with better theoretical properties concerning the variations of the function..
"""
mutable struct CPReward2 <: AbstractReward
    value::Float32
end

CPReward2(model::CPModel) = CPReward2(0)

"""
    set_reward!(::StepPhase, lh::LearnedHeuristic{DefaultReward, A}, model::CPModel, symbol::Union{Nothing, Symbol})
Change the "current_reward" attribute of the LearnedHeuristic at the StepPhase.
"""
function set_reward!(::Type{StepPhase}, lh::LearnedHeuristic{SR, CPReward2, A}, model::CPModel, symbol::Union{Nothing, Symbol}) where {
    SR <: AbstractStateRepresentation,
    A <: ActionOutput
}
    nothing
end

"""
    set_reward!(::Type{DecisionPhase}, lh::LearnedHeuristic{SR, CPReward2, A}, model::CPModel)
Change the current reward at the DecisionPhase. This is called right before making the next decision, so you know you have the very last state before the new decision
and every computation like fixPoints and backtracking has been done.
"""
function set_reward!(::Type{DecisionPhase}, lh::LearnedHeuristic{SR, CPReward2, A}, model::CPModel) where {
    SR <: AbstractStateRepresentation,
    A <: ActionOutput
}
    nothing
end

"""
    set_reward!(::Type{EndingPhase}, lh::LearnedHeuristic{SR, CPReward2, A}, model::CPModel, symbol::Union{Nothing, Symbol})
Increment the current reward at the EndingPhase. Called when the search is finished by an optimality proof or by a limit in term of nodes or
in terms of number of solution. This is useful to add some general results to the reward like the number of ndoes visited during the episode for instance.
"""
function set_reward!(::Type{EndingPhase}, lh::LearnedHeuristic{SR, CPReward2, A}, model::CPModel, symbol::Union{Nothing, Symbol}) where {
    SR <: AbstractStateRepresentation,
    A <: ActionOutput
}
    alpha = 0.5
    if symbol == :FoundSolution
        if isnothing(model.objective)
            lh.reward.value += 1
        else
            lh.reward.value += f(model.objective, alpha)
        end
    else
        q = length(branchable_variables(model))/nb_boundvariables(model)
        lh.reward.value += 0.5*(0.5/q^2 - q)
    end
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/MortarContact2D.jl/blob/master/LICENSE

using FEMBase, MortarContact2D, Test
using MortarContact2D: calculate_normals,
                       project_from_master_to_slave,
                       project_from_slave_to_master

X = Dict(
    1 => [0.0, 1.0],
    2 => [5/4, 1.0],
    3 => [2.0, 1.0],
    4 => [0.0, 1.0],
    5 => [3/4, 1.0],
    6 => [2.0, 1.0])
mel1 = Element(Seg2, [1, 2])
mel2 = Element(Seg2, [2, 3])
sel1 = Element(Seg2, [4, 5])
sel2 = Element(Seg2, [5, 6])
update!([mel1, mel2, sel1, sel2], "geometry", X)
master_elements = [mel1, mel2]
slave_elements = [sel1, sel2]
normals = calculate_normals(slave_elements, 0.0)
update!(slave_elements, "normal", 0.0 => normals)

X1 = [0.0, 1.0]
n1 = [0.0, 1.0]
xi2 = project_from_slave_to_master(mel1, X1, n1, 0.0)
@test isapprox(xi2, -1.0)

X1 = [0.75, 1.0]
n1 = [0.00, 1.0]
xi2 = project_from_slave_to_master(mel1, X1, n1, 0.0)
@test isapprox(xi2, 0.2)

X2 = mel1("geometry", xi2, 0.0)
@test isapprox(X2, [3/4, 1.0])

X2 = [0.0, 1.0]
xi1 = project_from_master_to_slave(sel1, X2, 0.0)
@test isapprox(xi1, -1.0)
X2 = [1.25, 1.00]
xi1 = project_from_master_to_slave(sel1, X2, 0.0)
X1 = sel1("geometry", xi1, 0.0)
@test isapprox(X1, [5/4, 1.0])

X = Dict(
    1 => [0.0, 0.0],
    2 => [0.0, 1.0],
    3 => [0.0, 1.0],
    4 => [0.0, 0.0])

sel1 = Element(Seg2, [1, 2])
mel1 = Element(Seg2, [3, 4])
update!([sel1, mel1], "geometry", X)
slave_elements = [sel1]
normals = calculate_normals(slave_elements, 0.0)
update!(slave_elements, "normal", 0.0 => normals)

X2 = mel1("geometry", (-1.0, ), 0.0)
xi = project_from_master_to_slave(sel1, X2, 0.0)
@test isapprox(xi, 1.0)

X2 = mel1("geometry", (1.0, ), 0.0)
xi = project_from_master_to_slave(sel1, X2, 0.0)
@test isapprox(xi, -1.0)

X1 = sel1("geometry", (-1.0, ), 0.0)
n1 = sel1("normal", (-1.0, ), 0.0)
xi = project_from_slave_to_master(mel1, X1, n1, 0.0)
@test isapprox(xi, 1.0)

X1 = sel1("geometry", (1.0, ), 0.0)
n1 = sel1("normal", (1.0, ), 0.0)
xi = project_from_slave_to_master(mel1, X1, n1, 0.0)
@test isapprox(xi, -1.0)

@test_throws Exception project_from_master_to_slave(sel1, [0.0, 0.0], 0.0; tol=0.0)
@test_throws Exception project_from_slave_to_master(mel1, [0.0, 1.0], [-1.0, 0.0], 0.0; tol=0.0)