text:
#
# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons
# Licensed under the MIT license. See LICENSE file in the project root for details.
#

module FMIBuild

using FMICore: FMU2, FMU2Component, fmi2ModelDescription, fmi2ValueReference, fmi2Component, fmi2ComponentEnvironment, fmi2Status, fmi2EventInfo
using FMICore: fmi2Type, fmi2TypeModelExchange, fmi2TypeCoSimulation
using FMICore: fmi2True, fmi2False
using FMICore: fmi2ModelDescriptionModelExchange, fmi2ModelDescriptionCoSimulation, fmi2VariableNamingConventionStructured
using FMICore: fmi2CausalityToString, fmi2VariabilityToString, fmi2InitialToString
#using FMIExport: fmi2SaveModelDescription

import PackageCompiler
import Pkg
import ZipFile
using EzXML
import Dates

# exports
export fmi2Save

"""
    fmi2Save(fmu::FMU2,
     fmu_path::String,
     fmu_src_file::Union{Nothing, String}=nothing;
     standalone=true,
     compress=false,
     cleanup=true,
     modelExchange=true,
     coSimulation=false,
     removeLibDependency=true,
     removeNoExportBlocks=true,
     pkg_comp_kwargs...)

Initiates the FMU building process.

The current package is detected, duplicated and extended by the FMI-functions. The resulting package is compiled and a suitable FMI-model-description is deployed. Finally, all files are zipped into a standard-compliant FMU that can be executed in a variety of tools without having Julia installed.

# Arguments
    - `fmu_path` path to the (future) FMU file, must end on `*.fmu`
    - `fmu_src_file` file with the FMU constructor, if `nothing` file is determined automatically and the file from the call to `fmi2Save` is assumed (default=`nothing`)

# Keyword arguments
    - `standalone` if the FMU should be build in standalone-mode, meaning without external dependencies to a Julia-Installation (default=`true`)
    - `compress` if the FMU archive should be compressed to save disk space. On the other hand, this may enlarge loading time (default=`false`)
    - `cleanup` if the unzipped FMU archive should be deleted after creation (default=`true`)
    - `modelExchange` if the FMU should support ME (model exchange) (default=`true`)
    - `coSimulation` *currently not supported* if the FMU should support CS (co simulation) (default=`false`)
    - `removeLibDependency` removes the FMIBuild.jl-dependency, so it will not be part of the resulting FMU (default=`true`)
    - `removeNoExportBlocks` removes the blocks marked with `### FMIBUILD_NO_EXPORT_BEGIN ###` and `### FMIBUILD_NO_EXPORT_END ###` from the `fmu_src_file`, so it will not be part of the resulting FMU (default=`true`)
"""
function fmi2Save(fmu::FMU2, fmu_path::String, fmu_src_file::Union{Nothing, String}=nothing;
    standalone=true,
    compress=false,
    cleanup=true,
    modelExchange=true,
    coSimulation=false,
    removeLibDependency=true,
    removeNoExportBlocks=true,
    resources::Union{Dict{String, String}, Nothing}=nothing,
    pkg_comp_kwargs...)

    # @assert fmi2Check(fmu) == true ["fmiBuild(...): FMU-Pre-Check failed. See messages above for further information."]

    # searching the source file ...
    if fmu_src_file == nothing
        stack = stacktrace(backtrace())
        for i in 1:length(stack)
            if !endswith("$(stack[i].file)", "FMIBuild.jl")
                fmu_src_file = "$(stack[i].file)"
                break
            end
        end
        @assert (fmu_src_file != nothing) ["fmi2Save(...): Cannot automatically determine `fmu_src_file`, please provide manually via keyword-argument."]
        @info "[Build FMU] Automatically determined build file at: `$(fmu_src_file)`."
    end

    @assert standalone == true ["fmiBuild(...): Currently, only `standalone=true` is supported."]

    pkg_dir = "$(@__DIR__)/../template"
    (fmu_name, fmu_ext) = splitext(basename(fmu_path))

    if modelExchange && coSimulation
        #pkg_dir = joinpath(pkg_dir, "ME_CS")
        @assert false ["fmiBuild(...): `modelExchange` and `coSimulation` is under development."]
    elseif modelExchange
        pkg_dir = joinpath(pkg_dir, "ME")

        if fmu.modelDescription.modelExchange == nothing
            fmu.modelDescription.modelExchange = fmi2ModelDescriptionModelExchange()
        end
        fmu.modelDescription.modelExchange.modelIdentifier = fmu_name

    elseif coSimulation
        #pkg_dir = joinpath(pkg_dir, "CS")
        @assert false ["fmiBuild(...): `coSimulation` only is under development."]
    else
        @assert false ["fmiBuild(...): At least one of `modelExchange` or `coSimulation` must be supported."]
    end

    pkg_dir = joinpath(pkg_dir, "FMU2")

    @assert fmu_ext != "fmu" ["fmiBuild(...): `fmu_path` must end with `.fmu`."]

    fmu_dir = dirname(fmu_path)

    ### dirs
    target_dir = joinpath(fmu_dir)
    # save evrything in a temporary directory
    if cleanup
       target_dir = mktempdir(; prefix="fmibuildjl_", cleanup=false)
    end
    mkpath(joinpath(target_dir, fmu_name))
    md_path = joinpath(target_dir, fmu_name, "modelDescription.xml")

    bin_dir = joinpath(target_dir, fmu_name, "binaries")
    libext = ""

    # checking architecture
    juliaArch = Sys.WORD_SIZE
    if juliaArch == 64
        if Sys.iswindows()
            bin_dir = joinpath(bin_dir, "win64")
            libext = "dll"
        elseif Sys.islinux()
            bin_dir = joinpath(bin_dir, "x86_64-linux")
            libext = "so"
        elseif Sys.isapple()
            bin_dir = joinpath(bin_dir, "x86_64-darwin")
            libext = "dylib"
        end
    elseif juliaArch == 32
        if Sys.iswindows()
            bin_dir = joinpath(bin_dir, "win32")
            libext = "dll"
        elseif Sys.islinux()
            # pass
        elseif Sys.isapple()
            # pass
        end
    end

    @assert length(libext) > 0 "fmiBuild(...): Unsupported target platform. Supporting Windows (64-, 32-bit), Linux (64-bit) and MacOS (64-bit). Please open an issue online if you need further architectures."

    mkpath(bin_dir)

    pkg_dir = replace(pkg_dir, "\\"=>"/")
    target_dir = replace(target_dir, "\\"=>"/")

    @info "[Build FMU] Generating package ..."
    source_pkf_dir = dirname(fmu_src_file)
    fmu_src_in_merge_dir = splitpath(fmu_src_file)[end]
    while !isfile(joinpath(source_pkf_dir, "Project.toml")) && length(source_pkf_dir) > 0
        pathcomp = splitpath(source_pkf_dir)
        source_pkf_dir = (length(pathcomp) > 1 ? joinpath(pathcomp[1:end-1]...) : "")
        fmu_src_in_merge_dir = joinpath(pathcomp[end], fmu_src_in_merge_dir)
    end
    @assert length(source_pkf_dir) > 0 ["fmiBuild(...): Connot find a package where this file is stored in. For FMU-Export, this source file needs to be inside of a package."]
    merge_dir = joinpath(target_dir, "merged_" * fmu_name)
    cp(source_pkf_dir, merge_dir; force=true)
    chmod(target_dir, 0o777; recursive=true)
    @info "[Build FMU] Source package is $(source_pkf_dir), deployed at $(merge_dir)"
    @info "[Build FMU] Relative src file path is $(fmu_src_in_merge_dir)"

    fmu_res = "$(@__DIR__)/../template/ME/FMU2/src/FMU2_content.jl"

    @info "[Build FMU] ... reading FMU template file at $(fmu_res)"
    f = open(fmu_res, "r")
    fmu_code = read(f, String);
    close(f)

    @info "[Build FMU] ... reading old FMU source file at $(joinpath(merge_dir, fmu_src_in_merge_dir))"
    f = open(joinpath(merge_dir, fmu_src_in_merge_dir), "r")
    cdata = read(f, String);
    close(f)

    # ToDo: This will fail badly on `fmi2Save(a, (b, c), d)` -> use Julia code parser

    if removeNoExportBlocks
        @info "[Build FMU] Removing `FMIBUILD_NO_EXPORT_*` blocks ..."
        cdata = replace(cdata, r"### FMIBUILD_NO_EXPORT_BEGIN ###(.|\n|\r)*### FMIBUILD_NO_EXPORT_END ###" => "")
        @info "[Build FMU] ... removed `FMIBUILD_NO_EXPORT_*` blocks."
    end

    # projectToml = joinpath(merge_dir, "Project.toml")
    # s = stat(projectToml)
    # println(Int(s.mode))
    # println(s.uid)
    # println(s.gid)
    # chmod(projectToml, 0o777)
    # println(Int(s.mode))
    # println(s.uid)
    # println(s.gid)

    @info "[Build FMU] Adding/removing dependencies ..."
    currentEnv = Base.active_project()
    preCompState = 1
    try
        preCompState = ENV["JULIA_PKG_PRECOMPILE_AUTO"]
    catch e
        preCompState = 1
    end
    ENV["JULIA_PKG_PRECOMPILE_AUTO"] = 0

    Pkg.activate(merge_dir)
    Pkg.add("FMICore")
    #Pkg.add(name="FMICore", version="0.7.1")
    @info "[Build FMU]    > Added FMICore"
    if removeLibDependency
        cdata = replace(cdata, r"(using|import) FMIBuild" => "")
        try
            Pkg.rm("FMIBuild")
            @info "[Build FMU]    > Removed FMIBuild"
        catch e
            @info "[Build FMU]    > Not used FMIBuild"
        end
    end

    Pkg.activate(currentEnv)
    #Pkg.resolve()
    ENV["JULIA_PKG_PRECOMPILE_AUTO"]=preCompState
    @info "[Build FMU] ... adding/removing dependencies done."

    @info "[Build FMU] ... generating new FMU source file at $(joinpath(merge_dir, fmu_src_in_merge_dir))"

    f = open(joinpath(merge_dir, fmu_src_in_merge_dir), "w")
    assertline = "" # "\n@assert FMIBUILD_FMU != nothing \"`FMIBUILD_FMU = nothing`, did you forget to mark the FMU instance to export with `FMIBUILD_FMU = myFMUInstance`?\"\n"

    cdata_start = "" # "\nBase.@ccallable function init_FMU()::Cvoid\n"
    cdata_end = "" # "\nnothing\nend\n"
    write(f, "module " * fmu_name * "\n" * fmu_code * cdata_start * cdata * cdata_end * assertline * "\nend");
    close(f)

    @info "[Build FMU] ... generating package done."

    @info "[Build FMU] Compiling FMU ..."
    PackageCompiler.create_library(merge_dir, joinpath(target_dir, "_" * fmu_name);
                                    lib_name=fmu_name,
                                    precompile_execution_file=[joinpath(merge_dir, fmu_src_in_merge_dir)], # "$(@__DIR__)/../template/ME/precompile/FMU2_generate.jl"
                                    precompile_statements_file=["$(@__DIR__)/../template/ME/precompile/FMU2_additional.jl"],
                                    incremental=false,
                                    filter_stdlibs=false,
                                    julia_init_c_file = "$(@__DIR__)/../template/ME/header/FMU2_init.c",
                                    header_files = ["$(@__DIR__)/../template/ME/header/FMU2_init.h"],
                                    force=true,
                                    pkg_comp_kwargs...)
    @info "[Build FMU] ... compiling FMU done."

    cp(joinpath(target_dir, "_" * fmu_name, "bin"), joinpath(bin_dir); force=true)
    cp(joinpath(target_dir, "_" * fmu_name, "share"), joinpath(bin_dir, "..", "share"); force=true)
    cp(joinpath(target_dir, "_" * fmu_name, "include"), joinpath(bin_dir, "..", "include"); force=true)

    if resources != nothing
        @info "[Build FMU] Adding resource files ..."
        mkdir(joinpath(bin_dir, "..", "..", "resources"))
        for (key, val) in resources
            cp(key, joinpath(bin_dir, "..", "..", "resources", val); force=true)
            @info "[Build FMU] \t $val"
        end
        @info "[Build FMU] ... adding resoruce files done."
    end

    @info "[Build FMU] Patching libjulia.$(libext) @ `$(bin_dir)`..."
    patchJuliaLib(joinpath(bin_dir, "libjulia.$(libext)"))
    @info "[Build FMU] ... patching libjulia.$(libext) done."

    @info "[Build FMU] Building model description ..."
    #buildModelDescription(md_path, fmu_name, fmu_src_file)
    fmi2SaveModelDescription(fmu.modelDescription, md_path)
    @info "[Build FMU] ... building model description done."

    @info "[Build FMU] Zipping FMU ..."
    # parse and zip directories
    zipfile = joinpath(target_dir, fmu_name * ".zip")
    zdir = ZipFile.Writer(zipfile)
    for (root, dirs, files) in walkdir(joinpath(target_dir, fmu_name))
        for file in files
            filepath = joinpath(root, file)
            f = open(filepath, "r")
            content = read(f, String)
            close(f)

            zippath = subtractPath(filepath, joinpath(target_dir, fmu_name) * "/")
            println("\t$(zippath)")
            zf = ZipFile.addfile(zdir, zippath; method=(compress ? ZipFile.Deflate : ZipFile.Store));
            write(zf, content)
        end
    end
    close(zdir)
    # Rename ZIP-file to FMU-file
    cp(joinpath(target_dir, fmu_name * ".zip"), joinpath(fmu_dir, fmu_name * ".fmu"); force=true)
    @info "[Build FMU] ... zipping FMU done."

    if cleanup
        @info "[Build FMU] Clean up ..."
        # Clean-up is done by saving in a temporary directory (which may be deleted by the OS)
        @info "[Build FMU] ... clean up done."
    end

    return true
