text:
module Sessions

using SHA, HTTP, Dates, Nullables
using Genie, Genie.Cookies, Genie.Loggers

mutable struct Session
  id::String
  data::Dict{Symbol,Any}
end
Session(id::String) = Session(id, Dict{Symbol,Any}())

export Session

const session_adapter_name = string(Genie.config.session_storage) * "SessionAdapter"
include("session_adapters/$session_adapter_name.jl")
Core.eval(@__MODULE__, Meta.parse("using .$session_adapter_name"))
const SessionAdapter = Core.eval(@__MODULE__, Meta.parse(session_adapter_name))

"""
    id() :: String
    id(req::Request) :: String
    id(req::Request, res::Response) :: String

Generates a unique session id.
"""
function id() :: String
  if ! isdefined(Genie, :SECRET_TOKEN)
    log("Session error", :warn)
    log("$(@__FILE__):$(@__LINE__)", :warn)

    if ! Genie.Configuration.is_prod()
      log("Generating temporary secret token", :warn)
      Core.eval(Genie, :(const SECRET_TOKEN = $(Genie.REPL.secret_token())))
    else
      error("Can't compute session id - please make sure SECRET_TOKEN is defined in config/secrets.jl")
    end
  end

  try
    Genie.SECRET_TOKEN * ":" * bytes2hex(sha1(string(Dates.now()))) * ":" * string(rand()) * ":" * string(hash(Genie)) |> sha256 |> bytes2hex
  catch ex
    log("Session error", :err)
    log("$(@__FILE__):$(@__LINE__)", :err)
    error("Can't compute session id - please make sure SECRET_TOKEN is defined in config/secrets.jl")
  end
end
function id(req::HTTP.Request) :: String
  ! isnull(Cookies.get(req, Genie.config.session_key_name)) &&
    ! isempty(Base.get(Cookies.get(req, Genie.config.session_key_name))) &&
      return Base.get(Cookies.get(req, Genie.config.session_key_name))

  id()
end
function id(req::HTTP.Request, res::HTTP.Response) :: String
  ! isnull(Cookies.get(res, Genie.config.session_key_name)) &&
    ! isempty(Base.get(Cookies.get(res, Genie.config.session_key_name))) &&
      return Base.get(Cookies.get(res, Genie.config.session_key_name))

  ! isnull(Cookies.get(req, Genie.config.session_key_name)) &&
    ! isempty(Base.get(Cookies.get(req, Genie.config.session_key_name))) &&
      return Base.get(Cookies.get(req, Genie.config.session_key_name))

  id()
end

"""
    start(session_id::String, req::Request, res::Response; options = Dict{String,String}()) :: Session
    start(req::Request, res::Response) :: Session

Initiates a session.
"""
function start(session_id::String, req::HTTP.Request, res::HTTP.Response; options = Dict{String,String}()) :: Session
  options = merge(Dict("Path" => "/", "HttpOnly" => true, "Expires" => Dates.today()), options)
  Cookies.set!(res, Genie.config.session_key_name, session_id, options)

  load(session_id)
end
function start(req::HTTP.Request, res::HTTP.Response) :: Session
  start(id(req, res), req, res)
end

"""
    set!(s::Session, key::Symbol, value::Any) :: Session

Stores `value` as `key` on the `Session` `s`.
"""
function set!(s::Session, key::Symbol, value::Any) :: Session
  s.data[key] = value

  s
end

"""
    get(s::Session, key::Symbol) :: Nullable

Returns the value stored on the `Session` `s` as `key`, wrapped in a `Nullable`.
"""
function get(s::Session, key::Symbol) :: Nullable
  return  if haskey(s.data, key)
            Nullable(s.data[key])
          else
            Nullable()
          end
end

"""
    get!!(s::Session, key::Symbol)

Attempts to read the value stored on the `Session` `s` as `key` - throws an exception if the `key` does not exist.
"""
function get!!(s::Session, key::Symbol)
  s.data[key]
end

"""
    unset!(s::Session, key::Symbol) :: Session

Removes the value stored on the `Session` `s` as `key`.
"""
function unset!(s::Session, key::Symbol) :: Session
  delete!(s.data, key)

  s
end

"""
    is_set(s::Session, key::Symbol) :: Bool

Checks wheter or not `key` exists on the `Session` `s`.
"""
function is_set(s::Session, key::Symbol) :: Bool
  haskey(s.data, key)
end

"""
    persist(s::Session) :: Session

Generic method for persisting session data - delegates to the underlying `SessionAdapter`.
"""
function persist(s::Session) :: Session
  SessionAdapter.write(s)

  s
end

"""
    load(session_id::String) :: Session

Loads session data from persistent storage - delegates to the underlying `SessionAdapter`.
"""
function load(session_id::String) :: Session
  session = SessionAdapter.read(session_id)

  if isnull(session)
    return Session(session_id)
  else
    return Base.get(session)
  end
end

"""
    session() :: Sessions.Session
    session(params::Dict{Symbol,Any}) :: Sessions.Session

Returns the `Session` object associated with the current HTTP request.
"""
function session() :: Sessions.Session
  session(Router._params_())
end
function session(params::Dict{Symbol,Any}) :: Sessions.Session
  if haskey(params, Genie.PARAMS_SESSION_KEY)
    return params[Genie.PARAMS_SESSION_KEY]
  else
    msg = "Invalid params Dict -- must have $(Genie.PARAMS_SESSION_KEY) key"
    log(msg, :err)
    error(msg)
  end
