text:
# This file contains the tests for the `PDF` module

module TestPDF

using PDFHighlights
using Test

# Print the header
println("\e[1;32mRUNNING\e[0m: TestPDF.jl")

const pdf = joinpath(@__DIR__, "..", "pdf", "TestPDF.pdf")

@testset "_concatenate" begin

    # Without halves of words
    @test PDFHighlights.Internal.PDF._concatenate(
        String["Highlight 1", "Highlight 2", "Highlight 3"],
        String["Comment 1", ".c1 Comment 2", ".c2 Comment 3"],
        Int32[1, 2, 3],
    ) ==
    (
        String["Highlight 1", "Highlight 2 Highlight 3"],
        String["Comment 1", "Comment 2 Comment 3"],
        Int32[1, 2],
    )

    # With halves of words
    @test PDFHighlights.Internal.PDF._concatenate(
        String["Highlight 1", "High-", "light 2"],
        String["Comment 1", ".c1 Comment 2", ".c2"],
        Int32[1, 2, 3],
    ) ==
    (
        String["Highlight 1", "Highlight 2"],
        String["Comment 1", "Comment 2"],
        Int32[1, 2],
    )

end

@testset "_get_authors_from_PDF" begin

    dir = joinpath(@__DIR__, "..")

    @test PDFHighlights.Internal.PDF._get_authors_from_PDF(dir) == String[
        "Pavel Sobolev",
    ]

end

@testset "_get_highlights_from_PDF" begin

    # With concatenation
    @test PDFHighlights.Internal.PDF._get_highlights_from_PDF(pdf) == String[
        "Highlight 1",
        "Highlight 2 Highlight 3",
        "Highlight 4",
        "Highhighlight 5",
        "6th Highhigh light-",
        "High light 7",
        "8th Highlight-",
    ]

    # Without concatenation
    @test PDFHighlights.Internal.PDF._get_highlights_from_PDF(pdf; concatenate=false) ==
    String[
        "Highlight 1",
        "Highlight 2",
        "Highlight 3",
        "High-",
        "light 4",
        "Highhighlight 5",
        "6th Highhigh light-",
        "High light 7",
        "8th Highlight-",
    ]

end

@testset "_get_titles_from_PDF" begin

    dir = joinpath(@__DIR__, "..")

    @test PDFHighlights.Internal.PDF._get_titles_from_PDF(dir) == String[
        "A dummy PDF for tests",
    ]

end

@testset "_sort! (lines)" begin

    lines = ["High", "high", "light"]
    quad_x_anchors = Float64[0.21, 0.15, 0.17]
    quad_yl_anchors = Float64[0.10, 0.10, 0.10]
    quad_yu_anchors = Float64[0.15, 0.12, 0.15]

    @test PDFHighlights.Internal.PDF._sort!(
        lines,
        quad_x_anchors,
        quad_yl_anchors,
        quad_yu_anchors,
    ) == ["high", "light", "High"]

    lines = ["High", "high", "light"]
    quad_x_anchors = Float64[0.21, 0.15, 0.17]
    quad_yl_anchors = Float64[0.10, 0.10, 0.17]
    quad_yu_anchors = Float64[0.15, 0.12, 0.20]

    @test PDFHighlights.Internal.PDF._sort!(
        lines,
        quad_x_anchors,
        quad_yl_anchors,
        quad_yu_anchors,
    ) == ["high", "High", "light"]

    lines = ["High", "high", "highlight", "light"]
    quad_x_anchors = Float64[0.21, 0.15, 0.17, 0.19]
    quad_yl_anchors = Float64[0.10, 0.10, 0.17, 0.17]
    quad_yu_anchors = Float64[0.15, 0.12, 0.20, 0.20]

    @test PDFHighlights.Internal.PDF._sort!(
        lines,
        quad_x_anchors,
        quad_yl_anchors,
        quad_yu_anchors,
    ) == ["high", "High", "highlight", "light"]

end

@testset "_sort! (highlights)" begin

    pages = Int32[1, 1, 1]

    highlights = ["High", "high", "light"]
    comments = ["Com", "com", "ment"]
    quad_x_anchors = Float64[0.21, 0.15, 0.17]
    quad_yl_anchors = Float64[0.10, 0.10, 0.10]
    quad_yu_anchors = Float64[0.12, 0.15, 0.15]

    @test PDFHighlights.Internal.PDF._sort!(
        highlights,
        comments,
        pages,
        quad_x_anchors,
        quad_yl_anchors,
        quad_yu_anchors,
    ) ==
    (
        ["high", "light", "High"],
        ["com", "ment", "Com"],
    )

    highlights = ["High", "high", "light"]
    comments = ["Com", "com", "ment"]
    quad_x_anchors = Float64[0.21, 0.15, 0.17]
    quad_yl_anchors = Float64[0.10, 0.10, 0.17]
    quad_yu_anchors = Float64[0.12, 0.15, 0.20]

    @test PDFHighlights.Internal.PDF._sort!(
        highlights,
        comments,
        pages,
        quad_x_anchors,
        quad_yl_anchors,
        quad_yu_anchors,
    ) ==
    (
        ["high", "High", "light"],
        ["com", "Com", "ment"],
    )

    pages = Int32[1, 1, 1, 1]

    highlights = ["High", "high", "highlight", "light"]
    comments = ["Com", "com", "comment", "ment"]
    quad_x_anchors = Float64[0.21, 0.15, 0.17, 0.19]
    quad_yl_anchors = Float64[0.10, 0.10, 0.17, 0.17]
    quad_yu_anchors = Float64[0.12, 0.15, 0.20, 0.20]

    @test PDFHighlights.Internal.PDF._sort!(
        highlights,
        comments,
        pages,
        quad_x_anchors,
        quad_yl_anchors,
        quad_yu_anchors,
    ) ==
    (
        ["high", "High", "highlight", "light"],
        ["com", "Com", "comment", "ment"],
    )

    pages = Int32[1, 1, 2, 2]

    highlights = ["High", "highlight", "light", "high"]
    comments = ["Com", "comment", "ment", "com"]
    quad_yl_anchors = Float64[0.08, 0.10, 0.17, 0.10]
    quad_yu_anchors = Float64[0.09, 0.15, 0.20, 0.12]

    @test PDFHighlights.Internal.PDF._sort!(
        highlights,
        comments,
        pages,
        quad_x_anchors,
        quad_yl_anchors,
        quad_yu_anchors,
    ) ==
    (
        ["High", "highlight", "high", "light"],
        ["Com", "comment", "com", "ment"],
    )

end

@testset "get_author_title" begin

    @test get_author_title(pdf) == ("Pavel Sobolev", "A dummy PDF for tests")

    @test_throws(
        PDFHighlights.Internal.Exceptions.FileDoesNotExist("oof"),
        get_author_title("oof"),
    )

    touch("oof")

    @test_throws(
        PDFHighlights.Internal.Exceptions.NotPDF("oof"),
        get_author_title("oof"),
    )

    isfile("oof") && rm("oof")

end

@testset "get_author" begin

    @test get_author(pdf) == "Pavel Sobolev"

end

@testset "get_authors_titles" begin

    dir = joinpath(@__DIR__, "..")

    @test get_authors_titles(dir) == (
        String["Pavel Sobolev",],
        String["A dummy PDF for tests",]
    )

    @test_throws(
        PDFHighlights.Internal.Exceptions.DirectoryDoesNotExist("oof"),
        get_authors_titles("oof"),
    )

end

@testset "get_highlights_comments_pages" begin

    # With concatenation
    @test get_highlights_comments_pages(pdf) ==
    (
        String[
            "Highlight 1",
            "Highlight 2 Highlight 3",
            "Highlight 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        String[
            "Comment 1",
            "Comment 2 Comment 3",
            "Comment 4",
            "",
            "",
            "",
            "",
        ],
        Int32[1, 2, 4, 6, 7, 8, 9],
    )

    # Without concatenation
    @test get_highlights_comments_pages(pdf; concatenate=false) ==
    (
        String[
            "Highlight 1",
            "Highlight 2",
            "Highlight 3",
            "High-",
            "light 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        String[
            "Comment 1",
            ".c1 Comment 2",
            ".c2 Comment 3",
            ".c1 Comment 4",
            ".c2",
            "",
            "",
            "",
            "",
        ],
        Int32[1:9...],
    )

    @test_throws(
        PDFHighlights.Internal.Exceptions.DoesNotExist("oof"),
        get_highlights_comments_pages("oof"),
    )

    touch("oof")

    @test_throws(
        PDFHighlights.Internal.Exceptions.NotPDF("oof"),
        get_highlights_comments_pages("oof"),
    )

    isfile("oof") && rm("oof")

    dir = joinpath(@__DIR__, "..")

    # With concatenation
    @test get_highlights_comments_pages(dir) ==
    (
        String[
            "Highlight 1",
            "Highlight 2 Highlight 3",
            "Highlight 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        String[
            "Comment 1",
            "Comment 2 Comment 3",
            "Comment 4",
            "",
            "",
            "",
            "",
        ],
        Int32[1, 2, 4, 6, 7, 8, 9],
    )

    # Without concatenation
    @test get_highlights_comments_pages(dir; concatenate=false) ==
    (
        String[
            "Highlight 1",
            "Highlight 2",
            "Highlight 3",
            "High-",
            "light 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        String[
            "Comment 1",
            ".c1 Comment 2",
            ".c2 Comment 3",
            ".c1 Comment 4",
            ".c2",
            "",
            "",
            "",
            "",
        ],
        Int32[1:9...],
    )

end

@testset "get_comments_pages" begin

    # With concatenation
    @test get_comments_pages(pdf; concatenate=true) ==
    (
        String["Comment 1", "Comment 2 Comment 3", "Comment 4", "", "", "", ""],
        Int32[1, 2, 4, 6, 7, 8, 9],
    )

    # Without concatenation
    @test get_comments_pages(pdf) ==
    (
        String[
            "Comment 1",
            ".c1 Comment 2",
            ".c2 Comment 3",
            ".c1 Comment 4",
            ".c2",
            "",
            "",
            "",
            "",
        ],
        Int32[1:9...],
    )

end

@testset "get_highlights_comments" begin

    # With concatenation
    @test get_highlights_comments(pdf) ==
    (
        String[
            "Highlight 1",
            "Highlight 2 Highlight 3",
            "Highlight 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        String[
            "Comment 1",
            "Comment 2 Comment 3",
            "Comment 4",
            "",
            "",
            "",
            "",
        ],
    )

    # Without concatenation
    @test get_highlights_comments(pdf; concatenate=false) ==
    (
        String[
            "Highlight 1",
            "Highlight 2",
            "Highlight 3",
            "High-",
            "light 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        String[
            "Comment 1",
            ".c1 Comment 2",
            ".c2 Comment 3",
            ".c1 Comment 4",
            ".c2",
            "",
            "",
            "",
            "",
        ],
    )

end

