text:
alphaOnlyMatchRegex = r"^[a-zA-Z0-9_]*$"

"""
$(SIGNATURES)
Returns the transaction for a given query.
NOTE: Must commit(transaction) after you're done.
"""
function _queryNeo4j(neo4jInstance::Neo4jInstance, query::String)
    loadtx = transaction(neo4jInstance.connection)
    result = loadtx(query; submit=true)
    if length(result.errors) > 0
        error(string(result.errors))
    end
    # Have to finish the transaction
    commit(loadtx)
    return result
end

"""
$(SIGNATURES)
Create a node and optionally specify a parent.
Note: Using symbols so that the labels obey Neo4j requirements
"""
function _createNode(neo4jInstance::Neo4jInstance, labels::Vector{String}, properties::Dict{String, Any}, parentNode::Union{Nothing, Neo4j.Node}, relationshipLabel::Symbol=:NOTHING)::Neo4j.Node
    createdNode = Neo4j.createnode(neo4jInstance.graph, properties)
    addnodelabels(createdNode, labels)
    parentNode == nothing && return createdNode
    # Otherwise create the relationship
    createrel(parentNode, createdNode, String(relationshipLabel))
    return createdNode
end

"""
$(SIGNATURES)
Returns the list of CloudGraph nodes that matches the Cyphon query.
The nodes of interest should be labelled 'node' because the query will use the return of id(node)
#Example
matchCondition = '(u:User)-[:ROBOT]->(n:Robot:NewRobot)-[:SESSION]->(node:Session)'.
Will return all sessions because of 'node:Session'.
If orderProperty is not ==, then 'order by n.{orderProperty} will be appended to the query'.
So can make orderProperty = label or id.
"""
function _getLabelsFromCyphonQuery(neo4jInstance::Neo4jInstance, matchCondition::String, orderProperty::String="")::Vector{Symbol}
    query = "match $matchCondition return distinct(node.label) $(orderProperty != "" ? "order by node.$orderProperty" : "")";
    result = _queryNeo4j(neo4jInstance, query)
    nodeIds = map(node -> node["row"][1], result.results[1]["data"])
    return Symbol.(nodeIds)
end

"""
$(SIGNATURES)
Get a node property - returns nothing if not found
"""
function _getNodeProperty(neo4jInstance::Neo4jInstance, nodeLabels::Vector{String}, property::String)
    query = "match (n:$(join(nodeLabels, ":"))) return n.$property"
    result = DistributedFactorGraphs._queryNeo4j(neo4jInstance, query)
    length(result.results[1]["data"]) != 1 && return 0
    length(result.results[1]["data"][1]["row"]) != 1 && return 0
    return result.results[1]["data"][1]["row"][1]
end

"""
$(SIGNATURES)
Set a node property - returns count of changed nodes.
"""
function _setNodeProperty(neo4jInstance::Neo4jInstance, nodeLabels::Vector{String}, property::String, value::String)
    query = "match (n:$(join(nodeLabels, ":"))) set n.$property = \"$value\" return count(n)"
    result = DistributedFactorGraphs._queryNeo4j(neo4jInstance, query)
    length(result.results[1]["data"]) != 1 && return 0
    length(result.results[1]["data"][1]["row"]) != 1 && return 0
    return result.results[1]["data"][1]["row"][1]
end

"""
$(SIGNATURES)
Get a node's tags
"""
function _getNodeTags(neo4jInstance::Neo4jInstance, nodeLabels::Vector{String})::Union{Nothing, Vector{String}}
    query = "match (n:$(join(nodeLabels, ":"))) return labels(n)"
    result = DistributedFactorGraphs._queryNeo4j(neo4jInstance, query)
    length(result.results[1]["data"]) != 1 && return nothing
    return result.results[1]["data"][1]["row"][1]
end

function _getNodeCount(neo4jInstance::Neo4jInstance, nodeLabels::Vector{String})::Int
    query = "match (n:$(join(nodeLabels, ":"))) return count(n)"
    result = DistributedFactorGraphs._queryNeo4j(neo4jInstance, query)
    length(result.results[1]["data"]) != 1 && return 0
    length(result.results[1]["data"][1]["row"]) != 1 && return 0
    return result.results[1]["data"][1]["row"][1]
end

"""
$(SIGNATURES)
Returns the list of CloudGraph nodes that matches the Cyphon query.
The nodes of interest should be labelled 'node' because the query will use the return of id(node)
#Example
matchCondition = '(u:User)-[:ROBOT]->(n:Robot:NewRobot)-[:SESSION]->(node:Session)'.
Will return all sessions because of 'node:Session'.
If orderProperty is not ==, then 'order by n.{orderProperty} will be appended to the query'.
So can make orderProperty = label or id.
"""
function _getNeoNodesFromCyphonQuery(neo4jInstance::Neo4jInstance, matchCondition::String, orderProperty::String="")::Vector{Neo4j.Node}
    # 2. Perform the transaction
    loadtx = transaction(neo4jInstance.connection)
    query = "match $matchCondition return id(node) $(orderProperty != "" ? "order by node.$orderProperty" : "")";
    nodes = loadtx(query; submit=true)
    if length(nodes.errors) > 0
        error(string(nodes.errors))
    end
    # 3. Format the result
    nodes = map(node -> getnode(neo4jInstance.graph, node["row"][1]), nodes.results[1]["data"])
    # Have to finish the transaction
    commit(loadtx)
    return nodes
end

"""
$(SIGNATURES)
Utility function to get a Neo4j node for a robot.
"""
function _getRobotNeoNodeForUser(neo4jInstance::Neo4jInstance, userId::String, robotId::String)::Neo4j.Node
    # 2. Perform the transactionfunction _getLatestPoseForSession(userId::String, robotId::String, sessionId::String)::Neo4j.Node
    nodes = _getNeoNodesFromCyphonQuery(neo4jInstance, "(u:USER:$userId)-[:ROBOT]->(node:ROBOT:$robotId)")
    if(length(nodes) != 1)
        throw("Expected a single robot with ID '$robotId', retrieved $(length(nodes)). Should only be a single node")
    end
    return nodes[1]
end

"""
$(SIGNATURES)
Utility function to get a user root node given an ID.
"""
function _getUserNeoNode(neo4jInstance::Neo4jInstance, userId::String)::Neo4j.Node
    # 2. Perform the transaction
    nodes = _getNeoNodesFromCyphonQuery(neo4jInstance, "(node:USER:$userId)")
    if(length(nodes) != 1)
        error("Expected one user node with labels $userId and USER, received $(length(nodes)).")
    end
    return nodes[1]
end

"""
$(SIGNATURES)
Utility function to get a Neo4j node sessions for a robot.
"""
function _getSessionNeoNodesForRobot(neo4jInstance::Neo4jInstance, userId::String, robotId::String)::Vector{Neo4j.Node}
    # 2. Perform the transaction
    nodes = _getNeoNodesFromCyphonQuery(neo4jInstance, "(u:USER:$userId)-[:ROBOT]->(r:ROBOT:$robotId)-[:SESSION]->(node:SESSION)", "id")
    return nodes
end

"""
$(SIGNATURES)
Bind the SESSION node to the inital variable.
Checks for existence.
"""
function _bindSessionNodeToInitialVariable(neo4jInstance::Neo4jInstance, userId::String, robotId::String, sessionId::String, initialVariableLabel::String)::Nothing
    # 2. Perform the transaction
    loadtx = transaction(neo4jInstance.connection)
    query = """
    match (session:SESSION:$userId:$robotId:$sessionId),
    (var:VARIABLE:$userId:$robotId:$sessionId
    {label: '$initialVariableLabel'})
    WHERE NOT (session)-[:VARIABLE]->()
    CREATE (session)-[:VARIABLE]->(var) return id(var)
    """;
    loadtx(query; submit=true)
    commit(loadtx)
    return nothing
end

# """
# $(SIGNATURES)
# Create a session for a robot.
# """
# function _createSessionForRobot(userId::String, robotId::String, newSession::SessionDetailsRequest)::SessionDetailsResponse
#     # 1. Make sure the session doesn't exist.
#     sessions = _getSessionNeoNodesForRobot(userId, robotId)
#     for session in sessions
#         sessionLabels = getnodelabels(session)
#         if(newSession.id in sessionLabels)
#             error("Session $newSession.id cannot be created for robot $robotId, it already exists")
#         end
#     end
#     # 1b. Get the parent robot Neo4j node.
#     robotNode = _getRobotNeoNodeForUser(userId, robotId)
#
#     # 2. Create the correct response.
#     props = JSON.parse(JSON.json(newSession))
#     props["createdTimestamp"] = string(Dates.now(Dates.UTC))
#     props["lastSolvedTimestamp"] = string(Dates.DateTime(0))
#     props["userId"] = userId
#     props["robotId"] = robotId
#     props["solveCount"] = 1
#     props["solveTimes"] = [0]
#     props["isSolverEnabled"] = true # If false then the pilotfish won't pick up on it.
#
#     # 3. Create the session
#     cloudGraph = Main.App.NaviConfig.cgConnection
#     createdSessionNode = Neo4j.createnode(cloudGraph.neo4j.graph, props)
#     addnodelabels(createdSessionNode, ["SESSION", newSession.id, userId, robotId])
#     createrel(robotNode, createdSessionNode, "SESSION") #From USER to ROBOT
#
#     # 4. Initialize the factor graph
#     # Now start up our factor graph.
#     newSession.initialPoseType
#     if newSession.initialPoseType != nothing
#         labels=[newSession.id, robotId, userId]
#         fg = Caesar.initfg(sessionname=newSession.id, cloudgraph=cloudGraph)
#         # TODO: Cache refactoring discussion.
#         fg.stateless = true # so that we don't update local elements
#
#         poseType = nothing
#         if newSession.initialPoseType != "" && newSession.shouldInitialize
#             poseType = eval(Meta.parse(newSession.initialPoseType))
#             # Initialize the factor graph and insert first pose.
#             nodesyms = initFactorGraph!(fg, labels=labels, ready=0, firstPoseType=poseType)
#             varSymbol  = nodesyms[1]
#             factSymbol = nodesyms[2]
#
#             # Create a relationship for session -> first pose
#             _bindSessionNodeToInitialVariable(userId, robotId, newSession.id, String(varSymbol))
#
#             # Set the node symbols to ready
#             _setReadyStatus(userId, robotId, newSession.id, string.(nodesyms), true)
#         end
#     end
#
#     # 5. Format the result
#     props["nodeCount"] = 0
#     props["userId"] = userId
#     props["robotId"] = robotId
#     props["links"] = Dict{String, String}("self" => "/api/v0/users/$userId/robots/$robotId/sessions/$(newSession.id)", "nodes" => "/api/v0/users/$userId/robots/$robotId/sessions/$(newSession.id)/nodes")
#     model = Unmarshal.unmarshal(SessionDetailsResponse, props)
#     return model
# end

# """
# $(SIGNATURES)
# Function to delete a session.
# """
# function _deleteSession(userId::String, robotId::String, sessionId::String)::Nothing
#     cloudGraph = Main.App.NaviConfig.cgConnection
#     loadtx = transaction(cloudGraph.neo4j.connection)
#     query = "match (n:$userId:$robotId:$sessionId) detach delete n";
#     nodes = loadtx(query; submit=true)
#     commit(loadtx)
#     return nothing
# end

"""
$(SIGNATURES)
Utility function to get a Neo4j node for a session.
"""
function _getSessionNeoNodeForRobot(userId::String, robotId::String, sessionId::String)::Neo4j.Node
    sessions = _getSessionNeoNodesForRobot(userId, robotId)
    for session in sessions
        sessionLabels = getnodelabels(session)
        if(sessionId in sessionLabels)
            return session
        end
    end
    error("Cannot find session $sessionId associated with robot $robotId and user $userId")
end

# """
# $(SIGNATURES)
# Returns all Synchrony nodes for a given session.
# """
# function _getSynchronyNodesForSession(userId::String, robotId::String, sessionId::String)::Vector{NodeModel}
#     # 2. Perform the transaction
#     cloudGraph = Main.App.NaviConfig.cgConnection
#     loadtx = transaction(cloudGraph.neo4j.connection)
#     query = "match (n:$userId:$robotId:$sessionId:VARIABLE) return id(n), n.label, n.MAP_est";
#     nodes = loadtx(query; submit=true)
#     # Have to finish the transaction
#     commit(loadtx)
#
#     if length(nodes.results[1]["data"]) == 0
#         return Vector{NodeModel}() # Empty
#     end
#
#     # 3. Format the result
#     # Convert the mapEsts
#     nodes = map(r -> NodeModel(r["row"][1], r["row"][2], r["row"][3] != nothing ? Float64.(r["row"][3]) : nothing, Dict{String, String}("self" => "/api/v0/users/$userId/robots/$robotId/sessions/$sessionId/nodes/$(r["row"][1])")), nodes.results[1]["data"])
#     return nodes
# end

