text:
module DocumenterEpub

import Markdown
using Dates: Dates, now, UTC

import Documenter
import Documenter.Documents
import Documenter.Writers
import Documenter.Writers.FormatSelector
import Documenter.Anchors
import Documenter.Selectors
import Documenter.Documenter
import Documenter.Utilities
using Documenter.Utilities.DOM: DOM, Tag, @tags
using Documenter.Utilities: Default
using Documenter.Utilities.MDFlatten

import NodeJS
import ZipFile
import EzXML
import ImageMagick
using FileIO

export EPUB

#####

abstract type EPUBFormat <: Documenter.Writers.FormatSelector end

Selectors.order(::Type{EPUBFormat}) = 4.0 # or a higher number.
Selectors.matcher(::Type{EPUBFormat}, fmt, _) = isa(fmt, EPUB)
Selectors.runner(::Type{EPUBFormat}, fmt, doc) = render(doc, fmt)

"""
    EPUB(;[color=false, lang="en"])

EPUB format Writer.
 - `color` specifies whether or not the code highlighting will be grayscale
or in color.
 - `lang` sets the language in the EPUB.
 - `snap_animations` if `true` replaces the animations with the first frame of it.
    (currently only works on .gif files)
"""
Base.@kwdef struct EPUB <: Documenter.Writer
    color::Bool = false
    lang::String = "en"
    snap_animations::Bool = true
end

"""
[`DocumenterEpub`](@ref)-specific globals that are passed to [`domify`](@ref) and
other recursive functions.
"""
mutable struct EPUBContext
    doc::Documents.Document
    settings::EPUB
    ids::Set{String}
    footnotes::Vector{Markdown.Footnote}
end

EPUBContext(doc, settings=EPUB(doc)) = EPUBContext(doc, settings, Set([]), [])

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

function toepub(doc)
    epub_dir = joinpath(doc.user.build, "epub")
    path = joinpath(doc.user.build, "$(doc.user.sitename).epub")
    w = ZipFile.Writer(path)
    for (root, dirs, files) in walkdir(epub_dir)
        for f in files
            p = joinpath(root, f)
            zf = ZipFile.addfile(w, relpath(p, epub_dir))
            write(zf, read(p))
            @debug "Zipping epub file: $p"
        end
    end
    close(w)
end

const OPFNS = [
    "ns" => "http://www.idpf.org/2007/opf",
    "dc" => "http://purl.org/dc/elements/1.1/"
    ]

# create functions operating on the content.opf
for f in (findfirst, findlast, findall)
    @eval $(Symbol("opf_" * string(f)))(xp, doc) = $f(xp, doc.root, OPFNS)
end

function update_content_metadata(doc;
    title,
    id=title,
    lang="en",
    modified_date=now(UTC),
    source=nothing)

    # identifier
    identifier_el = opf_findfirst("/ns:package/ns:metadata/dc:identifier", doc)
    identifier_el.content = id

    # modified date
    meta_el = opf_findfirst("/ns:package/ns:metadata/ns:meta", doc)
    meta_el.content = string(Dates.format(modified_date, "YYYY-mm-ddTHH:MM:SS"), "Z")

    # title
    title_el = opf_findfirst("/ns:package/ns:metadata/dc:title", doc)
    title_el.content = title

    # language
    language_el = opf_findfirst("/ns:package/ns:metadata/dc:language", doc)
    language_el.content = lang

    # source [optional]
    if source !== nothing
        metadata_el = opf_findfirst("/ns:package/ns:metadata", doc)
        addelement!(metadata_el, "dc:source", source)
    end
    return nothing
end

