text:
"""
Provides the [`render`](@ref) methods to write the documentation as HTML files
(`MIME"text/html"`).

# Page outline

The [`HTMLWriter`](@ref) makes use of the page outline that is determined by the
headings. It is assumed that if the very first block of a page is a level 1 heading,
then it is intended as the page title. This has two consequences:

1. It is then used to automatically determine the page title in the navigation menu
   and in the `<title>` tag, unless specified in the `.pages` option.
2. If the first heading is interpreted as being the page title, it is not displayed
   in the navigation sidebar.

# Default and custom assets

Documenter copies all files under the source directory (e.g. `/docs/src/`) over
to the compiled site. It also copies a set of default assets from `/assets/html/`
to the site's `assets/` directory, unless the user already had a file with the
same name, in which case the user's files overrides the Documenter's file.
This could, in principle, be used for customizing the site's style and scripting.

The HTML output also links certain custom assets to the generated HTML documents,
specfically a logo and additional javascript files.
The asset files that should be linked must be placed in `assets/`, under the source
directory (e.g `/docs/src/assets`) and must be on the top level (i.e. files in
the subdirectories of `assets/` are not linked).

For the **logo**, Documenter checks for the existence of `assets/logo.png`.
If that's present, it gets displayed in the navigation bar.

For **scripts**, every `assets/*.js` gets a `<script>` link in the `<head>` tag
of every page (except if it matches one of Documenter's default scripts;
the filtering is done in [`user_scripts`](@ref)).

Note that only javascript files are linked to the generated HTML. Any related CSS
must be loaded by the script. With jQuery this could be done with the following
snippet

```javascript
\$('head').append(\$('<link rel="stylesheet">').attr('href', documenterBaseURL + "/assets/<file>.css"))
```
"""
module HTMLWriter

using Compat

import ...Documenter:
    Anchors,
    Builder,
    Documents,
    Expanders,
    Formats,
    Documenter,
    Utilities,
    Writers

import ..Writers: Writer, render
import ...Utilities.DOM: DOM, Tag, @tags
using ...Utilities.MDFlatten

const requirejs_cdn = "https://cdnjs.cloudflare.com/ajax/libs/require.js/2.2.0/require.min.js"

"""
[`HTMLWriter`](@ref)-specific globals that are passed to [`domify`](@ref) and
other recursive functions.
"""
type HTMLContext
    doc :: Documents.Document
    logo :: Compat.String
    scripts :: Vector{Compat.String}
    documenter_js :: Compat.String
    documenter_css :: Compat.String
    search_js :: Compat.String
    search_index :: IOBuffer
    search_index_js :: Compat.String
    user_scripts :: Vector{Compat.String}
end
HTMLContext(doc) = HTMLContext(doc, "", [], "", "", "", IOBuffer(), "", [])

"""
Returns a page (as a [`Documents.Page`](@ref) object) using the [`HTMLContext`](@ref).
"""
getpage(ctx, path) = ctx.doc.internal.pages[path]
getpage(ctx, navnode::Documents.NavNode) = getpage(ctx, get(navnode.page))

function render(::Writer{Formats.HTML}, doc::Documents.Document)
    !isempty(doc.user.sitename) || error("HTML output requires `sitename`.")

    ctx = HTMLContext(doc)
    ctx.search_index_js = "search_index.js"

    ctx.documenter_css = copy_asset("documenter.css", doc)
    copy_asset("style.css", doc)

    let logo = joinpath("assets", "logo.png")
        if isfile(joinpath(doc.user.build, logo))
            ctx.logo = logo
        end
    end

    ctx.documenter_js = copy_asset("documenter.js", doc)
    ctx.search_js = copy_asset("search.js", doc)
    ctx.user_scripts = user_scripts(doc)

    for navnode in doc.internal.navlist
        render_page(ctx, navnode)
    end

    render_search(ctx)

    open(joinpath(doc.user.build, ctx.search_index_js), "w") do io
        println(io, "var documenterSearchIndex = {\"docs\": [\n")
        write(io, takebuf_string(ctx.search_index))
        println(io, "]}")
    end
end

"""
Copies an asset from Documenters `assets/html/` directory to `doc.user.build`.
Returns the path of the copied asset relative to `.build`.
"""
function copy_asset(file, doc)
    src = joinpath(Utilities.assetsdir(), "html", file)
    alt_src = joinpath(doc.user.source, "assets", file)
    dst = joinpath(doc.user.build, "assets", file)
    isfile(src) || error("Asset '$file' not found at $(abspath(src))")

    # Since user's alternative assets are already copied over in a previous build
    # step and they should override documenter's original assets, we only actually
    # perform the copy if <source>/assets/<file> does not exist. Note that checking
    # the existence of <build>/assets/<file> is not sufficient since the <build>
    # directory might be dirty from a previous build.
    if isfile(alt_src)
        Utilities.warn("Not copying '$src', provided by the user.")
    else
        ispath(dirname(dst)) || mkpath(dirname(dst))
        ispath(dst) && Utilities.warn("Overwriting '$dst'.")
        cp(src, dst, remove_destination=true)
    end
    normpath(joinpath("assets", file))
end

"""
Creates a list of `.js` files provided by the user under `assets`. Returns a list
of paths relative to the site root.
"""
function user_scripts(doc)
    documenter_assetsdir = joinpath(Utilities.assetsdir(), "html")
    local_assetsdir = joinpath(doc.user.source, "assets")
    isdir(local_assetsdir) || return []
    scripts = filter(readdir(local_assetsdir)) do file
        # we'll make sure that it's
        #   1. a file (i.e. not a directory),
        #   2. a script, determined by the extension
        #   3. not one of the default assets, which get handled separately
        isfile(joinpath(local_assetsdir, file)) &&
            last(splitext(file)) == ".js" &&
            !isfile(joinpath(documenter_assetsdir, file))
    end
    [joinpath("assets", script) for script in scripts]
end

# Page
# ------------------------------------------------------------------------------

"""
Constructs and writes the page referred to by the `navnode` to `.build`.
"""
function render_page(ctx, navnode)
    @tags html body

    page = getpage(ctx, navnode)

    head = render_head(ctx, navnode, ctx.user_scripts)
    navmenu = render_navmenu(ctx, navnode)
    article = render_article(ctx, navnode)

    htmldoc = DOM.HTMLDocument(
        html[:lang=>"en"](
            head,
            body(navmenu, article)
        )
    )
    open(Formats.extension(Formats.HTML, page.build), "w") do io
        print(io, htmldoc)
    end
end

function render_head(ctx, navnode, additional_scripts)
    @tags head meta link script title
    src = get(navnode.page)
    page_title = "$(mdflatten(pagetitle(ctx, navnode))) · $(ctx.doc.user.sitename) documentation"
    head(
        meta[:charset=>"UTF-8"],
        meta[:name => "viewport", :content => "width=device-width, initial-scale=1.0"],
        title(page_title),

        link[
            :href => relhref(src, ctx.documenter_css),
            :rel => "stylesheet",
            :type => "text/css"
        ],

        script("documenterBaseURL=\"$(relhref(src, "."))\""),
        script[
            :src => requirejs_cdn,
            Symbol("data-main") => relhref(src, ctx.documenter_js)
        ],
        map(additional_scripts) do s
            script[:src => relhref(src, s)]
        end
    )
end

## Search page
# ------------

function render_search(ctx)
    @tags article body h1 header hr html li nav p span ul
    navnode = Documents.NavNode("search", "Search", nothing)

    additional_scripts = vcat([ctx.search_js, ctx.search_index_js], ctx.user_scripts)
    head = render_head(ctx, navnode, additional_scripts)
    navmenu = render_navmenu(ctx, navnode)
    article = article(
        header(
            nav(ul(li("Search"))),
            hr()
        ),
        h1("Search"),
        p["#search-info"]("Number of results: ", span["#search-results-number"]("-")),
        ul["#search-results"]
    )

    htmldoc = DOM.HTMLDocument(
        html[:lang=>"en"](
            head,
            body(navmenu, article)
        )
    )
    open(Formats.extension(Formats.HTML, joinpath(ctx.doc.user.build, "search")), "w") do io
        print(io, htmldoc)
    end
end

# Navigation menu
# ------------------------------------------------------------------------------

function render_navmenu(ctx, navnode)
    @tags a form h1 img input nav
    src = get(navnode.page)
    navmenu = nav[".toc"]
    if !isempty(ctx.logo)
        push!(navmenu.nodes,
            a[:href => ""]( # TODO: link to github?
                img[
                    ".logo",
                    :src => relhref(src, ctx.logo),
                    :alt => "$(ctx.doc.user.sitename) logo"
                ]
            )
        )
    end
    push!(navmenu.nodes, h1(ctx.doc.user.sitename))
    push!(navmenu.nodes,
        form[".search", :action => relhref(src, "search.html")](
            input[
                "#search-query",
                :name => "q",
                :type => "text",
                :placeholder => "Search docs"
            ]
        )
    )
    push!(navmenu.nodes, navitem(ctx, navnode))
    navmenu
end

"""
[`navitem`](@ref) returns the lists and list items of the navigation menu.
It gets called recursively to construct the whole tree.
"""
navitem(ctx, current) = navitem(ctx, current, ctx.doc.internal.navtree)
navitem(ctx, current, nns::Vector) = DOM.Tag(:ul)(map(nn -> navitem(ctx, current, nn), nns))
function navitem(ctx, current, nn::Documents.NavNode)
    @tags ul li span a

    # construct this item
    title = mdconvert(pagetitle(ctx, nn))
    link = if isnull(nn.page)
        span[".toctext"](title)
    else
        a[".toctext", :href => navhref(nn, current)](title)
    end
    item = (nn === current) ? li[".current"](link) : li(link)

    # add the subsections (2nd level headings) from the page
    if nn === current && !isnull(nn.page)
        subs = collect_subsections(ctx.doc.internal.pages[get(nn.page)])
        internal_links = map(subs) do _
            istoplevel, anchor, text = _
            _li = istoplevel ? li[".toplevel"] : li[]
            _li(a[".toctext", :href => anchor](mdconvert(text)))
        end
        push!(item.nodes, ul[".internal"](internal_links))
    end

    # add the subsections, if any, as a single list
    if !isempty(nn.children)
        push!(item.nodes, navitem(ctx, current, nn.children))
    end

    item
end

# Article (page contents)
# ------------------------------------------------------------------------------

function render_article(ctx, navnode)
    @tags article header footer nav ul li hr span a

    header_links = map(Documents.navpath(navnode)) do nn
        title = mdconvert(pagetitle(ctx, nn))
        isnull(nn.page) ? li(title) : li(a[:href => navhref(nn, navnode)](title))
    end

    topnav = nav(ul(header_links))
    Utilities.unwrap(Utilities.url(ctx.doc.user.repo, getpage(ctx, navnode).source)) do url
        push!(topnav.nodes, a[".edit-page", :href => url](span[".fa"]("\uf09b"), " Edit on GitHub"))
    end
    art_header = header(topnav, hr())

    # build the footer with nav links
    art_footer = footer(hr())
    Utilities.unwrap(navnode.prev) do nn
        direction = span[".direction"]("Previous")
        title = span[".title"](mdconvert(pagetitle(ctx, nn)))
        link = a[".previous", :href => navhref(nn, navnode)](direction, title)
        push!(art_footer.nodes, link)
    end

    Utilities.unwrap(navnode.next) do nn
        direction = span[".direction"]("Next")
        title = span[".title"](mdconvert(pagetitle(ctx, nn)))
        link = a[".next", :href => navhref(nn, navnode)](direction, title)
        push!(art_footer.nodes, link)
    end

    pagenodes = domify(ctx, navnode)
    article["#docs"](art_header, pagenodes, art_footer)
end

## domify(...)
# ------------

"""
Converts recursively a [`Documents.Page`](@ref), `Base.Markdown` or Documenter
`*Node` objects into HTML DOM.
"""
function domify(ctx, navnode)
    page = getpage(ctx, navnode)
    sib = SearchIndexBuffer(ctx, navnode)
    ret = map(page.elements) do elem
        search_append(sib, elem)
        domify(ctx, navnode, page.mapping[elem])
    end
    search_flush(sib)
    ret
end

type SearchIndexBuffer
    ctx :: HTMLContext
    src :: Compat.String
    page :: Documents.Page
    loc :: Compat.String
    category :: Symbol
    title :: Compat.String
    page_title :: Compat.String
    buffer :: IOBuffer
    function SearchIndexBuffer(ctx, navnode)
        page_title = mdflatten(pagetitle(ctx, navnode))
        new(
            ctx,
            Formats.extension(Formats.HTML, get(navnode.page)),
            getpage(ctx, navnode),
            "",
            :page,
            page_title,
            page_title,
            IOBuffer()
        )
    end
end

function search_append(sib, node::Markdown.Header)
    search_flush(sib)
    sib.category = :section
    sib.title = mdflatten(node)
    a = sib.page.mapping[node]
    sib.loc = "$(a.id)-$(a.nth)"
end

search_append(sib, node) = mdflatten(sib.buffer, node)

function search_flush(sib)
    ref = isempty(sib.src) ? sib.src : "$(sib.src)#$(sib.loc)"
    text = takebuf_string(sib.buffer)
    println(sib.ctx.search_index, """
    {
        "location": "$(jsonescape(ref))",
        "page": "$(jsonescape(sib.page_title))",
        "title": "$(jsonescape(sib.title))",
        "category": "$(jsonescape(string(sib.category)))",
        "text": "$(jsonescape(text))"
    },
    """)
end

function jsonescape(s)
    s = replace(s, '\\', "\\\\")
    s = replace(s, '\n', "\\n")
    replace(s, '"', "\\\"")