"""
$(SIGNATURES)
Try get a Neo4j node ID from a node label.
"""
function _tryGetNeoNodeIdFromNodeLabel(neo4jInstance::Neo4jInstance, userId::String, robotId::String, sessionId::String, nodeLabel::Symbol)::Union{Nothing, Int}
    @debug "Looking up symbolic node ID where n.label = '$nodeLabel'..."
    nodes = _getNeoNodesFromCyphonQuery(neo4jInstance, "(node:$userId:$robotId:$sessionId) where exists(node.label) and node.label = \"$(string(nodeLabel))\"")
    if(length(nodes) != 1)
        return nothing
    end
    nodeId = nodes[1].id
    @debug "Found NeoNode ID = $nodeId..."
    return nodeId
end

"""
$(SIGNATURES)
Get a Neo4j node ID from a node label.
"""
function _getNeoNodeIdFromNodeLabel(neo4jInstance::Neo4jInstance, userId::String, robotId::String, sessionId::String, nodeLabel::Symbol)::Int
    nodeId = _tryGetNeoNodeIdFromNodeLabel(neo4jInstance, userId, robotId, sessionId, nodeLabel)
    if(nodeId == nothing)
        error("Cannot find a single symbolic node for label '$nodeLabel', received $(length(nodes)) nodes...")
    end
    return nodeId
end

"""
$(SIGNATURES)
Get a specific node given the Neo4j node ID.
"""
function _getSynchronyNodeFromId(userId::String, robotId::String, sessionId::String, nodeId::Union{Int, String})::NodeDetails
    cloudGraph = Main.App.NaviConfig.cgConnection

    if(typeof(nodeId) == String)
        nodeId = _getNeoNodeIdFromNodeLabel(userId, robotId, sessionId, nodeId)
    end

    node = CloudGraphs.get_vertex(cloudGraph, nodeId, false)

    # 3. Format the result
    rootPath = "/api/v0/users/$userId/robots/$robotId/sessions/$sessionId/nodes/$nodeId"
    node = NodeDetails(
        nodeId,
        haskey(node.properties, "label") ? node.properties["label"] : "",
        node.exVertexId,
        :softtype in fieldnames(typeof(node.packed)) ? string(typeof(node.packed.softtype)) : "",
        node.properties,
        node.packed,
        node.labels,
        Dict{String, String}("self" => rootPath, "data" => rootPath*"/data"))
    if !(robotId in node.labels) ||
        !(sessionId in node.labels) ||
        !(userId in node.labels)
        # Making sure user is asking for their session nodes.
        error("Node not found")
    end
    return(node)
end

"""
$(SIGNATURES)
Sets the ready status of the complete session.
"""
function _setSessionReadyStatus(userId::String, robotId::String, sessionId::String, ready::Bool)::Nothing
    # 2. Perform the transaction
    cloudGraph = Main.App.NaviConfig.cgConnection
    loadtx = transaction(cloudGraph.neo4j.connection)
    query = "match (n:$userId:$robotId:$sessionId) where exists(n.ready) set n.ready=$(ready ? 1 : 0) return count(n)";
    loadtx(query; submit=true)
    commit(loadtx)
    return nothing
end

"""
$(SIGNATURES)
Convert Vector{<:AbstractString} to square braced single quoted list of elements ready for Cypher query use"
"""
function _stringArrToCyphonList(stra::Vector{<:AbstractString})
    qs = "["
    for el in stra
        qs *= "'"*el*"', "
    end
    qs = chop(chop(qs))
    qs *= "]"
    return qs
end

"""
$(SIGNATURES)
Returns the BigDataElements for the node.
"""
function _getDataEntriesForNode(nodeId::Int)::Vector{BigDataElement}
    cloudGraph = Main.App.NaviConfig.cgConnection
    node = CloudGraphs.get_vertex(cloudGraph, nodeId, false)
    return node.bigData.dataElements
end

"""
$(SIGNATURES)
Sets the ready status of the complete session.
"""
function _setReadyStatus(userId::String, robotId::String, sessionId::String, nodeLabels::Vector{String}, ready::Bool)::Nothing
    # 2. Perform the transaction
    cloudGraph = Main.App.NaviConfig.cgConnection

    loadtx = transaction(cloudGraph.neo4j.connection)
    query = "match (n:$userId:$robotId:$sessionId) where exists(n.ready) and exists(n.label) and n.label in $(_stringArrToCyphonList(nodeLabels)) set n.ready=$(ready ? 1 : 0) return count(n)";
    loadtx(query; submit=true)
    commit(loadtx)
    return nothing
end

"""
$(SIGNATURES)
Validate all given inputs and return the results.
Error if any missing or blank.
Returns: A list of strings/ints/floats of all the inputs.
"""
function _validateHttpInputs(requiredInputs::Vector{Symbol}, params::Dict{Symbol, Any})::Vector{Any}
    rets = Vector{Any}()
    for input in requiredInputs
        input = haskey(params, input) ? params[input] : ""
        if(typeof(input) == String && isempty(lstrip(input)))
            error("'$(string(input))' is empty string.")
        end
        typeof(input) == String && !occursin(alphaOnlyMatchRegex, input) && error("$input is not alphanumeric, please only use alphanumeric and underscores for parameters.")
        push!(rets, input)
    end
    return rets
end

function readAndReturnFile(filename::String, mimeType::String="application/octet-stream")::HTTP.Response
    if !isfile(filename)
        error("The file to be returned does not exist - $filename")
    end

    s = open(filename, "r");
    data = read(s);
    close(s);

    headers = Dict{AbstractString, AbstractString}( "Server" => "Julia/$VERSION",
                    "Content-Type"     => mimeType,
                    "Date"             => Dates.format(now(Dates.UTC), Dates.RFC1123Format) )
    return HTTP.Response(200, headers; body=data)
end

function jsonReponse(jsonData::String; returnCode::Int=200)::HTTP.Response
    headers = Dict{AbstractString, AbstractString}( "Server" => "Julia/$VERSION",
                    "Content-Type"     => "application/json",
                    "Date"             => Dates.format(now(Dates.UTC), Dates.RFC1123Format) )
    return HTTP.Response(returnCode, headers; body=jsonData)
end

"""
Try to gracefully handle errors.
"""
function handleError(ex::Any)::HTTP.Response
    @show ex
    io = IOBuffer()
    showerror(io, ex, catch_backtrace())
    err = String(take!(io))
    @error "Error! $err"
    msg = (:msg in fieldnames(typeof(ex)) ? ex.msg : err)
    return jsonReponse(msg; returnCode=400)
end

"""
Push a graph change notification.
"""
function _addGraphChangeNotification(userId::String, robotId::String, sessionId::String, isWholeGraph::Bool)::Bool
    conn = Main.App.NaviConfig.redisConnection
    if !Redis.is_connected(conn)
        # Reinvigorate connection
        @warn "[AQP] Redis connection failed - reestablishing a connection!"
        conn = Redis.RedisConnection(host=conn.host, port=conn.port, db=conn.db, password=conn.password)
        Main.App.NaviConfig.redisConnection = conn
    end

    # 2. Perform the transaction
    @info "Pushing graph change notification..."
    gc = Main.App.NaviConfig.GraphChange(userId, robotId, sessionId, isWholeGraph)
    Main.App.NaviConfig.addGraphChange(conn, gc)
    return true
end


text:
import Base.:+
import Base.:-
import Base.:*
import Base.:/
import Base.:(==)
import Base.:<

export U1
export arg
export len

abstract type Group end

"""
    U1 <: Group

The symmetry group U(1), AKA the circle group, {z ∈ ℂ : |z| = 1}.

fields: θ.
"""
struct U1 <: Group
    θ::Float64
end

"""
    show(x)

Print a string representation of `x`.
"""
Base.show(io::IO, x::U1) = print(io, "eⁱᶿ : θ = $(x.θ)")

"""
    arg(x)

The phase or argument of `x` = eⁱᶿ.
"""
arg(x::U1) = x.θ

"""
    len(x)

The magnitude or length of `x` = eⁱᶿ.
"""
len(x::U1) = 1

## Unary Operators ##

"""
    +x

Unary plus, the identity operation on `x`.
"""
+(x::U1) = U1(+x.θ)

"""
    -x

Unary minus, map `x` to its additive inverse.
"""
-(x::U1) = U1(π - x.θ)

## Binary Operators ##

"""
    x * y

Multiply `x` times `y`.
"""
*(x::U1, y::U1) = U1(x.θ + y.θ)

"""
    x / y

Divide `x` over `y`.
"""
/(x::U1, y::U1) = U1(x.θ - y.θ)

## Numeric Comparisons ##

"""
    x == y

Check whether `x` and `y` are equal.
"""
==(x::U1, y::U1) = isapprox(x.θ, y.θ, atol=TOLERANCE)

"""
    x < y

Check whether `x` is less than `y`.
"""
<(x::U1, y::U1) = x.θ < y.θ

ext: jl


text:
abstract type AbstractRLEnvMDP{S, A} <: MDP{S, A} end
abstract type AbstractRLEnvPOMDP{S, A, O} <: POMDP{S, A, O} end

const AbstractRLEnvProblem = Union{AbstractRLEnvMDP, AbstractRLEnvPOMDP}

POMDPs.actions(m::AbstractRLEnvProblem) = RL.actions(m.env)
POMDPs.discount(m::AbstractRLEnvProblem) = m.discount

function POMDPs.actions(m::AbstractRLEnvProblem, s)
    if provided(RL.valid_actions, m.env)
        old = RL.state(m.env)
        RL.setstate!(m.env, s)
        va = RL.valid_actions(m.env)
        RL.setstate!(m.env, old)
        return va
    else
        return RL.actions(m.env)
    end
end

function POMDPs.initialstate(m::AbstractRLEnvProblem)
    return ImplicitDistribution(m) do m, rng
        RL.reset!(m.env)
        return RL.state(m.env)
    end
end

function POMDPs.isterminal(m::AbstractRLEnvProblem, s)
    old = RL.state(m.env)
    RL.setstate!(m.env, s)
    t = RL.terminated(m.env)
    RL.setstate!(m.env, old)
    return t
end

struct RLEnvMDP{E, S, A} <: AbstractRLEnvMDP{S, A}
    env::E
    discount::Float64
end

"""
    RLEnvMDP(env; discount=1.0)

Create an `MDP` by wrapping a `CommonRLInterface.AbstractEnv`. `state` and `setstate!` from `CommonRLInterface` must be provided, and the `POMDPs` generative model functionality will be provided.
"""
function RLEnvMDP(env; discount=1.0)
    S = infer_statetype(env)
    if S == Any
        @warn("State type inferred for $(typeof(env)) by looking at the return type of state(env) was Any. This could cause significant performance degradation.")
    end
    return RLEnvMDP{typeof(env), S, eltype(RL.actions(env))}(env, discount)
end

function POMDPs.gen(m::RLEnvMDP, s, a, rng)
    # rng is not currently used
    RL.setstate!(m.env, s)
    r = RL.act!(m.env, a)
    sp = RL.state(m.env)
    return (sp=sp, r=r)
end

"""
    RLEnvPOMDP(env; discount=1.0)

Create an `POMDP` by wrapping a `CommonRLInterface.AbstractEnv`. `state` and `setstate!` from `CommonRLInterface` must be provided, and the `POMDPs` generative model functionality will be provided.
"""
struct RLEnvPOMDP{E, S, A, O} <: AbstractRLEnvPOMDP{S, A, O}
    env::E
    discount::Float64
end

function RLEnvPOMDP(env; discount=1.0)
    S = infer_statetype(env)
    if S == Any
        @warn("State type inferred for $(typeof(env)) by looking at the return type of state(env) was Any. This could cause significant performance degradation.")
    end
    O = infer_obstype(env)
    if S == Any
        @warn("Observation type inferred for $(typeof(env)) by looking at the return type of observe(env) was Any. This could cause significant performance degradation.")
    end
    return RLEnvPOMDP{typeof(env), S, eltype(RL.actions(env)), O}(env, discount)
end

function POMDPs.gen(m::RLEnvPOMDP, s, a, rng)
    RL.setstate!(m.env, s)
    r = RL.act!(m.env, a)
    sp = RL.state(m.env)
    o = RL.observe(m.env)
    return (sp=sp, o=o, r=r)
