text:
#### Columnwise LU Solver

export ManyColumnLU, SingleColumnLU

abstract type AbstractColumnLUSolver <: AbstractSystemSolver end

"""
    ManyColumnLU()

This solver is used for systems that are block diagonal where each block is
associated with a column of the mesh.  The systems are solved using a
non-pivoted LU factorization.
"""
struct ManyColumnLU <: AbstractColumnLUSolver end

"""
    SingleColumnLU()

This solver is used for systems that are block diagonal where each block is
associated with a column of the mesh.  Moreover, each block is assumed to be
the same.  The systems are solved using a non-pivoted LU factorization.
"""
struct SingleColumnLU <: AbstractColumnLUSolver end

struct ColumnwiseLU{AT}
    A::AT
end

struct DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC, AT}
    data::AT
end
DGColumnBandedMatrix(
    A::DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC},
    data,
) where {D, P, NS, EH, EV, EB, SC} =
    DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC, typeof(data)}(data)
Base.eltype(A::DGColumnBandedMatrix) = eltype(A.data)
Base.size(A::DGColumnBandedMatrix) = size(A.data)
dimensionality(::DGColumnBandedMatrix{D}) where {D} = D
polynomialorders(::DGColumnBandedMatrix{D, P}) where {D, P} = P
# polynomialorder orders P is a tuple,
# last entry is always the vertical polynomial order
vertical_polynomialorder(::DGColumnBandedMatrix{D, P}) where {D, P} = P[end]
num_state(::DGColumnBandedMatrix{D, P, NS}) where {D, P, NS} = NS
num_horz_elem(::DGColumnBandedMatrix{D, P, NS, EH}) where {D, P, NS, EH} = EH
num_vert_elem(
    ::DGColumnBandedMatrix{D, P, NS, EH, EV},
) where {D, P, NS, EH, EV} = EV
elem_band(
    ::DGColumnBandedMatrix{D, P, NS, EH, EV, EB},
) where {D, P, NS, EH, EV, EB} = EB
single_column(
    ::DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC},
) where {D, P, NS, EH, EV, EB, SC} = SC
lower_bandwidth(A::DGColumnBandedMatrix) =
    (vertical_polynomialorder(A) + 1) * num_state(A) * elem_band(A) - 1
upper_bandwidth(A::DGColumnBandedMatrix) = lower_bandwidth(A)
Base.reshape(A::DGColumnBandedMatrix, args...) =
    DGColumnBandedMatrix(A, reshape(A.data, args...))
Adapt.adapt_structure(to, A::DGColumnBandedMatrix) =
    DGColumnBandedMatrix(A, adapt(to, A.data))

Base.@propagate_inbounds function Base.getindex(A::DGColumnBandedMatrix, I...)
    return A.data[I...]
end
Base.@propagate_inbounds function Base.setindex!(
    A::DGColumnBandedMatrix,
    val,
    I...,
)
    A.data[I...] = val
end

function prefactorize(op, solver::AbstractColumnLUSolver, Q, args...)
    dg = op.f!

    # TODO: can we get away with just passing the grid?
    A = banded_matrix(
        op,
        dg,
        similar(Q),
        similar(Q),
        args...;
        single_column = typeof(solver) <: SingleColumnLU,
    )

    band_lu!(A)

    ColumnwiseLU(A)
end

function linearsolve!(
    linop,
    clu::ColumnwiseLU,
    ::AbstractColumnLUSolver,
    Q,
    Qrhs,
    args...,
)
    A = clu.A
    Q .= Qrhs

    band_forward!(Q, A)
    band_back!(Q, A)
end

"""
    band_lu!(A)

"""
function band_lu!(A)
    device = array_device(A.data)

    nstate = num_state(A)
    Nq = polynomialorders(A) .+ 1
    @inbounds Nq_h = Nq[1]
    @inbounds Nqj = dimensionality(A) == 2 ? 1 : Nq[2]
    nhorzelem = num_horz_elem(A)

    groupsize = (Nq_h, Nqj)
    ndrange = (nhorzelem * Nq_h, Nqj)

    if single_column(A)
        # single column case
        #
        # TODO Would it be faster to copy the matrix to the host and factorize it
        # there?
        groupsize = (1, 1)
        ndrange = groupsize
        A = reshape(A, 1, 1, size(A)..., 1)
    end

    event = Event(device)
    event = band_lu_kernel!(device, groupsize)(
        A,
        ndrange = ndrange,
        dependencies = (event,),
    )
    wait(device, event)
end

function band_forward!(Q, A)
    device = array_device(Q)

    Nq = polynomialorders(A) .+ 1
    @inbounds Nq_h = Nq[1]
    @inbounds Nqj = dimensionality(A) == 2 ? 1 : Nq[2]
    nhorzelem = num_horz_elem(A)

    event = Event(device)
    event = band_forward_kernel!(device, (Nq_h, Nqj))(
        Q.data,
        A,
        ndrange = (nhorzelem * Nq_h, Nqj),
        dependencies = (event,),
    )
    wait(device, event)
end

function band_back!(Q, A)
    device = array_device(Q)

    Nq = polynomialorders(A) .+ 1
    @inbounds Nq_h = Nq[1]
    @inbounds Nqj = dimensionality(A) == 2 ? 1 : Nq[2]
    nhorzelem = num_horz_elem(A)

    event = Event(device)
    event = band_back_kernel!(device, (Nq_h, Nqj))(
        Q.data,
        A,
        ndrange = (nhorzelem * Nq_h, Nqj),
        dependencies = (event,),
    )
    wait(device, event)
end

"""
    banded_matrix(
        dg::DGModel,
        Q::MPIStateArray = MPIStateArray(dg),
        dQ::MPIStateArray = MPIStateArray(dg);
        single_column = false,
    )

Forms the banded matrices for each the column operator defined by the `DGModel`
dg.  If `single_column=false` then a banded matrix is stored for each column and
if `single_column=true` only the banded matrix associated with the first column
of the first element is stored. The bandwidth of the DG column banded matrix is
`p = q = (polynomialorder + 1) * nstate * eband - 1` with `p` and `q` being
the upper and lower bandwidths.

The banded matrices are stored in the LAPACK band storage format
<https://www.netlib.org/lapack/lug/node124.html>.

The banded matrices are returned as an arrays where the array type matches that
of `Q`. If `single_column=false` then the returned array has 5 dimensions, which
are:
- first horizontal column index
- second horizontal column index
- band index (-q:p)
- vertical DOF index with state `s`, vertical DOF index `k`, and vertical
  element `ev` mapping to `s + nstate * (k - 1) + nstate * nvertelem * (ev - 1)`
- horizontal element index

If the `single_column=true` then the returned array has 2 dimensions which are
the band index and the vertical DOF index.
"""
function banded_matrix(
    dg::DGModel,
    Q::MPIStateArray = MPIStateArray(dg),
    dQ::MPIStateArray = MPIStateArray(dg);
    single_column = false,
)
    banded_matrix(
        (dQ, Q) -> dg(dQ, Q, nothing, 0; increment = false),
        dg,
        Q,
        dQ;
        single_column = single_column,
    )
end

"""
    banded_matrix(
        f!,
        dg::DGModel,
        Q::MPIStateArray = MPIStateArray(dg),
        dQ::MPIStateArray = MPIStateArray(dg),
        args...;
        single_column = false,
    )

Forms the banded matrices for each the column operator defined by the linear
operator `f!` which is assumed to have the same banded structure as the
`DGModel` dg.  If `single_column=false` then a banded matrix is stored for each
column and if `single_column=true` only the banded matrix associated with the
first column of the first element is stored. The bandwidth of the DG column
banded matrix is `p = q = (polynomialorder + 1) * nstate * eband - 1` with
`p` and `q` being the upper and lower bandwidths.

The banded matrices are stored in the LAPACK band storage format
<https://www.netlib.org/lapack/lug/node124.html>.

The banded matrices are returned as an arrays where the array type matches that
of `Q`. If `single_column=false` then the returned array has 5 dimensions, which
are:
- first horizontal column index
- second horizontal column index
- band index (-q:p)
- vertical DOF index with state `s`, vertical DOF index `k`, and vertical
  element `ev` mapping to `s + nstate * (k - 1) + nstate * nvertelem * (ev - 1)`
- horizontal element index

If the `single_column=true` then the returned array has 2 dimensions which are
the band index and the vertical DOF index.

Here `args` are passed to `f!`.
"""
function banded_matrix(
    f!,
    dg::DGModel,
    Q::MPIStateArray = MPIStateArray(dg),
    dQ::MPIStateArray = MPIStateArray(dg),
    args...;
    single_column = false,
)
    # Initialize banded matrix data structure
    A = empty_banded_matrix(dg, Q; single_column = single_column)

    # Populate matrix with data
    update_banded_matrix!(
        A,
        f!,
        dg,
        Q,
        dQ,
        args...;
        single_column = single_column,
    )

    A
end

"""
    empty_banded_matrix(
        dg::DGModel,
        Q::MPIStateArray;
        single_column = false,
    )

Initializes an empty banded matrix stored in the LAPACK band storage format
<https://www.netlib.org/lapack/lug/node124.html>.
"""
function empty_banded_matrix(
    dg::DGModel,
    Q::MPIStateArray;
    single_column = false,
)
    bl = dg.balance_law
    grid = dg.grid
    topology = grid.topology
    @assert isstacked(topology)
    @assert typeof(dg.direction) <: VerticalDirection

    FT = eltype(Q.data)
    device = array_device(Q)

    nstate = number_states(bl, Prognostic())
    N = polynomialorders(grid)
    dim = dimensionality(grid)
    Nq = N .+ 1
    @inbounds begin
        Nq_h = Nq[1]
        Nqj = dim == 2 ? 1 : Nq[2]
        Nq_v = Nq[dim]
    end

    # p is lower bandwidth
    # q is upper bandwidth
    eband = number_states(bl, GradientFlux()) == 0 ? 1 : 2
    p = q = nstate * Nq_v * eband - 1

    nrealelem = length(topology.realelems)
    nvertelem = topology.stacksize
    nhorzelem = div(nrealelem, nvertelem)

    # first horizontal DOF index
    # second horizontal DOF index
    # band index -q:p
    # vertical DOF index
    # horizontal element index
    A = if single_column
        similar(Q.data, p + q + 1, Nq_v * nstate * nvertelem)
    else
        similar(Q.data, Nq_h, Nqj, p + q + 1, Nq_v * nstate * nvertelem, nhorzelem)
    end
    fill!(A, zero(FT))

    A = DGColumnBandedMatrix{
        dim,
        N,
        nstate,
        nhorzelem,
        nvertelem,
        eband,
        single_column,
        typeof(A),
    }(
        A,
    )

    A
end

"""
    update_banded_matrix!(
        A::DGColumnBandedMatrix,
        f!,
        dg::DGModel,
        Q::MPIStateArray = MPIStateArray(dg),
        dQ::MPIStateArray = MPIStateArray(dg),
        args...;
        single_column = false,
    )

Updates the banded matrices for each the column operator defined by the linear
operator `f!` which is assumed to have the same banded structure as the
`DGModel` dg.  If `single_column=false` then a banded matrix is stored for each
column and if `single_column=true` only the banded matrix associated with the
first column of the first element is stored. The bandwidth of the DG column
banded matrix is `p = q = (polynomialorder + 1) * nstate * eband - 1` with
`p` and `q` being the upper and lower bandwidths.

Here `args` are passed to `f!`.
"""
function update_banded_matrix!(
    A::DGColumnBandedMatrix,
    f!,
    dg::DGModel,
    Q::MPIStateArray = MPIStateArray(dg),
    dQ::MPIStateArray = MPIStateArray(dg),
    args...;
    single_column = false,
)
    bl = dg.balance_law
    grid = dg.grid
    topology = grid.topology
    @assert isstacked(topology)
    @assert typeof(dg.direction) <: VerticalDirection

    FT = eltype(Q.data)
    device = array_device(Q)

    nstate = number_states(bl, Prognostic())
    N = polynomialorders(grid)
    dim = dimensionality(grid)
    Nq = N .+ 1
    @inbounds begin
        Nq_h = Nq[1]
        Nqj = dim == 2 ? 1 : Nq[2]
        Nq_v = Nq[dim]
    end

    # p is lower bandwidth
    # q is upper bandwidth
    eband = number_states(bl, GradientFlux()) == 0 ? 1 : 2

    # Matrix bandwidth is determined by the vertical discretization
    p = q = nstate * Nq[end] * eband - 1

    nrealelem = length(topology.realelems)
    nvertelem = topology.stacksize
    nhorzelem = div(nrealelem, nvertelem)

    # loop through all DOFs in a column and compute the matrix column
    # loop only the first min(nvertelem, 2eband+1) elements
    # in each element loop, updating these columns correspond
    # to elements (ev :2eband+1 : nvertelem)
    for ev in 1:min(nvertelem, 2eband + 1)
        for s in 1:nstate
            for k in 1:Nq_v
                # Set a single 1 per column and rest 0
                event = Event(device)
                event = kernel_set_banded_data!(device, (Nq_h, Nqj, Nq_v))(
                    bl,
                    Val(dim),
                    Val(nstate),
                    Val(N),
                    Val(nvertelem),
                    Q.data,
                    k,
                    s,
                    ev,
                    1:nhorzelem,
                    1:nvertelem;
                    ndrange = (nvertelem * Nq_h, nhorzelem * Nqj, Nq_v),
                    dependencies = (event,),
                )
                wait(device, event)

                # Get the matrix column
                f!(dQ, Q, args...)

                # Store the banded matrix
                event = Event(device)
                event = kernel_set_banded_matrix!(device, (Nq_h, Nqj, Nq_v))(
                    A,
                    dQ.data,
                    k,
                    s,
                    ev,
                    1:nhorzelem,
                    (-eband):eband;
                    ndrange = ((2eband + 1) * Nq_h, nhorzelem * Nqj, Nq_v),
                    dependencies = (event,),
                )
                wait(device, event)
            end
        end
    end