end

# removes all leading elements of pathA that fit pathB
function subtractPath(pathA::String, pathB::String)
    pathA = replace(pathA, "\\"=>"/")
    pathB = replace(pathB, "\\"=>"/")

    if length(pathA) <= 0 || length(pathB) <= 0
        return pathA
    end

    while (length(pathA) > 0 && length(pathB) > 0) && (pathA[1] == pathB[1])
        pathA = pathA[2:end]
        pathB = pathB[2:end]
    end

    return pathA
end

# Thank you @staticfloat
# src: https://github.com/JuliaLang/PackageCompiler.jl/issues/658
# patches the compiled DLL, so the lib does not have "bad" relative paths like "../bin/", where "bin" is the directory of the DLL itself
function patchJuliaLib(libjulia_path)

    if !isfile(libjulia_path)
        error("Unable to open libjulia.* at $(libjulia_path)")
    end

    open(libjulia_path, read=true, write=true) do io
        # Search for `../bin/` string:
        needle = "../bin/"
        readuntil(io, needle)
        skip(io, -length(needle))
        libpath_offset = position(io)

        libpath = split(String(readuntil(io, UInt8(0))), ":")
        @info("Found embedded libpath", libpath, libpath_offset)

        # Get rid of `../bin/` prefix:
        libpath = map(libpath) do l
            if startswith(l, "../bin/")
                return l[8:end]
            end
        end

        @info("Filtered libpath", libpath)

        # Write filtered libpath out to the file, terminate with NULL.
        seek(io, libpath_offset)
        write(io, join(libpath, ":"))
        write(io, UInt8(0))
    end
end

function fmi2SaveModelDescription(md::fmi2ModelDescription, file_path::String)
    doc = XMLDocument()

    doc_root = ElementNode("fmiModelDescription")
    setroot!(doc, doc_root)
    link!(doc_root, AttributeNode("xmlns:xsi", "http://www.w3.org/2001/XMLSchema-instance"))

    # mandatory
    link!(doc_root, AttributeNode("fmiVersion", "2.0"))
    link!(doc_root, AttributeNode("modelName", md.modelName))
    link!(doc_root, AttributeNode("guid", "$(md.guid)"))

    # optional
    link!(doc_root, AttributeNode("generationTool", "FMIExport.jl (0.1.0) / FMIBuild.jl (0.1.0) by Tobias Thummerer, Lars Mikelsons"))
    if md.generationDateAndTime != nothing
        dateTimeString = ""
        if typeof(md.generationDateAndTime) == Dates.DateTime
            dateTimeString = Dates.format(md.generationDateAndTime, "yyyy-mm-dd") * "T" * Dates.format(md.generationDateAndTime, "HH:MM:SS") * "Z"
        elseif typeof(md.generationDateAndTime) == String
            dateTimeString = md.generationDateAndTime
        else
            @warn "fmi2SaveModelDescription(...): Unkown data type for field `generationDateAndTime`. Supported is `DateTime` and `String`, but given `$(md.generationDateAndTime)` (typeof `$(typeof(md.generationDateAndTime))`)."
        end
        link!(doc_root, AttributeNode("generationDateAndTime", dateTimeString))
    end
    if md.variableNamingConvention != nothing
        link!(doc_root, AttributeNode("variableNamingConvention", (md.variableNamingConvention == fmi2VariableNamingConventionStructured ? "structured" : "flat")))
    end
    if md.numberOfEventIndicators != nothing
        link!(doc_root, AttributeNode("numberOfEventIndicators", "$(md.numberOfEventIndicators)"))
    end

    if md.modelExchange != nothing
        me = ElementNode("ModelExchange")
        link!(doc_root, me)

        # mandatory
        link!(me, AttributeNode("modelIdentifier", md.modelExchange.modelIdentifier))

        # optional
        if md.modelExchange.canGetAndSetFMUstate != nothing
            link!(me, AttributeNode("canGetAndSetFMUstate", (md.modelExchange.canGetAndSetFMUstate ? "true" : "false")))
        end
        if md.modelExchange.canSerializeFMUstate != nothing
            link!(me, AttributeNode("canSerializeFMUstate", (md.modelExchange.canSerializeFMUstate ? "true" : "false")))
        end
        if md.modelExchange.providesDirectionalDerivative != nothing
            link!(me, AttributeNode("providesDirectionalDerivative", (md.modelExchange.providesDirectionalDerivative ? "true" : "false")))
        end
    end

    mv = ElementNode("ModelVariables")
    link!(doc_root, mv)
    for i in 1:length(md.modelVariables)
        link!(mv, CommentNode("Index=$(i)"))

        sv = md.modelVariables[i]
        sv_node = ElementNode("ScalarVariable")
        link!(mv, sv_node)

        # mandatory
        link!(sv_node, AttributeNode("name", sv.name))
        link!(sv_node, AttributeNode("valueReference", "$(sv.valueReference)"))

        # optional
        if sv.description != nothing
            link!(sv_node, AttributeNode("description", sv.description))
        end
        if sv.causality != nothing
            link!(sv_node, AttributeNode("causality", fmi2CausalityToString(sv.causality)))
        end
        if sv.variability != nothing
            link!(sv_node, AttributeNode("variability", fmi2VariabilityToString(sv.variability)))
        end
        if sv.initial != nothing
            link!(sv_node, AttributeNode("initial", fmi2InitialToString(sv.initial)))
        end
        if sv.canHandleMultipleSetPerTimeInstant != nothing
            link!(sv_node, AttributeNode("canHandleMultipleSetPerTimeInstant", (sv.canHandleMultipleSetPerTimeInstant ? "true" : "false")))
        end

        # Real
        if sv._Real != nothing
            r_node = ElementNode("Real")

            if sv._Real.start != nothing
                link!(r_node, AttributeNode("start", "$sv._Real.start)"))
            end
            if sv._Real.derivative != nothing
                link!(r_node, AttributeNode("derivative", "$(sv._Real.derivative)"))
            end

            # ToDo: Implement remaining attributes

            link!(sv_node, r_node)
        end

        # Integer
        if sv._Integer != nothing
            i_node = ElementNode("Integer")

            if sv._Integer.start != nothing
                link!(i_node, AttributeNode("start", "$sv._Integer.start)"))
            end

            # ToDo: Implement remaining attributes

            link!(sv_node, i_node)
        end

        # Boolean
        if sv._Boolean != nothing
            b_node = ElementNode("Boolean")

            if sv._Boolean.start != nothing
                link!(b_node, AttributeNode("start", "$sv._Boolean.start)"))
            end

            # ToDo: Implement remaining attributes

            link!(sv_node, b_node)
        end

        # String
        if sv._String != nothing
            s_node = ElementNode("String")

            if sv._String.start != nothing
                link!(s_node, AttributeNode("start", "$sv._String.start)"))
            end

            # ToDo: Implement remaining attributes

            link!(sv_node, s_node)
        end

        # _Enumeration
        if sv._Enumeration != nothing
            e_node = ElementNode("Enumeration")

            # if sv._Enumeration.start != nothing
            #     link!(e_node, AttributeNode("start", "$sv._Enumeration.start)"))
            # end

            # ToDo: Implement remaining attributes

            link!(sv_node, e_node)
        end
    end

    ms = ElementNode("ModelStructure")
    link!(doc_root, ms)

    if md.modelStructure.outputs != nothing
        outs = ElementNode("Outputs")
        link!(ms, outs)

        for i in 1:length(md.modelStructure.outputs)
            uk = md.modelStructure.outputs[i]

            uk_node = ElementNode("Unknown")
            link!(outs, uk_node)
            link!(uk_node, AttributeNode("index", "$(uk.index)"))
            if uk.dependencies != nothing
                link!(uk_node, AttributeNode("dependencies", ""))
                # ToDo: Actually save them!
            end
            if uk.dependenciesKind != nothing
                link!(uk_node, AttributeNode("dependenciesKind", ""))
                # ToDo: Actually save them!
            end
        end
    end

    if md.modelStructure.derivatives != nothing
        ders = ElementNode("Derivatives")
        link!(ms, ders)

        for i in 1:length(md.modelStructure.derivatives)
            uk = md.modelStructure.derivatives[i]

            uk_node = ElementNode("Unknown")
            link!(ders, uk_node)
            link!(uk_node, AttributeNode("index", "$(uk.index)"))
            if uk.dependencies != nothing
                link!(uk_node, AttributeNode("dependencies", ""))
                # ToDo: Actually save them!
            end
            if uk.dependenciesKind != nothing
                link!(uk_node, AttributeNode("dependenciesKind", ""))
                # ToDo: Actually save them!
            end
        end
    end

    if md.modelStructure.initialUnknowns != nothing
        inis = ElementNode("InitialUnknowns")
        link!(ms, inis)

        for i in 1:length(md.modelStructure.initialUnknowns)
            uk = md.modelStructure.initialUnknowns[i]

            uk_node = ElementNode("Unknown")
            link!(inis, uk_node)
            link!(uk_node, AttributeNode("index", "$(uk.index)"))
            if uk.dependencies != nothing
                link!(uk_node, AttributeNode("dependencies", ""))
                # ToDo: Actually save them!
            end
            if uk.dependenciesKind != nothing
                link!(uk_node, AttributeNode("dependenciesKind", ""))
                # ToDo: Actually save them!
            end
        end
    end

    # save modelDescription (with linebreaks)
    f = open(file_path, "w")
    prettyprint(f, doc)
    close(f)
end

end # module


text:
#=
Final parsing of preparsing structures

Numbers and variables will finaly
=#
include("ofMeshDictPreParser.jl")

abstract type OfEntity end
abstract type AbstractOfVar <: OfEntity end
abstract type AbstractOfVal <: OfEntity end

struct OfDictionary <: OfEntity
    name::String
    entries::Vector{OfEntity}
end

# array list of same entities with length
struct OfArray <: OfEntity
    length::Int64
    entries::Vector{OfEntity}
end

struct OfList<: OfEntity
    length::Int64
    entries::Vector{OfEntity}
end

struct OfVariable{T} <: AbstractOfVar
    name::String
    value::T
end

struct OfValue{T} <: AbstractOfVal
    value::T
end

"""
    dfs_parse(parse_entities::Vector{PreParsedEntity}, tokens::Vector{Token})

parses pre parsed vector structure to finally parsed file structure
"""
function dfs_parse(parse_entities::Vector{PreParsedEntity}, tokens::Vector{Token})

    of_entities = Vector{OfEntity}(undef, 0)

    for parse_elem in parse_entities
        if typeof(parse_elem) == PreParsedVariable
            push!(of_entities, parse_of_variable(parse_elem, tokens))
        elseif typeof(parse_elem) == PreParsedFNum
            push!(of_entities, OfValue(parse(Float64,parse_elem.string)))
        elseif typeof(parse_elem) == PreParsedINum
            push!(of_entities, OfValue(parse(Int64,parse_elem.string)))
        elseif typeof(parse_elem) == PreParsedDictonary
            mypush!(of_entities, OfDictionary(parse_elem.name,
                        dfs_parse(parse_elem.entries, tokens)))
        elseif typeof(parse_elem) == PreParsedArray
            mypush!(of_entities, OfArray(parse_elem.length,
                        dfs_parse(parse_elem.entries, tokens)))
        elseif typeof(parse_elem) == PreParsedList
            parsed_list = dfs_parse(parse_elem.entries, tokens)
            mypush!(of_entities, OfList(length(parsed_list), parsed_list))
        end
    end

    return of_entities
end

function mypush!(of_entities::Vector{OfEntity}, OfEntity::T) where T <: OfEntity
    push!(of_entities, OfEntity)
end