end

#####################
# Opaque: for when state and setstate! are not implemented
#####################

struct OpaqueRLEnvState
    age::BigInt
end

mutable struct OpaqueRLEnvMDP{E, A} <: AbstractRLEnvMDP{OpaqueRLEnvState, A}
    env::E
    age::BigInt
    discount::Float64
end

"""
    OpaqueRLEnvMDP(env; discount=1.0)

Wrap a `CommonRLInterface.AbstractEnv` in an `MDP` object. The state will be an `OpaqueRLEnvState` and only simulation will be supported.
"""
function OpaqueRLEnvMDP(env; discount::Float64=1.0)
    return OpaqueRLEnvMDP{typeof(env), eltype(RL.actions(env))}(env, 1, discount)
end

mutable struct OpaqueRLEnvPOMDP{E, A, O} <: AbstractRLEnvPOMDP{OpaqueRLEnvState, A, O}
    env::E
    age::BigInt
    discount::Float64
end

"""
    OpaqueRLEnvPOMDP(env; discount=1.0)

Wrap a `CommonRLInterface.AbstractEnv` in an `POMDP` object. The state will be an `OpaqueRLEnvState` and only simulation will be supported.
"""
function OpaqueRLEnvPOMDP(env, discount=1.0)
    return OpaqueRLEnvPOMDP{typeof(env), eltype(RL.actions(env)), typeof(RL.observe(env))}(env, 1, discount)
end

const OpaqueRLEnvProblem = Union{OpaqueRLEnvMDP, OpaqueRLEnvPOMDP}

function POMDPs.actions(m::OpaqueRLEnvProblem, s::OpaqueRLEnvState)
    if provided(RL.valid_actions, m.env) && s.age == m.age
        return RL.valid_actions(m.env)
    else
        return RL.actions(m.env)
    end
end

function POMDPs.initialstate(m::OpaqueRLEnvProblem)
    return ImplicitDistribution(m) do m, rng
        RL.reset!(m.env)
        m.age += 1
        return OpaqueRLEnvState(m.age)
    end
end

function POMDPs.isterminal(m::OpaqueRLEnvProblem, s::OpaqueRLEnvState)
    if s.age != m.age
        throw(OpaqueRLEnvState(m, s))
    end
    return RL.terminated(m.env)
end

function POMDPs.gen(m::OpaqueRLEnvMDP, s::OpaqueRLEnvState, a, rng)
    if s.age == m.age
        r = RL.act!(m.env, a)
        m.age += 1
        return (sp=OpaqueRLEnvState(m.age), r=r)
    else
        throw(OpaqueRLEnvState(m.env, m.age, s))
    end
end

function POMDPs.gen(m::OpaqueRLEnvPOMDP, s::OpaqueRLEnvState, a, rng)
    if s.age == m.age
        r = RL.act!(m.env, a)
        o = RL.observe(m.env)
        m.age += 1
        return (sp=OpaqueRLEnvState(m.age), o=o, r=r)
    else
        throw(OpaqueRLEnvState(m.env, m.age, s))
    end
end

function Base.convert(::Type{POMDP}, env::RL.AbstractEnv)
     if RL.provided(RL.state, env)
        s = RL.state(env)
        if RL.provided(RL.setstate!, env, s)
            return RLEnvPOMDP(env)
        end
    end
    return OpaqueRLEnvPOMDP(env)
end

function Base.convert(::Type{MDP}, env::RL.AbstractEnv)
    if RL.provided(RL.state, env)
        s = RL.state(env)
        if RL.provided(RL.setstate!, env, s)
            return RLEnvMDP(env)
        end
    end
    return OpaqueRLEnvMDP(env)
end

Base.convert(E::Type{<:RL.AbstractEnv}, m::AbstractRLEnvProblem) = convert(E, m.env)
Base.convert(::Type{RL.AbstractEnv}, m::AbstractRLEnvProblem) = m.env

struct OpaqueRLEnvStateError <: Exception
    env
    env_age::BigInt
    s::OpaqueRLEnvState
end

function Base.showerror(io::IO, e::OpaqueRLEnvStateError)
    print(io, "OpaqueRLEnvStateError: ")
    print(io, """An attempt was made to interact with the environment encapsulated in an `OpaqueRLEnv(PO)MDP` at a particular state, but the environment had been stepped forward, so it may be in a different state.

              Enironment age: $(e.env_age)
              State age: $(e.s.age)

              Suggestions: provide `CommonRLInterface.state(::$(typeof(e.env)))` and `CommonRLInterface.setstate!(::$(typeof(e.env)), s)` so that the environment can be converted to a `RLEnv(PO)MDP` instead of an `OpaqueRLEnv(PO)MDP`.
              """)
end

function infer_statetype(env)
    try
        return only(Base.return_types(RL.state, Tuple{typeof(env)}))
    catch ex
        @warn("""Unable to infer state type for $(typeof(env)) because of the following error:

              $(sprint(showerror, ex))

              Falling back to Any.
              """)
        return Any
    end
end

function infer_obstype(env)
    try
        return only(Base.return_types(RL.observe, Tuple{typeof(env)}))
    catch ex
        @warn("""Unable to infer observation type for $(typeof(env)) because of the following error:

              $(sprint(showerror, ex))

              Falling back to Any.
              """)
        return Any
    end
end


text:
#! /usr/bin/julia

# Rosetta Code, Arithmetic-geometric mean/Calculate Pi

function agm_step{T<:FloatingPoint}(x::T, y::T)
    (0.5*(x + y), sqrt(x*y))
end

function approx_pi_step{T<:FloatingPoint,U<:Integer}(x::T, y::T, z::T, n::U)
    (a, g) = agm_step(x, y)
    k = n + 1
    s = z + 2^(k+1)*(a^2 - g^2)
    return (a, g, s, k)
end

function approx_pi{T<:FloatingPoint}(a::T, g::T, s::T)
    4a^2/(1 - s)
end

prec = 512
set_bigfloat_precision(prec)
println("Approximating pi using ", prec, "-bit floats.")
println("   k     Error  Result")
a = big(1.0)
g = a/sqrt(big(2.0))
s = big(0.0)
k = 0
oldagpi = big(0.0)
for i in 1:100
    (a, g, s, k) = approx_pi_step(a, g, s, k)
    agpi = approx_pi(a, g, s)
    2eps(agpi) < abs(agpi-oldagpi) || break
    oldagpi = agpi
    err = pi - agpi
    print(@sprintf("  %2d ", i))
    print(@sprintf(" %9.1e", err))
    println(@sprintf(" %.60e", agpi))
end

max_stars_repo_name: MichaeLeroy/rosetta-code

max_issues_repo_name: MichaeLeroy/rosetta-code

max_forks_repo_name: MichaeLeroy/rosetta-code


text:
## install package for Gaia.Clustering

using Pkg

Pkg.add("PyCall")
Pkg.add("DataFrames")
Pkg.add("Statistics")
Pkg.add("Distributions")
Pkg.add("Random")
Pkg.add("CSV")
Pkg.add("PyPlot")
Pkg.add("Query")
Pkg.add("LsqFit")
Pkg.add("FFTW")
Pkg.add("StatsBase")
Pkg.add("Distances")
Pkg.add("JLD")
Pkg.build("HDF5")

println("## Package installation for GASS done...")

max_stars_repo_name: bosscha/GASS

max_issues_repo_name: bosscha/GASS

max_forks_repo_name: bosscha/GASS


text:
module VrpAnt

import Base.empty, Random.rand

# include("Way.jl")
# include("Solution.jl")
# include("Model.jl")
# include("Ant.jl")

# true if it has to stop, otherwise return false
function stopCondition!(model::Model, antList::Vector{Ant}, ind::Int64, stillSame::Int64, NCmax::Int64, stillSameMax::Int64)
    if ind > NCmax
        model.solution.state = NCEnding()
        return true
    elseif ind == 0
        return false
    else
        if stillSame >= stillSameMax
            model.solution.state = NoChangeEnding()
            return true
        elseif reduce( & , [ant.lengthMade ≈ antList[1].lengthMade for ant in antList])
            model.solution.state = AntEnding()
            return true
        else
            return false
        end
    end
end

function updateSolution!(model::Model, stillSame::Int64, antList::Vector{Ant})
    bestAnt::Ant = searchBestAnt(antList)

    if bestAnt.lengthMade < model.solution.length
        model.solution.path = bestAnt.way
        model.solution.length = bestAnt.lengthMade
        return 0
    else
        return stillSame+1
    end
end

function optimize(map::Map; m::Union{Nothing, Int64} = nothing, p::Float64 = 0.5, α::Real = 3, β::Real = 1, Q::Real = 100, NCmax::Int64 = 5000, stillSameMax::Int64 = 50, pheroInit::Float64 = 1.)
    model::Model = Model(map, pheroInit)
    m == nothing ? _optimize!(model, length(map.cities), p, α, β, Q, NCmax, stillSameMax) : _optimize!(model, m, p, α, β, Q, NCmax, stillSameMax)
    return model
end

function optimize!(model::Model; m::Union{Nothing, Int64} = nothing, p::Float64 = 0.5, α::Real = 3, β::Real = 1, Q::Real = 100, NCmax::Int64 = 5000, stillSameMax::Int64 = 50, pheroInit::Float64 = 1.)
    m == nothing ? _optimize!(model, length(map.cities), p, α, β, Q, NCmax, stillSameMax) : _optimize!(model, m, p, α, β, Q, NCmax, stillSameMax)
    return model
end

#Change about the model !!!!!!!!!

function _optimize!(model::Model, m::Int64, p::Float64, α::Real, β::Real, Q::Real, NCmax::Int64, stillSameMax::Int64)
    antList::Vector{Ant} = [Ant(model.map.cities) for loop = 1:m]
    ind::Int64 = 0
    stillSame::Int64 = 0

    while !stopCondition!(model, antList, ind, NCmax, stillSame, stillSameMax)

        if ind != 0
            sort!(antList, by = ant -> ant.lengthMade)

            for ant in antList[1:Int(ceil(0.1*m))]
                wayBack!(model, ant, Q)
            end
        end
        updatePhero!(model, p)
        for ant in antList
            round!(ant, model, α, β)
        end

        ind+=1
        stillSame = updateSolution!(model, stillSame, antList)
    end

    return model
end

export optimize!, Map, display, optimize, Model

end  # module


text:
# The 6D simplex
for N in 1:3:51
    facts("Simplex in $N dimensions") do
        N = 6
        m = Model()
        @variable(m, x[1:N] >= 0)
        @constraint(m, sum{x[i], i=1:N} == 1)

        v, r = get_extrema(m)

        @fact length(v) --> N
        @fact length(r) --> 0

        for k in 1:N
            ex = zeros(N)
            ex[k] = 1
            @fact ex in v --> true
        end
    end
end

# The 6D simplex, plus the origin
for N in 1:3:39
    facts("Simplex with the origin in $N dimensions") do
        m = Model()
        @variable(m, x[1:N] >= 0)
        @constraint(m, sum{x[i], i=1:N} ≤ 1)

        v, r = get_extrema(m)

        @fact length(v) --> N+1
        @fact length(r) --> 0

        @fact zeros(N) in v --> true
        for k in 1:N
            ex = zeros(N)
            ex[k] = 1
            @fact ex in v --> true
        end
    end
end


text:
# # Basics
#
# ## Packages
#
# We load the following packages into our environment:

using OptimalTransport
using Distances
using Plots
using PythonOT: PythonOT
using Tulip

using LinearAlgebra
using Random

Random.seed!(1234)

const POT = PythonOT

# ## Problem setup
#
# First, let us initialise two random probability measures $\mu$ (source measure)
# and $\nu$ (target measure) in 1D:

M = 200
μ = fill(1 / M, M)
μsupport = rand(M)

N = 250
ν = fill(1 / N, N)
νsupport = rand(N);

# Now we compute the quadratic cost matrix $C_{ij} = \| x_i - x_j \|_2^2$:

C = pairwise(SqEuclidean(), μsupport', νsupport'; dims=2);

# ## Exact optimal transport
#
# The **earth mover's distance** is defined as the optimal value of the
# Monge-Kantorovich problem
# ```math
# \inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle =
# \inf_{\gamma \in \Pi(\mu, \nu)} \sum_{i, j} \gamma_{ij} C_{ij},
# ```
# where $\Pi(\mu, \nu)$ denotes the set of couplings of $\mu$ and $\nu$, i.e., the set of joint
# distributions whose marginals are $\mu$ and $\nu$.
# If $C$ is the quadratic cost matrix, the earth mover's distance is known as the square of
# the **2-Wasserstein distance**.
#
# The function `emd` returns the optimal transport *plan* $\gamma$