end

"""
    banded_matrix_vector_product!(
        A,
        dQ::MPIStateArray,
        Q::MPIStateArray
    )

Compute a matrix vector product `dQ = A * Q` where `A` is assumed to be a matrix
created using the `banded_matrix` function.

This function is primarily for testing purposes.
"""
function banded_matrix_vector_product!(A, dQ::MPIStateArray, Q::MPIStateArray)
    device = array_device(Q)

    Nq = polynomialorders(A) .+ 1
    @inbounds begin
        Nq_h = Nq[1]
        Nqj = dimensionality(A) == 2 ? 1 : Nq[2]
        Nq_v = Nq[end]
    end
    nvertelem = num_vert_elem(A)
    nhorzelem = num_horz_elem(A)

    event = Event(device)
    event = kernel_banded_matrix_vector_product!(device, (Nq_h, Nqj, Nq_v))(
        dQ.data,
        A,
        Q.data,
        1:nhorzelem,
        1:nvertelem;
        ndrange = (nvertelem * Nq_h, nhorzelem * Nqj, Nq_v),
        dependencies = (event,),
    )
    wait(device, event)
end

using StaticArrays
using KernelAbstractions.Extras: @unroll

@doc """
    band_lu_kernel!(A)

This performs Band Gaussian Elimination (Algorithm 4.3.1 of Golub and Van
Loan).  The array `A` contains a band matrix for each vertical column.  For
example, `A[i, j, :, :, h]`, is the band matrix associated with the `(i, j)`th
degree of freedom in the horizontal element `h`.

Each `n` by `n` band matrix is assumed to have upper bandwidth `q` and lower
bandwidth `p` where `n = nstate * Nq_h * nvertelem` and `p = q = nstate * Nq *
eband - 1`.

Each band matrix is stored in the [LAPACK band storage](https://www.netlib.org/lapack/lug/node124.html).
For example the band matrix

    B = [b₁₁ b₁₂ 0   0   0
         b₂₁ b₂₂ b₂₃ 0   0
         b₃₁ b₃₂ b₃₃ b₃₄ 0
         0   b₄₂ b₄₃ b₄₄ b₄₅
         0   0   b₅₃ b₅₄ b₅₅]

is stored as

    B = [0   b₁₂ b₂₃ b₃₄ b₄₅
         b₁₁ b₂₂ b₃₃ b₄₄ b₅₅
         b₂₁ b₃₂ b₄₃ b₅₄ 0
         b₃₁ b₄₂ b₅₃ 0   0]

### Reference

 - [GolubVanLoan2013](@cite)

""" band_lu_kernel!
@kernel function band_lu_kernel!(A)
    @uniform begin
        Nq = polynomialorders(A) .+ 1
        @inbounds Nq_h = Nq[1]
        @inbounds Nq_v = Nq[end]
        nstate = num_state(A)
        nvertelem = num_vert_elem(A)
        n = nstate * Nq_v * nvertelem
        p, q = lower_bandwidth(A), upper_bandwidth(A)
    end

    h = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    @inbounds begin
        for v in 1:nvertelem
            for k in 1:Nq_v
                for s in 1:nstate
                    kk = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v

                    Aq = A[i, j, q + 1, kk, h]
                    for ii in 1:p
                        A[i, j, q + ii + 1, kk, h] /= Aq
                    end

                    for jj in 1:q
                        if jj + kk ≤ n
                            Ajj = A[i, j, q - jj + 1, jj + kk, h]
                            for ii in 1:p
                                A[i, j, q + ii - jj + 1, jj + kk, h] -=
                                    A[i, j, q + ii + 1, kk, h] * Ajj
                            end
                        end
                    end
                end
            end
        end
    end
end

@doc """
    band_forward_kernel!(b, LU)

This performs Band Forward Substitution (Algorithm 4.3.2 of Golub and Van
Loan), i.e., the right-hand side `b` is replaced with the solution of `L*x=b`.

The array `b` is of the size `(Nq * Nqj * Nq, nstate, nvertelem * nhorzelem)`.

The LU-factorization array `LU` contains a single band matrix or one
for each vertical column, see [`band_lu!`](@ref).

Each `n` by `n` band matrix is assumed to have upper bandwidth `q` and lower
bandwidth `p` where `n = nstate * Nq * nvertelem` and `p = q = nstate * Nq *
eband - 1`.

### Reference

 - [GolubVanLoan2013](@cite)

""" band_forward_kernel!
@kernel function band_forward_kernel!(b, LU)
    @uniform begin
        FT = eltype(b)
        nstate = num_state(LU)
        Nq = polynomialorders(LU) .+ 1
        @inbounds begin
            Nq_h = Nq[1]
            Nqj = dimensionality(LU) == 2 ? 1 : Nq[2]
            Nq_v = Nq[end]
        end
        nvertelem = num_vert_elem(LU)
        n = nstate * Nq_v * nvertelem
        eband = elem_band(LU)
        p, q = lower_bandwidth(LU), upper_bandwidth(LU)

        l_b = MArray{Tuple{p + 1}, FT}(undef)
    end

    h = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    @inbounds begin
        @unroll for v in 1:eband
            @unroll for k in 1:Nq_v
                @unroll for s in 1:nstate
                    ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)
                    ee = v + nvertelem * (h - 1)
                    ii = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v
                    l_b[ii] = nvertelem ≥ v ? b[ijk, s, ee] : zero(FT)
                end
            end
        end

        for v in 1:nvertelem
            @unroll for k in 1:Nq_v
                @unroll for s in 1:nstate
                    jj = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v

                    @unroll for ii in 2:(p + 1)
                        Lii =
                            single_column(LU) ? LU[ii + q, jj] :
                            LU[i, j, ii + q, jj, h]
                        l_b[ii] -= Lii * l_b[1]
                    end

                    ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)
                    ee = v + nvertelem * (h - 1)

                    b[ijk, s, ee] = l_b[1]

                    @unroll for ii in 1:p
                        l_b[ii] = l_b[ii + 1]
                    end

                    if jj + p < n
                        (idx, si) = fldmod1(jj + p + 1, nstate)
                        (vi, ki) = fldmod1(idx, Nq_v)

                        ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (ki - 1)
                        ee = vi + nvertelem * (h - 1)

                        l_b[p + 1] = b[ijk, si, ee]
                    end
                end
            end
        end
    end
end

@doc """
    band_back_kernel!(b, LU)

This performs Band Back Substitution (Algorithm 4.3.3 of Golub and Van
Loan), i.e., the right-hand side `b` is replaced with the solution of `U*x=b`.

The array `b` is of the size `(Nq * Nqj * Nq, nstate, nvertelem * nhorzelem)`.

The LU-factorization array `LU` contains a single band matrix or one
for each vertical column, see [`band_lu!`](@ref).

Each `n` by `n` band matrix is assumed to have upper bandwidth `q` and lower
bandwidth `p` where `n = nstate * Nq * nvertelem` and `p = q = nstate * Nq *
eband - 1`.

### Reference

 - [GolubVanLoan2013](@cite)

""" band_back_kernel!
@kernel function band_back_kernel!(b, LU)
    @uniform begin
        FT = eltype(b)
        nstate = num_state(LU)
        Nq = polynomialorders(LU) .+ 1
        @inbounds begin
            Nq_h = Nq[1]
            Nqj = dimensionality(LU) == 2 ? 1 : Nq[2]
            Nq_v = Nq[end]
        end
        nvertelem = num_vert_elem(LU)
        n = nstate * Nq_h * nvertelem
        q = upper_bandwidth(LU)
        eband = elem_band(LU)

        l_b = MArray{Tuple{q + 1}, FT}(undef)
    end

    h = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    @inbounds begin
        @unroll for v in nvertelem:-1:(nvertelem - eband + 1)
            @unroll for k in Nq_v:-1:1
                @unroll for s in nstate:-1:1
                    vi = eband - nvertelem + v
                    ii = s + (k - 1) * nstate + (vi - 1) * nstate * Nq_v

                    ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)
                    ee = v + nvertelem * (h - 1)

                    l_b[ii] = b[ijk, s, ee]
                end
            end
        end

        for v in nvertelem:-1:1
            @unroll for k in Nq_v:-1:1
                @unroll for s in nstate:-1:1
                    jj = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v

                    l_b[q + 1] /=
                        single_column(LU) ? LU[q + 1, jj] :
                        LU[i, j, q + 1, jj, h]

                    @unroll for ii in 1:q
                        Uii =
                            single_column(LU) ? LU[ii, jj] : LU[i, j, ii, jj, h]
                        l_b[ii] -= Uii * l_b[q + 1]
                    end

                    ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)
                    ee = v + nvertelem * (h - 1)

                    b[ijk, s, ee] = l_b[q + 1]

                    @unroll for ii in q:-1:1
                        l_b[ii + 1] = l_b[ii]
                    end

                    if jj - q > 1
                        (idx, si) = fldmod1(jj - q - 1, nstate)
                        (vi, ki) = fldmod1(idx, Nq_v)

                        ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (ki - 1)
                        ee = vi + nvertelem * (h - 1)

                        l_b[1] = b[ijk, si, ee]
                    end
                end
            end
        end
    end
end

### TODO: Document this
@kernel function kernel_set_banded_data!(
    bl::BalanceLaw,
    ::Val{dim},
    ::Val{nstate},
    ::Val{N},
    ::Val{nvertelem},
    Q,
    kin,
    sin,
    evin0,
    helems,
    velems,
) where {dim, N, nstate, nvertelem}
    @uniform begin
        FT = eltype(Q)

        Nq = N .+ 1
        @inbounds begin
            Nq_h = Nq[1]
            Nqj = dim == 2 ? 1 : Nq[2]
            Nq_v = Nq[end]
        end

        eband = number_states(bl, GradientFlux()) == 0 ? 1 : 2
    end

    ev, eh = @index(Group, NTuple)
    i, j, k = @index(Local, NTuple)

    @inbounds begin
        e = ev + (eh - 1) * nvertelem
        ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)
        @unroll for s in 1:nstate
            if k == kin && s == sin && ((ev - evin0) % (2eband + 1) == 0)
                Q[ijk, s, e] = 1
            else
                Q[ijk, s, e] = 0
            end
        end
    end
end

@kernel function kernel_set_banded_matrix!(
    A,
    dQ,
    kin,
    sin,
    evin0,
    helems,
    vpelems,
)
    @uniform begin
        FT = eltype(A)
        nstate = num_state(A)
        Nq = polynomialorders(A) .+ 1
        @inbounds begin
            Nq_h = Nq[1]
            Nq_v = Nq[end]
            Nqj = dimensionality(A) == 2 ? 1 : Nq[2]
        end
        nvertelem = num_vert_elem(A)
        p = lower_bandwidth(A)
        q = upper_bandwidth(A)

        eband = elem_band(A)
        eshift = elem_band(A) + 1
    end

    ep, eh = @index(Group, NTuple)
    ep = ep - eshift
    i, j, k = @index(Local, NTuple)

    for evin in evin0:(2eband + 1):nvertelem
        # sin, kin, evin are the state, vertical dof, and vert element we are
        # handling
        # column index of matrix
        jj = sin + (kin - 1) * nstate + (evin - 1) * nstate * Nq_v

        # one thread is launch for dof that might contribute to column jj's band
        @inbounds begin
            # ep is the shift we need to add to evin to get the element we need to
            # consider
            ev = ep + evin
            if 1 ≤ ev ≤ nvertelem
                e = ev + (eh - 1) * nvertelem
                ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)
                @unroll for s in 1:nstate
                    # row index of matrix
                    ii = s + (k - 1) * nstate + (ev - 1) * nstate * Nq_v
                    # row band index
                    bb = ii - jj
                    # make sure we're in the bandwidth
                    if -q ≤ bb ≤ p
                        if !single_column(A)
                            A[i, j, bb + q + 1, jj, eh] = dQ[ijk, s, e]
                        else
                            if (i, j, eh) == (1, 1, 1)
                                A[bb + q + 1, jj] = dQ[ijk, s, e]
                            end
                        end
                    end
                end
            end
        end
    end