end

domify(ctx, navnode, node) = mdconvert(node, Base.Markdown.MD())

function domify(ctx, navnode, anchor::Anchors.Anchor)
    @tags a
    aid = "$(anchor.id)-$(anchor.nth)"
    if isa(anchor.object, Markdown.Header)
        h = anchor.object
        DOM.Tag(Symbol("h$(Utilities.header_level(h))"))(
            a[".nav-anchor", :id => aid, :href => "#$aid"](mdconvert(h.text, h))
        )
    else
        a[".nav-anchor", :id => aid, :href => "#$aid"](domify(ctx, navnode, anchor.object))
    end
end

immutable ListBuilder
    es::Vector
end
ListBuilder() = ListBuilder([])

import Base: push!
function push!(lb::ListBuilder, level, node)
    @assert level >= 1
    if level == 1
        push!(lb.es, node)
    else
        if isempty(lb.es) || typeof(last(lb.es)) !== ListBuilder
            push!(lb.es, ListBuilder())
        end
        push!(last(lb.es), level-1, node)
    end
end

function domify(lb::ListBuilder)
    @tags ul li
    ul(map(e -> isa(e, ListBuilder) ? domify(e) : li(e), lb.es))
end

function domify(ctx, navnode, contents::Documents.ContentsNode)
    @tags a
    lb = ListBuilder()
    for (count, path, anchor) in contents.elements
        path = Formats.extension(ctx.doc.user.format, path)
        header = anchor.object
        url = string(path, '#', anchor.id, '-', anchor.nth)
        node = a[:href=>url](mdconvert(header.text))
        level = Utilities.header_level(header)
        push!(lb, level, node)
    end
    domify(lb)
end

function domify(ctx, navnode, index::Documents.IndexNode)
    @tags a code li ul
    lis = map(index.elements) do _
        object, doc, page, mod, cat = _
        page = Formats.extension(ctx.doc.user.format, page)
        url = string(page, "#", Utilities.slugify(object))
        li(a[:href=>url](code("$(object.binding)")))
    end
    ul(lis)
end

function domify(ctx, navnode, docs::Documents.DocsNodes)
    [domify(ctx, navnode, node) for node in docs.nodes]
end

function domify(ctx, navnode, node::Documents.DocsNode)
    @tags a code div section span

    # push to search index
    sib = SearchIndexBuffer(ctx, navnode)
    sib.loc = node.anchor.id
    sib.title = string(node.object.binding)
    sib.category = Utilities.doccat(node.object)
    mdflatten(sib.buffer, node.docstr)
    search_flush(sib)

    section[".docstring"](
        div[".docstring-header"](
            a[".docstring-binding", :id=>node.anchor.id, :href=>"#$(node.anchor.id)"](code("$(node.object.binding)")),
            " — ", # —
            span[".docstring-category"]("$(Utilities.doccat(node.object))"),
            "."
        ),
        domify_doc(ctx, navnode, node.docstr)
    )
end

function domify_doc(ctx, navnode, md::Markdown.MD)
    @tags a br
    if haskey(md.meta, :results)
        # The `:results` field contains a vector of `Docs.DocStr` objects associated with
        # each markdown object. The `DocStr` contains data such as file and line info that
        # we need for generating correct source links.
        map(zip(md.content, md.meta[:results])) do _
            markdown, result = _
            ret = Any[domify(ctx, navnode, Writers.MarkdownWriter.dropheaders(markdown))]
            # When a source link is available then print the link.
            Utilities.unwrap(Utilities.url(ctx.doc.internal.remote, ctx.doc.user.repo, result)) do url
                push!(ret, a[".source-link", :target=>"_blank", :href=>url]("source"))
                push!(ret, br())
            end
            ret
        end
    else
        # Docstrings with no `:results` metadata won't contain source locations so we don't
        # try to print them out. Just print the basic docstring.
        domify(ctx, navnode, Writers.MarkdownWriter.dropheaders(md))
    end
end

function domify(ctx, navnode, node::Documents.EvalNode)
    node.result === nothing ? DOM.Node[] : domify(ctx, navnode, node.result)
end

# nothing to show for MetaNodes, so we just return an empty list
domify(ctx, navnode, node::Documents.MetaNode) = DOM.Node[]

# Utilities
# ------------------------------------------------------------------------------

"""
Get the relative hyperlink between two [`Documents.NavNode`](@ref)s. Assumes that both
[`Documents.NavNode`](@ref)s have an associated [`Documents.Page`](@ref) (i.e. `.page` is not null).
"""
navhref(to, from) = Formats.extension(Formats.HTML, relhref(get(from.page), get(to.page)))

"""
Calculates a relative HTML link from one path to another.
"""
function relhref(from, to)
    pagedir = dirname(from)
    relpath(to, isempty(pagedir) ? "." : pagedir)
end

"""
Tries to guess the page title by looking at the `<h1>` headers and returns the
header contents of the first `<h1>` on a page as a `Nullable` (nulled if the algorithm
was unable to find any `<h1>` headers).
"""
function pagetitle(page::Documents.Page)
    title = Nullable{Any}()
    for element in page.elements
        if isa(element, Base.Markdown.Header{1})
            title = Nullable{Any}(element.text)
            break
        end
    end
    title
end

function pagetitle(ctx, navnode::Documents.NavNode)
    isnull(navnode.title_override) || return get(navnode.title_override)

    if !isnull(navnode.page)
        title = pagetitle(getpage(ctx, navnode))
        isnull(title) || return get(title)
    end

    "-"
end

"""
Returns an ordered list of tuples, `(toplevel, anchor, text)`, corresponding to level 1 and 2
headings on the `page`. Note that if the first header on the `page` is a level 1 header then
it is not included -- it is assumed to be the page title and so does not need to be included
in the navigation menu twice.
"""
function collect_subsections(page::Documents.Page)
    local sections = []
    for element in page.elements
        if isa(element, Base.Markdown.Header) && Utilities.header_level(element) < 3
            local toplevel = Utilities.header_level(element) === 1
            # Don't include the first header if it is `h1`.
            toplevel && isempty(sections) && continue
            local anchor = page.mapping[element]
            push!(sections, (toplevel, "#$(anchor.id)-$(anchor.nth)", element.text))
        end
    end
    return sections
end

# mdconvert
# ------------------------------------------------------------------------------

const md_block_nodes = [Markdown.MD, Markdown.BlockQuote]
fieldtype(Markdown.List, :ordered) == Int && push!(md_block_nodes, Markdown.List)
if isdefined(Base.Markdown, :Admonition) push!(md_block_nodes, Markdown.Admonition) end

"""
[`MDBlockContext`](@ref) is a union of all the Markdown nodes whose children should
be blocks. It can be used to dispatch on all the block-context nodes at once.
"""
typealias MDBlockContext Union{md_block_nodes...}

"""
Convert a markdown object to a `DOM.Node` object.

The `parent` argument is passed to allow for context-dependant conversions.
"""
mdconvert(md) = mdconvert(md, md)

mdconvert(text::AbstractString, parent) = DOM.Node(text)

mdconvert(vec::Vector, parent) = [mdconvert(x, parent) for x in vec]

mdconvert(md::Markdown.MD, parent) = Tag(:div)(mdconvert(md.content, md))

mdconvert(b::Markdown.BlockQuote, parent) = Tag(:blockquote)(mdconvert(b.content, b))

mdconvert(b::Markdown.Bold, parent) = Tag(:strong)(mdconvert(b.text, parent))

function mdconvert(c::Markdown.Code, parent::MDBlockContext)
    @tags pre code
    language = isempty(c.language) ? "none" : c.language
    pre(code[".language-$(language)"](c.code))
end
mdconvert(c::Markdown.Code, parent) = Tag(:code)(c.code)

mdconvert{N}(h::Markdown.Header{N}, parent) = DOM.Tag(Symbol("h$N"))(mdconvert(h.text, h))

mdconvert(::Markdown.HorizontalRule, parent) = Tag(:hr)()

mdconvert(i::Markdown.Image, parent) = Tag(:img)[:src => i.url, :alt => i.alt]

mdconvert(i::Markdown.Italic, parent) = Tag(:em)(mdconvert(i.text, i))

mdconvert(m::Markdown.LaTeX, ::MDBlockContext)   = Tag(:div)(string("\\[", m.formula, "\\]"))
mdconvert(m::Markdown.LaTeX, parent) = Tag(:span)(string('$', m.formula, '$'))

mdconvert(::Markdown.LineBreak, parent) = Tag(:br)()

const ABSURL_REGEX = r"^[[:alpha:]+-.]+://"
function mdconvert(link::Markdown.Link, parent)
    # TODO: fixing the extension should probably be moved to an earlier step
    if ismatch(ABSURL_REGEX, link.url)
        Tag(:a)[:href => link.url](mdconvert(link.text, link))
    else
        s = split(link.url, "#", limit = 2)
        path = first(s)
        path = endswith(path, ".md") ? Formats.extension(Formats.HTML, path) : path
        url = (length(s) > 1) ? "$path#$(last(s))" : Compat.String(path)
        Tag(:a)[:href => url](mdconvert(link.text, link))
    end
end

mdconvert(list::Markdown.List, parent) = (isordered(list) ? Tag(:ol) : Tag(:ul))(map(Tag(:li), mdconvert(list.items, list)))

mdconvert(paragraph::Markdown.Paragraph, parent) = Tag(:p)(mdconvert(paragraph.content, paragraph))

mdconvert(t::Markdown.Table, parent) = Tag(:table)(
    Tag(:tr)(map(_ -> Tag(:th)(mdconvert(_, t)), t.rows[1])),
    map(_ -> Tag(:tr)(map(__ -> Tag(:td)(mdconvert(__, _)), _)), t.rows[2:end])
)

mdconvert(expr::Union{Expr,Symbol}, parent) = string(expr)

# Only available on Julia 0.5.
if isdefined(Base.Markdown, :Footnote)
    mdconvert(f::Markdown.Footnote, parent) = footnote(f.id, f.text, parent)
    footnote(id, text::Void, parent) = Tag(:a)[:href => "#footnote-$(id)"]("[$id]")
    function footnote(id, text, parent)
        Tag(:div)[".footnote#footnote-$(id)"](
            Tag(:a)[:href => "#footnote-$(id)"](Tag(:strong)("[$id]")),
            mdconvert(text, parent),
        )
    end
end

if isdefined(Base.Markdown, :Admonition)
    function mdconvert(a::Markdown.Admonition, parent)
        @tags div
        div[".admonition.$(a.category)"](
            div[".admonition-title"](a.title),
            div[".admonition-text"](mdconvert(a.content, a))
        )
    end
end

if isdefined(Base.Markdown, :isordered)
    import Base.Markdown: isordered
else
    isordered(a::Markdown.List) = a.ordered::Bool
end

mdconvert(html::Documents.RawHTML, parent) = Tag(Symbol("#RAW#"))(html.code)

end


text:
using CartesianGP
using Base.Test

# Create some 2-input, 2-output goals, these should all be equivalent
# to one another.

g2_u = Goal(2, (0b0010, 0b0001))
g2_p = BasicPackedGoal(2, 2, 0b00010010)
g2_i = InterleavedPackedGoal(2, 2, 0b00000110)

@test convert(BasicPackedGoal, g2_u) == g2_p
@test convert(InterleavedPackedGoal, g2_u) == g2_i

@test convert(Goal, g2_p) == g2_u
@test convert(InterleavedPackedGoal, g2_p) == g2_i

@test convert(Goal, g2_i) == g2_u
@test convert(BasicPackedGoal, g2_i) == g2_p

@test Goal(2, (0b0000, 0b0000)) != g2_u
@test BasicPackedGoal(2, 2, 0b00000000) != g2_p
@test InterleavedPackedGoal(2, 2, 0b00000000) != g2_i

@test_throws ErrorException g2_u == g2_p
@test_throws ErrorException g2_p == g2_u
@test_throws ErrorException g2_p == g2_i
@test_throws ErrorException g2_i == g2_p
@test_throws ErrorException g2_u == g2_i
@test_throws ErrorException g2_i == g2_u

# Create some 3-input 3-output goals, these should all be equivalent
# to one another.

g3_uh = Goal(3, (0xC9, 0x0F, 0xA5))  # hexadecimal version
g3_ub = Goal(3, (0b11001001, 0b00001111, 0b10100101)) # equivalent binary version
g3_ph = BasicPackedGoal(3, 3, 0xA50FC9)
g3_io = InterleavedPackedGoal(3, 3, 0o51403627) # octal is most natural for this goal

@test g3_ub == g3_uh # check that binary and hex versions are the same
@test convert(BasicPackedGoal, g3_uh) == g3_ph
@test convert(Goal, g3_ph) == g3_uh
@test convert(InterleavedPackedGoal, g3_ph) == g3_io
@test convert(BasicPackedGoal, g3_io) == g3_ph
@test convert(Goal, g3_io) == g3_ub
@test convert(Goal, g3_ph) == g3_ub

# Create goals to test the composition of goals.

g3_inv = InterleavedPackedGoal(3, 3, 0o02753164)  # compositional inverse of g3_io
g3_ident = InterleavedPackedGoal(3, 3, 0o76543210) # compositional identity function
@test compose(g3_io, g3_inv) == g3_ident
@test compose(g3_inv, g3_io) == g3_ident
@test compose(g3_io, g3_ident) == g3_io
@test compose(g3_ident, g3_inv) == g3_inv

# Test reading a PLU file.

fname = "add1c.plu"
add1c_goal = Goal(3, (0x96, 0xE8))
add1c_plu_goal = read_plu(fname)

@test read_plu(fname) == convert(BasicPackedGoal, add1c_goal)

