text:
using LinearAlgebra
using Flux
using CUDA
import Flux.gpu, Flux.cpu
using Distributions

includet("./solvers.jl")
includet("./fixed_point_layer_general.jl")
includet("./output_layer.jl")

# Implicit, direct parametrizations
abstract type implicit_param end

# Explicit types that we simulate
mutable struct explicit_ren_cell
    ϕ
    A
    B1
    B2
    C1
    D11
    D12
    bx
    bv
end

function (implicit_cell::implicit_param)(xt, ut, solver)
    explicit_cell = explicit(implicit_cell)
    return explicit_cell(xt, ut, solver)
end

# Forward method of rnn-cell
function (model::explicit_ren_cell)(xt, ut, solver::ForwardSolver)
    # First solve fixed point equation for v
    b = model.C1 * xt + model.D12 * ut .+ model.bv
    f(z, x) = model.D11 * model.ϕ.(z) + x
    vt = fixed_point_layer(solver, f, b)
    wt = model.ϕ.(vt)

    # next state and current output
    xn = model.A * xt + model.B1 * wt + model.B2 * ut .+ model.bx
    return xn, (xt, wt)
end

function simulate(explicit_param::explicit_ren_cell, x0, u, solver::ForwardSolver)
    eval_cell = (x, u) -> explicit_param(x, u, solver)
    recurrent = Flux.Recur(eval_cell, x0)
    return unzip(recurrent.(u))
end

function simulate(cell::implicit_param, x0, u, solver::ForwardSolver)
    explicit_param = explicit(cell)
    return simulate(explicit_param, x0, u, solver)
end

function (exp_cell::explicit_ren_cell)(xt, ut, solver::OperatorSplitting, V)
    # First solve fixed point equation for v
    b = exp_cell.C1 * xt + exp_cell.D12 * ut .+ exp_cell.bv
    wt = fixed_point_layer(solver, exp_cell.ϕ, exp_cell.D11, b, V)
    # wt = exp_cell.ϕ.(vt)

    # next state and current output
    xn = exp_cell.A * xt + exp_cell.B1 * wt + exp_cell.B2 * ut .+ exp_cell.bx
    return xn, (xt, wt)
end

function (exp_cell::explicit_ren_cell)(xt, ut, solver::OperatorSplitting)
    # If operator splitting but V is not provided
    V = ((1 + solver.α) * I - solver.α * exp_cell.D11)
    return exp_cell(xt, ut, solver, V)
end

function simulate(explicit_param::explicit_ren_cell, x0, u, solver::OperatorSplitting, V)
    eval_cell = (x, u) -> explicit_param(x, u, solver, V)
    recurrent = Flux.Recur(eval_cell, x0)
    return unzip(recurrent.(u))
end

function simulate(explicit_param::explicit_ren_cell, x0, u, solver::OperatorSplitting)
    V = ((1 + solver.α) * I - solver.α * explicit_param.D11)
    return simulate(explicit_param, x0, u, solver, V)
end

function simulate(cell::implicit_param, x0, u, solver::OperatorSplitting)
    explicit_param = explicit(cell)
    return simulate(explicit_param, x0, u, solver)
end

## Stable REN cell
mutable struct implicit_ren{T} <: implicit_param
    ϕ
    λ::Union{Vector{T},CuVector{T}}
    V::Union{Matrix{T},CuMatrix{T}}
    S_1::Union{Matrix{T},CuMatrix{T}}
    S_2::Union{Matrix{T},CuMatrix{T}}
    B2::Union{Matrix{T},CuMatrix{T}}
    D12::Union{Matrix{T},CuMatrix{T}}
    bx::Union{Vector{T},CuVector{T}}
    bv::Union{Vector{T},CuVector{T}}
    ϵ::T
end
Flux.trainable(L::implicit_ren) = [L.λ, L.V, L.S_1, L.S_2, L.B2, L.D12, L.bx, L.bv]

function Flux.gpu(M::implicit_ren{T}) where T
    if T != Float32
        println("Moving type: ", T, " to gpu may not be supported. Try Float32!")
    end
    return implicit_ren{T}(M.ϕ, gpu(M.λ), gpu(M.V), gpu(M.S_1), gpu(M.S_2),
                           gpu(M.B2), gpu(M.D12), gpu(M.bx), gpu(M.bv), M.ϵ)
end

function Flux.cpu(M::implicit_ren{T}) where T
    return implicit_ren{T}(M.ϕ, cpu(M.λ), cpu(M.V), cpu(M.S_1), cpu(M.S_2),
                           cpu(M.B2), cpu(M.D12), cpu(M.bx), cpu(M.bv), M.ϵ)
end