function mypush!(of_entities::Vector{OfEntity},
                    entities::Vector{T}) where T <: OfEntity
    for e in entities
        push!(of_entities, e)
    end
end

"""
    parse_of_variable(pe::PreParsedEntity, tokens::Vector{Token})

Parsing of a PreParsedEntity to final type OfVariable, which completes parsing.
"""
# TODO: only works for numbers and strings by now
function parse_of_variable(pe::PreParsedEntity, tokens::Vector{Token})::OfVariable

    name = ""
    value = []

    if tokens[pe.id_token_from].type_id == TokenTypeID_Dict[:word]
        name = tokens[pe.id_token_from].string

        if ( (pe.id_token_to - pe.id_token_from) == 2 )
            if  (tokens[pe.id_token_from+1].type_id ==
                 TokenTypeID_Dict[:int_number])

                value = parse(Int64, tokens[pe.id_token_from+1].string)

            elseif (tokens[pe.id_token_from+1].type_id ==
                    TokenTypeID_Dict[:float_number])

                value = parse(Float64, tokens[pe.id_token_from+1].string)

            elseif (tokens[pe.id_token_from+1].type_id
                    == TokenTypeID_Dict[:word])

                value = tokens[pe.id_token_from+1].string

            elseif (tokens[pe.id_token_from+1].type_id
                == TokenTypeID_Dict[:string])

                value = tokens[pe.id_token_from+1].string

            else
                value = nothing
            end
        elseif ( ((pe.id_token_to - pe.id_token_from) == 3) &&
                 (tokens[pe.id_token_from+1].type_id ==
                        TokenTypeID_Dict[:minus]) )

            if  (tokens[pe.id_token_from+2].type_id ==
                TokenTypeID_Dict[:int_number])

                value = parse(Int64, tokens[pe.id_token_from+1].string *
                                        tokens[pe.id_token_from+2].string)

            elseif (tokens[pe.id_token_from+2].type_id ==
                    TokenTypeID_Dict[:float_number])

                 value = parse(Float64, tokens[pe.id_token_from+1].string *
                                        tokens[pe.id_token_from+2].string)
            else
                value =  nothing
            end

        else
            value = 0
        end
    else
        error("First Token of Variable should be a word!")
    end

    return OfVariable(name, value)
end

"""
    ofprint(ofvar::OfEntity, depth::Int64=0)

dfs print in case for single OfEntity
"""
function ofprint(ofvar::OfEntity, depth::Int64=0)
    ofprint([ofvar], depth)
end

"""
    ofprint(ofvars::Vector{OfEntity}, depth::Int64=0)

dfs print of Vector{OfEntity} to get a nicer view of finally parsed entities
"""
function ofprint(ofvars::Vector{OfEntity}, depth::Int64=0)
    for ofvar in ofvars
        if typeof(ofvar) <: AbstractOfVar
            println("\t"^depth,ofvar.name, ": ", ofvar.value)

        elseif typeof(ofvar) <: AbstractOfVal
            print(ofvar.value, " ")

        elseif typeof(ofvar) == OfDictionary
            println()
            println("\t"^depth,ofvar.name, " {")
            depth += 1
            ofprint(ofvar.entries, depth)
            depth -= 1
            println("\t"^depth,"}")
        elseif typeof(ofvar) == OfArray
            print("\t"^depth,ofvar.length, " (")
            depth += 1
            ofprint(ofvar.entries,depth)
            print(")\n")
            depth -= 1
        elseif typeof(ofvar) == OfList
            print("\t"^depth,"(")
            ofprint(ofvar.entries, depth)
            print(")\n")
        end
    end

end

"""
    get_dictionary(ofvars::Vector{OfEntity}, dict_name::string)

Returns Dictionary in first layer of ofvars with .name == dict_name
"""
function get_dictionary(ofvars::Vector{OfEntity}, dict_name::String)

    for f in ofvars
        if typeof(f) == OfDictionary && f.name == dict_name
            return f
        end
    end

    return nothing
end

"""
    get_variable(ofvars::Vector{OfEntity}, var_name::string)

Returns Variable in first layer of ofvars with .name == var_name
"""
function get_variable(ofvars::Vector{OfEntity}, var_name::String)

    for f in ofvars
        if typeof(f) <: AbstractOfVar && f.name == var_name
            return f
        end
    end

    return nothing
end

"""
    get_variable_value(ofvars::Vector{OfEntity}, var_name::string)

Returns Variable value in first layer of ofvars with .name == var_name
"""
function get_variable_value(ofvars::Vector{OfEntity}, var_name::String)

    var = get_variable(ofvars::Vector{OfEntity}, var_name::String)

    if !isnothing(var)
        return var.value
    else
        error("Variable ", var_name, " cannot be found")
    end

end

ext: jl

max_stars_repo_name: c-schubert/linearMeshIO

max_issues_repo_name: c-schubert/linearMeshIO

max_forks_repo_name: c-schubert/linearMeshIO


text:
using Documenter, LiteQTL

const src = "https://github.com/senresearch/LiteQTL.jl"
const dst = "https://senresearch.github.io/LiteQTL.jl/stable"

function main()
    ci = get(ENV, "CI", "") == "true"
    @info "Generating Documenter.jl site"
    DocMeta.setdocmeta!(LiteQTL, :DocTestSetup, :(using LiteQTL); recursive=true)

    makedocs(
        # Must have the following args:
        sitename = "LiteQTL.jl",
        authors = "Chelsea Trotter, Saunak Sen",
        repo = "$src/blob/{commit}{path}#{line}",
        format = Documenter.HTML(
            # Use clean URLs on CI
            prettyurls = ci,
            canonical = dst,
            analytics = "UA-154489943-2",
        ),
        modules = [LiteQTL],
        pages = Any[
            "Home" => "index.md",
            "Library" => "library.md"
        ],

        # Additional, not must-haves.
        doctest = ("doctest=only" in ARGS) ? :only : true,
        strict = !("strict=false" in ARGS),
    )

    if ci
        @info "Deploying to GitHub"
        deploydocs(
            # Must have the following args:
            repo = "github.com/senresearch/LiteQTL.jl.git",
            devurl = "dev",
            versions = ["stable" => "v^", "v#.#"],
            push_preview = true
        )
    end

end

isinteractive() || main()

# makedocs(;
#     modules=[LiteQTL],
#     format=Documenter.HTML(),
#     pages=[
#         "Home" => "index.md",
#     ],
#     repo="/blob/{commit}{path}#L{line}",
#     sitename="LiteQTL.jl",
#     authors="Chelsea Trotter, Sen Research",
#     assets=String[],
# )

# deploydocs(;
#     repo="github.com/senresearch/LiteQTL.jl.git",
#     devurl = "dev",
#     versions = ["stable" => "v^", "v#.#"],
# )


text:
@testset "Booleans" begin
    @test missing | true === true
    @test ismissing(missing | false)
    @test ismissing(missing | missing)
    @test true | missing === true
    @test ismissing(false | missing)

    @test ismissing(missing & true)
    @test missing & false === false
    @test ismissing(missing & missing)
    @test ismissing(true & missing)
    @test false & missing === false

    @test ismissing(missing ⊻ true)
    @test ismissing(missing ⊻ false)
    @test ismissing(missing ⊻ missing)
    @test ismissing(true ⊻ missing)
    @test ismissing(false ⊻ missing)

    @test any((@data [1, 2, missing]) .== 1) === true
    @test any((@data [missing, 1, 2]) .== 1) === true
    @test ismissing(any((@data [1, 2, missing]) .== 3))
    @test any((@data [1, 2, 3] ).== 4) === false

    @test ismissing(all((@data [1, 1, missing]) .== 1))
    @test ismissing(all((@data [missing, 1, 1]) .== 1))
    @test all((@data [1, 1, 1]) .== 1) === true
    @test all((@data [1, 2, 1]) .== 1) === false
end


text:
function demo_pulsed_solution()
  cur_T = 13.065

  cur_reactor = demo_pulsed_deck(cur_T)

  cur_reactor.n_bar = 0.7983
  cur_reactor.I_P = 19.6

  cur_reactor.R_0 = 9.072
  cur_reactor.B_0 = 5.667

  cur_reactor.P_F = 2037

  cur_reactor.f_BS = 0.348
  cur_reactor.f_CD = 0.096
  cur_reactor.f_IN = 0.557

  cur_reactor.beta_N = cur_reactor.max_beta_N
  cur_reactor.q_95 = 3.247
  cur_reactor.P_W = 1.05

  cur_reactor.volume = 2502

  cur_reactor
end


text:
module NeRCA

using LinearAlgebra
using Statistics
using Printf
using Sockets
using DrWatson
using Parameters
using DataStructures
using StaticArrays
using HDF5
using PGFPlotsX
using RecipesBase
using PlotThemes
using JuMP
using Ipopt
using IntervalArithmetic

import Base: +, -, *

export
    Position, Direction,
    ChannelID, DOMID, ToT, Floor, DU, HitTime,
    MCEventInfo, MCTrack, Hit, CalibratedHit, McHit, TimesliceHit,
    RecoTrack, NoRecoTrack,
    DAQEvent, read_io,
    is_mxshower, is_3dmuon, is_3dshower,
    read_indices, read_hits, read_tracks, read_mctracks,
    read_calibration, read_event_info,
    Calibration, calibrate,
    duhits, domhits, nfoldhits, triggered,
    svdfit, matrix, rows,
    CHClient, CHTag, CHPrefix, CHMessage, subscribe,
    ztplot,
    savefigs,
    most_frequent,
    @ip_str

include("constants.jl")
include("types.jl")
include("tools.jl")
include("io.jl")
include("math.jl")
include("hits.jl")
include("mc.jl")
include("fit.jl")
include("controlhost.jl")
include("plot.jl")

"""
    function calibrate(hits::Vector{T}, calibration::Calibration) where {T<:DAQHit}

Apply geometry and time calibration to given hits.
"""
function calibrate(hits::Vector{T}, calibration::Calibration) where {T<:DAQHit}
    calibrated_hits = Vector{CalibratedHit}()
    for hit in hits
        dom_id = hit.dom_id
        channel_id = hit.channel_id
        tot = hit.tot
        pos = calibration.pos[dom_id][channel_id+1]
        dir = calibration.dir[dom_id][channel_id+1]
        t0 = calibration.t0[dom_id][channel_id+1]
        t = hit.t + t0
        du = calibration.du[dom_id]
        floor = calibration.floor[dom_id]
        triggered = false
        if T === Hit
            triggered = hit.triggered
        end
        if T === DAQTriggeredHit
            triggered = hit.trigger_mask > 0
        end
        c_hit = CalibratedHit(channel_id, dom_id, du, floor, t, tot,
                              pos, dir, t0, triggered, Multiplicity(0,0))
        push!(calibrated_hits, c_hit)
    end
    calibrated_hits
end

"""
    function calibrate(mc_hits::Vector{T}, calibration::Calibration) where {T<:McHit}

Apply geometry and time calibration to given mc_hits.
"""
function calibrate(mc_hits::Vector{T},
                   calibration::Calibration,
                   event_info::Union{Nothing,MCEventInfo,DAQEventInfo} = nothing) where {T<:McHit}
    calibrated_hits = Vector{CalibratedHit}()
    if event_info != nothing
        mctime = make_mc_time_converter(event_info)
    else
        mctime = x->x
    end
    for hit in mc_hits
        omkey = calibration.omkeys[hit.pmt_id]
        dom_id = omkey.dom_id
        channel_id = omkey.channel_id
        tot = hit.a
        pos = calibration.pos[dom_id][channel_id+1]
        dir = calibration.dir[dom_id][channel_id+1]
        t0 = calibration.t0[dom_id][channel_id+1]
        t = mctime(hit.t)
        du = calibration.du[dom_id]
        floor = calibration.floor[dom_id]
        triggered = false
        if T === Hit
            triggered = hit.triggered
        end
        if T === DAQTriggeredHit
            triggered = hit.trigger_mask > 0
        end
        c_hit = CalibratedHit(channel_id, dom_id, du, floor, t, tot,
                              pos, dir, t0, triggered, Multiplicity(0,0))
        push!(calibrated_hits, c_hit)
    end
    calibrated_hits
end

# Utility
rows(x) = (x[i, :] for i in indices(x,1))

# Math
Base.angle(d1::Direction, d2::Direction) = acos(dot(d1/norm(d1), d2/norm(d2)))
Base.angle(a::T, b::T) where {T<:Union{CalibratedHit, PMT, MCTrack}} = Base.angle(a.dir, b.dir)
Base.angle(a::FieldVector{3}, b::Union{CalibratedHit, PMT, MCTrack}) = Base.angle(a, b.dir)
Base.angle(a::Union{CalibratedHit, PMT, MCTrack}, b::FieldVector{3}) = Base.angle(a.dir, b)