ext: jl


text:
using Documenter, JulzDummy

makedocs(
  modules = [JulzDummy],
  format = :html,
  sitename = "JulzDummy.jl",
  pages = Any[
    "Home" => "index.md",
  ]
)

deploydocs(
  repo = "github.com/djsegal/JulzDummy.jl.git",
  target = "build",
  deps = nothing,
  make = nothing
)


text:
rm("images", force=true, recursive=true)
rm("models", force=true, recursive=true)
module README; include("README.jl"); end
using Literate
Literate.markdown("README.jl", "..", documenter=false)


text:
### old definition

#=
# animals
abstract type Sheep <: AnimalSpecies end
abstract type Wolf <: AnimalSpecies end

# plants
abstract type Mushroom <: PlantSpecies end
abstract type Grass <: PlantSpecies end

export Grass, Sheep, Wolf, Mushroom

Base.show(io::IO, ::Type{Sheep}) = print(io,"🐑")
Base.show(io::IO, ::Type{Wolf}) = print(io,"🐺")

Base.show(io::IO,::Type{Mushroom}) = print(io,"🍄")
Base.show(io::IO, ::Type{Grass}) = print(io,"🌿")

function EcosystemCore.eat!(s::Animal{Sheep}, m::Plant{Mushroom}, w::World)
    if size(p)>0
        incr_energy!(s, -size(m)*Δenergy(s))
        m.size = 0
    end
end

function EcosystemCore.eat!(a::Animal{Wolf}, b::Animal{Sheep}, w::World)
    incr_energy!(a, energy(b)*Δenergy(a))
    kill_agent!(b,w)
end

function EcosystemCore.eat!(a::Animal{Sheep}, b::Plant{Grass}, w::World)
    incr_energy!(a, size(b)*Δenergy(a))
    b.size = 0
end

EcosystemCore.eats(::Animal{Sheep}, ::Plant{Mushroom}) = true
EcosystemCore.eats(::Animal{Sheep}, ::Plant{Grass}) = true
EcosystemCore.eats(::Animal{Wolf},::Animal{Sheep}) = true

EcosystemCore.mates(::Animal{S,Female}, ::Animal{S,Male}) where S<:Species = true
EcosystemCore.mates(::Animal{S,Male}, ::Animal{S,Female}) where S<:Species = true
EcosystemCore.mates(a::Agent, b::Agent) = false
=#

### new definition using macros from `ecosystem_macros.jl`
@plant  Grass       🌿
@plant  Broccoli    🥦
@plant  Mushroom    🍄
@animal Sheep       🐑
@animal Wolf        🐺
@animal Rabbit      🐇

@eats Rabbit [Grass => 0.5, Broccoli => 1.0]
@eats Sheep  [Grass => 0.5, Broccoli => 1.0, Mushroom => -1.0]
@eats Wolf   [Sheep => 0.9]


text:
push!(LOAD_PATH, joinpath(@__DIR__, "..", ".."))

using BenchmarkTools, Plots
using CUDA

include("bickleyjet_dg.jl")
include("bickleyjet_dg_reference.jl")

n1, n2 = 16, 16

Nqs = 2:7
volTs = Float64[]
volRs = Float64[]
volRGPUs = Float64[]
faceTs = Float64[]
faceRs = Float64[]
faceRGPUs = Float64[]

DA = CUDA.has_cuda_gpu() ? CuArray : Array

for Nq in Nqs
    # setup core
    mesh = Meshes.EquispacedRectangleMesh(domain, n1, n2)
    grid_topology = Topologies.GridTopology(mesh)
    quad = Spaces.Quadratures.GLL{Nq}()
    space = Spaces.SpectralElementSpace2D(grid_topology, quad)

    y0 = init_state.(Fields.coordinate_field(space), Ref(parameters))
    dydt = Fields.Field(similar(Fields.field_values(y0)), space)
    volume!(dydt, y0, (parameters,), 0.0)
    # TODO: move this to volume!
    dydt_data = Fields.field_values(dydt)
    dydt_data .= RecursiveApply.rdiv.(dydt_data, space.local_geometry.WJ)

    # setup reference
    X = coordinates(Val(Nq), n1, n2)
    y0_ref = init_st_ref(X, Val(Nq), parameters)
    dydt_ref = similar(y0_ref)
    tendency_states = init_tendency_states(n1, n2, Val(Nq))
    volume_ref!(
        dydt_ref,
        y0_ref,
        (n1, n2, parameters, Val(Nq), tendency_states),
        0.0,
    )

    # check equivalent
    @assert y0_ref ≈ reshape(parent(y0), (Nq, Nq, 4, n1, n2))
    @assert dydt_ref ≈ reshape(parent(dydt), (Nq, Nq, 4, n1, n2))

    # setup GPU reference
    if DA === CuArray
        y0_gpu = DA(y0_ref)
        dydt_gpu = similar(y0_gpu)
        volume_ref_cuda!(dydt_gpu, y0_gpu, parameters, Nq)
        # check equivalent
        @assert Array(y0_gpu) ≈ y0_ref
        @assert Array(dydt_gpu) ≈ dydt_ref
    end
    # run benchmarks
    @info("Benchmark volume!", Nq)
    push!(volTs, @belapsed volume!($dydt, $y0, ($parameters,), 0.0))
    @info("Benchmark volume_ref!", Nq)
    push!(
        volRs,
        @belapsed volume_ref!(
            $dydt_ref,
            $y0_ref,
            ($n1, $n2, $parameters, $(Val(Nq)), $tendency_states),
            0.0,
        )
    )
    if DA === CuArray
        @info("Benchmark volume_ref GPU!", Nq)
        push!(
            volRGPUs,
            @belapsed volume_ref_cuda!($dydt_gpu, $y0_gpu, $parameters, $Nq)
        )
    end

    # faces
    fill!(parent(dydt), 0.0)
    add_face!(dydt, y0, (parameters,), 0.0)
    # TODO: move this to volume!
    dydt_data = Fields.field_values(dydt)
    dydt_data .= RecursiveApply.rdiv.(dydt_data, space.local_geometry.WJ)

    fill!(dydt_ref, 0.0)
    add_face_ref!(dydt_ref, y0_ref, (n1, n2, parameters, Val(Nq)), 0.0)

    @assert dydt_ref ≈ reshape(parent(dydt), (Nq, Nq, 4, n1, n2))

    # setup GPU face reference
    if DA === CuArray
        fill!(dydt_gpu, 0.0)
        add_face_ref_cuda!(dydt_gpu, y0_gpu, parameters, Nq)
        @assert Array(y0_gpu) ≈ y0_ref
        @assert Array(dydt_gpu) ≈ dydt_ref
    end
    @info("Benchmark face!", Nq)
    push!(faceTs, @belapsed add_face!($dydt, $y0, ($parameters,), 0.0))
    @info("Benchmark face_ref!", Nq)
    push!(
        faceRs,
        @belapsed add_face_ref!(
            $dydt_ref,
            $y0_ref,
            ($n1, $n2, $parameters, $(Val(Nq))),
            0.0,
        )
    )
    if DA === CuArray
        @info("Benchmark face_ref GPU!", Nq)
        push!(
            faceRGPUs,
            @belapsed add_face_ref_cuda!($dydt_gpu, $y0_gpu, $parameters, $Nq)
        )
    end
end

using Plots
ENV["GKSwstype"] = "nul"

plt = plot(
    ylims = (0, Inf),
    xlabel = "Nq",
    ylabel = "Time (ms)",
    title = "Volume",
)
plot!(plt, Nqs, 1e3 .* volTs, label = "ClimaCore")
plot!(plt, Nqs, 1e3 .* volRs, label = "Reference")
if DA === CuArray
    plot!(plt, Nqs, 1e3 .* volRGPUs, label = "Reference GPU")
end

png(plt, joinpath(@__DIR__, "volume.png"))

plt =
    plot(ylims = (0, Inf), xlabel = "Nq", ylabel = "Time (ms)", title = "Face")
plot!(plt, Nqs, 1e3 .* faceTs, label = "ClimaCore")
plot!(plt, Nqs, 1e3 .* faceRs, label = "Reference")
if DA === CuArray
    plot!(plt, Nqs, 1e3 .* faceRGPUs, label = "Reference GPU")
end

png(plt, joinpath(@__DIR__, "face.png"))

function linkfig(figpath, alt = "")
    # buildkite-agent upload figpath
    # link figure in logs if we are running on CI
    if get(ENV, "BUILDKITE", "") == "true"
        artifact_url = "artifact://$figpath"
        print("\033]1338;url='$(artifact_url)';alt='$(alt)'\a\n")
    end
end

linkfig("benchmarks/bickleyjet/volume.png", "Volume benchmarks")
linkfig("benchmarks/bickleyjet/face.png", "Face benchmarks")


text:
module SSLTestBasics
using SimpleSDMLayers
using Test

M = rand(Bool, (3,5))
S = SimpleSDMPredictor(M, 0.0, 1.0, 0.0, 1.0)

@assert longitudes(S) == LinRange(0.1, 0.9, 5)
@assert latitudes(S) == LinRange(1/6, 5/6, 3)

M = rand(Bool, (4,3))
S = SimpleSDMPredictor(M, 0.2, 1.8, -1.0, 2.0)

@assert longitudes(S) == LinRange(S.left+stride(S,1), S.right-stride(S,1), size(S,2))
@assert latitudes(S) == LinRange(S.bottom+stride(S,2), S.top-stride(S,2), size(S,1))

end


text:
struct SplitZeroBridge{T, F<:MOI.Utilities.TypedLike{T}, G<:MOI.Utilities.TypedLike{Complex{T}}} <: MOI.Bridges.Constraint.AbstractBridge
    dimension::Int
    constraint::MOI.ConstraintIndex{F, MOI.Zeros}
    real_indices::Vector{Int}
    imag_indices::Vector{Int}
end
function _nonzero_indices(func::MOI.AbstractVectorFunction)
    return [i for (i, scalar_func) in enumerate(MOIU.scalarize(func)) if !iszero(scalar_func)]
end
function MOI.Bridges.Constraint.bridge_constraint(
    ::Type{SplitZeroBridge{T, F, G}}, model::MOI.ModelLike,
    f::G,
    set::MOI.Zeros
) where {T, F, G}
    real_part = real(f)
    imag_part = MOI.Utilities.operate(imag, T, f)
    real_indices = _nonzero_indices(real_part)
    imag_indices = _nonzero_indices(imag_part)
    func = MOIU.operate(
        vcat, T,
        MOIU.eachscalar(real_part)[real_indices],
        MOIU.eachscalar(imag_part)[imag_indices]
    )
    constraint = MOI.add_constraint(model, func, MOI.Zeros(length(real_indices) + length(imag_indices)))
    return SplitZeroBridge{T, F, G}(MOI.dimension(set), constraint, real_indices, imag_indices)
end

# We don't support `MOI.VectorOfVariables` as it would be a self-loop in the bridge graph
function MOI.supports_constraint(
    ::Type{SplitZeroBridge{T}}, ::Type{<:MOI.Utilities.TypedLike{Complex{T}}},
    ::Type{MOI.Zeros}) where T
    return true
end
MOIB.added_constrained_variable_types(::Type{<:SplitZeroBridge}) = Tuple{DataType}[]
function MOIB.added_constraint_types(::Type{SplitZeroBridge{T, F, G}}) where {T, F, G}
    return Tuple{DataType, DataType}[(F, MOI.Zeros)]
end
function MOI.Bridges.Constraint.concrete_bridge_type(
    ::Type{<:SplitZeroBridge{T}}, G::Type{<:MOI.Utilities.TypedLike},
    ::Type{MOI.Zeros}) where T
    F = MA.promote_operation(imag, G)
    return SplitZeroBridge{T, F, G}
end

# Attributes, Bridge acting as a model
function MOI.get(::SplitZeroBridge{T, F},
                 ::MOI.NumberOfConstraints{F, MOI.Zeros}) where {T, F}
    return 1
end
function MOI.get(bridge::SplitZeroBridge{T, F},
                 ::MOI.ListOfConstraintIndices{F, MOI.Zeros}) where {T, F}
    return [bridge.constraint]
end

# Indices
function MOI.delete(model::MOI.ModelLike, bridge::SplitZeroBridge)
    MOI.delete(model, bridge.constraint)
end

# Attributes, Bridge acting as a constraint
function MOI.supports(
    ::MOI.ModelLike,
    ::Union{MOI.ConstraintPrimalStart, MOI.ConstraintDualStart},
    ::Type{<:SplitZeroBridge})

    return true
end
function MOI.get(model::MOI.ModelLike, attr::Union{MOI.ConstraintPrimal, MOI.ConstraintPrimalStart, MOI.ConstraintDual, MOI.ConstraintDualStart},
                 bridge::SplitZeroBridge)
    values = MOI.get(model, attr, bridge.constraint)
    output = zeros(Complex{eltype(values)}, bridge.dimension)
    for (i, idx) in enumerate(bridge.real_indices)
        output[idx] = values[i]
    end
    for (i, idx) in enumerate(bridge.imag_indices)
        output[idx] = values[length(bridge.real_indices) + i] * im
    end
    return output
end
function MOI.set(model::MOI.ModelLike, attr::Union{MOI.ConstraintPrimalStart, MOI.ConstraintDualStart},
                 bridge::SplitZeroBridge{T}, value) where T
    input = Vector{T}(undef, length(bridge.real_indices) + length(bridge.imag_indices))
    for (i, idx) in enumerate(bridge.real_indices)
        input[i] = real(value[idx])
    end
    for (i, idx) in enumerate(bridge.imag_indices)
        input[length(bridge.real_indices) + i] = imag(value[idx])
    end
    MOI.set(model, attr, bridge.constraint, input)