end

@kernel function kernel_banded_matrix_vector_product!(dQ, A, Q, helems, velems)
    @uniform begin
        FT = eltype(A)
        nstate = num_state(A)
        Nq = polynomialorders(A) .+ 1
        @inbounds begin
            Nq_h = Nq[1]
            Nqj = dimensionality(A) == 2 ? 1 : Nq[2]
            Nq_v = Nq[end]
        end
        nvertelem = num_vert_elem(A)
        p = lower_bandwidth(A)
        q = upper_bandwidth(A)

        elo = div(q, Nq_v * nstate - 1)
        eup = div(p, Nq_v * nstate - 1)
    end

    ev, eh = @index(Group, NTuple)
    i, j, k = @index(Local, NTuple)

    # matrix row loops
    @inbounds begin
        e = ev + nvertelem * (eh - 1)
        @unroll for s in 1:nstate
            Ax = -zero(FT)
            ii = s + (k - 1) * nstate + (ev - 1) * nstate * Nq_v

            # banded matrix column loops
            @unroll for evv in max(1, ev - elo):min(nvertelem, ev + eup)
                ee = evv + nvertelem * (eh - 1)
                @unroll for kk in 1:Nq_v
                    ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (kk - 1)
                    @unroll for ss in 1:nstate
                        jj = ss + (kk - 1) * nstate + (evv - 1) * nstate * Nq_v
                        bb = ii - jj
                        if -q ≤ bb ≤ p
                            if !single_column(A)
                                Ax +=
                                    A[i, j, bb + q + 1, jj, eh] * Q[ijk, ss, ee]
                            else
                                Ax += A[bb + q + 1, jj] * Q[ijk, ss, ee]
                            end
                        end
                    end
                end
            end
            ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)
            dQ[ijk, s, e] = Ax
        end
    end
end

max_forks_repo_name: jm-c/CLIMA


text:
module TestSKCrossValidator

using Test
using Random
using AutoMLPipeline

function crossval_class(ppl,X,Y,folds,verbose)
  @test crossvalidate(ppl,X,Y,"accuracy_score",folds,verbose).mean > 0.80
  @test crossvalidate(ppl,X,Y,"balanced_accuracy_score",folds,verbose).mean > 0.80
  @test crossvalidate(ppl,X,Y,"cohen_kappa_score",folds,verbose).mean > 0.80
  @test crossvalidate(ppl,X,Y,"jaccard_score","weighted",folds,verbose).mean > 0.80
  @test crossvalidate(ppl,X,Y,"matthews_corrcoef",folds,verbose).mean > 0.80
  @test crossvalidate(ppl,X,Y,"hamming_loss",folds,verbose).mean < 0.1
  @test crossvalidate(ppl,X,Y,"zero_one_loss",folds,verbose).mean < 0.1
  @test crossvalidate(ppl,X,Y,"f1_score","weighted",folds,verbose).mean > 0.80
  @test crossvalidate(ppl,X,Y,"precision_score","weighted",folds,verbose).mean > 0.80
  @test crossvalidate(ppl,X,Y,"recall_score","weighted",folds,verbose).mean > 0.80
end

function crossval_reg(ppl,X,Y,folds,verbose)
  @test crossvalidate(ppl,X,Y,"mean_squared_error",folds,verbose).mean < 0.5
  @test crossvalidate(ppl,X,Y,"mean_squared_log_error",folds,verbose).mean < 0.5
  @test crossvalidate(ppl,X,Y,"mean_absolute_error",folds,verbose).mean < 0.5
  @test crossvalidate(ppl,X,Y,"median_absolute_error",folds,verbose).mean < 0.5
  @test crossvalidate(ppl,X,Y,"r2_score",folds,verbose).mean > 0.50
  @test crossvalidate(ppl,X,Y,"max_error",folds,verbose).mean < 0.7
  @test crossvalidate(ppl,X,Y,"mean_poisson_deviance",folds,verbose).mean < 0.7
  @test crossvalidate(ppl,X,Y,"mean_gamma_deviance",folds,verbose).mean < 0.7
  @test crossvalidate(ppl,X,Y,"mean_tweedie_deviance",folds,verbose).mean < 0.7
  @test crossvalidate(ppl,X,Y,"explained_variance_score",folds,verbose).mean > 0.50
end

function test_skcross_reg()
  data=getiris()
  X=data[:,1:3]
  Y=data[:,4]
  ppl1 = Pipeline(Dict(:machines=>[RandomForest()]))
  crossval_reg(ppl1,X,Y,10,false)
  ppl2 = Pipeline(Dict(:machines=>[VoteEnsemble()]))
  crossval_reg(ppl2,X,Y,10,false)
  cat = CatFeatureSelector()
  num = NumFeatureSelector()
  pca = SKPreprocessor("PCA")
  ptf = SKPreprocessor("PowerTransformer")
  rbc = SKPreprocessor("RobustScaler")
  ppl3=@pipeline ((cat + num) + (num |> ptf) + (num |> rbc) + (num |> pca)) |> VoteEnsemble()
  crossval_reg(ppl3,X,Y,10,false)
end
@testset "CrossValidator Regression" begin
  Random.seed!(123)
  test_skcross_reg()
end

function test_skcross_class()
  data=getiris()
  X=data[:,1:4]
  Y=data[:,5] |> Vector{String}
  ppl1 = Pipeline(Dict(:machines=>[RandomForest()]))
  crossval_class(ppl1,X,Y,10,false)
  ppl2 = Pipeline(Dict(:machines=>[VoteEnsemble()]))
  crossval_class(ppl2,X,Y,10,false)
  cat = CatFeatureSelector()
  num = NumFeatureSelector()
  pca = SKPreprocessor("PCA")
  ptf = SKPreprocessor("PowerTransformer")
  rbc = SKPreprocessor("RobustScaler")
  ppl3=@pipeline ((cat + num) + (num |> ptf) + (num |> rbc) + (num |> pca)) |> VoteEnsemble()
  crossval_class(ppl3,X,Y,10,false)
end
@testset "CrossValidator Classification" begin
  Random.seed!(123)
  test_skcross_class()
end

function test_crossval_options()
  data=getiris()
  X=data[:,1:4]
  Y=data[:,5] |> collect
  acc(x,y)=score(:accuracy,x,y)
  ppl1 = Pipeline(RandomForest())
  @test crossvalidate(ppl1,X,Y,"accuracy_score",10,false).mean > 0.90
  @test crossvalidate(ppl1,X,Y,"accuracy_score").mean > 0.90
  @test crossvalidate(ppl1,X,Y,"accuracy_score",10).mean > 0.90
  @test crossvalidate(ppl1,X,Y,"accuracy_score",false).mean > 0.90
  @test crossvalidate(ppl1,X,Y,"accuracy_score",verbose=false).mean > 0.90
  @test crossvalidate(ppl1,X,Y,metric=acc,verbose=false).mean > 0.90
  @test crossvalidate(ppl1,X,Y,metric=acc,nfolds=5).mean > 0.90
  @test crossvalidate(ppl1,X,Y,acc,5,true).mean > 0.90
end
@testset "CrossValidator Argument Options" begin
  Random.seed!(123)
  test_crossval_options()
end

end

ext: jl


text:
export Kernel
module Kernel

using CUDA
using GemmKernels
using GemmKernels.Tiling
using GPUifyLoops: @unroll
using StaticArrays