function matrix(v::Vector)
    m = length(v)
    n = length(v[1])

    M = zeros(m, n)
    i = 1
    for j ∈ 1:n
        for k ∈ 1:m
            M[i] = v[k][j]
            i += 1
        end
    end
    M
end

function svdfit(M)
    com = mean(M, dims=1)
    subtr = M .- com
    U, S, V = svd(subtr)
    return com[:], V[:, 1]
end

end # module


text:
dir_path = @__DIR__
push!(LOAD_PATH, dir_path * "/sim") # simulation
push!(LOAD_PATH, dir_path * "/vis") # visualization
push!(LOAD_PATH, dir_path * "/adv") # adversary policies
push!(LOAD_PATH, dir_path * "/dis") # discretization
push!(LOAD_PATH, dir_path * "/olm")
push!(LOAD_PATH, dir_path * "/pomdp")
push!(LOAD_PATH, dir_path * "/tap")
using sim
using vis
using adv
using dis
using olm
using pomdp
using tap

using Serialization
using Printf
using DelimitedFiles
using Statistics

include("update_info.jl")

function discretize_adv(world::World)
  pts_per_acc = 3
  pts_per_ste = 3
  min_acc = -15
  max_acc = 5
  min_ste = -1e-1
  max_ste = 1e-1
  ctrl_d = Discretization([pts_per_acc, pts_per_ste],
                          [min_acc, min_ste],
                          [max_acc, max_ste])
  return ctrl_d
end

# Scenario: Racer scenario
function main(use_tape::Bool=false)
  # load the graphical context (OpenGL handle, the graphical window, etc.)
  context = vis.setup()
  vis_scaling = 0.01

  # make the road
  path = make_figure8_path()
  road_width = 20.0
  road = Road(path, road_width,
              vis.make_road(context, path.X, path.Y, road_width))
  road.road.T = vis.scale_mat(vis_scaling, vis_scaling)

  # make the world
  world = World(road, [], vis_scaling)

	len_rd = path.S[end]

  # forward search agent ---------------------------------------------------- #
  x01 = [len_rd - 50.0; 25.0; 0.0]
  agent1 = Agent(1, copy(x01), vis.make_car(context))
	agent1.controller! = olm.controller_pofs!
  agent1.custom = [0.0, 0.0]
  b0 = fill(1 / length(pomdp.DRIVERS), length(pomdp.DRIVERS))
  b = b0
	ctrl_d = discretize_adv(world)
  # ------------------------------------------------------------------------- #

  v_track = 25.5
	p_track = 0.0
  x02 = [len_rd / 2 - 50.0; v_track; 0]
  agent2 = Agent(2, copy(x02), vis.make_car(context, [1.0, 0.0, 0.0]))
	agent2.controller! = pomdp.adv_controller!
	agent2.custom = [v_track, p_track, pomdp.NOTHING, ctrl_d, pomdp.WEAK]

  push!(world.agents, agent1)
  push!(world.agents, agent2)

  # make diagnostics render objects
  info1 = vis.InfoBox(context, 0.75, 0.75, vis_scaling)
  info2 = vis.InfoBox(context, -0.75, -0.25, vis_scaling)

  # main loop for rendering and simulation
  window = true
  h = 1e-2
  dt = 1 / 60

  tf = 4.0
  N = round(Int, tf / dt)
  D = 4
  global T = Array{Float64, 2}(undef, (N, D))
  global Cpofs = Array{Float64, 2}(undef, (N, D))
  global Cpomc = Array{Float64, 2}(undef, (N, D))

  do_vis = false
  for j in 1:D
    t = 0.0
    agent1.x = copy(x01)
    agent2.x = copy(x02)
    b = b0
    agent2.custom[5] = pomdp.WEAK
    for i in 1:N
      if do_vis == true
        tape = tap.Tape(world)
        tape.info_objs[1] = info1
        tape.info_objs[2] = info2
        tap.add_frame!(tape)
      end

      t1 = time_ns() / 1e9
      (u, ret) = olm.plan_pomc(agent1.x, b, agent1, agent2, world,
                               pomdp.reward, ctrl_d, j)
      t2 = time_ns() / 1e9
      Cpomc[i, j] = t2 - t1

      t1 = time_ns() / 1e9
      (u, ret) = olm.plan_pofs(agent1.x, b, agent1, agent2, world,
                               pomdp.reward, ctrl_d, j)
      t2 = time_ns() / 1e9
      Cpofs[i, j] = t2 - t1
      agent1.custom = u

      (o, _) = adv.replan_adv(agent2, world)
      agent2.custom[3] = o
      (c, cv) = predict_collision(agent2.x, agent1.x, world)
      b = pomdp.update_belief(b, o, c)

      ## advance one frame in time
      advance!(sim.default_dynamics!, agent1.x, Pair(agent1, world), 0.0, dt, h)
      advance!(agent2.dynamics!, agent2.x, Pair(agent2, world), 0.0, dt, h)

      ## visualize
      if do_vis == true
        diag = diagnostic(agent1, world, 0.0)
        (x, y, sx, sy, dx, u) = diag
        (_, cv) = predict_collision(agent1.x, agent2.x, world)
        tap.add_to_frame!(tape, agent1.x, diag, cv)

        diag = diagnostic(agent2, world, 0.0)
        (x, y, sx, sy, dx, u) = diag
        (_, cv) = predict_collision(agent2.x, agent1.x, world)
        tap.add_to_frame!(tape, agent2.x, diag, cv)
      end

      T[i, j] = t
      t += dt

      if mod(i, div(N, 10)) == 0
        println("$(j) -> $(t)")
      end

      if do_vis == true
        to_vis = tap.visualize_frame(tape, 1)
        window = vis.visualize(context, to_vis)
      end

    end
    println(agent1.x)
    println("Finished depth = $(j)")
  end

  DATApofs = fill(0.0, (N, 2 * D))
  DATApomc = fill(0.0, (N, 2 * D))
  for i in 1:D
    DATApofs[:, 2 * (i - 1) + 1] = T[:, i]
    DATApofs[:, 2 * (i - 1) + 2] = Cpofs[:, i]

    DATApomc[:, 2 * (i - 1) + 1] = T[:, i]
    DATApomc[:, 2 * (i - 1) + 2] = Cpomc[:, i]
  end
  writedlm("../data/pofs_data.txt", DATApofs)
  writedlm("../data/pomc_data.txt", DATApomc)

  AVGCpofs = map(i -> mean(DATApofs[:, 2 * (i - 1) + 2]), 1:D)
  AVGCpomc = map(i -> mean(DATApomc[:, 2 * (i - 1) + 2]), 1:D)

  writedlm("../data/pofs_avgc.txt", AVGCpofs)
  writedlm("../data/pomc_avgc.txt", AVGCpomc)
end


text:
module Propulsion

using LinearAlgebra
using StaticArrays, ComponentArrays, StructArrays, RecursiveArrayTools
using Unitful
using UnPack

using Flight.Plotting
import Flight.Plotting: plots

using Flight.Airdata
using Flight.Kinematics
using Flight.Dynamics

using Flight.Modeling
import Flight.Modeling: init_x0, init_y0, init_u0, init_d0, f_cont!, f_disc!

using Flight.Components
import Flight.Components: WrenchTrait, AngularMomentumTrait, get_wr_b, get_hr_b

export SimpleProp, Gearbox, ElectricMotor, Battery, CW, CCW
export EThruster, EThrusterU, EThrusterD, EThrusterY

@enum TurnSense begin
    CW = 1
    CCW = -1
end

################ EThruster Component ###################

Base.@kwdef struct SimpleProp
    kF::Float64 = 2e-3
    kM::Float64 = 5e-5
    J::Float64 = 0.5
end

function get_wrench(prop::SimpleProp, ω::Real, air::AirData) #air data just for interface demo
    @unpack kF, kM = prop
    F_ext_Os_s = kF * ω^2 * SVector(1,0,0)
    M_ext_Os_s = -tanh(ω/1.0) * kM * ω^2 * SVector(1,0,0) #choose ω_ref = 1.0
    Wrench(F = F_ext_Os_s, M = M_ext_Os_s)
end

Base.@kwdef struct Gearbox
    n::Float64 = 1.0 #gear ratio
    η::Float64 = 1.0 #efficiency
end

Base.@kwdef struct ElectricMotor #defaults from Hacker Motors Q150-4M-V2
    i₀::Float64 = 6.78
    R::Float64 = 0.004
    kV::Float64 = 125u"rpm*V" |> upreferred |> ustrip #rad/s/V, using the same value for kM
    J::Float64 = 0.003 #kg*m^2 #ballpark figure, assuming a cylinder
    α::TurnSense = CW
end

back_emf(m::ElectricMotor, ω::Real) = Int(m.α) * ω / m.kV
torque(m::ElectricMotor, i::Real, ω::Real) = (Int(m.α) * i - tanh(ω) * m.i₀) / m.kV
R(m::ElectricMotor) = m.R

Base.@kwdef struct Battery
    n_cells::Int64 = 14 #number of cells in series
    V_cell::Float64 = 4.2 #no-load cell voltage at C=Cmax (V)
    R_cell::Float64 = 5e-3 #internal cell resistance (Ω)
    Cmax::Float64 = 50000u"mA*hr" |> upreferred |> ustrip #capacity (Coulomb)
end

voltage_curve(b::Battery, charge_ratio::Real) = 1 #c = charge_ratio
voltage_open(b::Battery, charge_ratio::Real) = b.n_cells * b.V_cell * voltage_curve(b, charge_ratio)
R(b::Battery) = b.n_cells * b.R_cell
ċ(b::Battery, i::Real) = -i/b.Cmax

Base.@kwdef struct EThruster <: SystemDescriptor
    frame::FrameTransform = FrameTransform()
    battery::Battery = Battery()
    motor::ElectricMotor = ElectricMotor()
    gearbox::Gearbox = Gearbox()
    propeller::SimpleProp = SimpleProp()
end

const EThrusterXTemplate = ComponentVector(ω_shaft = 0.0, c_bat = 1.0)
const EThrusterX{T, D} = ComponentVector{T, D, typeof(getaxes(EThrusterXTemplate))} where {T,D}

#maybe disallow default values to avoid subtle bugs when failing to change a
#constructor call in user code after changing the struct definition
Base.@kwdef struct EThrusterY
    throttle::Float64 = 0
    ω_shaft::Float64 = 0
    ω_prop::Float64 = 0
    i::Float64 = 0
    c_bat::Float64 = 0
    wr_c::Wrench = Wrench()
    wr_b::Wrench = Wrench()
    hr_b::SVector{3,Float64} = zeros(SVector{3})
end

Base.@kwdef mutable struct EThrusterU
    throttle::Float64 = 0.0
end

Base.@kwdef mutable struct EThrusterD end

init_x0(::EThruster) = copy(EThrusterXTemplate)
# init_d0(::EThruster) = EThrusterD() #not really needed, defaults to nothing
init_u0(::EThruster) = EThrusterU()
init_y0(::EThruster) = EThrusterY()

################ EThruster System ###################

#the default System constructor works OK in this case

WrenchTrait(::System{EThruster}) = HasWrench()
AngularMomentumTrait(::System{EThruster}) = HasAngularMomentum()

get_wr_b(sys::System{EThruster}) = sys.y.wr_b
get_hr_b(sys::System{EThruster}) = sys.y.hr_b

f_disc!(sys::System{EThruster}) = false

function f_cont!(sys::System{EThruster}, kin::KinData, air::AirData)

    @unpack ẋ, x, y, u, params = sys #no need for subsystems
    @unpack frame, battery, motor, propeller, gearbox = params
    @unpack n, η = gearbox
    @unpack ω_shaft, c_bat = x

    throttle = u.throttle

    ω_prop = ω_shaft / n
    wr_c = get_wrench(propeller, ω_prop, air)
    wr_b = frame(wr_c)

    i = (throttle * voltage_open(battery, c_bat) - back_emf(motor, ω_shaft)) /
        (R(battery) + R(motor))

    M_eng_shaft = torque(motor, i, ω_shaft)
    M_air_prop = wr_c.M[1]

    hr_c = SVector(motor.J * ω_shaft + propeller.J * ω_prop, 0, 0)
    hr_b = frame.q(hr_c)

    ω_shaft_dot = (M_eng_shaft + M_air_prop/(η*n)) / (motor.J + propeller.J/(η*n^2))
    ẋ.ω_shaft = ω_shaft_dot
    ẋ.c_bat = ċ(battery, i)

    sys.y = EThrusterY(throttle, ω_shaft, ω_prop, i, c_bat, wr_c, wr_b, hr_b)

    return nothing

end