_build_xhtml_navcontent(ctx,nnv::Vector, iob) = foreach(nn -> _build_xhtml_navcontent(ctx, nn, iob), nnv)
function _build_xhtml_navcontent(ctx, nn, iob)
    @tags a span
    print(iob, "
")
    if !isempty(nn.children)
        title = mdflatten(pagetitle(ctx, nn))
        if nn.page !== nothing # hidden can have a target and children..
            page = get_url(nn.page)
            print(iob, a[:href => page](title))
            print(iob, "

")
        end
        print(iob, span(title))
        print(iob, "
")
        _build_xhtml_navcontent(ctx, nn.children, iob)
        print(iob, "
")
    else
        title = mdflatten(pagetitle(ctx, nn))
        page = get_url(nn.page)
        print(iob, a[:href => page](title))
    end
    print(iob, "
")
end
function _create_xhtml_nav(ctx, doc)
    iob = IOBuffer()
    print(iob, "
")
    _build_xhtml_navcontent(ctx, doc.internal.navtree, iob)
    print(iob, "
")
    String(take!(iob))
end
# depth first search for page
function firstpage(nn)
    nn.page !== nothing && return nn.page
    for nnc in nn.children
        fp = firstpage(nnc)
        fp !== nothing && return fp
    end
    return nothing
end
_build_ncx_nav_content(ctx,nnv::Vector, iob) =  foreach(nn -> _build_ncx_nav_content(ctx, nn, iob), nnv)
function _build_ncx_nav_content(ctx, nn, iob)
    @tags navLabel text content
    if !isempty(nn.children)
        page = get_url(firstpage(nn))
        title = mdflatten(pagetitle(ctx, nn))
        print(iob, """<navPoint id="navpoint-$(rand(Int))">""")
        print(iob, navLabel(text(title)))
        print(iob, content[:src => page]())
        _build_ncx_nav_content(ctx, nn.children, iob)
        print(iob, """</navPoint>""")
    else
        title = mdflatten(pagetitle(ctx, nn))
        page = get_url(nn.page)
        print(iob, """<navPoint id="navpoint-$(rand(Int))">""")
        print(iob, navLabel(text(title)))
        print(iob, content[:src => page]())
        print(iob, """</navPoint>""")
    end
end
function _create_ncx_nav(ctx, doc)
    iob = IOBuffer()
    print(iob, "<navMap>")
    _build_ncx_nav_content(ctx, doc.internal.navtree, iob)
    print(iob, "</navMap>")
    String(take!(iob))
end

function _contains_svg(root, pagefile)
    try
        page_as_xml = EzXML.readxml(joinpath(root, pagefile))
        return findfirst("//svg:svg", page_as_xml.root, ["svg" => "http://www.w3.org/2000/svg"]) !== nothing
    catch e
        @debug "Error while checking for svg tag" e
    end
    return false
end

# Copied from https://github.com/JuliaLang/julia/blob/acb7bd93fb2d5adbbabeaed9f39ab3c85495b02f/stdlib/Markdown/src/render/html.jl#L25-L31
const _htmlescape_chars = IdDict()
for ch in """"'`!\$%()=+{}[]<>&"""
    _htmlescape_chars[ch] = "#$(Int(ch));"
end

"""
    html_unescape

Internal function to reverse the escaping of some html code (in order to avoid
double escaping when pre-rendering with highlight, see issue 326).
"""
function html_unescape(cs::AbstractString)
    # this is a bit inefficient but whatever, `cs` shouldn't  be very long.
    for (ssfrom, ssto) in _htmlescape_chars
        cs = replace(cs, ssto => ssfrom)
    end
    return cs
end

# https://www.npmjs.com/package/mathjax
function prerender_mathjax(formula::String, display_mode::Bool)
    mathjax = joinpath(@__DIR__, "..", "res", "mathjax","node-main.js")
    rf = escape_string(replace(formula, '\'' => "\\prime"))
    svg_code = nothing
    cd(dirname(mathjax)) do
        svg_code= String(read(`$(NodeJS.nodejs_cmd()) -e """
            require('$(escape_string(mathjax))').init({
                loader: {load: ['input/tex', 'output/svg']}, svg: {
                    fontCache: 'none',
                    'localID':$(rand(1:100000))}
            }).then((MathJax) => {
                const svg = MathJax.tex2svg('$rf', {display: $(display_mode)});
                console.log(MathJax.startup.adaptor.innerHTML(svg));
            }).catch((err) => console.log(err.message));
            """`))
    end
    return svg_code
end

const HLJS_LANGS = readlines(joinpath(@__DIR__, "..", "res", "supported_langs.txt"))

"""
Takes a html string that may contain `<pre><code ... </code></pre>` blocks and use node and
highlight.js to pre-render them to HTML.
"""
function prerender_highlightjs(hs::String, lang::String)::String
    # select highlight js script (julia one stems from https://fredrikekre.se/posts/highlight-julia/)
    hljsfile = occursin("julia", lang) ? "julia.highlight.min.js" : "highlight.pack.js"
    hljsfile = abspath(joinpath(@__DIR__, "..", "res", hljsfile))

    # buffer to write the JS script
    inbuffer = IOBuffer()
    write(inbuffer, """const hljs = require('$(escape_string(hljsfile))');""")

    # un-escape code string
    cs = escape_string(html_unescape(hs))
    # add to content of jsbuffer
    write(inbuffer, """console.log("
" + hljs.highlight("$lang", "$cs").value + "
");""")

    outbuffer = IOBuffer()
    run(pipeline(`$(NodeJS.nodejs_cmd()) -e "$(String(take!(inbuffer)))"`, stdout=outbuffer))
    return String(take!(outbuffer))
end

"""
    safeid(page)
Sanitize `page` for use in the navigation files toc.xhtml and toc.ncx.
"""
safeid(page) = replace(replace(page, "-" => "_"), "/" => "_")

# create a manifest entry for a given 'pagefile'
function _create_manifest_item(root, pagefile)
    pagefile = normpath(pagefile)
    mediatypes = Dict(".jpg|.jpeg" => "image/jpeg",
                    ".gif" => "image/gif",
                    ".png" => "image/png",
                    ".svg" => "image/svg+xml",
                    ".css" => "text/css",
                    ".js" => "application/javascript",
                    ".woff2" => "font/woff2",
                    ".woff" => "font/woff",
                    ".ttf" => "application/font-sfnt",
                    ".ncx" => "application/x-dtbncx+xml")

    ext = lowercase(splitext(pagefile)[2])
    mediatype = get(mediatypes, ext, "application/xhtml+xml")

    manifest_item = EzXML.ElementNode("item")
    manifest_item["id"] = safeid(pagefile)
    manifest_item["href"] = pagefile
    manifest_item["media-type"] = mediatype
    if mediatype == "application/xhtml+xml" && endswith(pagefile, ".xhtml")
        if _contains_svg(root, pagefile)
            manifest_item["properties"] = "svg"
        end
    end
    if pagefile == "toc.xhtml"
        manifest_item["properties"] = "nav"
    end
    if pagefile == "content.opf"
        return nothing
    end
    return manifest_item
end

function _create_spine_item(pagefile)
    spine_item = EzXML.ElementNode("itemref")
    spine_item["idref"] = safeid(pagefile)
    return spine_item
end

const ALLOWED_EXT = (".xhtml", ".png", ".jpg", "jpeg", ".gif", ".md", ".svg", "")

function snap_animations(start)
    for (root, _, files) in walkdir(start)
        for f in files
            if lowercase(splitext(f)[2]) == ".gif"
                giffile = joinpath(root, f)
                dims,_ = ImageMagick.metadata(giffile)
                if length(dims) > 2
                    img = load(giffile)
                    save(giffile,img[:,:,1])
                end
            end
        end
    end
end

function render(doc::Documents.Document, settings::EPUB=EPUB())
    !isempty(doc.user.sitename) || error("EPUB output requires `sitename`.")

    if isempty(doc.blueprint.pages)
        error("Aborting HTML build: no pages under src/")
    elseif !haskey(doc.blueprint.pages, "index.md")
        @warn "Can't generate landing page (index.html): src/index.md missing" keys(doc.blueprint.pages)
    end
    @info "DocumenterEpub: rendering HTML pages for EPUB."

    # first do the html creation into doc.user.build
    ctx = EPUBContext(doc, settings)
    for page in keys(doc.blueprint.pages)
        idx = findfirst(nn -> nn.page == page, doc.internal.navlist)
        nn = (idx === nothing) ? Documents.NavNode(page, nothing, nothing) : doc.internal.navlist[idx]
        @debug "Rendering $(page) [$(repr(idx))]"
        render_page(ctx, nn)
    end

    # move the created files into from "<doc.user.build>/*" into "<doc.user.build>/epub/Content/*"
    tmp_html_build = mktempdir()
    mv(doc.user.build, tmp_html_build; force=true)
    epub_content_root = abspath(joinpath(doc.user.build, "epub", "Content"))
    mkpath(epub_content_root)
    mv(tmp_html_build, epub_content_root;force=true)

    # remove all non allowed files
    for (root, _, files) in walkdir(epub_content_root)
        for f in files
            if lowercase(splitext(f)[2]) ∉ ALLOWED_EXT
                rm(joinpath(root, f))
            end
        end
    end

    if settings.snap_animations
        snap_animations(epub_content_root)
    end

    # copy the epub template into the `epub_content_root`
    template_source = joinpath(@__DIR__, "..", "res", "template")
    template_target = dirname(epub_content_root)
    for (root, _, files) in walkdir(template_source)
        destpath = joinpath(template_target, relpath(root, template_source))
        isdir(destpath) || mkpath(destpath)

        for f in files
            src = joinpath(root, f)
            dest = joinpath(destpath, f)
            isfile(dest) && error("Destination $dest for as target of template already exists")
            cp(src, dest)
            chmod(dest, 0o664)
        end
    end

    # remove the unused css
    rm(joinpath(epub_content_root, settings.color ? "grayscale.css" : "github.css"))

    # create the coverpage
    open(joinpath(epub_content_root, "titlepage.xhtml"), "w") do io
        print(io,  render_cover_page(ctx))
    end

    # fill/adjust the epub3 relevant files
    # first the content.opf
    @info "EPUB: Writing meta data to content.opf"
    content_xml_path = joinpath(epub_content_root, "content.opf")
    content_xml_doc = EzXML.readxml(content_xml_path)
    content_xml_ns = EzXML.namespace(content_xml_doc.root)
    content_xml_manifest_node = findfirst("//x:manifest", content_xml_doc.root, ["x" => content_xml_ns])
    content_xml_spine_node = findfirst("//x:spine", content_xml_doc.root, ["x" => content_xml_ns])
    update_content_metadata(content_xml_doc;
        id=doc.user.sitename,
        title=doc.user.sitename,
        lang=settings.lang,
        source=isempty(doc.user.repo) ? nothing : doc.user.repo )

    # walk along the navlist and create manifest and spine item
    # the navlist can contain the pages multiple times, so only do it once for each page
    visited = Set{String}([])
    root = epub_content_root

    # add titlepage
    spine_item = _create_spine_item("titlepage.xhtml")
    manifest_item = _create_manifest_item(root, "titlepage.xhtml")
    if manifest_item !== nothing
        EzXML.link!(content_xml_manifest_node, manifest_item)
        EzXML.link!(content_xml_spine_node, spine_item)
    end

    for nn in doc.internal.navlist
        pagefile = replace(nn.page, r".md$" => ".xhtml")
        pagefile in visited && continue
        spine_item = _create_spine_item(pagefile)
        manifest_item = _create_manifest_item(root, pagefile)
        if manifest_item !== nothing
            EzXML.link!(content_xml_manifest_node, manifest_item)
            EzXML.link!(content_xml_spine_node, spine_item)
        end
        push!(visited, pagefile)
    end

    # there might be files that are not part of the navlist, walk through all files
    # and add any remaining to the manifest
    for (fileroot, _, files) in walkdir(epub_content_root)
        for file in files
            pagefile = normpath(replace(joinpath(relpath(fileroot, root), file), r".md$" => ".xhtml"))
            pagefile in visited && continue
            pagefile == "titlepage.xhtml" && continue # already added
            manifest_item = _create_manifest_item(root, pagefile)
            spine_item = nothing
            if endswith(pagefile, ".xhtml") || endswith(pagefile, ".svg")
                spine_item = _create_spine_item(pagefile)
            end
            if manifest_item !== nothing
                EzXML.link!(content_xml_manifest_node, manifest_item)
                if spine_item !== nothing
                    EzXML.link!(content_xml_spine_node, spine_item)
                end
            end
        end
    end

    # write the adjusted xml content into content.opf
    write(content_xml_path, content_xml_doc)

    # create the navigation file `toc.ncx` for epub2 compatibility
    @info "EPUB: Writing toc.ncx"
    toc_ncx_path = joinpath(epub_content_root, "toc.ncx")
    toc_ncx_doc = EzXML.readxml(toc_ncx_path)
    toc_ncx_ns = EzXML.namespace(toc_ncx_doc.root)
    toc_ncx_doctitle_node = findfirst("//x:docTitle", toc_ncx_doc.root, ["x" => toc_ncx_ns])
    EzXML.addelement!(toc_ncx_doctitle_node, "text", doc.user.sitename)
    toc_ncx_head_node = findfirst("//x:meta[@name='dtb:uid']", toc_ncx_doc.root, ["x" => toc_ncx_ns])
    toc_ncx_head_node["content"] = doc.user.sitename
    toc_ncx_navmap_content = EzXML.unlink!(EzXML.parsexml(_create_ncx_nav(ctx, doc)).root)
    EzXML.link!(toc_ncx_doc.root, toc_ncx_navmap_content)
    write(toc_ncx_path, toc_ncx_doc)

    # create the navigation file `toc.xhtml` for epub3
    @info "EPUB: Writing toc.xhtml"
    toc_xhtml_path = joinpath(epub_content_root, "toc.xhtml")
    toc_xhtml_doc = EzXML.readxml(toc_xhtml_path)
    toc_xhtml_ns = EzXML.namespace(toc_xhtml_doc.root)
    toc_xhtml_nav_node = findfirst("//x:nav", toc_xhtml_doc.root, ["x" => toc_xhtml_ns])
    toc_xhtml_nav_content = EzXML.unlink!(EzXML.parsexml(_create_xhtml_nav(ctx, doc)).root)
    EzXML.link!(toc_xhtml_nav_node, toc_xhtml_nav_content)
    write(toc_xhtml_path, toc_xhtml_doc)

    # zip the directory into a epub file
    toepub(doc)

    @info """
    TIP: Use https://github.com/pgaskin/kepubify to optimize the ebook for KOBO devices!
    """
end

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

## Cover page

function render_cover_page(ctx)
    @tags html div body head meta link script title h1 p

    return DOM.HTMLDocument(
        html[Symbol("lang") => ctx.settings.lang,
            Symbol("xml:lang") => ctx.settings.lang,
            :xmlns => "http://www.w3.org/1999/xhtml",
            Symbol("xmlns:epub") => "http://www.idpf.org/2007/ops",
            Symbol("xmlns:svg") => "http://www.w3.org/2000/svg"](
                head(
                    meta[:charset => "UTF-8"],
                    title(ctx.doc.user.sitename),
                    link[:href => "style.css", :rel => "stylesheet", :type => "text/css"](),
                ),
                body[:class => "titlepage",Symbol("epub:type") => "frontmatter titlepage"](
                    h1[:class => "titlepage"](uppercase(ctx.doc.user.sitename)),
                    p(ctx.doc.user.version),
                    p(ctx.doc.user.authors)
                ),
            )
        )
end

## Standard page
"""
Constructs and writes the page referred to by the `navnode` to `.build`.
"""
function render_page(ctx, navnode)
    @tags html div body
    page = getpage(ctx, navnode)
    head = render_head(ctx, navnode)
    article = render_article(ctx, navnode)
    htmldoc = render_html(ctx, head,  article)
    open_output(ctx, navnode) do io
        print(io,  htmldoc)
    end
end

## Rendering HTML elements
# ------------------------------------------------------------------------------

"""
Renders the main `<html>` tag.
"""
function render_html(ctx, head,  article)
    @tags html body div
    DOM.HTMLDocument(
    html[Symbol("lang") => ctx.settings.lang,
        Symbol("xml:lang") => ctx.settings.lang,
        :xmlns => "http://www.w3.org/1999/xhtml",
        Symbol("xmlns:epub") => "http://www.idpf.org/2007/ops",
        Symbol("xmlns:svg") => "http://www.w3.org/2000/svg"](
            head,
            body(article),
        )
    )
end

function render_head(ctx, navnode)
    @tags head meta link script title
    src = get_url(navnode)
    stylesheet = ctx.settings.color ? "github.css" : "grayscale.css"

    page_title = "$(mdflatten(pagetitle(ctx, navnode))) · $(ctx.doc.user.sitename)"

    prefix = !isempty(dirname(src)) ? relpath(".", dirname(src)) : "."
    head(
        meta[:charset => "UTF-8"],
        title(page_title),
        link[:href => joinpath(prefix, stylesheet), :rel => "stylesheet", :type => "text/css"](),
        link[:href => joinpath(prefix, "style.css"), :rel => "stylesheet", :type => "text/css"](),
        )

end

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

function find_image_asset(ctx, name)
    for ext in ["svg", "png", "webp", "gif", "jpg", "jpeg"]
        filename = joinpath("assets", "$(name).$(ext)")
        isfile(joinpath(ctx.doc.user.build, filename)) && return filename
    end
    return nothing
end

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

function render_article(ctx, navnode)
    @tags article section hr span a div p aside dl dt dd

    # Build the page itself (and collect any footnotes)
    empty!(ctx.footnotes)
    empty!(ctx.ids)
    art_body = article(domify(ctx, navnode))
    # Footnotes, if there are any
    if !isempty(ctx.footnotes)
        fnotes = map(ctx.footnotes) do f
            fid = "footnote-$(f.id)"
            if length(f.text) == 1 && first(f.text) isa Markdown.Paragraph
                dt(f.id),
                dd[:id => "$(fid)", Symbol("epub:type") => "footnote"](mdconvert(f.text[1].content))
            else
                dt(f.id),
                # passing an empty MD() as `parent` to give it block context
                dd[:id => "$(fid)", Symbol("epub:type") => "footnote"](mdconvert(f.text, Markdown.MD()))
            end
        end
        push!(art_body.nodes, hr())
        push!(art_body.nodes, section[Symbol("epub:type") => "footnotes"](dl(fnotes)))
    end
    return art_body
end

"""
    saneid(anchor)
Sanitize the `anchor`s id.
"""
saneid(a::Anchors.Anchor) = saneid(a.id, a.nth)
function saneid(id, nth=1)

    isempty(id) && return string("id", rand(1_10000))
    # use strict html 4 rules for ids: start with a letter followed by a subset of allowed chars
    if occursin(r"^[^a-zA-Z]", id) || occursin(r"[^a-zA-Z0-9-_:.]", id)
        # just hash it
        return string("id", hash(id) + nth - 1)
    else
        return nth == 1 ? id : string(id, "-", nth)
    end
end

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

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

function domify(ctx, navnode, node)
    fixlinks!(ctx, navnode, node)
    mdconvert(node, Markdown.MD(); footnotes=ctx.footnotes)
end

function domify(ctx, navnode, anchor::Anchors.Anchor)
    @tags a
    id = saneid(anchor)
    if id ∈ ctx.ids
        id = string("id-", rand(1:100000))
    end
    push!(ctx.ids, id)
    frag = '#' * id
    if isa(anchor.object, Markdown.Header)
        h = anchor.object
        fixlinks!(ctx, navnode, h)
        DOM.Tag(Symbol("h$(Utilities.header_level(h))"))[:id => id](mdconvert(h.text, h))
    else
        a[:href => frag](domify(ctx, navnode, anchor.object))
    end
end

struct 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 -> e isa ListBuilder ? li(domify(e)) : li(e), lb.es))
end

function domify(ctx, navnode, contents::Documents.ContentsNode)
    @tags a
    navnode_dir = dirname(navnode.page)
    navnode_url = get_url(navnode)
    lb = ListBuilder()
    for (count, path, anchor) in contents.elements
        path = joinpath(navnode_dir, path) # links in ContentsNodes are relative to current page
        path = relhref(navnode_url, get_url(path))
        header = anchor.object
        url = string(path, '#', saneid(anchor))
        node = a[:href => url](mdconvert(header.text; droplinks=true))
        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
    navnode_dir = dirname(navnode.page)
    navnode_url = get_url(navnode)
    lis = map(index.elements) do el
        object, doc, path, mod, cat = el
        path = joinpath(navnode_dir, path) # links in IndexNodes are relative to current page
        path = relhref(navnode_url, get_url(path))
        url = string(path, "#", 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 code article header span

    id = saneid(node.anchor)
    if id ∈ ctx.ids
        id = string("id-", rand(1:100000))
    end
    push!(ctx.ids, id)
    article[".docstring"](
        header(
            span[".docstring-binding", :id => 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 section div
    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 md
            markdown, result = md
            ret = section(div(domify(ctx, navnode, Writers.MarkdownWriter.dropheaders(markdown))))
            # When a source link is available then print the link.
            return 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.
        section(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[]

function domify(ctx, navnode, raw::Documents.RawNode)
    raw.name === :html ? Tag(Symbol("#RAW#"))(raw.text) : DOM.Node[]
end

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

"""
Opens the output file of the `navnode` in write node. If necessary, the path to the output
file is created before opening the file.
"""
function open_output(f, ctx, navnode, mode="w")
    path = joinpath(ctx.doc.user.build, get_url(navnode))
    isdir(dirname(path)) || mkpath(dirname(path))
    open(f, path, mode)
end

"""
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 `nothing`).
"""
navhref(ctx, to, from) = relhref(get_url(from), get_url(to))

"""
Calculates a relative HTML link from one path to another.
"""
function relhref(from, to)
    pagedir = dirname(from)
    # The regex separator replacement is necessary since otherwise building the docs on
    # Windows will result in paths that have `//` separators which break asset inclusion.
    replace(relpath(to, isempty(pagedir) ? "." : pagedir), r"[/\\]+" => "/")
end

"""
Returns the full path corresponding to a path of a `.md` page file. The the input and output
paths are assumed to be relative to `src/`.
"""
get_url(path::AbstractString) = string(splitext(path)[1], ".xhtml")

"""
Returns the full path of a [`Documents.NavNode`](@ref) relative to `src/`.
"""
get_url(navnode::Documents.NavNode) = get_url(navnode.page)

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

function pagetitle(ctx, navnode::Documents.NavNode)
    if navnode.title_override !== nothing
        # parse title_override as markdown
        md = Markdown.parse(navnode.title_override)
        # Markdown.parse results in a paragraph so we need to strip that
        if !(length(md.content) === 1 && isa(first(md.content), Markdown.Paragraph))
            error("Bad Markdown provided for page title: '$(navnode.title_override)'")
        end
        return first(md.content).content
    end

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

    "-"
end

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

const md_block_nodes = [
    Markdown.MD,
    Markdown.BlockQuote,
    Markdown.List,
    Markdown.Admonition,
]

"""
[`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.
"""
const 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; kwargs...) = mdconvert(md, md; kwargs...)

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

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

mdconvert(md::Markdown.MD, parent; kwargs...) = mdconvert(md.content, md; kwargs...)

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

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

function codelang(infostring::AbstractString)
    m = match(r"^\s*(\S*)", infostring)
    return m[1]
end
function mdconvert(c::Markdown.Code, parent::MDBlockContext; kwargs...)
    @tags pre code span
    language = codelang(c.language)
    if isempty(language) || startswith(language, "@") || language ∉ HLJS_LANGS
        pre(code[".hljs"](c.code))
    else
        rendered_code = prerender_highlightjs(c.code, string(language))
        Tag(Symbol("#RAW#"))(rendered_code)
    end
end

mdconvert(c::Markdown.Code, parent; kwargs...) = Tag(:code)(c.code)

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

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

function mdconvert(i::Markdown.Image, parent; kwargs...)
    @tags video img a

    # remove external links
    startswith(i.url, "http") && return i.url

    alt = i.alt === nothing ? i.url : isempty(i.alt) ? "none" : i.alt
    if occursin(r"\.(webm|mp4|ogg|ogm|ogv|avi)$", i.url)
        video[:src => i.url, :controls => "true", :title => alt](
            a[:href => i.url](alt)
        )
    else
        img[:src => i.url, :alt => alt]
    end
end

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

function mdconvert(m::Markdown.LaTeX, ::MDBlockContext; kwargs...)
    @tags div
    return div[:class => "mjx-container"](Tag(Symbol("#RAW#"))(prerender_mathjax(m.formula, true)))
end

function mdconvert(m::Markdown.LaTeX, parent; kwargs...)
    @tags div
    return Tag(Symbol("#RAW#"))(prerender_mathjax(m.formula, false))
end

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

function mdconvert(link::Markdown.Link, parent; droplinks=false, kwargs...)
    link_text = mdconvert(link.text, link; droplinks=droplinks, kwargs...)
    # remove external links
    droplinks = startswith(link.url, "http") || droplinks
    droplinks ? link_text : Tag(:a)[:href => link.url](link_text)
end

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

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

# For compatibility with versions before Markdown.List got the `loose field, Julia PR #26598
const list_has_loose_field = :loose in fieldnames(Markdown.List)
function mdconvert(paragraph::Markdown.Paragraph, parent::Markdown.List; kwargs...)
    content = mdconvert(paragraph.content, paragraph; kwargs...)
    return (list_has_loose_field && !parent.loose) ? content : Tag(:p)(content)
end

function mdconvert(t::Markdown.Table, parent; kwargs...)
    @tags table tr th td
    alignment_style = map(t.align) do align
        if align == :r
            "text-align: right"
        elseif align == :c
            "text-align: center"
        else
            "text-align: left"
        end
    end
    table(
        tr(map(enumerate(t.rows[1])) do (i, x)
        th[:style => alignment_style[i]](mdconvert(x, t; kwargs...))
    end),
        map(t.rows[2:end]) do x
        tr(map(enumerate(x)) do (i, y) # each cell in a row
            td[:style => alignment_style[i]](mdconvert(y, x; kwargs...))
        end)
    end
    )
end

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

function mdconvert(f::Markdown.Footnote, parent; footnotes=nothing, kwargs...)
    @tags sup a
    if f.text === nothing # => Footnote link
        return sup(a[Symbol("epub:type") => "noteref",:href => "#footnote-$(f.id)"]("[$(f.id)]"))
    elseif footnotes !== nothing # Footnote definition
        push!(footnotes, f)
    else # => Footnote definition, but nowhere to put it
        @error "Bad footnote definition."
    end
    return []
end

function mdconvert(a::Markdown.Admonition, parent; kwargs...)
    @tags header div span
    prefix = (a.category == "danger")  ? "\u26a1"  :
        (a.category == "warning") ? "\u26a0" :
        (a.category == "note")    ? "\u261b"    :
        (a.category == "info")    ? "\u2328"    :
        (a.category == "tip")     ? "\u270e" :
        (a.category == "compat")  ? "\u221e"  : "\u58e"
    colorclass =
        (a.category == "danger")  ? ".is-danger"  :
        (a.category == "warning") ? ".is-warning" :
        (a.category == "note")    ? ".is-info"    :
        (a.category == "info")    ? ".is-info"    :
        (a.category == "tip")     ? ".is-success" :
        (a.category == "compat")  ? ".is-compat"  : ".is-default"
    div[".admonition$(colorclass)",Symbol("epub:type") => "notice"](
        header[".admonition-header"](span[".icon"](prefix), " " * a.title),
        div[".admonition-body"](mdconvert(a.content, a; kwargs...))
    )
end

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

# Select the "best" representation for HTML output.
mdconvert(mo::Documents.MultiOutput, parent; kwargs...) =
    Base.invokelatest(mdconvert, mo.content, parent; kwargs...)
function mdconvert(d::Dict{MIME,Any}, parent; kwargs...)
    if haskey(d, MIME"text/html"())
        out = Documents.RawHTML(d[MIME"text/html"()])
    elseif haskey(d, MIME"image/svg+xml"())
        out = Documents.RawHTML(d[MIME"image/svg+xml"()])
    elseif haskey(d, MIME"image/png"())
        out = Documents.RawHTML(string(""))
    elseif haskey(d, MIME"image/webp"())
        out = Documents.RawHTML(string(""))
    elseif haskey(d, MIME"image/gif"())
        out = Documents.RawHTML(string(""))
    elseif haskey(d, MIME"image/jpeg"())
        out = Documents.RawHTML(string(""))
    elseif haskey(d, MIME"text/latex"())
        # If the show(io, ::MIME"text/latex", x) output is already wrapped in \[ ... \] or $$ ... $$, we
        # unwrap it first, since when we output Markdown.LaTeX objects we put the correct
        # delimiters around it anyway.
        latex = d[MIME"text/latex"()]
        # Make sure to match multiline strings!
        m_bracket = match(r"\s*\\\[(.*)\\\]\s*"s, latex)
        m_dollars = match(r"\s*\$\$(.*)\$\$\s*"s, latex)
        if m_bracket === nothing && m_dollars === nothing
            out = Utilities.mdparse(latex; mode=:single)
        else
            out = Markdown.LaTeX(m_bracket !== nothing ? m_bracket[1] : m_dollars[1])
        end
    elseif haskey(d, MIME"text/markdown"())
        out = Markdown.parse(d[MIME"text/markdown"()])
    elseif haskey(d, MIME"text/plain"())
        @tags pre code
        return pre(code[".hljs"](d[MIME"text/plain"()]))
    else
        error("this should never happen.")
    end
    return mdconvert(out, parent; kwargs...)
end

# Fallback
function mdconvert(x, parent; kwargs...)
    @debug "Strange inline Markdown node (typeof(x) = $(typeof(x))), falling back to repr()" x
    repr(x)
end

# fixlinks!
# ------------------------------------------------------------------------------

"""
Replaces URLs in `Markdown.Link` elements (if they point to a local `.md` page) with the
actual URLs.
"""
function fixlinks!(ctx, navnode, link::Markdown.Link)
    fixlinks!(ctx, navnode, link.text)
    Utilities.isabsurl(link.url) && return

    # anything starting with mailto: doesn't need fixing
    startswith(link.url, "mailto:") && return

    # links starting with a # are references within the same file -- there's nothing to fix
    # for such links
    startswith(link.url, '#') && return

    s = split(link.url, "#", limit=2)
    if Sys.iswindows() && ':' in first(s)
        @warn "invalid local link: colons not allowed in paths on Windows in $(Utilities.locrepr(navnode.page))" link = link.url
        return
    end
    path = normpath(joinpath(dirname(navnode.page), first(s)))

    if (endswith(path, ".md") && path in keys(ctx.doc.blueprint.pages)) ||
        (endswith(path, ".xhtml") && string(splitext(path)[1], ".md") in keys(ctx.doc.blueprint.pages))
        # make sure that links to different valid pages are correct
        path = relhref(get_url(navnode), get_url(path))
    elseif isfile(joinpath(ctx.doc.user.build, path))
        # update links to other files that are present in build/ (e.g. either user
        # provided files or generated by code examples)
        path = relhref(get_url(navnode), path)
    else
        @warn "invalid local link: unresolved path in $(Utilities.locrepr(navnode.page))" link.text link.url path
    end

    # Replace any backslashes in links, if building the docs on Windows
    path = replace(path, '\\' => '/')

    link.url = (length(s) > 1) ? "$path#$(saneid(last(s)))" : String(path)
end

function fixlinks!(ctx, navnode, img::Markdown.Image)
    Utilities.isabsurl(img.url) && return

    if Sys.iswindows() && ':' in img.url
        @warn "invalid local image: colons not allowed in paths on Windows in $(Utilities.locrepr(navnode.page))" link = img.url
        return
    end

    path = joinpath(dirname(navnode.page), img.url)
    if isfile(joinpath(ctx.doc.user.build, path))
        path = relhref(get_url(navnode), path)
        # Replace any backslashes in links, if building the docs on Windows
        img.url = replace(path, '\\' => '/')
    else
        @warn "invalid local image: unresolved path in $(Utilities.locrepr(navnode.page))" link = img.url
    end
end

fixlinks!(ctx, navnode, md::Markdown.MD) = fixlinks!(ctx, navnode, md.content)
function fixlinks!(ctx, navnode, a::Markdown.Admonition)
    fixlinks!(ctx, navnode, a.title)
    fixlinks!(ctx, navnode, a.content)
end
fixlinks!(ctx, navnode, b::Markdown.BlockQuote) = fixlinks!(ctx, navnode, b.content)
fixlinks!(ctx, navnode, b::Markdown.Bold) = fixlinks!(ctx, navnode, b.text)
fixlinks!(ctx, navnode, f::Markdown.Footnote) = fixlinks!(ctx, navnode, f.text)
fixlinks!(ctx, navnode, h::Markdown.Header) = fixlinks!(ctx, navnode, h.text)
fixlinks!(ctx, navnode, i::Markdown.Italic) = fixlinks!(ctx, navnode, i.text)
fixlinks!(ctx, navnode, list::Markdown.List) = fixlinks!(ctx, navnode, list.items)
fixlinks!(ctx, navnode, p::Markdown.Paragraph) = fixlinks!(ctx, navnode, p.content)
fixlinks!(ctx, navnode, t::Markdown.Table) = fixlinks!(ctx, navnode, t.rows)

fixlinks!(ctx, navnode, mds::Vector) = map(md -> fixlinks!(ctx, navnode, md), mds)
fixlinks!(ctx, navnode, md) = nothing

# TODO: do some regex-magic in raw HTML blocks? Currently ignored.
# fixlinks!(ctx, navnode, md::Documents.RawHTML) = ...

end


text:
using CornerSpaceRenorm, QuantumOptics
using LightGraphs, GraphPlot
using LinearAlgebra, IterativeSolvers, Random
using Test, InteractiveUtils

macro no_error(ex)
    quote
        try
            $(esc(ex))
            true
        catch
            false
        end
    end
end

@testset "CornerSpaceRenorm" begin
    @testset "Operators" begin
        lb = FockBasis(1)
        lH = randoperator(lb)
        tH = randoperator(lb)

        for pbc in [true, false]
            L = NdLattice((3,2,1);periodic=pbc)
            H1 = hamiltonian(L,lH,(1.,1.,1.),tH)
            H2 = hamiltonian(L,lH,(1.,tH))
            H3 = hamiltonian(L,lH,1.,tH)
            @test maximum(abs2.(H1.data .- H2.data)) ≈ 0. atol=eps(Float64)
            @test maximum(abs2.(H1.data .- H3.data)) ≈ 0. atol=eps(Float64)
        end
    end;

    @testset "Types" begin
        @testset "NdLattice" begin
            L = NdLattice((2,2))

            # Test method extensions
            # TO DO: extend tests to instances of all subtypes of Lattice
            @test eltype(L) == eltype(L.L)
            @test vertices(L) == vertices(L.L)
            @test edges(L) == edges(L.L)
            @test nv(L) == nv(L.L)
            @test ne(L) == ne(L.L)
            @test all([has_edge(L,v1,v2) == has_edge(L.L,v1,v2) for v1 in vertices(L), v2 in vertices(L)])
            @test all([has_edge(L,e) == has_edge(L.L,e) for e in edges(L)])
            @test all([has_vertex(L,v) == has_vertex(L.L,v) for v in collect(vertices(L)) ∪ collect(-10:-1)])
            @test all([inneighbors(L,v) == inneighbors(L.L,v) for v in vertices(L)])
            @test all([outneighbors(L,v) == outneighbors(L.L,v) for v in vertices(L)])
            @test is_directed(L) == false
            @test all(is_directed.(subtypes(Lattice)) .== false)

            function compare_lattices(L1,L2)
                @test typeof(L1) == typeof(L2)
                @test all([getfield(L1,f) == getfield(L2,f) for f in fieldnames(typeof(L1))])
            end

            # Test lattice merging
            L1 = NdLattice((2,3))
            L2 = NdLattice((2,1))
            L3_12 = union(L1,L2,2)
            L3_21 = union(L2,L1,2)
            compare_lattices(L3_12,L3_21)
            @test_throws DimensionMismatch union(L1,L2,1)

            L1 = NdLattice((3,2))
            L2 = NdLattice((1,2))
            L3_12 = union(L1,L2,1)
            L3_21 = union(L2,L1,1)
            @test typeof(L3_12) == typeof(L3_21)
            # TO DO: test boundaries
            @test LightGraphs.Experimental.could_have_isomorph(L3_12.L, L3_21.L)
            @test L3_12.shape[1] == L1.shape[1] + L2.shape[1]
            compare_lattices(L3_12, union(L1,L2,1))
            @test_throws DimensionMismatch union(L1,L2,2)

            Lpbc = NdLattice((2,2);periodic=true)
            Lobc = NdLattice((2,2);periodic=false)
            @test_throws AssertionError union(Lpbc,Lobc,1)
            @test_throws AssertionError union(Lpbc,Lobc,2)

            L = NdLattice((2,);periodic=true)
            LuL = union(L,L,1)
            L2 = NdLattice((4,);periodic=true)
            compare_lattices(LuL,L2)

            L = NdLattice((2,2);periodic=true)
            LuL = union(L,L,1)
            L2 = NdLattice((4,2);periodic=true)
            @test typeof(LuL) == typeof(L2)
            @test LightGraphs.Experimental.could_have_isomorph(LuL.L, L2.L)

            L = NdLattice((2,2);periodic=true)
            LuL = union(L,L,2)
            L2 = NdLattice((2,4);periodic=true)
            @test typeof(LuL) == typeof(L2)
            compare_lattices(LuL,L2)

            L = NdLattice((1,);periodic=true)
            for i in 1:3
                L = union(L,L,1)
            end
            compare_lattices(L, NdLattice((8,); periodic=true))

            # Test pbc_from_obc and obc_from_pbc
            Lo = NdLattice((5,5);periodic=false)
            Lp = NdLattice((5,5);periodic=true)
            compare_lattices(pbc_from_obc(Lo), Lp)
            compare_lattices(Lo, obc_from_pbc(Lp))

            Lo = NdLattice((1,1);periodic=false)
            Lo = union(Lo,Lo,1)
            Lo = union(Lo,Lo,1)
            Lo = union(Lo,Lo,1)
            Lo = union(Lo,Lo,2)
            Lo = union(Lo,Lo,2)
            Lo = union(Lo,Lo,2)
            Lp = NdLattice((8,8);periodic=true)
            compare_lattices(pbc_from_obc(Lo), Lp)
            compare_lattices(Lo, obc_from_pbc(Lp))

            # Dummy tests on plotting
            L = NdLattice((2,2);periodic=true)
            @test @no_error gplot(L)
            @test @no_error gplot(L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))))
            x = gplot(L)
            y = gplot(L.L)
            @test typeof(x) == typeof(y)
            x = gplot(L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))))
            y = gplot(L.L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))))
            @test all([typeof(getfield(x,f)) == typeof(getfield(y,f)) for f in fieldnames(typeof(x))])

            # Test pretty-printing
            L = NdLattice((2,2);periodic=true)
            @test @no_error repr(L)
            @test @no_error repr(MIME("text/plain"),L)
            @test @no_error repr(MIME("application/prs.juno.inline"),L)
            L = NdLattice((2,2);periodic=false)
            @test @no_error repr(L)
            @test @no_error repr(MIME("text/plain"),L)
            @test @no_error repr(MIME("application/prs.juno.inline"),L)
        end;

        @testset "NdSystem" begin
            # Local basis
            b_spin = SpinBasis(1//2)

            # Build some operators in the local basis
            sm = sigmam(b_spin)
            sp = sigmap(b_spin)
            sz = sigmaz(b_spin)
            sx = sigmax(b_spin)
            sy = sigmay(b_spin)

            # Build a set of local observables in the local basis
            lobs = Dict("sigmax"=>sx,"sigmay"=>sy,"sigmaz"=>sz)

            g = 1.
            gamma = 1.
            V = 2/2.

            # Test assertions in constructors
            L = NdLattice((4,))
            H = hamiltonian(L, g/2 * sx, V/4, sz)
            J = dissipators(L, [sqrt(2gamma) * sm])
            # Incompatible lattice
            L_err = NdLattice((5,))
            @test_throws AssertionError NdSystem(L_err, H, (V/4,), sz, J, lobs)
            # Number of tunnelling rates incompatible with number of dimensions
            @test_throws AssertionError NdSystem(L, H, (V/4,V/4), sz, J, lobs)
            # Incompatible operators
            @test_throws ErrorException NdSystem(L, H, (V/4,), destroy(FockBasis(1)), J, lobs)

            # Test constructor with missing observables
            s = NdSystem(L, H, (V/4,), sz, J)
            @test eltype(s.observables) == Dict{String,typeof(H)}

            # Broken: cannot construct a system with 1 site.

            function compare_systems(s1::AbstractSystem,s2::AbstractSystem)
                for f in fieldnames(typeof(s1))
                    if f == :lattice
                        L1 = getfield(s1,f)
                        L2 = getfield(s2,f)
                        for ff in fieldnames(typeof(L1))
                            @test getfield(L1,ff) == getfield(L2,ff)
                        end
                    else
                        @test getfield(s1,f) == getfield(s2,f)
                    end
                end
            end
            # Simpler constructor for homogeneous coupling
            L = NdLattice((2,1,1,1))
            H = hamiltonian(L, g/2 * sx, V/4, sz)
            J = dissipators(L, [sqrt(2gamma) * sm])
            s1 = NdSystem(L, H, V/4, sz, J, lobs)
            s2 = NdSystem(L, H, Tuple([V/4 for i in 1:4]), sz, J, lobs)
            compare_systems(s1,s2)
            s1 = NdSystem(L, H, V/4, sz, J)
            s2 = NdSystem(L, H, Tuple([V/4 for i in 1:4]), sz, J)
            compare_systems(s1,s2)

            # Test new constructor
            L = NdLattice((2,3); periodic=true)
            H = hamiltonian(L, g/2 * sx, V/4/2., sz)
            J = dissipators(L, [sqrt(gamma) * sm])
            s1 = NdSystem(L, H, (V/4,V/4), (sz,sz), J, lobs)
            s2 = NdSystem(L, H, V/4, sz, J, lobs)
            compare_systems(s1,s2)
            s3 = NdSystem(L, H, (V/4, V/4), sz, J, lobs)
            compare_systems(s1,s3)
            s4 = NdSystem(L, H, V/4, (sz,sz), J, lobs)
            compare_systems(s1,s4)

            # Test multiple coupling
            @test @no_error s1 = NdSystem(L, H, V/4, ([sx,sz],[sx,sz]), J, lobs)
            @test @no_error s2 = NdSystem(L, H, (V/4,V/4), ([sx,sz],[sx,sz]), J, lobs)
            compare_systems(s1,s2)

            # Test construction and merging
            # 1D
            for pbc in [true, false]
                L = NdLattice((4,); periodic=pbc)
                H = hamiltonian(L, g/2 * sx, V/4, sz)
                J = dissipators(L, [sqrt(2gamma) * sm])
                s = NdSystem(L, H, (V/4,), sz, J, lobs)
                ρ = steadystate.master(s)[2][end]
                sUs = merge(s,s,1,ρ,ρ,256)

                L2 = NdLattice((8,); periodic=pbc)
                H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                J2 = dissipators(L2, [sqrt(2gamma) * sm])
                s2 = NdSystem(L2, H2, (V/4,), sz, J2, lobs)

                @test maximum(abs2.(eigvals(sUs.H.data) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                @test maximum([maximum(abs2.(eigvals(sUs.J[i].data'sUs.J[i].data) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(sUs.J)]) ≈ 0. atol=eps(Float64)

                d = 1
                @test maximum([maximum(abs2.(eigvals(sUs.Htint[d][i].data) .- eigvals(Matrix(s2.Htint[d][i].data)))) for i in 1:length(sUs.Htint[d])]) ≈ 0. atol=eps(Float64)
                @test maximum([maximum(abs2.(eigvals(sUs.Htext[d][i].data) .- eigvals(Matrix(s2.Htext[d][i].data)))) for i in 1:length(sUs.Htext[d])]) ≈ 0. atol=eps(Float64)
            end

            # 2D
            for pbc in [true, false]
                L = NdLattice((2,2); periodic=pbc)
                H = hamiltonian(L, g/2 * sx, V/4, sz)
                J = dissipators(L, [sqrt(2gamma) * sm])
                s = NdSystem(L, H, (V/4, V/4), sz, J, lobs)
                ρ = steadystate.master(s)[2][end]
                sUs = merge(s,s,2,ρ,ρ,256)

                L2 = NdLattice((2,4); periodic=pbc)
                H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                J2 = dissipators(L2, [sqrt(2gamma) * sm])
                s2 = NdSystem(L2, H2, (V/4, V/4), sz, J2, lobs)

                @test maximum(abs2.(eigvals(sUs.H.data) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                @test maximum([maximum(abs2.(eigvals(sUs.J[i].data'sUs.J[i].data) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(sUs.J)]) ≈ 0. atol=eps(Float64)

                for d in 1:2
                    @test maximum([maximum(abs2.(eigvals(sUs.Htint[d][i].data) .- eigvals(Matrix(s2.Htint[d][i].data)))) for i in 1:length(sUs.Htint[d])]) ≈ 0. atol=eps(Float64)
                    @test maximum([maximum(abs2.(eigvals(sUs.Htext[d][i].data) .- eigvals(Matrix(s2.Htext[d][i].data)))) for i in 1:length(sUs.Htext[d])]) ≈ 0. atol=eps(Float64)
                end

                # Test pretty-printing
                @test @no_error repr(s)
                @test @no_error repr(MIME("text/plain"),s)
                @test @no_error repr(MIME("application/prs.juno.inline"),s)
                @test @no_error repr(sUs)
                @test @no_error repr(MIME("text/plain"),sUs)
                @test @no_error repr(MIME("application/prs.juno.inline"),sUs)
            end
            # TO DO: Test construction and merging of generic lattices

            # Dummy tests on plotting
            L = NdLattice((4,))
            H = hamiltonian(L, g/2 * sx, V/4, sz)
            J = dissipators(L, [sqrt(2gamma) * sm])
            s = NdSystem(L, H, (V/4,), sz, J)
            @test @no_error gplot(s);
            @test @no_error gplot(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            @test @no_error plot_system(s);
            @test @no_error plot_system(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            x = gplot(s);
            y = gplot(L.L);
            @test typeof(x) == typeof(y)
            x = plot_system(s);
            @test typeof(x) == typeof(y)
            x = gplot(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            y = gplot(L.L, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            @test all([typeof(getfield(x,f)) == typeof(getfield(y,f)) for f in fieldnames(typeof(x))])
            x = plot_system(s, Float64.(collect(1:nv(L))), Float64.(collect(1:nv(L))));
            @test all([typeof(getfield(x,f)) == typeof(getfield(y,f)) for f in fieldnames(typeof(x))])
        end;

        @testset "CornerBasis" begin
            # Test constructors
            b1 = SpinBasis(1//2)
            b2 = CompositeBasis([b1 for i in 1:10])
            bC = CornerBasis(b1, b2, 10)
            @test bC == CornerBasis([10],10,[[2], fill(2,10)])
            bC2 = CornerBasis(bC, bC, 10)
            @test bC2 == CornerBasis([10],10,vcat(bC.local_shapes, bC.local_shapes))

            # Test Base.== overload
            bC = CornerBasis(10)
            @test bC == deepcopy(bC)
            @test bC !== SpinBasis(1//2)
            @test bC !== FockBasis(1)
            @test bC !== GenericBasis(1)
            bC2 = CornerBasis([11],bC.M,bC.local_shapes)
            @test bC !== bC2
            bC2 = CornerBasis(bC.shape,bC.M+1,bC.local_shapes)
            @test bC !== bC2
            bC2 = CornerBasis(bC.shape,bC.M,vcat(bC.local_shapes, [[1]]))
            @test bC !== bC2

            # Test pretty-printing
            @test @no_error repr(bC)
            @test @no_error repr(MIME("text/plain"),bC)
            @test @no_error repr(MIME("application/prs.juno.inline"),bC)
        end;
    end;

    @testset "Corner algorithm" begin
        @testset "Core" begin
            # Test orthonormalisation
            N = 10
            U = diagm([i=>ones(ComplexF64,N-i) for i in 0:N]...)
            @test maximum(abs2.(CornerSpaceRenorm.cgs(U) .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)
            @test maximum(abs2.(CornerSpaceRenorm.mgs(U) .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)
            V = copy(U); CornerSpaceRenorm.cgs!(V)
            @test maximum(abs2.(V .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)
            V = copy(U); CornerSpaceRenorm.mgs!(V)
            @test maximum(abs2.(V .- Diagonal(ones(ComplexF64,N)))) ≈ 0. atol=eps(Float64)

            # Test maxk
            x = randperm(1000)
            idcs, y = CornerSpaceRenorm.maxk(x,100)
            # Test that indices match sorted values
            @test all([x[idcs[i]] == y[i] for i in 1:100])
            # Test that k sorted values are the largest
            @test all([val == y[i] for (i,val) in enumerate(1000:-1:901)])

            # Test max_prod_pairs
            x1 = 1:50; x2 = 1:50
            idcs, pairs = CornerSpaceRenorm.max_prod_pairs(x1,x2,10)
            @test all([x1[idx[1]] * x2[idx[2]] for idx in idcs] .== pairs)
            @test all(sort([x1[i] * x2[j] for i in 1:50, j in 1:50][:],rev=true)[1:10] .== pairs)

            # Test hermitianize
            b = GenericBasis(50)
            Op = randoperator(b)
            OpH = hermitianize(Op)
            @test ishermitian(OpH.data)
            @test maximum(abs2.(OpH.data .- (Op + dagger(Op)).data/2.)) ≈ 0. atol=eps(Float64)
            Op = randoperator(b);
            OpH = copy(Op);
            hermitianize!(OpH)
            @test ishermitian(OpH.data)
            @test maximum(abs2.(OpH.data .- (Op + dagger(Op)).data/2.)) ≈ 0. atol=eps(Float64)
        end;

        @testset "Merging" begin
            @testset "NdSystem" begin
                # Local basis
                b_spin = SpinBasis(1//2)

                # Build some operators in the local basis
                sm = sigmam(b_spin)
                sp = sigmap(b_spin)
                sz = sigmaz(b_spin)
                sx = sigmax(b_spin)
                sy = sigmay(b_spin)

                # Build a set of local observables in the local basis
                lobs = Dict("sigmax"=>sx,"sigmay"=>sy,"sigmaz"=>sz)

                g = 1.
                gamma = 1.
                V = 2/2.

                L1 = NdLattice((2,1))
                H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                J1 = dissipators(L1, [sqrt(2gamma) * sm])
                s1 = NdSystem(L1, H1, (V/4,V/4), sz, J1, lobs)
                ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                L2 = NdLattice((1,2))
                H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                J2 = dissipators(L2, [sqrt(2gamma) * sm])
                s2 = NdSystem(L2, H2, (V/4,V/4), sz, J2, lobs)
                ρ2 = steadystate.master(s2.H,s2.J)[2][end]

                @test_throws DimensionMismatch merge(s1, s2, 1, ρ1, ρ2, 2)
                @test_throws DimensionMismatch merge(s1, s2, 2, ρ1, ρ2, 2)

                L1_pbc = NdLattice((2,1);periodic=true)
                H1_pbc = hamiltonian(L1_pbc, g/2 * sx, V/4, sz)
                J1_pbc = dissipators(L1_pbc, [sqrt(2gamma) * sm])
                s1_pbc = NdSystem(L1_pbc, H1_pbc, (V/4,V/4), sz, J1_pbc, lobs)
                ρ1_pbc = steadystate.master(s1_pbc.H,s1_pbc.J)[2][end]

                @test_throws AssertionError merge(s1,s1_pbc,1,ρ1,ρ1_pbc,2)

                # Full space corner vs exact system
                # 1D
                for pbc in [true, false]
                    L1 = NdLattice((2,); periodic=pbc)
                    H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                    J1 = dissipators(L1, [sqrt(2gamma) * sm])
                    s1 = NdSystem(L1, H1, (V/4,), sz, J1, lobs)
                    ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                    s1 = merge(s1,s1,1,ρ1,ρ1,16)
                    ρ1 = steadystate.master(s1.H,s1.J;tol=1e-6)[2][end]

                    L2 = NdLattice((4,); periodic=pbc)
                    H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                    J2 = dissipators(L2, [sqrt(2gamma) * sm])
                    s2 = NdSystem(L2, H2, (V/4,), sz, J2, lobs)
                    ρ2 = steadystate.master(s2.H,s2.J;tol=1e-6)[2][end]

                    @test maximum(abs2.(eigvals(Matrix(s1.H.data)) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                    @test maximum([maximum(abs2.(eigvals(Matrix(s1.J[i].data's1.J[i].data)) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(s1.J)]) ≈ 0. atol=eps(Float64)
                    @test maximum(abs2.(eigvals(ρ1.data) .- eigvals(ρ2.data))) ≈ 0. atol=1e-5
                end

                # 2D
                for pbc in [true, false]
                    L1 = NdLattice((2,1); periodic=pbc)
                    H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                    J1 = dissipators(L1, [sqrt(2gamma) * sm])
                    s1 = NdSystem(L1, H1, (V/4,V/4), sz, J1, lobs)
                    ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                    s1 = merge(s1,s1,1,ρ1,ρ1,16)
                    ρ1 = steadystate.master(s1.H,s1.J;tol=1e-5)[2][end]

                    L2 = NdLattice((4,1); periodic=pbc)
                    H2 = hamiltonian(L2, g/2 * sx, V/4, sz)
                    J2 = dissipators(L2, [sqrt(2gamma) * sm])
                    s2 = NdSystem(L2, H2, (V/4,V/4), sz, J2, lobs)
                    ρ2 = steadystate.master(s2.H,s2.J;tol=1e-5)[2][end]

                    @test maximum(abs2.(eigvals(Matrix(s1.H.data)) .- eigvals(Matrix(s2.H.data)))) ≈ 0. atol=eps(Float64)
                    @test maximum([maximum(abs2.(eigvals(Matrix(s1.J[i].data's1.J[i].data)) .- eigvals(Matrix(s2.J[i].data's2.J[i].data)))) for i in 1:length(s1.J)]) ≈ 0. atol=eps(Float64)
                    @test maximum(abs2.(eigvals(ρ1.data) .- eigvals(ρ2.data))) ≈ 0. atol=1e-5
                end

                # Check returning the product density matrix when merging
                L1 = NdLattice((2,))
                H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
                J1 = dissipators(L1, [sqrt(2gamma) * sm])
                s1 = NdSystem(L1, H1, (V/4,), sz, J1, lobs)
                ρ1 = steadystate.master(s1.H,s1.J)[2][end]
                s1, ρ02 = merge(s1,s1,1,ρ1,ρ1,16;return_dm=true)
                @test tr(ρ02) ≈ 1.0 atol=1e-9
                @test opnorm(ρ02.data .- ρ02.data) ≈ 0.0 atol=1e-9
                @test all(eigvals(hermitianize(ρ02).data) .> 0.)
                tss = steadystate.master(s1.H,s1.J;tol=1e-5)[1][end]
                tss_init = steadystate.master(s1.H,s1.J;rho0=ρ02,tol=1e-5)[1][end]
                @test tss_init < tss

                # Merging different lattices
                Ls = NdLattice((3,))
                Hs = hamiltonian(Ls, g/2 * sx, V/4, sz)
                Js = dissipators(Ls, [sqrt(2gamma) * sm])
                ss = NdSystem(Ls, Hs, V/4, sz, Js, lobs)
                ρs = steadystate.master(ss)[2][end]

                t = 1.
                Δ = 1.
                κ = 1.
                F = 5e-1 * κ

                lb = FockBasis(1)
                Lb = NdLattice((3,),periodic=false)
                Hb = hamiltonian(Lb, -Δ * number(lb) + F * (create(lb) + destroy(lb)), t/2., destroy(lb))
                Jb = dissipators(Lb,[sqrt(κ) * destroy(lb)])
                sb = NdSystem(Lb,Hb,(t/2.,),destroy(lb),Jb)
                ρb = steadystate.master(sb)[2][end]

                @test_logs (:info,"The two input systems do not have equal tunnelling rates along all directions.") s_sb = merge(ss,sb,1,ρs,ρb,10)
            end;
        end;
    end;

    @testset "steadystate" begin
        # Local basis
        b_spin = SpinBasis(1//2)

        # Build some operators in the local basis
        sm = sigmam(b_spin)
        sp = sigmap(b_spin)
        sz = sigmaz(b_spin)
        sx = sigmax(b_spin)
        sy = sigmay(b_spin)

        # Build a set of local observables in the local basis
        lobs = Dict("sigmax"=>sx,"sigmay"=>sy,"sigmaz"=>sz)

        g = 1.
        gamma = 1.
        V = 2/2.

        L1 = NdLattice((6,))
        H1 = hamiltonian(L1, g/2 * sx, V/4, sz)
        J1 = dissipators(L1, [sqrt(2gamma) * sm])
        s1 = NdSystem(L1, H1, (V/4,), sz, J1, lobs)
        ρ1 = steadystate.master(s1.H,s1.J; tol=1e-5)[2][end]
        ρ2 = steadystate_bicg(s1)
        @test fidelity(ρ1, ρ2) ≈ 1. atol=1e-5
        ρ3, log = steadystate_bicg(s1; log=true)
        @test typeof(log) == ConvergenceHistory{true,Nothing}
        ρ4 = CornerSpaceRenorm.steadystate_bicg_LtL(s1)
        @test fidelity(ρ1, ρ4) ≈ 1. atol=1e-5
        ρ5, log = steadystate_bicg!(ρ3, s1; log=true)
        @test log.iters == 0
        ρ, log = steadystate_bicg(s1; log=true)
        @test typeof(log) == ConvergenceHistory{true,Nothing}
        ρ, log = CornerSpaceRenorm.steadystate_bicg_LtL(s1; log=true)
        @test typeof(log) == ConvergenceHistory{true,Nothing}
    end;

end;

ext: jl


text:
using PyPlot, PyCall, IJulia

"""
    t = delaunay(p)

Delaunay triangulation `t` of N x 2 node array `p`.
"""
function delaunay(p)
    tri = pyimport("matplotlib.tri")
    t = tri[:Triangulation](p[:,1], p[:,2])
    return Int64.(t[:triangles] .+ 1)
end

"""
    edges, boundary_indices = all_edges(t)

Find all unique edges in the triangulation `t` (ne x 2 array)
Second output is indices to the boundary edges.
"""
function all_edges(t)
    edges = vcat(t[:,[1,2]], t[:,[2,3]], t[:,[3,1]])
    edges = sortslices(sort(edges, dims=2), dims=1)
    dup = all(edges[2:end,:] - edges[1:end-1,:] .== 0, dims=2)[:]
    keep = .![false;dup]
    edges = edges[keep,:]
    dup = [dup;false]
    dup = dup[keep]
    return edges, findall(.!dup)
end

"""
    e = boundary_nodes(t)

Find all boundary nodes in the triangulation `t`.
"""
function boundary_nodes(t)
    edges, boundary_indices = all_edges(t)
    return unique(edges[boundary_indices,:][:])
end

"""
    tplot(p, t, u=nothing)

If `u` == nothing: Plot triangular mesh with nodes `p` and triangles `t`.
If `u` == solution vector: Plot filled contour color plot of solution `u`.
"""
function tplot(p, t, u=nothing)
    clf()
    axis("equal")
    if u == nothing
        tripcolor(p[:,1], p[:,2], t .- 1, 0*t[:,1],
                  cmap="Set3", edgecolors="k", linewidth=1)
    else
        tricontourf(p[:,1], p[:,2], t .- 1, u, 20)
    end
    draw()
end

"""
    inside = inpolygon(p, pv)

Determine if each point in the N x 2 node array `p` is inside the polygon
described by the NE x 2 node array `pv`.
"""
function inpolygon(p::Array{Float64,2}, pv::Array{Float64,2})
    path = pyimport("matplotlib.path")
    poly = path[:Path](pv)
    inside = [poly[:contains_point](p[ip,:]) for ip = 1:size(p,1)]
end

function remove_outside_tris(p, t, pv)
    pmid = dropdims(sum(p[t,:], dims=2), dims=2) / 3
    is_inside = inpolygon(pmid, pv)
    t = t[is_inside,:]
end

function triarea(p, t)
    d12 = @. p[t[:,2],:] - p[t[:,1],:]
    d13 = @. p[t[:,3],:] - p[t[:,1],:]
    @. abs(d12[:,1] * d13[:,2] - d12[:,2] * d13[:,1]) / 2
end

function remove_tiny_tris(p, t)
    t = t[triarea(p,t) .> 1e-14,:]
end

function circumcenter(p)
    dp1 = @. p[2,:] - p[1,:]
    dp2 = @. p[3,:] - p[1,:]

    mid1 = @. (p[1,:] + p[2,:]) / 2
    mid2 = @. (p[1,:] + p[3,:]) / 2

    s = [ -dp1[2] dp2[2]; dp1[1] -dp2[1]] \ (-mid1 .+ mid2)
    pc = @. mid1' + s[1] * [-dp1[2] dp1[1]]
end

function edge_midpoints(p, t)
    pmid = reshape(p[t,:] + p[t[:,[2,3,1]],:], :, 2) / 2
    pmid = unique(pmid, dims=1)
end

function pmesh(pv, hmax, nref)
    p = zeros(Float64, 0, 2)
    for i = 1:size(pv,1) - 1
        pp = pv[i:i+1,:]
        L = sqrt(sum(diff(pp, dims=1).^2, dims=2))[1]
        if L > hmax
            n = ceil(Int, L / hmax)
            ss = (0:n) / n
            pp = [1 .- ss ss] * pp
        end
        p = [p; pp[1:end-1,:]]
    end

    t = zeros(Int64, 0, 3)
    while true
        t = delaunay(p)
        t = remove_tiny_tris(p, t)
        t = remove_outside_tris(p, t, pv)
        # tplot(p,t), pause(1e-3)

        area = triarea(p, t)
        maxarea, ix = findmax(area)
        if maxarea < hmax^2 / 2
            break
        end
        pc = circumcenter(p[t[ix,:],:])
        p = [p; pc]
    end

    for iref = 1:nref
        p = [p; edge_midpoints(p, t)]
        t = delaunay(p)
        t = remove_tiny_tris(p, t)
        t = remove_outside_tris(p, t, pv)
        # tplot(p, t), pause(1e-3)
    end

    e = boundary_nodes(t)
    p, t, e
end

function pmesh_test()
    pv = [0 0; 1 0; .5 .5; 1 1; 0 1; 0 0]
    p,t,e = pmesh(pv, 0.2, 1);
    tplot(p,t)
end

using SparseArrays
using LinearAlgebra
function coefficients(RHS, p, t, k)
    LHS = [1 p[t[k,1],1] p[t[k,1],2]; 1 p[t[k,2],1] p[t[k,2],2]; 1 p[t[k,3],1] p[t[k,3],2]]
    coeff = LHS\RHS'
    return coeff'
end

function element(coeffmatrix, a, b)
    return coeffmatrix[a,2]*coeffmatrix[b,2] + coeffmatrix[a,3]*coeffmatrix[b,3]
end

function fempoi(p, t, e)
    areas = triarea(p,t)
    n = size(p)[1]
    A = spzeros(n, n); b = zeros(n);
    for k = 1:size(t)[1]
        cm = [coefficients([1 0 0] , p, t, k) ;
              coefficients([0 1 0] , p, t, k) ;
              coefficients([0 0 1] , p, t, k) ]
        insert = [element(cm, 1, 1) element(cm, 1, 2) element(cm, 1, 3);
                  element(cm, 2, 1) element(cm, 2, 2) element(cm, 2, 3);
                  element(cm, 3, 1) element(cm, 3, 2) element(cm, 3, 3)]
        insert *=  areas[k]

        A[t[k,:],t[k,:]] = A[t[k,:],t[k,:]]+insert
        #aa = areas[k]/3.0
        b[t[k,:]] = b[t[k,:]].+ areas[k]/3.0
    end
    for k = 1:size(e)[1]
        i=e[k]
        A[i,:].=0.0
        A[i,i]=1.0
        b[i]=0.0
    end
    dropzeros!(A)
    println(b)
    println(A)

    return A \ b
end

function poiconv(pv, hmax, nrefmax)
   pmax, tmax, emax = pmesh(pv, hmax, nrefmax)
   umax = fempoi(pmax, tmax, emax)
   errors = zeros(nrefmax)
   for ncur = 0:(nrefmax-1)
       pprox, tprox, eprox = pmesh(pv, hmax, ncur)
       uprox = fempoi(pprox, tprox, eprox)
       errorarray = maximum(abs.(umax[1:size(uprox)[1]] - uprox))
       #errors = [errors ; errorarray]
       errors[ncur+1]=errorarray
   end

   return errors
end


text:
#define ZE_MAKE_VERSION( _major, _minor )  (( _major << 16 )|( _minor & 0x0000ffff))
unmake_version(ver) = VersionNumber(Int(ver) >> 16, Int(ver) & 0x0000ffff)


text:
# Run tests for the inner solves

using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test,Random
using LinearAlgebra

Random.seed!(0);
@bench @testset "Inner Solves" begin
    dep=nep_gallery("dep0",200);
    n=size(dep,1);
    Random.seed!(0);

    nn=opnorm(compute_Mder(dep,0));
    errmeasure= (λ,v) -> norm(compute_Mlincomb(dep,λ,v))/nn;

    pnep=create_proj_NEP(dep);

    k = 5
    Q,R=qr(randn(n,k));
    Q = Matrix(Q)
    set_projectmatrices!(pnep,Q,Q)

    logger=ErrorLogger(50,100,0);
    λv,V = inner_solve(DefaultInnerSolver(), ComplexF64, pnep; σ=0.0+0im, neigs=3, tol=eps()*100, inner_logger=logger)
    verify_lambdas(3, pnep, λv, V,eps()*10000)
    E=logger.errs[:,1]; E=E[(!isnan).(E)];

    @test length(E) > 5 #Has done more than 5 iterations
    @test E[end] <eps()*10000 # Error measure fulfills stopping criteria

    logger=ErrorLogger(1,100,0);
    λv,V = inner_solve(NewtonInnerSolver(), ComplexF64, pnep; λv=[0.0,1.0] .+ 0im, V=ones(k, 2), tol=eps()*100, inner_logger=logger)
    verify_lambdas(2, pnep, λv, V,eps()*10000)
    E=logger.errs[:,1]; E=E[(!isnan).(E)];
    @test length(E) > 5 #Has done more than 5 iterations
    @test E[end] <eps()*10000 # Error measure fulfills stopping criteria

    #λv,V=inner_solve(SGIterInnerSolver(),pnep,λv=[0.0],j=1);
    #verify_lambdas(2, pnep, λv, V,eps()*10000)

    λv,V = inner_solve(IARChebInnerSolver(), ComplexF64, pnep; λv=[0,1,2,3] .+ 0.0im)
    verify_lambdas(10, pnep, λv, V, sqrt(eps()))

    # TODO: this test results in a "Rank drop" warning, and a third eigenvalue that's not converged
    λv,V = inner_solve(ContourBeynInnerSolver(), ComplexF64, pnep; λv=[0,1] .+ 0.0im, neigs=3)
    # @test minimum(svdvals(compute_Mder(pnep,λv[1]))) < eps()*100
    verify_lambdas(2, pnep, λv[1:2], V[:,1:2], sqrt(eps()))

    λv,V = inner_solve(NleigsInnerSolver(), ComplexF64, pnep; λv=[0,1,2,3] .+ 0.0im)
    @test size(λv,1) >= 5;
end


text:
module ExitWindow

using Gtk, Gtk.ShortNames
using ..Ahorn

function exitAhorn(widget=nothing, event=nothing)
    dialogText = "You might have unsaved changes in your map.\nPlease confirm that you want to exit the program."
    res = Ahorn.loadedState.lastSavedHash == hash(Ahorn.loadedState.side) || ask_dialog(dialogText, Ahorn.window)

    # For detecting unpredicted shutdowns
    Ahorn.persistence["currently_running"] = false

    Ahorn.saveConfig(Ahorn.config, true)
    Ahorn.saveConfig(Ahorn.persistence, true)

    # If called via Exit button instead of "X" on window
    if isa(widget, Ahorn.MenuItemsTypes) && res
        destroy(Ahorn.window)
    end

    return !res
end

end


text:
#
# Factors Access
#
# For ϕ[:A] syntax and such

# Column access (ϕ[:A]) returns the dimension for comparisons and stuff
Base.getindex(ϕ::Factor, dims::NodeNameUnion) = pattern(ϕ, dims)

# Index by number gets that ind
Base.getindex(ϕ::Factor, i::Int) = ϕ.potential[i]
Base.getindex(ϕ::Factor, I::Vararg{Int}) = ϕ.potential[I]

"""
    getindex(ϕ, a)

Get values with dimensions consistent with an assignment.
Colons select entire dimension.
"""
function Base.getindex(ϕ::Factor, a::Assignment)
    # more complicated logic than finding indices b/c also updating
    #  dimensions
    inds = _translate_index(ϕ, a)
    keep = inds .== Colon()
    new_dims = ϕ.dimensions[keep]
    @inbounds new_p = ϕ.potential[inds...]

    # in case array access returns a scalar and not an array
    # as always, the weird edge case to get a zero-dimensional array and not
    #  a scalar
    if ndims(new_p) == 0
        new_p = squeeze([new_p], 1)
    end

    return Factor(new_dims, new_p)
end
Base.getindex(ϕ::Factor, pair::Pair{NodeName}...) = Base.getindex(ϕ, Assignment(pair))

function Base.setindex!(ϕ::Factor, v, a::Assignment)
    @inbounds return ϕ.potential[_translate_index(ϕ, a)...] = v
end

@inline function _translate_index(ϕ::Factor, a::Assignment)
    inds = Array{Any}(ndims(ϕ))
    inds[:] = Colon()

    for (i, dim) in enumerate(ϕ.dimensions)
        if haskey(a, dim)
            ind = a[dim]

            if isa(ind, Colon)
                continue
            elseif isa(ind, Int)
                if ind < 1 || ind > size(ϕ, dim)
                    throw(BoundsError(dim, ind))
                end
            else
                throw(TypeError(:getindex, "Invalid state for dimension $dim",
                            Int, typeof(ind)))
            end

            inds[i] = ind
        end
    end

    return inds
end

Base.sub2ind(ϕ::Factor, i...) = sub2ind(size(ϕ), i...)
Base.ind2sub(ϕ::Factor, i) = ind2sub(size(ϕ), i)


text:
using PyPlot
using Jagot.plotting
plot_style("ggplot")

using PyCall
Cycler = pyimport("cycler")
plt.rc("axes", prop_cycle=plt.rcParams["axes.prop_cycle"])

using SphericalBesselFunctions

include(joinpath(@__DIR__, "../test/reference.jl"))

function simple_example()
    nx = 1000
    x = range(-12, stop=12, length=nx)

    nℓ = 10
    j, j′, y, y′ = bessels(x, nℓ)

    cfigure("simple example",figsize=(7,6)) do
        csubplot(221,nox=true) do
            plot(x, j)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            ylabel(L"j_n(x)")
        end
        csubplot(222,nox=true) do
            plot(x, y)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            ylabel(L"y_n(x)")
            axes_labels_opposite(:y)
        end
        csubplot(223) do
            plot(x, j′)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            xlabel(L"x")
            ylabel(L"j'_n(x)")
        end
        csubplot(224) do
            plot(x, y′)
            margins(0, 0.1)
            ylim(-0.5, 0.5)
            xlabel(L"x")
            ylabel(L"y'_n(x)")
            axes_labels_opposite(:y)
        end
    end
    savefig("docs/src/figures/simple-example.svg")
end

function accuracy()
    nx = 1001
    x = 10 .^ range(-1, stop=4, length=nx)
    nℓ = 105
    j, j′, y, y′ = @time bessels(x, nℓ)

    sel = 1:min(nℓ,104)

    nℓref = min(nℓ,length(sel))
    jref = zeros(nx, nℓref)
    j′ref = zeros(nx, nℓref)
    yref = zeros(nx, nℓref)
    y′ref = zeros(nx, nℓref)

    @time for (i,x) in enumerate(x)
        bessels_ref!(view(jref, i, :),
                     view(j′ref, i, :),
                     view(yref, i, :),
                     view(y′ref, i, :),
                     x)
    end

    j = j[:,sel]
    j′ = j′[:,sel]
    y = y[:,sel]
    y′ = y′[:,sel]
    jref = jref[:,sel]
    j′ref = j′ref[:,sel]
    yref = yref[:,sel]
    y′ref = y′ref[:,sel]

    cfigure("accuracy", figsize=(9,10)) do
        m,n = 4,3
        ks = reshape(1:m*n, n, m)'

        csubplot(m,n,ks[1,1],nox=true) do
            semilogx(x, j, rasterized=true)
            ylim(-0.5,0.5)
            ylabel(L"j_n(x)")
            title("SphericalBesselFunctions.jl")
        end
        csubplot(m,n,ks[2,1],nox=true) do
            semilogx(x, j′, rasterized=true)
            ylim(-0.5,0.5)
            ylabel(L"j_n'(x)")
        end
        csubplot(m,n,ks[3,1],nox=true) do
            semilogx(x, y, rasterized=true)
            ylim(-0.5,0.5)
            ylabel(L"y_n(x)")
        end
        csubplot(m,n,ks[4,1]) do
            semilogx(x, y′, rasterized=true)
            ylim(-0.5,0.5)
            xlabel(L"x")
            ylabel(L"y_n'(x)")
        end

        csubplot(m,n,ks[1,2],nox=true,noy=true) do
            semilogx(x, jref, rasterized=true)
            ylim(-0.5,0.5)
            title("SpecialFunctions.jl")
        end
        csubplot(m,n,ks[2,2],nox=true,noy=true) do
            semilogx(x, j′ref, rasterized=true)
            ylim(-0.5,0.5)
        end
        csubplot(m,n,ks[3,2],nox=true,noy=true) do
            semilogx(x, yref, rasterized=true)
            ylim(-0.5,0.5)
        end
        csubplot(m,n,ks[4,2],noy=true) do
            semilogx(x, y′ref, rasterized=true)
            ylim(-0.5,0.5)
            xlabel(L"x")
        end

        csubplot(m,n,ks[1,3],nox=true) do
            jerror = abs.(j-jref)
            loglog(x, jerror, rasterized=true)
            ylim(max(1e-19,minimum(jerror)), max(maximum(jerror),1e-4))
            title("Difference")
            axes_labels_opposite(:y)
        end
        csubplot(m,n,ks[2,3],nox=true) do
            j′error = abs.(j′-j′ref)
            loglog(x, j′error, rasterized=true)
            ylim(max(1e-19,minimum(j′error)), max(maximum(j′error),1e-4))
            axes_labels_opposite(:y)
        end
        csubplot(m,n,ks[3,3],nox=true) do
            yerror = abs.(y-yref)
            loglog(x, yerror, rasterized=true)
            ylim(max(1e-19,minimum(yerror)), max(maximum(yerror),1e-4))
            axes_labels_opposite(:y)
        end
        csubplot(m,n,ks[4,3]) do
            loglog(x, abs.(y′-y′ref), rasterized=true)
            xlabel(L"x")
            axes_labels_opposite(:y)
        end
    end
    savefig("docs/src/figures/accuracy.svg")
end

macro echo(expr)
    println(expr)
    :(@time $expr)
end

@info "Documentation plots"
mkpath("docs/src/figures")
@echo simple_example()
@echo accuracy()


text:
using FileToRigidTransform

using Test
@testset "Configuration" begin
    import FileToRigidTransform: specbyname,
                                read_dlm_line_after,
                                DevConfig,
                                devices_configuration_vector
    configfi = joinpath(@__DIR__, "..", "example", "config_template.txt")
    if !isfile(configfi)
        configfi = joinpath(@__DIR__, "example", "config_template.txt")
    end
    reserved = Set(["Surge", "Sway", "Heave", "Roll", "Pitch", "Yaw"])
    defined, lp = read_dlm_line_after(configfi)
    @test lp == 626
    @test intersect(reserved, defined) == reserved
    nextline, lastpos = read_dlm_line_after(configfi, startpos = lp)
    @test nextline ==  ["2", "4", "6", "8", "10", "12", "1", "1"]
    spec = specbyname(configfi)
    @test length(spec) == 8
    @test typeof(DevConfig("template", spec)) <: DevConfig
    # test would only work after configuration:
    #@test typeof(devices_configuration_vector()) <: Vector{DevConfig}
end
@testset "Interactive tests of show methods and loggers" begin
    import FileToRigidTransform: Bytevector, DevState
    ds = DevState([1,2,3],time(),true)
    dsprev = DevState([2,4,6],time(),true)
    import FileToRigidTransform.DevConfig
    devconf = DevConfig()
    accumulated = Vector{Int64}()
    ios = IOBuffer()
    import FileToRigidTransform: log_by_channel,
                                 log_by_bit,
                                 log_by_bit_accumulate_changes,
                                 minus_ds, specbyname, extract
    # Empty device configuration
    ds = DevState([1,2,3],time(),true)
    dsprev = DevState([2,4,6],time(),true)

    log_by_bit(ios, devconf, accumulated, dsprev, ds)
    @test_throws AssertionError log_by_channel(ios, devconf, accumulated, dsprev, ds)
    log_by_bit_accumulate_changes(ios, devconf, accumulated, dsprev, ds);

    # Template device configuration
    configfi = begin
        configfi = joinpath(@__DIR__, "..", "example", "config_template.txt")
        if !isfile(configfi)
            configfi = joinpath(@__DIR__, "example", "config_template.txt")
        end
        configfi
    end
    devconf = DevConfig("template", specbyname(configfi))
    ds = DevState(collect(1:13),time(),true)
    dsprev = DevState(collect(2:14),time(),true)
    log_by_bit(ios, devconf,  accumulated, dsprev, ds)
    log_by_channel(ios, devconf,  accumulated, dsprev, ds)
    log_by_bit_accumulate_changes(ios, devconf, accumulated, dsprev, ds);
    dsprev = DevState([],time(),true)
    log_by_bit(ios, devconf, accumulated, dsprev, ds)
    log_by_channel(ios, devconf, accumulated, dsprev, ds)
    log_by_bit_accumulate_changes(ios, devconf, accumulated, dsprev, ds)
end
using Test, FileToRigidTransform
@testset "bits manipulation" begin
    import FileToRigidTransform: drophighbits, bytestring, droplowbits
    @test bytestring(0b1100) == "0000 1100"
    @test drophighbits(7,0b11) == 1
    @test drophighbits(7,0b10) == 0
    @test drophighbits(5,0b1111) == 0b111
    @test drophighbits(5,0b1111) |> bytestring == "0000 0111"
    @test drophighbits(5,0b0100) |> bytestring == "0000 0100"
    droplowbits(1, UInt8(16)) == 8
    droplowbits(8, UInt8(255)) == 0
    droplowbits(1, UInt128(5)) == 2
end
using Test, FileToRigidTransform
@testset "Extract 1 byte or less from within one byte" begin
    import FileToRigidTransform: Bytevector, extract_type,
                                 extract, bytestring,
                                 droplowbits, drophighbits,
                                 bitlength_ceiling
    v = [0, 127, 255]
    bv = Bytevector(v)
    bytepos = 1
    bitpad = 0
    bitlength = 8
    bigendian = true
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 0
    bytepos = 2
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 127
    bytepos = 3
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 255
    bitpad = 1; bitlength = 7
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 127
    bitpad = 2; bitlength = 6
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 63
    bytepos = 2
    bitpad = 1
    bitlength = 7
    @test extract(bv, bytepos, 1, 7) == 127
    @test extract(bv, bytepos, 1, 6) == 63
    bitpad = 2; bitlength = 6
    @test extract(bv, bytepos, bitpad, bitlength) == 63
end

@testset "Extract 1 byte or less, crossing two source bytes" begin
    import FileToRigidTransform: Bytevector, extract_type,
                                 extract, bytestring,
                                 droplowbits, drophighbits,
                                 bitlength_ceiling
    v = [1, 192]
    bv = Bytevector(v)
    bytepos = 1
    bitpad = 7
    bitlength = 3
    bigendian = true
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 0b111
    bigendian = false
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 0b111
    v = [255, 255]
    bv = Bytevector(v)
    bigendian = true
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 0b111
    bigendian = false
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 0b111
    v = [0, 255, 0, 255]
    bv = Bytevector(v)
    bytepos = 1
    bitpad = 7
    bitlength = 11
    bigendian = true
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 0b00111111110
    bigendian = false
    @test extract(bv, bytepos, bitpad, bitlength, bigendian = bigendian) == 0b01111111100
end

using Test
@testset "extract more than 1 byte or less" begin
    import FileToRigidTransform: Bytevector, extract_type,
        extract, bytestring,
        droplowbits, drophighbits
    v = [255, 255, 255]
    bv = Bytevector(v)
    bytepos = 1; bitpad = 0; bitlength = 16;     bigendian = true
    @test extract(bv, bytepos, bitpad, bitlength, bigendian=bigendian) == 0b1111111111111111
    bigendian = false
    @test extract(bv, bytepos, bitpad, bitlength, bigendian=bigendian) == 0b1111111111111111
    bytepos = 1; bitpad = 1; bitlength = 16;     bigendian = true
    @test extract(bv, bytepos, bitpad, bitlength, bigendian=bigendian) == 0b1111111111111111
end


text:
# Testing Script for EMD module
using Base.Test

include("../src/EMDUtil.jl")

println("EMDUtil test started")
#test case for findExtrema function
y = Float64[1,3,1,2]
t = Float64[1,2,3,4]
max, min, tmax, tmin = findExtrema(y, t)

@test max[1] == 1
@test max[2] == 3
@test max[3] == 2

@test tmax[1] == 1
@test tmax[2] == 2
@test tmax[3] == 4

@test min[1] == 1
@test min[2] == 1
@test min[3] == 2

@test tmin[1] == 1
@test tmin[2] == 3
@test tmin[3] == 4


text:
# DARNN.jl
#
# Layers for a DA-RNN network
# Based on paper Quin et. al. "A Dual-Stage Attention-Based Recurrent Neural
# Network for Time Series Prediction" https://arxiv.org/abs/1704.02971

# Copy LSTM Code to prevent its special CUDA implementation
struct FALSTMCell{A,V,S}
	Wi::A
	Wh::A
	b::V
	state0::S
end

function FALSTMCell(in::Integer, out::Integer;
		init=Flux.glorot_uniform,
		initb=Flux.zeros32,
		init_state=Flux.zeros32)
	cell = FALSTMCell(init(out * 4, in), init(out * 4, out), initb(out * 4), (init_state(out, 1), init_state(out, 1)))
	cell.b[Flux.gate(out, 2)] .= 1
	return cell
end

function (m::FALSTMCell{A,V,<:NTuple{2,AbstractMatrix{T}}})((h, c), x::Union{AbstractVecOrMat{T},Flux.OneHotArray}) where {A,V,T}
	b, o = m.b, size(h, 1)
	g = m.Wi * x .+ m.Wh * h .+ b
	input = σ.(Flux.gate(g, o, 1))
	forget = σ.(Flux.gate(g, o, 2))
	cell = tanh.(Flux.gate(g, o, 3))
	output = σ.(Flux.gate(g, o, 4))
	c = forget .* c .+ input .* cell
	h′ = output .* tanh.(c)
	sz = size(x)
	return (h′, c), reshape(h′, :, sz[2:end]...)
end

Flux.@functor FALSTMCell
Recur(m::FALSTMCell) = Flux.Recur(m, m.state0)

Base.show(io::IO, l::FALSTMCell) =
	print(io, "FALSTMCell(", size(l.Wi, 2), ", ", size(l.Wi, 1) ÷ 4, ")")

# struct that stores layers and some extra information
mutable struct DARNNCell{A,B,C,D,E,F,W,X,Y,Z}
    # Encoder part
	encoder_lstm::A
	encoder_attn::B
    # Decoder part
	decoder_lstm::C
	decoder_attn::D
	decoder_fc::E
	decoder_fc_final::F
    # Index for original data etc
	encodersize::W
	decodersize::X
	orig_idx::Y
	poollength::Z
end

# Initialization function to obtain the right size of the hidden state
# after a `reset!` of the model.
function darnn_init(m::DARNNCell, x)
	s = size(x)
	m.encoder_lstm(simarray(x, s[1], s[4]))
	m.decoder_lstm(simarray(x, 1, s[4]))
	return nothing
end
simarray(x::Flux.CUDA.CuArray, size...) = Flux.CUDA.zeros(eltype(x), size...)
simarray(x, size...) = zeros(eltype(x), size...)

Flux.Zygote.@nograd darnn_init

# Dispatch on `trainable` and `reset!` to make standard Flux scripts work
Flux.trainable(m::DARNNCell) = (m.encoder_lstm, m.encoder_attn, m.decoder_lstm,
    m.decoder_attn, m.decoder_fc, m.decoder_fc_final)
Flux.reset!(m::DARNNCell) = foreach(Flux.reset!, (m.encoder_lstm, m.decoder_lstm))
Flux.@functor DARNNCell

"""
    DARNN(inp, encodersize, decodersize, poollength, orig_idx)

Create a DA-RNN layer based on the architecture described in
[Qin et. al.](https://arxiv.org/abs/1704.02971), as implemented for PyTorch
[here](https://github.com/Seanny123/da-rnn/blob/master/modules.py). `inp` specifies the
number of input features. `encodersize` defines the number of LSTM encoder layers, and
`decodersize` defines the number of LSTM decoder layers. `poolsize` gives the
length of the window for the pooled input data, and `orig_idx` defines the array
index where the original time series is stored in the input data,

Data is expected as array with dimensions `features x poolsize x 1 x data`, i.e.
for 1000 data points containing 31 features that have been windowed over 6
timesteps, `DARNN` expects an input size of `(31, 6, 1, 1000)`.

Takes the keyword arguments `init` and `bias` for the initialization of the
weight vector and bias of the linear layers.
"""
function DARNN(inp::Integer, encodersize::Integer, decodersize::Integer, poollength::Integer, orig_idx::Integer;
	init=Flux.glorot_uniform, bias=true)

	# Encoder part
	encoder_lstm = Seq(HiddenRecur(FALSTMCell(inp, encodersize)))
	encoder_attn = Chain(Dense(2 * encodersize + poollength, poollength, init=init, bias=bias),
	                    a -> tanh.(a),
	                    Dense(poollength, 1, init=init, bias=bias))
	# Decoder part
	decoder_lstm = Seq(HiddenRecur(FALSTMCell(1, decodersize)))
	decoder_attn = Chain(Dense(2 * decodersize + encodersize, encodersize, init=init, bias=bias),
	                    a -> tanh.(a),
	                    Dense(encodersize, 1, init=init, bias=bias))
	decoder_fc = Dense(encodersize + 1, 1)
	decoder_fc_final = Dense(decodersize + encodersize, 1, init=init, bias=bias)

	return DARNNCell(encoder_lstm, encoder_attn, decoder_lstm, decoder_attn, decoder_fc,
	 		  decoder_fc_final, encodersize, decodersize, orig_idx, poollength)
end

# model output
function (m::DARNNCell)(x)
	# Initialize after reset
	size(m.encoder_lstm.state, 1) == 1 && darnn_init(m, x)
	input_data = dropdims(x; dims=3)
	input_encoded = darnn_encoder(m, input_data)
	context = darnn_decoder(m, input_encoded, input_data)
	return m.decoder_fc_final(cat(m.decoder_lstm.state[1], context, dims=1))
end

function _encoder(m::DARNNCell, input_data, slice)
	hidden = m.encoder_lstm.state[1]
	cell = m.encoder_lstm.state[2]
	repeatmat = darnn_getones(input_data)

	x = cat(hidden .* repeatmat, cell .* repeatmat,  # CUDA compatible repeat(hidden, inner=(1,1,size(input_data,1)) etc.
			permutedims(input_data, [2,3,1]), dims=1)  |>
		a -> m.encoder_attn(reshape(a, (:, size(input_data, 1) * size(input_data, 3))))
	attn_weights = Flux.softmax(reshape(x, (size(input_data, 1), size(input_data, 3))))
	weighted_input = attn_weights .* slice
	_ = m.encoder_lstm(weighted_input)
	return m.encoder_lstm.state[1]
end

function darnn_encoder(m::DARNNCell, input_data::Array)
	sl = Slices(input_data, True(), False(), True())
	fun = s -> _encoder(m, input_data, s)
	return Align(map(fun, sl), True(), False(), True())
end

function darnn_encoder(m::DARNNCell, input_data::Flux.CUDA.CuArray)
	input_encoded = Flux.Zygote.Buffer(input_data, m.encodersize, m.poollength,
	 				size(input_data, 3))
	@inbounds for t in 1:m.poollength
	      input_encoded[:,t,:] = Flux.unsqueeze(_encoder(m, input_data, input_data[:,t,:]), 2)
	end
	return copy(input_encoded)
end

function darnn_decoder(m::DARNNCell, input_encoded, input_data)
	context = Flux.zeros32(m.encodersize, size(input_data, 3))
	@inbounds for t in 1:m.poollength
	      hidden = m.decoder_lstm.state[1]
		  cell = m.decoder_lstm.state[2]
		  repeatmat = darnn_getones(input_data, m.poollength)
	      x = cat(permutedims(hidden .* repeatmat, [1,3,2]),
	              permutedims(cell .* repeatmat, [1,3,2]),
	              input_encoded, dims=1) |>
	          a -> m.decoder_attn(reshape(a, (2 * m.decodersize + m.encodersize, :))) |>
	          a -> Flux.softmax(reshape(a, (m.poollength, :)))
	      context = dropdims(NNlib.batched_mul(input_encoded, Flux.unsqueeze(x, 2)), dims=2)
	      ỹ = m.decoder_fc(cat(context, input_data[m.orig_idx,t,:]', dims=1))
	      _ = m.decoder_lstm(ỹ)
	end
	return context
end

# pretty printing
function Base.show(io::IO, l::DARNNCell)
	print(io, "DARNN(", size(l.encoder_lstm.chain.cell.Wi, 2))
	print(io, ", ", l.encodersize)
	print(io, ", ", Int(size(l.decoder_lstm.chain.cell.Wi, 1) / 4))
	print(io, ", ", l.poollength)
	print(io, ", ", l.orig_idx)
	print(io, ")")
end

# Helper functions to replace `repeat` with CUDA friendly version
darnn_getones(x::Array, size) = Flux.ones32(1, 1, size)
darnn_getones(x::Flux.CUDA.CuArray, size) = Flux.CUDA.ones(1, 1, size)
darnn_getones(x) = darnn_getones(x, size(x, 1))
Flux.Zygote.@nograd darnn_getones

max_stars_repo_name: sdobber/FluxArchitectures

max_issues_repo_name: sdobber/FluxArchitectures

max_forks_repo_name: sdobber/FluxArchitectures


text:
#****************************************************************************
#                     Molecular Dynamics Potentials (MDP)
#                            CESMIX-MIT Project
#
#  Contributing authors: Ngoc-Cuong Nguyen (cuongng@mit.edu, exapde@gmail.com)
#****************************************************************************

function readJSON(filename, species)

    config = initializeconfig(0);

    mystr = read(filename, String);
    i1 = findfirst("{", mystr);
    mystr = mystr[i1[1]:end];
    dict2 = JSON.parse(mystr);
    Dataset = dict2["Dataset"]["Data"];

    nframe = 0;
    for j = 1:length(Dataset)
        Data = Dataset[j];

        nframe = nframe + 1;

        if haskey(Data, "NumAtoms")
            config.natom = hcat(config.natom, reshape([Data["NumAtoms"]],1,1))
            natom = Data["NumAtoms"];
        end
        if haskey(Data, "Lattice")
            d = hcat(Data["Lattice"]...);
            d = reshape(d, length(d), 1);
            if nframe == 1
                config.lattice = d;
            else
                config.lattice = hcat(config.lattice, d);
            end
            config.ncl = size(config.lattice,1);
        end
        if haskey(Data, "Stress")
            d = hcat(Data["Stress"]...);
            d = reshape(d, length(d), 1);
            if nframe == 1
                config.stress = d;
            else
                config.stress = hcat(config.stress, d);
            end
            config.ncs = size(config.stress,1);
        end
        if haskey(Data, "Energy")
            d = hcat(Data["Energy"]...);
            d = reshape(d, length(d), 1);
            if nframe == 1
                config.e = d;
            else
                config.e = hcat(config.e, d);
            end
            config.nce = size(config.e,1);
        end
        if haskey(Data, "Positions")
            d = hcat(Data["Positions"]...);
            if nframe == 1
                config.x = d;
            else
                config.x = hcat(config.x, d);
            end
            config.ncx = size(config.x,1);
            config.dim = size(config.x,1);
        end
        if haskey(Data, "Forces")
            d = hcat(Data["Forces"]...);
            if nframe == 1
                config.f = d;
            else
                config.f = hcat(config.f, d);
            end
            config.ncf = size(config.f,1);
        end
        if haskey(Data, "Velocities")
            d = hcat(Data["Velocities"]...);
            if nframe == 1
                config.v = d;
            else
                config.v = hcat(config.v, d);
            end
            config.ncf = size(config.v,1);
        end
        if haskey(Data, "Charges")
            d = hcat(Data["Charges"]...);
            d = reshape(d, 1, length(d));
            if nframe == 1
                config.q = d;
            else
                config.q = hcat(config.q, d);
            end
            config.ncq = size(config.q,1);
        end
        if haskey(Data, "Masses")
            d = hcat(Data["Masses"]...);
            d = reshape(d, 1, length(d));
            if nframe == 1
                config.mass = d;
            else
                config.mass = hcat(config.mass, d);
            end
            config.ncm = size(config.mass,1);
        end
        if haskey(Data, "Groups")
            d = hcat(Data["Groups"]...);
            d = reshape(d, 1, length(d));
            if nframe == 1
                config.group = d;
            else
                config.group = hcat(config.group, d);
            end
            config.ncg = size(config.group,1);
        end
        if haskey(Data, "Moves")
            d = hcat(Data["Moves"]...);
            d = reshape(d, 1, length(d));
            if nframe == 1
                config.move = d;
            else
                config.move = hcat(config.move, d);
            end
            config.nco = size(config.move,1);
        end
        if haskey(Data, "Tags")
            d = hcat(Data["Tags"]...);
            d = reshape(d, 1, length(d));
            if nframe == 1
                config.tags = d;
            else
                config.tags = hcat(config.tags, d);
            end
            config.nct = size(config.tags,1);
        end
        if haskey(Data, "Z")
            d = hcat(Data["Z"]...);
            d = reshape(d, 1, length(d));
            if nframe == 1
                config.Z = d;
            else
                config.Z = hcat(config.Z, d);
            end
            config.ncz = size(config.Z,1);
        end
        if haskey(Data, "AtomTypes")
            d = hcat(Data["AtomTypes"]...);
            d = reshape(d, 1, length(d));
            b = ones(1, length(d));
            for k=1:length(species)
                ind = findall(d[:] .== species[k]);
                b[ind] .= k;
            end
            if nframe == 1
                config.t = b;
            else
                config.t = hcat(config.t, b);
            end
            config.nct = size(config.t,1);
        end
    end

    config.nconfigs = nframe;
    config.natomall = sum(config.natom);

    return config

end

max_stars_repo_name: cesmix-mit/MDP

max_issues_repo_name: cesmix-mit/MDP

max_forks_repo_name: cesmix-mit/MDP