@testset "get_highlights_pages" begin

    # With concatenation
    @test get_highlights_pages(pdf) ==
    (
        String[
            "Highlight 1",
            "Highlight 2 Highlight 3",
            "Highlight 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        Int32[1, 2, 4, 6, 7, 8, 9],
    )

    # Without concatenation
    @test get_highlights_pages(pdf; concatenate=false) ==
    (
        String[
            "Highlight 1",
            "Highlight 2",
            "Highlight 3",
            "High-",
            "light 4",
            "Highhighlight 5",
            "6th Highhigh light-",
            "High light 7",
            "8th Highlight-",
        ],
        Int32[1:9...],
    )

end

@testset "get_comments" begin

    # With concatenation
    @test get_comments(pdf; concatenate=true) == String[
        "Comment 1",
        "Comment 2 Comment 3",
        "Comment 4",
        "",
        "",
        "",
        "",
    ]

    # Without concatenation
    @test get_comments(pdf) == String[
        "Comment 1",
        ".c1 Comment 2",
        ".c2 Comment 3",
        ".c1 Comment 4",
        ".c2",
        "",
        "",
        "",
        "",
    ]

end

@testset "get_pages" begin

    # With concatenation
    @test get_pages(pdf; concatenate=true) ==
    Int32[1, 2, 4, 6, 7, 8, 9]

    # Without concatenation
    @test get_pages(pdf) ==
    Int32[1:9...]

end

@testset "get_title" begin

    @test get_title(pdf) == "A dummy PDF for tests"

end

end


text:
@testset "Zlattices" begin
  G = matrix(ZZ, 2, 2, [2,1,1,2])
  B = matrix(ZZ, 1, 2, [1,0])
  L = Zlattice(B, gram=G)
  LL = orthogonal_sum(L,L)
  @test rank(LL) == 2*rank(L)
end

ext: jl


text:
"""
Simulated Binary Crossover (SBX).

See Deb&Agrawal "Simulated binary crossover for continuous search space", 1994, Complex Systems
"""
struct SimulatedBinaryCrossover <: CrossoverOperator{2,2}
    p::Float64      # probability to modify a dimension
    η::Float64      # distribution index
    η_exp::Float64  # pre-processed index

    function SimulatedBinaryCrossover(p::Number, η::Number)
        (0.0 <= p <= 1.0) || throw(ArgumentError("p must be within [0,1]"))
        η > 0.0 || throw(ArgumentError("η must be positive"))
        new(p, η, (1/(η+1)))
    end
    SimulatedBinaryCrossover(params::Parameters) =
        SimulatedBinaryCrossover(params[:SBX_p], params[:SBX_η])
end

const SBX_DefaultOptions = ParamsDict(
    :SBX_p => 0.2,
    :SBX_η => 3.0
)

# sample SPX spread factor β
function randbeta(xover::SimulatedBinaryCrossover)
    u = rand()
    return (u < 0.5 ? 2*u : 1/(2*(1-u)))^xover.η_exp
end

function apply!(xover::SimulatedBinaryCrossover,
                targets::AbstractVector{<:AbstractIndividual}, targetIndices::Vector{Int},
                pop, parentIndices)
    @assert length(targets) == 2
    @assert length(targetIndices) == 2
    @assert length(parentIndices) == 2
    p1ix, p2ix = parentIndices
    @inbounds for i in 1:numdims(pop)
        if rand() <= xover.p
            # crossover
            β = randbeta(xover)
            parent_avg = 0.5 * (pop[i,p1ix] + pop[i,p2ix])
            delta = 0.5 * (pop[i,p1ix] - pop[i,p2ix])
            if rand(Bool) delta = -delta end # swap the parents
            targets[1][i] = parent_avg + β * delta
            targets[2][i] = parent_avg - β * delta
        else
            # just copy the parents
            targets[1][i] = pop[i,p1ix]
            targets[2][i] = pop[i,p2ix]
        end
    end
    return targets
end


text:
module Topologies
import ..BrickMesh
import MPI
using DocStringExtensions

export AbstractTopology, BrickTopology, StackedBrickTopology,
    CubedShellTopology, StackedCubedSphereTopology, isstacked

export grid1d, SingleExponentialStretching, InteriorStretching

"""
    AbstractTopology{dim}

Represents the connectivity of individual elements, with local dimension `dim`.
"""
abstract type AbstractTopology{dim} end

"""
    BoxElementTopology{dim, T} <: AbstractTopology{dim}

The local topology of a larger MPI-distributed topology, represented by
`dim`-dimensional box elements.

This contains the necessary information for the connectivity elements of the
elements on the local process, along with "ghost" elements from neighbouring
processes.

# Fields

$(DocStringExtensions.FIELDS)
"""
struct BoxElementTopology{dim, T} <: AbstractTopology{dim}

  """
  MPI communicator for communicating with neighbouring processes.
  """
  mpicomm::MPI.Comm

  """
  Range of element indices
  """
  elems::UnitRange{Int64}

  """
  Range of real (aka nonghost) element indices
  """
  realelems::UnitRange{Int64}

  """
  Range of ghost element indices
  """
  ghostelems::UnitRange{Int64}

  """
  Ghost element to face is received; `ghostfaces[f,ge] == true` if face `f` of
  ghost element `ge` is received.
  """
  ghostfaces::BitArray{2}

  """
  Array of send element indices
  """
  sendelems::Array{Int64, 1}

  """
  Send element to face is sent; `sendfaces[f,se] == true` if face `f` of send
  element `se` is sent.
  """
  sendfaces::BitArray{2}

  """
  Element to vertex coordinates; `elemtocoord[d,i,e]` is the `d`th coordinate of
  corner `i` of element `e`

  !!! note
      currently coordinates always are of size 3 for `(x1, x2, x3)`
  """
  elemtocoord::Array{T, 3}

  """
  Element to neighboring element; `elemtoelem[f,e]` is the number of the element
  neighboring element `e` across face `f`.  If there is no neighboring element
  then `elemtoelem[f,e] == e`.
  """
  elemtoelem::Array{Int64, 2}

  """
  Element to neighboring element face; `elemtoface[f,e]` is the face number of
  the element neighboring element `e` across face `f`.  If there is no
  neighboring element then `elemtoface[f,e] == f`."
  """
  elemtoface::Array{Int64, 2}

  """
  element to neighboring element order; `elemtoordr[f,e]` is the ordering number
  of the element neighboring element `e` across face `f`.  If there is no
  neighboring element then `elemtoordr[f,e] == 1`.
  """
  elemtoordr::Array{Int64, 2}

  """
  Element to boundary number; `elemtobndy[f,e]` is the boundary number of face
  `f` of element `e`.  If there is a neighboring element then `elemtobndy[f,e]
  == 0`.
  """
  elemtobndy::Array{Int64, 2}

  """
  List of the MPI ranks for the neighboring processes
  """
  nabrtorank::Array{Int64, 1}

  """
  Range in ghost elements to receive for each neighbor
  """
  nabrtorecv::Array{UnitRange{Int64}, 1}

  """
  Range in `sendelems` to send for each neighbor
  """
  nabrtosend::Array{UnitRange{Int64}, 1}

  """
  original order in partitioning
  """
  origsendorder::Array{Int64,1}

  """
  boolean for whether or not this topology has a boundary
  """
  hasboundary::Bool
end

"""
    hasboundary(topology::AbstractTopology)

query function to check whether a topology has a boundary (i.e., not fully
periodic)
"""
hasboundary(topology::AbstractTopology) = topology.hasboundary

if VERSION >= v"1.2-"
  isstacked(::T) where {T<:AbstractTopology} = hasfield(T, :stacksize)
else
  isstacked(::T) where {T<:AbstractTopology} = Base.fieldindex(T, :stacksize,
                                                               false) > 0
end

"""
    BrickTopology{dim, T} <: AbstractTopology{dim}

A simple grid-based topology. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct BrickTopology{dim, T} <: AbstractTopology{dim}
  topology::BoxElementTopology{dim, T}
end
Base.getproperty(a::BrickTopology, p::Symbol) =
getproperty(getfield(a, :topology), p)

"""
    CubedShellTopology{T} <: AbstractTopology{2}

A cube-shell topology. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct CubedShellTopology{T} <: AbstractTopology{2}
  topology::BoxElementTopology{2, T}
end
Base.getproperty(a::CubedShellTopology, p::Symbol) =
getproperty(getfield(a, :topology), p)

"""
    StackedBrickTopology{dim, T} <: AbstractTopology{dim}

A simple grid-based topology, where all elements on the trailing dimension are
stacked to be contiguous. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct StackedBrickTopology{dim, T} <: AbstractTopology{dim}
  topology::BoxElementTopology{dim, T}
  stacksize::Int64
end
function Base.getproperty(a::StackedBrickTopology, p::Symbol)
  return p == :stacksize ? getfield(a, p) :
                           getproperty(getfield(a, :topology), p)
end

"""
    StackedCubedSphereTopology{3, T} <: AbstractTopology{3}

A cube-sphere topology. All elements on the same "vertical" dimension are
stacked to be contiguous. This is a convenience wrapper around
[`BoxElementTopology`](@ref).
"""
struct StackedCubedSphereTopology{T} <: AbstractTopology{3}
  topology::BoxElementTopology{3, T}
  stacksize::Int64
end
function Base.getproperty(a::StackedCubedSphereTopology, p::Symbol)
  return p == :stacksize ? getfield(a, p) :
                           getproperty(getfield(a, :topology), p)
end

""" A wrapper for the BrickTopology """
BrickTopology(mpicomm, Nelems::NTuple{N, Integer}; kw...) where N =
BrickTopology(mpicomm, map(Ne->0:Ne, Nelems); kw...)

"""
    BrickTopology{dim, T}(mpicomm, elemrange; boundary, periodicity)

Generate a brick mesh topology with coordinates given by the tuple `elemrange`
and the periodic dimensions given by the `periodicity` tuple.

The elements of the brick are partitioned equally across the MPI ranks based
on a space-filling curve.

By default boundary faces will be marked with a one and other faces with a
zero.  Specific boundary numbers can also be passed for each face of the brick
in `boundary`.  This will mark the nonperiodic brick faces with the given
boundary number.

# Examples

We can build a 3 by 2 element two-dimensional mesh that is periodic in the
\$x2\$-direction with
```jldoctest brickmesh

using CLIMA.Topologies
using MPI
MPI.Init()
topology = BrickTopology(MPI.COMM_SELF, (2:5,4:6);
                         periodicity=(false,true),
                         boundary=((1,2),(3,4)))
```
This returns the mesh structure for

             x2

              ^
              |
             6-  +-----+-----+-----+
              |  |     |     |     |
              |  |  3  |  4  |  5  |
              |  |     |     |     |
             5-  +-----+-----+-----+
              |  |     |     |     |
              |  |  1  |  2  |  6  |
              |  |     |     |     |
             4-  +-----+-----+-----+
              |
              +--|-----|-----|-----|--> x1
                 2     3     4     5

For example, the (dimension by number of corners by number of elements) array
`elemtocoord` gives the coordinates of the corners of each element.
```jldoctest brickmesh
julia> topology.elemtocoord
2×4×6 Array{Int64,3}:
[:, :, 1] =
 2  3  2  3
 4  4  5  5

[:, :, 2] =
 3  4  3  4
 4  4  5  5

[:, :, 3] =
 2  3  2  3
 5  5  6  6

[:, :, 4] =
 3  4  3  4
 5  5  6  6

[:, :, 5] =
 4  5  4  5
 5  5  6  6

[:, :, 6] =
 4  5  4  5
 4  4  5  5
```
Note that the corners are listed in Cartesian order.

The (number of faces by number of elements) array `elemtobndy` gives the
boundary number for each face of each element.  A zero will be given for
connected faces.
```jldoctest brickmesh
julia> topology.elemtobndy
4×6 Array{Int64,2}:
 1  0  1  0  0  0
 0  0  0  0  2  2
 0  0  0  0  0  0
 0  0  0  0  0  0
```
Note that the faces are listed in Cartesian order.

"""
function BrickTopology(mpicomm, elemrange;
                       boundary=ntuple(j->(1,1), length(elemrange)),
                       periodicity=ntuple(j->false, length(elemrange)),
                       connectivity=:face, ghostsize=1)

  if boundary isa Matrix
    boundary = tuple(mapslices(x -> tuple(x...), boundary, dims=1)...)
  end

  # We cannot handle anything else right now...
  @assert connectivity == :face
  @assert ghostsize == 1

  mpirank = MPI.Comm_rank(mpicomm)
  mpisize = MPI.Comm_size(mpicomm)
  topology = BrickMesh.brickmesh(elemrange, periodicity, part=mpirank+1,
                              numparts=mpisize, boundary=boundary)
  topology = BrickMesh.partition(mpicomm, topology...)
  origsendorder = topology[5]
  topology = BrickMesh.connectmesh(mpicomm, topology[1:4]...)

  dim = length(elemrange)
  T = eltype(topology.elemtocoord)
  return BrickTopology{dim, T}(BoxElementTopology{dim, T}(
              mpicomm, topology.elems, topology.realelems,
              topology.ghostelems, topology.ghostfaces, topology.sendelems,
              topology.sendfaces, topology.elemtocoord, topology.elemtoelem,
              topology.elemtoface, topology.elemtoordr, topology.elemtobndy,
              topology.nabrtorank, topology.nabrtorecv, topology.nabrtosend,
              origsendorder, !minimum(periodicity)))
end

""" A wrapper for the StackedBrickTopology """
StackedBrickTopology(mpicomm, Nelems::NTuple{N, Integer}; kw...) where N =
StackedBrickTopology(mpicomm, map(Ne->0:Ne, Nelems); kw...)

"""
    StackedBrickTopology{dim, T}(mpicomm, elemrange; boundary, periodicity)

Generate a stacked brick mesh topology with coordinates given by the tuple
`elemrange` and the periodic dimensions given by the `periodicity` tuple.

The elements are stacked such that the elements associated with range
`elemrange[dim]` are contiguous in the element ordering.

The elements of the brick are partitioned equally across the MPI ranks based
on a space-filling curve.  Further, stacks are not split at MPI boundaries.

By default boundary faces will be marked with a one and other faces with a
zero.  Specific boundary numbers can also be passed for each face of the brick
in `boundary`.  This will mark the nonperiodic brick faces with the given
boundary number.

# Examples

We can build a 3 by 2 element two-dimensional mesh that is periodic in the
\$x2\$-direction with
```jldoctest brickmesh

using CLIMA.Topologies
using MPI
MPI.Init()
topology = StackedBrickTopology(MPI.COMM_SELF, (2:5,4:6);
                                periodicity=(false,true),
                                boundary=((1,2),(3,4)))
```
This returns the mesh structure stacked in the \$x2\$-direction for

             x2

              ^
              |
             6-  +-----+-----+-----+
              |  |     |     |     |
              |  |  2  |  4  |  6  |
              |  |     |     |     |
             5-  +-----+-----+-----+
              |  |     |     |     |
              |  |  1  |  3  |  5  |
              |  |     |     |     |
             4-  +-----+-----+-----+
              |
              +--|-----|-----|-----|--> x1
                 2     3     4     5

For example, the (dimension by number of corners by number of elements) array
`elemtocoord` gives the coordinates of the corners of each element.
```jldoctest brickmesh
julia> topology.elemtocoord
2×4×6 Array{Int64,3}:
[:, :, 1] =
 2  3  2  3
 4  4  5  5

[:, :, 2] =
 2  3  2  3
 5  5  6  6

[:, :, 3] =
 3  4  3  4
 4  4  5  5

[:, :, 4] =
 3  4  3  4
 5  5  6  6

[:, :, 5] =
 4  5  4  5
 4  4  5  5

[:, :, 6] =
 4  5  4  5
 5  5  6  6
```
Note that the corners are listed in Cartesian order.

The (number of faces by number of elements) array `elemtobndy` gives the
boundary number for each face of each element.  A zero will be given for
connected faces.
```jldoctest brickmesh
julia> topology.elemtobndy
4×6 Array{Int64,2}:
 1  0  1  0  0  0
 0  0  0  0  2  2
 0  0  0  0  0  0
 0  0  0  0  0  0
```
Note that the faces are listed in Cartesian order.
"""
function StackedBrickTopology(mpicomm, elemrange;
                       boundary=ntuple(j->(1,1), length(elemrange)),
                       periodicity=ntuple(j->false, length(elemrange)),
                       connectivity=:face, ghostsize=1)

  if boundary isa Matrix
    boundary = tuple(mapslices(x -> tuple(x...), boundary, dims=1)...)
  end

  dim = length(elemrange)

  dim <= 1 && error("Stacked brick topology works for 2D and 3D")

  # Build the base topology
  basetopo = BrickTopology(mpicomm, elemrange[1:dim-1];
                     boundary=boundary[1:dim-1],
                     periodicity=periodicity[1:dim-1],
                     connectivity=connectivity,
                     ghostsize=ghostsize)

  # Use the base topology to build the stacked topology
  stack = elemrange[dim]
  stacksize = length(stack) - 1

  nvert = 2^dim
  nface = 2dim

  nreal = length(basetopo.realelems)*stacksize
  nghost = length(basetopo.ghostelems)*stacksize

  elems=1:(nreal+nghost)
  realelems=1:nreal
  ghostelems=nreal.+(1:nghost)

  sendelems = similar(basetopo.sendelems,
                      length(basetopo.sendelems)*stacksize)
  for i=1:length(basetopo.sendelems), j=1:stacksize
    sendelems[stacksize*(i-1) + j] = stacksize*(basetopo.sendelems[i]-1) + j
  end

  ghostfaces = similar(basetopo.ghostfaces, nface, length(ghostelems))
  ghostfaces .= false

  for i=1:length(basetopo.ghostelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f = 1:2(dim-1)
      ghostfaces[f, e] = basetopo.ghostfaces[f, i]
    end
  end

  sendfaces = similar(basetopo.sendfaces, nface, length(sendelems))
  sendfaces .= false

  for i=1:length(basetopo.sendelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f = 1:2(dim-1)
      sendfaces[f, e] = basetopo.sendfaces[f, i]
    end
  end

  elemtocoord = similar(basetopo.elemtocoord, dim, nvert, length(elems))

  for i=1:length(basetopo.elems), j=1:stacksize
    e = stacksize*(i-1) + j

    for v = 1:2^(dim-1)
      for d = 1:(dim-1)
        elemtocoord[d, v, e] = basetopo.elemtocoord[d, v, i]
        elemtocoord[d, 2^(dim-1) + v, e] = basetopo.elemtocoord[d, v, i]
      end

      elemtocoord[dim, v, e] = stack[j]
      elemtocoord[dim, 2^(dim-1) + v, e] = stack[j+1]
    end
  end

  elemtoelem = similar(basetopo.elemtoelem, nface, length(elems))
  elemtoface = similar(basetopo.elemtoface, nface, length(elems))
  elemtoordr = similar(basetopo.elemtoordr, nface, length(elems))
  elemtobndy = similar(basetopo.elemtobndy, nface, length(elems))

  for e=1:length(basetopo.elems)*stacksize, f=1:nface
    elemtoelem[f, e] = e
    elemtoface[f, e] = f
    elemtoordr[f, e] = 1
    elemtobndy[f, e] = 0
  end

  for i=1:length(basetopo.realelems), j=1:stacksize
    e1 = stacksize*(i-1) + j

    for f = 1:2(dim-1)
      e2 = stacksize*(basetopo.elemtoelem[f, i]-1) + j

      elemtoelem[f, e1] = e2
      elemtoface[f, e1] = basetopo.elemtoface[f, i]

      # We assume a simple orientation right now
      @assert basetopo.elemtoordr[f, i] == 1
      elemtoordr[f, e1] = basetopo.elemtoordr[f, i]
    end

    et = stacksize*(i-1) + j + 1
    eb = stacksize*(i-1) + j - 1
    ft = 2(dim-1) + 1
    fb = 2(dim-1) + 2
    ot = 1
    ob = 1

    if j == stacksize
      et = periodicity[dim] ? stacksize*(i-1) + 1 : e1
      ft = periodicity[dim] ? ft : 2(dim-1) + 2
    end
    if j == 1
      eb = periodicity[dim] ? stacksize*(i-1) + stacksize : e1
      fb = periodicity[dim] ? fb : 2(dim-1) + 1
    end

    elemtoelem[2(dim-1)+1, e1] = eb
    elemtoelem[2(dim-1)+2, e1] = et
    elemtoface[2(dim-1)+1, e1] = fb
    elemtoface[2(dim-1)+2, e1] = ft
    elemtoordr[2(dim-1)+1, e1] = ob
    elemtoordr[2(dim-1)+2, e1] = ot
  end

  for i=1:length(basetopo.elems), j=1:stacksize
    e1 = stacksize*(i-1) + j

    for f = 1:2(dim-1)
      elemtobndy[f, e1] = basetopo.elemtobndy[f, i]
    end

    bt = bb = 0

    if j == stacksize
      bt = periodicity[dim] ? bt : boundary[dim][2]
    end
    if j == 1
      bb = periodicity[dim] ? bb : boundary[dim][1]
    end

    elemtobndy[2(dim-1)+1, e1] = bb
    elemtobndy[2(dim-1)+2, e1] = bt
  end

  nabrtorank = basetopo.nabrtorank
  nabrtorecv =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtorecv[n])-1)+1,
                             stacksize*last(basetopo.nabrtorecv[n]))
                   for n = 1:length(nabrtorank)]
  nabrtosend =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtosend[n])-1)+1,
                             stacksize*last(basetopo.nabrtosend[n]))
                   for n = 1:length(nabrtorank)]

  T = eltype(basetopo.elemtocoord)

  StackedBrickTopology{dim, T}(
    BoxElementTopology{dim, T}(
      mpicomm, elems, realelems, ghostelems, ghostfaces, sendelems, sendfaces,
      elemtocoord, elemtoelem, elemtoface, elemtoordr, elemtobndy,
      nabrtorank, nabrtorecv, nabrtosend, basetopo.origsendorder, !minimum(periodicity)),
    stacksize)
end

"""
    CubedShellTopology(mpicomm, Nelem, T) <: AbstractTopology{dim}

Generate a cubed shell mesh with the number of elements along each dimension of
the cubes being `Nelem`. This topology actual creates a cube mesh, and the
warping should be done after the grid is created using the `cubedshellwarp`
function. The coordinates of the points will be of type `T`.

The elements of the shell are partitioned equally across the MPI ranks based
on a space-filling curve.

Note that this topology is logically 2-D but embedded in a 3-D space

# Examples

We can build a cubed shell mesh with 10 elements on each cube, total elements is
`10 * 10 * 6 = 600`, with
```jldoctest brickmesh
using CLIMA.Topologies
using MPI
MPI.Init()
topology = CubedShellTopology(MPI.COMM_SELF, 10, Float64)

# Typically the warping would be done after the grid is created, but the cell
# corners could be warped with...

# Shell radius = 1
x1, x2, x3 = ntuple(j->topology.elemtocoord[j, :, :], 3)
for n = 1:length(x1)
   x1[n], x2[n], x3[n] = Topologies.cubedshellwarp(x1[n], x2[n], x3[n])
end

# Shell radius = 10
x1, x2, x3 = ntuple(j->topology.elemtocoord[j, :, :], 3)
for n = 1:length(x1)
  x1[n], x2[n], x3[n] = Topologies.cubedshellwarp(x1[n], x2[n], x3[n], 10)
end
```
"""
function CubedShellTopology(mpicomm, Neside, T; connectivity=:face,
                            ghostsize=1)

  # We cannot handle anything else right now...
  @assert connectivity == :face
  @assert ghostsize == 1

  mpirank = MPI.Comm_rank(mpicomm)
  mpisize = MPI.Comm_size(mpicomm)

  topology = cubedshellmesh(Neside, part=mpirank+1, numparts=mpisize)

  topology = BrickMesh.partition(mpicomm, topology...)#topology(1), topology(2), topology(3), topology(4) ) #topology...)
  origsendorder = topology[5]
  dim, nvert = 3, 4
  elemtovert = topology[1]
  nelem = size(elemtovert, 2)
  elemtocoord = Array{T}(undef, dim, nvert, nelem)
  ind2vert = CartesianIndices((Neside+1, Neside+1, Neside+1))
  for e in 1:nelem
    for n = 1:nvert
      v = elemtovert[n, e]
      i, j, k = Tuple(ind2vert[v])
      elemtocoord[:, n, e] = (2*[i-1, j-1, k-1] .- Neside) / Neside
    end
  end

  topology = BrickMesh.connectmesh(mpicomm, topology[1], elemtocoord, topology[3],
                                topology[4]; dim = 2)

  CubedShellTopology{T}(
    BoxElementTopology{2, T}(
      mpicomm, topology.elems, topology.realelems,
      topology.ghostelems, topology.ghostfaces, topology.sendelems,
      topology.sendfaces, topology.elemtocoord, topology.elemtoelem,
      topology.elemtoface, topology.elemtoordr, topology.elemtobndy,
      topology.nabrtorank, topology.nabrtorecv, topology.nabrtosend, origsendorder, false))
end

"""
    cubedshellmesh(T, Ne; part=1, numparts=1)

Generate a cubed mesh with each of the "cubes" has an `Ne X Ne` grid of
elements.

The mesh can optionally be partitioned into `numparts` and this returns
partition `part`.  This is a simple Cartesian partition and further partitioning
(e.g, based on a space-filling curve) should be done before the mesh is used for
computation.

This mesh returns the cubed spehere in a flatten fashion for the vertex values,
and a remapping is needed to embed the mesh in a 3-D space.

The mesh structures for the cubes is as follows:

```
x2
   ^
   |
4Ne-           +-------+
   |           |       |
   |           |   6   |
   |           |       |
3Ne-           +-------+
   |           |       |
   |           |   5   |
   |           |       |
2Ne-           +-------+
   |           |       |
   |           |   4   |
   |           |       |
 Ne-   +-------+-------+-------+
   |   |       |       |       |
   |   |   1   |   2   |   3   |
   |   |       |       |       |
  0-   +-------+-------+-------+
   |
   +---|-------|-------|------|-> x1
       0      Ne      2Ne    3Ne
```

"""
function cubedshellmesh(Ne; part=1, numparts=1)
  dim = 2
  @assert 1 <= part <= numparts

  globalnelems = 6 * Ne^2

  # How many vertices and faces per element
  nvert = 2^dim # 4
  nface = 2dim  # 4

  # linearly partition to figure out which elements we own
  elemlocal = BrickMesh.linearpartition(prod(globalnelems), part, numparts)

  # elemen to vertex maps which we own
  elemtovert = Array{Int}(undef, nvert, length(elemlocal))
  elemtocoord = Array{Int}(undef, dim, nvert, length(elemlocal))

  nelemcube = Ne^dim # Ne^2

  etoijb = CartesianIndices((Ne, Ne, 6))
  bx = [0 Ne 2Ne  Ne  Ne  Ne]
  by = [0  0   0  Ne 2Ne 3Ne]

  vertmap = LinearIndices((Ne+1, Ne+1, Ne+1))
  for (le, e) = enumerate(elemlocal)
    i, j, blck = Tuple(etoijb[e])
    elemtocoord[1, :, le] = bx[blck] .+ [i-1 i i-1 i]
    elemtocoord[2, :, le] = by[blck] .+ [j-1 j-1 j j]

    for n = 1:4
      ix = i + mod(n-1, 2)
      jx = j + div(n-1, 2)
      # set the vertices like they are the face vertices of a cube
      if blck == 1
        elemtovert[n, le] = vertmap[   1   , Ne+2-ix,      jx]
      elseif blck == 2
        elemtovert[n, le] = vertmap[     ix,   1    ,      jx]
      elseif blck == 3
        elemtovert[n, le] = vertmap[Ne+1   ,      ix,      jx]
      elseif blck == 4
        elemtovert[n, le] = vertmap[     ix,      jx,  Ne+1  ]
      elseif blck == 5
        elemtovert[n, le] = vertmap[     ix, Ne+1   , Ne+2-jx]
      elseif blck == 6
        elemtovert[n, le] = vertmap[     ix, Ne+2-jx,    1   ]
      end
    end
  end

  # no boundaries for a shell
  elemtobndy = zeros(Int, nface, length(elemlocal))

  # no faceconnections for a shell
  faceconnections = Array{Array{Int, 1}}(undef, 0)

  (elemtovert, elemtocoord, elemtobndy, faceconnections, collect(elemlocal))
end

"""
    cubedshellwarp(a, b, c, R = max(abs(a), abs(b), abs(c)))

Given points `(a, b, c)` on the surface of a cube, warp the points out to a
spherical shell of radius `R` based on the equiangular gnomonic grid proposed by
Ronchi, Iacono, Paolucci (1996) <https://doi.org/10.1006/jcph.1996.0047>

```
@article{RonchiIaconoPaolucci1996,
  title={The ``cubed sphere'': a new method for the solution of partial
         differential equations in spherical geometry},
  author={Ronchi, C. and Iacono, R. and Paolucci, P. S.},
  journal={Journal of Computational Physics},
  volume={124},
  number={1},
  pages={93--114},
  year={1996},
  doi={10.1006/jcph.1996.0047}
}
```

"""
function cubedshellwarp(a, b, c, R = max(abs(a), abs(b), abs(c)))

  function f(sR, ξ, η)
    X, Y = tan(π * ξ / 4), tan(π * η / 4)
    x1 = sR / sqrt(X^2 + Y^2 + 1)
    x2, x3 = X * x1, Y * x1
    x1,x2,x3
  end

  fdim = argmax(abs.((a, b, c)))
  if fdim == 1 && a < 0
    # (-R, *, *) : Face I from Ronchi, Iacono, Paolucci (1996)
    x1,x2,x3 = f(-R, b/a, c/a)
  elseif fdim == 2 && b < 0
    # ( *,-R, *) : Face II from Ronchi, Iacono, Paolucci (1996)
    x2,x1,x3 = f(-R, a/b, c/b)
  elseif fdim == 1 && a > 0
    # ( R, *, *) : Face III from Ronchi, Iacono, Paolucci (1996)
    x1,x2,x3 = f(R, b/a, c/a)
  elseif fdim == 2 && b > 0
    # ( *, R, *) : Face IV from Ronchi, Iacono, Paolucci (1996)
    x2,x1,x3 = f(R, a/b, c/b)
  elseif fdim == 3 && c > 0
    # ( *, *, R) : Face V from Ronchi, Iacono, Paolucci (1996)
    x3,x2,x1 = f(R, b/c, a/c)
  elseif fdim == 3 && c < 0
    # ( *, *,-R) : Face VI from Ronchi, Iacono, Paolucci (1996)
    x3,x2,x1 = f(-R, b/c, a/c)
  else
    error("invalid case for cubedshellwarp: $a, $b, $c")
  end

  return x1, x2, x3
end

"""
    cubedshellunwarp(x1, x2, x3)

The inverse of [`cubedshellwarp`](@ref).
"""
function cubedshellunwarp(x1,x2,x3)

  function g(R, X, Y)
    ξ = atan(X) * 4 / pi
    η = atan(Y) * 4 / pi
    R, R*ξ, R*η
  end

  R = hypot(x1,x2,x3)
  fdim = argmax(abs.((x1, x2, x3)))

  if fdim == 1 && x1 < 0
    # (-R, *, *) : Face I from Ronchi, Iacono, Paolucci (1996)
    a,b,c = g(-R, x2/x1, x3/x1)
  elseif fdim == 2 && x2 < 0
    # ( *,-R, *) : Face II from Ronchi, Iacono, Paolucci (1996)
    b,a,c = g(-R, x1/x2, x3/x2)
  elseif fdim == 1 && x1 > 0
    # ( R, *, *) : Face III from Ronchi, Iacono, Paolucci (1996)
    a,b,c = g(R, x2/x1, x3/x1)
  elseif fdim == 2 && x2 > 0
    # ( *, R, *) : Face IV from Ronchi, Iacono, Paolucci (1996)
    b,a,c = g(R, x1/x2, x3/x2)
  elseif fdim == 3 && x3 > 0
    # ( *, *, R) : Face V from Ronchi, Iacono, Paolucci (1996)
    c,b,a = g(R, x2/x3, x1/x3)
  elseif fdim == 3 && x3 < 0
    # ( *, *,-R) : Face VI from Ronchi, Iacono, Paolucci (1996)
    c,b,a = g(-R, x2/x3, x1/x3)
  else
    error("invalid case for cubedshellunwarp: $a, $b, $c")
  end

  return a,b,c
end

"""
   StackedCubedSphereTopology(mpicomm, Nhorz, Rrange;
                              boundary=(1,1)) <: AbstractTopology{3}

Generate a stacked cubed sphere topology with `Nhorz` by `Nhorz` cells for each
horizontal face and `Rrange` is the radius edges of the stacked elements.  This
topology actual creates a cube mesh, and the warping should be done after the
grid is created using the `cubedshellwarp` function. The coordinates of the
points will be of type `eltype(Rrange)`. The inner boundary condition type is
`boundary[1]` and the outer boundary condition type is `boundary[2]`.

The elements are stacked such that the vertical elements are contiguous in the
element ordering.

The elements of the brick are partitioned equally across the MPI ranks based
on a space-filling curve. Further, stacks are not split at MPI boundaries.

# Examples

We can build a cubed sphere mesh with 10 x 10 x 5 elements on each cube, total
elements is `10 * 10 * 5 * 6 = 3000`, with
```jldoctest brickmesh
using CLIMA.Topologies
using MPI
MPI.Init()
Nhorz = 10
Nstack = 5
Rrange = Float64.(accumulate(+,1:Nstack+1))
topology = StackedCubedSphereTopology(MPI.COMM_SELF, Nhorz, Rrange)

x1, x2, x3 = ntuple(j->reshape(topology.elemtocoord[j, :, :],
                            2, 2, 2, length(topology.elems)), 3)
for n = 1:length(x1)
   x1[n], x2[n], x3[n] = Topologies.cubedshellwarp(x1[n], x2[n], x3[n])
end
```
Note that the faces are listed in Cartesian order.
"""
function StackedCubedSphereTopology(mpicomm, Nhorz, Rrange; boundary = (1, 1),
                                    connectivity=:face, ghostsize=1)
  T = eltype(Rrange)

  basetopo = CubedShellTopology(mpicomm, Nhorz, T; connectivity=connectivity,
                                   ghostsize=ghostsize)

  dim = 3
  nvert = 2^dim
  nface = 2dim
  stacksize = length(Rrange) - 1

  nreal = length(basetopo.realelems)*stacksize
  nghost = length(basetopo.ghostelems)*stacksize

  elems=1:(nreal+nghost)
  realelems=1:nreal
  ghostelems=nreal.+(1:nghost)

  sendelems = similar(basetopo.sendelems,
                      length(basetopo.sendelems)*stacksize)

  for i=1:length(basetopo.sendelems), j=1:stacksize
    sendelems[stacksize*(i-1) + j] = stacksize*(basetopo.sendelems[i]-1) + j
  end

  ghostfaces = similar(basetopo.ghostfaces, nface, length(ghostelems))
  ghostfaces .= false

  for i=1:length(basetopo.ghostelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f = 1:2(dim-1)
      ghostfaces[f, e] = basetopo.ghostfaces[f, i]
    end
  end

  sendfaces = similar(basetopo.sendfaces, nface, length(sendelems))
  sendfaces .= false

  for i=1:length(basetopo.sendelems), j=1:stacksize
    e = stacksize*(i-1) + j
    for f=1:2(dim-1)
      sendfaces[f, e] = basetopo.sendfaces[f, i]
    end
  end

  elemtocoord = similar(basetopo.elemtocoord, dim, nvert, length(elems))

  for i=1:length(basetopo.elems), j=1:stacksize
    # i is base element
    # e is stacked element
    e = stacksize*(i-1) + j

    # v is base vertex
    for v = 1:2^(dim-1)
      for d = 1:dim # dim here since shell is embedded in 3-D
        # v is lower stacked vertex
        elemtocoord[d, v, e] = basetopo.elemtocoord[d, v, i] * Rrange[j]
        # 2^(dim-1) + v is higher stacked vertex
        elemtocoord[d, 2^(dim-1) + v, e] = basetopo.elemtocoord[d, v, i] *
                                           Rrange[j+1]
      end
    end
  end

  elemtoelem = similar(basetopo.elemtoelem, nface, length(elems))
  elemtoface = similar(basetopo.elemtoface, nface, length(elems))
  elemtoordr = similar(basetopo.elemtoordr, nface, length(elems))
  elemtobndy = similar(basetopo.elemtobndy, nface, length(elems))

  for e=1:length(basetopo.elems)*stacksize, f=1:nface
    elemtoelem[f, e] = e
    elemtoface[f, e] = f
    elemtoordr[f, e] = 1
    elemtobndy[f, e] = 0
  end

  for i=1:length(basetopo.realelems), j=1:stacksize
    e1 = stacksize*(i-1) + j

    for f = 1:2(dim-1)
      e2 = stacksize*(basetopo.elemtoelem[f, i]-1) + j

      elemtoelem[f, e1] = e2
      elemtoface[f, e1] = basetopo.elemtoface[f, i]

      # since the basetopo is 2-D we only need to worry about two orientations
      @assert basetopo.elemtoordr[f, i] ∈ (1,2)
      #=
      orientation 1:
      2---3     2---3
      |   | --> |   |
      0---1     0---1
      same:
      (a,b) --> (a,b)

      orientation 3:
      2---3     3---2
      |   | --> |   |
      0---1     1---0
      reverse first index:
      (a,b) --> (N+1-a,b)
      =#
      elemtoordr[f, e1] = basetopo.elemtoordr[f, i] == 1 ? 1 : 3
    end

    # If top or bottom of stack set neighbor to self on respective face
    elemtoelem[2(dim-1)+1, e1] = j == 1         ? e1 : stacksize*(i-1) + j - 1
    elemtoelem[2(dim-1)+2, e1] = j == stacksize ? e1 : stacksize*(i-1) + j + 1

    elemtoface[2(dim-1)+1, e1] = j == 1         ? 2(dim-1) + 1 : 2(dim-1) + 2
    elemtoface[2(dim-1)+2, e1] = j == stacksize ? 2(dim-1) + 2 : 2(dim-1) + 1

    elemtoordr[2(dim-1)+1, e1] = 1
    elemtoordr[2(dim-1)+2, e1] = 1
  end

  # Set the top and bottom boundary condition
  for i=1:length(basetopo.elems)
    eb = stacksize*(i-1) + 1
    et = stacksize*(i-1) + stacksize

    elemtobndy[2(dim-1)+1, eb] = boundary[1]
    elemtobndy[2(dim-1)+2, et] = boundary[2]
  end

  nabrtorank = basetopo.nabrtorank
  nabrtorecv =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtorecv[n])-1)+1,
                             stacksize*last(basetopo.nabrtorecv[n]))
                   for n = 1:length(nabrtorank)]
  nabrtosend =
    UnitRange{Int}[UnitRange(stacksize*(first(basetopo.nabrtosend[n])-1)+1,
                             stacksize*last(basetopo.nabrtosend[n]))
                   for n = 1:length(nabrtorank)]

  StackedCubedSphereTopology{T}(
    BoxElementTopology{3, T}(
      mpicomm, elems, realelems, ghostelems, ghostfaces, sendelems,
      sendfaces, elemtocoord, elemtoelem, elemtoface, elemtoordr, elemtobndy,
      nabrtorank, nabrtorecv, nabrtosend, basetopo.origsendorder, true),
    stacksize)
end

"""
    grid1d(a, b[, stretch::AbstractGridStretching]; elemsize, nelem)

Discretize the 1D interval [`a`,`b`] into elements.
Exactly one of the following keyword arguments must be provided:
- `elemsize`: the average element size, or
- `nelem`: the number of elements.

The optional `stretch` argument allows stretching, otherwise the element sizes will be uniform.

Returns either a range object or a vector containing the element boundaries.
"""
function grid1d(a, b, stretch=nothing; elemsize=nothing, nelem=nothing)
  xor(nelem === nothing, elemsize === nothing) || error("Either `elemsize` or `nelem` arguments must be provided")
  if elemsize !== nothing
    nelem = round(Int,abs(b-a)/elemsize)
  end
  grid1d(a, b, stretch, nelem)
end
function grid1d(a, b, ::Nothing, nelem)
  range(a, stop=b, length=nelem+1)
end

# TODO: document these
abstract type AbstractGridStretching end

"""
    SingleExponentialStretching(A)

Apply single-exponential stretching: `A > 0` will increase the density of points at the lower boundary, `A < 0` will increase the density at the upper boundary.

# Reference
* "Handbook of Grid Generation" J. F. Thompson, B. K. Soni, N. P. Weatherill (Editors) RCR Press 1999, §3.6.1 Single-Exponential Function
"""
struct SingleExponentialStretching{T} <: AbstractGridStretching
  A::T
end
function grid1d(a::A, b::B, stretch::SingleExponentialStretching, nelem) where {A,B}
  F = float(promote_type(A,B))
  s = range(zero(F), stop=one(F), length=nelem+1)
  a .+ (b-a) .* expm1.(stretch.A .* s) ./ expm1(stretch.A)
end

struct InteriorStretching{T} <: AbstractGridStretching
  attractor::T
end
function grid1d(a::A, b::B, stretch::InteriorStretching, nelem) where {A,B}
  F = float(promote_type(A,B))
  coe = F(2.5)
  s = range(zero(F), stop=one(F), length=nelem+1)
  range(a, stop=b, length=nelem+1) .+ coe .* (stretch.attractor .- (b-a).*s) .* (1 .-s) .* s
end

end

max_stars_repo_name: emassoud/CLIMA

max_issues_repo_name: emassoud/CLIMA

max_forks_repo_name: emassoud/CLIMA


text:
__precompile__()

module FixedPointNumbers

using Base: reducedim_initarray

import Base: ==, <, <=, -, +, *, /, ~, isapprox,
             convert, promote_rule, show, showcompact, isinteger, abs, decompose,
             isnan, isinf, isfinite,
             zero, one, typemin, typemax, realmin, realmax, eps, sizeof, reinterpret,
             float, trunc, round, floor, ceil, bswap,
             div, fld, rem, mod, mod1, rem1, fld1, min, max, minmax,
             start, next, done, r_promote, reducedim_init, rand

if VERSION <= v"0.5.0-dev+755"
    macro pure(ex)
        nothing
    end
else
    using Base: @pure
end

using Compat

# T => BaseType
# f => Number of Bytes reserved for fractional part
abstract FixedPoint{T <: Integer, f} <: Real

export
    FixedPoint,
    Fixed,
    UFixed,
    Fixed16,
    UFixed8,
    UFixed10,
    UFixed12,
    UFixed14,
    UFixed16,
    # literal constructor constants
    uf8,
    uf10,
    uf12,
    uf14,
    uf16,
    # Functions
    scaledual

reinterpret(x::FixedPoint) = x.i

# comparison
=={T <: FixedPoint}(x::T, y::T) = x.i == y.i
 <{T <: FixedPoint}(x::T, y::T) = x.i  < y.i
<={T <: FixedPoint}(x::T, y::T) = x.i <= y.i
"""
    isapprox(x::FixedPoint, y::FixedPoint; rtol=0, atol=max(eps(x), eps(y)))

For FixedPoint numbers, the default criterion is that `x` and `y` differ by no more than `eps`, the separation between adjacent fixed-point numbers.
"""
function isapprox{T<:FixedPoint}(x::T, y::T; rtol=0, atol=max(eps(x), eps(y)))
    maxdiff = T(atol+rtol*max(abs(x), abs(y)))
    rx, ry, rd = reinterpret(x), reinterpret(y), reinterpret(maxdiff)
    abs(signed(widen1(rx))-signed(widen1(ry))) <= rd
end
function isapprox(x::FixedPoint, y::FixedPoint; rtol=0, atol=max(eps(x), eps(y)))
    isapprox(promote(x, y)...; rtol=rtol, atol=atol)
end

# predicates
isinteger{T,f}(x::FixedPoint{T,f}) = (x.i&(1<<f-1)) == 0

# traits
typemax{T<: FixedPoint}(::Type{T}) = T(typemax(rawtype(T)), 0)
typemin{T<: FixedPoint}(::Type{T}) = T(typemin(rawtype(T)), 0)
realmin{T<: FixedPoint}(::Type{T}) = typemin(T)
realmax{T<: FixedPoint}(::Type{T}) = typemax(T)

widen1(::Type{Int8})   = Int16
widen1(::Type{UInt8})  = UInt16
widen1(::Type{Int16})  = Int32
widen1(::Type{UInt16}) = UInt32
widen1(::Type{Int32})  = Int64
widen1(::Type{UInt32}) = UInt64
widen1(::Type{Int64})  = Int128
widen1(::Type{UInt64}) = UInt128
widen1(x::Integer) = x % widen1(typeof(x))

typealias ShortInts Union{Int8,UInt8,Int16,UInt16}

floattype{T<:ShortInts,f}(::Type{FixedPoint{T,f}}) = Float32
floattype{T,f}(::Type{FixedPoint{T,f}}) = Float64
floattype(x::FixedPoint) = floattype(supertype(typeof(x)))

include("fixed.jl")
include("ufixed.jl")
include("deprecations.jl")

eps{T<:FixedPoint}(::Type{T}) = T(one(rawtype(T)),0)
eps{T<:FixedPoint}(::T) = eps(T)
sizeof{T<:FixedPoint}(::Type{T}) = sizeof(rawtype(T))

# Promotions for reductions
const Treduce = Float64
r_promote{T}(::typeof(@functorize(+)), x::FixedPoint{T}) = Treduce(x)
r_promote{T}(::typeof(@functorize(*)), x::FixedPoint{T}) = Treduce(x)

reducedim_init{T<:FixedPoint}(f::typeof(@functorize(identity)),
                              op::typeof(@functorize(+)),
                              A::AbstractArray{T}, region) =
    reducedim_initarray(A, region, zero(Treduce))
reducedim_init{T<:FixedPoint}(f::typeof(@functorize(identity)),
                              op::typeof(@functorize(*)),
                              A::AbstractArray{T}, region) =
    reducedim_initarray(A, region, one(Treduce))

# TODO: rewrite this by @generated
for T in tuple(Fixed16, UF...)
    R = rawtype(T)
    @eval begin
        reinterpret(::Type{$R}, x::$T) = x.i
    end
end

for f in (:div, :fld, :fld1)
    @eval begin
        $f{T<:FixedPoint}(x::T, y::T) = $f(reinterpret(x),reinterpret(y))
    end
end
for f in (:rem, :mod, :mod1, :rem1, :min, :max)
    @eval begin
        $f{T<:FixedPoint}(x::T, y::T) = T($f(reinterpret(x),reinterpret(y)),0)
    end
end

# When multiplying by a float, reduce two multiplies to one.
# Particularly useful for arrays.
scaledual(Tdual::Type, x) = one(Tdual), x
scaledual{Tdual<:Number}(b::Tdual, x) = b, x
scaledual{T<:FixedPoint}(Tdual::Type, x::Union{T,AbstractArray{T}}) =
    convert(Tdual, 1/one(T)), reinterpret(rawtype(T), x)
scaledual{Tdual<:Number, T<:FixedPoint}(b::Tdual, x::Union{T,AbstractArray{T}}) =
    convert(Tdual, b/one(T)), reinterpret(rawtype(T), x)

# printing
function show{T,f}(io::IO, x::FixedPoint{T,f})
    shorttype = typeof(x)<:UFixed ? "UFixed" : "Fixed"
    print(io, shorttype, "{", T, ",", f, "}")
    print(io, "(")
    showcompact(io, x)
    print(io, ")")
end
const _log2_10 = 3.321928094887362
showcompact{T,f}(io::IO, x::FixedPoint{T,f}) = show(io, round(Float64(x), ceil(Int,f/_log2_10)))

@noinline function throw_converterror{T<:FixedPoint}(::Type{T}, x)
    n = 2^(8*sizeof(T))
    bitstring = sizeof(T) == 1 ? "an 8-bit" : "a $(8*sizeof(T))-bit"
    io = IOBuffer()
    showcompact(io, typemin(T)); Tmin = takebuf_string(io)
    showcompact(io, typemax(T)); Tmax = takebuf_string(io)
    throw(ArgumentError("$T is $bitstring type representing $n values from $Tmin to $Tmax; cannot represent $x"))
end

rand{T<:FixedPoint}(::Type{T}) = reinterpret(T, rand(rawtype(T)))
rand{T<:FixedPoint}(::Type{T}, sz::Dims) = reinterpret(T, rand(rawtype(T), sz))

end # module


text:
"""
    AdaptiveMeanPool(target_out)
Adaptive mean pooling layer. 'target_out' stands for the size of one layer(channel) of output.
Irrespective of Input size and pooling window size, it returns the target output dimensions.
"""
struct AdaptiveMeanPool{N}
  target_out::NTuple{N, Int}
end

function (m::AdaptiveMeanPool)(x)
  w = size(x, 1) - m.target_out[1] + 1
  h = size(x, 2) - m.target_out[2] + 1
  pdims = PoolDims(x, (w, h); padding = (0, 0), stride = (1, 1))
  return meanpool(x, pdims)
end

"""
    AdaptiveMaxPool(target_out)
Adaptive max pooling layer. 'target_out' stands for the size of one layer(channel) of output.
Irrespective of Input size and pooling window size, it returns the target output dimensions.
"""
struct AdaptiveMaxPool{N}
  target_out::NTuple{N, Int}
end

function (m::AdaptiveMaxPool)(x)
  w = size(x, 1) - m.target_out[1] + 1
  h = size(x, 2) - m.target_out[2] + 1
  pdims = PoolDims(x, (w, h); padding = (0, 0), stride = (1, 1))
  return maxpool(x, pdims)
end

"""
  flatten(x)

Flatten `x` in all dimensions except the last/batch dimension
"""
flatten(x) = reshape(x, :, size(x)[end])

"""
    Residual(model) == SkipConnection(model, +)

Helper for defining residual connections
"""
Residual(model) = SkipConnection(model, +)

function pixelshuffle(x::AbstractArray{T, 4}, ratios = (2, 2)) where T
    iH, iW, C, B = size(x)
    pH, pW = ratios
    oC, oH, oW = C ÷ (pH*pW), iH * pH, iW * pW

    x = reshape(x, iH, iW, pH, pW, oC, B)
    x = permutedims(x, (4, 1, 3, 2, 5, 6))  # pH, iH, pW, iW, oC, B
    x = reshape(x, oH, oW, oC, B)

    return x
end

"""
    PixelShuffle(scale, kernels_in[, kernels_out])

Pixel shuffle layer that upscales height and width of `x` by `scale`. Has reduced
checkerboard artifacts compared to `ConvTranspose`

Introduced in [Real-Time Single Image and Video Super-Resolution Using
an EfficientSub-Pixel Convolutional Neural Network](https://arxiv.org/pdf/1609.05158.pdf).

"""
struct PixelShuffle
    conv::Conv
    scales::Tuple{Int, Int}
end

Flux.@functor PixelShuffle

function PixelShuffle(scales::Tuple{Int, Int}, k_in, k_out = k_in)
    return PixelShuffle(
        Conv((1, 1), k_in => k_out * scales[1] * scales[2]),
        scales,
    )
end

PixelShuffle(scale::Int, k_in, k_out = k_in) = PixelShuffle((scale, scale), k_in, k_out)

function (ps::PixelShuffle)(x)
    pixelshuffle(ps.conv(x), ps.scales)
end


text:
function add_default_ocean_tides_readers(::Type{GravityFieldFactory})
    return jcall(GravityFieldFactory, "addDefaultOceanTidesReaders", void, ())
end

function add_default_potential_coefficients_readers(::Type{GravityFieldFactory})
    return jcall(GravityFieldFactory, "addDefaultPotentialCoefficientsReaders", void, ())
end

function add_ocean_tides_reader(::Type{GravityFieldFactory}, arg0::OceanTidesReader)
    return jcall(GravityFieldFactory, "addOceanTidesReader", void, (OceanTidesReader,), arg0)
end

function add_potential_coefficients_reader(::Type{GravityFieldFactory}, arg0::PotentialCoefficientsReader)
    return jcall(GravityFieldFactory, "addPotentialCoefficientsReader", void, (PotentialCoefficientsReader,), arg0)
end

function clear_ocean_tides_readers(::Type{GravityFieldFactory})
    return jcall(GravityFieldFactory, "clearOceanTidesReaders", void, ())
end

function clear_potential_coefficients_readers(::Type{GravityFieldFactory})
    return jcall(GravityFieldFactory, "clearPotentialCoefficientsReaders", void, ())
end

function configure_ocean_load_deformation_coefficients(::Type{GravityFieldFactory}, arg0::OceanLoadDeformationCoefficients)
    return jcall(GravityFieldFactory, "configureOceanLoadDeformationCoefficients", void, (OceanLoadDeformationCoefficients,), arg0)
end

function get_constant_normalized_provider(::Type{GravityFieldFactory}, arg0::jint, arg1::jint)
    return jcall(GravityFieldFactory, "getConstantNormalizedProvider", NormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end

function get_constant_unnormalized_provider(::Type{GravityFieldFactory}, arg0::jint, arg1::jint)
    return jcall(GravityFieldFactory, "getConstantUnnormalizedProvider", UnnormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end

function get_gravity_fields(::Type{GravityFieldFactory})
    return jcall(GravityFieldFactory, "getGravityFields", LazyLoadedGravityFields, ())
end

function get_normalized_provider(::Type{GravityFieldFactory}, arg0::UnnormalizedSphericalHarmonicsProvider)
    return jcall(GravityFieldFactory, "getNormalizedProvider", NormalizedSphericalHarmonicsProvider, (UnnormalizedSphericalHarmonicsProvider,), arg0)
end

function get_normalized_provider(::Type{GravityFieldFactory}, arg0::jdouble, arg1::jdouble, arg2::TideSystem, arg3::Vector{Vector{jdouble}}, arg4::Vector{Vector{jdouble}})
    return jcall(GravityFieldFactory, "getNormalizedProvider", NormalizedSphericalHarmonicsProvider, (jdouble, jdouble, TideSystem, Vector{Vector{jdouble}}, Vector{Vector{jdouble}}), arg0, arg1, arg2, arg3, arg4)
end

function get_normalized_provider(::Type{GravityFieldFactory}, arg0::jint, arg1::jint)
    return jcall(GravityFieldFactory, "getNormalizedProvider", NormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end

function get_ocean_load_deformation_coefficients(::Type{GravityFieldFactory})
    return jcall(GravityFieldFactory, "getOceanLoadDeformationCoefficients", OceanLoadDeformationCoefficients, ())
end

function get_ocean_tides_waves(::Type{GravityFieldFactory}, arg0::jint, arg1::jint)
    return jcall(GravityFieldFactory, "getOceanTidesWaves", List, (jint, jint), arg0, arg1)
end

function get_unnormalization_factors(::Type{GravityFieldFactory}, arg0::jint, arg1::jint)
    return jcall(GravityFieldFactory, "getUnnormalizationFactors", Vector{Vector{jdouble}}, (jint, jint), arg0, arg1)
end

function get_unnormalized_provider(::Type{GravityFieldFactory}, arg0::NormalizedSphericalHarmonicsProvider)
    return jcall(GravityFieldFactory, "getUnnormalizedProvider", UnnormalizedSphericalHarmonicsProvider, (NormalizedSphericalHarmonicsProvider,), arg0)
end

function get_unnormalized_provider(::Type{GravityFieldFactory}, arg0::jdouble, arg1::jdouble, arg2::TideSystem, arg3::Vector{Vector{jdouble}}, arg4::Vector{Vector{jdouble}})
    return jcall(GravityFieldFactory, "getUnnormalizedProvider", UnnormalizedSphericalHarmonicsProvider, (jdouble, jdouble, TideSystem, Vector{Vector{jdouble}}, Vector{Vector{jdouble}}), arg0, arg1, arg2, arg3, arg4)
end

function get_unnormalized_provider(::Type{GravityFieldFactory}, arg0::jint, arg1::jint)
    return jcall(GravityFieldFactory, "getUnnormalizedProvider", UnnormalizedSphericalHarmonicsProvider, (jint, jint), arg0, arg1)
end

function read_gravity_field(::Type{GravityFieldFactory}, arg0::jint, arg1::jint)
    return jcall(GravityFieldFactory, "readGravityField", PotentialCoefficientsReader, (jint, jint), arg0, arg1)
end


text:
using JuMP
import Ipopt; IpoptOptimizer = Ipopt.Optimizer

@with_kw struct MMDNumerical <: AbstractMMD
    positivity::Bool=true
    normalisation::Bool=true
end

function _estimate_ratio(mmd::MMDNumerical, Kdede, Kdenu)
    n_de, n_nu = size(Kdenu)
    model = Model(with_optimizer(IpoptOptimizer; print_level=0))
    @variable(model, r[1:n_de])
    @objective(model, Min, 1 / n_de ^ 2 * sum(r[i] * Kdede[i,j] * r[j] for i = 1:n_de, j=1:n_de) - 2 / (n_de * n_nu) * sum(r[i] * Kdenu[i,j] for i = 1:n_de, j=1:n_nu))
    if mmd.positivity
        @constraint(model, r .>= 0)
    end
    if mmd.normalisation
        @constraint(model, 1 / n_de * sum(r) == 1)
    end
    optimize!(model)
    return value.(r)
end


text:
export  HermiteMapComponent,
        αreg,
        EmptyHermiteMapComponent,
        getbasis,
        ncoeff,
        getcoeff,
        setcoeff!,
        clearcoeff!,
        getidx,
        active_dim,
        evaluate!,
        evaluate,
        log_pdf!,
        log_pdf,
        grad_x_log_pdf!,
        grad_x_log_pdf,
        hess_x_log_pdf!,
        hess_x_log_pdf,
        reduced_hess_x_log_pdf!,
        reduced_hess_x_log_pdf,
        mean_hess_x_log_pdf!,
        mean_hess_x_log_pdf,
        negative_log_likelihood!,
        negative_log_likelihood,
        hess_negative_log_likelihood!

"""
$(TYPEDEF)

## Fields

$(TYPEDFIELDS)
"""
struct HermiteMapComponent
    m::Int64
    Nψ::Int64
    Nx::Int64
    # IntegratedFunction
    I::IntegratedFunction
    # Regularization parameter
    α::Float64
end

# Define default regularization
const αreg = 1.0e-6

function HermiteMapComponent(I::IntegratedFunction; α::Float64=αreg)
    return HermiteMapComponent(I.m, I.Nψ, I.Nx, I, α)
end

function HermiteMapComponent(m::Int64, Nx::Int64, idx::Array{Int64,2}, coeff::Array{Float64,1}; α::Float64 = αreg, b::String="CstProHermiteBasis")
    Nψ = size(coeff,1)
    @assert size(coeff,1) == size(idx,1) "Wrong dimension"
    if b ∈ ["ProHermiteBasis"; "PhyHermiteBasis";
            "CstProHermiteBasis"; "CstPhyHermiteBasis";
            "CstLinProHermiteBasis"; "CstLinPhyHermiteBasis"]
        B = MultiBasis(eval(Symbol(b))(m), Nx)
    else
        error("The basis "*b*" is not defined.")
    end

    return HermiteMapComponent(m, Nψ, Nx, IntegratedFunction(ExpandedFunction(B, idx, coeff)), α)
end

function HermiteMapComponent(f::ExpandedFunction; α::Float64 = αreg)
    return HermiteMapComponent(f.m, f.Nψ, f.Nx, IntegratedFunction(f), α)
end

function HermiteMapComponent(m::Int64, Nx::Int64; α::Float64 = αreg, b::String="CstProHermiteBasis")
    Nψ = 1

    # m is the dimension of the basis
    if b ∈ ["ProHermiteBasis"; "PhyHermiteBasis";
            "CstProHermiteBasis"; "CstPhyHermiteBasis";
            "CstLinProHermiteBasis"; "CstLinPhyHermiteBasis"]
        B = MultiBasis(eval(Symbol(b))(m), Nx)
    else
        error("The basis "*b*" is not defined.")
    end

    idx = zeros(Int64, Nψ, Nx)
    coeff = zeros(Nψ)

    f = ExpandedFunction(B, idx, coeff)
    I = IntegratedFunction(f)
    return HermiteMapComponent(I; α = α)
end

function Base.show(io::IO, C::HermiteMapComponent)
    println(io,"Hermite map component of dimension "*string(C.Nx)*" with Nψ = "*string(C.Nψ)*" active features")
    # for i=1:B.m
    #     println(io, B[i])
    # end
end

"""
$(TYPEDSIGNATURES)

Returns the number of features of the `HermiteMapComponent` `C` to `coeff`.
"""
ncoeff(C::HermiteMapComponent) = C.Nψ

"""
$(TYPEDSIGNATURES)

Returns the coefficients of the `HermiteMapComponent` `C` to `coeff`.
"""
getcoeff(C::HermiteMapComponent)= C.I.f.coeff

"""
$(TYPEDSIGNATURES)

Set the coefficients of the `HermiteMapComponent` `C` to `coeff`.
"""
function setcoeff!(C::HermiteMapComponent, coeff::Array{Float64,1})
        @assert size(coeff,1) == C.Nψ "Wrong dimension of coeff"
        C.I.f.coeff .= coeff
end

"""
$(TYPEDSIGNATURES)

Set all the coefficients of the `HermiteMapComponent` `C` to zero.
"""
clearcoeff!(C::HermiteMapComponent) = fill!(C.I.f.coeff, 0.0)

"""
$(TYPEDSIGNATURES)

Returns the multi-indices of the features of the `HermiteMapComponent` `C`.
"""
getidx(C::HermiteMapComponent) = C.I.f.idx

"""
$(TYPEDSIGNATURES)

Returns the kind of basis of the `HermiteMapComponent` `C`.
"""
getbasis(C::HermiteMapComponent) = getbasis(C.I)

"""
$(TYPEDSIGNATURES)

Returns the active dimensions of the `HermiteMapComponent` `C`.
"""
active_dim(C::HermiteMapComponent) = C.I.f.dim

## Evaluate

"""
$(TYPEDSIGNATURES)

Evaluates in-place the `HermiteMapComponent` `C` for the ensemble matrix `X`.
"""
function evaluate!(out, C::HermiteMapComponent, X)
    @assert C.Nx==size(X,1) "Wrong dimension of the sample"
    @assert size(out,1) == size(X,2) "Dimensions of the output and the samples don't match"
    return evaluate!(out, C.I, X)
end

"""
$(TYPEDSIGNATURES)

Evaluates the `HermiteMapComponent` `C` for the ensemble matrix `X`.
"""
evaluate(C::HermiteMapComponent, X::Array{Float64,2}) =
    evaluate!(zeros(size(X,2)), C, X)

"""
$(TYPEDSIGNATURES)

Evaluates the `HermiteMapComponent` `C` at `x`.
"""
(C::HermiteMapComponent)(x::Array{Float64,1}) = evaluate(C, reshape(x, (size(x,1), 1)))

## Compute log_pdf

"""
$(TYPEDSIGNATURES)

Evaluates in-place the logarithm of the pullback of the reference density (standard normal) by the `HermiteMapComponent` `C`for the ensemble matrix `X`.
"""
function log_pdf!(result, cache, C::HermiteMapComponent, X)
    NxX, Ne = size(X)
    @assert C.Nx == NxX "Wrong dimension of the sample"
    @assert size(result, 1) == Ne "Wrong dimension of the sample"
    @assert size(cache, 1) == Ne "Wrong dimension of the sample"

    evaluate!(result, C, X)
    cache .= grad_xd(C.I, X)

    @avx for i=1:size(X,2)
        result[i] = log_pdf(result[i]) + log(cache[i])
    end

    return result
end

"""
$(TYPEDSIGNATURES)

Evaluates the logarithm of the pullback of the reference density (standard normal) by the `HermiteMapComponent` `C`for the ensemble matrix `X`.
"""
log_pdf(C::HermiteMapComponent, X) = log_pdf!(zeros(size(X,2)), zeros(size(X,2)), C, X)

## Compute grad_x_log_pdf

"""
$(TYPEDSIGNATURES)

Evaluates in-place the gradient of the logarithm of the pullback of the reference density (standard normal) by the `HermiteMapComponent` `C` for the ensemble matrix `X`.
"""
function grad_x_log_pdf!(result, cache, C::HermiteMapComponent, X)
    NxX, Ne = size(X)
    @assert C.Nx == NxX "Wrong dimension of the sample"
    @assert size(result) == (Ne, NxX) "Wrong dimension of the result"

    # Compute gradient of log η∘C(x_{1:k})
    evaluate!(cache, C, X)
    grad_x!(result, C.I, X)
    @avx @. result *= cache
    rmul!(result, -1.0)

    # Compute gradient of log ∂k C(x_{1:k})
    cache .= grad_xd(C.I.f, X)
    grad_x_logeval!(cache, C.I.g, cache)
    result += cache .* grad_x_grad_xd(C.I.f, X)
    return result
end

"""
$(TYPEDSIGNATURES)

Evaluates in-place the gradient of the logarithm of the pullback of the reference density (standard normal) by the `HermiteMapComponent` `C` for the ensemble matrix `X`.
"""
grad_x_log_pdf(C::HermiteMapComponent, X) = grad_x_log_pdf!(zeros(size(X,2), size(X,1)), zeros(size(X,2)), C, X)

## Compute hess_x_log_pdf

"""
$(TYPEDSIGNATURES)

Evaluates in-place the hessian of the logarithm of the pullback of the reference density (standard normal) by the `HermiteMapComponent` `C` for the ensemble matrix `X`.
"""
function hess_x_log_pdf!(result, dcache, cache, C::HermiteMapComponent, X)
    NxX, Ne = size(X)
    Nx = C.Nx
    @assert Nx == NxX "Wrong dimension of the sample"
    @assert size(result) == (Ne, NxX, NxX) "Wrong dimension of the result"

    # Compute hessian of log η∘C(x_{1:k}) with η the log pdf of N(O, I_n)
    evaluate!(cache, C, X)
    grad_x!(dcache, C.I, X)
    hess_x!(result, C.I, X)

    dim = active_dim(C)

    @inbounds for i=1:length(dim)
                for j=i:length(dim)
                dcachei = view(dcache,:,dim[i])
                dcachej = view(dcache,:,dim[j])
                resultij = view(result,:,dim[i],dim[j])
                resultji = view(result,:,dim[j], dim[i])
                @avx @. resultij = resultij * cache + dcachei * dcachej
                resultji .= resultij
        end
    end

    rmul!(result, -1.0)

    # Compute hessian of log ∂k C(x_{1:k})
    cache .= grad_xd(C.I.f, X)
    cached2log = vhess_x_logeval(C.I.g, cache)
    dcache .= grad_x_grad_xd(C.I.f, X)

    @inbounds for i=1:length(dim)
                for j=i:length(dim)
                dcachei = view(dcache,:,dim[i])
                dcachej = view(dcache,:,dim[j])
                resultij = view(result,:,dim[i],dim[j])
                resultji = view(result,:,dim[j], dim[i])
                @avx @. resultij += dcachei * dcachej * cached2log

                resultji .= resultij
        end
    end
    #
    grad_x_logeval!(cache, C.I.g, cache)
    result .+= hess_x_grad_xd(C.I.f, X) .* cache

    return result
end

"""
$(TYPEDSIGNATURES)

Evaluates the hessian of the logarithm of the pullback of the reference density (standard normal) by the `HermiteMapComponent` `C` for the ensemble matrix `X`.
"""
hess_x_log_pdf(C::HermiteMapComponent, X) = hess_x_log_pdf!(zeros(size(X,2), size(X,1), size(X,1)),
                                                     zeros(size(X,2), size(X,1)),
                                                     zeros(size(X,2)), C, X)

# This version outputs a result of size(Ne, active_dim(C), active_dim(C))

"""
$(TYPEDSIGNATURES)

Evaluates in-place the hessian of the logarithm of the pullback of the reference density (standard normal) by the map `M` for the ensemble matrix `X`.
This routine exploits the sparsity pattern of the map `M`.
"""
function reduced_hess_x_log_pdf!(result, dcache, cache, C::HermiteMapComponent, X)
    NxX, Ne = size(X)
    Nx = C.Nx

    dim = active_dim(C)
    dimoff = dim[dim .< Nx]

    @assert Nx == NxX "Wrong dimension of the sample"
    @assert size(result) == (Ne, length(dim), length(dim)) "Wrong dimension of the result"

    # Compute hessian of log η∘C(x_{1:k}) with η the log pdf of N(O, I_n)
    evaluate!(cache, C, X)
    reduced_grad_x!(dcache, C.I, X)
    reduced_hess_x!(result, C.I, X)

    dim = active_dim(C)

    @inbounds for i=1:length(dim)
        for j=i:length(dim)
             # dcachei = view(dcache,:,dim[i])
             # dcachej = view(dcache,:,dim[j])
             # resultij = view(result,:,dim[i],dim[j])
             # resultji = view(result,:,dim[j], dim[i])
             dcachei = view(dcache,:,i)
             dcachej = view(dcache,:,j)
             resultij = view(result,:,i,j)
             resultji = view(result,:,j,i)
             @. resultij = resultij * cache + dcachei * dcachej
             resultji .= resultij
         end
    end

    rmul!(result, -1.0)

    # Compute hessian of log ∂k C(x_{1:k})
    cache .= grad_xd(C.I.f, X)
    cached2log = vhess_x_logeval(C.I.g, cache)

    if Nx == dim[end] # check if the last component is an active dimension
        # dcache .= grad_x_grad_xd(C.I.f, X)
        # Clear dcache
        fill!(dcache, 0.0)
        reduced_grad_x_grad_xd!(dcache, C.I.f, X)
        @inbounds for i=1:length(dim)
             for j=i:length(dim)
                 # dcachei = view(dcache,:,dim[i])
                 # dcachej = view(dcache,:,dim[j])
                 # resultij = view(result,:,dim[i],dim[j])
                 # resultji = view(result,:,dim[j], dim[i])
                 dcachei = view(dcache,:,i)
                 dcachej = view(dcache,:,j)
                 resultij = view(result,:,i,j)
                 resultji = view(result,:,j, i)
                 @avx @. resultij += dcachei * dcachej * cached2log

                 resultji .= resultij
             end
         end
    end
    grad_x_logeval!(cache, C.I.g, cache)
    result .+= reduced_hess_x_grad_xd(C.I.f, X) .* cache

    return result
end

reduced_hess_x_log_pdf(C::HermiteMapComponent, X) = reduced_hess_x_log_pdf!(zeros(size(X,2), length(active_dim(C)), length(active_dim(C))),
                                                zeros(size(X,2), length(active_dim(C))), zeros(size(X,2)), C, X)

## negative_log_likelihood

# function negative_log_likelihood!(J, dJ, coeff, S::Storage{m, Nψ, k}, C::HermiteMapComponent{m, Nψ, k}, X::Array{Float64,2}) where {m, Nψ, k}

function negative_log_likelihood!(J, dJ, coeff, S::Storage, C::HermiteMapComponent, X)
    NxX, Ne = size(X)
    m = C.m
    Nx = C.Nx
    Nψ = C.Nψ
    @assert NxX == Nx "Wrong dimension of the sample X"
    @assert size(S.ψoff, 1) == Ne
    @assert size(S.ψoff, 2) == Nψ

    # Define δ (regularization term) - add small diagonal term
    δ = 1e-9

    # Output objective, gradient
    xlast = view(X,NxX,:)

    fill!(S.cache_integral, 0.0)

    # Integrate at the same time for the objective, gradient
    function integrand!(v::Vector{Float64}, t::Float64)
        repeated_grad_xk_basis!(S.cache_dcψxdt, S.cache_gradxd, C.I.f, t*xlast)

         # This computing is also reused in the computation of the gradient, no interest to skip it
        @avx @. S.cache_dcψxdt *= S.ψoff
        mul!(S.cache_dψxd, S.cache_dcψxdt, coeff)

        # This doesn't work here, if we sue this line we need to add ψoff[i,j] in the last multiplication
        # @avx @. S.cache_dψxd = (S.cache_dcψxdt * S.ψoff) *ˡ coeff

        # Integration for J
        vJ = view(v,1:Ne)
        evaluate!(vJ, C.I.g, S.cache_dψxd)

        # Integration for dcJ

        grad_x!(S.cache_dψxd, C.I.g, S.cache_dψxd)

        # @avx for j=2:Nψ+1
        #     for i=1:Ne
        #         v[(j-1)*Ne+i] = S.cache_dψxd[i]*S.cache_dcψxdt[i,j-1]
        #     end
        # end
        # The reshape version is faster than unrolling
        v[Ne+1:Ne+Ne*Nψ] .= reshape(S.cache_dψxd .* S.cache_dcψxdt , (Ne*Nψ))
        # v[Ne+1:Ne+Ne*Nψ] .= reshape(grad_x(C.I.g, S.cache_dψxd) .* S.cache_dcψxdt , (Ne*Nψ))
    end

    quadgk!(integrand!, S.cache_integral, 0.0, 1.0)#; rtol = 1e-3)#, order = 3)#; order = 9, rtol = 1e-10)

    # Multiply integral by xlast (change of variable in the integration)
    # @avx for j=1:Nψ+1
    #     for i=1:Ne
    #         S.cache_integral[(j-1)*Ne+i] *= xlast[i]
    #     end
    # end

    # Multiply integral by xlast (change of variable in the integration)
    @inbounds for j=1:Nψ+1
        @. S.cache_integral[(j-1)*Ne+1:j*Ne] *= xlast
    end

    # Add f(x_{1:d-1},0) i.e. (S.ψoff .* S.ψd0)*coeff to S.cache_integral
    @avx for i=1:Ne
        f0i = zero(Float64)
        for j=1:Nψ
            f0i += S.ψoffψd0[i,j]*coeff[j]
        end
        # Add also the regularization term
        S.cache_integral[i] += f0i + δ*xlast[i]

    end
    # @show δ

    # Store g(∂_{xk}f(x_{1:k})) in S.cache_g

    @avx for i=1:Ne
        prelogJi = zero(Float64)
        for j=1:Nψ
            prelogJi += S.ψoffdψxd[i,j]*coeff[j]
        end
        # Add also the regularization term
        S.cache_g[i] = prelogJi + δ
    end

    # Formatting for Optim.jl
    if dJ != nothing
        reshape_cacheintegral = reshape(S.cache_integral[Ne+1:end], (Ne, Nψ))
        fill!(dJ, 0.0)
        @inbounds for i=1:Ne
            for j=1:Nψ
            dJ[j] += gradlog_pdf(S.cache_integral[i])*(reshape_cacheintegral[i,j] + S.ψoffψd0[i,j]) + # dsoftplus(S.cache_g[i])*S.ψoff[i,j]*S.dψxd[i,j]*(1/softplus(S.cache_g[i]))
                     grad_x(C.I.g, S.cache_g[i])*S.ψoffdψxd[i,j]/C.I.g(S.cache_g[i])
            end
        end
        rmul!(dJ, -1/Ne)
        # Add derivative of the L2 penalty term ∂_c α ||c||^2 = 2 *α c
        dJ .+= 2*C.α*coeff
    end

    if J != nothing
        J = 0.0
        @avx for i=1:Ne
            J += log_pdf(S.cache_integral[i]) + log(C.I.g(S.cache_g[i]))
        end
        J *=(-1/Ne)
        J += C.α*norm(coeff)^2
        return J
    end
end

negative_log_likelihood(S::Storage, C::HermiteMapComponent, X) = (J, dJ, coeff) -> negative_log_likelihood!(J, dJ, coeff, S, C, X)

function hess_negative_log_likelihood!(J, dJ, d2J, coeff, S::Storage, C::HermiteMapComponent, X::Array{Float64,2})
    Nψ = C.Nψ
    Nx = C.Nx

    NxX, Ne = size(X)
    @assert NxX == Nx "Wrong dimension of the sample X"
    @assert size(S.ψoff, 1) == Ne
    @assert size(S.ψoff, 2) == Nψ

    # Output objective, gradient
    xlast = view(X, NxX, :)#)

    fill!(S.cache_integral, 0)

    dcdψouter = zeros(Ne, Nψ, Nψ)

    # Integrate at the same time for the objective, gradient
    function integrand!(v::Vector{Float64}, t::Float64)
        S.cache_dcψxdt .= repeated_grad_xk_basis(C.I.f, t*xlast)

        # @avx @. S.cache_dψxd = (S.cache_dcψxdt .* S.ψoff) *ˡ coeff
        S.cache_dcψxdt .*= S.ψoff
        mul!(S.cache_dψxd, S.cache_dcψxdt, coeff)

        @inbounds for i=1:Nψ
            for j=1:Nψ
                dcdψouter[:,i,j] = S.cache_dcψxdt[:,i] .* S.cache_dcψxdt[:, j]
            end
        end

        # Integration for J
        vJ = view(v,1:Ne)
        evaluate!(vJ, C.I.g, S.cache_dψxd)
        # v[1:Ne] .= C.I.g(S.cache_dψxd)

        # Integration for dcJ
        v[Ne+1:Ne+Ne*Nψ] .= reshape(grad_x(C.I.g, S.cache_dψxd) .* S.cache_dcψxdt , (Ne*Nψ))

        # Integration for d2cJ
        v[Ne + Ne*Nψ + 1: Ne + Ne*Nψ + Ne*Nψ*Nψ] .= reshape(hess_x(C.I.g, S.cache_dψxd) .* dcdψouter, (Ne*Nψ*Nψ))

    end

    quadgk!(integrand!, S.cache_integral, 0.0, 1.0)#; rtol = 1e-3)#; order = 9, rtol = 1e-10)

    # Multiply integral by xlast (change of variable in the integration)
    @inbounds for j=1:1 + Nψ# + Nψ*Nψ
        @. S.cache_integral[(j-1)*Ne+1:j*Ne] *= xlast
    end

    # Add f(x_{1:d-1},0) i.e. (S.ψoff .* S.ψd0)*coeff to S.cache_integral
    @avx for i=1:Ne
        f0i = zero(Float64)
        for j=1:Nψ
            f0i += (S.ψoff[i,j] * S.ψd0[i,j])*coeff[j]
        end
        S.cache_integral[i] += f0i
    end

    # Store g(∂_{xk}f(x_{1:k})) in S.cache_g
    @avx for i=1:Ne
        prelogJi = zero(Float64)
        for j=1:Nψ
            prelogJi += (S.ψoff[i,j] * S.dψxd[i,j])*coeff[j]
        end
        S.cache_g[i] = prelogJi
    end

    # Formatting to use with Optim.jl
    if dJ != nothing
        reshape_cacheintegral = reshape(S.cache_integral[Ne+1:Ne+Ne*Nψ], (Ne, Nψ))
        fil!(dJ, 0.0)#dJ .= zeros(Nψ)
        @inbounds for i=1:Ne
            # dJ .= zeros(Nψ)
            for j=1:Nψ
                dJ[j] += gradlog_pdf(S.cache_integral[i])*(reshape_cacheintegral[i,j] + S.ψoff[i,j]*S.ψd0[i,j])
                dJ[j] += grad_x(C.I.g, S.cache_g[i])*S.ψoff[i,j]*S.dψxd[i,j]/C.I.g(S.cache_g[i])
            end
            # @show i, dJ
        end
        rmul!(dJ, -1/Ne)
        # Add derivative of the L2 penalty term ∂_c α ||c||^2 = 2 *α c
        dJ .+= 2*C.α*coeff
    end

    if d2J != nothing
        reshape_cacheintegral = reshape(S.cache_integral[Ne+1:Ne+Ne*Nψ], (Ne, Nψ))
        reshape2_cacheintegral = reshape(S.cache_integral[Ne + Ne*Nψ + 1: Ne + Ne*Nψ + Ne*Nψ*Nψ], (Ne, Nψ, Nψ))

        fill!(d2J, 0.0)
        # d2J .= zeros(Nψ, Nψ)
        @inbounds for l=1:Ne
            # Exploit symmetry of the Hessian
            for j=1:Nψ
                for i=j:Nψ
                d2J[i,j] +=  reshape2_cacheintegral[l,i,j]*S.cache_integral[l]
                d2J[i,j] +=  (reshape_cacheintegral[l,i] + S.ψoff[l,i]*S.ψd0[l,i]) * (reshape_cacheintegral[l,j] + S.ψoff[l,j]*S.ψd0[l,j])
                d2J[i,j] -=  ( (S.ψoff[l,i]*S.dψxd[l,i]) * (S.ψoff[l,j]*S.dψxd[l,j])*(
                                hess_x(C.I.g, S.cache_g[l]) * C.I.g(S.cache_g[l]) -
                                grad_x(C.I.g, S.cache_g[l])^2))/C.I.g(S.cache_g[l])^2

                d2J[j,i] = d2J[i,j]
                end
            end
        end
        rmul!(d2J, 1/Ne)
        # Add derivative of the L2 penalty term ∂^2_c α ||c||^2 = 2 *α *I
        @inbounds for i=1:Nψ
            d2J[i,i] += 2*C.α*I
        end
        # d2J = Symmetric(d2J)
        # return d2J
    end

    if J != nothing
        J = 0.0
        @inbounds for i=1:Ne
            J += log_pdf(S.cache_integral[i]) + log(C.I.g(S.cache_g[i]))
        end
        J *=(-1/Ne)
        J += C.α*norm(coeff)^2
        return J
    end
    # return J, dJ, d2J
end

hess_negative_log_likelihood!(S::Storage, C::HermiteMapComponent, X::Array{Float64,2}) =
    (J, dJ, d2J, coeff) -> hess_negative_log_likelihood!(J, dJ, d2J, coeff, S, C, X)


text:
function conductive_thin_wall(;nx::Int=3, ny::Int=3, Re::Float64 = 10.0,
    Ha::Float64 = 10.0, c_w::Float64 = 0.1, α::Float64 = 1.0,
    f_u = nothing, resultsfile = nothing)

  L = 1.0
  N = Ha^2/Re
  K = Ha / (1-0.825*Ha^(-1/2)-Ha^(-1))
  ∂p∂z = -Re * K / L^3

  if f_u === nothing
    _f_u = VectorValue(0.0,0.0, -∂p∂z) * L/Re^2 # Assumed ν=L=1 by default
  else
    _f_u = f_u
  end
  g_u = VectorValue(0.0,0.0,0.0)
  g_j = VectorValue(0.0,0.0,0.0)
  B = VectorValue(0.0,1.0,0.0)

  # Discretization
  order = 2
  domain = (-1.0,1.0,-1.0,1.0,0.0,0.1)
  map(x) = VectorValue(sign(x[1])*(abs(x[1])*0.5)^0.5,
                       sign(x[2])*(abs(x[2])*0.5)^0.5, x[3]*sqrt(2)/2)*2/sqrt(2)

  dirichlet_tags_u = append!(collect(1:20),[23,24,25,26])
  dirichlet_tags_j = append!(collect(1:20),[25,26])
  neumann_j = [23,24]

  partition=(nx,ny,3)
  model = CartesianDiscreteModel(domain,partition;
    isperiodic=(false,false,true), map=map)

  labels = get_face_labeling(model)
  add_tag_from_tags!(labels,"dirichlet_u",dirichlet_tags_u)
  add_tag_from_tags!(labels,"dirichlet_j",dirichlet_tags_j)
  add_tag_from_tags!(labels,"neumann_j",neumann_j)

  Vu = FESpace(model, ReferenceFE(lagrangian,VectorValue{3,Float64},order);
      conformity=:H1, dirichlet_tags="dirichlet_u")

  Vp = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:P);
      conformity=:L2, constraint=:zeromean)

  Vj = FESpace(model, ReferenceFE(raviart_thomas,Float64,order-1);
      conformity=:Hdiv, dirichlet_tags="dirichlet_j")

  Vφ = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:Q);
      conformity=:L2)

  U = TrialFESpace(Vu,g_u)
  P = TrialFESpace(Vp)
  J = TrialFESpace(Vj,g_j)
  Φ = TrialFESpace(Vφ)

  Y = MultiFieldFESpace([Vu, Vp, Vj, Vφ])
  X = MultiFieldFESpace([U, P, J, Φ])

  trian = Triangulation(model)
  degree = 2*(order)
  dΩ = Measure(trian,degree)

  btrian_j = BoundaryTriangulation(model,tags=["neumann_j"])
  dΓ = Measure(btrian_j,degree)
  n = get_normal_vector(btrian_j)

  res_Ω(x,y) = InductionlessMHD.dimensionless_residual(x, y, Re, N, B, f_u, dΩ)
  jac_Ω(x,dx,y) = InductionlessMHD.dimensionless_jacobian(x, dx, y, Re, N, B, dΩ)
  res_Γ(x,y) = InductionlessMHD.dimensionless_conducting_wall(x, y, n, c_w, dΓ, α=α)
  jac_Γ(x,dx,y) = InductionlessMHD.dimensionless_conducting_wall(dx, y, n, c_w, dΓ, α=α)

  res(x,y) = res_Ω(x,y) + res_Γ(x,y)
  jac(x,dx,y) = jac_Ω(x,dx,y) + jac_Γ(x,dx,y)
  op  = FEOperator(res,jac,X,Y)

  nls = NLSolver(;
    show_trace=true, method=:newton, linesearch=BackTracking())
  solver = FESolver(nls)

  xh = solve(solver,op)

  if resultsfile != nothing
    uh, ph, jh, φh = xh
    writevtk(trian, resultsfile,
      cellfields=["uh"=>uh, "ph"=>ph, "jh"=>jh, "φh"=>φh])
  end

  (xh, trian, dΩ)
end


text:
module SpeechDatasets

# Standard names/labels for the data preparation
include("consts.jl")

include("corpora/alffa.jl")
include("corpora/babel.jl")
include("corpora/massdataset.jl")
include("corpora/mboshi.jl")
include("corpora/timit.jl")
include("corpora/yoruba.jl")

using .ALFFA
using .BABEL
using .MASSDATASET
using .MBOSHI
using .TIMIT
using .YORUBA

export ALFFA
export BABEL
export MASSDATASET
export MBOSHI
export TIMIT
export YORUBA

end # module

max_issues_repo_name: lucasondel/SpeechDatasets

max_forks_repo_name: lucasondel/SpeechDatasets


text:
using FromFile
using SymbolicRegression, Test
@from "test_params.jl" import maximum_residual

_inv(x::Float32)::Float32 = 1f0/x
X = rand(Float32, 5, 100) .+ 1
y = 1.2f0 .+ 2 ./ X[3, :]

options = SymbolicRegression.Options(
    binary_operators=(+, *),
    unary_operators=(_inv,),
    npopulations=8
)
hallOfFame = EquationSearch(X, y, niterations=8, options=options, numprocs=4)

dominating = calculateParetoFrontier(X, y, hallOfFame, options)

best = dominating[end]
# Test the score
@test best.score < maximum_residual / 10


text:
module Wigner

using StaticArrays

__L2syms = [:s, :p, :d, :f, :g, :h, :i, :k]
__syms2L = Dict( [sym => L-1 for (L, sym) in enumerate(__L2syms)]... )
get_orbsym(L::Integer)  = __L2syms[L+1]

# Index of entries in D matrix (sign included)
struct D_Index
	sign::Int64
	μ::Int64
	m::Int64
end

# Equation (1.1) - forms the covariant matrix D(Q)(indices only)
function rotation_D_matrix(L::Integer)
	if L<0
		error("Orbital type shall be represented as a positive integer!")
	end
    D = Array{D_Index}(undef, 2 * L + 1, 2 * L + 1)
    for i = 1 : 2 * L + 1
        for j = 1 : 2 * L + 1
            D[j,i] = D_Index(1, i - 1 - L, j - 1 - L);
        end
    end
	return D
end

# Equation (1.1) - forms the covariant matrix D(Q)(indices only)
rotation_D_matrix_ast(L) =
	[ D_Index((-1)^(i+j), -(i - 1 - L), -(j - 1 - L))
	   for i = 1:2*L+1, j = 1:2*L+1 ]
# function rotation_D_matrix_ast(L::Integer)
# 	if L<0
# 		error("Orbital type shall be represented as a positive integer!")
# 	end
#     D = Array{D_Index}(undef, 2 * L + 1, 2 * L + 1)
#     for i = 1 : 2 * L + 1
#         for j = 1 : 2 * L + 1
#             D[i,j] = D_Index((-1)^(i+j), -(i - 1 - L), -(j - 1 - L));
#         end
#     end
# 	return D
# end

# D_{μm}^l(Ang2Mat_zyz(α,β,γ))
function wigner_big_D(μ,m,l,α,β,γ)
	return (exp(-im*α*m) * wigner_d(m,μ,l,β)  * exp(-im*γ*μ))'
end

# Wigner small d, modified from
# https://github.com/cortner/SlaterKoster.jl/blob/
# 8dceecb073709e6448a7a219ed9d3a010fa06724/src/code_generation.jl#L73
function wigner_d(μ, m, l, β)
    fc1 = factorial(l+m)
    fc2 = factorial(l-m)
    fc3 = factorial(l+μ)
    fc4 = factorial(l-μ)
    fcm1 = sqrt(fc1 * fc2 * fc3 * fc4)

    cosb = cos(β / 2.0)
    sinb = sin(β / 2.0)

    p = m - μ
    low  = max(0,p)
    high = min(l+m,l-μ)

    temp = 0.0
    for s = low:high
       fc5 = factorial(s)
       fc6 = factorial(l+m-s)
       fc7 = factorial(l-μ-s)
       fc8 = factorial(s-p)
       fcm2 = fc5 * fc6 * fc7 * fc8
       pow1 = 2 * l - 2 * s + p
       pow2 = 2 * s - p
       temp += (-1)^(s+p) * cosb^pow1 * sinb^pow2 / fcm2
    end
    temp *= fcm1

    return temp
end

function mat2ang(Q)
	return mod(atan(Q[2,3],Q[1,3]),2pi), acos(Q[3,3]), mod(atan(Q[3,2],-Q[3,1]),2pi);
end

# Rotation D matrix
function rot_D(L::Integer, Q)
	Mat_D = zeros(ComplexF64, 2L + 1, 2L + 1);
	D = rotation_D_matrix(L);
	α, β, γ = mat2ang(Q);
	for i = 1 : 2L + 1
		for j = 1 : 2L + 1
			Mat_D[i,j] = wigner_big_D(D[i,j].μ, D[i,j].m, L, α, β, γ);
		end
	end
	return SMatrix{2L+1, 2L+1, ComplexF64}(Mat_D)
end

function rand_QD(L)
	rotz(α) = [cos(α) -sin(α) 0; sin(α) cos(α) 0; 0 0 1]
	roty(α) = [cos(α) 0 sin(α); 0 1 0;-sin(α) 0 cos(α)]
	Ang2Mat_zyz(α,β,γ) = rotz(α)*roty(β)*rotz(γ)

	α = 2pi*rand();
	β = pi*rand();
	γ = 2pi*rand();

	# construct the Q matrix
	Q = Ang2Mat_zyz(α,β,γ)
	Q = SMatrix{3,3}(Q)

	rand_ref = rand((-1,1))

	return rand_ref * Q, rand_ref^L * rot_D(L, Q)
end

function rand_QD(L1, L2)
	rotz(α) = [cos(α) -sin(α) 0; sin(α) cos(α) 0; 0 0 1]
	roty(α) = [cos(α) 0 sin(α); 0 1 0;-sin(α) 0 cos(α)]
	Ang2Mat_zyz(α,β,γ) = rotz(α)*roty(β)*rotz(γ)

	α = 2pi*rand();
	β = pi*rand();
	γ = 2pi*rand();

	# construct the Q matrix
	Q = Ang2Mat_zyz(α,β,γ)
	Q = SMatrix{3,3}(Q)

	rand_ref = rand((-1,1))

	return rand_ref * Q, rand_ref^(L1) * rot_D(L1, Q), rand_ref^(L2) * rot_D(L2, Q)
end

end


text:
using DSGE, JLD, Statistics

# Solve model
m = KrusellSmithCT()
global T, R, C, inverse_basis, basis = solve(m)
global T = Matrix(T)
global R = Matrix(R)
global C = Vector(C)
global inverse_basis = Matrix(inverse_basis)

# Get quarterly data using the truth to change from reduced to full basis
trials = load("many_fine_data_red_basis.jld", "trials")
ntrials = length(keys(trials))
global states = trials[1]
global data = zeros(Dims((1, (length(states) - 1) / 90)))

# created daily data -> transform into quarterly data for discrete observations
for i = 1:length(data)
    full_state = inverse_basis * states[1+ (i-1)*90]
    data[i] = full_state[m.endogenous_states[:output][1]]
end

# Initialize storage matrices
σ_vec = collect(range(0.005, stop = .009, length = 20))
σ_vec = sort([σ_vec; .007])
sim_freqs = [2; 3; 12]
loglik_mat = zeros(length(σ_vec), length(sim_freqs), ntrials)
max_σ_mat = zeros(ntrials, length(sim_freqs))
max_lik_mat = zeros(ntrials, length(sim_freqs))
true_lik_vec = zeros(length(sim_freqs))
abs_dist_mat = zeros(ntrials, length(sim_freqs))
st_dev_vec = zeros(length(sim_freqs))
max_abs_dist_vec = zeros(length(sim_freqs))

# Compute true logliik and compile code
for i = 1:length(sim_freqs)
    TT = Matrix{Float64}(I, size(T)) + T * 1/get_setting(m, :state_simulation_freq) # adjust transition equation to account for discretization
    TT, RR, CC = transform_transition_matrices(m, TT, Matrix(R), Vector(C); track_lag = false) # get expanded filter

    # Define measurement equation
    measure = measurement(m, T, R, C, inverse_basis)
    ZZ = measure.ZZ
    DD = measure.DD
    QQ = measure.QQ
    EE = measure.EE

    out = DSGE.kalman_filter(data, TT, RR, CC, QQ, ZZ, DD, EE)
    true_lik_vec[i] = sum(out[1])
end

for j = 1:length(sim_freqs)
    println("Sim_freq: ", sim_freqs[j], " states per quarter")
    update!(m.settings[:state_simulation_freq], Setting(:state_simulation_freq, sim_freqs[j]))

    for n = 1:length(keys(trials))
        global states = trials[n]
        global data = zeros(Dims((1, (length(states) - 1) / 90)))

        for i = 1:length(data)
            full_state = inverse_basis * states[1+ (i-1)*90]
            data[i] = full_state[m.endogenous_states[:output][1]]
        end

        for i = 1:length(σ_vec)
            # Recompute steady state solution
            m[:σ_tfp] = σ_vec[i]
            steadystate!(m)
            global T, R, C, inverse_basis, basis = solve(m)
            global inverse_basis = Matrix(inverse_basis)
            TT = Matrix{Float64}(I, size(T)) + T * 1/get_setting(m, :state_simulation_freq) # adjust transition equation to account for discretization
            global TT, RR, CC = transform_transition_matrices(m, TT, R, C; track_lag = false) # get expanded filter

            # Define measurement equation
            measure = measurement(m, T, R, C, inverse_basis)
            ZZ = measure.ZZ
            DD = measure.DD
            QQ = measure.QQ
            EE = measure.EE

            out = DSGE.kalman_filter(data, TT, RR, CC, QQ, ZZ, DD, EE)
            loglik_mat[i, j, n] = sum(out[1])
        end
        max_ind = findall(maximum(loglik_mat[:, j, n]) .== loglik_mat[:, j, n])[1]
        max_σ_mat[n, j] = σ_vec[max_ind]
        max_lik_mat[n, j] = maximum(loglik_mat[:, j, n])
    end
    abs_dist_mat[:, j] = max_σ_mat[:, j] .- .007
    st_dev_vec[j] = std(abs_dist_mat[:, j])
    max_abs_dist_vec[j] = maximum(abs.(abs_dist_mat[:, j]))

    println("Mean absolute distance")
    println(mean(abs_dist_mat[:, j]))
    println("\n")

    println("Standard deviation of error")
    println(st_dev_vec[j])
    println("\n")

    println("Max absolute error")
    println(max_abs_dist_vec[j])
    println("\n")
end

save("sim_results.jld", "abs_dist_mat", abs_dist_mat, "st_dev_vec", st_dev_vec, "max_abs_dist_vec", max_abs_dist_vec, "max_σ_mat", max_σ_mat, "true_σ", .007, "σ_vec", σ_vec, "max_lik_mat", max_lik_mat, "true_lik_vec", true_lik_vec)


text:
module RandomCircuit

using CircuitViz

g = Graph()

iter = 1 : 100
for k in iter
    add_vertex!(g)
end

prod = Iterators.product(iter, iter)
for (k, v) in prod
    if rand() > 0.99
        @info (k, v)
        add_edge!(g, k, v)
    end
end

compile!(g)

end # module


text:
using LazySets, MathematicalPredicates

function crane_specification()
    # initial set
    X0 = Hyperrectangle([2.5, 0.0, 0.0, 0.0, 0.0, 0.0],
                        [2.5, 0.0, 0.2, 0.1, 0.0, 0.0])

    # safety property: x1 ≥ -1.8
    property = is_contained_in(HalfSpace([-1., 0.], 1.8))

    # time horizon: 15 time units
    time_horizon = 15.0

    # specification
    O = Dict(:T => time_horizon, :property => property)

    return X0, O
end

max_stars_repo_name: JuliaReach/ReachabilityBenchmarks

max_issues_repo_name: JuliaReach/ReachabilityBenchmarks

max_forks_repo_name: JuliaReach/ReachabilityBenchmarks


text:
import UUIDs

"""Internal storage common to InfrastructureSystems types."""
struct InfrastructureSystemsInternal
    uuid::Base.UUID
end

"""Creates PowerSystemInternal with a UUID."""
InfrastructureSystemsInternal() = InfrastructureSystemsInternal(UUIDs.uuid4())

"""Gets the UUID for any PowerSystemType."""
function get_uuid(obj::InfrastructureSystemsType)::Base.UUID
    return obj.internal.uuid
end