function plots(t, data::AbstractVector{<:EThrusterY}; mode, save_path, kwargs...)

    @unpack wr_c, wr_b = StructArray(data)

    pd = Dict{String, Plots.Plot}()

    pd["01_wr_Oc_c"] = thplot(t, wr_c;
        plot_title = "Thruster Wrench [Thruster Frame]",
        wr_source = "thr", wr_frame = "c",
        kwargs...)

    pd["02_wr_Ob_b"] = thplot(t, wr_b;
        plot_title = "Thruster Wrench [Airframe]",
        wr_source = "thr", wr_frame = "b",
        kwargs...)

    save_plots(pd; save_path)

end

##################### Hand of God ##########################

# Base.@kwdef struct HandOfGod <: SystemDescriptor
#     frame::FrameTransform = FrameTransform()
#     enabled::Bool = false
# end

# Base.@kwdef struct HandOfGodY
#     wr_c::Wrench = Wrench()
#     wr_b::Wrench = Wrench()
#     hr_b::SVector{3,Float64} = zeros(SVector{3})
# end

# Base.@kwdef mutable struct HandOfGodU
#     F::MVector{3,Float64} = 0.0
#     M::MVector{3,Float64} = 0.0
# end

# #required to make EThruster compatible with System
# Base.@kwdef mutable struct EThrusterD end

# init_x0(::EThruster) = copy(EThrusterXTemplate)
# init_d0(::EThruster) = EThrusterD()
# init_u0(::EThruster) = EThrusterU()
# init_y0(::EThruster) = EThrusterY()