γ = emd(μ, ν, C, Tulip.Optimizer());

# whilst using `emd2` returns the optimal transport *cost*:

emd2(μ, ν, C, Tulip.Optimizer())

# ## Entropically regularised optimal transport
#
# We may add an entropy term to the Monge-Kantorovich problem to obtain the
# **entropically regularised** optimal transport problem
# ```math
# \inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle + \epsilon \Omega(\gamma),
# ```
# where $\Omega(\gamma) = \sum_{i, j} \gamma_{ij} \log(\gamma_{ij})$ is the negative
# entropy of the coupling $\gamma$ and $\epsilon$ controls the strength of the regularisation.
#
# This problem is *strictly convex* and admits a very efficient iterative scaling scheme for
# its solution known as the [Sinkhorn algorithm](https://doi.org/10.1214/aoms/1177703591).
#
# We compute the optimal entropically regularised transport plan:

ε = 0.01
γ = sinkhorn(μ, ν, C, ε);

# We can check that one obtains the same result with the Python Optimal Transport (POT) package:

γpot = POT.sinkhorn(μ, ν, C, ε)
norm(γ - γpot, Inf)

# We can compute the optimal cost (a scalar) of the entropically regularized optimal
# transport problem with `sinkhorn2`:

sinkhorn2(μ, ν, C, ε)

# ## Quadratically regularised optimal transport
#
# Instead of the common entropically regularised optimal transport problem, we can solve the
# [quadratically regularised optimal transport problem](https://doi.org/10.1007/s00245-019-09614-w)
# ```math
# \inf_{\gamma \in \Pi(\mu, \nu)} \langle \gamma, C \rangle + \epsilon \frac{\| \gamma \|_F^2}{2}.
# ```
# One property of the quadratically regularised optimal transport problem is that the
# resulting transport plan $\gamma$ is *sparse*. We take advantage of this and represent it as
# a sparse matrix.

quadreg(μ, ν, C, ε; maxiter=100);

# ## Stabilized Sinkhorn algorithm
#
# When $\epsilon$ is very small, we can use a
# [log-stabilised version of the Sinkhorn algorithm](https://doi.org/10.1137/16M1106018).

ε = 0.005
γ = sinkhorn_stabilized(μ, ν, C, ε; maxiter=5_000);

# Again we can check that the same result is obtained with the POT package:

γ_pot = POT.sinkhorn(μ, ν, C, ε; method="sinkhorn_stabilized", numItermax=5_000)
norm(γ - γ_pot, Inf)

# ## Stabilized Sinkhorn algorithm with $\epsilon$-scaling
#
# In addition to log-stabilisation, we can use [$\epsilon$-scaling](https://doi.org/10.1137/16M1106018):

γ = sinkhorn_stabilized_epsscaling(μ, ν, C, ε; maxiter=5_000);

# The POT package yields the same result:

γpot = POT.sinkhorn(μ, ν, C, ε; method="sinkhorn_epsilon_scaling", numItermax=5000)
norm(γ - γpot, Inf)

# ## Unbalanced optimal transport
#
# [Unbalanced optimal transport](https://doi.org/10.1090/mcom/3303) deals with general
# positive measures which do not necessarily have the same total mass. For unbalanced
# source and target marginals $\mu$ and $\nu$ and a cost matrix $C$, entropically
# regularised unbalanced optimal transport solves
# ```math
# \inf_{\gamma \geq 0} \langle \gamma, C \rangle + \epsilon \Omega(\gamma) + \lambda_1 \mathrm{KL}(\gamma 1 | \mu) + \lambda_2 \mathrm{KL}(\gamma^{\mathsf{T}} 1 | \nu),
# ```
# where $\epsilon$ controls the strength of the entropic regularisation, and $\lambda_1$ and
# $\lambda_2$ control how strongly we enforce the marginal constraints.
#
# We construct two random measures, now with different total masses:

M = 100
μ = fill(1 / M, M)
μsupport = rand(M)

N = 200
ν = fill(1 / M, N)
νsupport = rand(N);

# We compute the quadratic cost matrix:

C = pairwise(SqEuclidean(), μsupport', νsupport'; dims=2);

# Now we solve the corresponding unbalanced, entropy-regularised transport problem with
# $\epsilon = 0.01$ and $\lambda_1 = \lambda_2 = 1$:

ε = 0.01
λ = 1
γ = sinkhorn_unbalanced(μ, ν, C, λ, λ, ε);

# We check that the result agrees with POT:

γpot = POT.sinkhorn_unbalanced(μ, ν, C, ε, λ)
norm(γ - γpot, Inf)

# ## Plots
#
# ### Entropically regularised transport
#
# Let us construct source and target measures again:

μsupport = νsupport = range(-2, 2; length=100)
C = pairwise(SqEuclidean(), μsupport', νsupport'; dims=2)
μ = normalize!(exp.(-μsupport .^ 2 ./ 0.5^2), 1)
ν = normalize!(νsupport .^ 2 .* exp.(-νsupport .^ 2 ./ 0.5^2), 1)

plot(μsupport, μ; label=raw"$\mu$", size=(600, 400))
plot!(νsupport, ν; label=raw"$\nu$")

# Now we compute the entropically regularised transport plan:

γ = sinkhorn(μ, ν, C, 0.01)
heatmap(
    μsupport,
    νsupport,
    γ;
    title="Entropically regularised optimal transport",
    size=(600, 600),
)

# ### Quadratically regularised transport
#
# Notice how the "edges" of the transport plan are sharper if we use quadratic regularisation
# instead of entropic regularisation:

γquad = quadreg(μ, ν, C, 5; maxiter=100);
heatmap(
    μsupport,
    νsupport,
    γquad;
    title="Quadratically regularised optimal transport",
    size=(600, 600),
)

# ### Sinkhorn barycenters
#
# For a collection of discrete probability measures $\{\mu_i\}_{i=1}^N \subset \mathcal{P}$,
# cost matrices $\{C_i\}_{i=1}^N$, and positive weights $\{\lambda_i\}_{i=1}^N$ summing to $1$,
# the entropically regularised **barycenter** in $\mathcal{P}$ is the discrete probability
# measure $\mu$ that solves
# ```math
# \inf_{\mu \in \mathcal{P}} \sum_{i = 1}^N \lambda_i \operatorname{OT}_{\epsilon}(\mu, \mu_i)
# ```
# where $\operatorname{OT}_\epsilon(\mu, \mu_i)$ denotes the entropically regularised
# optimal transport cost with marginals $\mu$ and $\mu_i$, cost matrix $C$, and entropic
# regularisation parameter $\epsilon$.
#
# We set up two measures and compute the weighted barycenters. We choose weights
# $\lambda_1 \in \{0.25, 0.5, 0.75\}$.

support = range(-1, 1; length=250)
mu1 = normalize!(exp.(-(support .+ 0.5) .^ 2 ./ 0.1^2), 1)
mu2 = normalize!(exp.(-(support .- 0.5) .^ 2 ./ 0.1^2), 1)

plt = plot(; size=(800, 400), legend=:outertopright)
plot!(plt, support, mu1; label=raw"$\mu_1$")
plot!(plt, support, mu2; label=raw"$\mu_2$")