end

end


text:
include("Manifold.jl")
const MANIFOLD_TYPE = "manifolds.manifold"
AetherlangObject(m::ManifoldsManifold) = AetherlangObject{ManifoldsManifold}(m, Ref{String}(MANIFOLD_TYPE))

function manifolds_manifold(line::Ref{UInt}, ns::Namespace, args...)::AetherlangObject
    exp::Vector{Token} = Token[(haskey(ns, t) ? ns[t] : t) for t in aetherlang_tokenize(args[end].dataref[2:end-1])]
    AetherlangObject(ManifoldsManifold([args[argi].dataref for argi in 1:(length(args)-1)], exp))
end

# make manifold slicing
function AE_section(line::Ref{UInt}, ns::Namespace, m::AetherlangObject{ManifoldsManifold}, args...)::AetherlangObject
    if length(args) % 2 != 0
        throw(AetherlangError("Wrong number of arguments given to manifolds.section"))
    end
    nm::ManifoldsManifold = deepcopy(m.dataref)
    i::Int = 1
    while i <= length(args)
        popat!(nm.vars, indexin([args[i].dataref], nm.vars)[1])
        nm.exp = Token[t==args[i].dataref ? args[i+1] : t for t in nm.exp]
        i += 2
    end
    AetherlangObject(nm)
end

# overload from base
function AE_contains(line::Ref{UInt}, ns::Namespace, args::AetherlangObject, m::AetherlangObject{ManifoldsManifold})::AetherlangObject
    nns::Namespace = copy(ns)
    for i::Int in 1:length(m.dataref.vars)
        nns[m.dataref.vars[i]] = AetherlangObject(args.dataref[i])
    end
    aetherlang_eval!(m.dataref.exp, nns, Ref{UInt}(0))
end

manifolds_namespace_modify = Namespace(
    "manifold" => AetherlangObject(manifolds_manifold),
    "section" => AetherlangObject(AE_section)
)

ext: jl

max_stars_repo_name: Aetherlang/aetherlang

max_issues_repo_name: gojakuch/aetherlang

max_forks_repo_name: gojakuch/aetherlang


text:
module Dispatch
export TopMetaTable, MetaMethodTable, MetaMethod, MetaMethodError,
       getmethod, newmethod!, removemethod!, prefermethod!, prefermethod_over!,
       @macromethod, @metafunction, @metadestruct, @metadispatch, @metamodule

include("Structure.jl")
include("BaseImplementations.jl")
include("TableManipulation.jl")
include("TopMetaTables.jl")
include("Applications.jl")
include("Reflection.jl")
include("MetaModule.jl")

using .Structure
using .TableManipulation
using .Applications
using .MetaModule

end


text:
"""
This module is an adapter to the AWS Lex Runtime. It provides
application-specific types and definitions for values returned by that service.

The main entrypoint for this module is the `PostText` method which returns a
`LexResponse`.
"""
module LexRuntime

using AWS: @service
using AWS: AWSConfig
using AWS: global_aws_config
using AWS.AWSServices: lex_runtime_service

export LexResponse

include("./DialogStates.jl")
include("./Intents.jl")

const botAlias = "alpha"
const botName = "TrialScheduler"

"""
A structured representation of the AWS Lex service's response. Data which isn't
used by this application is not included.
"""
struct LexResponse{DS <: DialogState, I <: Intent}
    dialog_state::DS
    intent::I
    slots::AbstractDict{String, String}
    message::Union{Some{String}, Nothing}
end

"Build a Lex Response from the raw return values from the lex service."
function build_response(raw::AbstractDict)::LexResponse
    @debug "$raw"

    message = raw["message"] === nothing ?
        nothing :
        Some(raw["message"])

    slots = raw["slots"] === nothing ?
        Dict{String, String}() :
        Dict{String, String}( k => string(v) for (k, v) in raw["slots"] )

    dialog_state = get(DialogStateNames, raw["dialogState"], Failed())
    intent = get(IntentNames, raw["intentName"], NothingIntent())

    LexResponse(dialog_state, intent, slots, message)
end

"Post a user's message to the conversational lex bot."
function PostText(user_message::String, userId::String)::LexResponse
    lex_runtime_service(
        "POST",
        "/bot/$(botName)/alias/$(botAlias)/user/$(userId)/text",
        Dict{String, Any}("inputText"=>user_message);
        aws_config=global_aws_config()
    ) |> build_response
end

end


text:
include("ir_dump.jl")
export ir_dump

include("fld.jl")
export fld_read


text:
# play with julia, using laplace 3d kernel eval.
# Barnett 9/18/18.  After Steven Johnson tweaking, 2/20/20.

#= comment block
=#

# How to run:
# mutlithreading: must start julia with
# JULIA_NUM_THREADS=8 julia

# then, to run:  include("lap3dkernel.jl")

using LinearAlgebra
using Printf
using BenchmarkTools
using Base.Threads
# this is useless, apparently:
#using Devectorize

function eye(n)
    return Matrix(1.0I,n,n)
end

function lap3dcharge(y,q,x)       # vec over targs
    nt = size(x,2)
    pot = zeros(1,nt)    # note zeros(nt), col vec, fails to add r later
    prefac = 1/(4*pi)
    for j in eachindex(q)           # loop over srcs
        r2 = sum((x .- y[:,j]).^2,dims=1)      # sq dist
        r = sqrt.(r2)
        pot +=  prefac * q[j] ./ r
    end
    return pot