# TODO: This is currently not great as it does not account for dissipativity in Htild construction
function implicit_ren{T}(nu, nx, nv; nl=relu, ϵ=0.01, bx_scale=0.0, bv_scale=1.0) where T
    glorot_normal(n, m) = convert.(T, randn(n, m) / sqrt(n + m))
    E = Matrix{T}(I, nx, nx)
    F = Matrix{T}(I, nx, nx)
    P = Matrix{T}(I, nx, nx)
    B1 = zeros(T, nx, nv)
    B2 = glorot_normal(nx, nu)

    C1 = zeros(T, nv, nx)
    D11 = glorot_normal(nv, nv)
    D12 = zeros(T, nv, nu)

    λ = rand(T, nv)
    Λ = Diagonal(exp.(λ))
    H22 = 2Λ - D11 - D11'
    Htild = [(E + E' - P) -C1' F';
             -C1 H22 B1'
             F B1  P] + ϵ * I

    S_1 = glorot_normal(nx, nx)
    S_2 = glorot_normal(nv, nv)

    V = Matrix{T}(cholesky(Htild).U) # H = V'*V

    bv = convert.(T, bv_scale * randn(nv) / sqrt(nv))
    bx = convert.(T, bx_scale * rand(nx) / sqrt(nx))

    return implicit_ren{T}(nl, λ, V, S_1, S_2, B2, D12, bx, bv, ϵ)
end
implicit_ren(nu, nx, nv; nl=relu, ϵ=0.01, bx_scale=0.0, bv_scale=1.0) = implicit_ren{Float64}(nu, nx, nv; nl=nl, ϵ=ϵ, bx_scale=bx_scale, bv_scale=bv_scale)

function explicit(model::implicit_ren)
    nx = size(model.B2, 1)
    nu = size(model.B2, 2)
    nv = length(model.λ)

    H = model.V' * model.V + model.ϵ * I

    # For some reason taking view doesn't work with CUDA
    H11 = H[1:nx, 1:nx]
    H22 = H[ nx + 1:nx + nv, nx + 1:nx + nv]
    H33 = H[ nx + nv + 1:2nx + nv, nx + nv + 1:2nx + nv]
    H21 = H[nx + 1:nx + nv, 1:nx]
    H31 = H[ nx + nv + 1:2nx + nv, 1:nx]
    H32 = H[ nx + nv + 1:2nx + nv, nx + 1:nx + nv]

    # Implicit model parameters
    S_1 = (model.S_1 - model.S_1') / 2
    S_2 = (model.S_2 - model.S_2') / 2

    B2 = model.B2
    D12 = model.D12

    # There appears to be a bug in the zygote adjoint for diagm
    Λᵢ = exp.(model.λ)
    Λ = Diagonal(Λᵢ)

    P = H33
    E = (H11 + P + S_1) / 2
    F = H31

    # equilibrium network stuff
    C1 = - (H21)
    B1 = H32
    D11 = Λ - H22 / 2  - S_2

    # Construct explicit rnn model.
    # Use \bb to differentiate from implicit model
    𝔸 = E \ F
    𝔹_1 = E \ B1
    𝔹_2 = E \ B2

    ℂ_1 = (1 ./ Λᵢ) .* C1
    𝔻_11 = (1 ./ Λᵢ) .* D11
    𝔻_12 = (1 ./ Λᵢ) .* D12

    bx = model.bx
    bv = model.bv

    return explicit_ren_cell(model.ϕ, 𝔸, 𝔹_1, 𝔹_2, ℂ_1, 𝔻_11, 𝔻_12, bx, bv)
end

# Model seq2seq dissipative REN
mutable struct dissipative_ren{T} <: implicit_param
    nu
    nx
    nv
    ny
    implicit_ren::implicit_ren{T}
    output::output{T}
    Q::Matrix{T}
    S::Matrix{T}
    R::Matrix{T}
end

# return paramaters for recurrent and output layers
Flux.trainable(M::dissipative_ren) = (Flux.trainable(M.implicit_ren)..., Flux.trainable(M.output)...)

function dissipative_ren{T}(nu, nx, nv, ny, Q, S, R; nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0) where T
    ren_params = implicit_ren{T}(nu, nx, nv; nl=nl, ϵ, bx_scale, bv_scale)
    output_params = output{T}(nu, nx, nv, ny)
    return dissipative_ren{T}(nu, nx, nv, ny, ren_params, output_params, Q, S, R)
end

function init_state(model::dissipative_ren{T}, batches) where T
    return typeof(model.implicit_ren.V)(zeros(T, model.nx, batches))
end

function init_state(model::implicit_ren{T}, batches) where T
    nx = size(model.B2, 1)
    return typeof(model.V)(zeros(T, nx, batches))
end

function (model::dissipative_ren)(x0, u, solver::Solver)
    cell = explicit(model)
    # eval_cell = (x, u) -> cell(x, u, solver)
    x, w = simulate(cell, x0, u, solver)
    y = model.output.(x, w, u)
    return y, x[end]
end

function (model::dissipative_ren)(x0, u, solver::OperatorSplitting)
    cell = explicit(model)
    D11 = cell.D11
    V = ((1 + solver.α) * I - solver.α * D11)

    # eval_cell = (xt, ut) -> cell(xt, ut, solver, V) # Capture solver, V
    x, w = simulate(cell, x0, u, solver, V)
    y = model.output.(x, w, u)
    return y, x[end]
end

# Common forms of dissipativity
function bounded_ren(T::DataType, nu, nx, nv, ny, γ;
                        nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0)
    R = Matrix{T}(γ * I, nu, nu)
    S = zeros(T, nu, ny)
    Q = Matrix{T}(-I / γ, ny, ny)
    return dissipative_ren{T}(nu, nx, nv, ny, Q, S, R; nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0)
end

function stable_ren(T::DataType, nu, nx, nv, ny;
                        nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0)
    return bounded_ren(T, nu, nx, nv, ny, Inf; nl=nl, ϵ=ϵ, bx_scale=bx_scale, bv_scale=bv_scale)
end

function passive_ren(T::DataType, nu, nx, nv, ny;
                     η=0.0, nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0)

    R = zeros(T, nu, nu)
    S = Matrix(I / 2, nu, ny)
    Q = Matrix(-η * I, nu, nu)
    return dissipative_ren(nu, nx, nv, ny, Q, S, R; nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0)
end

# default type is Float64
bounded_ren(nu, nx, nv, ny, γ; nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0) =
                    bounded_ren(Float64, nu, nx, nv, ny, γ;
                            nl=nl, ϵ=ϵ, bx_scale=bx_scale, bv_scale=bv_scale)

stable_ren(nu, nx, nv, ny; nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0) =
                    stable_ren(Float64, nu, nx, nv, ny;
                                nl=nl, ϵ=ϵ, bx_scale=bx_scale, bv_scale=bv_scale)

passive_ren(nu, nx, nv, ny; η=0.0, nl=relu, ϵ=0.001, bx_scale=0.0, bv_scale=1.0) =
                    passive_ren(Float64, nu, nx, nv, ny;
                                        η=η, nl=nl, ϵ=ϵ, bx_scale=bx_scale, bv_scale=bv_scale)

# Construct and explicit
function explicit(model::dissipative_ren)
    nx = model.nx
    nu = model.nu
    ny = model.ny
    nv = model.nv

    # dissipation parameter
    Q = model.Q
    S = model.S
    R = model.R

    # Implicit model parameters
    S_1 = (model.implicit_ren.S_1 - model.implicit_ren.S_1') / 2
    S_2 = (model.implicit_ren.S_2 - model.implicit_ren.S_2') / 2

    C2 = model.output.C2
    D21 = model.output.D21
    D22 = model.output.D22

    B2 = model.implicit_ren.B2
    D12 = model.implicit_ren.D12

    # There appears to be a bug in the zygote adjoint for diagm
    Λᵢ = exp.(model.implicit_ren.λ)
    Λ = Diagonal(Λᵢ)

    # RHS of dissipation inequality
    Γ1 = [C2'; D21'; zeros(nx, ny)] * Q * [C2 D21 zeros(ny, nx)]  # possibly transpose zeros
    Γ2 = [(C2' * S'); (D21' * S' - D12); B2] * inv(R) * [(S * C2) (S * D21 - D12') B2']

    H = model.implicit_ren.V' * model.implicit_ren.V - Γ1 + Γ2 + model.implicit_ren.ϵ * I
    H11 = H[1:nx, 1:nx]
    H22 = H[ nx + 1:nx + nv, nx + 1:nx + nv]
    H33 = H[ nx + nv + 1:2nx + nv, nx + nv + 1:2nx + nv]
    H21 = H[nx + 1:nx + nv, 1:nx]
    H31 = H[ nx + nv + 1:2nx + nv, 1:nx]
    H32 = H[ nx + nv + 1:2nx + nv, nx + 1:nx + nv]

    # Extract parameters
    P = H33
    E = (H11 + P + S_1) / 2
    F = H31

    # equilibrium network stuff
    C1 = -H21
    B1 = H32
    D11 = Λ - H22 / 2  - S_2

    # Construct explicit rnn model. Use \bb font to differentiate
    A = E \ F
    𝔹_1 = E \ B1
    𝔹_2 = E \ B2

    ℂ_1 = (1 ./ Λᵢ) .* C1
    𝔻_11 = (1 ./ Λᵢ) .* D11
    𝔻_12 = (1 ./ Λᵢ) .* D12

    bx = model.implicit_ren.bx
    bv = model.implicit_ren.bv

    return explicit_ren_cell(model.implicit_ren.ϕ, A, 𝔹_1, 𝔹_2, ℂ_1, 𝔻_11, 𝔻_12, bx, bv)
end

function check_equilibrium_solve(model, solver::Solver; batches=10)
    explicit_model = explicit(model)

    nx, nu = size(model.B2)
    x0 = randn(nx, batches)
    u = randn(nu, batches)

    xn, (xt, wt) = explicit_model(x0, u, solver)

    bias = explicit_model.C1 * x0 + explicit_model.D12 * u .+ explicit_model.bv

    err = norm(wt - explicit_model.ϕ.(explicit_model.D11 * wt + bias))
    println("Checking tolerance of explicit solve: |w - ϕ(Dw + b)| = ", err)
    println("     Relative tolerance: |w - ϕ(Dw + b)| / |w| = ", err / norm(wt))
end

function check_lmi(model::implicit_ren)
    nx = size(model.B2, 1)
    nu = size(model.B2, 2)
    nv = length(model.λ)

    H = model.V' * model.V + model.ϵ * I

    # For some reason taking view doesn't work with CUDA
    H11 = H[1:nx, 1:nx]
    H22 = H[ nx + 1:nx + nv, nx + 1:nx + nv]
    H33 = H[ nx + nv + 1:2nx + nv, nx + nv + 1:2nx + nv]
    H21 = H[nx + 1:nx + nv, 1:nx]
    H31 = H[ nx + nv + 1:2nx + nv, 1:nx]
    H32 = H[ nx + nv + 1:2nx + nv, nx + 1:nx + nv]

    # Implicit model parameters
    S_1 = (model.S_1 - model.S_1') / 2
    S_2 = (model.S_2 - model.S_2') / 2

    B2 = model.B2
    D12 = model.D12

    # There appears to be a bug in the zygote adjoint for diagm
    Λᵢ = exp.(model.λ)
    Λ = Diagonal(Λᵢ)

    P = H33
    E = (H11 + P + S_1) / 2
    F = H31

    # equilibrium network stuff
    C1 = - (H21)
    B1 = H32
    D11 = Λ - H22 / 2  - S_2

    lmi = [(E + E' - P) (-C1') F';
    (-C1) (2Λ - D11 - D11') B1';
    F B1 P]
end

function check_lmi(model::dissipative_ren)
    nx = model.nx
    nu = model.nu
    ny = model.ny
    nv = model.nv

    # dissipation parameter
    Q = model.Q
    S = model.S
    R = model.R

    # Implicit model parameters
    S_1 = (model.implicit_ren.S_1 - model.implicit_ren.S_1') / 2
    S_2 = (model.implicit_ren.S_2 - model.implicit_ren.S_2') / 2

    C2 = model.output.C2
    D21 = model.output.D21
    D22 = model.output.D22

    B2 = model.implicit_ren.B2
    D12 = model.implicit_ren.D12

    # RHS of dissipation inequality
    Γ1 = [C2'; D21'; zeros(nx, ny)] * Q * [C2 D21 zeros(ny, nx)]  # possibly transpose zeros
    Γ2 = [(C2' * S'); (D21' * S' - D12); B2] * inv(R) * [(S * C2) (S * D21 - D12') B2']

    H = model.implicit_ren.V' * model.implicit_ren.V - Γ1 + Γ2 + model.implicit_ren.ϵ * I
    H11 = H[1:nx, 1:nx]
    H22 = H[ nx + 1:nx + nv, nx + 1:nx + nv]
    H33 = H[ nx + nv + 1:2nx + nv, nx + nv + 1:2nx + nv]
    H21 = H[nx + 1:nx + nv, 1:nx]
    H31 = H[ nx + nv + 1:2nx + nv, 1:nx]
    H32 = H[ nx + nv + 1:2nx + nv, nx + 1:nx + nv]

    # Extract parameters
    Λᵢ = exp.(model.implicit_ren.λ)
    Λ = Diagonal(Λᵢ)
    D11 = Λ - H22 / 2  - S_2
    P = H33
    E = (H11 + P + S_1) / 2
    F = H31
    C1 = -H21
    B1 = H32

    lmi = [(E + E' - P + C2' * Q * C2) (C2' * Q * D21 - C1') (C2' * S') F';
            (D21' * Q * C2 - C1) (2Λ - D11 - D11' + D21' * Q * D21) (D21' * S' - D12) B1';
            (S * C2) (S * D21 - D12') R B2';
            F B1 B2 P]
end

function test_REN()

    nx =  5
    # Check equilibrium solve accuracy
    model = implicit_ren{Float64}(5, 5, 10; nl=tanh)
    solver = PeacemanRachford(;tol=1E-8)
    nPoints = 100

    x0 = init_state(model, nPoints)
    utrain = randn(nx, nPoints)
    ytrain = sin.(utrain)

    xn, (xt, wt) = model(x0, utrain, solver)

    exp_cell = explicit(model)
    b = exp_cell.C1 * x0 + exp_cell.D12 * utrain .+ exp_cell.bv
    err = norm(wt - exp_cell.ϕ.(exp_cell.D11 * wt + b))

    if err > 1E-6
        println("Forward solve doesnt seem to be working well")
    end

    # check gradient calculation via FD
    sample_loss(y1, y2) = norm(y1 - y2)^2
    L() = mean(sample_loss.(ytrain, model(x0, utrain, solver)[1]))
    fd_test_grads(L, Flux.params(model))  # param gradients
    fd_test_grads(L, Flux.Params([x0, utrain]))  # input gradients

    # Try simple noiseless regression task
    model = implicit_ren(1, 1, 10; ϵ=1.0, nl=tanh)

    nPoints = 100
    utrain = randn(1, nPoints)
    ytrain = (utrain.^3)
    x0 = init_state(model, nPoints)

    loss() = norm(ytrain - model(x0, utrain, solver)[1])
    ps = Flux.Params(Flux.trainable(model))

    opt = Flux.Optimise.ADAM(1E-3)
    for k in 1:1000
        train_loss, back = Zygote.pullback(loss, ps)

        # calculate gradients and update loss
        ∇J = back(one(train_loss))
        update!(opt, ps, ∇J)

        printfmt("Iteration: {1:2d}\tTraining loss: {2:1.2E}\n", k, train_loss)
    end

    exp_cell = explicit(model)
    ϕ = exp_cell.ϕ
    W = exp_cell.D11
    b = exp_cell.C1 * x0 + exp_cell.D12 * utrain .+ exp_cell.bv

    xn, (xt, wt) = model(x0, utrain, solver)
    err = norm(wt - ϕ.(W * wt + b))
    zeq = wt

    # Run to test if gradient is wrong
    fd_test_grads(loss, Flux.params(model))  # param gradients
    fd_test_grads(loss, Flux.Params([x0, utrain]))  # input gradients

    yest, (xt, wt) = model(x0, utrain)

    plot(utrain', ytrain'; seriestype=:scatter)
    plot!(utrain', yest'; seriestype=:scatter)

end

max_stars_repo_name: imanchester/REN

max_issues_repo_name: imanchester/REN

max_forks_repo_name: imanchester/REN


text:
using HorizonSideRobots
robot = Robot(animate = true)

ext: jl

max_stars_repo_name: TheXlebushek/mirea-progs

max_issues_repo_name: TheXlebushek/mirea-progs

max_forks_repo_name: TheXlebushek/mirea-progs


text:
module DynMultiply

export dynmultiply, dynmultiply_plan, dynmultiplystr

import Base.getindex, Base.setindex!

# The table contains at index N + ... + (N - l + 2) + s the entry
# corresponding to a subsequence with start index s and length l
type DynMultTable
    data::Array{Int, 1}
    N::Int
    DynMultTable(N) = new(zeros(Int, div(N*(N+1), 2)), N)
end

function getdataindex(table::DynMultTable, s::Int, l::Int)
    return (l-1)*table.N - div((l-1)*(l-2), 2) + s
end

# Accessor functions
function getindex(table::DynMultTable, s::Int, l::Int)
    return table.data[getdataindex(table, s, l)]
end

function setindex!(table::DynMultTable, val::Int, s::Int, l::Int)
    table.data[getdataindex(table, s, l)] = val
end

# Multiplication cost for multiplying group s, ..., j-1 and
# j, ..., f-1
function multiplycost(arrays, s::Int, j::Int, f::Int)
    size(arrays[s], 1) * size(arrays[j], 1) * size(arrays[f-1], 2)
end

# Compute multiplication order of matrices
function dynmultiply_plan(arrays...)
    N = length(arrays)
    costs = DynMultTable(N)
    plan = DynMultTable(N)
    for s=1:N
        costs[s, 1] = 0
    end

    # calculate costs and subsequence lengths
    for l= 2:N          # sequence length
        for s=1:N-l+1   # start index
            mincost = typemax(Int)
            l1opt = 0
            for l1=1:l-1    # first sequence length
                multcost = multiplycost(arrays, s, s+l1, s+l)
                cost = costs[s, l1] + costs[s+l1, l-l1] + multcost

                if cost < mincost
                    mincost = cost
                    l1opt = l1
                end
            end

            costs[s, l] = mincost
            plan[s, l] = l1opt
        end
    end

    return plan
end

# Multiply arrays in optimal order
function dynmultiply(arrays...)
    plan = dynmultiply_plan(arrays...)
    N = length(arrays)
    return dynmultiply_part(arrays, 1, N, plan)
end

# Multiply subsequence s, ..., s+l corresponding to plan
function dynmultiply_part(arrays, s::Int, l::Int, plan)
    if l == 1
        return arrays[s]
    else
        l1 = plan[s, l]
        return dynmultiply_part(arrays, s, l1, plan) * dynmultiply_part(arrays, s+l1, l-l1, plan)
    end
end

# Create string representation of optimal multiplication order
function dynmultiplystr(arrays...)
    plan = dynmultiply_plan(arrays...)
    N = length(arrays)
    return dynmultiplystr_part(1, N, plan)
end

# Create string representation of multiplication order of subsequence
# s, ..., s+l corresponding to plan
function dynmultiplystr_part(s::Int, l::Int, plan)
    if l == 1
        return "A$s"
    else
        l1 = plan[s, l]
        return "(" * dynmultiplystr_part(s, l1, plan) * "*" * dynmultiplystr_part(s+l1, l-l1, plan) * ")"
    end
end

end


text:
using Plots

dt = 0.01
g = 9.80665
s = 2
a = 1.2

function lander(x_0, v_0)
    x = x_0
    v = v_0
    t = 0
    xs = []
    ts = []
    while (x > 0)
        f = t >= s ? a : 0
        v += (g-f*g)*dt
        x -= v*dt
        t += dt
        push!(xs, x)
        push!(ts, t)
    end

    (xs, ts)
end
xs, ts = lander(100, 0)

plot(ts, xs, label="lander",xlabel="t",ylabel="x", color=:cornflowerblue)
savefig("lander.png")


text:
# A grid can be 0D (a node), 1D (a line of nodes), 2D (a surface of nodes) or 3D (a surface of nodes with values at each node).

"""
An abstract space type. Your grid type should have the following fields: `dimensions` (Tuple{Integer, Integer, Integer}), agent_positions (Array{Array{Integer}}), and a `grid` named space.

`agent_positions` should always be a list of lists that accept `Integers`, i.e. agent ids.
"""
abstract type AbstractSpace end

function grid0D()
  g = Agents.Graph(1)
end

"""
A path graph. A 1D grid that can optionally be toroidal (a ring).
"""
function grid1D(length::Integer; periodic=false)
  g = LightGraphs.path_graph(length)
  if periodic
    add_edge!(g, 1, length)
  end
  return g
end

"""
A regular 2D grid (SimpleGraph object) where each node is at most connected to four neighbors. It can optionally be toroidal.
"""
function grid2D(x::Integer, y::Integer; periodic::Bool=false)
  g = LightGraphs.grid([x, y], periodic=periodic)
end

"""
A regular 3D grid (`SimpleGraph` object) where each node is at most connected to 6 neighbors. It can optionally be toroidal.
"""
function grid3D(x::Integer, y::Integer, z::Integer; periodic=false, triangular=false)
  if periodic
    if triangular
      g = grid2D_Moore(x, y, periodic=true)
    else
      g = grid2D(x, y, periodic=true)
    end
  else
    if triangular
      g = grid2D_Moore(x, y, periodic=false)
    else
      g = grid2D(x, y)
    end
  end
  gp = deepcopy(g)
  gv = nv(gp)
  for layer in 2:z
    factor = layer-1
    for newnode in 1:gv
      newnodeid = newnode + (factor*gv)
      connect_to = newnodeid - gv
      if newnodeid > nv(g)
        add_vertex!(g)
      end
      add_edge!(g, newnodeid, connect_to)
      for nn in neighbors(gp, newnode)
        newneighbor = nn + (factor*gv)
        if newneighbor > nv(g)
          add_vertex!(g)
        end
        add_edge!(g, newnodeid, newneighbor)
      end
    end
  end
  return g
end

"""
A regular 2D grid (`SimpleGraph` object) where each node connects to its orthogonal and diagonal neighbors. It can optionally be toroidal
"""
function grid2D_Moore(xdim::Integer, ydim::Integer; periodic=false)
  g = LightGraphs.grid([xdim, ydim], periodic=periodic)
  for x in 1:xdim
    for y in 1:ydim
      nodeid = coord_to_vertex((x, y), (xdim, ydim))
      connect_to = []
      if y == ydim
        if x == 1
          if periodic
            tp = (xdim, 1); push!(connect_to, tp)
            tp = (2, ydim-1); push!(connect_to, tp)
          else
            tp = (2, ydim-1); push!(connect_to, tp)
          end
        elseif x == xdim
          if periodic
            tp = (1, 1); push!(connect_to, tp)
            tp = (xdim-1, ydim-1); push!(connect_to, tp)
          else
            tp = (xdim-1, ydim-1); push!(connect_to, tp)
          end
        else
          if periodic
            tp = (x-1, 1); push!(connect_to, tp)
            tp = (x+1, 1); push!(connect_to, tp)
            tp = (x-1, y-1); push!(connect_to, tp)
            tp = (x+1, y-1); push!(connect_to, tp)
          else
            tp = (x-1, y-1); push!(connect_to, tp)
            tp = (x+1, y-1); push!(connect_to, tp)
          end
        end
      elseif y == 1
        if x == 1
          if periodic
            tp = (xdim, ydim); push!(connect_to, tp)
            tp = (2, y+1); push!(connect_to, tp)
          else
            tp = (2, y+1); push!(connect_to, tp)
          end
        elseif x == xdim
          if periodic
            tp = (1, y); push!(connect_to, tp)
            tp = (xdim-1, y+1); push!(connect_to, tp)
          else
            tp = (xdim-1, y+1); push!(connect_to, tp)
          end
        else
          if periodic
            tp = (x-1, y+1); push!(connect_to, tp)
            tp = (x+1, y+1); push!(connect_to, tp)
            tp = (x-1, ydim); push!(connect_to, tp)
            tp = (x+1, ydim); push!(connect_to, tp)
          else
            tp = (x-1, y+1); push!(connect_to, tp)
            tp = (x+1, y+1); push!(connect_to, tp)
          end
        end
      elseif y != 1 && y != ydim && x == 1
        if periodic
          tp = (x+1, y+1); push!(connect_to, tp)
          tp = (x+1, y-1); push!(connect_to, tp)
          tp = (xdim, y+1); push!(connect_to, tp)
          tp = (xdim, y-1); push!(connect_to, tp)
        else
          tp = (x+1, y+1); push!(connect_to, tp)
          tp = (x+1, y-1); push!(connect_to, tp)
        end
      elseif y != 1 && y != ydim && x == xdim
        if periodic
          tp = (x-1, y+1); push!(connect_to, tp)
          tp = (x-1, y-1); push!(connect_to, tp)
          tp = (1, y+1); push!(connect_to, tp)
          tp = (1, y-1); push!(connect_to, tp)
        else
          tp = (x-1, y+1); push!(connect_to, tp)
          tp = (x-1, y-1); push!(connect_to, tp)
        end
      else
          tp = (x+1, y+1); push!(connect_to, tp)
          tp = (x-1, y-1); push!(connect_to, tp)
          tp = (x+1, y-1); push!(connect_to, tp)
          tp = (x-1, y+1); push!(connect_to, tp)
      end

      for pp in connect_to
        add_edge!(g, nodeid, coord_to_vertex((pp[1], pp[2]), (xdim, ydim)))
      end
    end
  end
  return g
end

"""
    grid(x::Integer, y::Integer, z::Integer, periodic=false, Moore=false)

Return a grid (`SimpleGraph` object) based on its dimensions. `x`, `y`, and `z` are the dimensions of the grid. If all dimensions are 1, it will return a 0D space, where all agents are in the same position. If `x` is more than 1, but `y` and `z` are 1, it will return a 1D grid. If `x` and `y` are more than 1, and `z=1`, it will return a 2D regular grid.

* `periodic=true` will create toroidal grids.
* `Moore=true` will return a regular grid in which each node is connected to its diagonal neighbors. If `false`, each node will only connect to its orthogonal neighbors.
"""
function grid(x::Integer, y::Integer, z::Integer, periodic::Bool=false, Moore::Bool=false)
  if x < 1 || y < 1 || z < 1
    throw("x, y, z each can be minimum 1.")
  end
  if x ==1 && y == 1 && z == 1
    g = grid0D()
  elseif x > 1 && y == 1 && z == 1
    g = grid1D(x, periodic=periodic)
  elseif x > 1 && y > 1 && z == 1
    g = grid(x, y, periodic, Moore)
  elseif x > 1 && y > 1 && z > 1
    g = grid3D(x, y, z)
  else
    throw("Invalid grid dimensions! If only one dimension is 1, it should be `z`, if two dimensions are 1, they should be `y` and `z`.")
  end
  return g
end

function grid(x::Integer, y::Integer, periodic::Bool=false, Moore::Bool=false)
  if Moore
    g = grid2D_Moore(x, y, periodic=periodic)
  else
    g = grid2D(x, y, periodic=periodic)
  end
  return g
end

"""
    grid(dims::Tuple{Integer, Integer, Integer}, periodic=false, Moore=false)

Return a grid (`SimpleGraph` object) based on its dimensions. `x`, `y`, and `z` are the dimensions of the grid. If all dimensions are 1, it will return a 0D space, where all agents are in the same position. If `x` is more than 1, but `y` and `z` are 1, it will return a 1D grid. If `x` and `y` are more than 1, and `z=1`, it will return a 2D regular grid.

* `periodic=true` will create toroidal grids.
* `Moore=true` will return a regular grid in which each node is connected to its diagonal neighbors. If `false`, each node will only connect to its orthogonal neighbors.
"""
function grid(dims::Tuple{Integer, Integer, Integer}, periodic::Bool=false, Moore::Bool=false)
  grid(dims[1], dims[2], dims[3], periodic, Moore)
end

"""
    grid(dims::Tuple{Integer, Integer}, periodic=false, Moore=false)

Return a grid (`SimpleGraph` object) based on its dimensions. `x`, `y` are the dimensions of the grid. If all dimensions are 1, it will return a 0D space, where all agents are in the same position. If `x` is more than 1, but `y` is 1, it will return a 1D grid.

* `periodic=true` will create toroidal grids.
* `Moore=true` will return a regular grid in which each node is connected to its diagonal neighbors. If `false`, each node will only connect to its orthogonal neighbors.
"""
function grid(dims::Tuple{Integer,Integer}, periodic::Bool=false, Moore::Bool=false)
  grid(dims[1], dims[2], periodic, Moore)
end

function grid()
  g = grid0D()
  return g
end

"""
    gridsize(dims::Tuple{Integer, Integer, Integer})

Returns the size of a grid with dimenstions `dims`.
"""
function gridsize(dims::Tuple{Integer, Integer, Integer})
  dims[1] * dims[2] * dims[3]
end

"""
    gridsize(dims::Tuple{Integer, Integer})

Returns the size of a grid with dimenstions `dims`.
"""
function gridsize(dims::Tuple{Integer, Integer})
  dims[1] * dims[2]
end

function gridsize(x::Integer, y::Integer, z::Integer)
  gridsize((x,y,z))
end

function gridsize(x::T) where T<:SimpleGraph
  return nv(obj1)
end

"""
    gridsize(model::AbstractModel)

Returns the size of the grid in the model
"""
function gridsize(model::AbstractModel)
  gridsize(model.space.dimensions)
end

"""
    move_agent!(agent::AbstractAgent, pos::Tuple, model::AbstractModel)

Adds `agentID` to a new position in the grid and removes it from the old position. Also updates the agent to represent the new position. `pos` is tuple of x, y, z (only if its a 3D space) coordinates of the grid node. If `pos` is not given, the agent is moved to a random position on the grid.
"""
function move_agent!(agent::AbstractAgent, pos::Tuple, model::AbstractModel)
  # node number from x, y, z coordinates
  nodenumber = coord_to_vertex(pos, model)
  move_agent!(agent, nodenumber, model)
end

"""
    move_agent!(agent::AbstractAgent, pos::Integer, model::AbstractModel)

Adds `agentID` to a new position in the grid and removes it from the old position. Also updates the agent to represent the new position. `pos` is an integer showing the number of the node on the grid node. If `pos` is not given, the agent is moved to a random position on the grid.
"""
function move_agent!(agent::AbstractAgent, pos::Integer, model::AbstractModel)
  push!(model.space.agent_positions[pos], agent.id)
  # remove agent from old position
  if typeof(agent.pos) <: Tuple
    oldnode = coord_to_vertex(agent.pos, model)
    splice!(model.space.agent_positions[oldnode], findfirst(a->a==agent.id, model.space.agent_positions[oldnode]))
    agent.pos = vertex_to_coord(pos, model)  # update agent position
  else
    splice!(model.space.agent_positions[agent.pos], findfirst(a->a==agent.id, model.space.agent_positions[agent.pos]))
    agent.pos = pos
  end
end

function move_agent!(agent::AbstractAgent, model::AbstractModel)
  nodenumber = rand(1:nv(model.space.space))
  move_agent!(agent, nodenumber, model)
  return agent.pos
end

"""
    move_agent_single!(agent::AbstractAgent, model::AbstractModel)

Moves agent to a random nodes on the grid while respecting a maximum of one agent per node. If there are no empty nodes, the agent wont move.

Return the agent's new position.
"""
function move_agent_single!(agent::AbstractAgent, model::AbstractModel)
  empty_cells = [i for i in 1:length(model.space.agent_positions) if length(model.space.agent_positions[i]) == 0]
  if length(empty_cells) > 0
    random_node = rand(empty_cells)
    move_agent!(agent, random_node, model)
  end
  return agent.pos
end

"""
    add_agent!(agent::AbstractAgent, pos::Tuple{Integer, Integer, Integer}, model::AbstractModel)

Adds the agent to the `pos` in the space and to the list of agents. `pos` is tuple of x, y, and z (only if its a 3D space) coordinates of the grid node. If `pos` is not given, the agent is added to a random position.
"""
function add_agent!(agent::AbstractAgent, pos::Tuple, model::AbstractModel)
  # node number from x, y, z coordinates
  nodenumber = coord_to_vertex(pos, model)
  add_agent!(agent, nodenumber, model)
end

"""
    add_agent!(agent::AbstractAgent, pos::Integer, model::AbstractModel)

Adds the agent to the `pos` in the space and to the list of agents. `pos` is the node number of the space. If `pos` is not given, the agent is added to a random position.
"""
function add_agent!(agent::AbstractAgent, pos::Integer, model::AbstractModel)
  push!(model.space.agent_positions[pos], agent.id)
  push!(model.agents, agent)
  if typeof(agent.pos) <: Integer
    agent.pos = pos
  elseif typeof(agent.pos) <: Tuple
    agent.pos = vertex_to_coord(pos, model)  # update agent position
  else
    throw("Unknown type of agent.pos.")
  end
end

"""
    add_agent!(agent::AbstractAgent, model::AbstractModel)
Adds agent to a random node in the space and to the list of agents.

Returns the agent's new position.
"""
function add_agent!(agent::AbstractAgent, model::AbstractModel)
  nodenumber = rand(1:nv(model.space.space))
  add_agent!(agent, nodenumber, model)
  return agent.pos
end

"""
    add_agent_single!(agent::AbstractAgent, model::AbstractModel)

Adds agent to a random node in the space while respecting a maximum one agent per node. It does not do anything if there are no empty nodes.

Returns the agent's new position.
"""
function add_agent_single!(agent::AbstractAgent, model::AbstractModel)
  empty_cells = [i for i in 1:length(model.space.agent_positions) if length(model.space.agent_positions[i]) == 0]
  if length(empty_cells) > 0
    random_node = rand(empty_cells)
    add_agent!(agent, random_node, model)
  end
  return agent.pos
end

"""
    find_empty_nodes_coords(model::AbstractModel)

Returns the coordinates of empty nodes on the model grid.
"""
function find_empty_nodes_coords(model::AbstractModel)
  empty_cells = find_empty_nodes(model::AbstractModel)
  empty_cells_coord = [vertex_to_coord(i, model) for i in empty_cells]
end

"""
    find_empty_nodes(model::AbstractModel)

Returns the IDs of empty nodes on the model space.
"""
function find_empty_nodes(model::AbstractModel)
  empty_cells = [i for i in 1:length(model.space.agent_positions) if length(model.space.agent_positions[i]) == 0]
  return empty_cells
end

"""
    pick_empty(model)

Returns the ID of a random empty cell. Returns 0 if there are no empty cells
"""
function pick_empty(model)
  empty_cells = find_empty_nodes(model)
  if length(empty_cells) == 0
    return 0
  else
    random_node = rand(empty_cells)
    return random_node
  end
end

"""

Returns true if the cell at `cell_id` is empty.
"""
function is_empty(cell_id::Integer, model::AbstractModel)
  if length(model.space.agent_positions[cell_id]) == 0
    return true
  else
    return false
  end
end

"""
    empty_nodes(model::AbstractArray)

Returns true if there are empty nodes, otherwise returns false.
"""
function empty_nodes(model::AbstractArray)
  ee = false
  for el in model.space.agent_positions
    if length(el) == 0
      return true
    end
  end
  return ee
end

"""
    coord_to_vertex(coord::Tuple{Integer, Integer, Integer}, model::AbstractModel)

Returns the node number from x, y, z coordinates.
"""
function coord_to_vertex(coord::Tuple, model::AbstractModel)
  dims = model.space.dimensions
  coord_to_vertex(coord, dims)
end

function coord_to_vertex(agent::AbstractAgent, model::AbstractModel)
  coord_to_vertex(agent.pos, model)
end

function coord_to_vertex(x::Integer, y::Integer, z::Integer, model::AbstractModel)
  coord_to_vertex((x,y,z), model)
end

function coord_to_vertex(x::Integer, y::Integer, model::AbstractModel)
  coord_to_vertex((x,y), model)
end

function coord_to_vertex(x::Integer, y::Integer, z::Integer,dims::Tuple{Integer, Integer, Integer})
  coord_to_vertex((x,y,z), dims)
end

function coord_to_vertex(x::Integer, y::Integer,dims::Tuple{Integer, Integer})
  coord_to_vertex((x,y), dims)
end

"""
    coord_to_vertex(coord::Tuple{T, T, T}, dims::Tuple{Integer, Integer, Integer}) where T<: Integer

Returns node number from x, y, z coordinates.
"""
function coord_to_vertex(coord::Tuple{T, T, T}, dims::Tuple{Integer, Integer, Integer}) where T<: Integer
  if (dims[2] == 1 && dims[3] == 1) || (dims[1] == 1 && dims[3] == 1) || (dims[1] == 1 && dims[2] == 1) # 1D grid
    nodeid = maximum(coord[1])
  elseif dims[1] > 1 && dims[2] > 1 && dims[3] == 1  # 2D grid
    nodeid = coord_to_vertex((coord[1], coord[2]), (dims[1], dims[2]))
  else # 3D grid
    nodeid = (coord[2] * dims[1]) - (dims[1] - coord[1])  # (y * xlength) - (xlength - x)
    nodeid = nodeid + ((dims[1]*dims[2]) * (coord[3]-1))
  end
  return nodeid
end

function coord_to_vertex(coord::Tuple{Integer,Integer}, dims::Tuple{Integer,Integer})
  nodeid = (coord[2] * dims[1]) - (dims[1] - coord[1])  # (y * xlength) - (xlength - x)
  return nodeid
end

"""
    vertex_to_coord(vertex::Integer, model::AbstractModel)

Returns the coordinates of a node given its number on the graph.
"""
function vertex_to_coord(vertex::Integer, model::AbstractModel)
  dims = model.space.dimensions
  vertex_to_coord(vertex, dims)
end

"""
    vertex_to_coord(vertex::Integer, dims::Tuple{Integer, Integer, Integer})

Returns the coordinates of a node given its number on a 3D grid.
"""
function vertex_to_coord(vertex::T, dims::Tuple{Integer, Integer, Integer}) where T<:Integer
  if dims[1] > 1 && dims[2] == 1 && dims[3] == 1  # 1D grid
    coord = (vertex, T(1))
  elseif dims[1] > 1 && dims[2] > 1 && dims[3] == 1  # 2D grid
    coord = vertex_to_coord(vertex, (dims[1], dims[2]))
    coord = (coord[1], coord[2], T(1))
  elseif dims[1] > 1 && dims[2] > 1 && dims[3] > 1  # 3D grid
    gridbasesize = dims[1]*dims[2]
    zcoord = ceil(T, vertex/gridbasesize)
    vertex2d = vertex - ((zcoord-T(1)) * gridbasesize)
    coord2d = vertex_to_coord(vertex2d, (dims[1], dims[2]))
    coord = (T(coord2d[1]), T(coord2d[2]), zcoord)
  else
    error("Wrong coords!")
  end
  return coord
end

"""
    vertex_to_coord(vertex::T, dims::Tuple{Integer,Integer}) where T<: Integer

Returns the coordinates of a node given its number on a 2D grid.
"""
function vertex_to_coord(vertex::T, dims::Tuple{Integer,Integer}) where T<: Integer
  x = T(vertex % dims[1])
  if x == 0
    x = T(dims[1])
  end
  y = ceil(T, vertex/dims[1])
  coord = (x, y)
  return coord
end

"""
    get_node_contents(agent::AbstractAgent, model::AbstractModel)

Returns all agents' ids in the same node as the `agent`.
"""
function get_node_contents(agent::AbstractAgent, model::AbstractModel)
  get_node_contents(agent.pos, model)
end

"""
    get_node_contents(coords::Tuple, model::AbstractModel)

Returns the id of agents in the node at `coords`
"""
function get_node_contents(coords::Tuple, model::AbstractModel)
  node_number = coord_to_vertex(coords, model)
  get_node_contents(node_number, model)
end

"""
    get_node_contents(node_number::Integer, model::AbstractModel)

Returns the id of agents in the node at `node_number`
"""
function get_node_contents(node_number::Integer, model::AbstractModel)
  ns = model.space.agent_positions[node_number]
  return ns
end

"""
    id_to_agent(id::Integer, model::AbstractModel)

Returns an agent given its ID.
"""
function id_to_agent(id::Integer, model::AbstractModel)
  agent_index = findfirst(a-> a.id==id, model.agents)
  agent = model.agents[agent_index]
  return agent
end

"""
    node_neighbors(agent::AbstractAgent, model::AbstractModel)

Returns neighboring node coords/numbers of the node on which the agent resides. If agent `pos` is recorded an integer, the function will return node numbers of the neighbors. If the agent `pos` is a tuple, the function will return the coordinates of neighbors on a grid.
"""
function node_neighbors(agent::AbstractAgent, model::AbstractModel)
  node_neighbors(agent.pos, model)
end

"""
    node_neighbors(node_number::Integer, model::AbstractModel)

Returns neighboring node IDs of the node with `node_number`.
"""
function node_neighbors(node_number::Integer, model::AbstractModel)
  nn = neighbors(model.space.space, node_number)
  return nn
end

"""
    node_neighbors(node_coord::Tuple, model::AbstractModel)

Returns neighboring node coords of the node with `node_coord`.
"""
function node_neighbors(node_coord::Tuple, model::AbstractModel)
  node_number = coord_to_vertex(node_coord, model)
  nn = node_neighbors(node_number, model)
  nc = [vertex_to_coord(i, model) for i in nn]
  return nc
end

"""
    node_neighbors(node_number::Integer, model::AbstractModel, radius::Integer)

Returns a list of neighboring cells to the node `node_number` within the `radius`.
"""
function node_neighbors(node_number::Integer, model::AbstractModel, radius::Integer)
  neighbor_cells = Set(node_neighbors(node_number, model))
  included_cells = Set()
  for rad in 2:radius
    templist = Array{Integer}(undef, 0)
    for nn in neighbor_cells
      if !in(nn, included_cells)
        newns = node_neighbors(nn, model)
        for newn in newns
          push!(templist, newn)
        end
      end
    end
    for tt in templist
      push!(neighbor_cells, tt)
    end
  end
  nlist = collect(neighbor_cells)
  j = findfirst(a-> a==node_number, nlist)
  if j != nothing
    splice!(nlist, j)
  end
  return nlist
end

"""
    Node_iter(model::AbstractModel)

An iterator that returns node coordinates, if the graph is a grid, or otherwise node numbers, and the agents in each node.
"""
struct Node_iter
  model::AbstractModel
  length::Integer
  postype::DataType
end

Node_iter(model::AbstractModel) = Node_iter(model, length(model.space.agent_positions), typeof(model.agents[1].pos))

Base.length(iter::Node_iter) = iter.length

function Base.iterate(iter::Node_iter, state=1)
  if state > iter.length
      return nothing
  end

  cellcontent = iter.model.space.agent_positions[state]
  nagents = length(cellcontent)
  if nagents == 0
    element = (Integer[], Array{AbstractArray}(undef,0))
  else
    if iter.postype <: Tuple
      pp = vertex_to_coord(state, iter.model)
      agentlist = Array{AbstractAgent}(undef, nagents)
      for n in 1:nagents
        agentlist[n] = id_to_agent(cellcontent[n], iter.model)
      end
      element = (pp, agentlist)
    else
      pp = state
      agentlist = Array{AbstractAgent}(undef, nagents)
      for n in 1:nagents
        agentlist[n] = id_to_agent(cellcontent[n], iter.model)
      end
      element = (pp, agentlist)
    end
  end

  return (element, state+1)
end


text:
### routines for approximating functions.

function smoothmin(x::Vector{T})::T where T <: Real
    return smoothextrema(x,convert(T,-999.0))
end

function smoothmax(x::Vector{T})::T where T <: Real
    return smoothextrema(x,convert(T,999.0))
end

function smoothextrema(x::Vector{T}, α::T) where T <: Real
    u = collect( log(x[i]) + x[i]*α for i = 1:length(x) )
    v = collect( x[i]*α for i = 1:length(x) )
    return exp( StatsFuns.logsumexp(u) - StatsFuns.logsumexp(v) )
    #return sum( x[i]*exp(x[i]*α) for i = 1:length(x))/sum( exp(x[i]*α) for i = 1:length(x))
end
# ## test code.
# α = 999.0
# D = 5
# x = randn(D)
# t = rand()
# y = t.*x
# smax_x = smoothextrema(x,α)
# smax_y = smoothextrema(y,α)
# println("smooth max of x is ", smax_x)
# println("max of x is ", maximum(x))
# println("smooth max of y is ", smax_y)
# println("max of y is ", maximum(y))
#
# smin_x = smoothextrema(x,-α)
# smin_y = smoothextrema(y,-α)
# println("smooth min of x is ", smin_x)
# println("min of x is ", minimum(x))
# println("smooth min of y is ", smin_y)
# println("min of y is ", minimum(y))

max_stars_repo_name: RoyCCWang/RWJuliaUtilities

max_issues_repo_name: RoyCCWang/RWJuliaUtilities

max_forks_repo_name: RoyCCWang/RWJuliaUtilities


text:
# Get the callback URL with the provider parameter.
callback_url(p::AbstractString) =
    string(CONFIG["server_url"], ROUTES[:CALLBACK], "?provider=", HTTP.escapeuri(p))

# Get a query string parameter from a request.
getquery(r::HTTP.Request, key::AbstractString, default="") =
    get(HTTP.queryparams(HTTP.URI(r.target)), key, default)

# Get a cookie from a request.
function getcookie(r::HTTP.Request, key::AbstractString, default="")
    cookies = HTTP.cookies(r)
    ind = findfirst(c -> c.name == key, cookies)
    return ind === nothing ? default : cookies[ind].value
end

tplpath(tpl::AbstractString) = joinpath(@__DIR__, "templates", tpl)
const INDEX_TPL = tplpath("index.tpl")
const SELECT_TPL = tplpath("select.tpl")

# Return an HTML response.
html(body::AbstractString) = html(200, body)
function html(status::Int, body::AbstractString)
    registry = REGISTRY[].url
    doc = render_from_file(
              INDEX_TPL,
              route_index=ROUTES[:INDEX],
              registry_url=registry,
              docs_url=DOCS,
              body=body,
          )
    return HTTP.Response(status, ["Content-Type" => "text/html"]; body=doc)
end

json(; kwargs...) = json(200; kwargs...)
json(status::Int; kwargs...) = HTTP.Response(status, ["Content-Type" => "text/json"]; body=JSON.json(kwargs))

# Parse an HTML form.
function parseform(s::AbstractString)
    # In forms, '+' represents a space.
    pairs = split(replace(s, "+" => " "), "&")
    return Dict(map(p -> map(strip ∘ HTTP.unescapeuri, split(p, "=")), pairs))
end


text:
tol=0.00000001
b = [0.5]
c = [1.0, 0.3]

@assert maximum_value(b) - 0.5 < tol
@assert maximum_value(c) - 1.0 < tol

@assert algebraic_sum(b) - 0.5 < tol
@assert algebraic_sum(c) - 1.0 < tol

@assert bounded_sum(b) - 0.5 < tol
@assert bounded_sum(c) - 1.0 < tol

@assert drastic_sum(b) - 0.5 < tol
@assert drastic_sum(c) - 1.0 < tol

@assert einstein_sum(b) - 0.5 < tol
@assert einstein_sum(c) - 1.0 < tol

@assert hamacher_sum(b) - 0.5 < tol
@assert hamacher_sum(c) - 1.0 < tol


text:
abstract type PSRKUNIFACModel <: UNIFACModel end

struct PSRKUNIFAC{c<:EoSModel} <: PSRKUNIFACModel
    components::Array{String,1}
    icomponents::UnitRange{Int}
    groups::GroupParam
    params::UNIFACParam
    puremodel::Vector{c}
    absolutetolerance::Float64
    references::Array{String,1}
end

@registermodel PSRKUNIFAC
export PSRKUNIFAC

function PSRKUNIFAC(components; puremodel=PR,
    userlocations=String[],
     verbose=false)
    groups = GroupParam(components, ["Activity/UNIFAC/PSRK/PSRK_groups.csv"]; verbose=verbose)

    params = getparams(groups, ["Activity/UNIFAC/PSRK/PSRK_like.csv", "Activity/UNIFAC/PSRK/PSRK_unlike.csv"]; userlocations=userlocations,  asymmetricparams=["A","B","C"], ignore_missing_singleparams=["A","B","C"], verbose=verbose)
    A  = params["A"]
    B  = params["B"]
    C  = params["C"]
    R  = params["R"]
    Q  = params["Q"]
    icomponents = 1:length(components)

    init_puremodel = [puremodel([groups.components[i]]) for i in icomponents]
    packagedparams = UNIFACParam(A,B,C,R,Q)
    references = String[]
    model = PSRKUNIFAC(components,icomponents,groups,packagedparams,init_puremodel,1e-12,references)
    return model
end


text:
# Load Julia packages (libraries) needed.

using Pkg, DrWatson

#@quickactivate "StstisticalRethinkingStan"
using StanSample
using StatisticalRethinking

using NamedTupleTools

function quap(sm::SampleModel)
  s = read_samples(sm, :dataframe)
  ntnames = (:coef, :vcov, :converged, :distr, :params)
  n = Symbol.(names(s))
  coefnames = tuple(n...,)
  p = quap(s)
  c = [mean(p[k]) for k in n]
  cvals = reshape(c, 1, length(n))
  coefvalues = tuple(cvals...,)
  v = Statistics.covm(Array(s), cvals)

  ntvalues = tuple(
    namedtuple(coefnames, coefvalues),
    v, true, MvNormal(c, v), n
  )

  namedtuple(ntnames, ntvalues)
end

q = quap(m5_5s)
r = rand(q.distr, 4000)'

mean(r, dims=1) |> display


text:
now = time()
seconds = now % 60 # 60 seconds in a minute
now -= seconds
# to minutes
now = div(now, 60)
minutes = now % 60 # 60 minutes in an hour
now -= minutes
# to hours
now = div(now, 60)
hours = now % 24 # 24 hours in a day
now -= hours
# to days
now = div(now, 24)
days = now % 365 # ± 365 days in a year


text:
using Strategems, Temporal, Indicators, Base.Dates
using Base.Test

# define universe and gather data
assets = ["EOD/AAPL", "EOD/MCD", "EOD/JPM", "EOD/MMM", "EOD/XOM"]
universe = Universe(assets)
gather!(universe)

# define indicator and parameters
function fun(x::TS; args...)::TS
    close_prices = cl(x)
    moving_average = sma(close_prices; args...)
    output = [close_prices moving_average]
    output.fields = [:Close, :MA]
    return output
end

#TODO: define method for when only one parameter is needed where it automatically puts things in vectors
# (so that you could call this by doing `ParameterSet(:n, 50)`)
indicator = Indicator(fun, ParameterSet([:n], [50]))

# define signals
signals = Dict{Symbol,Signal}(:GoLong => Signal(:(Close ↑ MA)),
                              :GoShort => Signal(:(Close ↓ MA)))

# define trading rules
rules = Dict{Symbol,Rule}(:EnterLong => Rule(:GoLong, :(buy,asset,100)),
                          :EnterShort => Rule(:GoShort, :(sell,asset,100)))

# construct and test the strategy
strat = Strategy(universe, indicator, signals, rules)

#TODO: print diagnostic information while running backtest
backtest!(strat, px_trade=:Open, px_close=:Close)


text:
# # Building
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/Building.ipynb)
#
#md # !!! note "Overview"
#md #     System type: Affine system\
#md #     State dimension: 48\
#md #     Application domain: Mechanical Engineering
#
# ## Model description
#
# The model corresponds to a building (the Los Angeles University Hospital) with
# 8 floors each having 3 degrees of freedom, namely displacements in $x$ and $y$
# directions, and rotation [^ASG00]. Such 24 variables evolve according to
# ```math
#   M\ddot{q}(t) + C\dot{q}(t) + Kq(t) = vu(t),
# ```
# where $u(t)$ is the input. This system can be put into a traditional state-space
# form of order 48 by defining $x = (q, \dot{q})^T$. We are interested in the
# motion of the first coordinate $q_1(t)$, hence we choose $v = (1, 0, \ldots, 0)^T$
# and the output $y(t) = \dot{q}_1(t) = x_{25}(t)$. In canonical form,
#
# ```math
#   \begin{array}{lcl}
#   \dot{x}(t) &=& Ax(t) + Bu(t),\qquad u(t) \in \mathcal{U} \\
#   y(t) &=& C x(t)
#   \end{array}
# ```
# where $x(t) \in \mathbb{R}^{48}$ is the state vector, $\mathcal{U} \in \subseteq \mathbb{R}$
# is the input set, and $A \in \mathbb{R}^{48 × 48}$ and $B ∈ \mathbb{R}^{48 × 1}$ are matrices
# given in the file `building.jld2`. Here $y(t)$ is the output with
# $C \in \mathbb{R}^{1 × 48}$ is the projection onto the coordinate 25.
#
# There are two versions of this benchmark:
#
# - **Time-varying inputs:** The inputs can change arbitrarily over time: $\forall t: u(t)\in \mathcal{U}$.
#
# - **Constant inputs:** The inputs are uncertain only in their initial value, and
#    constant over time: ``u(0)\in \mathcal{U}``, ``\dot u (t)= 0``.
#
# In both cases the input set $\mathcal{U}$ is the interval $[0.8, 1.0]$ is,
# and the initial states are taken from Table 2.2 in [^TLT16].

using ReachabilityAnalysis, SparseArrays, JLD2

LazySets.set_ztol(Float64, 1e-14)

const x25 = [zeros(24); 1.0; zeros(23)]
const x25e = vcat(x25, 0.0)

examples_dir = normpath(@__DIR__, "..", "..", "..", "examples")
building_path = joinpath(examples_dir, "Building", "building.jld2")

function building_BLDF01()
    @load building_path A B
    n = size(A, 1)
    U = Interval(0.8, 1.0)
    S = @system(x' = Ax + Bu, u ∈ U, x ∈ Universe(n))

    ## initial states
    center_X0 = [fill(0.000225, 10); fill(0.0, 38)]
    radius_X0 = [fill(0.000025, 10); fill(0.0, 14); 0.0001; fill(0.0, 23)]
    X0 = Hyperrectangle(center_X0, radius_X0)

    prob_BLDF01 = InitialValueProblem(S, X0)
end

using ReachabilityAnalysis: add_dimension

function building_BLDC01()
    @load building_path A B
    n = size(A, 1)
    U = Interval(0.8, 1.0)

    ## initial states
    center_X0 = [fill(0.000225, 10); fill(0.0, 38)]
    radius_X0 = [fill(0.000025, 10); fill(0.0, 14); 0.0001; fill(0.0, 23)]
    X0 = Hyperrectangle(center_X0, radius_X0)

    Ae = add_dimension(A)
    Ae[1:n, end] = B
    prob_BLDC01 = @ivp(x' = Ae * x, x(0) ∈ X0 × U)
end

# ## Reachability settings
#
# The verification goal is to check whether the displacement ``y_1`` of the top
# floor of the building remains below a given bound. In addition to the safety
# specification from the original benchmark, there are two UNSAT instances that
# serve as sanity checks to ensure that the model and the tool work as intended.
# But there is a caveat: In principle, verifying an UNSAT instance only makes
# sense formally if a witness is provided (counter-example, under-approximation,
# etc.). Since most of the participating tools do not have this capability, we
# run the tools with the same accuracy settings on an SAT-UNSAT pair of
# instances. The SAT instance demonstrates that the over-approximation is not
# too coarse, and the UNSAT instance indicates that the over-approximation is
# indeed conservative.
#
# - **BDS01:** Bounded time, safe property: For all ``t \in [0, 20]``,
#    ``y_1(t) \leq 5.1\cdot 10^{-3}``. This property is assumed to be satisfied.
#
# - **BDU01:** Bounded time, unsafe property: For all ``t \in [0, 20]``,
#    ``y_1(t) \leq 4\cdot 10^{-3}``. This property is assumed to be violated.
#    Property BDU01 serves as a sanity check. A tool should be run with the same
#    accuracy settings on BLDF01-BDS01 and BLDF01-BDU01, returning UNSAT on the
#    former and SAT on the latter.
#
# - **BDU02:** Bounded time, unsafe property: The forbidden states are
#    ``\{ y_1(t) \leq -0.78\cdot 10^{-3} \wedge t = 20\}``. This property is
#    assumed to be violated for BLDF01 and satisfied for BLDC01. Property BDU02
#    serves as a sanity check to confirm that time-varying inputs are taken into
#    account. A tool should be run with the same accuracy settings on BLDF01-BDU02
#    and BLDC01-BDU02, returning UNSAT on the former and SAT on the latter.

# ## Results

# For the discrete-time analysis we use a step size of ``0.01``.

using Plots

# ### BLDF01

prob_BLDF01 = building_BLDF01()

# #### Dense time

sol_BLDF01_dense = solve(prob_BLDF01, T=20.0, alg=LGG09(δ=0.004, vars=(25), n=48));

plot(sol_BLDF01_dense, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

# Safety properties

ρ(x25, sol_BLDF01_dense)
#-
ρ(x25, sol_BLDF01_dense) <= 5.1e-3 # BLDF01 - BDS01
#-
ρ(x25, sol_BLDF01_dense) <= 4e-3 # BLDF01 - BDU01
#-
ρ(x25, sol_BLDF01_dense(20.0))
#-
ρ(x25, sol_BLDF01_dense(20.0)) <= -0.78e-3 # BLDF01 - BDU02

# #### Discrete time

sol_BLDF01_discrete = solve(prob_BLDF01, T=20.0, alg=LGG09(δ=0.01, vars=(25), n=48, approx_model=NoBloating()));

plot(sol_BLDF01_discrete, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

# Safety properties

ρ(x25, sol_BLDF01_discrete)
#-
ρ(x25, sol_BLDF01_discrete) <= 5.1e-3 # BLDF01 - BDS01
#-
ρ(x25, sol_BLDF01_discrete)
#-
ρ(x25, sol_BLDF01_discrete) <= 4e-3 # BLDF01 - BDU01
#-
ρ(x25, sol_BLDF01_discrete(20.0))
#-
ρ(x25, sol_BLDF01_discrete(20.0)) <= -0.78e-3 # BLDF01 - BDU02

# ### BLDC01

prob_BLDC01 = building_BLDC01()

# #### Dense time

sol_BLDC01_dense = solve(prob_BLDC01, T=20.0, alg=LGG09(δ=0.005, vars=(25), n=49))

plot(sol_BLDC01_dense, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

# Safety properties

ρ(x25e, sol_BLDC01_dense)
#-
ρ(x25e, sol_BLDC01_dense) <= 5.1e-3 # BLDC01 - BDS01
#-
ρ(x25e, sol_BLDC01_dense) <= 4e-3 # BLDC01 - BDU01
#-
ρ(x25, sol_BLDF01_discrete(20.0))
#-
ρ(x25e, sol_BLDC01_dense(20.0)) <= -0.78e-3 # BLDC01 - BDU02

# #### Discrete time

sol_BLDC01_discrete = solve(prob_BLDC01, T=20.0, alg=LGG09(δ=0.01, vars=(25), n=49, approx_model=NoBloating()))

plot(sol_BLDC01_discrete, vars=(0, 25), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0, xlab="t", ylab="x25")

# Safety properties

ρ(x25e, sol_BLDC01_discrete)
#-
ρ(x25e, sol_BLDC01_discrete) <= 5.1e-3 # BLDC01 - BDS01
#-
ρ(x25e, sol_BLDC01_discrete)
#-
ρ(x25e, sol_BLDC01_discrete) <= 4e-3 # BLDC01 - BDU01
#-
ρ(x25e, sol_BLDC01_discrete(20.0))
#-
ρ(x25e, sol_BLDC01_discrete(20.0)) <= -0.78e-3 # BLDC01 - BDU02

# ## References

# [^ASG00]: Antoulas, Athanasios C., Danny C. Sorensen, and Serkan Gugercin. A survey of model reduction methods for large-scale systems. 2000.

# [^TLT16]: Tran, Hoang-Dung, Luan Viet Nguyen, and Taylor T. Johnson. *Large-scale linear systems from order-reduction (benchmark proposal).* 3rd Applied Verification for Continuous and Hybrid Systems Workshop (ARCH), Vienna, Austria. 2016.


text:
function bisection_indices(right_value, left_value=1, is_first_call=true)
  ( right_value - left_value == 1 ) && return []

  mid_value = (left_value + right_value) ÷ 2

  cur_indices = []
  push!(cur_indices, mid_value)

  this_list = bisection_indices(mid_value, left_value, false)
  that_list = bisection_indices(right_value, mid_value, false)

  append!(cur_indices, collect(Base.Iterators.flatten(zip(this_list, that_list))))

  this_length = length(this_list)
  that_length = length(that_list)

  if this_length != that_length
    @assert abs(this_length - that_length) == 1
    if this_length > that_length
      push!(cur_indices, this_list[end])
    else
      push!(cur_indices, that_list[end])
    end
  end

  unique(cur_indices)
end


text:
abstract type AbstractWikJ end
abstract type AbstractWikJDiagonal <: AbstractWikJ end
abstract type AbstractWikJCommute <: AbstractWikJ end
abstract type AbstractWikJGeneral <: AbstractWikJ end

struct WikJDiagonal_oop <: AbstractWikJDiagonal end

mutable struct WikJDiagonal_iip{WikJType} <: AbstractWikJDiagonal
    WikJ::WikJType
    WikJDiagonal_iip(ΔW) = new{typeof(ΔW)}(false .* ΔW .* ΔW)
end

struct WikJCommute_oop <: AbstractWikJCommute end

mutable struct WikJCommute_iip{WikJType} <: AbstractWikJCommute
    WikJ::WikJType
    function WikJCommute_iip(ΔW)
        WikJ = false .* vec(ΔW) .* vec(ΔW)'
        new{typeof(WikJ)}(WikJ)
    end
end

struct WikJGeneral_oop <: AbstractWikJGeneral
    m_seq::Matrix{Int}
    function WikJGeneral_oop(ΔW)
        m = length(ΔW)
        M = div(m*(m-1),2)
        m_seq = Matrix{Int}(undef, M, 2)
        k = 1
        for i in 1:length(ΔW)
          for j in i+1:length(ΔW)
            m_seq[k,1] = i
            m_seq[k,2] = j
            k += 1
          end
        end
        new(m_seq)
    end
end

mutable struct WikJGeneral_iip{rateNoiseElTypeNoUnits, WikJType} <: AbstractWikJGeneral
    WikJ::WikJType
    WikJ2::WikJType
    WikJ3::WikJType
    m_seq::Matrix{Int}
    vec_ζ::Vector{rateNoiseElTypeNoUnits}
    vec_η::Vector{rateNoiseElTypeNoUnits}
    Gp₁::Vector{rateNoiseElTypeNoUnits}
    Gp₂::Vector{rateNoiseElTypeNoUnits}
    Aᵢ::Vector{rateNoiseElTypeNoUnits}
end

function WikJGeneral_iip(ΔW)
    WikJ = false .* vec(ΔW) .* vec(ΔW)'
    WikJ2 = false .* vec(ΔW) .* vec(ΔW)'
    WikJ3 = false .* vec(ΔW) .* vec(ΔW)'
    m = length(ΔW)
    M = div(m*(m-1),2)
    m_seq = Array{Int}(undef, M, 2)
    k = 1
    for i in 1:length(ΔW)
      for j in i+1:length(ΔW)
        m_seq[k,1] = i
        m_seq[k,2] = j
        k += 1
      end
    end
    vec_ζ = zero(ΔW)
    vec_η = zero(ΔW)
    Gp₁ = false .* Array{eltype(ΔW)}(undef, M)
    Gp₂ = false .* Array{eltype(ΔW)}(undef, M)
    Aᵢ = zero(ΔW)
    WikJGeneral_iip{eltype(ΔW), typeof(WikJ)}(WikJ, WikJ2, WikJ3, m_seq, vec_ζ, vec_η, Gp₁, Gp₂, Aᵢ)
end

function get_WikJ(ΔW,prob,alg)
    if isinplace(prob)
        if typeof(ΔW) <: Number || is_diagonal_noise(prob)
          return WikJDiagonal_iip(ΔW)
      elseif alg.ii_approx isa IICommutative
          return WikJCommute_iip(ΔW)
        else
          return WikJGeneral_iip(ΔW)
        end
    else
        if typeof(ΔW) <: Number || is_diagonal_noise(prob)
          return WikJDiagonal_oop()
        elseif alg.ii_approx isa IICommutative
          return WikJCommute_oop()
        else
          return WikJGeneral_oop(ΔW)
        end
    end
end

function get_iterated_I!(dt, dW, Wik::WikJDiagonal_oop, C=1)
    WikJ = 1//2 .* dW .* dW
    WikJ
end

function get_iterated_I!(dt, dW, Wik::WikJDiagonal_iip, C=1)
    @unpack WikJ = Wik
    if typeof(dW) <: Number
        Wik.WikJ = 1//2 .* dW .^ 2
    else
        @.. WikJ = 1//2*dW^2
    end
    return nothing
end

function get_iterated_I!(dt, dW, Wik::WikJCommute_oop, C=1)
    WikJ = 1//2 .* vec(dW) .* vec(dW)'
    WikJ
end

function get_iterated_I!(dt, dW, Wik::WikJCommute_iip, C=1)
    @unpack WikJ = Wik
    mul!(WikJ,vec(dW),vec(dW)')
    @.. WikJ *= 1//2
    return nothing
end

"""

    get_iterated_I!(dt, dW, Wik::WikJGeneral_oop, C=1)

This function calculates WikJ, a mxm Array for a m dimensional general noise problem, which is a approximation
to the second order iterated integrals.

WikJDiagonal and WikJCommute use the properties of respective noises to simplify the calculations.
While the calculation for General Noise case is taken from section 4 of [SDELab: A Package for solving stochastic differential
equations in MATLAB](https://doi.org/10.1016/j.cam.2006.05.037) and SDELAB2(https://github.com/tonyshardlow/SDELAB2)
which is the Implementation of SDELab in Julia.
```math
    𝒜ᵖ = (Iₘ² - Pₘ)Kₘᵀ Δt/(2π) √(𝑎ₚ) √(Σ∞) Gp₁
```

```math
    √(Σ∞) = (Σ∞ + 2αIₘ)/(√2 * (1 + α))
```

let the combined operators be,
```math
    F = Kₘ(Iₘ² - Pₘ)(Iₘ ⨂ W(Δt)W(Δt)ᵀ)(Iₘ² - Pₘ)Kₘᵀ
```

```math
    Σ∞ = 2Iₘ + (2/Δt)F
```

See the paper for further details of specific operators.
Here we've only shown in which order these are implemented in this code.

From above we can see:

```math
    Δt/(2π) √(𝑎ₚ) √(Σ∞) Gp₁ = Δt/π √(𝑎ₚ) (√(Σ∞)/2 Gp₁)
```

let,
```math
    Oper2(Gp₁) = (√(Σ∞)/2 Gp₁) = (Iₘ/√2 + F/(√2 * (1 + α) * Δt))(Gp₁)
    Oper2(Gp₁) = Iₘ*Gp₁/√2 + F(Gp₁/(√2*(1+α)*Δt))
    Oper2(Gp₁) = Gp₁/√2 + F(Gp₁/(√2*(1+α)*Δt))
```

we have,

    Gp₂ = Gp₁/(sqrt(2)*(1+α)*dt)

```math
    𝒜ᵖ = (Iₘ² - Pₘ)Kₘᵀ Δt/π √(𝑎ₚ) Oper2(Gp₁/√2 + F(Gp₂))
    𝒜ᵖ = √(𝑎ₚ)*Δt/π * (Iₘ² - Pₘ)Kₘᵀ Oper2(Gp₁/√2 + F(Gp₂))
```
In the code we have

```math
    WikJ2 = (Iₘ² - Pₘ)Kₘᵀ(Oper2(Gp₁))
```

"""
function get_iterated_I!(dt, dW, Wik::WikJGeneral_oop, C=1)
    @unpack m_seq = Wik
    m      = length(dW)
    M      = div(m*(m-1),2)
    sum_dW² = dot(dW,dW)

    WikJ = dW*dW'
    eltWikJ = eltype(WikJ)
    WikJ2 =zeros(eltWikJ, size(WikJ))
    Gp₁ = randn(eltWikJ,M)
    α = convert(eltWikJ,sqrt(1 + sum_dW²/dt))
    Gp₂ = Gp₁/convert(eltWikJ, sqrt(2)*(1+α)*dt)
    #operator (Iₘ² - Pₘ)Kₘᵀ
    for i in 1:M
        WikJ2[m_seq[i,1], m_seq[i,2]] = Gp₂[i]
        WikJ2[m_seq[i,2], m_seq[i,1]] = -Gp₂[i]
    end

    #operator (Iₘ X W*Wᵀ)
    WikJ2 = WikJ*WikJ2

    #operator Kₘ(Iₘ² - Pₘ)
    WikJ2 = WikJ2 - WikJ2'
    for i in 1:M
        Gp₂[i] = WikJ2[m_seq[i,1], m_seq[i,2]]
    end
    Gp₁ = Gp₁/convert(eltWikJ,sqrt(2))+Gp₂

    #operator (Iₘ² - Pₘ)Kₘᵀ
    for i in 1:M
        WikJ2[m_seq[i,1], m_seq[i,2]] = Gp₁[i]
        WikJ2[m_seq[i,2], m_seq[i,1]] = -Gp₁[i]
    end

    WikJ *= 1//2
    𝑎ₚ = (π^2)/6
    p = Int(floor((1/(C*π))*sqrt(M/(24*dt))*sqrt(m + 4*sum_dW²/dt) + 1))
    Aᵢ = false .* vec(dW)   # Aᵢ is vector of aᵢ₀
    for r in 1:p
        𝑎ₚ -= (1/r^2)
        var = convert(eltWikJ, sqrt(dt/(2*π*r)))
        vec_ζ = randn(eltWikJ,m)*var
        vec_η = randn(eltWikJ,m)*var
        WikJ += (vec_ζ*vec_η' - vec_η*vec_ζ')
        Aᵢ -=  convert(eltWikJ, 2/(sqrt(π*r)))*vec_ζ
    end

    WikJ -= 1//2*(dW*Aᵢ' - Aᵢ*dW')
    WikJ += convert(eltWikJ, sqrt(𝑎ₚ)*dt/π)*WikJ2
    WikJ
end

"""

    get_iterated_I!(dt, dW, Wik::WikJGeneral_iip, C=1)

This function calculates WikJ, a mxm Array for a m dimensional general noise problem, which is a approximation
to the second order iterated integrals.

WikJDiagonal and WikJCommute use the properties of respective noises to simplify the calculations.
While the calculation for General Noise case is taken from section 4 of [SDELab: A Package for solving stochastic differential
equations in MATLAB](https://doi.org/10.1016/j.cam.2006.05.037) and SDELAB2(https://github.com/tonyshardlow/SDELAB2)
which is the Implementation of SDELab in Julia.
```math
    𝒜ᵖ = (Iₘ² - Pₘ)Kₘᵀ Δt/(2π) √(𝑎ₚ) √(Σ∞) Gp₁
```

```math
    √(Σ∞) = (Σ∞ + 2αIₘ)/(√2 * (1 + α))
```

let the combined operators be,
```math
    F = Kₘ(Iₘ² - Pₘ)(Iₘ ⨂ W(Δt)W(Δt)ᵀ)(Iₘ² - Pₘ)Kₘᵀ
```

```math
    Σ∞ = 2Iₘ + (2/Δt)F
```

See the paper for further details of specific operators.
Here we've only shown in which order these are implemented in this code.

From above we can see:

```math
    Δt/(2π) √(𝑎ₚ) √(Σ∞) Gp₁ = Δt/π √(𝑎ₚ) (√(Σ∞)/2 Gp₁)
```

let,
```math
    Oper2(Gp₁) = (√(Σ∞)/2 Gp₁) = (Iₘ/√2 + F/(√2 * (1 + α) * Δt))(Gp₁)
    Oper2(Gp₁) = Iₘ*Gp₁/√2 + F(Gp₁/(√2*(1+α)*Δt))
    Oper2(Gp₁) = Gp₁/√2 + F(Gp₁/(√2*(1+α)*Δt))
```

we have,

    Gp₂ = Gp₁/(sqrt(2)*(1+α)*dt)

```math
    𝒜ᵖ = (Iₘ² - Pₘ)Kₘᵀ Δt/π √(𝑎ₚ) Oper2(Gp₁/√2 + F(Gp₂))
    𝒜ᵖ = √(𝑎ₚ)*Δt/π * (Iₘ² - Pₘ)Kₘᵀ Oper2(Gp₁/√2 + F(Gp₂))
```
In the code we have

```math
    WikJ2 = (Iₘ² - Pₘ)Kₘᵀ(Oper2(Gp₁))
```

"""
function get_iterated_I!(dt, dW, Wik::WikJGeneral_iip, C=1)
    @unpack WikJ, WikJ2, WikJ3, m_seq, vec_ζ, vec_η, Gp₁, Gp₂, Aᵢ = Wik

    m      = length(dW)
    M      = div(m*(m-1),2)

    sum_dW² = dot(dW,dW) #zero(eltype(dW))
    # mul!(sum_dW²,dW', dW)

    Gp₁ .= randn(M)
    α = sqrt(1 + sum_dW²/dt)
    @.. Gp₂ = Gp₁/(sqrt(2)*(1+α)*dt)

    #operator (Iₘ² - Pₘ)Kₘᵀ
    for i in 1:M
        WikJ2[m_seq[i,1], m_seq[i,2]] = Gp₂[i]
        WikJ2[m_seq[i,2], m_seq[i,1]] = -Gp₂[i]
    end

    #operator (Iₘ X W*Wᵀ)
    mul!(WikJ,dW,dW')
    mul!(WikJ3,WikJ,WikJ2)

    #operator Kₘ(Iₘ² - Pₘ)
    @.. WikJ2 = WikJ3 - WikJ3'
    for i in 1:M
        Gp₂[i] = WikJ2[m_seq[i,1], m_seq[i,2]]
    end
    @.. Gp₁ = Gp₁/sqrt(2) + Gp₂

    #operator (Iₘ² - Pₘ)Kₘᵀ
    for i in 1:M
        WikJ2[m_seq[i,1], m_seq[i,2]] = Gp₁[i]
        WikJ2[m_seq[i,2], m_seq[i,1]] = -Gp₁[i]
    end

    @.. WikJ *= 1//2
    𝑎ₚ = (π^2)/6
    p = Int(floor((1/(C*π))*sqrt(M/(24*dt))*sqrt(m + 4*sum_dW²/dt) + 1))
    Aᵢ .= false .* vec(dW)    # Aᵢ is vector of aᵢ₀
    for r in 1:p
        𝑎ₚ -= (1/r^2)
        var = sqrt(dt/(2*π*r))
        vec_ζ .= randn(m) .* var
        vec_η .= randn(m) .* var
        mul!(WikJ3, vec_ζ, vec_η')
        @.. WikJ += WikJ3 - WikJ3'
        @.. Aᵢ -= (2/sqrt(π*r))*vec_ζ
    end
    mul!(WikJ3, dW, Aᵢ')
    @.. WikJ -= 1//2*(WikJ3 - WikJ3')
    @.. WikJ += (sqrt(𝑎ₚ)*dt/π)*WikJ2
    return nothing
end


text:
# Raw memory management
# Copied from CUDAdrv.jl

# NOTE: Mostly implemented for compatibility and ease-of-use, however we're
# currently not nearly as featureful as the original, which should change
# where possible.

export Mem

module Mem

using ..HSARuntime
import HSARuntime: check, get_region, hsa_memory_allocate, hsa_memory_free

## buffer type

struct Buffer
    ptr::Ptr{Cvoid}
    bytesize::Int

    agent::HSAAgent
end

Base.unsafe_convert(::Type{Ptr{T}}, buf::Buffer) where {T} = convert(Ptr{T}, buf.ptr)

function view(buf::Buffer, bytes::Int)
    bytes > buf.bytesize && throw(BoundsError(buf, bytes))
    return Mem.Buffer(buf.ptr+bytes, buf.bytesize-bytes, buf.agent)
end

## refcounting

const refcounts = Dict{Buffer, Int}()

function refcount(buf::Buffer)
    get(refcounts, Base.unsafe_convert(Ptr{Cvoid}, buf), 0)
end

"""
    retain(buf)

Increase the refcount of a buffer.
"""
function retain(buf::Buffer)
    refcount = get!(refcounts, buf, 0)
    refcounts[buf] = refcount + 1
    return
end

"""
    release(buf)

Decrease the refcount of a buffer. Returns `true` if the refcount has dropped to 0, and
some action needs to be taken.
"""
function release(buf::Buffer)
    haskey(refcounts, buf) || error("Release of unmanaged $buf")
    refcount = refcounts[buf]
    @assert refcount > 0 "Release of dead $buf"
    refcounts[buf] = refcount - 1
    return refcount==1
end

## memory info

"""
    info()

Returns a tuple of two integers, indicating respectively the free and total amount of memory
(in bytes) available for allocation on the agent.
"""
function info()
    free_ref = Ref{Csize_t}()
    total_ref = Ref{Csize_t}()
    # FIXME: I'm not sure HSA has an API for this...
    return convert(Int, free_ref[]), convert(Int, total_ref[])
end

"""
    free()

Returns the free amount of memory (in bytes), available for allocation on the agent.
"""
free() = info()[1]

"""
    total()

Returns the total amount of memory (in bytes), available for allocation on the agent.
"""
total() = info()[2]

"""
    used()

Returns the used amount of memory (in bytes), allocated on the agent.
"""
used() = total()-free()

## generic interface (for documentation purposes)

"""
Allocate linear memory on the device and return a buffer to the allocated memory. The
allocated memory is suitably aligned for any kind of variable. The memory will not be freed
automatically, use [`free(::Buffer)`](@ref) for that.
"""
function alloc end

"""
Free device memory.
"""
function free end

"""
Initialize device memory with a repeating value.
"""
function set! end

"""
Upload memory from host to device.
Executed asynchronously on `queue` if `async` is true.
"""
function upload end
@doc (@doc upload) upload!

"""
Download memory from device to host.
Executed asynchronously on `queue` if `async` is true.
"""
function download end
@doc (@doc download) download!

"""
Transfer memory from device to device.
Executed asynchronously on `queue` if `async` is true.
"""
function transfer end
@doc (@doc transfer) transfer!

## pointer-based

"""
    alloc(bytes::Integer)

Allocate `bytesize` bytes of fine-grained memory.
"""
function alloc(agent::HSAAgent, bytesize::Integer)
    bytesize == 0 && return Buffer(C_NULL, 0, agent)

    ptr_ref = Ref{Ptr{Cvoid}}()
    region = get_region(agent, :finegrained)
    check(hsa_memory_allocate(region[], bytesize, ptr_ref))
    return Buffer(ptr_ref[], bytesize, agent)
end
alloc(bytesize) = alloc(get_default_agent(), bytesize)

function free(buf::Buffer)
    if buf.ptr != C_NULL
        check(hsa_memory_free(buf.ptr))
    end
    return
end

"""
    set!(buf::Buffer, value::T, len::Integer)

Initialize device memory by copying the value `val` for `len` times.
"""
function set!(buf::Buffer, value::T, len::Integer) where T<:Unsigned
    ccall(:memset, Cvoid,
          (Ptr{Cvoid}, T, Csize_t),
          buf.ptr, value, len*sizeof(T))
end

"""
    upload!(dst::Buffer, src, nbytes::Integer)

Upload `nbytes` memory from `src` at the host to `dst` on the device.
"""
function upload!(dst::Buffer, src::Ref, nbytes::Integer)
    ccall(:memcpy, Cvoid,
          (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t),
          dst.ptr, src, nbytes)
end

"""
    download!(dst::Ref, src::Buffer, nbytes::Integer)

Download `nbytes` memory from `src` on the device to `dst` on the host.
"""
function download!(dst::Ref, src::Buffer, nbytes::Integer)
    ccall(:memcpy, Cvoid,
          (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t),
          dst, src.ptr, nbytes)
end

"""
    transfer!(dst::Buffer, src::Buffer, nbytes::Integer)

Transfer `nbytes` of device memory from `src` to `dst`.
"""
function transfer!(dst::Buffer, src::Buffer, nbytes::Integer)
    ccall(:memcpy, Cvoid,
          (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t),
          dst.ptr, src.ptr, nbytes)
end

## array based

"""
    alloc(src::AbstractArray)

Allocate space to store the contents of `src`.
"""
function alloc(src::AbstractArray)
    return alloc(sizeof(src))
end

"""
    upload!(dst::Buffer, src::AbstractArray)

Upload the contents of an array `src` to `dst`.
"""
function upload!(dst::Buffer, src::AbstractArray)
    upload!(dst, Ref(src, 1), sizeof(src))
end

"""
    upload(src::AbstractArray)::Buffer

Allocates space for and uploads the contents of an array `src`, returning a Buffer.
"""
function upload(src::AbstractArray)
    dst = alloc(src)
    upload!(dst, src)
    return dst
end

"""
    download!(dst::AbstractArray, src::Buffer)

Downloads memory from `src` to the array at `dst`. The amount of memory downloaded is
determined by calling `sizeof` on the array, so it needs to be properly preallocated.
"""
function download!(dst::AbstractArray, src::Buffer)
    ref = Ref(dst, 1)
    download!(ref, src, sizeof(dst))
    return
end

## type based

function check_type(::Type{Buffer}, T)
    if isa(T, UnionAll) || T.abstract || !isconcretetype(T)
        throw(ArgumentError("cannot represent abstract or non-leaf object"))
    end
    Base.datatype_pointerfree(T) || throw(ArgumentError("cannot handle non-ptrfree objects"))
    sizeof(T) == 0 && throw(ArgumentError("cannot represent singleton objects"))
end

"""
    alloc(T::Type, [count::Integer=1])

Allocate space for `count` objects of type `T`.
"""
function alloc(::Type{T}, count::Integer=1) where {T}
    check_type(Buffer, T)

    return alloc(sizeof(T)*count)
end

"""
    download(::Type{T}, src::Buffer, [count::Integer=1])::Vector{T}

Download `count` objects of type `T` from the device at `src`, returning a vector.
"""
function download(::Type{T}, src::Buffer, count::Integer=1) where {T}
    dst = Vector{T}(undef, count)
    download!(dst, src)
    return dst
end

end


text:
using Images

render(env::Env) = env |> get_obs |> normedview .|> Gray

function draw_rectangle!(screen_array::Array{UInt8, 2}, x, y, w, h, factor)
    X, Y, W, H = (x, y, w, h) .|> floor .|> Int .|>(num) -> *(num, factor)

    max_h, max_w = size(screen_array)

    for i=max(X, 1):min(X+W, max_w)
        for j=max(Y, 1):min(Y+H, max_h)
            screen_array[j, i] = UInt8(255)
        end
    end
end

get_score(env::Env) = println("AI = $(env.player_a.score) | HUMAN = $(env.player_b.score)")

# Converts the present state information into a 2D UInt8 dot matrix.
function get_obs(env::Env)
    factor = 6

    player_a = env.player_a
    player_b = env.player_b
    ball = env.ball
    arena = env.arena

    screen = arena.dims + 2arena.margin

    screen_array = zeros(UInt8, Int(screen.y * factor), Int(screen.x*factor))

    # the walls
    draw_rectangle!(screen_array, 0, 0, screen.x, arena.margin.y, factor)
    draw_rectangle!(screen_array, 0, (screen.y - arena.margin.y), screen.x, arena.margin.y, factor)

    # the paddles
    draw_rectangle!(screen_array, player_a.position.x,
                                  player_a.position.y + arena.margin.y,
                                  arena.margin.x, player_a.length, factor)

    draw_rectangle!(screen_array, player_b.position.x + arena.margin.x,
                                  player_b.position.y + arena.margin.y,
                                  arena.margin.x, player_b.length, factor)

    # the ball
    draw_rectangle!(screen_array, ball.position.x + arena.margin.x - ball.radius,
                                  ball.position.y + arena.margin.y - ball.radius,
                                  2ball.radius, 2ball.radius, factor)
    return screen_array
end

function update_paddle_position(player::Player, arena::Arena, dir)
    dummy = player.position.y + dir * player.velocity
    dummy = clamp(dummy, 0, arena.dims.y - player.length)
    Vector_(player.position.x, dummy)
end


text:
# Example for basic 2D modeling:
# The receiver positions and the source wavelets are the same for each of the four experiments.
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#

using JUDI.TimeModeling, PyPlot, JLD, SegyIO, LinearAlgebra

# Load Marmousi (4 km x 2 km)
if !isfile("marmousi_small.jld")
    run(`wget ftp://slim.gatech.edu/data/SoftwareRelease/Imaging.jl/CompressiveLSRTM/marmousi_small.jld`)
end
M = load("marmousi_small.jld")

# Setup info and model structure
model = Model(M["n"], M["d"], M["o"], M["m"])
model0 = Model(M["n"], M["d"], M["o"], M["m0"])
dm = vec(M["dm"])

# Setup receivers
nsrc = 10
nxrec = 761 # 10 m receiver spacing
xrec = range(100f0, stop=3900f0, length=nxrec)
yrec = 0f0
zrec = range(210f0, stop=210f0, length=nxrec)

# receiver sampling and recording time
timeR = 2500f0
dtR = 4f0

# Set up receiver structure
rec_geometry = Geometry(xrec, yrec, zrec; dt=dtR, t=timeR, nsrc=nsrc)

## Set up source geometry (cell array with source locations for each shot)
xsrc = convertToCell(range(200f0, stop=3800f0, length=nsrc))
ysrc = convertToCell(range(0f0, stop=0f0, length=nsrc))
zsrc = convertToCell(range(20f0, stop=20f0, length=nsrc))

# source sampling and number of time steps
timeS = 2500f0
dtS = 1f0

# Set up source structure
src_geometry = Geometry(xsrc,ysrc,zsrc;dt=dtS,t=timeS)

# Set up source
wavelet = ricker_wavelet(src_geometry.t[1], src_geometry.dt[1], 0.030)  # 25 Hz peak frequency
q = judiVector(src_geometry, wavelet)

# Set up info structure for linear operators
ntComp = get_computational_nt(q.geometry, rec_geometry, model0)
info = Info(prod(model0.n), nsrc, ntComp)

#################################################################################################

opt = Options(isic=true, optimal_checkpointing=true)

# Setup operators
Pr = judiProjection(info, rec_geometry)
F = judiModeling(info, model; options=opt)
F0 = judiModeling(info, model0; options=opt)
Ps = judiProjection(info, src_geometry)
J = judiJacobian(Pr*F0*Ps', q)

# Time-domain RTM
d_lin = J*dm
rtm_time_5 = adjoint(J[5])*d_lin[5]
rtm_time = adjoint(J)*d_lin

# Frequency-domain RTM
J.options.optimal_checkpointing = false
J.options.dft_subsampling_factor = 8
q_dist = generate_distribution(q)
J.options.frequencies = Array{Any}(undef, nsrc)
for j=1:nsrc
    J.options.frequencies[j] = select_frequencies(q_dist; fmin=0.002, fmax=0.05, nf=20)
end
rtm_freq = adjoint(J)*d_lin
rtm_freq_5 = adjoint(J[5])*d_lin[5]

######################################### Plots ####################################################

n = model0.n
rtm_time_5 = reshape(rtm_time_5, n)[:, 41:end]
rtm_time = reshape(rtm_time, n)[:, 41:end]
rtm_freq_5 = reshape(rtm_freq_5, n)[:, 41:end]
rtm_freq = reshape(rtm_freq, n)[:, 41:end]

rtm_time_5 /= norm(rtm_time_5, Inf)
rtm_time /= norm(rtm_time, Inf)
rtm_freq_5 /= norm(rtm_freq_5, Inf)
rtm_freq /= norm(rtm_freq, Inf)

figure(figsize=(6.66, 6))

s1=subplot(3, 2, 1)
im1 = imshow(adjoint(model0.m[:, 41:end]),cmap="gray", extent=(0.0, 4.0, 2.0, 0.2))
s1[:tick_params]("both", labelsize=8)
xlabel("Lateral position [km]", size=8);
ylabel("Depth [km]", size=8)
t1 = title("a)", fontweight="bold", size=10, loc="left")
t1["set_position"]([-.18, 1])

s2=subplot(3, 2, 2)
im1 = imshow(adjoint(reshape(dm, n)[:, 41:end]), vmin=-1e-1, vmax=1e-1,cmap="gray", extent=(0.0, 4.0, 2.0, 0.2))
s2[:tick_params]("both", labelsize=8)
xlabel("Lateral position [km]", size=8);
ylabel("Depth [km]", size=8)
t2 = title("b)", fontweight="bold", size=10, loc="left")
t2["set_position"]([-.18, 1])

s3=subplot(3, 2, 3)
im1 = imshow(adjoint(rtm_time_5), vmin=-2e-2, vmax=2e-2,cmap="gray", extent=(0.0, 4.0, 2.0, 0.2))
s3[:tick_params]("both", labelsize=8)
xlabel("Lateral position [km]", size=8);
ylabel("Depth [km]", size=8)
t3 = title("c)", fontweight="bold", size=10, loc="left")
t3["set_position"]([-.18, 1])

s4=subplot(3, 2, 4)
im1 = imshow(adjoint(rtm_time), vmin=-5e-2, vmax=5e-2,cmap="gray", extent=(0.0, 4.0, 2.0, 0.2))
s4[:tick_params]("both", labelsize=8)
xlabel("Lateral position [km]", size=8);
ylabel("Depth [km]", size=8)
t4 = title("d)", fontweight="bold", size=10, loc="left")
t4["set_position"]([-.18, 1])

s5=subplot(3, 2, 5)
im1 = imshow(adjoint(rtm_freq_5), vmin=-8e-2, vmax=8e-2,cmap="gray", extent=(0.0, 4.0, 2.0, 0.2))
s5[:tick_params]("both", labelsize=8)
xlabel("Lateral position [km]", size=8);
ylabel("Depth [km]", size=8)
t5 = title("e)", fontweight="bold", size=10, loc="left")
t5["set_position"]([-.18, 1])

s6=subplot(3, 2, 6)
im1 = imshow(adjoint(rtm_freq), vmin=-1.5e-1, vmax=1.5e-1,cmap="gray", extent=(0.0, 4.0, 2.0, 0.2))
s6[:tick_params]("both", labelsize=8)
xlabel("Lateral position [km]", size=8);
ylabel("Depth [km]", size=8)
t6 = title("f)", fontweight="bold", size=10, loc="left")
t6["set_position"]([-.18, 1])
tight_layout()

savefig("figure2", dpi=300, format="png")
run(`mogrify -trim figure2`)


text:
using Random, Plots, Distributions, Statistics, StatsBase

# Example solution for exercise 5.12 using Julia 1.4.1
# I only implemented the 1st race track, but it is trival to modify the
# race track data for the 2nd one.
#
# By @burmecia at GitHub, 29th April 2020
#
# Another useful solution in Python:
# https://towardsdatascience.com/solving-racetrack-in-reinforcement-learning-using-monte-carlo-control-bdee2aa4f04e

# number of race track rows and columns
rows, cols = 32, 17

ε = 0.1
γ = 1

# action can be horizontal and vertical with value 0, 1 or -1
# represented by 1 to 9
actions = [
    (0, 0),
    (0, 1),
    (0, -1),
    (1, 0),
    (1, 1),
    (1, -1),
    (-1, 0),
    (-1, 1),
    (-1, -1),
]
act_len = length(actions)

# 5 velocity values 0 to 4, represented by 1 to 5
vel_len = 5

# initialise Q, C and π
# state is 4-tuple: (row, col, velocity_horizontal, verlocity_vertical)
Q = rand(rows, cols, vel_len, vel_len, act_len) .* 400 .- 500
C = zeros(Float64, rows, cols, vel_len, vel_len, act_len)
π = ones(Int64, rows, cols, vel_len, vel_len)
for r = 1:rows, c = 1:cols, h = 1:vel_len, v = 1:vel_len
    π[r, c, h, v] = argmax(Q[r, c, h, v, :])
end

# set up the 1st race track map, origin (1,1) is at the bottom left,
# boundaries are marked with 1
track = zeros(Int8, rows, cols)
track[32, 1:3] .= 1
track[31, 1:2] .= 1
track[30, 1:2] .= 1
track[29, 1] = 1
track[1:18, 1] .= 1
track[1:10, 2] .= 1
track[1:3, 3] .= 1
track[1:26, 10:end] .= 1
track[26, 10] = 0
start_cols = 4:9  # start line columns
fin_cells = Set([
    (27, cols),
    (28, cols),
    (29, cols),
    (30, cols),
    (31, cols),
    (32, cols),
])  # finish cells

# valid actions for each velocity combination
# Both velocity components are restricted to be nonnegative and less than 5,
# and they cannot both be zero.
valid_acts = [
    map(
        x -> x[1],
        filter(
            a ->
                (h + a[2][1]) in 0:4 &&
                (v + a[2][2]) in 0:4 &&
                !((h + a[2][1]) == 0 && (v + a[2][2]) == 0),
            collect(enumerate(actions)),
        ),
    ) for h = 0:vel_len-1, v = 0:vel_len-1
]

# pre-allocated state, action, probility tracjectory array
S = Array{Tuple{Int64,Int64,Int64,Int64}}(undef, 10^6)
A = Array{Int64}(undef, 10^6)
B = Array{Float64}(undef, 10^6)

function make_trajectory(ε; noise=true)
    t = 1

    # start state
    S[t] = (1, rand(start_cols), 1, 1)

    while true
        s = S[t]

        # get all valid actions for current velocity
        acts = valid_acts[s[3], s[4]]
        num_acts = length(acts)

        # choose next action using ε-greedy policy if πa is valid
        # otherwise choose randomly, and save its probility
        πa = π[s...]
        πa_valid = πa in acts
        if rand() >= ε
            if πa_valid
                a = πa
                b = 1 - ε + ε / num_acts
            else
                a = sample(acts)
                b = 1 / num_acts
            end
        else
            a = sample(acts)
            b = (πa_valid ? ε : 1) / num_acts
        end

        # add some noise
        # with probability 0.1 at each time step the velocity increments are
        # both zero, independently of the intended increments
        if noise && rand() < 0.1
            a = 1
            b = 0.1
        end

        A[t] = a
        B[t] = b
        act = actions[a]

        # next state
        vel = (s[3] + act[1], s[4] + act[2])
        next_s = (s[1] + vel[2] - 1, s[2] + vel[1] - 1, vel[1], vel[2])

        # check if car hits finish line
        path = Set([
            (min(s[1] + i, next_s[1]), min(s[2] + i, next_s[2]))
            for i = 0:maximum(vel)
        ])
        if !isempty(intersect(path, fin_cells))
            return t
        end

        #@show s, a, act, b, next_s, path
        if !checkbounds(Bool, track, next_s[1], next_s[2]) ||
           track[next_s[1], next_s[2]] == 1
            # if car hits boundary, go back to start
            s = (1, rand(start_cols), 1, 1)
        else
            # otherwise go to the next state
            s = next_s
        end

        # append state to tracjectory
        t += 1
        S[t] = s
    end
end

function run_episode(T)
    G, W, R = 0.0, 1.0, -1

    for t = T:-1:1
        s = S[t]
        st = CartesianIndex(Tuple(collect(s)))
        sa = CartesianIndex(Tuple([Tuple(st)..., A[t]]))

        G = γ * G + R
        C[sa] += W
        Q[sa] += W * (G - Q[sa]) / C[sa]

        acts = valid_acts[s[3], s[4]]
        π[st] = acts[argmax(Q[st, :][acts])]
        if A[t] != π[st]
            return t
        end

        W /= B[t]
    end

    0
end

# print several trajectories following the optimal policy
function output_trajectories()
    for i = 1:3
        T = make_trajectory(0.0, noise=false)
        println("\noptimal trajectory #$(i):")
        println("S: ", S[1:T])
        println("A: ", A[1:T])
        println("R: ", -1 * T)
    end
end

function main()
    episode_num = 10^5
    rewards = []

    for i = 1:episode_num
        T = make_trajectory(ε)
        t = run_episode(T)
        println("episode $(i): $(T), $(t), $(T-t)")

        if i % 9 == 0
            T = make_trajectory(0.0)
            push!(rewards, -1 * T)
        end
    end

    plot(rewards)

    output_trajectories()
end

# uncomment the main function below start learning
# take a cup of coffee, it will run for several minutes

#main()


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

using Printf
using Statistics
using CUDA

using Oceananigans
using Oceananigans.Advection
using Oceananigans.AbstractOperations
using Oceananigans.OutputWriters
using Oceananigans.Grids
using Oceananigans.Fields
using Oceananigans.Utils: prettytime

using GeophysicalDissipation.Bickley
using GeophysicalDissipation.DiskTimeSerieses: DiskTimeSeries

oceananigans_prefix(Nh, advection) = "oceananigans_bickley_Nh$(Nh)_$(typeof(advection).name.wrapper)"

function run(; Nh = 128,
               output_time_interval = 2,
               stop_time = 200,
               arch = CPU(),
               ν = 0,
               advection = WENO5())

    experiment_name = oceananigans_prefix(Nh, advection)

    grid = RegularCartesianGrid(size=(Nh, Nh, 1),
                                x = (-2π, 2π), y=(-2π, 2π), z=(0, 1),
                                topology = (Periodic, Periodic, Bounded))

    model = IncompressibleModel(architecture = arch,
                                 timestepper = :RungeKutta3,
                                   advection = advection,
                                        grid = grid,
                                     tracers = :c,
                                     closure = IsotropicDiffusivity(ν=ν, κ=ν),
                                    buoyancy = nothing)

    # ** Initial conditions **
    #
    # u, v: Large-scale jet + vortical perturbations
    #    c: Sinusoid

    # Parameters
    ϵ = 0.1 # perturbation magnitude
    ℓ = 0.5 # Gaussian width
    k = 0.5 # Sinusoidal wavenumber

    # Total initial conditions
    uᵢ(x, y, z) = Bickley.U(y) + ϵ * Bickley.ũ(x, y, ℓ, k)
    vᵢ(x, y, z) = ϵ * Bickley.ṽ(x, y, ℓ, k)
    cᵢ(x, y, z) = Bickley.C(y, grid.Ly)

    set!(model, u=uᵢ, v=vᵢ, c=cᵢ)

    progress(sim) = @info(@sprintf("Iter: %d, time: %.1f, Δt: %.3f, max|u|: %.2f",
                                   sim.model.clock.iteration, sim.model.clock.time,
                                   sim.Δt.Δt, maximum(abs, u.data.parent)))

    wizard = TimeStepWizard(cfl=0.1, Δt=1e-4, max_change=1.1, max_Δt=10.0)

    simulation = Simulation(model, Δt=wizard, stop_time=stop_time,
                            iteration_interval=10, progress=progress)

    # Output: primitive fields + computations
    u, v, w, c = merge(model.velocities, model.tracers)

    ζ = ComputedField(∂x(v) - ∂y(u))

    outputs = merge(model.velocities, model.tracers, (ζ=ζ,))

    function init_grid_and_fields!(file, model)
        file["serialized/grid"] = model.grid

        for (i, field) in enumerate(outputs)
            field_name = keys(outputs)[i]
            file["timeseries/$field_name/meta/location"] = location(field)
        end

        return nothing
    end

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, outputs,
                                schedule = TimeInterval(output_time_interval),
                                init = init_grid_and_fields!,
                                prefix = experiment_name,
                                field_slicer = nothing,
                                force = true)

    @info "Running a simulation of an unstable Bickley jet with $(Nh)² degrees of freedom..."

    start_time = time_ns()

    run!(simulation)

    run_time = (time_ns() - start_time) * 1e-9

    @info "The simulation with Nh = $Nh and advection = $advection ran for $(prettytime(run_time))"
    @show cost = run_time / (model.clock.iteration * Nh^2)

    return experiment_name
end

function visualize(experiment_name)
    @info "Making a fun movie about an unstable Bickley jet..."

    filepath = experiment_name * ".jld2"

    ζ_timeseries = DiskTimeSeries(:ζ, filepath)
    c_timeseries = DiskTimeSeries(:c, filepath)

    grid = c_timeseries.grid

    xζ, yζ, zζ = nodes(ζ_timeseries)
    xc, yc, zc = nodes(c_timeseries)

    anim = @animate for (i, iteration) in enumerate(c_timeseries.iterations)

        @info "    Plotting frame $i from iteration $iteration..."

        ζ = ζ_timeseries[i]
        c = c_timeseries[i]
        t = ζ_timeseries.times[i]

        ζi = interior(ζ)[:, :, 1]
        ci = interior(c)[:, :, 1]

        kwargs = Dict(
                      :aspectratio => 1,
                      :linewidth => 0,
                      :colorbar => :none,
                      :ticks => nothing,
                      :clims => (-1, 1),
                      :xlims => (-grid.Lx/2, grid.Lx/2),
                      :ylims => (-grid.Ly/2, grid.Ly/2)
                     )

        ζ_plot = heatmap(xζ, yζ, clamp.(ζi, -1, 1)'; color = :balance, kwargs...)
        c_plot = heatmap(xc, yc, clamp.(ci, -1, 1)'; color = :thermal, kwargs...)

        ζ_title = @sprintf("ζ at t = %.1f", t)
        c_title = @sprintf("c at t = %.1f", t)

        plot(ζ_plot, c_plot, title = [ζ_title c_title], size = (4000, 2000))
    end

    gif(anim, experiment_name * ".gif", fps = 8)

    return nothing
end

experiment_name = run(Nh=32, arch=CPU(), advection=WENO5(), stop_time=200)
visualize(experiment_name)


text:
"""
    UniversalWrapper(::AbstractMatrix)

UniversalWrapper generalizes the well-known wrappers `*Triangular, Symmetric, Adjoint,
Diagonal, ...` in a more flexible way.
Each wrapper is considered as an operation performed on a subset of the set of square
matrices, applied to such a matrix (called parent).
The set of operators with the standard composition form a semi-group.
It is the aim of this structure to represent all combinations in one object.
The types are distinguished by a big amount of type parameters to enable easy dispatch
on all properties.

Type parameters are a triple of integers {Up,Di,Lo}  (for upper, diagonal, lower)

Up - relies to triangle above and excluding diagonal
Di - relies to diagonal elements of parent matrix
Lo - relies to triangle below and excluding diagonal

Values
Up/Lo:
-  0 constant zero (ignore A[i,j])
-  1 reproduce A[i,j]
- -1 replace A[i,j] by conj.(A[i,j])
-  2 replace A[i,j] by transpose(A[j,i])
- -2 replace A[i,j] by adjoint(A[j,i])
Di:
if A[i,i] isa AbstractMatrix convert to Wrapper with same Up,Di,Lo.
otherwise:
-  0 constant one
-  1 A[i,i]
- -1 conj(A[i,i])
-  2 real(A[i,i])
"""
struct UniversalWrapper{Tv,P<:AbstractMatrix{Tv},Up,Di,Lo} <: AbstractMatrix{Tv}
   parent::P
   function UniversalWrapper{Tv,P,Up,Di,Lo}(p::P) where {Tv,P<:AbstractMatrix{Tv},Up,Di,Lo}
       check_updilo(Up, Di, Lo) || throw(ArgumentError("Up/Lo = $Up/$Lo invalid"))
       new{Tv,P,Up,Di,Lo}(p)
   end
end

function UniversalWrapper{Tv1,P1,Up1,Di1,Lo1}(p::UniversalWrapper{Tv,P,Up2,Di2,Lo2}) where {Tv1,P1,Up1,Di1,Lo1,Tv,P,Up2,Di2,Lo2}
    Up, Di, Lo = compose(Tv, Up1, Di1, Lo1, Up2, Di2, Lo2)
    UniversalWrapper{Tv,P,Up,Di,Lo}(parent(p))
end

@inline function universal(Up1, Di1, Lo1, ::Type{UniversalWrapper{Tv,P,Up2,Di2,Lo2}}) where {Tv,P,Up2,Di2,Lo2}
    Up, Di, Lo = compose(Tv, Up1, Di1, Lo1, Up2, Di2, Lo2)
    UniversalWrapper{Tv,P,Up,Di,Lo}
end
@inline universal(Up1, Di1, Lo1, p::UniversalWrapper{Tv,P,Up2,Di2,Lo2}) where {Tv,P,Up2,Di2,Lo2} = universal(Up1, Di1, Lo1, typeof(p))(parent(p))

@inline universal(Up, Di, Lo, a::AbstractMatrix) = universal(Up, Di, Lo, UniversalWrapper(a))

@inline universal(Up, Di, Lo, T::Type) = universal(Up, Di, Lo, UniversalWrapper{T,Matrix{T},1,1,1})

@inline universal(t::Tuple{Int,Int,Int}, x) = universal(t..., x)

transpose(p::UniversalWrapper) = universal(2, 1, 2, p)
adjoint(p::UniversalWrapper) = universal(-2, -1, -2, p)
unituppertria(p::UniversalWrapper) = universal(1, 0, 0, p)
unitlowertria(p::UniversalWrapper) = universal(0, 0, 1, p)
uppertria(p::UniversalWrapper) = universal(1, 1, 0, p)
lowertria(p::UniversalWrapper) = universal(0, 1, 1, p)
uppersymm(p::UniversalWrapper) = universal(1, 1, 2, p)
lowersymm(p::UniversalWrapper) = universal(2, 1, 1, p)
upperherm(p::UniversalWrapper) = universal(1, 2, -2, p)
lowerherm(p::UniversalWrapper) = universal(-2, 2, 1, p)
conjugate(p::UniversalWrapper) = universal(-1, -1, -1, p)
diagonal(p::UniversalWrapper) = universal(0, 1, 0, p)
Base.one(p::UniversalWrapper) = universal(0, 0, 0, p)

for gen in (:transpose, :adjoint, :unituppertria, :unitlowertria, :uppertria, :lowertria,
            :uppersymm, :lowersymm, :upperherm, :lowerherm, :conjugate)

    @eval ($gen)(p::AbstractMatrix) = ($gen)(UniversalWrapper(p))
end

function_to_wrapper(f::Function, T::Type=Complex) = typeof(f(UniversalWrapper(zeros(T,2,2))))

UniversalWrapper(p::UniversalWrapper) = p
UniversalWrapper(p::P) where {Tv,P<:AbstractMatrix{Tv}} = UniversalWrapper{Tv,P,1,1,1}(p)
UniversalWrapper(p::Transpose) = transpose(_upp(p))
UniversalWrapper(p::Adjoint) = adjoint(_upp(p))
UniversalWrapper(p::UnitUpperTriangular) = unituppertria(_upp(p))
UniversalWrapper(p::UnitLowerTriangular) = unitlowertria(_upp(p))
UniversalWrapper(p::UpperTriangular) = uppertria(_upp(p))
UniversalWrapper(p::LowerTriangular) = lowertria(_upp(p))
UniversalWrapper(p::Symmetric) = p.uplo == 'U' ?  uppersymm(_upp(p)) : lowersymm(_upp(p))
UniversalWrapper(p::Hermitian) = p.uplo == 'U' ?  upperherm(_upp(p)) : lowerherm(_upp(p))

_upp(p) = UniversalWrapper(parent(p))

Base.parent(A::UniversalWrapper) = A.parent
Base.size(A::UniversalWrapper) = size(parent(A))

function Base.getindex(U::UniversalWrapper{Tv,<:Any,Up,Di,Lo}, i::Integer,j::Integer) where {Tv,Up,Di,Lo}
    A = parent(U)
    if i < j
        Up ==  1 ? A[i,j] :
        Up == -1 ? conj.(A[i,j]) :
        Up ==  2 ? transpose(A[j,i]) :
        Up == -2 ? adjoint(A[j,i]) :
        zero(A[i,j])
    elseif i > j
        Lo ==  1 ? A[i,j] :
        Lo == -1 ? conj.(A[i,j]) :
        Lo ==  2 ? transpose(A[j,i]) :
        Lo == -2 ? adjoint(A[j,i]) :
        zero(A[i,j])
    else # i == j
        Tv <: AbstractArray && return convert_element(U, A[i,i])
        Di ==  1 ? A[i,i] :
        Di == -1 ? conj(A[i,i]) :
        Di ==  2 ? real(A[i,i]) :
        one(A[i,i])
    end
end

function check_updilo(Up::Integer, Di::Integer, Lo::Integer)
    true
    # Up == 0 || Lo == 0 || Up + Lo != 0
    # Up == 0 || Lo == 0 || Di != 2 || abs(Up+Lo) == 1 || throw(ArgumentError("Up/Di/Lo = $Up/$Di/$Lo invalid"))
    # Di != 0 || Up == 0 || Lo == 0 || throw(ArgumentError("Up/Di/Lo = $Up/$Di/$Lo invalid"))
end

function convert_element(u::UniversalWrapper{<:Any,<:Any,Up,Di,Lo}, a::X) where {Up,Di,Lo,Tv,X<:UniversalWrapper{Tv}}
    universal(Up, Di, Lo, a)
end
function convert_element(u::UniversalWrapper{<:Any,<:Any,Up,Di,Lo}, a::X) where {Up,Di,Lo,Tv,X<:AbstractMatrix{Tv}}
    unwrap(UniversalWrapper{Tv,X,Up,Di,Lo}(a))
end
iswrsparse(::Type{<:UniversalWrapper{<:Any,P,Up,Di,Lo}}) where {P,Up,Di,Lo} = iswrsparse(P)

@inline function compose_ul(Up1, Up2, Lo2)
    Up1 ==  1 ? Up2 :
    Up1 == -1 ? -Up2 :
    Up1 ==  2 ? cominv(Lo2) :
    Up1 == -2 ? -cominv(Lo2) :
    0
end
@inline function compose_di(Di1, Di2)
    Di1 ==  1 ? Di2 :
    Di1 == -1 ? (Di2 == 2 ? Di2 : -Di2) :
    Di1 ==  2 ? Di2 == 0 ? 0 : Di1 :
    0
end
@inline cominv(Up) = sign(Up) * 3 - Up

function compose(::Type{<:Real}, arg...)
    Up, Di, Lo = compose(arg...)
    abs(Up), min(abs(Di), 1), abs(Lo)
end
compose(::Type, args...) = compose(args...)

@inline function compose(Up1, Di1, Lo1, Up2, Di2, Lo2)
    compose_ul(Up1, Up2, Lo2), compose_di(Di1, Di2), compose_ul(Lo1, Lo2, Up2)
end

import Base.∘
function ∘(::Type{<:UniversalWrapper{T,S,Up1,Di1,Lo1}},::Type{<:UniversalWrapper{T,S,Up2,Di2,Lo2}}) where {T,S,Up1,Di1,Lo1,T2,Up2,Di2,Lo2}

    UniversalWrapper{T,S,compose(T, Up1, Di1, Lo1, Up2, Di2, Lo2)...}
end

parameters(::Type{<:UniversalWrapper{<:Any,<:Any,Up,Di,Lo}}) where {Up,Di,Lo} = (Up, Di, Lo)

all_parameters(::Type) = [ (Up, Di, Lo) for Up = -2:2 for Di = -1:2 for Lo = -2:2 if check_updilo(Up, Di, Lo)]
all_parameters(::Type{<:Real}) = [ (Up, Di, Lo) for Up = 0:2 for Di = 0:1 for Lo = 0:2 if check_updilo(Up, Di, Lo)]

function generate(g::Union{Set,AbstractVector}, T::Type = Complex)
    g = Set([ x isa Function ? function_to_wrapper(x, T) : x for x in g])
    res = copy(g)
    n = 0
    while length(res) != n
        n = length(res)
        res2 = copy(res)
        for v in res2, u in g
            uv =  u ∘ v
            if !(uv in res)
                # println("$u ∘ $v = $uv")
                push!(res, uv)
            end
        end
    end
    collect(res)
end

function norm_conj(U::Type{<:UniversalWrapper{T,S,Up,Di,Lo}}) where {T,S,Up,Di,Lo}
   V = universal(-1,-1,-1, U)
   Up2, Di2, Lo2 = parameters(V)
   cmp(u, d, l) = (count((u,d,l) .< 0), count((u,d,l) .== -1), d, u, l)
   cmp(Up, Di, Lo) < cmp(Up2, Di2, Lo2) ? U : V
end

check_well_defined() = isempty(not_well_defined())
function not_well_defined()
    res = []
    A = UniversalWrapper([1+1im 2+2im; 3+3im 4+4im])
    for updilo2 in all_parameters(Complex)
        U2A = universal(updilo2..., A)
        MU2A = universal(1, 1, 1, Matrix(U2A))
        for updilo1 in all_parameters(Complex)
            if universal(updilo1..., U2A) != universal(updilo1..., MU2A)
                push!(res, (updilo1, updilo2))
            end
        end
    end
    res
end

# check associativity of composition (move to test)
check_assoc() = isempty(non_associatives())
function non_associatives()
    res = []
    for Up1 in -2:2, Lo1 in -2:2
        for Up2 in -2:2, Lo2 in -2:2
            for Up3 in -2:2, Lo3 in -2:2
                if compose(compose(Up1, 0, Lo1, Up2, 0, Lo2)..., Up3, 0, Lo3) !=
                   compose(Up1, 0, Lo1, compose(Up2, 0, Lo2, Up3, 0, Lo3)...)
                    push!(res, ((Up1,0,Lo1), (Up2,0,Lo2), (Up3,0,Lo3)))
                end
            end
        end
    end
    for Di1 in -1:2, Di2 in -1:2, Di3 in -1:2
        if compose(compose(0, Di1, 0, 0, Di2, 0)..., 0, Di3, 0) !=
           compose(0, Di1, 0, compose(0, Di2, 0, 0, Di3, 0)...)
           push!(res, ((0,Di1,0), (0,Di2,0), (0,Di3,0)))
        end
    end
    res
end

check_conjugate_commutative() = relation_empty(!commutes, universal(-1,-1,-1,Complex))
function relation(f::Function, u::Type{<:UniversalWrapper})
    [v for v in universal.(all_parameters(Complex), Complex) if f(u, v) ]
end

relation_empty(f::Function, u::Type{<:UniversalWrapper}) = isempty(relation(f, u))

function commutator()
    [v for v in universal.(all_parameters(Complex), Complex) if relation_empty(!commutes, v)]
end

function inverses(u::Type{<:UniversalWrapper})
    e = universal(1, 1, 1, Complex)
    [v for v in universal.(all_parameters(Complex), Complex) if  v ∘ u == e == u ∘ v]
end

function invertibles()
    [v for v in universal.(all_parameters(Complex), Complex) if !isempty(inverses(v))]
end

function commutes(u::Type{<:UniversalWrapper}, v::Type{<:UniversalWrapper})
    u ∘ v == v ∘ u
end


text:
include("libavcodec_h.jl")

include("avcodec.jl")
#include("avdct.jl")
#include("d3d11va.jl")
#include("dirac.jl")
#include("dv_profile.jl")
#include("qsv.jl")
#include("vaapi.jl")
#include("vdpau.jl")
#include("version.jl")
#include("videotoolbox.jl")
#include("vorbis_parser.jl")
#include("xvmc.jl")

max_stars_repo_name: asafasban/AAVideoIO

max_issues_repo_name: asafasban/AAVideoIO

max_forks_repo_name: asafasban/AAVideoIO