function matmul_singlestage(a, b, c, d,
                          transf_gl2sh_a, transf_gl2sh_b, transf_gl2sh_c, transf_sh2gl_d,
                          transf_sh2rf_a, transf_sh2rf_b, transf_sh2rf_c, transf_rf2sh_d,
                          epilogue,
                          conf::Type{GemmKernels.Config{MATMUL_SHAPE, BLOCK_SHAPE, WARPS_PER_BLOCK, MEM_A_WARP, MEM_A_THREAD, MEM_B_WARP, MEM_B_THREAD, MEM_CD_WARP, MEM_CD_THREAD, COMPUTE_WARP, COMPUTE_OP_SHAPE, GLOBAL_A_LAYOUT, GLOBAL_B_LAYOUT, GLOBAL_C_LAYOUT, GLOBAL_D_LAYOUT, SHARED_A_LAYOUT, SHARED_B_LAYOUT, SHARED_C_LAYOUT, SHARED_D_LAYOUT, OPERATOR, IS_A_COL_MAJOR, IS_B_COL_MAJOR}}) where {MATMUL_SHAPE, BLOCK_SHAPE, WARPS_PER_BLOCK, MEM_A_WARP, MEM_A_THREAD, MEM_B_WARP, MEM_B_THREAD, MEM_CD_WARP, MEM_CD_THREAD, COMPUTE_WARP, COMPUTE_OP_SHAPE, GLOBAL_A_LAYOUT, GLOBAL_B_LAYOUT, GLOBAL_C_LAYOUT, GLOBAL_D_LAYOUT, SHARED_A_LAYOUT, SHARED_B_LAYOUT, SHARED_C_LAYOUT, SHARED_D_LAYOUT, OPERATOR, IS_A_COL_MAJOR, IS_B_COL_MAJOR}
    # Calculate the number of fragments needed to fully cover a warp tile
    NUM_FRAGMENTS_M = COMPUTE_WARP.M ÷ COMPUTE_OP_SHAPE.M
    NUM_FRAGMENTS_N = COMPUTE_WARP.N ÷ COMPUTE_OP_SHAPE.N

    # Constants
    block_i = (blockIdx().x - 1) * BLOCK_SHAPE.M
    block_j = (blockIdx().y - 1) * BLOCK_SHAPE.N

    warpId = (threadIdx().x - 1) ÷ 32 + 1
    laneId = (threadIdx().x - 1) % 32 + 1

    gemm_sz = Tile(MATMUL_SHAPE)
    block_tile = Tile(BLOCK_SHAPE)

    # (1) Cooperatively load a BLOCK_SHAPE.M x BLOCK_SHAPE.N tile of C from global to shared memory within one threadblock
    shmem_c = @cuDynamicSharedMem(Layout.eltype(SHARED_C_LAYOUT), Layout.physical_size(SHARED_C_LAYOUT, block_tile.MN.size))

    @unroll for warp_tile = parallellise(block_tile.MN, Tile(MEM_CD_WARP), warpId, WARPS_PER_BLOCK)
        @unroll for thread_tile = parallellise(warp_tile, Tile(MEM_CD_THREAD), laneId, 32)
            x = Layout.load(GLOBAL_C_LAYOUT, c, translate_base(thread_tile, (M = block_i, N = block_j)))
            x = transf_gl2sh_c(x, thread_tile)
            Layout.store!(SHARED_C_LAYOUT, shmem_c, x, thread_tile)
        end
    end

    sync_threads()

    # (2) Load a COMPUTE_WARP.M x COMPUTE_WARP.N tile of C from shared memory into registers
    warp_tile = subdivide(block_tile.MN, Tile(COMPUTE_WARP).MN, warpId, WARPS_PER_BLOCK)

    c_frags = MArray{Tuple{NUM_FRAGMENTS_M, NUM_FRAGMENTS_N}, Operator.fragtype_accum(OPERATOR, SHARED_C_LAYOUT)}(undef)

    @unroll for i = 1 : NUM_FRAGMENTS_M
        @unroll for j = 1 : NUM_FRAGMENTS_N
            tile = translate_offset(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
            @inbounds c_frags[i, j] = transf_sh2rf_c(Operator.load_c(OPERATOR, SHARED_C_LAYOUT, shmem_c, tile), tile)
        end
    end

    sync_threads()

    # (3) Compute a BLOCK_SHAPE.M x BLOCK_SHAPE.N x BLOCK_SHAPE.K matrix product within one threadblock
    shmem_a = @cuDynamicSharedMem(Layout.eltype(SHARED_A_LAYOUT), Layout.physical_size(SHARED_A_LAYOUT, block_tile.MK.size))
    shmem_b = @cuDynamicSharedMem(Layout.eltype(SHARED_B_LAYOUT), Layout.physical_size(SHARED_B_LAYOUT, block_tile.KN.size),
                                    length(shmem_a) * sizeof(Layout.eltype(SHARED_A_LAYOUT)))

    @unroll for block_k = 0 : block_tile.size.K : gemm_sz.size.K - 1
        if Layout.threadblock_condition(GLOBAL_A_LAYOUT, GLOBAL_B_LAYOUT, block_i, block_j, block_k, block_tile)
            # (3.1) Cooperatively load a BLOCK_SHAPE.M x BLOCK_SHAPE.K tile of A from global to shared memory within one threadblock
            @unroll for warp_tile = parallellise(block_tile.MK, Tile(MEM_A_WARP), warpId, WARPS_PER_BLOCK, IS_A_COL_MAJOR)
                @unroll for thread_tile = parallellise(warp_tile, Tile(MEM_A_THREAD), laneId, 32, IS_A_COL_MAJOR)
                    x = Layout.load(GLOBAL_A_LAYOUT, a, translate_base(thread_tile, (M = block_i, K = block_k)))
                    x = transf_gl2sh_a(x, thread_tile)
                    Layout.store!(SHARED_A_LAYOUT, shmem_a, x, thread_tile)
                end
            end

            # (3.2) Cooperatively load a BLOCK_SHAPE.K x BLOCK_SHAPE.N tile of B from global to shared memory within one threadblock
            @unroll for warp_tile = parallellise(block_tile.KN, Tile(MEM_B_WARP), warpId, WARPS_PER_BLOCK, IS_B_COL_MAJOR)
                @unroll for thread_tile = parallellise(warp_tile, Tile(MEM_B_THREAD), laneId, 32, IS_B_COL_MAJOR)
                    x = Layout.load(GLOBAL_B_LAYOUT, b, translate_base(thread_tile, (K = block_k, N = block_j)))
                    x = transf_gl2sh_b(x, thread_tile)
                    Layout.store!(SHARED_B_LAYOUT, shmem_b, x, thread_tile)
                end
            end

            sync_threads()

            # (3.3) Calculate a COMPUTE_WARP.M x COMPUTE_WARP.N tile of D, using a COMPUTE_WARP.M x COMPUTE_WARP.N x COMPUTE_WARP.K operation
            @unroll for warp_tile = parallellise(block_tile, Tile(COMPUTE_WARP), warpId, WARPS_PER_BLOCK)
                # (3.3.1) Load a COMPUTE_WARP.M x COMPUTE_WARP.K tile of A from shared memory into registers
                a_frags = MArray{Tuple{NUM_FRAGMENTS_M}, Operator.fragtype_a(OPERATOR, SHARED_A_LAYOUT)}(undef)

                @unroll for i = 1 : NUM_FRAGMENTS_M
                    a_tile = translate_offset(warp_tile.MK, (M = (i-1)*COMPUTE_OP_SHAPE.M, K = 0))
                    @inbounds a_frags[i] = transf_sh2rf_a(Operator.load_a(OPERATOR, SHARED_A_LAYOUT, shmem_a, a_tile), a_tile)
                end

                # (3.3.2) Load a COMPUTE_WARP.K x COMPUTE_WARP.N tile of B from shared memory into registers
                b_frags = MArray{Tuple{NUM_FRAGMENTS_N}, Operator.fragtype_b(OPERATOR, SHARED_B_LAYOUT)}(undef)

                @unroll for j = 1 : NUM_FRAGMENTS_N
                    b_tile = translate_offset(warp_tile.KN, (K = 0, N = (j-1)*COMPUTE_OP_SHAPE.N))
                    @inbounds b_frags[j] = transf_sh2rf_b(Operator.load_b(OPERATOR, SHARED_B_LAYOUT, shmem_b, b_tile), b_tile)
                end

                # (3.3.3) Compute a COMPUTE_WARP.M x COMPUTE_WARP.N x COMPUTE_WARP.K matrix product within one warp
                @unroll for i = 1 : NUM_FRAGMENTS_M
                    @unroll for j = 1 : NUM_FRAGMENTS_N
                        @inbounds c_frags[i, j] = Operator.mma(OPERATOR, a_frags[i], b_frags[j], c_frags[i, j])
                    end
                end
            end
        end

        sync_threads()
    end

    # (4) Store the COMPUTE_WARP.M x COMPUTE_WARP.N tile of D from registers to shared memory
    shmem_d = @cuDynamicSharedMem(Layout.eltype(SHARED_D_LAYOUT), Layout.physical_size(SHARED_D_LAYOUT, block_tile.MN.size))

    warp_tile = subdivide(block_tile.MN, Tile(COMPUTE_WARP).MN, warpId, WARPS_PER_BLOCK)

    @unroll for i = 1 : NUM_FRAGMENTS_M
        @unroll for j = 1 : NUM_FRAGMENTS_N
            tile = translate_offset(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
            Operator.store_d(OPERATOR, SHARED_D_LAYOUT, shmem_d, transf_rf2sh_d(c_frags[i, j], tile), tile)
        end
    end

    sync_threads()

    # (5) Run the epilogue
    epilogue(d, shmem_d, transf_sh2gl_d, conf)

    return
end

function matmul_pipelined(a, b, c, d,
                          transf_gl2sh_a, transf_gl2sh_b, transf_gl2sh_c, transf_sh2gl_d,
                          transf_sh2rf_a, transf_sh2rf_b, transf_sh2rf_c, transf_rf2sh_d,
                          epilogue,
                          conf::Type{GemmKernels.Config{MATMUL_SHAPE, BLOCK_SHAPE, WARPS_PER_BLOCK, MEM_A_WARP, MEM_A_THREAD, MEM_B_WARP, MEM_B_THREAD, MEM_CD_WARP, MEM_CD_THREAD, COMPUTE_WARP, COMPUTE_OP_SHAPE, GLOBAL_A_LAYOUT, GLOBAL_B_LAYOUT, GLOBAL_C_LAYOUT, GLOBAL_D_LAYOUT, SHARED_A_LAYOUT, SHARED_B_LAYOUT, SHARED_C_LAYOUT, SHARED_D_LAYOUT, OPERATOR, IS_A_COL_MAJOR, IS_B_COL_MAJOR}}) where {MATMUL_SHAPE, BLOCK_SHAPE, WARPS_PER_BLOCK, MEM_A_WARP, MEM_A_THREAD, MEM_B_WARP, MEM_B_THREAD, MEM_CD_WARP, MEM_CD_THREAD, COMPUTE_WARP, COMPUTE_OP_SHAPE, GLOBAL_A_LAYOUT, GLOBAL_B_LAYOUT, GLOBAL_C_LAYOUT, GLOBAL_D_LAYOUT, SHARED_A_LAYOUT, SHARED_B_LAYOUT, SHARED_C_LAYOUT, SHARED_D_LAYOUT, OPERATOR, IS_A_COL_MAJOR, IS_B_COL_MAJOR}
    # Calculate the number of fragments needed to fully cover a warp tile
    NUM_FRAGMENTS_M = COMPUTE_WARP.M ÷ COMPUTE_OP_SHAPE.M
    NUM_FRAGMENTS_N = COMPUTE_WARP.N ÷ COMPUTE_OP_SHAPE.N

    # Constants
    block_i = (blockIdx().x - 1) * BLOCK_SHAPE.M
    block_j = (blockIdx().y - 1) * BLOCK_SHAPE.N

    warpId = (threadIdx().x - 1) ÷ 32 + 1
    laneId = (threadIdx().x - 1) % 32 + 1

    gemm_sz = Tile(MATMUL_SHAPE)
    block_tile = Tile(BLOCK_SHAPE)

    # (1) Cooperatively load a BLOCK_SHAPE.M x BLOCK_SHAPE.N tile of C from global to shared memory within one threadblock
    shmem_c = @cuDynamicSharedMem(Layout.eltype(SHARED_C_LAYOUT), Layout.physical_size(SHARED_C_LAYOUT, block_tile.MN.size))

    @unroll for warp_tile = parallellise(block_tile.MN, Tile(MEM_CD_WARP), warpId, WARPS_PER_BLOCK)
        @unroll for thread_tile = parallellise(warp_tile, Tile(MEM_CD_THREAD), laneId, 32)
            x = Layout.load(GLOBAL_C_LAYOUT, c, translate_base(thread_tile, (M = block_i, N = block_j)))
            x = transf_gl2sh_c(x, thread_tile)
            Layout.store!(SHARED_C_LAYOUT, shmem_c, x, thread_tile)
        end
    end

    sync_threads()

    # (2) Load a COMPUTE_WARP.M x COMPUTE_WARP.N tile of C from shared memory into registers
    warp_tile = subdivide(block_tile.MN, Tile(COMPUTE_WARP).MN, warpId, WARPS_PER_BLOCK)

    c_frags = MArray{Tuple{NUM_FRAGMENTS_M, NUM_FRAGMENTS_N}, Operator.fragtype_accum(OPERATOR, SHARED_C_LAYOUT)}(undef)

    @unroll for i = 1 : NUM_FRAGMENTS_M
        @unroll for j = 1 : NUM_FRAGMENTS_N
            tile = translate_offset(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
            @inbounds c_frags[i, j] = transf_sh2rf_c(Operator.load_c(OPERATOR, SHARED_C_LAYOUT, shmem_c, tile), tile)
        end
    end

    sync_threads()

    # (3) Compute a BLOCK_SHAPE.M x BLOCK_SHAPE.N x BLOCK_SHAPE.K matrix product within one threadblock
    shmem_a = @cuDynamicSharedMem(Layout.eltype(SHARED_A_LAYOUT), Layout.physical_size(SHARED_A_LAYOUT, block_tile.MK.size))
    shmem_b = @cuDynamicSharedMem(Layout.eltype(SHARED_B_LAYOUT), Layout.physical_size(SHARED_B_LAYOUT, block_tile.KN.size),
                                    length(shmem_a) * sizeof(Layout.eltype(SHARED_A_LAYOUT)))

    # Sizes of a_fragment and b_fragment
    a_frag_i = (block_tile.size.M * block_tile.size.K) ÷ (MEM_A_WARP.M * MEM_A_WARP.K * WARPS_PER_BLOCK)
    a_frag_j = (MEM_A_WARP.M * MEM_A_WARP.K) ÷ (MEM_A_THREAD.M * MEM_A_THREAD.K * 32)
    b_frag_i = (block_tile.size.K * block_tile.size.N) ÷ (MEM_B_WARP.K * MEM_B_WARP.N * WARPS_PER_BLOCK)
    b_frag_j = (MEM_B_WARP.K * MEM_B_WARP.N) ÷ (MEM_B_THREAD.K * MEM_B_THREAD.N * 32)

    a_fragment = MArray{Tuple{a_frag_i, a_frag_j}, Layout.fragtype(GLOBAL_A_LAYOUT, MEM_A_THREAD)}(undef)
    b_fragment = MArray{Tuple{b_frag_i, b_frag_j}, Layout.fragtype(GLOBAL_B_LAYOUT, MEM_B_THREAD)}(undef)

    a_frags = MArray{Tuple{2, NUM_FRAGMENTS_M}, Operator.fragtype_a(OPERATOR, SHARED_A_LAYOUT)}(undef)
    b_frags = MArray{Tuple{2, NUM_FRAGMENTS_N}, Operator.fragtype_b(OPERATOR, SHARED_B_LAYOUT)}(undef)

    warp_tile_mn = subdivide(block_tile, Tile(COMPUTE_WARP), warpId, WARPS_PER_BLOCK)

    # ld.global(0 : BLOCK_SHAPE.K)
    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.MK, Tile(MEM_A_WARP), warpId, WARPS_PER_BLOCK, IS_A_COL_MAJOR))
        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_A_THREAD), laneId, 32, IS_A_COL_MAJOR))
            @inbounds a_fragment[i, j] = Layout.load(GLOBAL_A_LAYOUT, a, translate_base(thread_tile, (M = block_i, K = 0)))
        end
    end

    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.KN, Tile(MEM_B_WARP), warpId, WARPS_PER_BLOCK, IS_B_COL_MAJOR))
        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_B_THREAD), laneId, 32, IS_B_COL_MAJOR))
            @inbounds b_fragment[i, j] = Layout.load(GLOBAL_B_LAYOUT, b, translate_base(thread_tile, (K = 0, N = block_j)))
        end
    end

    # st.shared()
    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.MK, Tile(MEM_A_WARP), warpId, WARPS_PER_BLOCK, IS_A_COL_MAJOR))
        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_A_THREAD), laneId, 32, IS_A_COL_MAJOR))
            @inbounds x = transf_gl2sh_a(a_fragment[i, j], thread_tile)
            Layout.store!(SHARED_A_LAYOUT, shmem_a, x, thread_tile)
        end
    end

    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.KN, Tile(MEM_B_WARP), warpId, WARPS_PER_BLOCK, IS_B_COL_MAJOR))
        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_B_THREAD), laneId, 32, IS_B_COL_MAJOR))
            @inbounds x = transf_gl2sh_b(b_fragment[i, j], thread_tile)
            Layout.store!(SHARED_B_LAYOUT, shmem_b, x, thread_tile)
        end
    end

    sync_threads()

    # ld.shared(0 : COMPUTE_OP_SHAPE.K, stage = 1)
    warp_tile = translate_offset(warp_tile_mn, (M = 0, N = 0, K = 0))

    @unroll for i = 1 : NUM_FRAGMENTS_M
        a_tile = translate_offset(warp_tile.MK, (M = (i-1)*COMPUTE_OP_SHAPE.M, K = 0))
        @inbounds a_frags[1, i] = transf_sh2rf_a(Operator.load_a(OPERATOR, SHARED_A_LAYOUT, shmem_a, a_tile), a_tile)
    end

    @unroll for j = 1 : NUM_FRAGMENTS_N
        b_tile = translate_offset(warp_tile.KN, (K = 0, N = (j-1)*COMPUTE_OP_SHAPE.N))
        @inbounds b_frags[1, j] = transf_sh2rf_b(Operator.load_b(OPERATOR, SHARED_B_LAYOUT, shmem_b, b_tile), b_tile)
    end

    # ld.global(BLOCK_SHAPE.K : 2 * BLOCK_SHAPE.K)
    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.MK, Tile(MEM_A_WARP), warpId, WARPS_PER_BLOCK, IS_A_COL_MAJOR))
        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_A_THREAD), laneId, 32, IS_A_COL_MAJOR))
            @inbounds a_fragment[i, j] = Layout.load(GLOBAL_A_LAYOUT, a, translate_base(thread_tile, (M = block_i, K = block_tile.size.K)))
        end
    end

    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.KN, Tile(MEM_B_WARP), warpId, WARPS_PER_BLOCK, IS_B_COL_MAJOR))
        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_B_THREAD), laneId, 32, IS_B_COL_MAJOR))
            @inbounds b_fragment[i, j] = Layout.load(GLOBAL_B_LAYOUT, b, translate_base(thread_tile, (K = block_tile.size.K, N = block_j)))
        end
    end

    @unroll 1 for block_k = 0 : block_tile.size.K : gemm_sz.size.K - 1
        @unroll for (i, warp_k) = enumerate(0 : COMPUTE_OP_SHAPE.K : block_tile.size.K - 1)
            cur_stage = mod1(i, 2)
            nxt_stage = mod1(i + 1, 2)

            if i == block_tile.size.K ÷ COMPUTE_OP_SHAPE.K # last iteration
                sync_threads()

                # st.shared()
                @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.MK, Tile(MEM_A_WARP), warpId, WARPS_PER_BLOCK, IS_A_COL_MAJOR))
                    @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_A_THREAD), laneId, 32, IS_A_COL_MAJOR))
                        @inbounds x = transf_gl2sh_a(a_fragment[i, j], thread_tile)
                        Layout.store!(SHARED_A_LAYOUT, shmem_a, x, thread_tile)
                    end
                end

                @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.KN, Tile(MEM_B_WARP), warpId, WARPS_PER_BLOCK, IS_B_COL_MAJOR))
                    @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_B_THREAD), laneId, 32, IS_B_COL_MAJOR))
                        @inbounds x = transf_gl2sh_b(b_fragment[i, j], thread_tile)
                        Layout.store!(SHARED_B_LAYOUT, shmem_b, x, thread_tile)
                    end
                end

                sync_threads()

                # avoid out of bounds access for global memory
                if block_k < (gemm_sz.size.K - 2 * block_tile.size.K)
                    # ld.global(block_k + 2 * BLOCK_SHAPE.K : block_k + 3 * BLOCK_SHAPE.K)
                    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.MK, Tile(MEM_A_WARP), warpId, WARPS_PER_BLOCK, IS_A_COL_MAJOR))
                        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_A_THREAD), laneId, 32, IS_A_COL_MAJOR))
                            @inbounds a_fragment[i, j] = Layout.load(GLOBAL_A_LAYOUT, a, translate_base(thread_tile, (M = block_i, K = block_k + 2 * block_tile.size.K)))
                        end
                    end

                    @unroll for (i, warp_tile) = enumerate(parallellise(block_tile.KN, Tile(MEM_B_WARP), warpId, WARPS_PER_BLOCK, IS_B_COL_MAJOR))
                        @unroll for (j, thread_tile) = enumerate(parallellise(warp_tile, Tile(MEM_B_THREAD), laneId, 32, IS_B_COL_MAJOR))
                            @inbounds b_fragment[i, j] = Layout.load(GLOBAL_B_LAYOUT, b, translate_base(thread_tile, (K = block_k + 2 * block_tile.size.K, N = block_j)))
                        end
                    end
                end
            end

            # ld.shared((warp_k + COMPUTE_OP_SHAPE.K) % BLOCK_SHAPE.K, stage = nxt_stage)
            warp_tile = translate_offset(warp_tile_mn, (M = 0, N = 0, K = (warp_k + COMPUTE_OP_SHAPE.K) % block_tile.size.K))

            @unroll for i = 1 : NUM_FRAGMENTS_M
                a_tile = translate_offset(warp_tile.MK, (M = (i-1)*COMPUTE_OP_SHAPE.M, K = 0))
                @inbounds a_frags[nxt_stage, i] = transf_sh2rf_a(Operator.load_a(OPERATOR, SHARED_A_LAYOUT, shmem_a, a_tile), a_tile)
            end

            @unroll for j = 1 : NUM_FRAGMENTS_N
                b_tile = translate_offset(warp_tile.KN, (K = 0, N = (j-1)*COMPUTE_OP_SHAPE.N))
                @inbounds b_frags[nxt_stage, j] = transf_sh2rf_b(Operator.load_b(OPERATOR, SHARED_B_LAYOUT, shmem_b, b_tile), b_tile)
            end

            # mma(cur_stage)
            @unroll for i = 1 : NUM_FRAGMENTS_M
                @unroll for j = 1 : NUM_FRAGMENTS_N
                    @inbounds c_frags[i, j] = Operator.mma(OPERATOR, a_frags[cur_stage, i], b_frags[cur_stage, j], c_frags[i, j])
                end
            end
        end

        sync_threads()
    end

    # (4) Store the COMPUTE_WARP.M x COMPUTE_WARP.N tile of D from registers to shared memory
    shmem_d = @cuDynamicSharedMem(Layout.eltype(SHARED_D_LAYOUT), Layout.physical_size(SHARED_D_LAYOUT, block_tile.MN.size))

    warp_tile = subdivide(block_tile.MN, Tile(COMPUTE_WARP).MN, warpId, WARPS_PER_BLOCK)

    @unroll for i = 1 : NUM_FRAGMENTS_M
        @unroll for j = 1 : NUM_FRAGMENTS_N
            tile = translate_offset(warp_tile, (M = (i-1)*COMPUTE_OP_SHAPE.M, N = (j-1)*COMPUTE_OP_SHAPE.N))
            Operator.store_d(OPERATOR, SHARED_D_LAYOUT, shmem_d, transf_rf2sh_d(c_frags[i, j], tile), tile)
        end
    end

    sync_threads()

    # (5) Run the epilogue
    epilogue(d, shmem_d, transf_sh2gl_d, conf)

    return