end

function lap3dcharge_devec(y,q,x)       # unwrap i,j.   C-style coding, SIMD
    nt = size(x,2)
    ns = size(y,2)
    pot = zeros(1,nt)    # note zeros(nt), col vec, fails to add r later
    prefac = 1/(4*pi)
@inbounds   for i = 1:nt      # targs
        @simd for j = 1:ns          # srcs
            #r2ij = sum((x[:,i] - y[:,j]).^2)      # sq dist - terrible!
            r2ij = (x[1,i]-y[1,j])^2+ (x[2,i]-y[2,j])^2+ (x[3,i]-y[3,j])^2
            rij = sqrt(r2ij)
            pot[i] +=  prefac * q[j] / rij
        end
    end
    return pot
end

function lap3dcharge_devec_par(y,q,x)   # multi-threaded version of above
    nt = size(x,2)
    pot = zeros(nt)    # note zeros(nt), col vec, fails to add r later
    prefac = 1/(4*pi)
    @threads for i in eachindex(pot)      # targs.
        @simd for j in eachindex(q)          # srcs
            @inbounds r2ij = (x[1,i]-y[1,j])^2+ (x[2,i]-y[2,j])^2+ (x[3,i]-y[3,j])^2
            #rij = sqrt(r2ij)
            @inbounds pot[i] +=  prefac * q[j] / sqrt(r2ij) #rij
        end
    end
    return pot
end

ns = 10000
nt = 10000
x = rand(3,nt)
y = rand(3,ns)
q = randn(ns)    # charges
t = @elapsed lap3dcharge(y,q,x)    # discards return value
t = @elapsed lap3dcharge(y,q,x)    # discards return value - is already compiled
#pot, t = @timed lap3dcharge(y,q,x)      # gets the time in secs
@printf("targ-vec: %d src-targ pairs in %.3g s: %.3g Gpair/s\n",ns*nt,t,ns*nt/t/1e9)
# 0.06 Gpair/s, ie 5x slower than devec

lap3dcharge_devec(y,q,x)   # compile it?
t = @elapsed lap3dcharge_devec(y,q,x)    # discards return value
t = @elapsed lap3dcharge_devec(y,q,x)    # discards return value
@printf("devec: %d src-targ pairs in %.3g s: %.3g Gpair/s\n",ns*nt,t,ns*nt/t/1e9)
# best, single-threaded,  0.33 Gpair/s   (py+numba was 1.3 since multithreaded)

t = @elapsed lap3dcharge_devec_par(y,q,x)    # discards return value
t = @elapsed lap3dcharge_devec_par(y,q,x)    # discards return value
@printf("devec par: %d src-targ pairs in %.3g s: %.3g Gpair/s\n",ns*nt,t,ns*nt/t/1e9)
# 1.26 Gpair/s, matches py+numba

# S. Johnson chat: use benchmark tools.
# @btime lap3dcharge_devec_par($y,$q,$x);
# the $ is a btime thing: means figure out the type of args before benchmarking
# ; suppresses output - but actually doesn't.

max_stars_repo_name: magland/floatingspeed

max_issues_repo_name: magland/floatingspeed

max_forks_repo_name: magland/floatingspeed


text:
"""
An implementation of "Polychronization: computation with spikes"
Izhikevich EM. Neural Comput. 2006 Feb;18(2):245-82. doi: 10.1162/089976606775093882. PMID: 16378515.
https://pubmed.ncbi.nlm.nih.gov/16378515/
"""
using FastSpike
using LinearAlgebra: I
using Plots
using ProgressMeter

const time = 10_000
const Nₑ = 800
const Nᵢ = 200
const N = Nₑ + Nᵢ
const p = 0.1
#Izhikevich neurons parameters
a = [0.02 * ones(1, Nₑ) 0.1 * ones(1, Nᵢ)]
b = 0.2
c = -65
d = [8 * ones(1, Nₑ) 2 * ones(1, Nᵢ)]
# form adjacency matrix
adjacency = rand(N, N) .< p
adjacency -= adjacency .* (zeros(N, N) + I) # remove recurrent connections
# Inhibitory/Excitetory weights: excitatory: 0 → 10, inhibitory: -5 (rigid)
min_weight, max_weight = EI(
    Nₑ, Nᵢ;
    E_E=(0, 10), E_I=(0, 10), I_E=(-5, -5), I_I=(-5, -5),
    adjacency=adjacency
)
# Define Network
net = DelayNetwork(
    Izhikevich(1, a, b, c, d, 30.0),
    STDP(0.1, 0.12, 20, 20; min_weight=min_weight, max_weight=max_weight)
)
neurons = add_group!(net, N)
# connections in neural group
weights = [6 * ones(Nₑ, Nₑ + Nᵢ); -5 * ones(Nᵢ, Nₑ + Nᵢ)]
weights .*= adjacency
delay = rand(1:20, N, N)
delay .*= adjacency
connect!(net, neurons, neurons, weights, adjacency, delay)
# Define Monitor to record network activities
monitor = Monitor()
# Define training function
function train(time)
    spiked_neurons = []
    @showprogress 1 "training " for t = 1:time
        random_input = zeros(1, N)
        random_input[rand(1:N)] = 20
        run!(net; input_voltage=random_input)
        record!(monitor, net)
        append!(spiked_neurons, [(t, i[2]) for i in findall(net.spikes)])
        if t % 1000 == 0
            plot(
                scatter(Tuple.(spiked_neurons), markersize=1, c=:black, title="time: $(t÷1000)s",),
                histogram(net.weight[1:800, :][convert(Matrix{Bool}, net.adjacency)[1:800, :]], title="Excitatory weight histogram");
                layout=grid(4, 1, heights=[0.75, 0.25])
            )
            spiked_neurons = []
        end
    end
    return