end


text:
# define cross product of fields of 3-vectors of type TA<:AbstractArray with size(TA,1)==3
using Tullio
import LinearAlgebra: cross
function _cross_x(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3_x[ix,iy] := v1[2,ix,iy] * v2[3,ix,iy] - v1[3,ix,iy] * v2[2,ix,iy]
end
function _cross_y(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3_y[ix,iy] := v1[1,ix,iy] * v2[3,ix,iy] - v1[3,ix,iy] * v2[1,ix,iy]
end
function _cross_z(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3_z[ix,iy] := v1[1,ix,iy] * v2[2,ix,iy] - v1[2,ix,iy] * v2[1,ix,iy]
end
function _cross(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3[i,ix,iy] := v1[mod(i-2),ix,iy] * v2[mod(i-1),ix,iy] - v1[mod(i-1),ix,iy] * v2[mod(i-2),ix,iy] (i in 1:3)
end
function _sum_cross_x(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3_x := v1[2,ix,iy] * v2[3,ix,iy] - v1[3,ix,iy] * v2[2,ix,iy]
end
function _sum_cross_y(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3_y := v1[1,ix,iy] * v2[3,ix,iy] - v1[3,ix,iy] * v2[1,ix,iy]
end
function _sum_cross_z(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3_z := v1[1,ix,iy] * v2[2,ix,iy] - v1[2,ix,iy] * v2[1,ix,iy]
end
function _sum_cross(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,3},TA2<:AbstractArray{<:Number,3}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3[i] := v1[mod(i-2),ix,iy] * v2[mod(i-1),ix,iy] - v1[mod(i-1),ix,iy] * v2[mod(i-2),ix,iy] (i in 1:3)
end
function _cross_x(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3_x[ix,iy,iz] := v1[2,ix,iy,iz] * v2[3,ix,iy,iz] - v1[3,ix,iy,iz] * v2[2,ix,iy,iz]
end
function _cross_y(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3_y[ix,iy,iz] := v1[1,ix,iy,iz] * v2[3,ix,iy,iz] - v1[3,ix,iy,iz] * v2[1,ix,iy,iz]
end
function _cross_z(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3_z[ix,iy,iz] := v1[1,ix,iy,iz] * v2[2,ix,iy,iz] - v1[2,ix,iy,iz] * v2[1,ix,iy,iz]
end
function _cross(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio v3[i,ix,iy,iz] := v1[mod(i-2),ix,iy,iz] * v2[mod(i-1),ix,iy,iz] - v1[mod(i-1),ix,iy,iz] * v2[mod(i-2),ix,iy,iz] (i in 1:3)
end
function _sum_cross_x(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3_x := v1[2,ix,iy,iz] * v2[3,ix,iy,iz] - v1[3,ix,iy,iz] * v2[2,ix,iy,iz]
end
function _sum_cross_y(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3_y := v1[1,ix,iy,iz] * v2[3,ix,iy,iz] - v1[3,ix,iy,iz] * v2[1,ix,iy,iz]
end
function _sum_cross_z(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3_z := v1[1,ix,iy,iz] * v2[2,ix,iy,iz] - v1[2,ix,iy,iz] * v2[1,ix,iy,iz]
end
function _sum_cross(v1::TA1,v2::TA2) where {TA1<:AbstractArray{<:Number,4},TA2<:AbstractArray{<:Number,4}}
        @assert size(v1,1)==3
        @assert size(v2,1)==3
        @tullio sum_v3[i] := v1[mod(i-2),ix,iy,iz] * v2[mod(i-1),ix,iy,iz] - v1[mod(i-1),ix,iy,iz] * v2[mod(i-2),ix,iy,iz] (i in 1:3)
end

using ChainRulesCore: @thunk, NoTangent
using FillArrays
function ChainRulesCore.rrule(::typeof(_cross),v₁,v₂)
        v₃ = _cross(v₁,v₂)
        function _cross_pullback(v̄₃)
                return NoTangent(), @thunk(_cross(conj(v₂),v̄₃)), @thunk(_cross(v̄₃,conj(v₁)))
        end
        return v₃, _cross_pullback
end

function ChainRulesCore.rrule(::typeof(_sum_cross),v₁::AbstractArray{<:Number,N},v₂::AbstractArray{<:Number,N}) where N
        sum_v₃ = _sum_cross(v₁,v₂)
        function _sum_cross_pullback(sum_v̄₃)
                return NoTangent(), @thunk(_cross(conj(v₂),repeat(sum_v̄₃,outer=(1,size(v₂)[2:N]...)))) , @thunk(_cross(repeat(sum_v̄₃,outer=(1,size(v₁)[2:N]...)),conj(v₁)))
        end
        return sum_v₃, _sum_cross_pullback
end

function ChainRulesCore.rrule(::typeof(_sum_cross_x),v₁::AbstractArray{<:Number,N},v₂::AbstractArray{<:Number,N}) where N
        sum_v₃_x = _sum_cross_x(v₁,v₂)
        function _sum_cross_x_pullback(sum_v̄₃_x)
                return NoTangent(), @thunk(_cross(conj(v₂),repeat([sum_v̄₃_x,0.,0.],outer=(1,size(v₂)[2:N]...)))) , @thunk(_cross(repeat([sum_v̄₃_x,0.,0.],outer=(1,size(v₁)[2:N]...)),conj(v₁)))
        end
        return sum_v₃_x, _sum_cross_x_pullback
end

function ChainRulesCore.rrule(::typeof(_sum_cross_y),v₁::AbstractArray{<:Number,N},v₂::AbstractArray{<:Number,N}) where N
        sum_v₃_y = _sum_cross_y(v₁,v₂)
        function _sum_cross_y_pullback(sum_v̄₃_y)
                return NoTangent(), @thunk(_cross(conj(v₂),repeat([0.,sum_v̄₃_y,0.],outer=(1,size(v₂)[2:N]...)))) , @thunk(_cross(repeat([0.,sum_v̄₃_y,0.],outer=(1,size(v₁)[2:N]...)),conj(v₁)))
        end
        return sum_v₃_y, _sum_cross_y_pullback
end

function ChainRulesCore.rrule(::typeof(_sum_cross_z),v₁::AbstractArray{<:Number,N},v₂::AbstractArray{<:Number,N}) where N
        sum_v₃_z = _sum_cross_z(v₁,v₂)
        function _sum_cross_z_pullback(sum_v̄₃_z)
                return NoTangent(), @thunk(_cross(conj(v₂),repeat([0.,0.,sum_v̄₃_z],outer=(1,size(v₂)[2:N]...)))) , @thunk(_cross(repeat([0.,0.,sum_v̄₃_z],outer=(1,size(v₁)[2:N]...)),conj(v₁)))
        end
        return sum_v₃_z, _sum_cross_z_pullback
end

## tests
v₁ = randn(Float64,3,20,20)
v₂ = randn(Float64,3,20,20)
v₃, v₃_pb = pullback(_cross,v₁,v₂)
Δ = randn(Float64,3,20,20)
v₃_pb(Δ)

sum_v₃, sum_v₃_pb = pullback((vv1,vv2)->sum(_cross(vv1,vv2)),v₁,v₂)
Δ = 1.0
v̄₁,v̄₂ = sum_v₃_pb(Δ)
v̄₁_FD, v̄₂_FD   =	FiniteDifferences.grad(central_fdm(3,1),(vv1,vv2)->sum(_cross(vv1,vv2)),v₁,v₂)
@assert v̄₁_FD ≈ v̄₁
@assert v̄₂_FD ≈ v̄₂

sum_v₃, sum_v₃_pb = pullback(_sum_cross,v₁,v₂)
Δ = [1.0, 1.0, 1.0]
v̄₁,v̄₂ = sum_v₃_pb(Δ)
v̄₁_FD, v̄₂_FD   =	FiniteDifferences.grad(central_fdm(3,1),(vv1,vv2)->sum(_sum_cross(vv1,vv2)),v₁,v₂)
@assert v̄₁_FD ≈ v̄₁
@assert v̄₂_FD ≈ v̄₂

sum_v₃_z, sum_v₃_z_pb = pullback(_sum_cross_z,v₁,v₂)
Δ = 1.0
v̄₁,v̄₂ = sum_v₃_z_pb(Δ)
v̄₁_FD, v̄₂_FD   =	FiniteDifferences.grad(central_fdm(3,1),(vv1,vv2)->_sum_cross_z(vv1,vv2),v₁,v₂)
@assert v̄₁_FD ≈ v̄₁
@assert v̄₂_FD ≈ v̄₂

##
v₁ = randn(ComplexF64,3,20,20)
v₂ = randn(ComplexF64,3,20,20)
v₃, v₃_pb = pullback(_cross,v₁,v₂)
Δ = randn(ComplexF64,3,20,20)
v₃_pb(Δ)

sum_v₃, sum_v₃_pb = pullback((vv1,vv2)->abs2(sum(_cross(vv1,vv2))),v₁,v₂)
Δ = Complex(1.0)
v̄₁,v̄₂ = sum_v₃_pb(Δ)
v̄₁_FD, v̄₂_FD   =	FiniteDifferences.grad(central_fdm(3,1),(vv1,vv2)->abs2(sum(_cross(vv1,vv2))),v₁,v₂)
@assert v̄₁_FD ≈ v̄₁
@assert v̄₂_FD ≈ v̄₂

# sum_v₃, sum_v₃_pb = pullback(_sum_cross,v₁,v₂)
sum_v₃, sum_v₃_pb = pullback((vv1,vv2)->abs2(sum(_sum_cross(vv1,vv2))),v₁,v₂)
Δ = Complex(1.0) # Complex.([1.0, 1.0, 1.0])
v̄₁,v̄₂ = sum_v₃_pb(Δ)
v̄₁_FD, v̄₂_FD   =	FiniteDifferences.grad(central_fdm(3,1),(vv1,vv2)->abs2(sum(_sum_cross(vv1,vv2))),v₁,v₂)
@assert v̄₁_FD ≈ v̄₁
@assert v̄₂_FD ≈ v̄₂

# sum_v₃_z, sum_v₃_z_pb = pullback(_sum_cross_z,v₁,v₂)
sum_v₃_z, sum_v₃_z_pb = pullback((vv1,vv2)->abs2(_sum_cross_z(vv1,vv2)),v₁,v₂)
Δ = Complex(1.0)
v̄₁,v̄₂ = sum_v₃_z_pb(Δ)
v̄₁_FD, v̄₂_FD   =	FiniteDifferences.grad(central_fdm(3,1),(vv1,vv2)->abs2(_sum_cross_z(vv1,vv2)),v₁,v₂)
@assert v̄₁_FD ≈ v̄₁
@assert v̄₂_FD ≈ v̄₂


text:
using Distributed, StatsBase

addprocs(4, exeflags="--project")
@everywhere begin

    using TransitionalMCMC, Distributions, LinearAlgebra
    Ndims = 5
    # Prior Bounds
    lb, ub  = -10, 10

    mean1 = ones(Ndims) * 5
    mean2 = ones(Ndims) * (-5)

    cov = 1* Matrix(I, Ndims, Ndims)

    # Prior Density and sampler
    logprior(x) = sum(logpdf(Uniform.(lb, ub), x))
    priorRnd(Nsamples) = rand(Uniform(lb, ub), Nsamples, Ndims)

    # Log Likelihood
    logLik(x) = log.(pdf(MvNormal(mean1, cov), x) + pdf(MvNormal(mean2, cov), x))

end

Nsamples = 1000

@time samps, acc = tmcmc(logLik, logprior, priorRnd, Nsamples, 5, 2)

rmprocs(workers())


text:
function StatsBase.fit(::Type{LogNormal}, m::AbstractMoments)
    # https://en.wikipedia.org/wiki/Log-normal_distribution
    n_moments(m) >= 2 || error("Need mean and variance to estimate lognormal")
    γ = 1+var(m)/mean(m)^2
    μ = log(mean(m)/sqrt(γ))
    σ = sqrt(log(γ))
    return LogNormal(μ,σ)
end

function StatsBase.fit(::Type{LogNormal}, lower::QuantilePoint, upper::QuantilePoint)
    #length(qset) == 2 || error("only implemented yet for exactly two quantiles.")
    #qset_log = [QuantilePoint(qp, q = log(qp.q)) for qp in qset]
    lower_log = QuantilePoint(lower, q = log(lower.q))
    upper_log = QuantilePoint(upper, q = log(upper.q))
    DN = fit(Normal, lower_log, upper_log)
    LogNormal(params(DN)...)
end

function fit_mean_quantile(::Type{LogNormal}, mean::Real, qp::QuantilePoint)
    # solution of
    # (1) mean = exp(mu + sigma^2/2)
    # (2) upper = mu + sigmaFac sigma
    # see R packaage lognorm inst/doc/coefLognorm.Rmd for derivation
    sigmaFac = quantile(Normal(),qp.p)
    m = log(mean)
    discr = sigmaFac^2 - 2*(log(qp.q) - m)
    (discr > 0) || error("Cannot fit LogNormal with quantile $(qp) and mean $(mean).")
    sigma = sigmaFac > sqrt(discr) ? (sigmaFac - sqrt(discr)) : sigmaFac + sqrt(discr)
    mu = m - sigma^2/2
    LogNormal(mu, sigma)
end

function fit_mode_quantile(::Type{LogNormal}, mode::Real, qp::QuantilePoint)
    # solution of
    # (1) mle = exp(mu - sigma^2)
    # (2) upper = mu + sigmaFac sigma
    # see R packaage lognorm inst/doc/coefLognorm.Rmd for derivation
    sigmaFac = quantile(Normal(),qp.p)
    m = log(mode)
    discr = sigmaFac^2/4 + (log(qp.q) - m)
    (discr > 0) || error("Cannot fit LogNormal with quantile $(qp) and mode $(mode).")
    root_discr = sqrt(discr)
    mfh = -sigmaFac/2
    sigma = mfh > root_discr ? (mfh - root_discr) : (mfh + root_discr)
    #sigma = mfh + root_discr
    mu = m + sigma^2
    LogNormal(mu, sigma)
end


text:
using StaticArrays: identity_perm
using Base: decode_overlong

include("../utils.jl")
include("group_utils.jl")

@testset "General group tests" begin
    @testset "Not implemented operation" begin
        G = GroupManifold(NotImplementedManifold(), NotImplementedOperation())
        @test repr(G) ==
              "GroupManifold(NotImplementedManifold(), NotImplementedOperation())"
        p = [1.0, 2.0]
        X = [2.0, 3.0]
        eg = Identity(G)
        @test repr(eg) === "Identity(NotImplementedOperation)"
        @test number_eltype(eg) == Bool
        @test !is_group_manifold(NotImplementedManifold(), NotImplementedOperation())
        @test is_identity(G, eg) # identity transparent
        @test_throws MethodError identity_element(G) # but for a NotImplOp there is no concrete id.
        @test isapprox(G, eg, eg)
        @test !isapprox(G, Identity(AdditionOperation()), eg)
        @test !isapprox(G, Identity(AdditionOperation()), eg)
        @test !isapprox(
            G,
            Identity(AdditionOperation()),
            Identity(MultiplicationOperation()),
        )
        @test_throws DomainError is_point(G, Identity(AdditionOperation()), true)
        @test is_point(G, eg)
        @test_throws MethodError is_identity(G, 1) # same error as before i.e. dispatch isapprox works
        @test Manifolds.check_size(G, eg) === nothing
        @test Manifolds.check_size(
            Manifolds.EmptyTrait(),
            MetricManifold(NotImplementedManifold(), EuclideanMetric()),
            eg,
        ) isa DomainError

        @test Identity(NotImplementedOperation()) === eg
        @test Identity(NotImplementedOperation) === eg
        @test !is_point(G, Identity(AdditionOperation()))
        @test !isapprox(G, eg, Identity(AdditionOperation()))
        @test !isapprox(G, Identity(AdditionOperation()), eg)

        @test NotImplementedOperation(NotImplementedManifold()) === G
        @test (NotImplementedOperation())(NotImplementedManifold()) === G

        @test_throws ErrorException hat(Rotations(3), eg, [1, 2, 3])
        @test_throws ErrorException hat!(Rotations(3), randn(3, 3), eg, [1, 2, 3])
        # If you force it, you get a not that readable MethodError
        @test_throws MethodError hat(
            GroupManifold(Rotations(3), NotImplementedOperation()),
            eg,
            [1, 2, 3],
        )

        @test_throws ErrorException vee(Rotations(3), eg, [1, 2, 3])
        @test_throws ErrorException vee!(Rotations(3), randn(3), eg, [1, 2, 3])
        @test_throws MethodError vee(
            GroupManifold(Rotations(3), NotImplementedOperation()),
            eg,
            [1, 2, 3],
        )

        @test_throws MethodError inv!(G, p, p)
        @test_throws MethodError inv!(G, p, eg)
        @test_throws MethodError inv(G, p)

        # no function defined to return the identity array representation
        @test_throws MethodError copyto!(G, p, eg)

        @test_throws MethodError compose(G, p, p)
        @test compose(G, p, eg) == p
        xO = deepcopy(p)
        compose!(G, p, eg, p)
        @test xO == p
        compose!(G, p, p, eg)
        @test xO == p
        @test_throws MethodError compose!(G, p, p, p)
        @test_throws MethodError compose!(G, p, eg, eg)

        @test_throws MethodError translate(G, p, p)
        @test_throws MethodError translate(G, p, p, LeftAction())
        @test_throws MethodError translate(G, p, p, RightAction())
        @test_throws MethodError translate!(G, p, p, p)
        @test_throws MethodError translate!(G, p, p, p, LeftAction())
        @test_throws MethodError translate!(G, p, p, p, RightAction())

        @test_throws MethodError inverse_translate(G, p, p)
        @test_throws MethodError inverse_translate(G, p, p, LeftAction())
        @test_throws MethodError inverse_translate(G, p, p, RightAction())
        @test_throws MethodError inverse_translate!(G, p, p, p)
        @test_throws MethodError inverse_translate!(G, p, p, p, LeftAction())
        @test_throws MethodError inverse_translate!(G, p, p, p, RightAction())

        @test_throws MethodError translate_diff(G, p, p, X)
        @test_throws MethodError translate_diff(G, p, p, X, LeftAction())
        @test_throws MethodError translate_diff(G, p, p, X, RightAction())
        @test_throws MethodError translate_diff!(G, X, p, p, X)
        @test_throws MethodError translate_diff!(G, X, p, p, X, LeftAction())
        @test_throws MethodError translate_diff!(G, X, p, p, X, RightAction())

        @test_throws MethodError inverse_translate_diff(G, p, p, X)
        @test_throws MethodError inverse_translate_diff(G, p, p, X, LeftAction())
        @test_throws MethodError inverse_translate_diff(G, p, p, X, RightAction())
        @test_throws MethodError inverse_translate_diff!(G, X, p, p, X)
        @test_throws MethodError inverse_translate_diff!(G, X, p, p, X, LeftAction())
        @test_throws MethodError inverse_translate_diff!(G, X, p, p, X, RightAction())

        @test_throws MethodError exp_lie(G, X)
        @test_throws MethodError exp_lie!(G, p, X)
        # no transparency error, but _log_lie missing
        @test_throws MethodError log_lie(G, p)
        @test_throws MethodError log_lie!(G, X, p)
    end

    @testset "Action direction" begin
        @test switch_direction(LeftAction()) == RightAction()
        @test switch_direction(RightAction()) == LeftAction()

        @test Manifolds._action_order(1, 2, LeftAction()) === (1, 2)
        @test Manifolds._action_order(1, 2, RightAction()) === (2, 1)
    end

    @testset "Addition operation" begin
        G = GroupManifold(NotImplementedManifold(), Manifolds.AdditionOperation())
        test_group(
            G,
            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
            [],
            [[1.0, 2.0]];
            test_exp_lie_log=false, # there is no identity element so log/exp on Lie do not work
            test_one_arg_identity_element=false,
        )

        p = [1.0, 2.0]
        X = [3.0, 4.0]
        ge = Identity(G)
        @test number_eltype(ge) == Bool
        y = allocate(p)
        copyto!(G, y, ge)
        @test y ≈ zero(p)
        @test ge - p == -p
        @test p - ge === p
        @test ge - ge === ge
        @test ge + p ≈ p
        @test p + ge ≈ p
        @test ge + ge === ge
        @test -(ge) === ge
        @test +(ge) === ge
        @test ge * 1 === ge
        @test 1 * ge === ge
        @test ge * ge === ge
        @test inv(G, p) ≈ -p
        @test inv(G, ge) === ge
        @test compose(G, p, p) ≈ p + p
        @test compose(G, p, ge) ≈ p
        @test compose(G, ge, p) ≈ p
        @test compose(G, ge, ge) == ge
        compose!(G, y, p, p)
        @test y ≈ p + p
        compose!(G, y, p, ge)
        @test y ≈ p
        compose!(G, y, ge, p)
        @test y ≈ p
        @test exp_lie(G, X) === X
        @test log_lie(G, p) === p
    end

    @testset "Multiplication operation" begin
        G = GroupManifold(NotImplementedManifold(), Manifolds.MultiplicationOperation())
        test_group(
            G,
            [[2.0 1.0; 3.0 4.0], [3.0 2.0; 4.0 5.0], [4.0 3.0; 5.0 6.0]],
            [],
            [[1.0 2.0; 3.0 4.0]];
            test_exp_lie_log=false, # no identity available as array
            test_one_arg_identity_element=false,
        )

        p = [2.0 1.0; 2.0 3.0]
        ge = Identity(G)
        @test number_eltype(ge) == Bool
        @test copyto!(G, ge, ge) === ge
        y = allocate(p)
        identity_element!(G, y)
        @test_throws DimensionMismatch identity_element!(G, [1, 2, 3])
        @test y ≈ one(p)
        @test one(ge) === ge
        @test transpose(ge) === ge
        @test det(ge) == 1
        @test ge * p ≈ p
        @test p * ge ≈ p
        @test ge * ge === ge
        @test inv(G, ge) === ge
        @test *(ge) === ge

        @test p / ge ≈ p
        @test ge \ p ≈ p
        @test ge / ge === ge
        @test ge \ ge === ge
        @test ge / p ≈ inv(G, p)
        @test p \ ge ≈ inv(G, p)
        y = allocate(p)
        @test LinearAlgebra.mul!(y, p, ge) === y
        @test y ≈ p
        y = allocate(p)
        @test LinearAlgebra.mul!(y, ge, p) === y
        @test y ≈ p
        y = allocate(p)
        @test LinearAlgebra.mul!(y, ge, ge) === y
        @test y ≈ one(y)

        @test inv(G, p) ≈ inv(p)
        @test inv(G, ge) === ge
        z = allocate(p)
        copyto!(G, z, p)
        z2 = allocate(p)
        copyto!(G.manifold, z2, p)
        @test z == z2
        X = zeros(2, 2)
        Y = allocate(X)
        copyto!(G, Y, p, X)
        Y2 = allocate(X)
        copyto!(G.manifold, Y2, p, X)
        @test Y == Y2

        @test compose(G, p, p) ≈ p * p
        @test compose(G, p, ge) ≈ p
        @test compose(G, ge, p) ≈ p
        @test compose(G, ge, ge) == ge
        compose!(G, y, p, p)
        @test y ≈ p * p
        compose!(G, y, p, ge)
        @test y ≈ p
        compose!(G, y, ge, p)
        @test y ≈ p
        X = [1.0 2.0; 3.0 4.0]
        @test exp_lie!(G, y, X) === y
        @test_throws MethodError exp_lie!(G, y, :a)
        @test y ≈ exp(X)
        Y = allocate(X)
        @test log_lie!(G, Y, y) === Y
        @test Y ≈ log(y)

        q2 = SVDMPoint(2 * Matrix{Float64}(I, 3, 3))
        mul!(q2, ge, ge)
        qT = SVDMPoint(Matrix{Float64}(I, 3, 3))
        @test isapprox(FixedRankMatrices(3, 3, 3), q2, qT)
    end

    @testset "Identity on Group Manifolds" begin
        G = TranslationGroup(3)
        e = Identity(G)
        @test get_vector_lie(G, ones(3), DefaultOrthogonalBasis()) == ones(3)
        @test e - e == e
        @test ones(3) + e == ones(3)
        e_add = Identity(AdditionOperation)
        e_mul = Identity(MultiplicationOperation)
        @test e_add * e_mul === e_add
        @test e_mul * e_add === e_add
        @test mul!(e_mul, e_mul, e_mul) === e_mul
    end
end

struct NotImplementedAction <: AbstractGroupAction{LeftAction} end

@testset "General group action tests" begin
    @testset "Not implemented operations" begin
        A = NotImplementedAction()
        p = [1.0, 2.0]
        a = [1.0, 2.0]
        X = [1.0, 2.0]

        @test_throws ErrorException apply(A, a, p)
        @test_throws ErrorException apply!(A, p, a, p)
        @test_throws ErrorException inverse_apply(A, a, p)
        @test_throws ErrorException inverse_apply!(A, p, a, p)
        @test_throws ErrorException apply_diff(A, a, p, X)
        @test_throws ErrorException apply_diff!(A, X, p, a, X)
        @test_throws ErrorException inverse_apply_diff(A, a, p, X)
        @test_throws ErrorException inverse_apply_diff!(A, X, p, a, X)
        @test_throws ErrorException compose(A, a, a)
        @test_throws ErrorException compose!(A, a, a, a)
        @test_throws ErrorException optimal_alignment(A, p, p)
        @test_throws ErrorException optimal_alignment!(A, a, p, p)
    end
end


text:
@testset "Model Methods: Fast Base (Particle)" begin
    s = ContactImplicitMPC.get_simulation("particle", "flat_3D_lc", "flat_lc")
	model = s.model

    # Setup variables
    T = Float64
    nq = model.dim.q
    q0s = rand(nq)
    q̇0s = rand(nq) # TODO: Symbolics not happy with SizedArray

    # Testing fast methods
    @test norm(ContactImplicitMPC.M_fast(model, q0s) - ContactImplicitMPC.M_func(model, q0s), Inf) < 1.0e-8
	@test norm(ContactImplicitMPC.B_fast(model, q0s) - ContactImplicitMPC.B_func(model, q0s), Inf) < 1.0e-8
    @test norm(ContactImplicitMPC.A_fast(model, q0s) - ContactImplicitMPC.A_func(model, q0s), Inf) < 1.0e-8
    @test norm(ContactImplicitMPC.C_fast(model, q0s, q̇0s) - ContactImplicitMPC.C_func(model, q0s, q̇0s), Inf) < 1.0e-8
end

@testset "Model Methods: Fast Dynamics (Particle)" begin
	s = ContactImplicitMPC.get_simulation("particle", "flat_3D_lc", "flat_lc")
    model = s.model
	env = s.env

	# Setup variables
    T = Float64
    nq = model.dim.q
    nu = model.dim.u
    nw = model.dim.w
    nc = model.dim.c
	nb = nc * friction_dim(s.env)

	hs = 0.1 * ones(1)
    q0s = rand(nq)
    q1s = rand(nq)
    u1s = rand(nu)
    u1s = rand(nu)
    w1s = rand(nw)
    γ1s = rand(nc)
    b1s = rand(nb)
	λ1s = rand(nc * dim(env))
    q2s = rand(nq)

    # Testing dynamics methods
    @test norm(ContactImplicitMPC.d_fast(model, hs, q0s, q1s, u1s, w1s, λ1s, q2s) -
        ContactImplicitMPC.dynamics(model, hs, q0s, q1s, u1s, w1s, λ1s, q2s), Inf) < 1.0e-8
end


text:
using SnoopCompile, Test, Pkg

@testset "ColorTypes" begin
    success, line, modsym = SnoopCompile.parse_call("Tuple{typeof(ColorTypes.to_top), Type{C}} where C<:(ColorTypes.Colorant{T, N} where N where T)")
    @test success
    @test modsym == :ColorTypes

    @info "Beginning @snoopc test on ColorTypes. Note failures in ColorTypes' tests do not count as failures here."
    mktempdir() do tmpdir
        tmpfile = joinpath(tmpdir, "colortypes_compiles.csv")
        tmpdir2 = joinpath(tmpdir, "precompile")
        tmpfile2 = joinpath(tmpdir, "userimg_ColorTypes.jl")

        ### Log the compiles (in a separate process)
        SnoopCompile.@snoopc tmpfile begin
            using ColorTypes, Pkg
            include(joinpath(dirname(dirname(pathof(ColorTypes))), "test", "runtests.jl"))
        end

        ### Parse the compiles and generate precompilation scripts
        let data = SnoopCompile.read(tmpfile)
            # Use these two lines if you want to create precompile functions for
            # individual packages
            pc = SnoopCompile.parcel(reverse!(data[2]))
            SnoopCompile.write(tmpdir2, pc)

            # Use these two lines if you want to add to your userimg.jl
            pc = SnoopCompile.format_userimg(reverse!(data[2]))
            SnoopCompile.write(tmpfile2, pc)
        end

        function notisempty(filename, minlength)
            @test isfile(filename)
            @test length(readlines(filename)) >= minlength
            nothing
        end

        notisempty(joinpath(tmpdir2, "precompile_ColorTypes.jl"), 100)
        notisempty(joinpath(tmpdir2, "precompile_FixedPointNumbers.jl"), 2)
        notisempty(tmpfile2, 100)
    end
end


text:
mutable struct MPPI_Logger
    trajectories::Vector{Matrix{Float64}}
    traj_costs::Vector{Float64}
    traj_weights::Vector{Float64}
end

struct MPPI_Policy_Params{M<:AbstractWeightMethod}
    num_samples::Int
    horizon::Int
    λ::Float64
    α::Float64
    U₀::Vector{Float64}
    ss::Int
    as::Int
    cs::Int
    weight_method::M
    log::Bool
end

"""
MPPI_Policy_Params(env::AbstractEnv, type::Symbol; kwargs...)
    Construct the mppi policy parameter struct
kwargs:
    - num_samples::Int = 50,
    - horizon::Int = 50,
    - λ::Float64 = 1.0,
    - α::Float64 = 1.0,
    - U₀::Vector{Float64} = [0.0],
    - cov_mat::Union{Matrix{Float64},Vector{Float64}} = [1.0],
    - weight_method::Symbol = :IT,
    - elite_threshold::Float64 = 0.8,
    - rng::AbstractRNG = Random.GLOBAL_RNG,
    - log::Bool = false,
"""
function MPPI_Policy_Params(env::AbstractEnv, type::Symbol;
    num_samples::Int = 50,
    horizon::Int = 50,
    λ::Float64 = 1.0,
    α::Float64 = 1.0,
    U₀::Vector{Float64} = [0.0],
    cov_mat::Union{Matrix{Float64},Vector{Float64}} = [1.0],
    weight_method::Symbol = :IT,
    elite_threshold::Float64 = 0.8,
    rng::AbstractRNG = Random.GLOBAL_RNG,
    log::Bool = false,
)

    ss = length(env.state)                    # State space size
    as = action_space_size(action_space(env)) # Action space size
    cs = as * horizon                         # Control size (number of actions per sample)

    if length(U₀) == as
        U₀ = repeat(U₀, horizon)
    end
    length(U₀) == cs || error("U₀ must be length of action space or control space")

    if type == :mppi
        repeat_num = 1
        check_size = as
    elseif type == :gmppi
        repeat_num = horizon
        check_size = cs
    else
        error("Incorrect type for MPPPI")
    end

    if size(cov_mat)[1] == as
        cov_mat = block_diagm(cov_mat, repeat_num)
    end
    size(cov_mat)[1] == check_size || error("Covariance matrix size problem")
    size(cov_mat)[1] == size(cov_mat)[2] || error("Covriance must be square")
    Σ = cov_mat

    if weight_method == :IT
        weight_m = Information_Theoretic(λ)
    elseif cost_method == :CE
        n = round(Int, num_samples*(1-elite_threshold))
        weight_m = Cross_Entropy(elite_threshold, n)
    else
        error("No cost method implemented for $weight_method")
    end

    log_traj = [Matrix{Float64}(undef, (horizon, ss)) for _ in 1:num_samples]
    log_traj_costs = Vector{Float64}(undef, num_samples)
    log_traj_weights = Vector{Float64}(undef, num_samples)
    mppi_logger = MPPI_Logger(log_traj, log_traj_costs, log_traj_weights)

    params = MPPI_Policy_Params(
        num_samples, horizon, λ, α, U₀, ss, as, cs,
        weight_m, log
    )
    return params, U₀, Σ, rng, mppi_logger
end

#######################################
# MPPI
#######################################
mutable struct MPPI_Policy{R<:AbstractRNG} <: AbstractPathIntegralPolicy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    rng::R
    logger::MPPI_Logger
end

function MPPI_Policy(env::AbstractEnv; kwargs...)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :mppi; kwargs...)
    return MPPI_Policy(params, env, U₀, Σ, rng, mppi_logger)
end

function (pol::MPPI_Policy)(env::AbstractEnv)
    K, T = pol.params.num_samples, pol.params.horizon
    as, cs = pol.params.as, pol.params.cs
    trajectory_cost, E = calculate_trajectory_costs(pol, env)

    # Compute weights based on weight method
    weights = compute_weights(pol.params.weight_method, trajectory_cost)
    weights = reshape(weights, K, 1)

    # Weight the noise based on the calcualted weights
    weighted_noise = zeros(Float64, cs)
    for t ∈ 1:T
        for k ∈ 1:K
            weighted_noise[((t-1)*as+1):(t*as)] += weights[k] .* E[k,t]
        end
    end
    weighted_controls = pol.U + weighted_noise
    control = get_controls_roll_U!(pol, weighted_controls)

    if pol.params.log
        pol.logger.traj_costs = trajectory_cost
        pol.logger.traj_weights = vec(weights)
    end

    return control
end

function calculate_trajectory_costs(pol::MPPI_Policy, env::AbstractEnv)
    K, T = pol.params.num_samples, pol.params.horizon
    as = pol.params.as
    γ = pol.params.λ*(1-pol.params.α)

    # Get samples for which our trajectories will be defined
    P = Distributions.MvNormal(pol.Σ)
    E = rand(pol.rng, P, K, T)
    Σ_inv = Distributions.invcov(P)

    trajectory_cost = zeros(Float64, K)
    Threads.@threads for k ∈ 1:K
        sim_env = copy(env) # Slower, but allows for multi threading
        for t ∈ 1:T
            Eᵢ = E[k,t]
            uₜ = pol.U[((t-1)*as+1):(t*as)]
            Vₜ = uₜ + Eᵢ
            control_cost = γ * uₜ'*Σ_inv*Eᵢ
            model_controls = get_model_controls(action_space(sim_env), Vₜ)
            sim_env(model_controls)
            # Subtrating based on "reward", Adding based on "cost"
            trajectory_cost[k] = trajectory_cost[k] - reward(sim_env) + control_cost
            if pol.params.log
                pol.logger.trajectories[k][t, :] = sim_env.state
            end
        end
    end
    return trajectory_cost, E
end

#######################################
# GMPPI Policies
#######################################
function (pol::AbstractGMPPI_Policy)(env::AbstractEnv)
    cs = pol.params.cs
    trajectory_cost, E, weights = calculate_trajectory_costs(pol, env)

    # Weight the noise based on the calcualted weights
    weighted_noise = zeros(Float64, cs)
    for rᵢ ∈ 1:cs
        weighted_noise[rᵢ] = weights' * E[rᵢ,:]
    end
    weighted_controls = pol.U + weighted_noise
    control = get_controls_roll_U!(pol, weighted_controls)

    if pol.params.log
        pol.logger.traj_costs = trajectory_cost
        pol.logger.traj_weights = weights
    end
    return control
end

function simulate_model(pol::AbstractGMPPI_Policy, env::AbstractEnv,
    E::Matrix{Float64}, Σ_inv::Matrix{Float64}, U_orig::Vector{Float64},
    n::Int=1,
)
    K, T = pol.params.num_samples, pol.params.horizon
    as = pol.params.as
    γ = pol.params.λ*(1-pol.params.α)

    trajectory_cost = zeros(Float64, K)
    Threads.@threads for k ∈ 1:K
        sim_env = copy(env) # Slower, but allows for multi threading
        Vₖ = pol.U + E[:,k]
        control_cost = γ * U_orig'*Σ_inv*(Vₖ .- U_orig)
        model_controls = get_model_controls(action_space(sim_env), Vₖ, T)
        trajectory_cost[k] = rollout_model(sim_env, T, model_controls, pol, k, n)
        trajectory_cost[k] += control_cost  # Adding based on "cost"
    end
    return trajectory_cost
end

"""
GMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    Generalized MPPI policy strcut
"""
mutable struct GMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    rng::R
    logger::MPPI_Logger
end

"""
GMPPI_Policy(env::AbstractEnv; kwargs...)
    - env::AbstractEnv
kwargs are passed to MPPI_Policy_params
"""
function GMPPI_Policy(env::AbstractEnv; kwargs...)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :gmppi; kwargs...)
    return GMPPI_Policy(params, env, U₀, Σ, rng, mppi_logger)
end

function calculate_trajectory_costs(pol::GMPPI_Policy, env::AbstractEnv)
    K = pol.params.num_samples

    # Get samples for which our trajectories will be defined
    P = Distributions.MvNormal(pol.Σ)
    E = rand(pol.rng, P, K)
    Σ_inv = Distributions.invcov(P)

    # Use the samples to simulate our model to get the costs
    trajectory_cost = simulate_model(pol, env, E, Σ_inv, pol.U)
    weights = compute_weights(pol.params.weight_method, trajectory_cost)
    return trajectory_cost, E, weights
end

"""
CEMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    Cross-Entropy version of MPOPI
"""
mutable struct CEMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    opt_its::Int
    ce_elite_threshold::Float64
    Σ_estimation_method::Union{LinearShrinkage, SimpleCovariance}
    rng::R
    logger::MPPI_Logger
end

"""
CEMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    ce_elite_threshold::Float64 = 0.8,
    Σ_est::Symbol = :mle,
    kwargs...
kwargs passed to MPPI_Policy_Params

Options for Σ_est
    - :mle = maximum liklihood estimation
    - :lw = Lediot & Wolf (http://www.ledoit.net/honey.pdf)
    - :ss = Schaffer & Strimmer (https://strimmerlab.github.io/)
    - :rblw = Rao-Blackwell estimator (https://arxiv.org/pdf/0907.4698.pdf)
    - :oas = Oracle-Approximating (https://arxiv.org/pdf/0907.4698.pdf)
https://mateuszbaran.github.io/CovarianceEstimation.jl/dev/man/lshrink/
"""
function CEMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    ce_elite_threshold::Float64 = 0.8,
    Σ_est::Symbol = :mle,
    kwargs...
)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :gmppi; kwargs...)
    if Σ_est == :mle
        Σ_est_method = SimpleCovariance()
    elseif  Σ_est == :lw
        Σ_est_method = LinearShrinkage(DiagonalUnequalVariance(), :lw)
    elseif Σ_est == :ss
        Σ_est_method = LinearShrinkage(DiagonalUnequalVariance(), :ss)
    elseif Σ_est == :rblw
        Σ_est_method = LinearShrinkage(DiagonalCommonVariance(), :rblw)
    elseif Σ_est == :oas
        Σ_est_method = LinearShrinkage(DiagonalCommonVariance(), :oas)
    else
        error("CEMPPI_Policy - Not a valid Σ estimation method")
    end
    pol = CEMPPI_Policy(params, env, U₀,
              Σ, opt_its, ce_elite_threshold,
              Σ_est_method, rng, mppi_logger,
          )
    return pol
end

function calculate_trajectory_costs(pol::CEMPPI_Policy, env::AbstractEnv)
    K = pol.params.num_samples
    N = pol.opt_its
    m_elite = round(Int, K*(1-pol.ce_elite_threshold))

    # Initial covariance of distribution
    U_orig = pol.U
    Σ′ = pol.Σ

    trajectory_cost = zeros(Float64, K)
    # Optimize sample distribution and get trajectory costs
    for n ∈ 1:N
        # Get samples for which our trajectories will be defined
        P = Distributions.MvNormal(Σ′)
        E = rand(pol.rng, P, K)
        Σ_inv = Distributions.invcov(P)

        # Use the samples to simulate our model to get the costs
        trajectory_cost = simulate_model(pol, env, E, Σ_inv, U_orig)
        if n < N
            # Select elite samples, fit new distribution
            order = sortperm(trajectory_cost)
            elite = E[:, order[1:m_elite]]

            elite_traj_cost = trajectory_cost[order[1:m_elite]]
            if maximum(abs.(diff(elite_traj_cost, dims=1))) < 10e-3
                break
            end

            # (μ′, Σ′) = StatsBase.mean_and_cov(elite, 2)
            # Σ′ = Σ′ + + 10e-9*I
            # pol.U = pol.U + vec(μ′)

            # Transposing elite based on needed format (n x p)
            Σ′ = cov(pol.Σ_estimation_method, elite') + 10e-9*I
            pol.U = pol.U + vec(mean(elite, dims=2))
        end
    end
    E = E .+ (pol.U - U_orig)
    pol.U = U_orig
    weights = compute_weights(pol.params.weight_method, trajectory_cost)
    return trajectory_cost, E, weights
end

"""
CMAMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    Covariance Matrix Adaptation version of MPOPI
"""
mutable struct CMAMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    opt_its::Int
    σ::Float64
    m_elite::Int
    ws::Vector{Float64}
    μ_eff::Float64
    cσ::Float64
    dσ::Float64
    cΣ::Float64
    c1::Float64
    cμ::Float64
    E::Float64
    rng::R
    logger::MPPI_Logger
end

"""
CMAMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    σ::Float64 = 1.0,
    elite_perc_threshold::Float64 = 0.8,
    kwargs...
kwargs passed to MPPI_Policy_Params
"""
function CMAMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    σ::Float64 = 1.0,
    elite_perc_threshold::Float64 = 0.8,
    kwargs...
)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :gmppi; kwargs...)
    m = params.num_samples
    n = params.cs
    m_elite = round(Int, (1.0-elite_perc_threshold)*m)
    ws = log((m+1)/2) .- log.(1:m)
    ws[1:m_elite] ./= sum(ws[1:m_elite])
    μ_eff = 1 / sum(ws[1:m_elite].^2)
    cσ = (μ_eff + 2)/(n + μ_eff + 5)
    dσ = 1 + 2max(0, sqrt((μ_eff-1)/(n+1))-1) + cσ
    cΣ = (4 + μ_eff/n)/(n + 4 + 2μ_eff/n)
    c1 = 2/((n+1.3)^2 + μ_eff)
    cμ = min(1-c1, 2*(μ_eff-2+1/μ_eff)/((n+2)^2 + μ_eff))
    ws[m_elite+1:end] .*= -(1 + c1/cμ)/sum(ws[m_elite+1:end])
    E = n^0.5*(1-1/(4n)+1/(21*n^2))

    pol = CMAMPPI_Policy(params, env, U₀, Σ, opt_its, σ, m_elite,
        ws, μ_eff, cσ, dσ, cΣ, c1, cμ, E, rng, mppi_logger)
    return pol
end

function calculate_trajectory_costs(pol::CMAMPPI_Policy, env::AbstractEnv)
    K = pol.params.num_samples
    N = pol.opt_its
    cs = pol.params.cs
    σ = pol.σ
    m_elite = pol.m_elite
    ws, μ_eff, cσ, dσ = pol.ws, pol.μ_eff, pol.cσ, pol.dσ
    cΣ, c1, cμ, E_cma = pol.cΣ, pol.c1, pol.cμ, pol.E

    # Initial covariance of distribution
    U_orig = pol.U
    Σ = pol.Σ

    pσ, pΣ = zeros(pol.params.cs), zeros(pol.params.cs)
    trajectory_cost = zeros(Float64, K)
    # Optimize sample distribution and get trajectory costs
    for n ∈ 1:N
        # Get samples for which our trajectories will be defined
        if N > 1
            P = Distributions.MvNormal(σ^2*Σ)
        else
            P = Distributions.MvNormal(Σ)
        end
        Σ_inv = Distributions.invcov(P)
        E = rand(pol.rng, P, K)

        # Use the samples to simulate our model to get the costs
        trajectory_cost = simulate_model(pol, env, E, Σ_inv, U_orig)

        if n < N
            # Reorder, select elite samples, fit new distribution
            order = sortperm(trajectory_cost)
            elite_E = E[:, order[1:m_elite]]

            elite_traj_cost = trajectory_cost[order[1:m_elite]]
            if maximum(abs.(diff(elite_traj_cost, dims=1))) < 10e-3
                break
            end

            # selection and mean update
            δs = elite_E/σ
            δw = zeros(Float64, cs)
            for rᵢ ∈ 1:cs
                δw[rᵢ] = ws[1:m_elite]' * elite_E[rᵢ,:]
            end
            pol.U += σ*δw

            # step-size control
            C = Σ^-0.5
            pσ = (1-cσ)*pσ + sqrt(cσ*(2-cσ)*μ_eff)*C*δw
            σ *= exp(cσ/dσ * (norm(pσ)/E_cma - 1))

            # covariance adaptation
            hσ = Int(norm(pσ)/sqrt(1-(1-cσ)^(2n)) < (1.4+2/(cs+1))*E_cma)
            pΣ = (1-cΣ)*pΣ + hσ*sqrt(cΣ*(2-cΣ)*μ_eff)*δw

            temp_sum = 0
            for ii in 1:K
                if ws[ii] ≥ 0
                    w0 = ws[ii]
                else
                    w0 = n*ws[ii]/norm(C*δs[order[ii]])^2
                end
                temp_sum += w0*δs[order[ii]]*δs[order[ii]]'
            end

            Σ = (1-c1-cμ)*Σ + c1*(pΣ*pΣ' + (1-hσ)*cΣ*(2-cΣ)*Σ) .+ cμ*temp_sum
            Σ = triu(Σ)+triu(Σ,1)' # enforce symmetry
        end
    end
    E = E .+ (pol.U - U_orig)
    pol.U = U_orig
    weights = compute_weights(pol.params.weight_method, trajectory_cost)
    return trajectory_cost, E, weights
end

"""
μAISMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    Simple, mean-only M-PMC AIS version of MPOPI
"""
mutable struct μAISMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    opt_its::Int
    λ_ais::Float64
    rng::R
    logger::MPPI_Logger
end

"""
μAISMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    λ_ais::Float64 = 20.0,
    kwargs...
kwargs passed to MPPI_Policy_Params
"""
function μAISMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    λ_ais::Float64 = 20.0,
    kwargs...
)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :gmppi; kwargs...)
    return μAISMPPI_Policy(params, env, U₀, Σ, opt_its, λ_ais, rng, mppi_logger)
end

"""
    calculate_trajectory_costs(policy::μAISMPPI_Policy, env::AbstractEnv)
Simple AIS strategy in which only the new mean is adapted at each
iteration. New samples are then taken from the new distribution.
"""
function calculate_trajectory_costs(pol::μAISMPPI_Policy, env::AbstractEnv)
    K = pol.params.num_samples
    N = pol.opt_its
    weight_method = Information_Theoretic(pol.λ_ais)

    U_orig = pol.U
    P = Distributions.MvNormal(pol.Σ)
    Σ_inv = Distributions.invcov(P)

    trajectory_cost = Vector{Float64}(undef, K)
    ws = Vector{Float64}(undef, K)
    # Optimize sample distribution and get trajectory costs
    for n ∈ 1:N
        E = rand(pol.rng, P, K)
        trajectory_cost = simulate_model(pol, env, E, Σ_inv, U_orig)
        if n < N
            ws = compute_weights(weight_method, trajectory_cost)
            pw = StatsBase.ProbabilityWeights(ws)
            (μ′, Σ′) = StatsBase.mean_and_cov(E, pw, 2)
            pol.U = pol.U + vec(μ′)
        else
            ws = compute_weights(pol.params.weight_method, trajectory_cost)
        end
    end
    E = E .+ (pol.U - U_orig)
    pol.U = U_orig
    return trajectory_cost, E, ws
end

"""
μΣAISMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    Mean and Covariance M-PMC AIS with one distribution version of MPOPI
"""
mutable struct μΣAISMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    opt_its::Int
    λ_ais::Float64
    rng::R
    logger::MPPI_Logger
end

"""
μΣAISMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    λ_ais::Float64 = 20.0,
    kwargs...
    kwargs passed to MPPI_Policy_Params
"""
function μΣAISMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    λ_ais::Float64 = 20.0,
    kwargs...
)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :gmppi; kwargs...)
    return μΣAISMPPI_Policy(params, env, U₀, Σ, opt_its, λ_ais, rng, mppi_logger)
end

"""
    calculate_trajectory_costs(policy::μΣAISMPPI_Policy, env::AbstractEnv)
Simple AIS strategy in which the new mean and covariance are adapted at each
iteration. New samples are then taken from the new distribution.
"""
function calculate_trajectory_costs(pol::μΣAISMPPI_Policy, env::AbstractEnv)
    K = pol.params.num_samples
    N = pol.opt_its
    weight_method = Information_Theoretic(pol.λ_ais)

    # Initial covariance of distribution
    U_orig = pol.U
    Σ′ = pol.Σ

    trajectory_cost = Vector{Float64}(undef, K)
    ws = Vector{Float64}(undef, K)
    # Optimize sample distribution and get trajectory costs
    for n ∈ 1:N
        # Get samples for which our trajectories will be defined
        P = Distributions.MvNormal(Σ′)
        E = rand(pol.rng, P, K)
        Σ_inv = Distributions.invcov(P)

        # Use the samples to simulate our model to get the costs
        trajectory_cost = simulate_model(pol, env, E, Σ_inv, U_orig)
        if n < N
            ws = compute_weights(weight_method, trajectory_cost)
            pw = StatsBase.ProbabilityWeights(ws)
            (μ′, Σ′) = StatsBase.mean_and_cov(E, pw, 2)
            Σ′ = Σ′ + + 10e-9*I
            pol.U = pol.U + vec(μ′)
        else
            ws = compute_weights(pol.params.weight_method, trajectory_cost)
        end
    end
    E = E .+ (pol.U - U_orig)
    pol.U = U_orig
    return trajectory_cost, E, ws
end

"""
PMCMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    PMC with multinomial resampling AIS with one distribution version of MPOPI
"""
mutable struct PMCMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    opt_its::Int
    λ_ais::Float64
    rng::R
    logger::MPPI_Logger
end

"""
PMCMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    λ_ais::Float64 = 20.0,
    kwargs...
kwargs passed to MPPI_Policy_Params
"""
function PMCMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    λ_ais::Float64 = 20.0,
    kwargs...
)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :gmppi; kwargs...)
    return PMCMPPI_Policy(params, env, U₀, Σ, opt_its, λ_ais, rng, mppi_logger)
end

"""
    calculate_trajectory_costs(policy::PMCMPPI_Policy, env::AbstractEnv)
Generic PMC strategy.
O Cappé, A Guillin, J. M Marin & C. P Robert (2004)
Population Monte Carlo, Journal of Computational and Graphical Statistics,
13:4, 907-929, DOI: 10.1198/106186004X12803
"""
function calculate_trajectory_costs(pol::PMCMPPI_Policy, env::AbstractEnv)
    K = pol.params.num_samples
    N = pol.opt_its
    weight_method = Information_Theoretic(pol.λ_ais)

    # Initial covariance of distribution
    U_orig = pol.U
    Σ′ = pol.Σ

    trajectory_cost = Vector{Float64}(undef, K)
    ws = Vector{Float64}(undef, K)
    # Optimize sample distribution and get trajectory costs
    for n ∈ 1:N
        # Get samples for which our trajectories will be defined
        P = Distributions.MvNormal(Σ′)
        E = rand(pol.rng, P, K)
        Σ_inv = Distributions.invcov(P)

        # Use the samples to simulate our model to get the costs
        trajectory_cost = simulate_model(pol, env, E, Σ_inv, U_orig)
        if n < N
            ws = compute_weights(weight_method, trajectory_cost)
            resample_cat_dist = Categorical(ws)
            resample_idxs = rand(pol.rng, resample_cat_dist, K)
            E′ = E[:, resample_idxs]
            (μ′, Σ′) = StatsBase.mean_and_cov(E′, 2)
            Σ′ = Σ′ + + 10e-9*I
            pol.U = pol.U + vec(μ′)
        else
            ws = compute_weights(pol.params.weight_method, trajectory_cost)
        end
    end
    E = E .+ (pol.U - U_orig)
    pol.U = U_orig
    return trajectory_cost, E, ws
end

"""
NESMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    Natual evolution strategy version of MPOPI
"""
mutable struct NESMPPI_Policy{R<:AbstractRNG} <: AbstractGMPPI_Policy
    params::MPPI_Policy_Params
    env::AbstractEnv
    U::Vector{Float64}
    Σ::Matrix{Float64}
    A::Matrix{Float64}
    opt_its::Int
    step_factor::Float64
    rng::R
    logger::MPPI_Logger
end

"""
NESMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    step_factor::Float64 = 0.01,
    kwargs...
kwargs passed to MPPI_Policy_Params
"""
function NESMPPI_Policy(env::AbstractEnv;
    opt_its::Int = 10,
    step_factor::Float64 = 0.01,
    kwargs...
)
    params, U₀, Σ, rng, mppi_logger = MPPI_Policy_Params(env, :gmppi; kwargs...)
    A = sqrt(Σ)
    pol = NESMPPI_Policy(params, env, U₀, Σ, A, opt_its, step_factor, rng, mppi_logger)
    return pol
end

function calculate_trajectory_costs(pol::NESMPPI_Policy, env::AbstractEnv)
    K = pol.params.num_samples
    N = pol.opt_its

    # Initial covariance and principal matrix
    U_orig = pol.U
    Σ′ = pol.Σ
    A′ = pol.A

    trajectory_cost = zeros(Float64, K)
    # Optimize sample distribution and get trajectory costs
    for n ∈ 1:N
        # Get samples for which our trajectories will be defined
        P = Distributions.MvNormal(Σ′)
        E = rand(pol.rng, P, K)
        Σ_inv = Distributions.invcov(P)

        # Use the samples to simulate our model to get the costs
        trajectory_cost = simulate_model(pol, env, E, Σ_inv, U_orig)
        if maximum(abs.(diff(trajectory_cost, dims=1))) < 10e-3
            break
        end

        if n < N
            ∇μlog_p_x = zeros(Float64, size(pol.U))
            ∇Alog_p_x = zeros(Float64, size(pol.Σ))
            for k ∈ 1:K
                ∇μlog_p_x .+= Σ_inv * E[:,k] .* trajectory_cost[k]
                ∇Σlog_p_x = 1/2*Σ_inv*E[:,k]*E[:,k]'*Σ_inv - 1/2*Σ_inv
                ∇Alog_p_x += A′*(∇Σlog_p_x + ∇Σlog_p_x') * trajectory_cost[k]
            end
            A′ -= pol.step_factor/K .* ∇Alog_p_x ./ K
            Σ′ = A′' * A′
            pol.U -= pol.step_factor/K .* ∇μlog_p_x
        end
    end
    E = E .+ (pol.U - U_orig)
    pol.U = U_orig
    weights = compute_weights(pol.params.weight_method, trajectory_cost)
    return trajectory_cost, E, weights
end

max_stars_repo_name: sisl/MPOPIS

max_issues_repo_name: sisl/MPOPIS

max_forks_repo_name: sisl/MPOPIS


text:
@testset "nextjournal" begin
    @testset "A first example" begin
        @polyvar x y

        # define f
        f = [x^2 + 2y, y^2 - 2]

        # solve f
        result = solve(f)
        @test nfinite(result) == 4
    end
    @testset "Start systems" begin
        @polyvar x y

        f = [x^2 + 2y, y^2 - 2]
        g = [x^2 - 1, y^2 - 1]

        start_sols = [[1, 1], [1, -1], [-1, 1], [-1, -1]]

        result = solve(g, f, start_sols)
        @test nfinite(result) == 4
        @test length(real_solutions(result)) == 2
    end
    @testset "Groups of variables" begin
        @polyvar x y z
        J = 3 * x^3 * y + y^2 * z^2 - 2x * y - x * 4 * z^3
        g = x^4 + y^4 + z^4 - 1

        # Introduce auxillary variable for Lagrangian
        @polyvar λ

        # define Lagrangian
        L = J - λ * g
        import DynamicPolynomials: differentiate
        ∇L = differentiate(L, [x, y, z, λ])

        # Now we solve the polynomial system ∇L = 0
        result = solve(∇L; show_progress = false)
        @test nfinite(result) == 108

        result = solve(∇L, variable_groups = [(x, y, z), (λ,)], show_progress = false)
        @test nfinite(result) == 108

        reals = real_solutions(result)
        minval, minindex = findmin(map(s -> J(s[1:3]), reals))
        minarg = reals[minindex][1:3]
        @test minarg ≈ [
            0.6893448348668392
            0.19072105130305433
            0.9376180557378104
        ] || minarg ≈ -[0.6893448348668392; 0.19072105130305433; 0.9376180557378104]
    end
    @testset "6R inverse problem" begin
        # initialize the variables
        @polyvar z[1:6, 1:3] p[1:3]
        α = randn(5)
        a = randn(9)

        # define the system of polynomials
        f = [z[i, :] ⋅ z[i, :] for i = 2:5]
        g = [z[i, :] ⋅ z[i+1, :] for i = 1:5]
        h =
            sum(a[i] .* (z[i, :] × z[i+1, :]) for i = 1:3) +
            sum(a[i+4] .* z[i, :] for i = 2:5)
        F′ = [f .- 1; g .- cos.(α); h .- p]

        # assign values to z₁ and z₆
        z₁ = normalize!(randn(3))
        z₆ = normalize!(randn(3))
        F = [subs(f, z[1, :] => z₁, z[6, :] => z₆, p => [1, 1, 0]) for f in F′]

        result = solve(F, show_progress = false)
        @test nfinite(result) == 16

        variable_groups = [[z[2, :]; z[4, :]], [z[3, :]; z[5, :]]]
        result = solve(F; variable_groups = variable_groups, show_progress = false)
        @test nfinite(result) == 16

        p_rand = randn(ComplexF64, 3)
        F_rand = [subs(f, z[1, :] => z₁, z[6, :] => z₆, p => p_rand) for f in F′]
        R_rand = solve(F_rand, variable_groups = variable_groups, show_progress = false)

        F̂ = [subs(f, z[1, :] => z₁, z[6, :] => z₆) for f in F′]
        q = [2, 3, 4]
        result = solve(
            F̂,
            solutions(R_rand);
            parameters = p, start_parameters = p_rand, target_parameters = q,
        )
        @test nfinite(result) == 16
    end
    @testset "Monodromy Method" begin
        @polyvar a[1:3] μ[1:3] σ²[1:3]

        m0 = a[1] + a[2] + a[3]
        m1 = a[1] * μ[1] + a[2] * μ[2] + a[3] * μ[3]
        m2 = a[1] * (μ[1]^2 + σ²[1]) + a[2] * (μ[2]^2 + σ²[2]) + a[3] * (μ[3]^2 + σ²[3])
        m3 = a[1] * (μ[1]^3 + 3 * σ²[1] * μ[1]) +
            a[2] * (μ[2]^3 + 3 * σ²[2] * μ[2]) +
            a[3] * (μ[3]^3 + 3 * σ²[3] * μ[3])
        m4 = a[1] * (μ[1]^4 + 6 * σ²[1] * μ[1]^2 + 3 * σ²[1]^2) +
            a[2] * (μ[2]^4 + 6 * σ²[2] * μ[2]^2 + 3 * σ²[2]^2) +
            a[3] * (μ[3]^4 + 6 * σ²[3] * μ[3]^2 + 3 * σ²[3]^2)
        m5 = a[1] * (μ[1]^5 + 10 * σ²[1] * μ[1]^3 + 15 * μ[1] * σ²[1]^2) +
            a[2] * (μ[2]^5 + 10 * σ²[2] * μ[2]^3 + 15 * μ[2] * σ²[2]^2) +
            a[3] * (μ[3]^5 + 10 * σ²[3] * μ[3]^3 + 15 * μ[3] * σ²[3]^2)
        m6 = a[1] * (μ[1]^6 + 15 * σ²[1] * μ[1]^4 + 45 * μ[1]^2 * σ²[1]^2 + 15 * σ²[1]^3) +
            a[2] * (μ[2]^6 + 15 * σ²[2] * μ[2]^4 + 45 * μ[2]^2 * σ²[2]^2 + 15 * σ²[2]^3) +
            a[3] * (μ[3]^6 + 15 * σ²[3] * μ[3]^4 + 45 * μ[3]^2 * σ²[3]^2 + 15 * σ²[3]^3)
        m7 = a[1] * (
            μ[1]^7 +
                21 * σ²[1] * μ[1]^5 +
                105 * μ[1]^3 * σ²[1]^2 +
                105 * μ[1] * σ²[1]^3
        ) +
            a[2] * (
            μ[2]^7 +
                21 * σ²[2] * μ[2]^5 +
                105 * μ[2]^3 * σ²[2]^2 +
                105 * μ[2] * σ²[2]^3
        ) +
            a[3] * (
            μ[3]^7 +
                21 * σ²[3] * μ[3]^5 +
                105 * μ[3]^3 * σ²[3]^2 +
                105 * μ[3] * σ²[3]^3
        )
        m8 = a[1] * (
            μ[1]^8 +
                28 * σ²[1] * μ[1]^6 +
                210 * μ[1]^4 * σ²[1]^2 +
                420 * μ[1]^2 * σ²[1]^3 +
                105 * σ²[1]^4
        ) +
            a[2] * (
            μ[2]^8 +
                28 * σ²[2] * μ[2]^6 +
                210 * μ[2]^4 * σ²[2]^2 +
                420 * μ[2]^2 * σ²[2]^3 +
                105 * σ²[2]^4
        ) +
            a[3] * (
            μ[3]^8 +
                28 * σ²[3] * μ[3]^6 +
                210 * μ[3]^4 * σ²[3]^2 +
                420 * μ[3]^2 * σ²[3]^3 +
                105 * σ²[3]^4
        )

        f = [m0, m1, m2, m3, m4, m5, m6, m7, m8]

        p₀ = [1; -1; 3; -5.5; 22.45; -50.75; 243.325; -635.725; 3420.7375]

        @test bezout_number(f - p₀) == 362880

        a₀ = randn(ComplexF64, 3)
        a₀ = a₀ / sum(a₀)
        μ₀ = rand(ComplexF64, 3)
        σ²₀ = rand(ComplexF64, 3)
        start_sol = [a₀; μ₀; σ²₀]

        q₀ = [m([a; μ; σ²] => start_sol) for m in f]
        @polyvar p[1:9]
        R = monodromy_solve(
            f - p,
            start_sol,
            q₀;
            parameters = p, target_solutions_count = 1350, show_progress = false,
        )
        R2 = solve(
            f - p,
            solutions(R);
            parameters = p,
            start_parameters = q₀,
            target_parameters = p₀,
            precision = PRECISION_ADAPTIVE,
            show_progress = false,
        )

        all_real_sols = real_solutions(R2)
        true_real_solutions = filter(s -> all(s[7:9] .> -1e-3), all_real_sols)

        S₃ = SymmetricGroup(3)
        relabeling =
            GroupActions(v -> map(p -> (v[1:3][p]..., v[4:6][p]..., v[7:9][p]...), S₃))
        mults = multiplicities(true_real_solutions, group_action = relabeling)
        @test length(mults) ≥ 1

        R_with_group_action = monodromy_solve(
            f - p,
            start_sol,
            q₀;
            parameters = p,
            group_action = relabeling,
            show_progress = false,
            target_solutions_count = 225,
        )

        @test all(length.([relabeling(s) for s in solutions(R_with_group_action)]) .== 6)
    end
end


text:
using Plots

percentile = [ 50,  66,   75,   80,  90,    95,   98,   99,   100]
jul        = [106, 115, 1796, 3366, 4839, 4885, 6652, 6681,  6704]
nodejs     = [610, 737, 1458, 1632, 1701, 1898, 1945, 1952,  2059]
apache     = [ 89,  95,   99,  102,  151,  424, 1075, 1085, 21092]

data = hcat(jul./1000, nodejs./1000, apache./1000)

p = plot(percentile, data, title="Request Latency", xlabel="Percentile of requests", ylabel="Latency in seconds", label=["Julia", "Nodejs", "Apache"])
display(p)
savefig(p, "C10k.png")


text:
using DifferentialEquations: include

#############################################################
# Running Simulations
#############################################################

#Include KG_Module to run simulation
include("../../src/KG_Module.jl")

#Helper function
function Kepler_Frequency( Rbh , R)
    return sqrt(Rbh/R^3)
end

function radius_parameters( Rbh)
    println("BH radius  : " , Rbh)
    println("LightRing  : " , 3*Rbh/2)
    println("ISCO radius: " , 3*Rbh)
    println("Safe Orbit : " , 3*Rbh)
end

#####################
# SIMULATION
#####################

#Setup masses
p0 = Param(

    #Temporal variables
    tmax           = 60000.0     ,
    out_every_t    = 2.5        ,
    deltat         = 0.05       ,

    #Spacial variables
    xmin           =  -30.0     ,
    xmax           =  30.0      ,
    xnodes         =  256       ,

    #Spacial variables
    ymin           =  -30.0     ,
    ymax           =  30.0      ,
    ynodes         =   256      ,

    #Initial Conditions
    A0              =   3.5     ,
    σ               =   1.5     ,
    r0              =   17.0    ,
    ω               =   0.5     ,
    m               =   2       ,

    #Klein Gordon Mass
    μ               =   0.0     ,

    #Dissipation terms
    dissipation     =  :true    ,

    R_orbit         = 3.34716475041       ,
    R1              = 1.5       ,
    R2              = 0.0,#1.5       ,

    Ω               = Kepler_Frequency( 1.5 , 3.34716475041) , # Kepler's law
    alpha1          = 10.0       ,
    alpha2          = 0.0, #10.0       ,

    #Periodic BC?
    Boundaries      = :radial    , # :square :periodic

    #Data variables
    folder          = "/2020_07_01/"  ,       #must have a bar before the folder
    fname           = "default_name"

)

#Binary BH with a certain orbital radius
p1 = Param(     p0,

                R1              = 1.5       ,
                R2              = 1.5       ,
                R_orbit         = 4.5       ,

                alpha1          = 10.0          ,
                alpha2          = 10.0           ,

                Ω               = Kepler_Frequency( 1.5 , 4.5) , # Kepler's law

                fname = "bbh_R1_1.5_Alpha_10.0_Rorbit_4.5_Rout_27.0_tmax_50000"
            )

p2 = Param(     p0,

            R1              = 1.5       ,
            R2              = 0.0       ,
            R_orbit         = 0.0       ,

            alpha1          = 10.0          ,
            alpha2          = 0.0           ,

            Ω               = Kepler_Frequency( 1.5 , 4.5) , # Kepler's law

            fname = "sbh_R1_1.5_Alpha_10.0_Omega_0.13_Rout_27.0_tmax_50000"
        )

#s = [p1, p2]

#for x in s
#    solve_wave_equation_2D(x)
#end

#include("../../src/KG_Plot_Module.jl")

#slices_observer(p2)
#slices_observer("./data/2020_07_01_bhb_sbh_comparison/default_name4.h5")