end

end


text:
using LinearAlgebraicRepresentation
LAR = LinearAlgebraicRepresentation
using ViewerGL
GL = ViewerGL

function stresstest2d()
	function input_collection(data::Array)::LAR.LAR
		assembly = LAR.Struct(data)
		return LAR.struct2lar(assembly)
	end

	V,(_,EV,FV) = LAR.cuboidGrid([4,1],true);
	W,(_,EW,FW) = LAR.cuboidGrid([1,5],true);
	mycircle(r,n) = LAR.circle(r)([n])
	data2d1 = (V,EV)
	data2d2 = LAR.Struct([ LAR.t(2,2), LAR.r(pi/3), LAR.t(-1.5,-2.5), (W,EW) ])
	data2d3 = LAR.Struct([ LAR.t(2,2), mycircle(2.5,160) ])
	data2d4 = LAR.Struct([ LAR.t(3.5,3.5), mycircle(.25,160) ])
	data2d5a = LAR.Struct([ LAR.t(5,3.5), mycircle(.75,160) ])
	data2d5 = LAR.Struct([ LAR.t(5,3.5), mycircle(.5,160) ])
	data2d6 = LAR.Struct([ LAR.t(5,3.5), mycircle(.25,160) ])
	data = [ data2d1, data2d2, data2d3, data2d4, data2d5,  data2d5a, data2d6 ]
	model2d = input_collection( [ LAR.Struct(data), LAR.t(-pi/6,pi/3), LAR.Struct(data) ] )
	V,EV = model2d
	GL.VIEW([ GL.GLLines(V,EV, GL.COLORS[1]) ]);

	model = Lar.Model(V);
	Lar.addModelCells!(model, 1, EV)

	model = Lar.Arrangement.planar_arrangement(model)

	EVs = Lar.FV2EVs(model)
	FVs = Lar.triangulate2D(model)
	GL.VIEW(GL.GLExplode(model.G, FVs, 1.2,1.2,1.2,99,1));
	GL.VIEW(GL.GLExplode(model.G, EVs, 1.2,1.2,1.2,99,1));
	GL.VIEW(GL.GLExplode(model.G, EVs, 1.2,1.2,1.2, 1,1));
end


text:
randu((a, b)) = a + rand() * (b-a)
randw(window) = (randu(window.x), randu(window.y))

function area(window::NamedTuple{(:x,:y)})
    -(window.x...) * -(window.y...)
end

function edgetrans(window::NamedTuple{(:x,:y)}, d)
    d1, d2 = abs.(d)
    x1, x2 = window.x
    y1, y2 = window.y
    (x2-x1)*(y2-y1) / ((x2-x1-d1)*(y2-y1-d2))
end

function bdist(window::NamedTuple{(:x,:y)}, xy)
    x, y = getx(xy), gety(xy)
    (x1,x2), (y1, y2) = window.x, window.y
    min(x-x1, y-y1, x2-x, y2-y)
end

""" inside(point, window)

Is point inside or on the boundary of window.
"""
function inside(T::Type{<:Union{Bool, Int}}, point, window::NamedTuple{(:x,:y)})
    x1, x2 = window.x
    y1, y2 = window.y
    x, y = getx(point), gety(point)
    if T==Bool
        x1 <= x <= x2 && y1 <= y <= y2
    else
        x1 < x < x2 && y1 < y < y2 && return 1
        x1 <= x <= x2 && y1 <= y <= y2 && return 0
        -1
    end
end
inside(point, window) = inside(Bool, point, window)
inside(window) = Base.Fix2(inside, window)


text:
"""
      predictEnsemble( predictions::Matrix{Float64}, weights::Vector{Float64} )

      perform bayesian ensemble prediction.
      #Arguments
      - `predictions::Matrix{Float64}`: each column is the prediction of one hypothesis.
      - `weights::Vector{Float64}`:     mixing weights.
      #Return
      - `Vector{Float64}`:              prediction y.
  """
  function predictEnsemble( predictions::Matrix{Float64}, weights::Vector{Float64} )
    return predictions * weights;
  end

  """
  predictEnsemble( predictions::Vector{Matrix{Float64}}, weights::Vector{Float64} )

  perform bayesian ensemble prediction.
  #Arguments
  - `predictions::Vector{Matrix{Float64}}`: each matrix is the prediction of one hypothesis.
  - `weights::Vector{Float64}`:             mixing weights.
  #Return
  - `Vector{Float64}`:                      prediction y.
  """
  function predictEnsemble( predictions::Vector{Matrix{Float64}}, weights::Vector{Float64} )
    len = size( weights, 1 );
    res = zeros( Float64, size( predictions[1], 1 ), size( predictions[1], 2 ) );
    for i in 1:1:len
      @inbounds res .+= weights[i] .* predictions[i];
    end
    return res
  end


text:
using LinearAlgebra
using Distributions

function mvnormal_loss(mu, s, x)
    d = MvNormal(mu, s)
    return logpdf(d, x)
end