end
# Run
train(time)


text:
using MatrixPolynomials, LinearAlgebra
using Test

x = variable("x")

@testset "column nullities" begin
    J1 = [x -1]
    e, d, Ar, Br = colNullities(J1)
    @test e == [0, 1]
    @test colIndices(J1) == [1]
    @test isempty(d)
    @test isempty(Ar)
    @test isempty(Br)

    J2 = [x -1 0;
          0  x -1]

    e, d, Ar, Br = colNullities(J2)
    @test e == [0, 0, 1]
    @test colIndices(J2) == [2]
    @test isempty(d)
    @test isempty(Ar)
    @test isempty(Br)

    J3 = [zeros(2, 2) J2]
    e, d, Ar, Br = colNullities(J3)
    @test e == [2, 0, 1]
    @test colIndices(J3) == [0, 0, 2]
    @test isempty(d)
    @test isempty(Ar)
    @test isempty(Br)

    J = [x]
    e, d, Ar, Br = colNullities(J)
    @test isempty(e)
    @test isempty(d)
    @test Ar == hcat([0])
    @test Br == hcat([1])

    J = [-1]
    e, d, Ar, Br = colNullities(J)
    @test isempty(e)
    @test d == [1]
    @test isempty(Ar)
    @test isempty(Br)

    J = [x -1 0 0 0;
        0 0 -1 0 0;
        0 0 0 x-1 0]

    e, d, Ar, Br = colNullities(J)
    @test e == [1, 1]
    @test colIndices(J) == [0, 1]
    @test d == [1]

    J = [x -1 0 0 0 0 0;
         0 0 x -1 0 0 0
         0 0 0 0 -1 0 0;
         0 0 0 0 0 x-1 0]

    e, d, Ar, Br = colNullities(J)
    @test e == [1, 2]
    @test colIndices(J) == [0, 1, 1]
    @test d == [1]

    J = [x -1 0 0 0 0  0 0  0 0 0;
         0 0 x -1 0 0  0 0  0 0 0;
         0 0 0 0 -1 0  0 0  0 0 0;
         0 0 0 0 0 x-1 0 0  0 0 0;
         0 0 0 0 0 0   0 0 -1 0 0;
         0 0 0 0 0 0   0 0  x 0 0;
         0 0 0 0 0 0   0 0  0 -1 0
         0 0 0 0 0 0   0 0  0  x -1]

    e, d, Ar, Br = colNullities(J)
    @test e == [2, 2]
    @test colIndices(J) == [0, 0, 1, 1]
    @test d == [1, 1]
end

@testset "row nullities" begin
    J1 = [-1;x]
    e, d, Ar, Br = rowNullities(J1)
    @test e == [0, 1]
    @test rowIndices(J1) == [1]
    @test isempty(d)
    @test isempty(Ar)
    @test isempty(Br)

    J2 = [-1 0;
        x -1;
        0 x]

    e, d, Ar, Br = rowNullities(J2)
    @test e == [0, 0, 1]
    @test rowIndices(J2) == [2]
    @test isempty(d)
    @test isempty(Ar)
    @test isempty(Br)

    J3 = [zeros(2, 2); J2]
    e, d, Ar, Br = rowNullities(J3)
    @test e == [2, 0, 1]
    @test rowIndices(J3) == [0, 0, 2]
    @test isempty(d)
    @test isempty(Ar)
    @test isempty(Br)

    J = [x]
    e, d, Ar, Br = rowNullities(J)
    @test isempty(e)
    @test isempty(d)
    @test Ar == hcat([0])
    @test Br == hcat([1])

    J = [-1]
    e, d, Ar, Br = rowNullities(J)
    @test isempty(e)
    @test d == [1]
    @test isempty(Ar)
    @test isempty(Br)

    J = [-1 0 0;
        x 0 0;
        0 -1 0;
        0 0 x-1;
        0 0 0]

    e, d, Ar, Br = rowNullities(J)
    @test e == [1, 1]
    @test rowIndices(J) == [0, 1]
    @test d == [1]

    J = [-1 0 0 0;
        x 0 0 0;
        0 -1 0 0;
        0 x 0 0;
        0 0 -1 0;
        0 0 0 0]

    e, d, Ar, Br = rowNullities(J)
    @test e == [1, 2]
    @test rowIndices(J) == [0, 1, 1]
    @test d == [1]

    J = [-1 0 0 0 0;
          x 0 0 0 0;
          0 -1 0 0 0;
          0 x 0 0 0;
          0 0 -1 0 0;
          0 0 0 -1 0;
          0 0 0 x -1;
          0 0 0 0 0]

    e, d, Ar, Br = rowNullities(J)
    @test e == [1, 2]
    @test rowIndices(J) == [0, 1, 1]
    @test d == [1, 1]
end