mu = hcat(mu1, mu2)
C = pairwise(SqEuclidean(), support'; dims=2)
for λ1 in (0.25, 0.5, 0.75)
    λ2 = 1 - λ1
    a = sinkhorn_barycenter(mu, C, 0.01, [λ1, λ2], SinkhornGibbs())
    plot!(plt, support, a; label="\$\\mu \\quad (\\lambda_1 = $λ1)\$")
end
plt


text:
@testset "UpperConfidneceBound" begin
    s = UCBSelector(2)
    values = [1, 2]
    @test s(values) == 2
    @test s(values) == 1
    @test s(values) == 2
    @test s(values) == 2
    @test s(values) == 1
end


text:
function CelestialBodyPointed(arg0::Frame, arg1::PVCoordinatesProvider, arg2::Vector3D, arg3::Vector3D, arg4::Vector3D)
    return CelestialBodyPointed((Frame, PVCoordinatesProvider, Vector3D, Vector3D, Vector3D), arg0, arg1, arg2, arg3, arg4)
end

function get_attitude(obj::CelestialBodyPointed, arg0::FieldPVCoordinatesProvider, arg1::FieldAbsoluteDate, arg2::Frame)
    return jcall(obj, "getAttitude", FieldAttitude, (FieldPVCoordinatesProvider, FieldAbsoluteDate, Frame), arg0, arg1, arg2)
end

function get_attitude(obj::CelestialBodyPointed, arg0::PVCoordinatesProvider, arg1::AbsoluteDate, arg2::Frame)
    return jcall(obj, "getAttitude", Attitude, (PVCoordinatesProvider, AbsoluteDate, Frame), arg0, arg1, arg2)
end


text:
branin(x::Vector; kwargs...) = branin(x[1], x[2]; kwargs...)
branin(x1, x2; a = 1, b = 5.1/(4π^2), c = 5/π, r = 6, s = 10, t = 1/(8π),
       noiselevel = 0) =
    a * (x2 - b*x1^2 + c*x1 - r)^2 + s*(1 - t)*cos(x1) + s + noiselevel * randn()

minima(::typeof(branin)) = [[-π, 12.275], [π, 2.275], [9.42478, 2.475]], 0.397887

euclidean(x, y) = √sum((x .- y).^2)
function regret(opt, func)
    mins, fmin = minima(func)
    (observed_dist = minimum(map(m -> euclidean(m, opt.observed_optimizer), mins)),
     observed_regret = abs(opt.observed_optimum - fmin),
     model_dist = minimum(map(m -> euclidean(m, opt.model_optimizer), mins)),
     model_regret = abs(opt.model_optimum - fmin))
end

@testset "branin" begin
    Random.seed!(123)
    for ac in [ProbabilityOfImprovement(), ExpectedImprovement(),
               UpperConfidenceBound(), ThompsonSamplingSimple(), MutualInformation()]
        println("testing on branin with $ac")
        opt = BOpt(x -> branin(x, noiselevel = 0),
                   ElasticGPE(2, mean = MeanConst(-10.),
                              kernel = SEArd([0., 0.], 5.),
                              logNoise = -2., capacity = 3000),
                   ac,
                   MAPGPOptimizer(every = 50, noisebounds = [-4, 3],
                                 kernbounds = [[-1, -1, 0], [4, 4, 10]],
                                 maxeval = 40),
                   [-5., 0.], [10., 15.], maxiterations = 200,
                   sense = Min,
                   verbosity = Silent)
        boptimize!(opt)
        reg = regret(opt, branin)
        @test reg.observed_regret < 0.05
    end
end


text:
using Documenter
using Land

# define default docs pages
pages = Any[
    "Home" => "index.md",
    "Tips" => "tips.md" ,
    "APIs" => [
              "CanopyLayers"    => "submodules/CanopyLayers.md"         ,
              "Photosynthesis"  => "submodules/Photosynthesis.md"       ,
              "PlantHydraulics" => "submodules/PlantHydraulics.md"      ,
              "StomataModels"   => "submodules/StomataModels.md"        ,
              "SPAC"            => "submodules/SoilPlantAirContinuum.md",
              "Land"            => "submodules/Land.md"                 ,
    ],
]

# add example pages
generated_dir = joinpath(@__DIR__, "src", "generated")
rm(generated_dir, force = true, recursive = true)
mkpath(generated_dir)
include("examples.jl")

# format the docs
mathengine = MathJax(Dict(
    :TeX => Dict(
        :equationNumbers => Dict(:autoNumber => "AMS"),
        :Macros => Dict(),
    ),
))

format = Documenter.HTML(
    prettyurls = get(ENV, "CI", nothing) == "true",
    mathengine = mathengine,
    collapselevel = 1,
    assets = ["assets/favicon.ico"]
)

# build the docs
makedocs(
    sitename = "Land",
    format = format,
    clean = false,
    modules = [Land],
    pages = pages,
)

# deploy the docs to Github gh-pages
deploydocs(
    repo = "github.com/CliMA/Land.git",
    target = "build",
    devbranch = "main",
    push_preview = true,
)

max_stars_repo_name: bbuman/Land

max_issues_repo_name: bbuman/Land

max_forks_repo_name: bbuman/Land


text:
# This script actually solves the WachterDisasterRisk model with a risk-adjusted linearization
# and times the methods, if desired
using BenchmarkTools, RiskAdjustedLinearizations, Test
include(joinpath(dirname(@__FILE__), "..", "..", "examples", "rbc_cc", "rbc_cc.jl"))

# Set up
m_rbc_cc = RBCCampbellCochraneHabits()
m = rbc_cc(m_rbc_cc, 0)

# Solve!
solve!(m; algorithm = :relaxation)

sssout = JLD2.jldopen(joinpath(dirname(@__FILE__), "..", "reference", "rbccc_sss_iterative_output.jld2"), "r")

@test isapprox(sssout["z_rss"], m.z, atol=1e-4)
@test isapprox(sssout["y_rss"], m.y, atol=1e-4)
@test isapprox(sssout["Psi_rss"], m.Ψ, atol=1e-4)


text:
theta0=rand(2)
dirm,sirm,dfrm,sfrm=classicaltoys()

data=rand(dirm,50)
llk=loglikelihood(dirm,data,theta0)

tol=1e-5
for model in (dirm,sirm,dfrm,sfrm)
	coef!(model,theta0)
	@test checkdp(model) <tol
	@test norm(llk-loglikelihood(model,data)) <tol
end


text:
using Printf
using LinearAlgebra
using Random

include("../common/PWGrid_v02.jl")
include("../common/ortho_gram_schmidt.jl")
include("../common/wrappers_fft.jl")
include("../common/gen_lattice_pwscf.jl")
include("EnergiesT.jl")
include("PotentialsT.jl")
include("gen_dr.jl")
include("init_pot_harm_3d.jl")
include("op_K.jl")
include("op_V_loc.jl")
include("op_H.jl")
include("calc_rho.jl")
include("calc_grad.jl")
include("calc_Energies.jl")
include("KS_solve_Emin_sd.jl")
include("KS_solve_Emin_cg.jl")
include("Poisson_solve.jl")
include("LDA_VWN.jl")
include("Kprec.jl")

include("diag_lobpcg.jl")
include("KS_solve_SCF.jl")

include("andersonmix.jl")
include("KS_solve_SCF_andersonmix.jl")

include("pulaymix.jl")
include("KS_solve_SCF_pulaymix.jl")

include("get_ub_lb_lanczos.jl")
include("KS_solve_ChebySCF.jl")
include("chebyfilt.jl")
include("norm_matrix_induced.jl")

function test_main( Ns; method="SCF" )

    LatVecs = gen_lattice_sc(6.0)

    pw = PWGrid( Ns, LatVecs )

    Ω  = pw.Ω
    r  = pw.r
    G  = pw.gvec.G
    G2 = pw.gvec.G2
    Npoints = prod(Ns)
    Ngwx = pw.gvecw.Ngwx

    @printf("Ns   = (%d,%d,%d)\n", Ns[1], Ns[2], Ns[3])
    @printf("Ngwx = %d\n", Ngwx)

    actual = Npoints/Ngwx
    theor = 1/(4*pi*0.25^3/3)
    @printf("Compression: actual, theor: %f , %f\n", actual, theor)

    #
    # Generate array of distances
    #
    center = sum(LatVecs,dims=2)/2
    dr = gen_dr( r, center )
    #
    # Setup potential
    #
    V_ionic = init_pot_harm_3d( pw, dr )
    println("sum(Vpot)*Ω/Npoints = ", sum(V_ionic)*Ω/Npoints)
    #
    Nstates = 4
    Focc = 2.0*ones(Nstates)

    if method == "SD"
        # This will need a lot of optimization steps
        psi, Energies, Potentials = KS_solve_Emin_sd( pw, V_ionic, Focc, Nstates, NiterMax=10 )
    elseif method == "CG"
        psi, Energies, Potentials = KS_solve_Emin_cg( pw, V_ionic, Focc, Nstates, NiterMax=1000 )
    elseif method == "SCF"
        Energies, Potentials, psi, evals = KS_solve_SCF( pw, V_ionic, Focc, Nstates, β=0.5 )
    elseif method == "SCF_andersonmix"
        Energies, Potentials, psi, evals = KS_solve_SCF_andersonmix( pw, V_ionic, Focc, Nstates, β=0.5 )
    elseif method == "SCF_pulaymix"
        Energies, Potentials, psi, evals = KS_solve_SCF_pulaymix( pw, V_ionic, Focc, Nstates, β=0.5 )
    elseif method == "ChebySCF"
        Energies, Potentials, psi, evals = KS_solve_ChebySCF( pw, V_ionic, Focc, Nstates, β=0.1 )
    else
        println("ERROR: unknown method: ", method)
        exit()
    end

    #psi, Energies, Potentials = KS_solve_Emin_cg( pw, V_ionic, Focc, Nstates, NiterMax=1000 )

    if method == "SD" || method == "CG"
        Y = ortho_gram_schmidt(psi)
        mu = Y' * op_H( pw, Potentials, Y )
        evals, evecs = eigen(mu)
        Psi = Y*evecs
    end

    print_Energies(Energies)
    for ist = 1:Nstates
        @printf("State # %d, Energy = %18.10f ===\n", ist, real(evals[ist]))
    end
end

#@time test_main( [30, 30, 30], method="SD" )
@time test_main( [30, 30, 30], method="CG" )
#@time test_main( [30, 30, 30], method="SCF" )
#@time test_main( [30, 30, 30], method="SCF_andersonmix" )
#@time test_main( [30, 30, 30], method="SCF_pulaymix")
#@time test_main( [30, 30, 30], method="ChebySCF" )


text:
module legacyODEDSL

using Base.Test
using DataStructures
using Docile
using Lexicon
using Calculus
import JSON

import Calculus.differentiate, Base.min, Base.max, Base.deepcopy

# In order to enable docstring-style documentation of Julia code

@document

# Include the source files
include("utils.jl")
include("datatypes.jl")
import .DataTypes
dt = DataTypes
include("parser.jl")
include("conversion_reaction.jl")
include("conversion_ode.jl")
include("symbolic.jl")
include("checks.jl")
include("translate.jl")
include("translate_rcpp.jl")

end


text:
using PSplines
using Test

@testset "PSplines.jl" begin
    include("test_MatricesBuilding.jl")
    include("test_types.jl")
    include("test_utils.jl")
end


text:
using Random
import Future
import Base.Threads: @threads
#macro threads(ex) :($(esc(ex))) end

"""
   function to create a population for GAState or ga
"""
function initializepop(model::GAModel,npop::Integer,
                       nelites::Integer, baserng,sort=true)
    # each thread gets its own auxiliary scratch space
    nthreads = Threads.nthreads()
    rngs = let m = MersenneTwister(1)
                [m; accumulate(Future.randjump, fill(big(10)^20, nthreads-1), init=m)]
            end;
    aux = map(i -> genauxga(model), 1:nthreads)
    if npop > 0
        # initialize population
        pop1 = randcreature(model,aux[1],rngs[1])
        pop = Vector{typeof(pop1)}(undef, npop)
        pop[1] = pop1
        @threads for i = 2:npop
            threadid = Threads.threadid()
            pop[i] = randcreature(model,aux[threadid],rngs[threadid])
        end
        if sort
            sort!(pop,by=fitness,rev=true,
                  alg=PartialQuickSort(max(1,nelites)))
        end
    else
        pop = nothing
    end
    return pop,aux,rngs
end

mutable struct GAState
    model::GAModel
    pop::Vector
    ngen::Integer
    curgen::Integer
    npop::Integer
    elite_fraction::Real
    crossover_params
    mutation_params
    print_fitness_iter::Real
    save_creature_iter::Real
    save_state_iter::Integer
    file_name_prefix::AbstractString
    baserng::AbstractRNG
end
function GAState(model::GAModel;
                 ngen=10,
                 npop=100,
                 elite_fraction=0.01,
                 crossover_params=nothing,
                 mutation_params=Dict(:rate=>0.1),
                 print_fitness_iter=1,
                 save_creature_iter=0,
                 save_state_iter=0,
                 file_name_prefix="gamodel",
                 baserng=Random.GLOBAL_RNG)
    0 <= elite_fraction <= 1 || error("elite_fraction bounds")
    nelites = Int(floor(elite_fraction*npop))
    pop,aux,rngs = initializepop(model, npop, nelites, baserng)
    return GAState(model,pop,ngen,0,npop,
                   elite_fraction,
                   crossover_params,
                   mutation_params,
                   print_fitness_iter,
                   save_creature_iter,
                   save_state_iter,
                   file_name_prefix,
                   baserng)
end

"""
       Saves ga state to file
       Doesn't support GAModels or GACreatures containing functions (e.g. CoordinateModel)
       since JLD doesn't support saving functions
"""
function savegastate(file_name_prefix::AbstractString, curgen::Integer, state::GAState)
    filename = "$(file_name_prefix)_state_$(curgen).jld"
    println("Saving state to file $filename")
    save(filename,"state",state)
end

function loadgastate(filename::AbstractString)
    println("Load state from file $filename")
    load(filename, "state")
end

"""
    ga function
    x in each generation, the following is done
        - select parents from all creatures in population
        - create children using crossover
        - replace non-elites in population with children
        - mutate all creatures (both elites and children) in population
        - logging
            x saves state every save_state_iter iterations to file
                - restart using state = loadgastate(filename) & ga(state)
            x outputs creature every save_creature_iter iterations to file
            x prints fitness value every print_fitness_iter iterations to screen
        - go to next generation
            """
function ga(state::GAState)
    # load from state
    model = state.model
    pop = state.pop
    ngen = state.ngen
    curgen = state.curgen
    npop = state.npop
    elite_fraction = state.elite_fraction
    crossover_params = state.crossover_params == nothing ? "Nothing" : state.crossover_params
    mutation_params = state.mutation_params
    print_fitness_iter = state.print_fitness_iter
    save_creature_iter = state.save_creature_iter
    save_state_iter = state.save_state_iter
    file_name_prefix = state.file_name_prefix
    baserng = state.baserng

    nelites = Int(floor(elite_fraction*npop))
    nchildren = npop-nelites
    children = deepcopy(pop[nelites+1:end])
    # initialize auxiliary space
    _,aux,rngs = initializepop(model, 0, 0, baserng)

    println("Running genetic algorithm with
            population size $npop,
            generation number $ngen,
            elite fraction $elite_fraction,
            children created $nchildren,
            crossover_params $crossover_params,
            mutation_params $mutation_params,
            printing fitness every $print_fitness_iter iteration(s),
            saving creature to file every $save_state_iter iteration(s),
            saving state every $save_state_iter iteration(s),
            with file name prefix $file_name_prefix.")

    # main loop:
    # 1. select parents
    # 2. crossover parents & create children
    # 4. replace non-elites in current generation with children
    # 3. mutate population
    # 5. sort population
    for outer curgen = curgen+1:ngen
        # crossover. uses multi-threading when available
        parents = selection(pop, nchildren, rngs[1])
        @threads for i = 1:nchildren
            threadid = Threads.threadid()
            p1,p2 = parents[i]
            children[i] = crossover(children[i], pop[p1], pop[p2],
                                    model, crossover_params, curgen,
                                    aux[threadid], rngs[threadid])
        end
        # moves children and elites to current pop
        for i = 1:nchildren
            ip = nelites+i
            # swapping instead of deepcopy-ing
            pop[ip], children[i] = children[i], pop[ip]
        end
        # mutate pop (including elites; except for the most elite creature
        # in order to preserve monotonocity wrt best fitness)
        @threads for i = 2:npop
            threadid = Threads.threadid()
            pop[i] = mutate(pop[i], model, mutation_params, curgen,
                            aux[threadid], rngs[threadid])
        end
        sort!(pop,by=fitness,rev=true,alg=PartialQuickSort(max(1,nelites)))

        if print_fitness_iter>0 && mod(curgen,print_fitness_iter)==0
            printfitness(curgen, pop[1])
        end

        if save_creature_iter>0 && mod(curgen,save_creature_iter)==0
            savecreature(file_name_prefix, curgen, pop[1], model)
        end

        if save_state_iter>0 && mod(curgen,save_state_iter)==0
            state = GAState(model,pop,ngen,curgen,npop,
                            elite_fraction,
                            crossover_params,
                            mutation_params,
                            print_fitness_iter,
                            save_creature_iter,
                            save_state_iter,
                            file_name_prefix,
                            baserng)
            savegastate(file_name_prefix, curgen, state)
        end
    end
    state.curgen = curgen
    pop[1]
end


text:
function animate_hist(pomdp::UrbanPOMDP, hist, overlays::Vector{SceneOverlay} = SceneOverlay[IDOverlay()], cam = FitToContentCamera(0.))
    duration = n_steps(hist)*pomdp.ΔT
    fps = Int(1/pomdp.ΔT)
    function render_hist(t, dt)
        stateindex = Int(floor(t/dt)) + 1
        scene = state_hist(hist)[stateindex]
        return AutoViz.render(scene, pomdp.env, overlays, cam = cam,  car_colors=get_colors(scene))
    end
    return duration, fps, render_hist
end

function animate_hist(pomdp::UrbanPOMDP, state_hist, lidar_hist, overlays::Vector{SceneOverlay} = SceneOverlay[IDOverlay()], cam = FitToContentCamera(0.))
    duration = length(state_hist)*pomdp.ΔT
    fps = Int(1/pomdp.ΔT)
    function render_hist(t, dt)
        stateindex = Int(floor(t/dt)) + 1
        scene = state_hist[stateindex]
        overlays_t = [overlays; LidarOverlay(lidar_hist[stateindex])]
        return AutoViz.render(scene, pomdp.env, overlays_t, cam = cam,  car_colors=get_colors(scene))
    end
    return duration, fps, render_hist
end

function animate_scenes(scenes::Vector{Scene}, models, env::UrbanEnv; overlays::Vector{SceneOverlay} = SceneOverlay[], cam=FitToContentCamera(0.),  sim_dt=0.1)
    duration =length(scenes)*sim_dt
    fps = Int(1/sim_dt)
    # cam = FitToContentCamera(0.)
    function render_rec(t, dt)
        frame_index = Int(floor(t/dt)) + 1
        disp = SceneOverlay[overlays...]
        push!(disp, TextOverlay(text = ["Step: $frame_index"],
                                                           font_size=20,
                                                           pos=VecE2(-15,-10),
                                                           incameraframe=true))
        for veh in scenes[frame_index]
            if veh.id == EGO_ID || veh.def.class == AgentClass.PEDESTRIAN
                continue
            end
            route = models[frame_index][veh.id].navigator.route
            blk_on, blk_right = turn_map(env, (route[1].tag, route[end].tag))
            push!(disp, BlinkerOverlay(on=blk_on, right=blk_right, veh=veh))
        end
        return AutoViz.render(scenes[frame_index], env, disp, cam=cam)
    end
    return duration, fps, render_rec
end

max_stars_repo_name: MaximeBouton/AutomotivePOMDPs

max_issues_repo_name: MaximeBouton/AutomotivePOMDPs

max_forks_repo_name: MaximeBouton/AutomotivePOMDPs


text:
#NNaux.jl
#Type definitions and useful functions
#TODO: define -?

#put into NNaux.jl?
#########Types and associated utility functions ########
########################################################
#how to make the size of the array modular?
type NNWeights{d}
  x::Dict{Int,Array{Float64,d}}
end

type NeuralNetwork
  #size::Array{Int,1}
  #weights::Array{Float64,2}#dict layer number ->array{float,1}?
  weights::NNWeights
  #biases::Array{Float64,2}
  biases::NNWeights
  neuron::Function
  gradient::Function
  output_neuron::Function
  output_gradient::Function
  c::Float64 #regularization term
  reg::String #type of regularization

  #lam1::Float64 #L1 regularization term
  #lam2::Float64 #L2 regularization term
  #alternatively, have a arbitrary "c" parameter that gets related to one of the regularization types: L1, L2, maxnorm, dropout
  yu::Array{Float64,1}
  yl::Array{Float64,1}
  xu::Array{Float64,1}
  xs::Array{Float64,1}
  verbose::Bool
end

function getindex(x::NNWeights,i::Int)
  return x.x[i]
end

function setindex!(x::NNWeights,y::Array{Float64,2},i::Int)
  x.x[i] = y
end

function setindex!(x::NNWeights,y::Array{Float64,1},i::Int)
  x.x[i] = y
end

function length(x::NNWeights)
  return length(x.x)
end

function add!(x::NNWeights,y::NNWeights)
  #check this vvvv
  L = length(x.x)
  assert(L == length(y.x))
  #assert keys are equal?
  assert(length(setdiff(keys(x.x),keys(y.x))) == 0)

  for i = 1:L
    x.x[i] += y.x[i]
  end

  return x
end

function *(x::NNWeights,a::Real)
  y = zero(x)
  L = length(x.x)
  for i = 1:L
    y.x[i] = x.x[i]*a
  end
  return y
end

function *(a::Real,x::NNWeights)
  return x*a
end

function +(x::NNWeights, y::NNWeights)
  L = length(x.x)
  assert(L == length(y.x))
  #assert keys are equal? Do I even need this assertion?--does it even make that much perf difference?
  assert(length(setdiff(keys(x.x),keys(y.x))) == 0)
  z = zero(x)
  for i = 1:L
    z.x[i] = x.x[i] + y.x[i]
  end
  return z
end

function +(x::NNWeights,y::Float64)
  return x + one(x)*y
end

function +(y::Float64,x::NNWeights)
  return x + y
end

function mult!(x::NNWeights,a::Real)
  L = length(x.x)
  for i = 1:L
    x.x[i] *= a
  end
  return x
end

function zero(x::NNWeights)
  #make a zero NNweight of the same size
  L = length(x.x)
  d = length(size(x.x[1]))
  y = Dict{Int,Array{Float64,d}}([i=>zeros(size(x.x[i])) for i = 1:L])
  return NNWeights{d}(y)
end

function zero(sizes)
  #i want the type of the parameter to be an array of integer tuples of length 1 or 2
  L = length(sizes)
  d = length(sizes[1])
  x = Dict{Int,Array{Float64,d}}([i=>zeros(sizes[i]...) for i = 1:L])
  return NNWeights{d}(x)
end

function one(x::NNWeights)
  #make a zero NNweight of the same size
  L = length(x.x)
  d = length(size(x.x[1]))
  y = Dict{Int,Array{Float64,d}}([i=>ones(size(x.x[i])) for i = 1:L])
  return NNWeights{d}(y)
end

function one(sizes)
  #i want the type of the parameter to be an array of integer tuples of length 1 or 2
  L = length(sizes)
  d = length(sizes[1])
  x = Dict{Int,Array{Float64,d}}([i=>ones(sizes[i]...) for i = 1:L])
  return NNWeights{d}(x)
end

function randweights(sizes)
  #i want the type of the parameter to be an array of integer tuples of length 1 or 2
  #only for initialization of parameters
  L = length(sizes)
  d = length(sizes[1])
  x = Dict{Int,Array{Float64,d}}([i=>(rand(sizes[i]...)-0.5)./sqrt(2.0/sizes[i][2]) for i = 1:L])
  return NNWeights{d}(x)
end

function norm(x::NNWeights)
  L = length(x.x)
  val = 0
  for i = 1:L
    val += sum([j*j for j in x.x[i][:]])#vecnorm(x.x[i])
  end
  return val
end

function ^(x::NNWeights,i::Int)
  y = zero(x)
  L = length(x.x)
  for i = 1:L
    y.x[i] = x.x[i].*x.x[i]
  end
  return y
end

function /(x::NNWeights,y::NNWeights)
  z = zero(x)
  L = length(x.x)
  for i = 1:L
    z.x[i] = x.x[i]./y.x[i]
  end
  return z
end

function sqrt(x::NNWeights)
  y = zero(x)
  L = length(x.x)
  for i = 1:L
    y.x[i] = sqrt(x.x[i])
  end
  return y
end
#######################################################################


text:
using NumericalMethodsforEngineers

a = [16. 4. 8.; 4. 5. -4.; 8. -4. 22.]
b = [4., 2., 5.]

a = sparse(a)

F = cholfact(a)
upper = F[:U]
upper |> display

lower = F[:L]
lower |> display

y = lower \ b

c = upper \ y
println("Solution Vector: \n $c")

@test c == a \ b

a = [16. 4. 8.; 4. 5. -4.; 8. -4. 22.]
b = [4., 2., 5.]

a[2, 2] += 1.0e20
b[2] = a[2, 2] * 5.0

a = sparse(a)

F = cholfact(a)
upper = F[:U]
upper |> display

lower = F[:L]
lower |> display

y = lower \ b

c = upper \ y
println("Solution Vector: \n $c")

@test c == a \ b


text:
# Configuration parameters

mutable struct Config
    # A minimum pixel radius to be included around each source.
    min_radius_pix::Float64

    function Config()
        config = new()
        config.min_radius_pix = 8.0
        config
    end
end


text:
include("fibonacci_app.jl")
app = fibonacci_app(true)
run_server(app)


text:
using BeamformerRecipes
using YAML

"""
    cal_solution(redis, arrayname, unixtime) -> nants, nchan, ants, cals

Fetch the most recent telstate calibration solution for subarray `arrayname`
from `redis` as of `unixtime`.  `nants` and `nchan` are integers used in the
dimensioning of the data arrays, `ants` is a Vector of antenna names (in F
engine order), and `cals` is a `Dict{Symbol, Any}` suitable for splatting into
the keyword arguments of the `CalInfo` constructor.
"""
function cal_solution(redis, arrayname, unixtime)
    # zrevrangebyscore returns an OrderedSet
    rkeys = zrevrangebyscore(redis, "$arrayname:cal_solutions:index", unixtime, 0, "limit", 0, 1)
    if length(rkeys) == 0
        @warn "no cal solutions found on or before $unixtime, using fallback hack" _module=nothing _file=nothing
        rkeys = ["array_1:cal_solutions:20220117T035342Z"]
    end
    length(rkeys) <= 0 && error("no cal solutions found")
    length(rkeys)  > 1 && error("too many cal solutions found")
    @info "using cal solution $(rkeys[1])"
    strcals = hgetall(redis, rkeys[1])

    # Parse integers
    nants = parse(Int, strcals["nants"])
    nchan = parse(Int, strcals["nchan"])
    ants = YAML.load(strcals["antenna_list"])
    # TODO Fix upstream and in redis
    # Workaround for nants bug
    if nants > 64
        nants ÷= 8
    end

    # Create Arrays and copy data into them
    # TODO Fix cal_K type upstream and in redis
    if sizeof(strcals["cal_K"]) == nants * 2 #=npols=# * sizeof(Float64) * 2 #=complex=#
        cal_K = Array{ComplexF64}(undef, nants, 2)
    else
        cal_K = Array{Float64}(undef, nants, 2)
    end
    @assert sizeof(cal_K) == sizeof(strcals["cal_K"])
    unsafe_copyto!(pointer(cal_K), Ptr{eltype(cal_K)}(pointer(strcals["cal_K"])), length(cal_K))

    cal_B = Array{ComplexF64}(undef, nants, 2, nchan)
    @assert sizeof(cal_B) == sizeof(strcals["cal_B"])
    unsafe_copyto!(pointer(cal_B), Ptr{ComplexF64}(pointer(strcals["cal_B"])), length(cal_B))

    cal_G = Array{ComplexF64}(undef, nants, 2)
    @assert sizeof(cal_G) == sizeof(strcals["cal_G"])
    unsafe_copyto!(pointer(cal_G), Ptr{ComplexF64}(pointer(strcals["cal_G"])), length(cal_G))

    cal_all = Array{ComplexF64}(undef, nants, 2, nchan)
    @assert sizeof(cal_all) == sizeof(strcals["cal_all"])
    unsafe_copyto!(pointer(cal_all), Ptr{ComplexF64}(pointer(strcals["cal_all"])), length(cal_all))

    # Replace all NaN values with 0
    cal_K[isnan.(cal_K)] .= 0
    cal_B[isnan.(cal_B)] .= 0
    cal_G[isnan.(cal_G)] .= 0
    cal_all[isnan.(cal_all)] .= 0

    cals = Dict{Symbol, Any}()
    cals[:cal_K] = real(cal_K) * 1e9 # seconds to nanoseconds
    cals[:cal_B] = cal_B
    cals[:cal_G] = cal_G
    cals[:cal_all] = cal_all
    cals[:refant] = strcals["refant"]

    # (Un-)conjugate cal_all antenna data upto refant
    refant_idx = findfirst(==(cals[:refant]), ants)
    @views conj!(cals[:cal_all][1:refant_idx-1, :, :])

    nants, nchan, ants, cals
end


text:
module Data

include("structs.jl")

using Flux
using Flux: onehotbatch
using Flux.Data: DataLoader
using Plots
using ScikitLearn
@sk_import datasets: make_moons

export get_moon_data_loader, get_nmist_data_loader

"""
    get_moon_data_loader(n = 300, offset = 0.5, shuffle = true, seed = 1337)

Returns a Flux.DataLoader with moon data for the specified parameters.
Supports creation of random moons and shuffling of the data points.
"""
function get_moon_data_loader(; n::Int64 = 300, offset::Float64 = 0.5, shuffle::Bool = true, seed::Int64 = rand((1, 2^31)), args...)
    args = Args(; args...)

    # offset: linearly separable => 1.0; not lin. sep. => 0.5
	x, y = generate_moons(n, offset = offset, seed = seed)
	x, y = transpose(x), onehotbatch(y, 0:1)

	# create data loader
	return DataLoader((x, y), batchsize = args.batchsize, shuffle = shuffle)
end

"""
    get_nmist_data_loader(shuffle = true)

Returns a Flux.DataLoader with MNIST data for the specified parameters.
"""
function get_nmist_data_loader(; shuffle::Bool = true, args...)
    args = Args(; args...)

	x, y = MLDatasets.MNIST.traindata(Float32)

	# reshape Data in order to flatten each image into a linear array
    # encode labels with one-hot
	x, y = Flux.flatten(x), onehotbatch(y, 0:9)

	# create data loader
	return DataLoader((x, y), batchsize = args.batchsize, shuffle = shuffle)
end

"""
    generate_moons_from_publication(offset = 0.5, coord_down_scale = 1.0, show_plot = false)

Returns the exact same moons that were used on the demo website
of the publication authors at https://mohammadpz.github.io/GradientStarvation.html
"""
function generate_moons_from_publication(; offset::Float64 = 0.0, coord_down_scale::Float64 = 1.0, show_plot::Bool = false)
    moons = hcat(moon_class_1(offset, coord_down_scale), moon_class_2(offset, coord_down_scale))
    labels = [repeat([0], 150); repeat([1], 150)]

    if show_plot
        display(scatter(moons[1,:], moons[2,:], c = labels))
    end

    return moons, labels
end

"""
    generate_moons(n = 300, noise = 0.42, offset = 0.5, rotation = 90, seed = 42, show_plot = true)

Wrapper function for scikitlearn make_moons function. Additionally features the options to rotate and offset the
moons or plot the generated data.
"""
function generate_moons(n::Int64; noise::Float64 = 0.09, offset::Float64 = 0.0, rotation::Int64 = 90, seed::Int64 = rand((1, 2^31)), show_plot::Bool = false)
    X, y = make_moons(n_samples = n, noise = noise, random_state = seed)
    p1 = scatter(X[:,1], X[:,2], c = y, title = "generated state")

    if offset != 0.0
        X = apply_moon_offset(X, y, offset)
    end
    p2 = scatter(X[:,1], X[:,2], c = y, title = "offsetted state")

    if rotation != 0.0
        θ = deg2rad(rotation)
        R = [cos(θ) -sin(θ); sin(θ) cos(θ)]
        X *= R
        # move center closer to (0, 0)
        if offset != 0.0
            X[:,2] .+= offset/2
        end
    end
    p3 = scatter(X[:,1], X[:,2], c = y, title = "rotated state")

    if show_plot
        display(plot(p1, p2, p3, layout = (1, 3)))
    end

    return X, y
end

"""
    apply_moon_offset(X, y, 0.5)

Separates both moons by the specified offset.
"""
function apply_moon_offset(X, y, offset)
    # use offset as total offset, thus devide by number of labels; assume 2 labels
    offset /= 2
    moons = []

    for label in unique(y)
        push!(moons, findall(l -> l == label, y))
    end

    max_y = [maximum(X[m, 2]) for m in moons]

    if max_y[1] > max_y[2]
        X[moons[1], 2] .+= offset
        X[moons[2], 2] .-= offset
    else
        X[moons[1], 2] .-= offset
        X[moons[2], 2] .+= offset
    end

    return X
end

"""
    moon_class_1(0.0, 1.0)

Returns the hardcoded coordinates of the first moon from the author's demo website
at https://mohammadpz.github.io/GradientStarvation.html
"""
function moon_class_1(offset::Float64 = 0.0, coord_down_scale::Float64 = 1.0)
    x = [
        -2.717835893494339 - offset,
        -2.0025756608302556 - offset,
        -2.568929506049188 - offset,
        -2.270602886384911 - offset,
        -2.5015948773601426 - offset,
        -2.330151863221954 - offset,
        -0.7415685155965577 - offset,
        -2.0839366412304012 - offset,
        -2.8528165357526083 - offset,
        -1.147008103324011 - offset,
        -2.8784004738595304 - offset,
        -0.8677959198389527 - offset,
        -1.7589836417118123 - offset,
        -2.998388218477527 - offset,
        -2.122794047965966 - offset,
        -0.6950154023659438 - offset,
        -2.6642421732719845 - offset,
        -0.4107292059512281 - offset,
        -2.5650150317912725 - offset,
        -0.03957624387021477 - offset,
        -1.459975933353292 - offset,
        -1.7826178769431 - offset,
        -0.0063910086450723325 - offset,
        -3.113867547530851 - offset,
        -2.8910961520154546 - offset,
        -2.3718952252880245 - offset,
        -1.3897222631802355 - offset,
        -2.1375714939983066 - offset,
        -0.9994610567775488 - offset,
        -0.7776520784930578 - offset,
        -0.353388277001343 - offset,
        -0.02895203726210091 - offset,
        -0.6726983784145278 - offset,
        -2.2946962578019665 - offset,
        -1.863557466465306 - offset,
        -0.26730788088522667 - offset,
        -2.498303344433606 - offset,
        -0.9859628310147427 - offset,
        -2.3015390566521345 - offset,
        -2.182161422047994 - offset,
        -0.6476638289513104 - offset,
        -1.9828161461009488 - offset,
        -2.759334618718058 - offset,
        -2.5751009200744224 - offset,
        -2.2754802700838868 - offset,
        -2.894718819023576 - offset,
        -2.3942266106035563 - offset,
        -2.3597155702364154 - offset,
        -0.8619147842448072 - offset,
        -1.5585362847484685 - offset,
        -2.2329166647636507 - offset,
        -0.2731017075516541 - offset,
        -0.4781390413985451 - offset,
        -1.7901443397126593 - offset,
        -1.1433045993548665 - offset,
        -2.0898117037306303 - offset,
        -1.0837510824272951 - offset,
        -2.6928615859081817 - offset,
        -2.649706340336663 - offset,
        -2.6375589115090947 - offset,
        -0.7346260030476295 - offset,
        -2.42417820501776 - offset,
        -0.7008726348089543 - offset,
        -0.3491363431388116 - offset,
        -2.6855954583467803 - offset,
        -0.9335914608244527 - offset,
        -1.1282308435392019 - offset,
        -1.5571341099841778 - offset,
        -2.0968912603614944 - offset,
        -0.730628023393316 - offset,
        -0.4951985004108488 - offset,
        -2.902350613309882 - offset,
        -1.7885697248681875 - offset,
        -1.4751061664049347 - offset,
        -1.5390777424462194 - offset,
        -2.8762322285242092 - offset,
        -0.4940392357287259 - offset,
        -0.012054114845436309 - offset,
        -3.032053951860437 - offset,
        -2.374125805765452 - offset,
        -2.4025184090990277 - offset,
        -1.0969861251226607 - offset,
        -2.7140378394308975 - offset,
        -1.753366456813175 - offset,
        -2.350980683560984 - offset,
        -0.4481758551592373 - offset,
        -0.4837435351770224 - offset,
        -2.6861383109374612 - offset,
        -2.2799772528992794 - offset,
        -2.0013629328184592 - offset,
        -2.9883132223545816 - offset,
        -2.1215585709864007 - offset,
        -2.421553486436626 - offset,
        -0.5889598185224101 - offset,
        -2.7964229143440953 - offset,
        -2.1055240403942315 - offset,
        -2.550214662670338 - offset,
        -2.60314101539995 - offset,
        -0.7509004504010787 - offset,
        -1.1097198256588752 - offset,
        -1.2853090821878264 - offset,
        -0.06881151445833164 - offset,
        -1.888037104217837 - offset,
        -2.655692934648788 - offset,
        -0.0712617816418619 - offset,
        -2.237386070712279 - offset,
        -2.624697844350373 - offset,
        -1.6932598229929476 - offset,
        -3.1292392711691432 - offset,
        -1.9302970612705348 - offset,
        -0.05552000023091433 - offset,
        -2.213859883403356 - offset,
        -2.2855940763118876 - offset,
        -1.5360545910906862 - offset,
        -1.8423162009755345 - offset,
        -1.2540315350699642 - offset,
        -2.4100284289768377 - offset,
        -1.2316448752106524 - offset,
        -1.009437911641195 - offset,
        -2.2341929394003626 - offset,
        -1.5936721603442507 - offset,
        -1.4301284921136954 - offset,
        -2.133956374401404 - offset,
        -2.5419546529030743 - offset,
        -2.4706384450766503 - offset,
        -1.4458442971552867 - offset,
        -0.675298309924629 - offset,
        -0.1307759125127804 - offset,
        -2.1724261559379694 - offset,
        -3.0301757418373185 - offset,
        -2.223984637645856 - offset,
        -0.8960548947689856 - offset,
        -1.6879441985156765 - offset,
        -1.5656539231364017 - offset,
        -1.1715302401075987 - offset,
        -2.1977928180739728 - offset,
        -2.3569595232747393 - offset,
        -1.812798093076864 - offset,
        -1.0133105775192217 - offset,
        -0.8510029404821573 - offset,
        -1.4057618612620275 - offset,
        -0.8426327453397771 - offset,
        -1.919390352513409 - offset,
        -2.7504997471435866 - offset,
        -1.463129347902804 - offset,
        -2.7820769325267394 - offset,
        -1.9904715999050417 - offset,
        -2.0374228816770796 - offset,
        -1.9573778533077693 - offset,
        -1.0107773780231675 - offset
    ]
    y = [
        0.6793610085529345,
        -0.7764598781122054,
        1.3030331041149927,
        1.696103804129755,
        1.8362701766489786,
        0.5295968005263707,
        3.5042711192603146,
        -0.7231575356785352,
        2.277743746352307,
        3.1121491784480697,
        1.4770243521157966,
        -1.0110529899688063,
        -0.21940659318914466,
        1.2783156815611991,
        -0.1791718798868415,
        3.884218615064042,
        1.9439974190429932,
        3.505002569928259,
        0.5755380178279407,
        -1.210873365508656,
        -1.0329923799436405,
        -1.0582442212859733,
        3.6142846309958205,
        1.5445673874139447,
        1.6469163448206863,
        0.2540079910444923,
        -1.258967648201369,
        2.794453234058449,
        -1.149826801806625,
        4.1149852157058895,
        -1.4081025494488113,
        3.5570091502314423,
        3.713458077324815,
        0.9623315035843127,
        2.792146626757914,
        -1.205590370248826,
        1.3344598403007284,
        3.057031649768471,
        2.578789613141873,
        -0.6543242980905074,
        -1.2100425673015185,
        2.9197086781500277,
        0.7696706532935776,
        1.6626933401609023,
        0.42253538652727796,
        1.4631522384301066,
        2.942466273998301,
        1.2364964850094777,
        3.152143289401243,
        3.6228480325523322,
        -0.6233110354243981,
        -1.3580082036894057,
        -1.305693088622172,
        -0.3935866480543121,
        3.111337871342845,
        2.410394433635401,
        -0.7759536147862817,
        1.3340538265678492,
        2.2303142906955555,
        0.3209555538663683,
        3.89346947511081,
        0.7649583168967733,
        -1.0802855176420891,
        -1.2568494543447684,
        2.4944104537407057,
        -0.679082202208783,
        3.693884882642762,
        -0.9569209611232933,
        -0.3002820219513763,
        -1.2603402702853281,
        -0.729087176112734,
        0.16046671198639534,
        2.8847324649501402,
        3.4142846215369556,
        2.9919434674688086,
        0.9143476976637405,
        -1.0038621010673856,
        3.7992457338282373,
        0.3803212691710048,
        2.577260760939388,
        1.9925005915231553,
        3.0431663584149904,
        1.7300408024972684,
        -0.46725570847050424,
        -0.25749919862715076,
        -1.493604284214328,
        4.396876304605113,
        1.5310130579985914,
        0.11093029172240809,
        -0.6139231414423276,
        0.9208951171502733,
        0.6712563554792417,
        2.618298719018602,
        3.6804690431500036,
        0.48814363461182453,
        1.7996187499638172,
        1.1625285801183705,
        0.7389373388724324,
        -1.2211299676858136,
        -0.8873468825484077,
        3.4776733678417986,
        3.575493958460944,
        2.3985492607241343,
        1.71557735708596,
        3.749934520345652,
        0.3575930800531485,
        2.1434683067151066,
        -0.9049053449496425,
        0.10465507554354464,
        3.1473437369575374,
        3.8373893523193847,
        -0.2809886987751162,
        -0.0809670563060226,
        -1.0586183593166278,
        2.9422791248401476,
        3.3776803834259663,
        3.153656143587366,
        3.5152991121388566,
        -1.281563505634356,
        1.7966322608299223,
        3.523561811455531,
        -0.7139026560764259,
        2.830553038963384,
        1.756811334089956,
        2.7361898523968753,
        -0.9036159267450711,
        -0.8009280139118677,
        -1.3972949916010728,
        2.4380037963785406,
        -0.05613139501335074,
        -0.9752089176334047,
        3.3757400103199693,
        3.2402795631690253,
        3.1005177839870184,
        -1.0528879963506719,
        0.11441934583356184,
        2.5168873649410255,
        2.960577394399894,
        -0.8705594922460475,
        3.573688045781485,
        -0.8907098511435746,
        3.294176954534472,
        -0.08783059329234524,
        2.292660484192841,
        3.4597453566827,
        0.253560604420727,
        -0.39806607018651047,
        2.5778821384858093,
        -0.7882243058017496,
        -0.9895456553950488
    ]

    return vcat(transpose.((x, y))...) / coord_down_scale
end

"""
    moon_class_2(0.0, 1.0)

Returns the hardcoded coordinates of the second moon from the author's demo website
at https://mohammadpz.github.io/GradientStarvation.html
"""
function moon_class_2(offset::Float64 = 0.0, coord_down_scale::Float64 = 1.0)
    x = [
        1.2605697631874289 + offset,
        0.4428841761287988 + offset,
        2.213537711353039 + offset,
        2.2246476075616277 + offset,
        0.934302193080736 + offset,
        0.8703714425371154 + offset,
        2.6652808441973175 + offset,
        2.5868192990813768 + offset,
        2.5485659972780015 + offset,
        1.3501341878455395 + offset,
        1.3394525649634517 + offset,
        1.5703488118666078 + offset,
        2.578906600693658 + offset,
        2.6017725256967363 + offset,
        0.4855765554512216 + offset,
        1.620317196388419 + offset,
        0.6107465335988917 + offset,
        2.705677633402291 + offset,
        2.529619115397467 + offset,
        1.7205080921848972 + offset,
        2.025034051353637 + offset,
        0.9454991520471547 + offset,
        2.4474158496753415 + offset,
        0.44018725848007784 + offset,
        2.3578626209991933 + offset,
        2.305287296588136 + offset,
        2.1078203674554112 + offset,
        2.5672512309425377 + offset,
        0.76210430597206 + offset,
        0.9009064185658802 + offset,
        0.8790969921498528 + offset,
        2.834290242153208 + offset,
        1.827620864188383 + offset,
        0.9106326766024155 + offset,
        2.4986568505695774 + offset,
        1.329902276945841 + offset,
        0.30026706132425485 + offset,
        1.8709094482262436 + offset,
        2.6457300611721415 + offset,
        0.903254228522091 + offset,
        1.6008172396371303 + offset,
        3.2153091597614525 + offset,
        2.5679011437621266 + offset,
        2.359601929920219 + offset,
        2.0630126331781993 + offset,
        0.9482707644288992 + offset,
        1.318717891036563 + offset,
        2.685806352629474 + offset,
        1.8516766580167963 + offset,
        2.356710102702996 + offset,
        2.1684918727788336 + offset,
        1.5800357609867945 + offset,
        1.033126269584137 + offset,
        1.872794212999521 + offset,
        0.5235449122291833 + offset,
        0.7008785749342308 + offset,
        2.553663974985213 + offset,
        2.7411100225680163 + offset,
        2.0343480716575852 + offset,
        0.5340990535248212 + offset,
        2.4353923087262066 + offset,
        2.697340999301913 + offset,
        2.6521161619740408 + offset,
        0.8016750570795748 + offset,
        2.0888346508237094 + offset,
        0.8304001118980417 + offset,
        2.1076008673731175 + offset,
        0.3589801618753179 + offset,
        1.1229934118941256 + offset,
        1.3365490905708357 + offset,
        1.8354193619448524 + offset,
        2.718680142657365 + offset,
        1.114668177329429 + offset,
        2.4373081838238693 + offset,
        2.414593845611318 + offset,
        0.12290616957133377 + offset,
        2.432773932914067 + offset,
        2.703247300810249 + offset,
        0.9329542112336406 + offset,
        1.8387052373191828 + offset,
        2.192531994043116 + offset,
        0.9347932451321164 + offset,
        2.6506791700072787 + offset,
        1.4357888060935096 + offset,
        2.5074984951495582 + offset,
        1.452676080486036 + offset,
        2.1649169298468043 + offset,
        2.096317932113875 + offset,
        1.1626855473307327 + offset,
        0.0698015672586913 + offset,
        2.263884734272171 + offset,
        2.452097505628506 + offset,
        2.318232969547698 + offset,
        1.5708028678519623 + offset,
        1.7619872579964682 + offset,
        0.48741766950630805 + offset,
        2.2860590634348394 + offset,
        2.355561815069757 + offset,
        2.624438667886823 + offset,
        1.854267558549543 + offset,
        2.3213276658541218 + offset,
        2.745584480726928 + offset,
        0.11463689662493518 + offset,
        1.2268104662722987 + offset,
        0.4883558476641737 + offset,
        0.302830362550608 + offset,
        2.8418136908650427 + offset,
        2.4110065851996403 + offset,
        1.1883288319221876 + offset,
        0.4231560219611501 + offset,
        2.7332667742376247 + offset,
        2.4548799817800466 + offset,
        1.4492552216872536 + offset,
        0.29812737679036055 + offset,
        0.6024251082982855 + offset,
        2.4071077716684606 + offset,
        2.545623250801896 + offset,
        1.1447541103640393 + offset,
        1.6300111419299077 + offset,
        2.2144073814114598 + offset,
        1.741764282945776 + offset,
        2.584934801845699 + offset,
        2.94429370319911 + offset,
        0.04396728231852283 + offset,
        1.511277836871522 + offset,
        1.5078872546350166 + offset,
        2.2830235983573757 + offset,
        1.141833418913716 + offset,
        2.31153512617248 + offset,
        0.7930014449666454 + offset,
        1.7855431618709825 + offset,
        1.2496515034524522 + offset,
        2.494095987642121 + offset,
        0.99530344900322 + offset,
        1.6423078459022569 + offset,
        1.2698194368721236 + offset,
        2.3092312380669213 + offset,
        1.508197474828362 + offset,
        0.5282536325342755 + offset,
        2.8592120295683343 + offset,
        1.964949523096605 + offset,
        2.6235858300049006 + offset,
        1.2938945593806481 + offset,
        2.2949841542633416 + offset,
        2.347836965953938 + offset,
        2.3781640679457445 + offset,
        2.0047220333928424 + offset,
        0.5906909597410415 + offset,
        1.2986453274833423 + offset,
        2.667821411223028 + offset
    ]
    y = [
        1.0458976448210704,
        -3.9067277248882855,
        0.10622312924614075,
        -3.131178556347437,
        1.01212401276388,
        1.14213909231317,
        -1.499740679978164,
        -2.3757729401741026,
        -2.122671490779725,
        -3.502953002856569,
        -3.256081109078368,
        -3.1749520646753595,
        -2.155913453889749,
        -0.2395214626274446,
        0.9755334524010564,
        -3.471619676384854,
        0.9470748554316055,
        -1.8958980011427715,
        -1.750911067902624,
        0.6014389134684253,
        -0.328098369799075,
        -3.4065530755818263,
        -0.570775344749403,
        1.2583085822944136,
        0.37952520933453554,
        -1.3523455959069988,
        -2.9400776134322326,
        -0.931511238718386,
        0.9834249491123324,
        -3.5025716943059186,
        -4.248785250853195,
        -0.6773777546863811,
        -3.146614604919824,
        1.6069751824711314,
        0.12949305293031732,
        0.8261796714411879,
        -3.569524692879105,
        0.7005037664902201,
        -2.425440866187578,
        -3.738092107397064,
        -3.2452624729025747,
        -1.368470502749851,
        0.07815030521334698,
        -2.430664435412014,
        -0.3937883060575643,
        0.9090318588982524,
        -3.2871385464785274,
        -1.0695984302841914,
        0.24434407871756103,
        -2.25954235705039,
        -0.7439150005837414,
        -3.0397095698565275,
        0.9359563676278917,
        1.0115209086189823,
        -3.3783557249585123,
        0.9594224070734308,
        -0.10771629916989434,
        -1.7184897675482067,
        -2.708724275370595,
        1.6742944306376044,
        -2.2521686832796735,
        -0.3155915671715026,
        -1.3799617068714918,
        -3.7029324960425725,
        0.6295951027167603,
        1.3620703160533432,
        -2.6320189409595427,
        1.2235177121091534,
        -3.3223373231734046,
        0.7410339357717708,
        0.5026708223105636,
        -1.2659346176335473,
        1.1018873624037109,
        -1.7793888239691904,
        -0.6335066629431028,
        0.7436758886663175,
        -2.271922352203361,
        -1.5681088934930953,
        1.387487453946818,
        -0.18205302519419486,
        0.7509498665446229,
        1.1304461649366861,
        0.2598239003357525,
        -3.638173831524141,
        -0.3246237440748867,
        -3.846851795839017,
        0.12598502378608517,
        -2.974937206560519,
        -3.611991741390597,
        -3.6894705097537193,
        -2.792874862433873,
        -2.998549361725964,
        -2.520540198719894,
        0.6171082788895208,
        -3.1598996224476084,
        1.5195493259285597,
        -1.9538102221406346,
        -2.061630645113948,
        -0.45792310107331247,
        0.07907581142523268,
        -0.10696635485443987,
        -1.710269012859103,
        -3.60107422436533,
        1.1092227889389914,
        -3.7822078275343984,
        -4.21797263754794,
        -1.719272711473141,
        -2.58454577962155,
        -3.436277286035127,
        0.9974824865464593,
        -1.2960555247563033,
        -2.098906157582724,
        -3.2980549567958297,
        -4.140712584756083,
        -3.6668594006359907,
        -1.9331827888832867,
        -1.6997537944929542,
        1.2288964589068012,
        -3.280361193419007,
        -2.7915168449133825,
        0.716151381796484,
        -0.4617926598386114,
        -0.7751800733034873,
        1.5080467962694248,
        -3.601524555444186,
        0.4038862638181287,
        -2.6698628730744787,
        -3.534713771642897,
        -1.5036639589034162,
        0.9984399344555597,
        0.3775513743392528,
        -3.61494166755954,
        -0.6561789952747799,
        -3.9953414791058472,
        0.8817193140291141,
        1.2639173654508489,
        0.6932869760752849,
        -3.4031870404877793,
        -3.6982857231300157,
        -2.4260474342934755,
        -3.3108893959920254,
        -0.7863371260473055,
        0.7098721965355509,
        0.21815437229061818,
        -2.078371734274718,
        -1.1400343437416685,
        -1.2769894466119114,
        1.6438700737392897,
        -3.0090552847647456,
        0.2317409067300455
    ]

    return vcat(transpose.((x, y))...) / coord_down_scale
end

end


text:
using VectorSphericalWaves
using StaticArrays
using ComplexOperations

import VectorSphericalWaves.M_mn_wave_SeparateRealImag_SMatrix
m, n, kr_r, kr_i, θ, ϕ = 1, 1, 1e5, 1e2, 0.2, 0.3
kr = complex(kr_r, kr_i)

M_mn_wave_SeparateRealImag(m, n, kr_r, kr_i, θ, ϕ, "regular");
N_mn_wave_SeparateRealImag(m, n, kr_r, kr_i, θ, ϕ, "regular");

M_mn_wave_SeparateRealImag_SMatrix(m, n, kr_r, kr_i, θ, ϕ, "regular");
N_mn_wave_SeparateRealImag_SMatrix(m, n, kr_r, kr_i, θ, ϕ, "regular");

# BigFloat
m, n, kr_r, kr_i, θ, ϕ = 1, 1, BigFloat(1e6), BigFloat(1e2), BigFloat(0.2), BigFloat(0.3);
kr = complex(kr_r, kr_i);

M_mn_wave(m, n, kr, θ, ϕ; kind="regular");
N_mn_wave(m, n, kr, θ, ϕ; kind="regular");