@testset "distributions" begin
    # diagonal covariance matrix
    # values for comparison obtained from PyTorch
    mu = zeros(2)
    s = [1.0 0.0; 0.0 1.0]
    x = [-0.5, 0.5]
    _, g = grad(mvnormal_loss, mu, s, x)
    @test g[1] == [-0.5, 0.5]
    @test g[2] == [-0.375 -0.125; -0.125 -0.375]
    @test g[3] == [0.5, -0.5]
    # full covariance matrix
    mu = zeros(2)
    s = [1.0 0.5; 0.5 1.0]
    x = [-0.5, 0.5]
    _, g = grad(mvnormal_loss, mu, s, x)
    @test isapprox(g[1], [-1, 1])
    @test isapprox(g[2], [-0.166667 -0.166667; -0.166667 -0.166667], rtol=1e-5)
    @test isapprox(g[3], [1, -1])
end


text:
import Pkg
Pkg.instantiate() # downloads all dependencies for the current project

Pkg.add("RDatasets")

using DataFrames, CSV, RDatasets, Statistics

df = CSV.read("../externals/Core.Math.Data/data/Pejcic_318.csv", copycols = true)

println(Statistics.cor(df.ATV, df.ATT))

# https://docs.julialang.org/en/v1/stdlib/Statistics/ #


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "Telescope"
version = v"0.2.0"

# Collection of sources required to build this package
sources = [
    GitSource("https://github.com/jhigginbotham64/Telescope.git",
              "0b88f367a1fe7905ca87b440b9d58ae1b0a43e8b"),
]

# Bash recipe for building across all platforms
script = raw"""
export CFLAGS="-I${includedir}"
export CXXFLAGS="-I${includedir}"
cd $WORKSPACE/srcdir/Telescope*
git submodule update --init --recursive
mkdir build && cd build
cmake -DCMAKE_INSTALL_PREFIX=${prefix} \
    -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
    -DCMAKE_BUILD_TYPE=Release \
    -S .. \
    -B .
make -j${nproc}
make install
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = supported_platforms()
filter!(p -> (arch(p) != "armv6l") && !(Sys.isapple(p) && arch(p) == "aarch64"), platforms)
platforms = expand_cxxstring_abis(platforms)

# The products that we will ensure are always built
products = [
    FileProduct("include/telescope.h", :telescope_h),
    LibraryProduct("libtelescope", :libtelescope),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    BuildDependency("Bullet_Physics_SDK_jll"),
    BuildDependency("OpenGLMathematics_jll"),
    Dependency("SDL2_jll"),
    Dependency("SDL2_image_jll"),
    Dependency("SDL2_mixer_jll"),
    Dependency("SDL2_net_jll"),
    Dependency("SDL2_ttf_jll"),
    Dependency("Shaderc_jll"),
    BuildDependency("Vulkan_Headers_jll"),
]

build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; julia_compat="1.6", preferred_gcc_version=v"8.1.0")


text:
function commitmentcost(m::JuMP.AbstractModel, devices::Array{T,1}, device_formulation::Type{D}, system_formulation::Type{S}) where {T <: PowerSystems.ThermalGen, D <: AbstractThermalCommitmentForm, S <: PM.AbstractPowerFormulation}

    on_th = m[:on_th]
    start_th = m[:start_th]
    stop_th = m[:stop_th]

    name_index = m[:on_th].axes[1]

    commit_cost = JuMP.AffExpr()

    for (ix, name) in enumerate(name_index)
        if name == devices[ix].name

                if devices[ix].econ.startupcost > 0

                    JuMP.add_to_expression!(commit_cost,sum(start_th[name,:])*devices[ix].econ.startupcost)

                end

                if devices[ix].econ.shutdncost > 0

                    JuMP.add_to_expression!(commit_cost,sum(stop_th[name,:])*devices[ix].econ.shutdncost)

                end

                if devices[ix].econ.fixedcost > 0

                    JuMP.add_to_expression!(commit_cost,sum(on_th[name,:])*devices[ix].econ.fixedcost)

                end
        else
            @error "Bus name in Array and variable do not match"
        end
    end

    return commit_cost

end


text:
"""
    loglike_beta(nij, A, cij)

Profile likelihood for the frequency response vectors.

*Reference*: Based on Equation (9) of Stompor et al., MNRAS, 392, 216 (2009)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `cij::Array{<:AbstractFloat,2}`: symmetric covariance matrix with the dimention of `(nν, nν)`.
"""
function loglike_beta(
    nij::Array{T,2},
    A::Array{T,2},
    cij::Array{T,2},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2] || size(cij)[1] ≠ size(cij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] || size(nij)[1] ≠ size(cij)[1]
        throw(DimensionMismatch("dimensions of input matrices do not match"))
    else
        N = Symmetric(nij)
        M = Symmetric(cij)
    end
    Ninv = inv(N)
    lnlike = 0.5 * tr(Ninv * A * inv(A' * Ninv * A) * A' * Ninv * M)
end

"""
    loglike_beta(nij, A, d)

Profile likelihood for the frequency response vectors.

*Reference*: Equation (9) of Stompor et al., MNRAS, 392, 216 (2009)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `d::Array{<:AbstractFloat,1}`: data vector for a given pixel (or any other appropriate domain). The number of elements is `nν`.
"""
function loglike_beta(
    nij::Array{T,2},
    A::Array{T,2},
    d::Array{T,1},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] || size(nij)[1] ≠ length(d)
        throw(DimensionMismatch("dimensions of matrices/vector do not match"))
    else
        M = Symmetric(nij)
    end
    x = M \ d
    s = inv(A' * inv(M) * A) * A'x
    lnlike = 0.5 * x'A * s
end

"""
    loglike_beta_grad(nij, A, dAdβ, d)

Derivative of the profile likelihood for the frequency response vectors with respect to a foreground parameter

*Reference*: Equation (A1) of Errard et al., PRD, 84, 063005 (2011)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `dAdβ::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the derivative of the frequency response with respect to a foreground parameter.
    - E.g., ``dAdβ = [zeros(nν) dsynch/dβs zeros(nν)]`` where `zeros(nν)` for CMB and dust because they do not depend on the synchrotron index `βs`.
- `d::Array{<:AbstractFloat,1}`: data vector for a given pixel (or any other appropriate domain). The number of elements is `nν`.
"""
function loglike_beta_grad(
    nij::Array{T,2},
    A::Array{T,2},
    dAdβ::Array{T,2},
    d::Array{T,1},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] ||
           size(nij)[1] ≠ length(d) ||
           size(A)[1] ≠ size(dAdβ)[1] ||
           size(A)[2] ≠ size(dAdβ)[2]
        throw(DimensionMismatch("dimensions of matrices/vector do not match"))
    else
        M = Symmetric(nij)
    end
    x = M \ d
    s = inv(A' * inv(M) * A) * A'x
    lnlike_grad = s' * dAdβ' * (x - M \ (A * s))
end

"""
    loglike_beta_hess(nij, A, dAdβI, dAdβJ, d)

Hessian of the profile likelihood for the frequency response vectors with respect to foreground parameters

*Reference*: Equation (5) of Errard et al., PRD, 84, 063005 (2011)

# Arguments
- `nij::Array{<:AbstractFloat,2}`: symmetric noise covariance matrix with the dimention of `(nν, nν)` where `nν` is the number of frequency bands.
- `A::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the frequency response, for `nc` components in sky.
    - E.g., ``A = [a B]`` where `a = ones(nν)` for CMB and `B` is a `nν`-by-`nc-1` matrix for the frequency response of foreground components.
- `dAdβI::Array{<:AbstractFloat,2}` and `dAdβJ::Array{<:AbstractFloat,2}`: `nν`-by-`nc` matrix of the derivative of the frequency response with respect to a foreground parameter.
    - E.g., ``dAdβI = [zeros(nν) dsynch/dβs zeros(nν)]`` where `zeros(nν)` for CMB and dust because they do not depend on the synchrotron index `βs`.
- `d::Array{<:AbstractFloat,1}`: data vector for a given pixel (or any other appropriate domain). The number of elements is `nν`.
"""
function loglike_beta_hess(
    nij::Array{T,2},
    A::Array{T,2},
    dAdβI::Array{T,2},
    dAdβJ::Array{T,2},
    d::Array{T,1},
) where {T<:AbstractFloat}
    if size(nij)[1] ≠ size(nij)[2]
        throw(DimensionMismatch("covariance matrix must be a square matrix"))
    elseif size(nij)[1] ≠ size(A)[1] ||
           size(nij)[1] ≠ length(d) ||
           size(A)[1] ≠ size(dAdβI)[1] ||
           size(A)[2] ≠ size(dAdβI)[2] ||
           size(A)[1] ≠ size(dAdβJ)[1] ||
           size(A)[2] ≠ size(dAdβJ)[2]
        throw(DimensionMismatch("dimensions of matrices/vector do not match"))
    else
        M = Symmetric(nij)
    end
    x = M \ d
    AtMinvAinv = inv(A' * inv(M) * A)
    s = AtMinvAinv * A'x
    u = M \ (dAdβJ * s)
    v = M \ (A * AtMinvAinv * A'u)
    lnlike_hess = tr(dAdβI' * (v - u) * s')
end


text:
"""
    molar_mass(c::AbstractEoSComponent)

Return the molar mass of a component.
"""
function molar_mass end

"""
    carbon_number(c::AbstractEoSComponent)

Return the number of carbons atoms in the hydrocarbon chain.
"""
function carbon_number end

"""
    name(c::AbstractEoSComponent)

Return the component species name.
"""
function name end

"""
    describe(c::AbstractEoSComponent)

Return `Dict` of parameters. Useful for logging.
"""
function describe end

"""
    load(::Type{T}; name::AbstractString, databases) where {T<:AbstractEoSComponent}

Load component `::T` by its `name` by joining the data on this species in `databases`.
"""
function load end

Base.length(m::AbstractEoSMixture) = m.number_of_components
Base.show(io::IO, x::AbstractEoSMixture) = print(io, "$(typeof(x))($(name(x)))")

components(m::AbstractEoSMixture) = m.components
name(m::AbstractEoSMixture) = join(map(name, components(m)), " + ")
describe(m::AbstractEoSMixture) = Dict{String,Any}("noparameters" => NaN)
load(::Type{T}; names, physdb) where {T<:AbstractEoSMixture} = error("NotImpemented")

function describe(x::BrusilovskyEoSComponent)
    return Dict{String,Any}(
        "data structure" => repr(x),
        "name" => name(x),
        "critical pressure [Pa]" => x.Pc,
        "critical temperature [K]" => x.Tc,
        "pitzer acentric factor" => x.acentric_factor,
        "molar mass [kg mol⁻¹]" => x.molar_mass,
        "number of carbons atoms" => x.carbon_number,
        "eos" => "brusilovsky",
        "eos param: ac [?]" => x.ac,
        "eos param: b [?]"  => x.b,
        "eos param: c [?]"  => x.c,
        "eos param: d [?]"  => x.d,
    )
end

max_stars_repo_name: vvpisarev/CubicEoS

max_issues_repo_name: vvpisarev/CubicEoS

max_forks_repo_name: vvpisarev/CubicEoS


text:
export gali
using OrdinaryDiffEq, LinearAlgebra
#####################################################################################
#                               Continuous GALI                                     #
#####################################################################################
"""
    gali(ds::DynamicalSystem, tmax, k::Int | Q0; kwargs...) -> GALI_k, t
Compute ``\\text{GALI}_k`` [1] for a given `k` up to time `tmax`.
Return ``\\text{GALI}_k(t)`` and time vector ``t``.

The third argument, which sets the order of `gali`, can be an integer `k`, or
a matrix with its columns being the deviation vectors (then
`k = size(Q0)[2]`). In the first case random orthonormal vectors are chosen.

## Keyword Arguments
* `threshold = 1e-12` : If `GALI_k` falls below the `threshold`
  iteration is terminated.
* `dt = 1` : Time-step between deviation vector normalizations. For continuous
  systems this is approximate.
* `u0` : Initial state for the system. Defaults to `get_state(ds)`.
* `diffeq...` : Keyword arguments propagated into `init` of DifferentialEquations.jl.
  See [`trajectory`](@ref) for examples. Only valid for continuous systems.

## Description
The Generalized Alignment Index,
``\\text{GALI}_k``, is an efficient (and very fast) indicator of chaotic or regular
behavior type in ``D``-dimensional Hamiltonian systems
(``D`` is number of variables). The *asymptotic* behavior of
``\\text{GALI}_k(t)`` depends critically on
the type of orbit resulting
from the initial condition. If it is a chaotic orbit, then
```math
\\text{GALI}_k(t) \\sim
\\exp\\left[\\sum_{j=1}^k (\\lambda_1 - \\lambda_j)t \\right]
```
with ``\\lambda_j`` being the `j`-th Lyapunov exponent
(see [`lyapunov`](@ref), [`lyapunovs`](@ref)).
If on the other hand the orbit is regular, corresponding
to movement in ``d``-dimensional torus with `` 1 \\le d \\le D/2``
then it holds
```math
\\text{GALI}_k(t) \\sim
    \\begin{cases}
      \\text{const.}, & \\text{if} \\;\\; 2 \\le k \\le d  \\; \\; \\text{and}
      \\; \\;d > 1 \\\\
      t^{-(k - d)}, & \\text{if} \\;\\;  d < k \\le D - d \\\\
      t^{-(2k - D)}, & \\text{if} \\;\\;  D - d < k \\le D
    \\end{cases}
```