@testset "eigenvalues computations" begin
    J = [x-1 0;
         0   x-2]
    @test eig(J) ≈ [1, 2]
    @test isempty(infMultiplicities(J))

    A = [x^2 1 0; 0 0 0;0 0 x]
    @test eig(A) ≈ [0, Inf]
    @test infMultiplicities(A) == [1]

    A = [1 0 0;0 x 0;0 0 x^2-x]
    @test eig(A) ≈ [0, 0, 1, Inf, Inf, Inf]
    @test infMultiplicities(A) == [1, 2]

    J = [x -1 0 0 0 0  0 0  0 0 0;
    0 0 x -1 0 0  0 0  0 0 0;
    0 0 0 0 -1 0  0 0  0 0 0;
    0 0 0 0 0 x-1 0 0  0 0 0;
    0 0 0 0 0 0   0 0 -1 0 0;
    0 0 0 0 0 0   0 0  x 0 0;
    0 0 0 0 0 0   0 0  0 -1 0
    0 0 0 0 0 0   0 0  0  x -1]
    @test eig(J) ≈ [1, Inf, Inf, Inf]
    @test infMultiplicities(J) == [1, 2]

    A = [x-1.1 0 0
         0    -1 0
         0     x -1]
    @test eig(A) ≈ [1.1, Inf, Inf]
    @test infMultiplicities(A) == [2]

    A = zeros(Pol{Int}, (14, 16))
    A[1,3:4] = [x -1]
    A[2:3, 5:7] = [x -1 0;0 x -1]
    A[5:8, 8:10] = [-1 0 0;x -1 0;0 x -1;0 0 x]
    A[9, 11] = -1
    A[10:11, 12:13] = [-1 0;x -1]
    A[12, 14] = x-2
    A[13:14, 15:16] = [x-3 0;-1 x-3]

    @test eig(A) ≈ [2, 3, 3, Inf, Inf, Inf]
    @test colIndices(A) == [0, 0, 1, 2]
    @test rowIndices(A) == [0, 3]
    @test infMultiplicities(A) == [1, 2]

    F = svd(rand(size(A)...), full=true)
    A1 = F.U*A*F.Vt
    @test eig(A1, 1e-6) ≈ [2, 3, 3, Inf, Inf, Inf]
end


text:
import Functors
import Flux

using Setfield: @set

#####
# ActorCritic
#####

export ActorCritic

"""
    ActorCritic(;actor, critic, optimizer=ADAM())
The `actor` part must return logits (*Do not use softmax in the last layer!*), and the `critic` part must return a state value.
"""
Base.@kwdef struct ActorCritic{A,C,O}
    actor::A
    critic::C
end

Functors.@functor ActorCritic

#####
# GaussianNetwork
#####

export GaussianNetwork

"""
    GaussianNetwork(;pre=identity, μ, logσ, min_σ=0f0, max_σ=Inf32, normalizer = tanh)

Returns `μ` and `logσ` when called.  Create a distribution to sample from using
`Normal.(μ, exp.(logσ))`. `min_σ` and `max_σ` are used to clip the output from
`logσ`. Actions are normalized according to the specified normalizer function.
"""
Base.@kwdef struct GaussianNetwork{P,U,S,F}
    pre::P = identity
    μ::U
    logσ::S
    min_σ::Float32 = 0.0f0
    max_σ::Float32 = Inf32
    normalizer::F = tanh
end

GaussianNetwork(pre, μ, logσ, normalizer=tanh) = GaussianNetwork(pre, μ, logσ, 0.0f0, Inf32, normalizer)

Functors.@functor GaussianNetwork

"""
This function is compatible with a multidimensional action space. When outputting an action, it uses the `normalizer` function to normalize it elementwise.

- `rng::AbstractRNG=Random.GLOBAL_RNG`
- `is_sampling::Bool=false`, whether to sample from the obtained normal distribution.
- `is_return_log_prob::Bool=false`, whether to calculate the conditional probability of getting actions in the given state.
"""
function (model::GaussianNetwork)(rng::AbstractRNG, s; is_sampling::Bool=false, is_return_log_prob::Bool=false)
    x = model.pre(s)
    μ, raw_logσ = model.μ(x), model.logσ(x)
    logσ = clamp.(raw_logσ, log(model.min_σ), log(model.max_σ))
    if is_sampling
        σ = exp.(logσ)
        z = Zygote.ignore() do
            noise = randn(rng, Float32, size(μ))
            model.normalizer.(μ .+ σ .* noise)
        end
        if is_return_log_prob
            logp_π = sum(normlogpdf(μ, σ, z) .- (2.0f0 .* (log(2.0f0) .- z .- softplus.(-2.0f0 .* z))), dims=1)
            return z, logp_π
        else
            return z
        end
    else
        return μ, logσ
    end
end

"""
    (model::GaussianNetwork)(rng::AbstractRNG, state, action_samples::Int)
Sample `action_samples` actions from each state. Returns a 3D tensor with dimensions (action_size x action_samples x batch_size).
`state` must be 3D tensor with dimensions (state_size x 1 x batch_size). Always returns the logpdf of each action along.
"""
function (model::GaussianNetwork)(rng::AbstractRNG, s, action_samples::Int)
    x = model.pre(s)
    μ, raw_logσ = model.μ(x), model.logσ(x)
    logσ = clamp.(raw_logσ, log(model.min_σ), log(model.max_σ))

    σ = exp.(logσ)
    z = Zygote.ignore() do
        noise = randn(rng, Float32, (size(μ, 1), action_samples, size(μ, 3))...)
        model.normalizer.(μ .+ σ .* noise)
    end
    logp_π = sum(normlogpdf(μ, σ, z) .- (2.0f0 .* (log(2.0f0) .- z .- softplus.(-2.0f0 .* z))), dims=1)
    return z, logp_π