# function PropulsionGroup(nt::NamedTuple{L, T}  where {L, T<:NTuple{N,AbstractThruster} where {N}})
#     PropulsionGroup{nt}()
# end
#= #interestingly, this does not work:
PropulsionGroup(nt::NamedTuple{L, NTuple{N, T}  where {L,N,T<:NTuple{N,
Propulsion.AbstractThruster}}) = PropulsionGroup{nt}()

#the reason is that NamedTuple, unlike Tuple (and therefore NTuple) is
#NOT covariant. that is:
#(EThruster(), EThruster()) isa NTuple{N, AbstractThruster} where {N} == true
#however:
#(a=EThruster(), b=NThruster) isa NamedTuple{L, NTuple{N, AbstractThruster} where {N} == false
=#
end #module


text:
###############################################
"""
    poly = tilepolygon(tile::VisClaw.AMRGrid)

generate a polygon data of a tile
"""
function tilepolygon(tile::VisClaw.AMRGrid)
    x1, x2, y1, y2 = VisClaw.getlims(tile)
    ll = GeometricalPredicates.Point(x1, y1)
    lr = GeometricalPredicates.Point(x2, y1)
    ur = GeometricalPredicates.Point(x2, y2)
    ul = GeometricalPredicates.Point(x1, y2)
    poly = GeometricalPredicates.Polygon(ll, lr, ur, ul)
    return poly
end

###############################################
"""
    coarsegridmask!(tiles::Vector{VisClaw.AMRGrid})
    coarsegridmask!(amrs::VisClaw.AMR)

replace values at coarser grids (lower levels) into NaN
"""
function coarsegridmask!(tiles::Vector{VisClaw.AMRGrid})
    ## return if single tile
    length(tiles)==1 && (return tiles)
    ## levels
    level_tiles = getfield.(tiles, :AMRlevel)
    maxlevel = maximum(level_tiles)

    for tl in tiles
        tl.AMRlevel == maxlevel && (continue)

        ## generate point data for all points in the target tile
        x_target, y_target = VisClaw.meshtile(tl)
        cellp = GeometricalPredicates.Point.(x_target, y_target)
        ## get the corners of the target tile
        xl, xr, yb, yt = VisClaw.getlims(tl)

        ## find tiles: one level finer
        ind_fine = tl.AMRlevel+1 .== level_tiles
        tile_fine = tiles[ind_fine]
        nfine = length(tile_fine)
        isempty(tile_fine) && (continue)

        ## fine tiles: inside of the target tiles
        ind_inside = trues(nfine)
        for j = 1:nfine
            x1, x2, y1, y2 = VisClaw.getlims(tile_fine[j])
            if x2 < xl || xr < x1 || y2 < yb || yt < y1
                ind_inside[j] = false
            end
        end
        tile_fine = tile_fine[ind_inside]
        nfine = length(tile_fine)
        isempty(tile_fine) && (continue)

        ## find grid where finer grids are assigned
        for j = 1:nfine
            poly = VisClaw.tilepolygon(tile_fine[j])
            inside = [GeometricalPredicates.inpolygon(poly, cellp[irow, jcol]) for irow=1:tl.my, jcol=1:tl.mx]
            if isa(tl, VisClaw.Velocity)
                tl.u[inside] .= NaN
                tl.v[inside] .= NaN
                tl.vel[inside] .= NaN
            elseif isa(tl, VisClaw.Storm)
                tl.u[inside] .= NaN
                tl.v[inside] .= NaN
                tl.slp[inside] .= NaN
            elseif isa(tl, VisClaw.SurfaceHeight)
                tl.eta[inside] .= NaN
            end
        end
    end
    return tiles
end
###############################################
function coarsegridmask!(amrs::VisClaw.AMR)
    amrs.amr = map(coarsegridmask!, amrs.amr)
    return amrs
end
###############################################


text:
"""
Package: Forecast

    co2(query = "paper")

Return dataset with atmospheric Carbon Dioxide Dry Air Mole Fractions from quasi-continuous measurements at Mauna Loa, Hawaii.

K.W. Thoning, A.M. Crotwell, and J.W. Mund (2020), Atmospheric Carbon Dioxide Dry Air Mole Fractions from continuous measurements at Mauna Loa, Hawaii, Barrow, Alaska, American Samoa and South Pole. 1973-2019, Version 2020-08 National Oceanic and Atmospheric Administration (NOAA), Global Monitoring Laboratory (GML), Boulder, Colorado, USA https://doi.org/10.15138/yaf1-bk21 FTP path: ftp://aftp.cmdl.noaa.gov/data/greenhouse_gases/co2/in-situ/surface/

# Arguments
- `query`: Defaults to "paper" and accepts the following values:
    - "full": Full original dataset from 1973 to 2020
    - "paper": Subset used in "STL: A Seasonal-Trend Decomposition Procedure Based on Loess" from Cleveland et. al.

# Returns
Dataframe containing information based on the parameter `query`

# Examples
```julia-repl
julia> co2()
[ Info: Dataset used in Cleveland et al. paper
4612×2 DataFrame
  Row │ date        co2
      │ Date        Float64?
──────┼────────────────────────
    1 │ 1974-05-17      333.38
    2 │ 1974-05-18      333.11
    3 │ 1974-05-19      333.46
   [...]
```
"""
function co2(query::String = "paper")

    @assert query in ["full","paper"]

    path = art_path("co2")

    co2_df = open(joinpath(path, "co2.csv")) do file
        CSV.read(file,DataFrame)
    end

    if query == "full"
        @info "Full dataset from 1973 to 2020"
        return co2_df
    end

    if query == "paper"
        @info "Dataset used in Cleveland et al. paper"

        # recode missing values
        co2_df.value = replace(co2_df.value, -999.99 => missing)

        dates_co2 = Dates.Date(1973,1,1):Dates.Day(1):Dates.Date(2019,12,31)
        co2_sdf = DataFrame([dates_co2, co2_df.value],[:date,:co2])

        # paper dates
        # using 17 may instead 17 april (as the paper claims)
        # since there seems to be no data for april in 1974
        dates_co2_stl = Dates.Date(1974,5,17):Dates.Day(1):Dates.Date(1986,12,31)

        # revome leap year day
        dates_co2_stl = filter(dates_co2_stl) do x
            (Dates.isleapyear(x) & (Dates.month(x) == 2)) ? Dates.day(x) != 29 : true
        end

        return @subset(co2_sdf, in.(:date,  Ref(dates_co2_stl)))

    end

end


text:
module TestPackage

f(x) = 2x

export f

end

max_stars_repo_name: SciML/TestPackage

max_issues_repo_name: SciML/TestPackage

max_forks_repo_name: SciML/TestPackage


text:
include("extra_modes.jl")

function after_normal(b::Buffer)
    clear_arg(b)
    clamp!(b)
    set_boundaries(b)
end

function delete_char(b::Buffer)
    setundo(b)
    deleteat(b, b.cursor.pos, parse_n(b))
    after(b)
end

function join_arg(b::Buffer)
    setundo(b)
    n = parse_n(b.args)
    join_lines(b, (b.cursor.pos[1],b.cursor.pos[1]+n))
    after(b)
end
function insert(b::Buffer)
    setmode(b,insert_mode)
    if x(b) == 0
        b.cursor.pos[2] = 1
    end
end
function inserta(b::Buffer)
    b.cursor.pos[2] += 1
    setmode(b, insert_mode)
end
function insert_beginning(b::Buffer) #different from vim on white non-empty lines
    b.cursor.pos[2] = nextword_pos(line(b), 1)
    insert(b)
end
insert_end(b::Buffer) = (move_eol(b); inserta(b))

function enter_cmdmode(b::Buffer)
    setmode(b,command_mode)
end

function start_replay(b::Buffer)
    setmode(b, replay_mode)
end

function enter_deletemode(b::Buffer)
    #This should be in a before() function
    #Also unsure if this is a good way of storing vars
    b.state[:lx] = xs(b)
    b.state[:ly] = ys(b)
    setmode(b, delete_mode)
end

function enter_yankmode(b::Buffer)
    #This should be in a before() function
    #Also unsure if this is a good way of storing vars
    b.state[:lx] = xs(b)
    b.state[:ly] = ys(b)
    setmode(b, yank_mode)
end

function enter_registermode(b::Buffer)
    setmode(b, register_clipboardmode )
end

function enter_replacemode(b::Buffer)
    setmode(b, replace_mode)
end

function pastea_register(b::Buffer)
    s = b.state[:register][1]
    pastea(b, s)
    b.state[:register] = "\""
end
function paste_register(b::Buffer)
    s = b.state[:register][1]
    paste(b, s)
    b.state[:register] = "\""
end
const normal_actions = merge(movements,
                       Dict('i'    => insert,
                            'a'    => inserta,
                            'I'    => insert_beginning,
                            'A'    => insert_end,
                            '\e'   => clear_arg,
                            ':'    => enter_cmdmode,
                            'x'    => delete_char,
                            'r'    => enter_replacemode,
                            'J'    => join_arg,
                            'q'    => start_record,
                            '@'    => start_replay,
                            'd'    => enter_deletemode,
                            'y'    => enter_yankmode,
                            '"'    => enter_registermode,
                            'p'    => pastea_register,
                            'P'    => paste_register,
                            'u'    => undo,
                            '\x12' => redo,
                            'Z'    => quit,
                           )
                      )

function nmap(c::Char, f)
    normal_actions[c] = f
end
function nmap(c::Char, s::String)
    normal_actions[c] = b -> replay(b, s)
end
normal_mode = Mode("normal", normal_actions)


text:
#TODO: non-matchin g blocks
isblockbanded(A::AbstractTriangular) =
    isblockbanded(parent(A))
isbandedblockbanded(A::AbstractTriangular) =
    isbandedblockbanded(parent(A))
function blockbandwidths(A::Union{UpperTriangular,UnitUpperTriangular})
    P = parent(A)
    if hasmatchingblocks(P)
        (min(0,blockbandwidth(P,1)), blockbandwidth(P,2))
    else
        blockbandwidths(P)
    end
end
function blockbandwidths(A::Union{LowerTriangular,UnitLowerTriangular})
    P = parent(A)
    if hasmatchingblocks(P)
        (blockbandwidth(P,1), min(0,blockbandwidth(P,2)))
    else
        blockbandwidths(P)
    end
end
subblockbandwidths(A::AbstractTriangular) = subblockbandwidths(parent(A))


text:
# Test layers and data/model movements on and off the GPU
# Add tests for layers and their gradients on the GPU
# Most of the forward passes should be fine being applied
# to bitstype objects, but this gives higher coverage for our use-cases
# Check that getting the gradients does not throw

# generic movement tests
@testset "Basic GPU Movement" begin
  @test gradient(x -> sum(gpu(x)), rand(3,3)) isa Tuple
  @test gradient(x -> sum(cpu(x)), gpu(rand(3,3))) isa Tuple
end

# TODO: These layers get into scalar indexing
# `AlphaDropout` throws a compilation error on GPUs,
# whereas, the rest are scalar indexing issues.
const BROKEN_LAYERS = Union{DepthwiseConv,
                            AlphaDropout}

function gpu_gradtest(name::String, layers::Vector, x_cpu, args...;
            setmode=false, test_cpu=true, rtol=1e-5, atol=1e-5)
  @testset "$name GPU grad tests" begin
    for layer in layers
      @testset "$layer GPU grad test" begin
        l_cpu = layer(args...)
        if l_cpu isa BROKEN_LAYERS
          l_gpu, x_gpu = l_cpu |> gpu, x_cpu |> gpu
          @test_broken gradient(() -> sum(l_gpu(x_gpu)), Flux.params(l_gpu)) isa Flux.Zygote.Grads
        else
          gpu_autodiff_test(l_cpu, x_cpu,
              test_equal=test_cpu, rtol=rtol, atol=atol)
          if setmode
            testmode!(l_cpu)
            gpu_autodiff_test(l_cpu, x_cpu,
              test_equal=test_cpu, rtol=rtol, atol=atol)
          end
        end
      end
    end
  end
end

# Just to give testset in gradtest meaningful labels
ConvNoBias(args...) = Conv(args...; bias=false)
ConvTransposeNoBias(args...) = ConvTranspose(args...; bias=false)
CrossCorNoBias(args...) = CrossCor(args...; bias=false)
DepthwiseConvNoBias(args...) = DepthwiseConv(args...; bias=false)
r = rand(Float32, 28, 28, 1, 1)
conv_layers = [Conv, ConvNoBias, ConvTranspose, ConvTransposeNoBias, CrossCor, CrossCorNoBias, DepthwiseConv, DepthwiseConvNoBias]
gpu_gradtest("Conv", conv_layers, r, (2,2), 1=>3)

pooling_layers = [MaxPool, MeanPool]
gpu_gradtest("Pooling", pooling_layers, r, (2,2))

adaptive_pooling_layers = [AdaptiveMaxPool, AdaptiveMeanPool]
gpu_gradtest("AdaptivePooling", adaptive_pooling_layers, r, (7,7))

dropout_layers = [Dropout, AlphaDropout]
gpu_gradtest("Dropout", dropout_layers, r, 0.5f0; test_cpu=false, setmode=true) # dropout is not deterministic

layer_norm = [i -> LayerNorm(i; affine=false), i -> LayerNorm(i; affine=true)]
gpu_gradtest("LayerNorm 1", layer_norm, rand(Float32, 8, 8, 3, 4), 8)
gpu_gradtest("LayerNorm 2", layer_norm, rand(Float32, 8, 8, 3, 4), (8,8))
gpu_gradtest("LayerNorm 3", layer_norm, rand(Float32, 5, 4), 5)

batch_norm = [BatchNorm]
gpu_gradtest("BatchNorm 3d", batch_norm, rand(Float32, 8, 8, 8, 3, 4), 3, setmode=false) # bug in CUDA.jl with gradient in testmode
gpu_gradtest("BatchNorm 2d", batch_norm, rand(Float32, 8, 8, 3, 4), 3, setmode=false) # bug in CUDA.jl with gradient in testmode
gpu_gradtest("BatchNorm 1d", batch_norm, rand(Float32, 8, 3, 4), 3, setmode=false) # bug in CUDA.jl with gradient in testmode
gpu_gradtest("BatchNorm fullyconn", batch_norm, rand(Float32, 5,4), 5, setmode=false)

instancenorm = [i -> InstanceNorm(i; affine=false), i -> InstanceNorm(i; affine=true)]
gpu_gradtest("InstanceNorm 3d", instancenorm, rand(Float32, 8, 8, 8, 3, 4), 3, setmode=true)
gpu_gradtest("InstanceNorm 2d", instancenorm, rand(Float32, 8, 8, 3, 4), 3, setmode=true)
gpu_gradtest("InstanceNorm 1d", instancenorm, rand(Float32, 8, 3, 4), 3, setmode=true)

groupnorm = [(i, j) -> GroupNorm(i, j; affine=false), (i, j) -> GroupNorm(i, j; affine=true)]
gpu_gradtest("GroupNorm 3d", groupnorm, rand(Float32, 8, 8, 8, 12, 4), 12, 3, setmode=true)
gpu_gradtest("GroupNorm 2d", groupnorm, rand(Float32, 8, 8, 12, 4), 12, 3, setmode=true)
gpu_gradtest("GroupNorm 1d", groupnorm, rand(Float32, 8, 3, 12, 4), 12, 3, setmode=true)

upsample = [x -> Upsample(scale=x)]
gpu_gradtest("Upsample 2d", upsample, rand(Float32, 3, 4, 2, 3), (2,2))
gpu_gradtest("Upsample 1d", upsample, rand(Float32, 3, 4, 2, 3), (2,))

pixelshuffle = [PixelShuffle]
gpu_gradtest("PixelShuffle 2d", pixelshuffle, rand(Float32, 3, 4, 18, 3), 3)
gpu_gradtest("PixelShuffle 1d", pixelshuffle, rand(Float32, 3, 18, 3), 3)

@testset "function layers" begin
  x = rand(Float32, 3,3)
  gpu_autodiff_test(x -> sum(Flux.normalise(x; dims=1)), x)
  gpu_autodiff_test(x -> sum(Flux.normalise(x; dims=2)), x)
  gpu_autodiff_test(x -> sum(Flux.normalise(x)), x)
end

@testset "BatchNorm mix stuff" begin
  m_cpu = BatchNorm(2)
  m_gpu = m_cpu |> gpu
  x_cpu = rand(Float32, 3, 2, 2)
  x_gpu = x_cpu |> gpu

  ## In :auto mode, track statistics only in gradient contest
  μ_cpu = copy(m_cpu.μ)
  m_cpu(x_cpu)
  @test m_cpu.μ ≈ μ_cpu
  gradient(() -> sum(m_cpu(x_cpu)), Flux.params(m_cpu))
  @test !(m_cpu.μ ≈ μ_cpu)

  μ_gpu = copy(m_gpu.μ)
  m_gpu(x_gpu)
  @test m_gpu.μ ≈ μ_gpu
  gradient(() -> sum(m_gpu(x_gpu)), Flux.params(m_gpu))
  @test !(m_gpu.μ ≈ μ_gpu)

  @test Array(m_gpu.μ) ≈ m_cpu.μ

  ## In testmode, never track statistics
  testmode!(m_cpu)
  μ_cpu = copy(m_cpu.μ)
  m_cpu(x_cpu)
  @test m_cpu.μ ≈ μ_cpu
  gradient(() -> sum(m_cpu(x_cpu)), Flux.params(m_cpu))
  @test m_cpu.μ ≈ μ_cpu

  testmode!(m_gpu)
  μ_gpu = copy(m_gpu.μ)
  m_gpu(x_gpu)
  @test m_gpu.μ ≈ μ_gpu
  gradient(() -> sum(m_gpu(x_gpu)), Flux.params(m_gpu))
  @test m_gpu.μ ≈ μ_gpu

  ## In trainmode, always track statistics
  trainmode!(m_cpu)
  μ_cpu = copy(m_cpu.μ)
  m_cpu(x_cpu)
  @test !(m_cpu.μ ≈ μ_cpu)
  μ_cpu = copy(m_cpu.μ)
  gradient(() -> sum(m_cpu(x_cpu)), Flux.params(m_cpu))
  @test !(m_cpu.μ ≈ μ_cpu)

  trainmode!(m_gpu)
  μ_gpu = copy(m_gpu.μ)
  m_gpu(x_gpu)
  @test !(m_gpu.μ ≈ μ_gpu)
  μ_gpu = copy(m_gpu.μ)
  gradient(() -> sum(m_gpu(x_gpu)), Flux.params(m_gpu))
  @test !(m_gpu.μ ≈ μ_gpu)

  ## No errors if input type mistmatch
  x_cpu = rand(Float64, 3, 2, 2)
  x_gpu = x_cpu |> gpu
  m_cpu(x_cpu)
  gradient(() -> sum(m_cpu(x_cpu)), Flux.params(m_cpu))
  m_gpu(x_gpu)
  gradient(() -> sum(m_gpu(x_gpu)), Flux.params(m_gpu))
end

@testset "Zeros mapped for $cl" for cl in (Conv, ConvTranspose, CrossCor, DepthwiseConv)
  l = cl((2,2), 1=>3, bias = false) |> gpu
  ip = zeros(Float32, 28,28,1,1) |> gpu
  if l isa BROKEN_LAYERS
    @test_broken sum(l(ip)) ≈ 0.f0
    @test_broken gradient(() -> sum(l(ip)), Flux.params(l)) isa Flux.Zygote.Grads
  else
    @test sum(l(ip)) ≈ 0.f0
    gs = gradient(() -> sum(l(ip)), Flux.params(l))
    @test l.bias ∉ gs.params
  end
end

@testset "Dense with Zeros bias" begin
  l = Dense(ones(Float32, 4,3), Flux.Zeros()) |> gpu
  ip = zeros(Float32, 3, 7) |> gpu

  @test sum(l(ip)) ≈ 0.f0
  gs = gradient(() -> sum(l(ip)), Flux.params(l))
  @test l.b ∉ gs.params
end

@testset "Two-streams Bilinear" begin
  x = zeros(Float32,10,9) |> gpu
  y = zeros(Float32,2,9) |> gpu
  b = Flux.Bilinear(10, 2, 3) |> gpu
  @test size(b(x,y)) == (3,9)
  @test sum(abs2, b(x,y)) ≈ 0f0
  gs_gpu = gradient(() -> sum(abs2.(b(x, y))), params(b))
  b_cpu, x_cpu, y_cpu = b |> cpu, x |> cpu, y |> cpu
  gs_cpu = gradient(() -> sum(abs2.(b_cpu(x_cpu, y_cpu))), params(b_cpu))
  for (pgpu, pcpu) in zip(params(b), params(b_cpu))
    @test gs_cpu[pcpu] ≈ Array(gs_gpu[pgpu])
  end
end

@testset "Parallel" begin
  @testset "zero sum" begin
    input = randn(10, 10, 10, 10) |> gpu
    layer_gpu = Parallel(+, zero, identity) |> gpu
    @test layer_gpu(input) == input
    @test layer_gpu(input) isa Flux.CUDA.CuArray
  end

  @testset "vararg input" begin
    inputs = (randn(10), randn(5), randn(4)) .|> gpu
    layer = Parallel(+, Dense(10, 2), Dense(5, 2), Dense(4, 2)) |> gpu
    @test size(layer(inputs)) == (2,)
  end

  @testset "gradient" begin
    input_cpu = randn(10, 10, 10, 10)
    input_gpu = input_cpu |> gpu
    layer_cpu = Parallel(+, x -> zero(x), identity)
    layer_gpu = layer_cpu |> gpu
    gs_cpu = gradient(() -> sum(abs2.(layer_cpu(input_cpu))), params(layer_cpu))
    gs_gpu = gradient(() -> sum(abs2.(layer_gpu(input_gpu))), params(layer_gpu))
    for (pgpu, pcpu) in zip(params(layer_cpu), params(layer_gpu))
      @test gs_cpu[pcpu] ≈ gs_gpu[pgpu]
    end
  end
end


text:
using Random
using DataStructures

@testset "quicksort" begin

import CUDA.Quicksort: flex_lt, find_partition,
        partition_batches_kernel, consolidate_batch_partition, bubble_sort

@testset "integer functions" begin
    @test flex_lt(1, 2, false, isless, identity) == true
    @test flex_lt(1, 2, true, isless, identity) == true
    @test flex_lt(2, 2, false, isless, identity) == false
    @test flex_lt(2, 2, true, isless, identity) == true
    @test flex_lt(3, 2, false, isless, identity) == false
    @test flex_lt(3, 2, true, isless, identity) == false

    @test find_partition([1, 2, 2, 3, 4, 1, 2, 2, 3, 4], 3, 0, 5, false, isless, identity) == 4
    @test find_partition([1, 2, 2, 3, 4, 1, 2, 2, 3, 4], 3, 5, 10, false, isless, identity) == 9
    @test find_partition([1, 2, 2, 3, 4, 1, 2, 2, 3, 4], 3, 0, 5, true, isless, identity) == 3
    @test find_partition([1, 2, 2, 3, 4, 1, 2, 2, 3, 4], 3, 5, 10, true, isless, identity) == 8
end

function test_batch_partition(T, N, lo, hi, seed, lt=isless, by=identity)
    my_range = lo + 1 : hi
    Random.seed!(seed)
    original = rand(T, N)
    A = CuArray(original)

    pivot = rand(original[my_range])
    block_N, block_dim = -1, -1

    kernel = @cuda launch=false partition_batches_kernel(A, pivot, lo, hi, true, lt, by)

    get_shmem(threads) = threads * (sizeof(Int) + sizeof(T))
    config = launch_configuration(kernel.fun, shmem=threads->get_shmem(threads), max_threads=1024)

    threads = prevpow(2, config.threads)
    blocks = ceil(Int, (hi - lo) ./ threads)
    block_N = blocks
    block_dim = threads
    @assert block_dim >= 32 "This test assumes block size can be >= 32"

    kernel(A, pivot, lo, hi, true, lt;
           threads=threads, blocks=(1,blocks), shmem=get_shmem(threads))
    synchronize()

    post_sort = Array(A)

    sort_match = true

    for block in 1:block_N
        block_range = lo + 1 + (block - 1) * block_dim: min(hi, lo + block * block_dim)
        temp = original[block_range]
        # this shows that batch partitioning is a stable sort where key for each value v is
        # whether v > or <= pivot
        expected_sort = vcat(filter(x -> x < pivot, temp), filter(x -> x >= pivot, temp))
        sort_match &= post_sort[block_range] == expected_sort
    end

    @test sort_match
end

@testset "batch partition" begin
    test_batch_partition(Int8, 10000, 2000, 6000, 0)
    test_batch_partition(Int8, 10000, 2000, 6000, 1)
    test_batch_partition(Int8, 10000000, 0, 10000000, 0)
    test_batch_partition(Int8, 10000000, 5000, 500000, 0)
    test_batch_partition(Int8, 10000, 0, 10000, 0)
    test_batch_partition(Int8, 10000, 2000, 6000, 0)
    test_batch_partition(Int8, 10000, 2000, 6000, 1)
    test_batch_partition(Int8, 10000000, 0, 10000000, 0)
    test_batch_partition(Int8, 10000000, 5000, 500000, 0)

    test_batch_partition(Float32, 10000, 0, 10000, 0)
    test_batch_partition(Float32, 10000, 2000, 6000, 0)
    test_batch_partition(Float32, 10000, 2000, 6000, 1)
    test_batch_partition(Float32, 10000000, 0, 10000000, 0)
    test_batch_partition(Float32, 10000000, 5000, 500000, 0)
    test_batch_partition(Float32, 10000, 0, 10000, 0)
    test_batch_partition(Float32, 10000, 2000, 6000, 0)
    test_batch_partition(Float32, 10000, 2000, 6000, 1)
    test_batch_partition(Float32, 10000000, 0, 10000000, 0)
    test_batch_partition(Float32, 10000000, 5000, 500000, 0)
end

function test_consolidate_kernel(vals, pivot, my_floor, L, b_sums, dest, parity, lt, by)
    i = threadIdx().x
    p = consolidate_batch_partition(vals, pivot, my_floor, L, b_sums, parity, lt, by)
    if i == 1
        dest[1] = p
    end
    return nothing
end

function test_consolidate_partition(T, N, lo, hi, seed, block_dim, lt=isless, by=identity)
    # assuming partition_batches works, we can validate consolidate by
    # checking that together they partition a large domain
    my_range = lo + 1 : hi
    Random.seed!(seed)
    original = rand(T, N)
    A = CuArray(original)
    pivot = rand(original[my_range])

    threads = blocks = -1
    sums = CuArray(zeros(Int, ceil(Int, hi - lo / block_dim)))

    kernel = @cuda launch=false partition_batches_kernel(A, pivot, lo, hi, true, lt, by)

    get_shmem(threads) = threads * (sizeof(Int) + sizeof(T))
    config = launch_configuration(kernel.fun, shmem=threads->get_shmem(threads), max_threads=1024)

    threads = isnothing(block_dim) ? prevpow(2, config.threads) : block_dim
    blocks = ceil(Int, (hi - lo) ./ threads)

    kernel(A, pivot, lo, hi, true, lt, by; threads=threads, blocks=(1,blocks), shmem=get_shmem(threads))
    synchronize()
    dest = CuArray(zeros(Int, 1))

    @cuda threads=threads test_consolidate_kernel(A, pivot, lo, hi - lo, sums, dest, true, lt, by)
    synchronize()

    partition = Array(dest)[1]
    temp = original[my_range]
    post_sort = Array(A)
    # consolidation is a highly unstable sort (again, by pivot comparison as the key) so we
    # compare by counting each element
    cc(x) = x |> counter |> collect |> sort
    @test cc(original) == cc(post_sort)
    @test all(post_sort[lo + 1 : partition] |> cc .== filter(x -> x < pivot, temp) |> cc)
    @test all(post_sort[partition + 1 : hi] |> cc .== filter(x -> x >= pivot, temp) |> cc)
end

@testset "consolidate partition" begin
    test_consolidate_partition(Int8, 10000, 0, 10000, 0, 16)
    test_consolidate_partition(Int8, 10000, 0, 10000, 0, 32)
    test_consolidate_partition(Int8, 10000, 0, 10000, 0, 64)
    test_consolidate_partition(Int8, 10000, 9, 6333, 0, 16)
    test_consolidate_partition(Int8, 10000, 9, 6333, 0, 32)
    test_consolidate_partition(Int8, 10000, 9, 6333, 0, 64)
    test_consolidate_partition(Int8, 10000, 129, 9999, 0, 16)
    test_consolidate_partition(Int8, 10000, 129, 9999, 0, 32)
    test_consolidate_partition(Int8, 10000, 129, 9999, 0, 64)
    test_consolidate_partition(Int8, 10000, 0, 10000, 1, 16)
    test_consolidate_partition(Int8, 10000, 0, 10000, 2, 32)
    test_consolidate_partition(Int8, 10000, 0, 10000, 3, 64)
    test_consolidate_partition(Int8, 10000, 9, 6333, 4, 16)
    test_consolidate_partition(Int8, 10000, 9, 6333, 5, 32)
    test_consolidate_partition(Int8, 10000, 9, 6333, 6, 64)
    test_consolidate_partition(Int8, 10000, 129, 9999, 7, 16)
    test_consolidate_partition(Int8, 10000, 129, 9999, 8, 32)
    test_consolidate_partition(Int8, 10000, 129, 9999, 9, 64)
    test_consolidate_partition(Int8, 10000, 3329, 9999, 10, 16)
    test_consolidate_partition(Int8, 10000, 3329, 9999, 11, 32)
    test_consolidate_partition(Int8, 10000, 3329, 9999, 12, 64)
end

function init_case(T, f, N::Integer)
    a = map(x -> T(f(x)), 1:N)
    c = CuArray(a)
    a, c
end

function init_case(T, f, N::Tuple)
    a = map(f, rand(T, N...))
    c = CuArray(a)
    a, c
end

"""
Tests if `c` is a valid sort of `a`
"""
function check_equivalence(a::Vector, c::Vector; kwargs...)
    counter(a) == counter(c) && issorted(c; kwargs...)
end

"""
Tests if `c` is a valid sort of `a`
"""
function check_equivalence(a::Array, c::Array; dims, kwargs...)
    @assert size(a) == size(c)
    nd = ndims(c)
    k = dims
    sz = size(c)

    1 <= k <= nd || throw(ArgumentError("dimension out of range"))

    remdims = ntuple(i -> i == k ? 1 : size(c, i), nd)
    v(a, idx) = view(a, ntuple(i -> i == k ? Colon() : idx[i], nd)...)
    all(counter(v(a, idx)) == counter(v(c, idx)) && issorted(v(c, idx); kwargs...)
        for idx in CartesianIndices(remdims))
end

"""
`T` - Element type to test
`N` - Either an integer for a vector length, or a tuple for array dimension
`f` - For a vector, fill with, for each index i, `T(f(i))`. Facilitates testing orderings
      For an array, fill with `f(rand(T))`. Facilitates testing distributions
"""
function check_sort!(T, N, f=identity; kwargs...)
    original_arr, device_arr = init_case(T, f, N)
    sort!(device_arr; kwargs...)
    host_result = Array(device_arr)
    check_equivalence(original_arr, host_result; kwargs...)
end

function check_sort(T, N, f=identity; kwargs...)
    original_arr, device_arr = init_case(T, f, N)
    host_result = Array(sort(device_arr; kwargs...))
    check_equivalence(original_arr, host_result; kwargs...)
end

# FIXME: these tests hang when running under compute-sanitizer on CUDA 11.2 with -g2
@not_if_sanitize @testset "interface" begin
    # pre-sorted
    @test check_sort!(Int, 1000000)
    @test check_sort!(Int32, 1000000)
    @test check_sort!(Float64, 1000000)
    @test check_sort!(Float32, 1000000)
    @test check_sort!(Int32, 1000000; rev=true)
    @test check_sort!(Float32, 1000000; rev=true)

    # reverse sorted
    @test check_sort!(Int32, 1000000, x -> -x)
    @test check_sort!(Float32, 1000000, x -> -x)
    @test check_sort!(Int32, 1000000, x -> -x; rev=true)
    @test check_sort!(Float32, 1000000, x -> -x; rev=true)

    @test check_sort!(Int, 10000, x -> rand(Int))
    @test check_sort!(Int32, 10000, x -> rand(Int32))
    @test check_sort!(Int8, 10000, x -> rand(Int8))
    @test check_sort!(Float64, 10000, x -> rand(Float64))
    @test check_sort!(Float32, 10000, x -> rand(Float32))
    @test check_sort!(Float16, 10000, x -> rand(Float16))

    # non-uniform distributions
    @test check_sort!(UInt8, 100000, x -> round(255 * rand() ^ 2))
    @test check_sort!(UInt8, 100000, x -> round(255 * rand() ^ 3))

    # more copies of each value than can fit in one block
    @test check_sort!(Int8, 4000000, x -> rand(Int8))

    # multiple dimensions
    @test check_sort!(Int32, (4, 50000, 4); dims=2)
    @test check_sort!(Int32, (4, 4, 50000); dims=3, rev=true)

    # various sync depths
    for depth in 0:4
        CUDA.limit!(CUDA.LIMIT_DEV_RUNTIME_SYNC_DEPTH, depth)
        @test check_sort!(Int, 100000, x -> rand(Int))
    end

    # using a `by` argument
    @test check_sort(Float32, 100000; by=x->abs(x - 0.5))
    @test check_sort!(Float32, (100000, 4); by=x->abs(x - 0.5), dims=1)
    @test check_sort!(Float32, (4, 100000); by=x->abs(x - 0.5), dims=2)
    @test check_sort!(Float64, 400000; by=x->8*x-round(8*x))
    @test check_sort!(Float64, (100000, 4); by=x->8*x-round(8*x), dims=1)
    @test check_sort!(Float64, (4, 100000); by=x->8*x-round(8*x), dims=2)
    # target bubble sort by using sub-blocksize input:
    @test check_sort!(Int, 200; by=x->x % 2)
    @test check_sort!(Int, 200; by=x->x % 3)
    @test check_sort!(Int, 200; by=x->x % 4)
end

end


text:
using OptimalApplication
using DataFrames
using PrettyTables
using Plots
using Statistics: median, quantile

function plot_params()
    Ts = 2.0 .^ (-4:2:4)
    rs = (2.0 .^ (-5:0))

    n_markets = 1000

    sizerangelog2 = (3, 9)

    sizes = round.(
        Int,
        2.0 .^ (
            sizerangelog2[1] .+ (sizerangelog2[2] - sizerangelog2[1]) * rand(n_markets)
        )
    )

    rats = zeros(length(rs), length(Ts), n_markets)

    @time for k in 1:n_markets
        mkt = VariedCostsMarket(sizes[k])
        opt = optimalportfolio_dynamicprogram(mkt)[2]

        Threads.@threads for j in 1:length(Ts)
            for i in 1:length(rs)
                rats[i, j, k] = optimalportfolio_simulatedannealing(
                    mkt;
                    temp=Ts[j],
                    red=rs[i]
                )[2]
            end
        end

        rats[:, :, k] /= optimalportfolio_dynamicprogram(mkt)[2]
    end

    xts = 2 .^ range(sizerangelog2...)

    pl = plot(
        layout=grid(length(rs), length(Ts)),
        link=:all,
        xscale=:log10,
        xticks=(xts, xts),
        legend=nothing,
        size=(1200, 1200),
        xlim=2 .^ (sizerangelog2 .+ (-0.5, 0.5)),
    )

    avgs = zeros(length(rs), length(Ts))
    for j in 1:length(Ts), i in 1:length(rs)
        avgs[i, j] = quantile(rats[i, j, :], 0.1)
    end
    # [sum(rats[i, j, :]) / length(rats[i, j, :]) for i in 1:length(rs), j in 1:length(Ts)]
    best_avg = maximum(avgs)
    @show best_avg

    for j in 1:length(Ts), i in 1:length(rs)
        hline!(pl[i, j], [avgs[i, j]], c=:navy)
        scatter!(
            pl[i, j],
            sizes,
            rats[i, j, :],
            c=avgs[i, j] ≥ best_avg ? :olivedrab : :crimson,
            msc=:auto,
            alpha=0.7,
            title="T = $(round(Ts[j], digits=2)), r = $(round(rs[i], digits=4))",
            titlefontsize=12,
        )
        if i < length(rs)
            plot!(pl[i, j], xticks=(xts, fill("", length(xts))))
        end
        if j > 1
            plot!(pl[i, j], yticks=(Plots.get_ticks(pl[1, 1], :y)[1], fill("", length(xts))))
        end
    end

    return pl, rats, sizes
end

pl, rats, sizes = plot_params()
display(pl)
# savefig("benchmark/annealingparams-plot.png")
# savefig("benchmark/annealingparams-plot.pdf")

max_stars_repo_name: maxkapur/OptimalApplication

max_issues_repo_name: maxkapur/OptimalApplication

max_forks_repo_name: maxkapur/OptimalApplication


text:
# Computes the average undiscounted perfromance of a policy
function average_performance(t, p, Neps)
    vals = zeros(Neps)
    for i = 1:Neps
        hist = simulate(HistoryRecorder(), t, p)
        vals[i] = undiscounted_reward(hist)
    end
    mean(vals), std(vals) / sqrt(Neps)
end

# Converts a policy to a table that maps states to actions
function policy_to_table(mdp, policy)
    if mdp isa AdversarialGridworldMDP
        lengths = (mdp.size[1], mdp.size[2], mdp.size[1], mdp.size[2])
        policy_table = Array{Symbol, 4}(undef, lengths...)
        for ijk in CartesianIndices(lengths)
            policy_table[ijk.I...] = action(policy, S(ijk.I...))
        end
    else
        lengths = (mdp.size[1], mdp.size[2], mdp.size[1], mdp.size[2], mdp.size[1], mdp.size[2])
        policy_table = Array{Symbol}(undef, lengths...)
        for ijk in CartesianIndices(lengths)
            policy_table[ijk.I...] = action(policy, S2(ijk.I...))[1]
        end
    end

    policy_table
end

# Gets a vector of action values for a given state using POMDPs.value(pol, s, a)
action_values(pol, s) = [value(pol, s, a) for a in actions(pol.mdp)]

# Converts a policy to a table that maps state,action pairs to value
# relies on POMDPs.value(pol, s, a)
function value_to_table(policy)
    mdp = policy.mdp
    lengths = (mdp.size[1],mdp.size[2], mdp.size[1], mdp.size[2])
    policy_table = Array{Float64, 5}(undef, lengths..., length(actions(mdp)))
    for ijk in CartesianIndices(lengths)
        policy_table[ijk.I..., :] .= action_values(policy, S(ijk.I...))
    end
    policy_table
end

# Solves the mdp for the optimal policy using LocalApproximation value iteration
function solve_for_policy(mdp; n_generative_samples = 10, max_iterations = 100, verbose = true)
    Nx, Ny = mdp.size
    if mdp isa AdversarialGridworldMDP
        grid = RectangleGrid(1:Nx, 1:Ny, 1:Nx, 1:Ny)
    else
        grid = RectangleGrid(1:Nx, 1:Ny, 1:Nx, 1:Ny, 1:Nx, 1:Ny)
    end
    solver = LocalApproximationValueIterationSolver(
                    LocalGIFunctionApproximator(grid),
                    is_mdp_generative = true,
                    n_generative_samples = n_generative_samples,
                    verbose = verbose,
                    max_iterations = max_iterations,
                    belres = 1e-6)
    solve(solver, mdp)
end

# Generates a gridworld_with_adversary task
function generate_task(;map_rng = MersenneTwister(0),
                        task_rng = MersenneTwister(0),
                        N_rewards = 4,
                        N_walls = 10,
                        Nx = 10,
                        Ny = 10,
                        max_reward = 100,
                        max_penalty = 100,
                        tprob_train = 1,
                        tprob_test = 0.7,
                        n_generative_samples_train = 10,
                        max_iterations_train = 100,
                        verbose = true,
                        return_orig = false,
                        type = AdversarialGridworldMDP)
    rewards = Dict{GWPos, Float64}()
    perturbed_rewards = Dict{GWPos, Float64}()
    for i = 1:N_rewards
        pos = GWPos(rand(map_rng, 1:Nx), rand(map_rng, 1:Ny))
        rewards[pos] = 1
        perturbed_rewards[pos] = rand(task_rng, 1:max_reward)
    end
    perturbed_penalty = rand(task_rng, 1:max_penalty)

    walls = GWPos[]
    while length(walls) < N_walls
        pos = GWPos(rand(map_rng, 1:Nx), rand(map_rng, 1:Ny))
        if !(pos in keys(rewards))
            push!(walls, pos)
        end
    end

    # Create the mdp that will be solved for the policy (perturbed rewards)
    mdp = type(rewards = perturbed_rewards,
                             walls = walls,
                             tprob = tprob_train,
                             failure_penalty = perturbed_penalty)
    # Solve for the ego policy and store it as an easily accessible table
    policy = solve_for_policy(mdp,
                    n_generative_samples = n_generative_samples_train,
                    max_iterations = max_iterations_train,
                    verbose = verbose)
    policy_table = policy_to_table(mdp, policy)

    # Construct the adversarial mdp with the solved ego policy with a reward that
    # is consistent across tasks
    adv_mdp = type(rewards = rewards,
                                 walls = walls,
                                 agent_gets_action = :adversary,
                                 tprob = 1,
                                 ego_policy = (s, rng) -> rand(rng) < tprob_test ? policy_table[s...] : rand(rng, actions(mdp)),
                                 failure_penalty = 1)
     if return_orig
         return mdp, policy_table, adv_mdp
     end
     adv_mdp
end

# Constructs (and optionally saves) a set of tasks using a specified random seed
function generate_task_set(seed, N; tprob_test = 0.7, folder = nothing, save = false, lzeros = 2, keep_map = true, verbose = false)
    try mkdir(folder) catch end
    task_rng = MersenneTwister(seed)
    map_rng = MersenneTwister(seed)
    tasks = []
    for i=1:N
        println("Generating task ", i, " in ", folder)
        mdp, policy_table, advmdp = generate_task(task_rng = task_rng, map_rng = keep_map ? MersenneTwister(0) : map_rng, verbose = verbose, tprob_test = tprob_test, return_orig = true)
        push!(tasks, advmdp)
        if save
            subfolder = string(folder, "/task_", lpad(i, lzeros, "0"),"/")
            try mkdir(subfolder) catch end
            serialize(string(subfolder, "mdp"), mdp)
            serialize(string(subfolder, "policy_table"), policy_table)
            serialize(string(subfolder, "advmdp"), advmdp)
            draw(PDF(string(subfolder, "mdp.pdf")), POMDPModelTools.render(mdp, initialstate(mdp)))
            draw(PDF(string(subfolder, "advmdp.pdf")), POMDPModelTools.render(advmdp, initialstate(advmdp)))
        end
    end
    tasks
end

# Solves for (and optionally saves) the optimal policy of a set of tasks
function solve_for_policies(tasks; folder = nothing, save = false, lzeros = 2)
    try mkdir(folder) catch end
    policies = []
    for i in 1:length(tasks)
        t = tasks[i]
        pol = solve_for_policy(t, verbose = false)
        push!(policies, pol)
        if save
            name = string(folder, "/policy_", lpad(i, lzeros, "0"), ".jls")
            println("writing ", name)
            serialize(name, pol)
        end
    end
    policies
end

# converts (and optionally saves) a set of policies to tables
function convert_policies_to_tables(policies; folder = nothing, save = false, lzeros = 2)
    try mkdir(folder) catch end
    tables = []
    for i in 1:length(policies)
        tab = value_to_table(policies[i])
        push!(tables, tab)
        if save
            name = string(folder, "/table_", lpad(i, lzeros, "0"), ".jls")
            println("writing ", name)
            serialize(name, tab)
        end
    end
    tables
end

# Deserializes a folder of items
function load_items(folder)
    files = readdir(folder)
    items = []
    for f in files
        i = deserialize(string(folder,"/",f))
        push!(items, i)
    end
    return items
end


text:
KitML.create_neural_closure(10, 5)

nn = Chain(Dense(4, 4, tanh), Dense(4, 1))
KitML.neural_closure(nn, randn(Float32, 4))

fnn = FastChain(FastDense(4, 4, tanh), FastDense(4, 1))
p = initial_params(fnn)
KitML.neural_closure(fnn, p, randn(Float32, 4))

cd(@__DIR__)
model = KitML.load_model("tfmodel.h5"; mode = :tf)
KitML.neural_closure(model, randn(Float32, 1, 4))


text:
module test_jive_runtests_get_all_files

using Test
using Jive

dir = normpath(@__DIR__, "..") # Jive/test/jive/
skip = String[]
targets = ["//./pipe", "USE_REVISE", "USE_PLOTPANE", "runtests/get_all_files"]
(all_files, start_idx) = Jive.get_all_files(dir, skip, targets)
@test all_files == ["runtests/get_all_files.jl"]

dir = normpath(@__DIR__, ".") # Jive/test/jive/runtests/
skip = String[]
targets = ["//./pipe", "USE_REVISE", "USE_PLOTPANE", "get_all_files"]
(all_files, start_idx) = Jive.get_all_files(dir, skip, targets)
@test all_files == ["get_all_files.jl"]

end # module test_jive_runtests_get_all_files


text:
module MR
using Distances, StatsBase

export Isomap, mds, procrustes, IOMR, MR_Missing

function Kmin(D,K)
    (n,n) = size(D)
    for i = 1:n
        M = findKmin(D[i,:], K+1)
        for j = 1:n
            if D[i,j] > M
                D[i,j] = Inf
            end
        end
    end

    for i = 1:n
        for j = 1:n
            if D[i,j] != Inf
                D[j,i] = D[i,j]
            end
        end
    end

    return D
end

function findKmin(A,k)
    A = Set(A)
    for i = 1:k-1
        m = minimum(A)
        A = delete!(A,m)
        if isempty(A)
            return m
        end
    end

    return minimum(A)
end

function apsp(G)
    (n,n) = size(G)
    U = Inf*ones(size(G))
    for i = 1:n
        U[i,i] = 0
    end

    for i = 1:n
        for j = 1:n
            if G[i,j] != Inf
                U[i,j] = G[i,j]
            end
        end
    end

    for k = 1:n
        for i=1:n
            for j=1:i-1
                if U[i,j] > U[i,k] + U[k,j]
                    U[i,j] = U[i,k]+U[k,j]
                    U[j,i] = U[i,j]
                end
            end
        end
    end

    return U
end

function mds(D, d; center=true,align=true)
    n = size(D, 1)

    P = eye(n) - ones(n, n) / n
    K = -0.5*(P*(D.^2)*P')
    K = (K+K')/2

    e, V = eig(K)
    idx = sortperm(e, rev=true)[1:d]

    if (d > sum(e .> 0))
        println("Using some eigenvaectors for negative eigenvalues")
    end

    X =  V[:,idx]*diagm(sqrt.(e[idx]))

    if center == true
        X = X .- mean(X,1)
    end

    if align == true
        for i =1:d
            j = indmax(abs.(X[:,i]))
            X[:,i] .*= sign.(X[j,i])
        end
    end

    return X
end

function procrustes(X, Y)
    X₀ = X - ones(size(X)[1],1)*mean(X, 1)
    Y₀ = Y - ones(size(X)[1],1)*mean(Y, 1)
    Z = X₀'*Y₀

    U,S,V = svd(Z, thin=true)
    Q = V*U'
    α = trace(Z*Q)/trace(Z)
    Ya = α*(Y₀*Q) .+ mean(X, 1)

    return Ya
end

function Isomap(X,k,d)
    D = pairwise(Euclidean(1e-12),X',X')
    Dmin = Kmin(D,k)
    Dmani = apsp(Dmin)
    Xhat = mds(Dmani,d)
    return Xhat
end

function IOMR(Dpert)
    n = size(Dpert)[1]

    Dnew = copy(Dpert)

    for k = 1:n
        for i = 1:n
            for j = 1:i-1
                if i!=j && i!=k && j!=k
                    if Dnew[i,j] > Dnew[i,k]+Dnew[j,k]
                        Dnew[i,k] = Dnew[i,j]-Dnew[j,k]
                        Dnew[k,i] = Dnew[i,j]-Dnew[j,k]
                    end
                end
            end
        end
    end

    return Dnew-Dpert
end

function sumInf(x)
    n = length(x)
    s = 0.0
    for i = 1:n
        if (x[i] < Inf)
            s += x[i]
        end
    end

    return s
end

function distance(X)
    (n,m) = size(X)
    D = zeros(n,n)
    for i =1:n
        for j = 1:i-1
            b = (X[i,:]-X[j,:]).^2
            D[i,j] = sumInf(b)^0.5
            D[j,i] = (D[i,j])
        end
    end

    return D
end

function MR_Missing(X,Q,k)
    (n,m) = size(Q)
    for i = 1:n
        for j = 1:m
            if Q[i,j] == 0
                X[i,j] = Inf
            end
        end
    end

    Dmiss = distance(X)
    Dfixed = Dmiss + IOMR(Dmiss)
    Dmin = Kmin(Dfixed,k)

    return Dmin
end

end