Traditionally, if ``\\text{GALI}_k(t)`` does not become less than
the `threshold` until `tmax`
the given orbit is said to be chaotic, otherwise it is regular.

Our implementation is not based on the original paper, but rather in
the method described in [2], which uses the product of the singular values of ``A``,
a matrix that has as *columns* the deviation vectors.

## Performance Notes
This function uses a [`tangent_integrator`](@ref).
For loops over initial conditions and/or
parameter values one should use the low level method that accepts
an integrator, and `reinit!` it to new initial conditions.
See the "advanced documentation" for info on the integrator object.
The low level method is
```
ChaosTools.gali(tinteg, tmax, dt, threshold)
```

## References

[1] : Skokos, C. H. *et al.*, Physica D **231**, pp 30–54 (2007)

[2] : Skokos, C. H. *et al.*, *Chaos Detection and Predictability* - Chapter 5
(section 5.3.1 and ref. [85] therein), Lecture Notes in Physics **915**,
Springer (2016)
"""
gali(ds::DS, tmax::Real, k::Int; kwargs...) =
    gali(ds, tmax, orthonormal(dimension(ds), k); kwargs...)

function gali(ds::DS{IIP, S, D}, tmax::Real, Q0::AbstractMatrix;
    threshold = 1e-12, dt = 1, u0 = get_state(ds),
    diffeq...) where {IIP, S, D}

    size(Q0)[1] != D && throw(ArgumentError(
    "Deviation vectors must have first dimension equal to the dimension of the "*
    "system, $D."))
    2 ≤ size(Q0)[2] ≤ D || throw(ArgumentError(
    "The order of GALI_k must be 2 ≤ k ≤ $D."))

    # Create tangent integrator:
    if typeof(ds) <: DDS
        tinteg = tangent_integrator(ds, Q0; u0 = u0)
    else
        tinteg = tangent_integrator(ds, Q0; u0 = u0, diffeq...)
    end

    ST = stateeltype(ds)
    TT = typeof(ds.t0)
    gal::Vector{ST}, tvec::Vector{TT} = gali(tinteg, tmax, dt, threshold)
    return gal, tvec
end

function gali(tinteg, tmax, dt, threshold)

    rett = [tinteg.t]
    gali_k = [one(eltype(tinteg.u))]
    t0 = tinteg.t

    while tinteg.t < tmax + t0
        step!(tinteg, dt)
        # Normalize deviation vectors
        normalize_deviations!(tinteg)
        # Calculate singular values:
        zs = LinearAlgebra.svd(get_deviations(tinteg)).S
        push!(gali_k, prod(zs))
        push!(rett, tinteg.t)

        if gali_k[end] < threshold
            break
        end
    end
    return gali_k, rett
end

#####################################################################################
#                            Normalize Deviation Vectors                            #
#####################################################################################
# Metaprogramming ontributed by @saschatimme
function normalize_impl(::Type{SMatrix{D, K, T, DK}}) where {D, K, T, DK}
    exprs = []
    for j = 1:K
        c_j = Symbol("c", j)
        push!(exprs, :($c_j = normalize(A[:, $j])))
    end

    ops = Expr[]
    for j=1:K, i=1:D
        c_j = Symbol("c", j)
        push!(ops, :($c_j[$i]))
    end

    Expr(:block,
        exprs...,
        Expr(:call, SMatrix{D, K, T, D*K}, ops...)
        )
end
@generated function normalize_devs(A::SMatrix)
    normalize_impl(A)
end

# I AM SURE THE FOLLOWING CAN BE SHORTENED USING UNIONS!!!

# OOP version (either cont or disc)
function normalize_deviations!(
    tinteg::ODEIntegrator{Alg, S}) where {Alg, S<:SMatrix}
    norms = normalize_devs(get_deviations(tinteg))
    set_deviations!(tinteg, norms)
    return
end
function normalize_deviations!(
    tinteg::TDI{false})
    norms = normalize_devs(get_deviations(tinteg))
    set_deviations!(tinteg, norms)
    return
end

# IIP
function normalize_inplace!(A)
    for i in 1:size(A)[2]
        LinearAlgebra.normalize!(view(A, :, i))
    end
end
function normalize_deviations!(tinteg::Union{
        ODEIntegrator{Alg, S},
        MDI{Alg, S}}) where {Alg, S<:Matrix}
    A = get_deviations(tinteg)
    normalize_inplace!(A)
    # no reason to call set_deviations, because
    # get_deviations always returns views for IIP
    u_modified!(tinteg, true)
    return
end
function normalize_deviations!(tinteg::TDI{true})
    A = get_deviations(tinteg)
    normalize_inplace!(A)
    return
end


text:
struct NNRODE{C,W,O,P,K} <: NeuralPDEAlgorithm
    chain::C
    W::W
    opt::O
    initθ::P
    autodiff::Bool
    kwargs::K
end
function NNRODE(chain,W,opt=Optim.BFGS(),init_params = nothing;autodiff=false,kwargs...)
    if init_params === nothing
        if chain isa FastChain
            initθ = DiffEqFlux.initial_params(chain)
        else
            initθ,re  = Flux.destructure(chain)
        end
    else
        initθ = init_params
    end
    NNRODE(chain,W,opt,initθ,autodiff,kwargs)
end

function DiffEqBase.solve(
    prob::DiffEqBase.AbstractRODEProblem,
    alg::NeuralPDEAlgorithm,
    args...;
    dt,
    timeseries_errors = true,
    save_everystep=true,
    adaptive=false,
    abstol = 1f-6,
    verbose = false,
    maxiters = 100)

    DiffEqBase.isinplace(prob) && error("Only out-of-place methods are allowed!")

    u0 = prob.u0
    tspan = prob.tspan
    f = prob.f
    p = prob.p
    t0 = tspan[1]

    #hidden layer
    chain  = alg.chain
    opt    = alg.opt
    autodiff = alg.autodiff
    Wg = alg.W
    #train points generation
    ts = tspan[1]:dt:tspan[2]
    initθ = alg.initθ

    if chain isa FastChain
        #The phi trial solution
        if u0 isa Number
            phi = (t,W,θ) -> u0 + (t-tspan[1])*first(chain(adapt(typeof(θ),[t,W]),θ))
        else
            phi = (t,W,θ) -> u0 + (t-tspan[1])*chain(adapt(typeof(θ),[t,W]),θ)
        end
    else
        _,re  = Flux.destructure(chain)
        #The phi trial solution
        if u0 isa Number
            phi = (t,W,θ) -> u0 + (t-t0)*first(re(θ)(adapt(typeof(θ),[t,W])))
        else
            phi = (t,W,θ) -> u0 + (t-t0)*re(θ)(adapt(typeof(θ),[t,W]))
        end
    end

    if autodiff
        # dfdx = (t,W,θ) -> ForwardDiff.derivative(t->phi(t,θ),t)
    else
        dfdx = (t,W,θ) -> (phi(t+sqrt(eps(t)),W,θ) - phi(t,W,θ))/sqrt(eps(t))
    end

    function inner_loss(t,W,θ)
        sum(abs,dfdx(t,W,θ) - f(phi(t,W,θ),p,t,W))
    end
    Wprob = NoiseProblem(Wg,tspan)
    Wsol = solve(Wprob;dt=dt)
    W  = NoiseGrid(ts , Wsol.W)
    function loss(θ)
        sum(abs2,inner_loss(ts[i],W.W[i],θ) for i in 1:length(ts)) # sum(abs2,phi(tspan[1],θ) - u0)
    end

    cb = function (p,l)
        Wprob = NoiseProblem(Wg,tspan)
        Wsol = solve(Wprob;dt=dt)
        W  = NoiseGrid(ts , Wsol.W)
        verbose && println("Current loss is: $l")
        l < abstol
    end
    res = DiffEqFlux.sciml_train(loss, initθ, opt; cb = cb, maxiters=maxiters, alg.kwargs...)

    #solutions at timepoints
    noiseproblem = NoiseProblem(Wg,tspan)
    W = solve(noiseproblem;dt=dt)
    if u0 isa Number
        u = [(phi(ts[i],W.W[i],res.minimizer)) for i in 1:length(ts)]
    else
        u = [(phi(ts[i],W.W[i],res.minimizer)) for i in 1:length(ts)]
    end

    sol = DiffEqBase.build_solution(prob,alg,ts,u,W = W,calculate_error = false)
    DiffEqBase.has_analytic(prob.f) && DiffEqBase.calculate_solution_errors!(sol;timeseries_errors=true,dense_errors=false)
    sol
end #solve


text:
"""
	makecpt(cmd0::String="", arg1=nothing; kwargs...)

Make static color palette tables (CPTs).

Full option list at [`makecpt`]($(GMTdoc)makecpt.html)