end

function (model::GaussianNetwork)(state; is_sampling::Bool=false, is_return_log_prob::Bool=false)
    model(Random.GLOBAL_RNG, state; is_sampling=is_sampling, is_return_log_prob=is_return_log_prob)
end

function (model::GaussianNetwork)(state, action_samples::Int)
    model(Random.GLOBAL_RNG, state, action_samples)
end

function (model::GaussianNetwork)(state, action)
    x = model.pre(state)
    μ, raw_logσ = model.μ(x), model.logσ(x)
    logσ = clamp.(raw_logσ, log(model.min_σ), log(model.max_σ))
    σ = exp.(logσ)
    logp_π = sum(normlogpdf(μ, σ, action) .- (2.0f0 .* (log(2.0f0) .- action .- softplus.(-2.0f0 .* action))), dims=1)
    return logp_π
end

#####
# CovGaussianNetwork
#####

export CovGaussianNetwork

"""
    CovGaussianNetwork(;pre=identity, μ, Σ, normalizer = tanh)

Returns `μ` and `Σ` when called where μ is the mean and Σ is a covariance
matrix. Unlike GaussianNetwork, the output is 3-dimensional.  μ has dimensions
(action_size x 1 x batch_size) and Σ has dimensions (action_size x action_size x
batch_size).  The Σ head of the `CovGaussianNetwork` should not directly return
a square matrix but a vector of length `action_size x (action_size + 1) ÷ 2`.
This vector will contain elements of the uppertriangular cholesky decomposition
of the covariance matrix, which is then reconstructed from it.  Sample from
`MvNormal.(μ, Σ)`. Actions are normalized elementwise according to the specified
normalizer function.
"""
mutable struct CovGaussianNetwork{P,U,S,F}
    pre::P
    μ::U
    Σ::S
    normalizer::F
end

CovGaussianNetwork(pre, m, s) = CovGaussianNetwork(pre, m, s, tanh)

Functors.@functor CovGaussianNetwork

"""
    (model::CovGaussianNetwork)(rng::AbstractRNG, state; is_sampling::Bool=false, is_return_log_prob::Bool=false)

This function is compatible with a multidimensional action space. When
outputting a sampled action, it uses the `normalizer` function to normalize it
elementwise.  To work with covariance matrices, the outputs are 3D tensors.  If
sampling, return an actions tensor with dimensions (action_size x action_samples
x batch_size) and logp_π (1 x action_samples x batch_size) If not, returns μ
with dimensions (action_size x 1 x batch_size) and L, the lower triangular of
the cholesky decomposition of the covariance matrix, with dimensions
(action_size x action_size x batch_size) The covariance matrices can be
retrieved with `Σ = Flux.stack(map(l -> l*l', eachslice(L, dims=3)),3)`

- `rng::AbstractRNG=Random.GLOBAL_RNG`
- `is_sampling::Bool=false`, whether to sample from the obtained normal
  distribution.
- `is_return_log_prob::Bool=false`, whether to calculate the conditional
  probability of getting actions in the given state.
"""
function (model::CovGaussianNetwork)(rng::AbstractRNG, state; is_sampling::Bool=false, is_return_log_prob::Bool=false)
    batch_size = size(state, 3)
    x = model.pre(state)
    μ, cholesky_vec = model.μ(x), model.Σ(x)
    da = size(μ, 1)
    L = vec_to_tril(cholesky_vec, da)

    if is_sampling
        z = Zygote.ignore() do
            noise = randn(rng, eltype(μ), da, 1, batch_size)
            model.normalizer.(Flux.stack(map(.+, eachslice(μ, dims=3), eachslice(L, dims=3) .* eachslice(noise, dims=3)), 3))
        end
        if is_return_log_prob
            logp_π = mvnormlogpdf(μ, L, z)
            return z, logp_π
        else
            return z
        end
    else
        return μ, L
    end
end

"""
    (model::CovGaussianNetwork)(rng::AbstractRNG, state::AbstractMatrix; is_sampling::Bool=false, is_return_log_prob::Bool=false)

Given a Matrix of states, will return actions, μ and logpdf in matrix format. The batch of Σ remains a 3D tensor.
"""
function (model::CovGaussianNetwork)(rng::AbstractRNG, state::AbstractMatrix; is_sampling::Bool=false, is_return_log_prob::Bool=false)
    output = model(rng, Flux.unsqueeze(state, 2); is_sampling=is_sampling, is_return_log_prob=is_return_log_prob)
    if output isa Tuple && is_sampling
        dropdims(output[1], dims=2), dropdims(output[2], dims=2)
    elseif output isa Tuple
        dropdims(output[1], dims=2), output[2] #can't reduce the dims of the covariance tensor
    else
        dropdims(output, dims=2)
    end
end

"""
    (model::CovGaussianNetwork)(rng::AbstractRNG, state, action_samples::Int)

Sample `action_samples` actions given `state` and return the `actions,
logpdf(actions)`.  This function is compatible with a multidimensional action
space. When outputting a sampled action, it uses the `normalizer` function to
normalize it elementwise.  The outputs are 3D tensors with dimensions
(action_size x action_samples x batch_size) and (1 x action_samples x
batch_size) for `actions` and `logdpf` respectively.
"""
function (model::CovGaussianNetwork)(rng::AbstractRNG, state, action_samples::Int)
    batch_size = size(state, 3) #3
    x = model.pre(state)
    μ, cholesky_vec = model.μ(x), model.Σ(x)
    da = size(μ, 1)
    L = vec_to_tril(cholesky_vec, da)
    z = Zygote.ignore() do
        noise = randn(rng, eltype(μ), da, action_samples, batch_size)
        model.normalizer.(Flux.stack(map(.+, eachslice(μ, dims=3), eachslice(L, dims=3) .* eachslice(noise, dims=3)), 3))
    end
    logp_π = mvnormlogpdf(μ, L, z)
    return z, logp_π
end

function (model::CovGaussianNetwork)(state::AbstractArray, args...; kwargs...)
    model(Random.GLOBAL_RNG, state, args...; kwargs...)
end

"""
    (model::CovGaussianNetwork)(state, action)

Return the logpdf of the model sampling `action` when in `state`.  State must be
a 3D tensor with dimensions (state_size x 1 x batch_size).  Multiple actions may
be taken per state, `action` must have dimensions (action_size x
action_samples_per_state x batch_size) Returns a 3D tensor with dimensions (1 x
action_samples_per_state x batch_size)
"""
function (model::CovGaussianNetwork)(state::AbstractArray, action::AbstractArray)
    da = size(action, 1)
    x = model.pre(state)
    μ, cholesky_vec = model.μ(x), model.Σ(x)
    L = vec_to_tril(cholesky_vec, da)
    logp_π = mvnormlogpdf(μ, L, action)
    return logp_π
end

"""
If given 2D matrices as input, will return a 2D matrix of logpdf. States and
actions are paired column-wise, one action per state.
"""
function (model::CovGaussianNetwork)(state::AbstractMatrix, action::AbstractMatrix)
    output = model(Flux.unsqueeze(state, 2), Flux.unsqueeze(action, 2))
    return dropdims(output, dims=2)
end

"""
Transform a vector containing the non-zero elements of a lower triangular da x da matrix into that matrix.
"""
function vec_to_tril(cholesky_vec, da)
    batch_size = size(cholesky_vec, 3)
    c2idx(i, j) = ((2da - j) * (j - 1)) ÷ 2 + i #return the position in cholesky_vec of the element of the triangular matrix at coordinates (i,j)
    function f(j) #return a slice (da x 1 x batchsize) containing the jth columns of the lower triangular cholesky decomposition of the covariance
        tc_diag = softplus.(cholesky_vec[c2idx(j, j):c2idx(j, j), :, :])
        tc_other = cholesky_vec[c2idx(j, j)+1:c2idx(j + 1, j + 1)-1, :, :]
        zs = Flux.Zygote.ignore() do
            zs = similar(cholesky_vec, da - size(tc_other, 1) - 1, 1, batch_size)
            zs .= zero(eltype(cholesky_vec))
            return zs
        end
        [zs; tc_diag; tc_other]
    end
    return mapreduce(f, hcat, 1:da)
end

#####
# DuelingNetwork
#####

export DuelingNetwork

"""
    DuelingNetwork(;base, val, adv)

Dueling network automatically produces separate estimates of the state value function network and advantage function network. The expected output size of val is 1, and adv is the size of the action space.
"""
Base.@kwdef struct DuelingNetwork{B,V,A}
    base::B
    val::V
    adv::A
end

Flux.@functor DuelingNetwork

function (m::DuelingNetwork)(state)
    x = m.base(state)
    val = m.val(x)
    return val .+ m.adv(x) .- mean(m.adv(x), dims=1)
end

#####
# PerturbationNetwork
#####

export PerturbationNetwork

"""
    PerturbationNetwork(;, ϕ)

Perturbation network outputs an adjustment to an action in the range [-ϕ, ϕ] to increase the diversity of seen actions.

# Keyword arguments

- `base`, a Flux based DNN model.
- `ϕ::Float32 = 0.05f0`
"""

Base.@kwdef struct PerturbationNetwork{N}
    base::N
    ϕ::Float32 = 0.05f0
end

Flux.@functor PerturbationNetwork

"""
This function accepts `state` and `action`, and then outputs actions after disturbance.
"""
function (model::PerturbationNetwork)(state, action)
    x = model.base(vcat(state, action))
    x = model.ϕ * tanh.(x)
    clamp.(x + action, -1.0f0, 1.0f0)
end

#####
# VAE (Variational Auto-Encoder)
#####

export VAE

"""
    VAE(;encoder, decoder, latent_dims)
"""
Base.@kwdef struct VAE{E,D}
    encoder::E
    decoder::D
    latent_dims::Int
end

Flux.@functor VAE

function (model::VAE)(rng::AbstractRNG, state, action)
    μ, logσ = model.encoder(vcat(state, action))
    σ = exp.(logσ)
    z = μ .+ σ .* randn(rng, Float32, size(μ))
    u = decode(model, state, z)
    return u, μ, σ
end

function (model::VAE)(state, action)
    return model(Random.GLOBAL_RNG, state, action)
end

function decode(rng::AbstractRNG, model::VAE, state, z=nothing; is_normalize::Bool=true)
    if z === nothing
        z = clamp.(randn(rng, Float32, (model.latent_dims, size(state)[2:end]...)), -0.5f0, 0.5f0)
    end
    a = model.decoder(vcat(state, z))
    if is_normalize
        a = tanh.(a)
    end
    return a
end

function decode(model::VAE, state, z=nothing; is_normalize::Bool=true)
    decode(Random.GLOBAL_RNG, model, state, z; is_normalize)