- **A** | **alpha** | **transparency** :: [Type => Str]

    Sets a constant level of transparency (0-100) for all color slices.
    ($(GMTdoc)makecpt.html#a)
- $(GMT.opt_C)
- **D** | **bg** | **background** :: [Type => Str | []]			`Arg = [i|o]`

    Select the back- and foreground colors to match the colors for lowest and highest
    z-values in the output CPT.
    ($(GMTdoc)makecpt.html#d)
- **E** | **data_levels** :: [Type => Int | []]		`Arg = [nlevels]`

    Implies reading data table(s) from file or arrays. We use the last data column to
    determine the data range
    ($(GMTdoc)makecpt.html#e)
- **F** | **force_rgb** :: [Type => Str | []]		`Arg = [R|r|h|c][+c]]`

    Force output CPT to written with r/g/b codes, gray-scale values or color name.
    ($(GMTdoc)makecpt.html#f)
- **G** | **truncate** :: [Type => Str]              `Arg = zlo/zhi`

    Truncate the incoming CPT so that the lowest and highest z-levels are to zlo and zhi.
    ($(GMTdoc)makecpt.html#g)
- **I** | **inverse** | **reverse** :: [Type => Str]	`Arg = [c][z]`

    Reverse the sense of color progression in the master CPT.
    ($(GMTdoc)makecpt.html#i)
- **M** | **overrule_bg** :: [Type => Bool]

    Overrule background, foreground, and NaN colors specified in the master CPT with the values of
    the parameters COLOR_BACKGROUND, COLOR_FOREGROUND, and COLOR_NAN.
    ($(GMTdoc)makecpt.html#m)
- **N** | **no_bg** | **nobg** :: [Type => Bool]

    Do not write out the background, foreground, and NaN-color fields.
    ($(GMTdoc)makecpt.html#n)
- **Q** | **log** :: [Type => Bool | Str]			`Arg = [i|o]`

    Selects a logarithmic interpolation scheme [Default is linear].
    ($(GMTdoc)makecpt.html#q)
- **S** | **auto** :: [Type => Bool | Str]			`Arg = [mode]`

    Determine a suitable range for the -T option from the input table(s) (or stdin).
    ($(GMTdoc)makecpt.html#s)
- **T** | **range** :: [Type => Str]			`Arg = [min/max/inc[+b|l|n]|file|list]`

    Defines the range of the new CPT by giving the lowest and highest z-value and interval.
    ($(GMTdoc)makecpt.html#t)
- **W** | **wrap** | **categorical** :: [Type => Bool | Str | []]      `Arg = [w]`

    Do not interpolate the input color table but pick the output colors starting at the
    beginning of the color table, until colors for all intervals are assigned.
    ($(GMTdoc)makecpt.html#w)
- **Z** | **continuous** :: [Type => Bool]

    Creates a continuous CPT [Default is discontinuous, i.e., constant colors for each interval].
    ($(GMTdoc)makecpt.html#z)
- $(GMT.opt_bi)
- $(GMT.opt_di)
- $(GMT.opt_i)
"""
function makecpt(cmd0::String="", arg1=nothing; kwargs...)

	length(kwargs) == 0 && return monolitic("makecpt", cmd0, arg1)	# Monolithic mode

	d = KW(kwargs)
	cmd = parse_common_opts(d, "", [:V_params])

    # If file name sent in, read it and compute a tight -R if this was not provided
	cmd, arg1, opt_R, = read_data(d, cmd0, cmd, arg1, " ")
	cmd, arg1, = add_opt_cpt(d, cmd, [:C :color :cmap], 'C', 0, arg1)
	cmd = parse_these_opts(cmd, d, [[:A :alpha :transparency], [:D :bg :background], [:F :force_rgb],
	                      [:G :truncate], [:I :inverse :reverse], [:M :overrule_bg], [:N :no_bg :nobg], [:Q :log], [:S :auto], [:T :range], [:W :wrap :categorical], [:Z :continuous]])

	if ((val = find_in_dict(d, [:E :data_levels])[1]) !== nothing)
		if (arg1 === nothing)  error("E option requires that a data table is provided as well")
		else                   cmd *= " -E" * arg2str(val)
		end
	end

	global IamModern
	if (IamModern && ((val = find_in_dict(d, [:H :getcpt])[1]) === nothing))  cmd *= " -H"  end
	if (haskey(d, :cptname))  cmd *=  " > " * d[:cptname]  end
	if (dbg_print_cmd(d, cmd) !== nothing)  return cmd  end
	global current_cpt = gmt("makecpt " * cmd, arg1)
end

# ---------------------------------------------------------------------------------------------------
# Version to use with the -E option
makecpt(arg1; kw...) = makecpt("", arg1; kw...)


text:
"""
`datenum(d::Dates.DateTime)` converts a Julia DateTime to a MATLAB style DateNumber. Also see: `toordinal`
"""
function datenum(dt::DateTime)
    dnum = Dates.datetime2epochms(dt)/(1000*60*60*24) + 1
    return dnum
end

"""
`datenum(v...)` does `datenum(DateTime(v...))`.
"""
function datenum(v...)
    dt = DateTime(v...)
    return datenum(dt)
end

"""
`toordinal(dt::DateTime)`converts a Julia DateTime to a python style DateNumber. Also see `datenum`.
"""
function toordinal(dt::DateTime)
    return datenum(dt) - 366
end

"""
`toordinal(v...)` does `toordinal(DateTime(v...))`.
"""
function toordinal(v...)
    dt = DateTime(v...)
    return toordinal(dt)
end

"""
`chkdatetime(v...)` use `try ... catch ...` to check if a vector `[yyyy, mm, dd, hh, MM, ss]` is a valid datetime. It returns `false` if it is not an legal date vector array.
"""
function chkdatetime(v...)
    dtisvalid = true
    try
        DateTime(v...)
    catch
        dtisvalid = false
    end
    return dtisvalid
end


text:
### File for testing the bifurcation functions.

#Tests on the standard networks.
for rn in reaction_networks_standard
    p_vals = map(amp->amp*rand(length(rn.params)),[1.,10.,100.])
    for p in p_vals
        p_bif = rand(rn.params)
        p_range = (0.4:0.4:2.0)*p[1]
        bif = bifurcation_diagram(rn,p,p_bif,p_range)
        for i = 1:5
            p[findfirst(p_bif.==rn.params)] = p_range[i]
            for j = 1:length(bif.fixed_points[i])
                bif.stabilities[i][j][1] ? (tend = 10) : (tend = -10)
                end_point = solve(ODEProblem(rn,bif.fixed_points[i][j],(0.,tend),p),Rosenbrock23())[end]
                @test maximum(end_point-bif.fixed_points[i][j])<0.0001
            end
        end
    end
end

#Tests on the networks containing hill functions.
for rn in reaction_networks_hill
    p_vals = map(amp->rand(1:amp,length(rn.params)),[3,5,10])
    for p in p_vals
        p_bif = rand(rn.params)
        p_range = 1:5
        bif = bifurcation_diagram(rn,p,p_bif,p_range)
        for i = 1:5
            p[findfirst(p_bif.==rn.params)] = p_range[i]
            for j = 1:length(bif.fixed_points[i])
                bif.stabilities[i][j][1] ? (tend = 10) : (tend = -10)
                end_point = solve(ODEProblem(rn,bif.fixed_points[i][j],(0.,tend),p),Rosenbrock23())[end]
                @test maximum(end_point-bif.fixed_points[i][j])<0.0001
            end
        end
    end
end

#Tests on the networks requiring fixed concentrations.
for i = 1:length(reaction_networks_fixed_conc)
    p_vals = map(amp->amp*rand(length(reaction_networks_fixed_conc[i].params)),[1.,10.,100.])
    for p in p_vals
        p_bif = rand(reaction_networks_fixed_conc[i].params)
        p_range = (0.4:0.4:2.0)*p[1]
        bif = bifurcation_diagram(reaction_networks_fixed_conc[i],p,p_bif,p_range,fixed_conc[i]...)
        for I = 1:5
            p[findfirst(p_bif.==reaction_networks_fixed_conc[i].params)] = p_range[I]
            for j = 1:length(bif.fixed_points[I])
                fps = steady_states(reaction_networks_fixed_conc[i],p,fixed_conc[i]...,stability_type=true)
                bif.stabilities[I][j][1] ? (tend = 10) : (tend = -10)
                end_point = solve(ODEProblem(reaction_networks_fixed_conc[i],bif.fixed_points[I][j],(0.,tend),p),Rosenbrock23())[end]
                @test maximum(end_point-bif.fixed_points[I][j])<0.0001
            end
        end
    end
end

#Test that the functionality to allow multiplication of original values workd properly (only tested on standard networks to avoid making tests run for to long).
for rn in reaction_networks_standard
    p_vals = map(amp->amp*rand(length(rn.params)),[1.,10.,100.])
    for p in p_vals
        p_bif = rn.params[1]
        p[1] = 2.0
        p_range_1 = [1.0,1.5,2.0,3.0,4.0]
        p_range_2 = [0.5,0.75,1.0,1.5,2.0]
        bif_1 = bifurcation_diagram(rn,p,p_bif,p_range_1)
        bif_2 = bifurcation_diagram(rn,p,p_bif,p_range_2,mult=true)
        @test bif_1.param_values==bif_2.param_values
        for i = 1:5
            fps1 = sort(bif_1.fixed_points[i])
            fps2 = sort(bif_1.fixed_points[i])
            @test length(fps1)==length(fps2)
            for j = 1:length(fps1)
                @test maximum(fps1[j]-fps2[j])<0.0001
            end
        end
    end
end

#Test to ensure the 2d bifurcation diagrams are found properly (only tested on standard networks to avoid making tests run for to long).
for rn in reaction_networks_standard
    p_vals = map(amp->amp*rand(length(rn.params)),[1.,10.,100.])
    for p in p_vals
        p_bif_1 = rand(rn.params)
        p_range_1 = (1.0:2.0)*p[1]
        p_bif_2 = rand(setdiff(rn.params,Set([p_bif_1])))
        p_range_2 = (1.0:2.0)*p[1]
        bif = bifurcation_diagram_2d(rn,p,p_bif_1,p_bif_2,p_range_1,p_range_2)
        for i1 = 1:2, i2 = 1:2
            p[findfirst(p_bif_1.==rn.params)] = p_range_1[i1]
            p[findfirst(p_bif_2.==rn.params)] = p_range_2[i2]
            for j = 1:length(bif.fixed_points[i1,i2])
                bif.stabilities[i1,i2][j][1] ? (tend = 10) : (tend = -10)
                end_point = solve(ODEProblem(rn,bif.fixed_points[i1,i2][j],(0.,tend),p),Rosenbrock23())[end]
                @test maximum(end_point-bif.fixed_points[i1,i2][j])<0.0001
            end
        end
    end
end

max_stars_repo_name: TorkelE/BRNEquilibrate

max_issues_repo_name: TorkelE/BRNEquilibrate

max_forks_repo_name: TorkelE/BRNEquilibrate


text:
using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))
Pkg.instantiate()

# ## setup
using Dojo
using IterativeLQR
using LinearAlgebra

# ## system
gravity = -9.81
timestep = 0.05
friction_coefficient = 0.8
damper = 5.0
spring = 0.0
env = get_environment(:quadruped,
    representation=:minimal,
    timestep=timestep,
    contact_body=false,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    damper=damper,
    spring=spring)

# ## template
include(joinpath(@__DIR__, "../../environments/quadruped/methods/template.jl"))

# ## visualizer
open(env.vis)

# ## dimensions
n = env.num_states
m = env.num_inputs

# ## reference trajectory
N = 2
initialize!(env.mechanism, :quadruped)
xref = quadruped_trajectory(env.mechanism,
    r=0.05,
    z=0.29;
    Δx=-0.04,
    Δfront=0.10,
    N=10,
    Ncycles=N)
zref = [minimal_to_maximal(env.mechanism, x) for x in xref]
visualize(env, xref)

# ## gravity compensation
## TODO: solve optimization problem instead
mech = get_mechanism(:quadruped,
    timestep=timestep,
    gravity=gravity,
    friction_coefficient=friction_coefficient,
    damper=damper,
    spring=spring)

initialize!(mech, :quadruped)
storage = simulate!(mech, 1.0,
    record=true,
    verbose=false)

visualize(mech, storage,
    vis=env.vis)
ugc = gravity_compensation(mech)
u_control = ugc[6 .+ (1:12)]

# ## horizon
T = N * (21 - 1) + 1

# ## model
dyn = IterativeLQR.Dynamics(
    (y, x, u, w) -> dynamics(y, env, x, u, w),
    (dx, x, u, w) -> dynamics_jacobian_state(dx, env, x, u, w),
    (du, x, u, w) -> dynamics_jacobian_input(du, env, x, u, w),
    n, n, m)

model = [dyn for t = 1:T-1]

# ## rollout
x1 = xref[1]
ū = [u_control for t = 1:T-1]
x̄ = IterativeLQR.rollout(model, x1, ū)
visualize(env, x̄)

# ## objective
qt = [0.3; 0.05; 0.05; 0.01 * ones(3); 0.01 * ones(3); 0.01 * ones(3); fill([0.2, 0.001], 12)...]
ots = [(x, u, w) -> transpose(x - xref[t]) * Diagonal(timestep * qt) * (x - xref[t]) + transpose(u) * Diagonal(timestep * 0.01 * ones(m)) * u for t = 1:T-1]
oT = (x, u, w) -> transpose(x - xref[end]) * Diagonal(timestep * qt) * (x - xref[end])

cts = IterativeLQR.Cost.(ots, n, m)
cT = IterativeLQR.Cost(oT, n, 0)
obj = [[cts for t = 1:T-1]..., cT]

# ## constraints
function goal(x, u, w)
    Δ = x - xref[end]
    return Δ[collect(1:3)]
end

cont = IterativeLQR.Constraint()
conT = IterativeLQR.Constraint(goal, n, 0)
cons = [[cont for t = 1:T-1]..., conT]

# ## solver
s = IterativeLQR.solver(model, obj, cons,
    opts=Options(verbose=false,
        linesearch=:armijo,
        α_min=1.0e-5,
        obj_tol=1.0e-3,
        grad_tol=1.0e-3,
        max_iter=100,
        max_al_iter=5,
        ρ_init=1.0,
        ρ_scale=10.0))
IterativeLQR.initialize_controls!(s, ū)
IterativeLQR.initialize_states!(s, x̄)

# ## solve
@time IterativeLQR.solve!(s)

# ## solution
x_sol, u_sol = IterativeLQR.get_trajectory(s)
@show IterativeLQR.eval_obj(s.m_data.obj.costs, s.m_data.x, s.m_data.u, s.m_data.w)
@show s.s_data.iter[1]
@show norm(goal(s.m_data.x[T], zeros(0), zeros(0)), Inf)

# ## visualize
vis= Visualizer()
open(env.vis)
x_view = [[x_sol[1] for t = 1:15]..., x_sol..., [x_sol[end] for t = 1:15]...]
visualize(env, x_view)

set_camera!(env.vis,
    cam_pos=[0.0, -3.0, 2.0],
    zoom=3.0)


text:
kronprod = kron

kronpow(A::Matrix,n::Integer) = n==1? A : kronprod([A for _ in 1:n]...)

function kronsum(terms...)
    Ms,Ns = zip([size(term) for term in terms]...)
    Ms == Ns || throw(DimensionMismatch("term not square"))
    D = [eye(M) for M in Ms]
    return sum([kronprod(setindex!(copy(D),term,i)...) for (i,term) in enumerate(terms)])
end

#function kronrepeatedsum(A::Matrix,n::Integer)
#    B = A
#    for k=1:n
#        B = kronsum(B,A)
#    end
#    return B
#end


text:
using EventSimulation
using Test
using Random

tests = ["core",
         "prioritytime",
         "queue",
         "resource",
         "../examples/mms_example",
         "../examples/mm1_example",
        ]

println("Running tests:")

for t in tests
    tfile = string(t, ".jl")
    println(" * $(tfile) ...")
    include(tfile)
end