end

function vae_loss(model::VAE, state, action)
    u, μ, σ = model(state, action)
    recon_loss = Flux.Losses.mse(u, action)
    kl_loss = -0.5f0 * mean(1.0f0 .+ log.(σ .^ 2) .- μ .^ 2 .- σ .^ 2)
    return recon_loss, kl_loss
end

#####
# TwinNetwork
#####

export TwinNetwork

Base.@kwdef mutable struct TwinNetwork{S,T}
    source::S
    target::T
    sync_freq::Int = 1
    ρ::Float32 = 0.0f0
    n_optimise::Int = 0
end

TwinNetwork(x; kw...) = TwinNetwork(; source=x, target=deepcopy(x), kw...)

Functors.functor(x::TwinNetwork) = (; source=x.source), y -> @set x.source = y.source

(model::TwinNetwork)(x) = model.source(x)

function RLBase.optimise!(A::Approximator{<:TwinNetwork}, gs)
    Flux.Optimise.update!(A.optimiser, Flux.params(A), gs)
    M = A.model
    M.n_optimise += 1

    if M.n_optimise % M.sync_freq == 0
        # polyak averaging
        for (dest, src) in zip(Flux.params(M.target), Flux.params(M.source))
            dest .= M.ρ .* dest .+ (1 - M.ρ) .* src
        end
    end
end


text:
module PlutoUtils

end


text:
import PyPlot: subplots

type VStyle
    fill::Union{RGB, Void}
    function VStyle(;fill=nothing)
        new(fill)
    end
end
style2mplkw(s::VStyle) = Dict(:fill => s.fill)

##Note: Mathematica and Matplotlib don't have nested scopes for the viewports, you can lay
## them over each other, but not inside one another. I am going to go this route, the
## full nested version is too complex, and I am not sure it buys me much.
type Viewport # a Viewport is not a grob, and therefore can't contain another Viewport
    grobs::Array{Grob}
    # to start I only support npc coordinates for viewports, this needs to change so I can
    # have inheritence for sub-viewports
    xy::Array{Float64}
    width::Float64
    height::Float64
    clip::Symbol
    style::VStyle # should likely be a special viewport style
    # justification
    # xscale
    # yscale
    # angle
    # layout
    # layout.pos.row
    # layout.pos.col
    # name # I think I should use id like SVG
    function Viewport(grobs; xy=[0.0, 0.0], width=1.0, height=1.0, clip=:on, style=VStyle())
        new(grobs, xy, width, height, clip, style)
    end
end

# a Viewport needs to know the Canvas it is in, so the render has an extra argument,
# I think I might need to change the name of render, to something more descriptive
function render!(f, v::Viewport)
    ax = f.add_axes(v.xy[1], v.xy[2], v.xy[1] + width, v.xy[2] + height)
    _addgrob2mpl!(ax, v.grobs)
    return nothing
end

#TODO: need to allow for changes to attributes, but not yet sure what the api shoudl be
Base.push!(vp::Viewport, grob::Grob) = push!(vp.grobs, grob)

type Canvas
    grobs::Array{Grob}
    style::Style
    function Canvas(grobs; style=Style(fill=NC"white"))
        new(grobs, style)
    end
    # function Canvas(vps::Array{Viewport}; style=Style(fill=NC"white"))
    #     new(vps)
    # end
end

function _addgrob2mpl!(ax, grobs::Array{Grob})
    for g in grobs
        p = render(g)
        # assume that we have either lines or patches
        if pyisinstance(p, mlines.Line2D)
            ax[:add_line](p)
        else
            ax[:add_patch](p)
        end
    end
    return nothing
end

function render(go::Canvas)
    # using collections removes all styling, style must be passed into the PatchCollection
    # patches = []
    # for g in go.grobs
    #     push!(patches, render(g))
    # end
    #pcol = mcol.PatchCollection(patches)
    (fig, ax) = subplots()
    _addgrob2mpl!(ax, go.grobs)
    # apply style to axes
    ax[:patch][:set_facecolor](go.style.fill)
    ax[:relim]() # resize the axes limits
    ax[:autoscale_view]() # make sure everyting is visable
    ax[:set_aspect]("equal") # TODO: make this an option
    return fig
end

##TODO: not really sure this is a container, and not just a style, but it is an attribute
# of an AxesView
type Spine
    # this needs to deal with inward/outward styles as well as special "center", or
    # the different coordinate locations
    position::Vector{Float64}
    #bounds
    #visible::Bool # do I want this, is there a more consistent way to have this on/off
    #style::Style
end

type Spines
    top::Spine
    bottom::Spine
    left::Spine
    right::Spine
end

type AxesTicks
end

type Axes
    spines::Spines
    #ticks::AxesTicks
    #labels
end

type AxesView
    grobs::Array{Grob}
    axes::Axes
    style::Style
end

max_stars_repo_name: gabrielgellner/OdinSon

max_issues_repo_name: gabrielgellner/OdinSon

max_forks_repo_name: gabrielgellner/OdinSon


text:
macro shmem(T, Dims)
    dims = Dims.args
    esc(quote
        if $iscpu(__DEVICE)
            $MArray{Tuple{$(dims...)}, $T}(undef)
        else
            @cuStaticSharedMem($T, $Dims)
        end
    end)
end


text:
using FLuxSymbolEquations
using Test

@testset "FLuxSymbolEquations.jl" begin
    # Write your tests here.
    @test greet()=="Hello World!"

end