text:
abstract type GridBuild{q <: QuadratureRule} end
abstract type aPrioriBuild{q, p, F} <: GridBuild{q} end
abstract type AdaptiveBuild{q, p, F, A, B} <: GridBuild{q} end
struct AdaptiveRaw{q, p, F, A, B, d} <: AdaptiveBuild{q, p, F, A, B}
  F_cache::Dict{SVector{p,Int}, Float64}
  cache::Dict{SVector{p,Int}, SVector{d,Float64}}
  node_cache::Dict{SVector{p,Int}, SVector{p,Float64}}
  baseline_cache::Dict{SVector{p,Int}, Float64}
  NeighborError::Dict{SVector{p,Int}, Float64}
  Neighbors::Dict{SVector{p, Int}, Δprod{p}}
  Grid::NestedGrid{p, q}
  e::SVector{p, SVector{p,Int}}
  f::F
  l::Int
  μ::A
  U::B
end
struct Adaptive{q, p, F, A, B, T} <: AdaptiveBuild{q, p, F, A, B}
  cache::Dict{SVector{p, Int}, T}
  F_cache::Dict{SVector{p,Int}, Float64}
  node_cache::Dict{SVector{p,Int}, SVector{p,Float64}}
  baseline_cache::Dict{SVector{p,Int}, Float64}
  NeighborError::Dict{SVector{p,Int}, Float64}
  Neighbors::Dict{SVector{p, Int}, Δprod{p}}
  Grid::NestedGrid{p, q}
  e::SVector{p, SVector{p,Int}}
  f::F
  l::Int
  μ::A
  U::B
end
struct SmolyakRaw{q, p, F} <: aPrioriBuild{q, p, F} end
struct Smolyak{q, p, F, T} <: aPrioriBuild{q, p, F} end
RawBuild{q, p, F} = Union{SmolyakRaw{q, p, F}, AdaptiveRaw{q, p, F}}
CacheBuild{q, p, F} = Union{Adaptive{q, p, F}, Smolyak{q, p, F}}
SmolyakBuild{q, p, F} = Union{SmolyakRaw{q, p, F}, Smolyak{q, p, F, T} where T}
default(::Type{AB}) where {AB <: AdaptiveBuild} = 1_000
default(::Type{AP}) where {q, AP <: aPrioriBuild{q}} = default(q)

function Base.sizehint!(ab::AdaptiveBuild, n::Int)
    sizehint!(ab.cache, n)
    sizehint!(ab.F_cache, n)
    sizehint!(ab.node_cache, n)
    sizehint!(ab.baseline_cache, n)
    sizehint!(ab.Grid.grid.w, n)
end

initialize_e(::Type{Val{p}}) where p = SVector{p}(svectors(eye(Int, p), Val{p}))

function add_neighbors!(ab::AdaptiveBuild{q,p,F} where {q,F}, Δ_prod::SVector{p,Int}) where p
  for i ∈ 1:p
    Δ_prod_i = Δ_prod + ab.e[i]
    if haskey(ab.NeighborError, Δ_prod_i) || Δ_prod_i[i] > ab.l
      continue
    end
    neighbor_valid = true
    for j ∈ 1:p
      if ( Δ_prod[j] == 1 ) || ( i == j )
        continue
      elseif !haskey(ab.Grid.Δ_prods, Δ_prod_i - ab.e[j])
        neighbor_valid = false
        break
      end
    end
    if neighbor_valid
      add_neighbor!(ab, Δ_prod_i)
    end
  end
end
function add_neighbor!(ab::AdaptiveBuild{q,p,F} where {q,F}, Δ_prod::SVector{p,Int}) where p
  ab.Neighbors[Δ_prod] = calc_Δ_prod(ab.Grid, Δ_prod)
  ab.NeighborError[Δ_prod] = Δ_prod_error!(ab, Δ_prod)
end

function construct!(ab::AdaptiveBuild{q,p,F} where {q,F}, n::Int) where p
    sizehint!(ab, n)
    Δ_prod = SVector{p}(ones(Int, p))
    ab.Grid.Δ_prods[Δ_prod] = calc_Δ_prod!(ab.Grid, Δ_prod)
    while length(ab.Grid) < n
        add_neighbors!(ab, Δ_prod)
        Δ_prod = reduce((i,v) ->  i[2] > v[2] ? i : v, ab.NeighborError)[1]
        ab.Grid.Δ_prods[Δ_prod] = pop!(ab.Neighbors, Δ_prod)
        expand_grid!(ab.Grid, Δ_prod)
        delete!(ab.NeighborError, Δ_prod)
    end
end

function Δ_prod_error!(ab::AdaptiveBuild{q,p,F} where {q,F}, Δ_prod::SVector{p,Int}) where p
  out = 0.0
  for (i, w) ∈ ab.Neighbors[Δ_prod].w
    out += w * eval_f!(ab, i)
  end
  abs(out)
end
function get_node(ab::AdaptiveBuild{GenzKeister,p,F} where {F}, i::SVector{p,Int}) where p
  get!(() -> get_node.(ab.Grid, i), ab.node_cache, i)
end
function get_node(ab::AdaptiveBuild{KronrodPatterson,p,F} where {F}, i::SVector{p,Int}) where p
    get!(() -> calc_node(ab, i), ab.node_cache, i)
end
function calc_node(ab::AdaptiveBuild{KronrodPatterson}, i::SVector{p,Int}) where p
    snode = get_node.(ab.Grid, i)
    if !haskey(ab.baseline_cache, i)
        ab.baseline_cache[i] = sum(log_sigmoid_jacobian, snode)
    end
    sigmoid.(snode)
end
function transformed_node(ab::AdaptiveRaw, i::SVector{p,Int}, v::SVector{p,Float64}) where p
    get!(() -> ab.μ + ab.U * v, ab.cache, i)
end

function eval_f!(ab::AdaptiveBuild, i::SVector)
  get!(() -> cache_f(ab, i), ab.F_cache, i)
end
function cache_f(ab::AdaptiveRaw, i::SVector)
  node = get_node(ab, i)
  exp( ab.f(transformed_node(ab, i, node)) + get_baseline!(ab, i, node) )
end
function cache_f(ab::Adaptive, i::SVector)
  node = get_node(ab, i)
  res, ab.cache[i] = ab.f(ab.μ + ab.U * node)
  exp( res + get_baseline!(ab, i, node) )
end
@inline function get_baseline!(ab::AdaptiveBuild{KronrodPatterson}, i::SVector, ::AbstractVector)
    get!( () -> sum(log_sigmoid_jacobian, get_node.(ab.Grid, i)), ab.baseline_cache, i )
end
@inline function get_baseline!(ab::AdaptiveBuild{GenzKeister}, i::SVector, x::AbstractVector)
    get!( () -> sum(abs2, x), ab.baseline_cache, i )
end
@inline sigmoid(x::Real) = log( (1 + x)/(1 - x) )
@inline log_sigmoid_jacobian(x::Real) = - log( 1 - x^2 )

@generated function smolyak!(Grid::NestedGrid{p,q}, l::Int) where {p,q<:NestedQuadratureRule}
  quote
    eval(parse("j_"*string($p+1)*" = "*string(l)))
    eval(parse("s_"*string($p+1)*" = 0"))
    @nloops $p i d -> begin
      1:j_{d+1}
    end d -> begin
      s_d = s_{d+1} + i_d - 1
      j_d = l - s_d
    end begin
      Δ = SVector{p}((@ntuple $p i))
      Grid.Δ_prods[Δ] = calc_Δ_prod!(Grid, Δ)
    end
  end
end

@generated function l2!(Grid::NestedGrid{p,q}, l::Int) where {p,q<:NestedQuadratureRule}
  quote
    eval(parse("j_"*string($p+1)*" = "*string(l)))
    eval(parse("s_"*string($p+1)*" = 0"))
    r2 = (l-1)^2
    @nloops $p i d -> begin
      1:j_{d+1}
    end d -> begin
      s_d = s_{d+1} + (i_d - 1)^2
      j_d = floor(Int, √(r2 - s_d) )
    end begin
      Δ = SVector{p}((@ntuple $p i))
      Grid.Δ_prods[Δ] = calc_Δ_prod!(Grid, Δ)
    end
  end
end

@generated function lx!(Grid::NestedGrid{p,q}, l::Int, x::Real) where {p,q<:NestedQuadratureRule}
  quote
    eval(parse("j_"*string($p+1)*" = "*string(l)))
    eval(parse("s_"*string($p+1)*" = 0"))
    rx = (l-1)^x
    inv_x = 1/x
    @nloops $p i d -> begin
      1:j_{d+1}
    end d -> begin
      s_d = s_{d+1} + (i_d - 1)^x
      j_d = floor(Int, (rx - s_d)^inv_x )
    end begin
      Δ = SVector{p}((@ntuple $p i))
      Grid.Δ_prods[Δ] = calc_Δ_prod!(Grid, Δ)
    end
  end
end


text:
# Autogenerated wrapper script for VMEC_jll for x86_64-w64-mingw32-libgfortran4
export libvmec_mkl, libvmec_openblas

using OpenBLAS_jll
using MKL_jll
using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("VMEC")
JLLWrappers.@declare_library_product(libvmec_mkl, "libvmec_mkl-0.dll")
JLLWrappers.@declare_library_product(libvmec_openblas, "libvmec_openblas-0.dll")
function __init__()
    JLLWrappers.@generate_init_header(OpenBLAS_jll, MKL_jll, CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libvmec_mkl,
        "bin\\libvmec_mkl-0.dll",
        nothing,
    )

    JLLWrappers.@init_library_product(
        libvmec_openblas,
        "bin\\libvmec_openblas-0.dll",
        nothing,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()

ext: jl


text:
tt = 0.:dt:T
# Generate Data

Random.seed!(2)
# generate one long path
T_long = 10.0
tt_long = 0.:dt:T_long
W_long = sample(tt_long, Wiener())
X_long = solve(Euler(), x0, W_long, P)
# write long forward to csv file
f = open(outdir*"longforward_nclar.csv","w")
head = "time, component, value \n"
write(f, head)
longpath = [Any[tt_long[j], d, X_long.yy[j][d]] for d in 1:3, j in 1:length(X_long) ][:]
writedlm(f, longpath, ',')
close(f)

# simulate forwards, on the shorter interval
Random.seed!(3)
W = sample(tt, Wiener())
X = solve(Euler(), x0, W, P)
XX = [X]
samples = 100
# draw more paths
for j in 2:samples
    W = sample(tt, Wiener())
    X = solve(Euler(), x0, W, P)
    push!(XX, X)
end

# write forwards to csv file
f = open(outdir*"forwards_nclar.csv","w")
head = "samplenr, time, component, value \n"
write(f, head)
iterates = [Any[i, tt[j], d, XX[i].yy[j][d]] for d in 1:3, j in 1:length(X), i in 1:samples ][:]
writedlm(f, iterates, ',')
close(f)

# simulate forwards, on the shorter interval
# Random.seed!(3)
# W = sample(tt, Wiener())
# X = solve(Euler(), x0, W, P)
# XX = []
# samples = 30
# s = 0
# # draw more paths
# while true
#     sample!(W, Wiener())
#     solve!(Euler(), X, x0, W, P)
#     if norm(L*X.yy[end] .- v) < 0.01
#         push!(XX, copy(X))
#         s += 1
#         s % 10 == 0 && println(".")
#         s >= samples && break
#     end
# end
#
# # write 'forward-bridges' to csv file
# f = open(outdir*"bridges_nclar-"*obs_scheme*".csv","w")
# head = "samplenr, time, component, value \n"
# write(f, head)
# iterates = [Any[i, tt[j], d, XX[i].yy[j][d]] for d in 1:3, j in 1:length(X), i in 1:samples ][:]
# writedlm(f, iterates, ',')
# close(f)
#
# tt = τ(T).(0.:dt:T)


text:
using LinearAlgebra
using FermiCG
using Printf
using Test

atoms = []
push!(atoms,Atom(1,"H",[0,0,0]))
push!(atoms,Atom(2,"H",[0,0,1]))
push!(atoms,Atom(3,"H",[0,0,2]))
push!(atoms,Atom(4,"H",[0,0,3]))
push!(atoms,Atom(5,"H",[0,0,4]))
push!(atoms,Atom(6,"H",[0,0,5]))
#basis = "6-31g"
basis = "sto-3g"

mol     = Molecule(0,1,atoms)
mf = FermiCG.pyscf_do_scf(mol,basis)
ints = FermiCG.pyscf_build_ints(mf.mol,mf.mo_coeff);
e_fci, d1_fci, d2_fci = FermiCG.pyscf_fci(ints,3,3)
# @printf(" FCI Energy: %12.8f\n", e_fci)

function compute_npairs(d2)
    npairs = 0
    for i in 1:size(d2)[1]
        for j in 1:size(d2)[1]
            npairs += d2[i,i,j,j]
        end
    end
    println(" NPairs:", round(npairs,digits=3))
end

compute_npairs(d2_fci)

FermiCG.pyscf_write_molden(mol,basis,mf.mo_coeff,filename="scf.molden")

# Cl = FermiCG.localize(mf.mo_coeff,"boys", mf)
# FermiCG.pyscf_write_molden(mol,"6-31g",Cl,filename="boys.molden")
C = mf.mo_coeff
rdm_mf = C[:,1:2] * C[:,1:2]'
Cl = FermiCG.localize(mf.mo_coeff,"lowdin",mf)
FermiCG.pyscf_write_molden(mol,basis,Cl,filename="lowdin.molden")
S = FermiCG.get_ovlp(mf)
U =  C' * S * Cl
println(" Build Integrals")
flush(stdout)
ints = FermiCG.orbital_rotation(ints,U)
println(" done.")
flush(stdout)
# ints = FermiCG.pyscf_build_ints(mf.mol,Cl);
# println(" done.")
# flush(stdout)

clusters    = [(1:2),(3:4),(5:6)]
#clusters    = [(1:4),(5:8),(9:12)]
init_fspace = [(1,1),(1,1),(1,1)]

clusters = [Cluster(i,collect(clusters[i])) for i = 1:length(clusters)]
display(clusters)

rdm1 = zeros(size(ints.h1))
rdm1a = rdm_mf*.5
rdm1b = rdm_mf*.5

#FermiCG.cmf_ci(ints, clusters, init_fspace, rdm1)
e_cmf, U = FermiCG.cmf_oo(ints, clusters, init_fspace, rdm1, verbose=0, gconv=1e-6)

@test isapprox(e_cmf, -3.205983033016, atol=1e-10)
C_cmf = Cl*U

FermiCG.pyscf_write_molden(mol,basis,C_cmf,filename="cmf.molden")

max_stars_repo_name: nmayhall-vt/FermiCG

max_issues_repo_name: nmayhall-vt/FermiCG

max_forks_repo_name: nmayhall-vt/FermiCG


text:
include("test_common.jl")

@testset "import" begin
    # Not support
    # node"import path from 'path'"
    path = require("path")
    @test path == node"require('path')"
    @test path.resolve(".") == abspath(".")[1:end-1]
    node"const fs = require('fs')"

    cd(@__DIR__)
    isfile("package.json") || NPM.init("-y")
    # Test with `canvas` that is a npm package with native dependencies.
    NPM.is_installed("canvas") || NPM.install("canvas")
    canvas_module = require("canvas")
    canvas = canvas_module.createCanvas(10, 10)
    ctx = canvas.getContext("2d")
    ctx.strokeStyle = "rgba(0,0,0,0.5)"
    ctx.beginPath()
    ctx.lineTo(1, 7)
    ctx.lineTo(10, 3)
    ctx.stroke()
    @test canvas.toDataURL() == "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAYAAACNMs+" *
        "9AAAABmJLR0QA/wD/AP+gvaeTAAAAV0lEQVQYla3PsQ1AAABE0aehYQAliQEkJHaxroROYgUWEJ1Ko1AgCr/+" *
        "l7vjb4KPTvYmhijRYLsTY1SosaDDfBXTM11gQo/12h+hRYIBI/anPblvx37kANwUCq3OvPHeAAAAAElFTkSuQmCC"
end


text:
using StatisticalNonlocality, LinearAlgebra, FFTW, SparseArrays, JLD2
import StatisticalNonlocality: chebyshev, fourier_nodes, fourier_wavenumbers
import StatisticalNonlocality: droprelativezeros!

function construct_four_state_filename(parameters, M, N, dirichlet)
    (; U, γ, ω, κ) = parameters
    base_name = "four_state_"
    descriptor = String[]
    push!(descriptor, "U_" * string(U) * "_")
    push!(descriptor, "γ_" * string(γ) * "_")
    push!(descriptor, "ω_" * string(ω) * "_")
    push!(descriptor, "κ_" * string(κ) * "_")
    push!(descriptor, "M_" * string(M) * "_")
    push!(descriptor, "N_" * string(N) * "_")
    if dirichlet
        push!(descriptor, "dirichlet" * ".jld2")
    else
        push!(descriptor, "neumann" * ".jld2")
    end
    return base_name * prod(descriptor)
end

# Helper functions

makelocal(E) = Array(Diagonal(sum(E, dims = 2)[:]))

function localdiffusivity(E, N, M)
    local localdiff = makelocal(E)
    localdiff = [localdiff[i, i] for i = 1:size(E)[1]]
    local localdiff = reshape(localdiff, (N, M + 1))
    return localdiff
end

function avglocaldiffusivity(E, N, M)
    local localdiff = localdiffusivity(E, N, M)
    local localdiff = sum(localdiff, dims = 1)[:] ./ N # zero'th mode
    return localdiff
end

function grabdiagonal(A)
    MM = minimum(size(A))
    diagA = zeros(MM)
    for i in 1:MM
        diagA[i] = A[i, i]
    end
    return diagA
end

function advection_operator(ψ, ∂x, ∂z)
    return Diagonal(∂z * ψ[:]) * ∂x + Diagonal(-∂x * ψ[:]) * ∂z
end

# Main Function
function four_state(parameters; M = 8 * 8, N = 8, filename = nothing, dirichlet = false)
    (; U, γ, ω, κ) = parameters
    U₀ = U
    if isnothing(filename)
        filename = construct_four_state_filename(parameters, M, N, dirichlet)
    end

    Dz, z = chebyshev(M)
    a, b = 0, 2π
    k = fourier_wavenumbers(N, L = b - a)
    x = fourier_nodes(N, a = a, b = b)

    ℱ = fft(I + zeros(N, N), 1)
    ℱ⁻¹ = ifft(I + zeros(N, N), 1)
    Dx = real.(ℱ⁻¹ * Diagonal(im .* k) * ℱ)
    x = reshape(x, (N, 1))
    z = reshape(z, (1, M + 1))

    ψ¹ = U₀ * sin.(x) .* cos.(π / 2 * z)
    ψ² = U₀ * cos.(x) .* cos.(π / 2 * z)
    ψ³ = -U₀ * sin.(x) .* cos.(π / 2 * z)
    ψ⁴ = -U₀ * cos.(x) .* cos.(π / 2 * z)

    ∂z = kron(Dz, I + zeros(N, N))
    ∂x = kron(I + zeros(M + 1, M + 1), Dx)

    zlifted = sparse(kron(Diagonal(z[:]), I + zeros(N, N)))

    A¹ = advection_operator(ψ¹, ∂x, ∂z)
    A² = advection_operator(ψ², ∂x, ∂z)

    Δ = ∂x^2 + ∂z^2

    # Create the block linear operator
    L = [(γ*I-κ*Δ) -ω*I 0.5*A¹; ω*I (γ*I-κ*Δ) -0.5*A²; A¹ -A² (2*γ*I-κ*Δ)]

    # Grab Boundary Indices
    zlifted = sparse(kron(Diagonal(z[:]), I + zeros(N, N)))
    zlifted = [zlifted 0*I 0*I; 0*I zlifted 0*I; 0*I 0*I zlifted]
    ∂Ω¹ = zlifted.rowval[zlifted.nzval.==z[1]]

    # Impose Boundary conditions at z[1]
    for (i, ∂i) in enumerate(∂Ω¹)
        if dirichlet
            L[∂i, :] .= 0.0
            L[∂i, ∂i] = 1.0
        else
            L[∂i, :] .= 0.0
            blockindex = div((i - 1), N)
            inds = (blockindex*N*(M+1)+1):((blockindex+1)*N*(M+1))
            L[∂i, inds] .= ∂z[∂Ω¹[(i-1)%N+1], :]
        end
    end

    # Impost boundary conditions at z[end]
    ∂Ω² = zlifted.rowval[zlifted.nzval.==z[end]]
    for (i, ∂i) in enumerate(∂Ω²)
        if dirichlet
            L[∂i, :] .= 0.0
            L[∂i, ∂i] = 1.0
        else
            L[∂i, :] .= 0.0
            blockindex = div((i - 1), N)
            inds = (blockindex*N*(M+1)+1):((blockindex+1)*N*(M+1))
            L[∂i, inds] .= ∂z[∂Ω²[(i-1)%N+1], :]
        end
    end

    # Find the Greens Function
    G = qr(L) \ I

    # Grab Components of the Velocity field
    u¹ = Diagonal(∂z * ψ¹[:])
    u² = Diagonal(∂z * ψ²[:])
    v¹ = Diagonal(-∂x * ψ¹[:])
    v² = Diagonal(-∂x * ψ²[:])
    # Account for Boundary conditions in U and V
    # Can be thought of as an operator that masks the values on boundaries
    # Technically the mask is applied to the entire rhs, but since we want the operator
    # we just applied it here.
    # This is equivalent since diagonal matrices commute and masking is idempotent
    U = -0.5 * [u¹; u²; 0 * I]
    Uᵀ = [u¹ u² 0 * I]
    V = -0.5 * [v¹; v²; 0 * I]
    Vᵀ = [v¹ v² 0 * I]

    ∂Ω = vcat(∂Ω¹, ∂Ω²)
    for ∂i in ∂Ω
        U[∂i, :] .= 0.0
        V[∂i, :] .= 0.0
    end

    EF¹¹ = Uᵀ * G * U # flux in x due to gradients in x
    EF¹² = Uᵀ * G * V # flux in x due to gradients in z
    EF²¹ = Vᵀ * G * U # flux in z due to gradients in x
    EF²² = Vᵀ * G * V # flux in z due to gradients in z

    E = [EF¹¹ EF¹²; EF²¹ EF²²]

    EF¹¹ *= -1.0
    EF¹² *= -1.0
    EF²¹ *= -1.0
    EF²² *= -1.0

    # always overwrite
    data_directory = "data"
    mkpath(data_directory)
    filepath = data_directory * "/" * filename
    if isfile(filepath)
        rm(filepath)
    end

    file = jldopen(filepath, "a+")

    # diffusivity
    groupname = "diffusivity"
    JLD2.Group(file, groupname)
    file[groupname]["K11"] = EF¹¹
    file[groupname]["K12"] = EF¹²
    file[groupname]["K21"] = EF²¹
    file[groupname]["K22"] = EF²²

    # parameters
    groupname = "parameters"
    JLD2.Group(file, groupname)
    file[groupname]["γ"] = γ
    file[groupname]["ω"] = ω
    file[groupname]["κ"] = κ
    file[groupname]["U"] = U₀

    # grid
    groupname = "grid"
    JLD2.Group(file, groupname)
    file[groupname]["x"] = x
    file[groupname]["z"] = z

    # transition matrix
    groupname = "transition"
    JLD2.Group(file, groupname)

    T = [-1 1 0 0; 0 -1 1 0; 0 0 -1 1; 1 0 0 -1]
    sT = (T + T') / 2
    aT = (T - T') / 2

    file[groupname]["T"] = γ * sT + ω * aT

    # streamfunction
    groupname = "streamfunction"
    JLD2.Group(file, groupname)
    file[groupname]["ψ¹"] = ψ¹
    file[groupname]["ψ²"] = ψ²
    file[groupname]["ψ³"] = ψ³
    file[groupname]["ψ⁴"] = ψ⁴

    # velocities
    groupname = "velocities"
    JLD2.Group(file, groupname)
    file[groupname]["u¹"] = reshape(grabdiagonal(u¹), size(ψ¹))
    file[groupname]["u²"] = reshape(grabdiagonal(u²), size(ψ¹))
    file[groupname]["v¹"] = reshape(grabdiagonal(v¹), size(ψ¹))
    file[groupname]["v²"] = reshape(grabdiagonal(v²), size(ψ¹))

    # local diffusivities
    groupname = "localdiffusivity"
    JLD2.Group(file, groupname)
    file[groupname]["κ11"] = localdiffusivity(EF¹¹, N, M)
    file[groupname]["κ12"] = localdiffusivity(EF¹², N, M)
    file[groupname]["κ21"] = localdiffusivity(EF²¹, N, M)
    file[groupname]["κ22"] = localdiffusivity(EF²², N, M)

    # zonally averaged local diffusivities
    groupname = "averagelocaldiffusivity"
    JLD2.Group(file, groupname)
    file[groupname]["κ11"] = avglocaldiffusivity(EF¹¹, N, M)
    file[groupname]["κ12"] = avglocaldiffusivity(EF¹², N, M)
    file[groupname]["κ21"] = avglocaldiffusivity(EF²¹, N, M)
    file[groupname]["κ22"] = avglocaldiffusivity(EF²², N, M)

    close(file)
    return filename
end

max_stars_repo_name: sandreza/StatisticalNonlocality

max_issues_repo_name: sandreza/StatisticalNonlocality

max_forks_repo_name: sandreza/StatisticalNonlocality


text:
using Test
import IntervalLapper
const IL = IntervalLapper
const Iv = IL.Interval{Int}

function setup_nonoverlapping()
	data = map(x -> Iv(x, x + 10, 0), 0:20:100)
	IL.Lapper(data)
end

function setup_overlapping()
	data = map(x -> Iv(x, x + 15, 0), 0:10:100)
	IL.Lapper(data)
end

function setup_badlapper()
	data = [
	    Iv(70, 120, 0), # max_len = 50
            Iv(10, 15, 0),
            Iv(10, 15, 0), # exact overlap
            Iv(12, 15, 0), # inner overlap
            Iv(14, 16, 0), # overlap end
            Iv(40, 45, 0),
            Iv(50, 55, 0),
            Iv(60, 65, 0),
            Iv(68, 71, 0), # overlap start
            Iv(70, 75, 0),
	]
	IL.Lapper(data)
end

function setup_single()
	data = [Iv(10, 35, 0)]
	IL.Lapper(data)
end

@testset "Query Stop Interval Start" begin
	lapper = setup_nonoverlapping()
	bits = IL.Bits(lapper.intervals)
	cursor = Ref(1)
	@test nothing == Base.iterate(IL.find(lapper, 30, 35))
	@test nothing == Base.iterate(IL.seek(lapper, 30, 35, cursor))
	@test IL.count(bits, 30, 35) == length(collect(IL.find(lapper, 30, 35)))
end

# Test that a query that overlaps the start of an interval returns that interval
@testset "Query Overlaps Interval Start" begin
	lapper = setup_nonoverlapping()
	cursor = Ref(1)
	expected = Iv(20, 30, 0)
	@test expected == Base.iterate(IL.find(lapper, 15, 25))[1]
	@test expected == Base.iterate(IL.seek(lapper, 15, 25, cursor))[1]
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 15, 25) == length(collect(IL.find(lapper, 15, 25)))
end

# Test that a query that overlaps the stop of an interval returns that interval
@testset "Query Overlaps Interval Stop" begin
	lapper = setup_nonoverlapping()
	cursor = Ref(1)
	expected = Iv(20, 30, 0)
	@test expected == Base.iterate(IL.find(lapper, 25, 35))[1]
	@test expected == Base.iterate(IL.seek(lapper, 25, 35, cursor))[1]
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 25, 35) == length(collect(IL.find(lapper, 25, 35)))
end

# Test that a query that is enveloped by interval returns interval<Paste>
@testset "Interval Envelops Query" begin
	lapper = setup_nonoverlapping()
	cursor = Ref(1)
	expected = Iv(20, 30, 0)
	@test expected == Base.iterate(IL.find(lapper, 22, 27))[1]
	@test expected == Base.iterate(IL.seek(lapper, 22, 27, cursor))[1]
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 22, 27) == length(collect(IL.find(lapper, 22, 27)))
end

# Test that a query that envolops an interval returns that interval
@testset "Query Envelops Interval" begin
	lapper = setup_nonoverlapping()
	cursor = Ref(1)
	expected = Iv(20, 30, 0)
	@test expected == Base.iterate(IL.find(lapper, 15, 35))[1]
	@test expected == Base.iterate(IL.seek(lapper, 15, 35, cursor))[1]
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 15, 35) == length(collect(IL.find(lapper, 15, 35)))
end

@testset "Overlapping Intervals" begin
	lapper = setup_overlapping()
	cursor = Ref(1)
	e1 = Iv(0, 15, 0)
	e2 = Iv(10, 25, 0)

	@test [e1, e2] == collect(IL.find(lapper, 8, 20))
	@test [e1, e2] == collect(IL.seek(lapper, 8, 20, cursor))
	@test 2 == length(collect(IL.find(lapper, 8, 20)))
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 8, 20) == length(collect(IL.find(lapper, 8, 20)))
end

@testset "Merge Overlaps" begin
	lapper = setup_badlapper()
	expected = [
	    Iv( 10,  16,  0),
            Iv( 40,  45,  0),
            Iv( 50,  55,  0),
            Iv( 60,  65,  0),
            Iv( 68,  120,  0), # max_len = 50
	]
	IL.merge_overlaps!(lapper)
	@test expected == lapper.intervals
end

@testset "Lapper Coverage" begin
	lapper = setup_badlapper()
	before = IL.coverage(lapper)
	IL.merge_overlaps!(lapper)
	after = IL.coverage(lapper)
	@test before == after

	lapper = setup_nonoverlapping()
	coverage = IL.coverage(lapper)
	@test coverage == 60
end

@testset "Interval Intersections" begin
	i1 = Iv(70, 120, 0)
	i2 = Iv(10, 15, 0)
	i3 = Iv( 10,  15,  0) # exact overlap
        i4 = Iv( 12,  15,  0) # inner overlap
        i5 = Iv( 14,  16,  0) # overlap end
        i6 = Iv( 40,  50,  0)
        i7 = Iv( 50,  55,  0)
        i8 = Iv( 60,  65,  0)
        i9 = Iv( 68,  71,  0) # overlap start
        i10 = Iv( 70,  75,  0)

	@test IL.intersectlen(i2, i3) == 5 # exact match
	@test IL.intersectlen(i2, i4) == 3 # inner intersect
	@test IL.intersectlen(i2, i5) == 1 #  end intersect
	@test IL.intersectlen(i9, i10) == 1 # start intersect
	@test IL.intersectlen(i7, i8) == 0 # no intersect
	@test IL.intersectlen(i6, i7) == 0 # no intersect stop = start
	@test IL.intersectlen(i1, i10) == 5 # inner intersect at start
end

@testset "Union and Itersect" begin
	data1 = [
	     Iv( 70,  120,  0), # max_len = 50
            Iv( 10,  15,  0), # exact overlap
            Iv( 12,  15,  0), # inner overlap
            Iv( 14,  16,  0), # overlap end
            Iv( 68,  71,  0), # overlap start
	]
	data2 = [
		Iv( 10,  15,  0),
		Iv( 40,  45,  0),
		Iv( 50,  55,  0),
		Iv( 60,  65,  0),
		Iv( 70,  75,  0),
	]

	lapper1 = IL.Lapper(data1)
	lapper2 = IL.Lapper(data2)

	# Should be the same either way it's calculated
	@testset "Non-merged-lappers" begin
		@testset "Lapper1 vs Lapper2" begin
			union, intersect = IL.union_and_intersect(lapper1, lapper2)
			@test intersect == 10
			@test union == 73
		end
		@testset "Lapper2 vs Lapper1" begin
			union, intersect = IL.union_and_intersect(lapper2, lapper1)
			@test intersect == 10
			@test union == 73
		end
	end

	# Should still be the same
	@testset "Merged-Lappers" begin
		IL.merge_overlaps!(lapper1)
		IL.merge_overlaps!(lapper2)
		cov1 = IL.coverage(lapper1)
		cov2 = IL.coverage(lapper2)
		@testset "Lapper1 vs Lapper2" begin
			union, intersect = IL.union_and_intersect(lapper1, lapper2, cov1, cov2)
			@test intersect == 10
			@test union == 73
		end
		@testset "Lapper2 vs Lapper1" begin
			union, intersect = IL.union_and_intersect(lapper2, lapper1, cov2, cov1)
			@test intersect == 10
			@test union == 73
		end
	end
end

@testset "Find Overlaps In Large Intervals" begin
        data1 = [
            Iv( 0,  8,  0),
            Iv( 1,  10,  0),
            Iv( 2,  5,  0),
            Iv( 3,  8,  0),
            Iv( 4,  7,  0),
            Iv( 5,  8,  0),
            Iv( 8,  8,  0),
            Iv( 9,  11,  0),
            Iv( 10,  13,  0),
            Iv( 100,  200,  0),
            Iv( 110,  120,  0),
            Iv( 110,  124,  0),
            Iv( 111,  160,  0),
            Iv( 150,  200,  0),
        ]
        lapper = IL.Lapper(data1);
	found = collect(IL.find(lapper, 8, 11))
        @test found == [
            Iv( 1,  10,  0),
            Iv( 9,  11,  0),
            Iv( 10,  13,  0),
        ]
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 8, 11) == length(collect(IL.find(lapper, 8, 11)))

	cursor = Ref(1)
	found = collect(IL.seek(lapper, 8, 11, cursor))
        @test found == [
            Iv( 1,  10,  0),
            Iv( 9,  11,  0),
            Iv( 10,  13,  0),
        ]

	found = collect(IL.find(lapper, 145, 151))
	@test found == [
            Iv( 100,  200,  0),
            Iv( 111,  160,  0),
            Iv( 150,  200,  0),
        ]
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 145, 151) == length(collect(IL.find(lapper, 145, 151)))

	cursor = Ref(1)
	found = collect(IL.seek(lapper, 145, 151, cursor))
	@test found == [
            Iv( 100,  200,  0),
            Iv( 111,  160,  0),
            Iv( 150,  200,  0),
        ]
end

@testset "Depth Sanity" begin
        data1 = [
            Iv( 0,  10,  0),
            Iv( 5,  10,  0)
        ]
        lapper = IL.Lapper(data1);
	found = collect(IL.depth(lapper))
        @test found == [
                   IL.Interval( 0,  5,  1),
                   IL.Interval( 5,  10,  2)
        ]
end

@testset "Depth Harder" begin
	data1 = [
            Iv( 1,  10,  0),
            Iv( 2,  5,  0),
            Iv( 3,  8,  0),
            Iv( 3,  8,  0),
            Iv( 3,  8,  0),
            Iv( 5,  8,  0),
            Iv( 9,  11,  0),
            Iv( 15,  20,  0),
        ]
        lapper = IL.Lapper(data1);
	found = collect(IL.depth(lapper))
        @test found == [
                   IL.Interval( 1,  2,  1),
                   IL.Interval( 2,  3,  2),
                   IL.Interval( 3,  8,  5),
                   IL.Interval( 8,  9,  1),
                   IL.Interval( 9,  10,  2),
                   IL.Interval( 10,  11,  1),
                   IL.Interval( 15,  20,  1),
        ]
end

@testset "Depth Hard" begin
        data1 = [
            Iv( 1,  10,  0),
            Iv( 2,  5,  0),
            Iv( 3,  8,  0),
            Iv( 3,  8,  0),
            Iv( 3,  8,  0),
            Iv( 5,  8,  0),
            Iv( 9,  11,  0),
        ]
        lapper = IL.Lapper(data1);
	found = collect(IL.depth(lapper))
        @test found == [
                   IL.Interval( 1,  2,  1),
                   IL.Interval( 2,  3,  2),
                   IL.Interval( 3,  8,  5),
                   IL.Interval( 8,  9,  1),
                   IL.Interval( 9,  10,  2),
                   IL.Interval( 10,  11,  1),
        ]
end

#=
# Bug tests - these are tests that came from real life
=#

# Test that it's not possible to induce index out of bounds by pushing the
# cursor past the end of the lapper
@testset "Seek Over Len" begin
	lapper = setup_nonoverlapping();
	single = setup_single();
	cursor = Ref(1)
	count = 0
	for interval in lapper.intervals
		for o_interval in IL.seek(single, interval.start, interval.stop, cursor)
			count += 1
		end
	end
end

# Test that if lower_bound puts us before the first match, we still return match
@testset "Find Over Behind First Match" begin
        lapper = setup_badlapper();
	e1 = Iv( 50,  55,  0)
	found = Base.iterate(IL.find(lapper, 50, 55))[1];
        @test found == e1
end

# When there is a very long interval that spans many little intervals, test that the
# little intervals still get returned properly
@testset "Bad Skips" begin
        data = [
            Iv(25264912,  25264986,  0),
            Iv(27273024,  27273065	,  0),
            Iv(27440273,  27440318	,  0),
            Iv(27488033,  27488125	,  0),
            Iv(27938410,  27938470	,  0),
            Iv(27959118,  27959171	,  0),
            Iv(28866309,  33141404	,  0),
        ]
        lapper = IL.Lapper(data)

	found = collect(IL.find(lapper, 28974798, 33141355))
	@test found == [
            Iv(28866309,  33141404	,  0),
        ]
	bits = IL.Bits(lapper.intervals)
	@test IL.count(bits, 28974798, 33141355) == length(collect(IL.find(lapper, 28974798, 33141355)))
end


text:
module RNGPool

import Random123.Threefry4x

const RNG = Threefry4x{UInt64, 20}

let
  engines::Vector{RNG} = Vector{RNG}(undef, 0)
  global function getRNG()
    @inbounds return engines[Threads.threadid()]
  end
  global function getRNG(i::Int64)
    @inbounds return engines[i]
  end
  global function setRNGs(v::Int64)
    Threads.@threads for i = 1:length(engines)
      seed = (0,0,0,v+i)
      @inbounds engines[Threads.threadid()] = Threefry4x(seed)
    end
  end
  ## happens at runtime to avoid false sharing
  global function initializeRNGs()
    engines = Vector{RNG}(undef, Threads.nthreads())
    Threads.@threads for i = 1:length(engines)
      @inbounds engines[Threads.threadid()] = Threefry4x()
    end
  end
end

function __init__()
  initializeRNGs()
end

export RNG, getRNG, setRNGs

end


text:
# ----------------------------------------------------
# --- Find.jl
# ----------------------------------------------------
# Methods to scan for SDR devices, based on backend specifications or parameters

"""
Scan interface and returns the founded SDR
# --- Syntax
sdr = scan()
sdr = scan(backend;key...)
# Input parameter
If the function is called without parameters il will search for all avaliable backends such as UHDBindings and AdalmPluto. Otherwise the search will be limited to the desired backend
The optionnal arguments are the one supported by UHDBindings and AdalmPluto. See `uhd_find_devices()` in UHDBindings and `scan` function in AdalmPluto
# Keywords
- args : String used in UHD backend to specify USRP IP address. Example: scan(:uhd;args="addr=192.168.10.16")
- backend : Sring used in Pluto backend to specify the interface used ("local", "xml", "ip", "usb")
"""
function scan(backend::Union{Nothing,Vector{Symbol}}=nothing;key...)
    # --- If call w/o argument we search for all potential backends
    # Note that we can not search for SDROverNetwork, RadioSims and RTLSDR
    # TODO => Scan methods for RTLSDR ?
    if isnothing(backend)
        backend = [:uhd;:pluto;:rtlsdr]
    end
    allStr = String[];
    for b in backend
        if b == :uhd
            println("----------------------------")
            println("--- Scan for UHD devices ---")
            println("----------------------------")
            # ----------------------------------------------------
            # --- UHD Find device call
            # ----------------------------------------------------
            # --- Restrict keywords to uhd_find_devices
            key = parseKeyword(key,[:args])
            # scan keyword is uhd_find_devices parameter so we should handle empty case
            (isempty(key)) && (key[:args] = "")
            # --- Call scanner
            e  = UHDBindings.uhd_find_devices(key[:args])
            # --- Return the direct IP address based on the str cal
            for eN in e
                eM = eN[findfirst("_addr=",eN)[end] .+ (1:13)]
                push!(allStr,eM)
            end
        elseif b == :pluto
            println("------------------------------")
            println("--- Scan for Pluto devices ---")
            println("------------------------------")
            # ----------------------------------------------------
            # --- Scan call
            # ----------------------------------------------------
            # --- Restrict keywords to pluto
            key = parseKeyword(key,[:backend])
            if (isempty(key))
                # By default we look for all available backends
                backend = ["usb","local","ip","xml"]
            else
                # Focus on given backend
                backend = [key[:backend]]
            end
            for b in backend
                # --- Call scanner
                e = AdalmPluto.scan(b)
                # --- Push in Vector of string
                # AdalmPluto answer "" and this corresponds to nothing interessting. We push in the vector only if what we had was not empty
                (!isempty(e)) && (push!(allStr,e))
            end
        elseif b == :radiosim
            # ----------------------------------------------------
            # --- Radiosims backend
            # ----------------------------------------------------
            # No need to worry, we always have the simulated backend
            # Returns an arbitrary String with description
            println("------------------------------")
            println("--- Scan for RadioSims    ---")
            println("------------------------------")
            push!(allStr,"RadioSim  backend is always there !")
        elseif b == :rtlsdr
            # ----------------------------------------------------
            # --- RTLSDR scanning mode
            # ----------------------------------------------------
            println("------------------------------")
            println("--- Scan for RTL-SDR       ---")
            println("------------------------------")
            RTLSDRBindings.scan()
            push!(allStr,"RTL-SDR USB dongle")
        end
    end
    return allStr
end
scan(backend::Symbol;key...) = scan([backend];key...)


text:
using Documenter

# Print `@debug` statements (https://github.com/JuliaDocs/Documenter.jl/issues/955)
if haskey(ENV, "GITHUB_ACTIONS")
    ENV["JULIA_DEBUG"] = "Documenter"
end

using AbstractMCMC
using Random

DocMeta.setdocmeta!(
    AbstractMCMC,
    :DocTestSetup,
    :(using AbstractMCMC);
    recursive=true,
)

makedocs(;
    sitename="AbstractMCMC",
    format=Documenter.HTML(),
    modules=[AbstractMCMC],
    pages=[
        "Home" => "index.md",
        "api.md",
        "design.md",
    ],
    strict=true,
    checkdocs=:exports,
)

deploydocs(;
    repo="github.com/TuringLang/AbstractMCMC.jl.git", push_preview=true
)


text:
#=
Problem 1
If we list all the natural numbers below 10 that are multiples of 3 or 5, we get 3, 5, 6 and 9. The sum of these multiples is 23.
Find the sum of all the multiples of 3 or 5 below 1000.
=#
sum = 0
for i = 1:999
  if (i % 3 == 0 || i % 5 == 0)
    sum += i
  end
end

println("The sum is: $(sum)")

max_stars_repo_name: mebrunet/project-euler

max_issues_repo_name: mebrunet/project-euler

max_forks_repo_name: mebrunet/project-euler


text:
using CDSAPI
using ZipFile, GZip, Tar
using GRIB, NetCDF

using Test

# list of tests
testfiles = [
    "py2ju.jl",
    "retrieve.jl",
]

@testset "CDSAPI.jl" begin
    for testfile in testfiles
        include(testfile)
    end
end


text:
################################################################################
## Linear Regression
##   y ~ N(b0 + b1 * x, s2)
##   b0, b1 ~ N(0, 1000)
##   s2 ~ invgamma(0.001, 0.001)
################################################################################

using Mamba

## Data
data = Dict{Symbol, Any}(
  :x => [1, 2, 3, 4, 5],
  :y => [1, 3, 3, 3, 5]
)

## Log-transformed unnormalized joint posterior for b0, b1, and log(s2)
logf = function(x::DenseVector)
   b0 = x[1]
   b1 = x[2]
   logs2 = x[3]
   (-0.5 * length(data[:y]) - 0.001) * logs2 -
     (0.5 * sumabs2(data[:y] - b0 - b1 * data[:x]) + 0.001) / exp(logs2) -
     0.5 * b0^2 / 1000 - 0.5 * b1^2 / 1000
end

## Log-transformed unnormalized full conditional densities for the model
## parameters beta and log(s2) defined below in the MCMC simulation
logf_beta(x) = logf([x; logs2])
logf_logs2(x) = logf([beta; x])

## MCMC simulation
n = 10000
burnin = 1000
sim = Chains(n, 3, names = ["b0", "b1", "s2"])
beta = AMWGVariate([0.0, 0.0], 1.0, logf_beta)
logs2 = SliceMultivariate([0.0], 5.0, logf_logs2)
for i in 1:n
  sample!(beta, adapt = (i <= burnin))
  sample!(logs2)
  sim[i, :, 1] = [beta; exp(logs2)]
end
describe(sim)


text:
export UpdateElo

"""
    UpdateElo <: UpdateRule

 Update using standard Elo where K is constant, but distribution can be set (with constant parameters)
     Elo2 has logistic distribution, but allows varying K

## Parameters
* `r0::Float64`: default rating
* `K::Float64`: the value of the gain parameter (use Elo2 for variable K)
* `dist::ContinuousUnivariateDistribution`: the model to use for performance as a function of rating, e.g, Normal or Logistic
* the distribution can have its own parameters, e.g., `β::Float64`, the spread parameter for the Logistic

```
"""
struct UpdateElo <: UpdateRule
    r0::Float64
    K::Float64
    dist::ContinuousUnivariateDistribution
    function UpdateElo(r0, K, dist)
        @check_args(UpdateElo, K >= zero(K))
        @check_args(UpdateElo, r0 >= zero(r0))
        new(r0, K, dist)
    end
end
# defaults imply the distribution is  Logistic(μ=0.0, θ=100.0)
#   https://en.wikipedia.org/wiki/Logistic_distribution
#   https://juliastats.github.io/Distributions.jl/latest/univariate.html#Continuous-Distributions-1
#      mean=mu=0.0, sdev=π θ / √3
default_θ = 400 / log(10) # this results in the standard values used for Elo in Chess and many other cases
UpdateElo(;r0=1500.0, K=32.0, θ=default_θ) = UpdateElo( r0, K, Logistic(0.0, θ) )

function update_info( rule::UpdateElo )
    info = Dict(
                :name => "Elo",
                :reference => "\"Whos's #1\", Langville and Meyer, p.?? and ??",
                :computation => "simultaneous",
                :state_model => "recursive",
                :input => "outcome",
                :output => "probabilistic",
                :model => "single",
                :ties => true,
                :factors => true,
                :parameters => ["r0, default rating", "K, gain", "dist(), performance model"],
                :record => false
                )
    return info
end
update_info( ::Type{UpdateElo} ) =  update_info( UpdateElo() )

function update_ratings( rule::UpdateElo,
                         input_ratings::RatingsList,
                         input_competitions::DataFrame)
    n = size(input_competitions,1)
    m = length( input_ratings.players )
    I = player_indexes( input_ratings.players )

    # set initial ratings if they aren't already
    old_r = copy( input_ratings.ratings )
    for p in input_ratings.players
        if !haskey( old_r, p )
            old_r[p] = rule.r0
        end
    end
    new_r = copy(old_r) # start with the old ratings, in case one isn't modified at all

    # perform update for each outcome
    d = input_competitions # just an abbreviation
    for i=1:n
        pA = d[i,PlayerA]
        pB = d[i,PlayerB]
        old_rA = old_r[ pA ]
        old_rB = old_r[ pB ]
        (ΔrA, ΔrB) = update( rule, old_rA, old_rB, d[i,Outcome], missing, missing)
        new_r[ pA ] += ΔrA
        new_r[ pB ] += ΔrB
    end

    # output ratings list
    output_ratings = RatingsList(input_ratings.players, new_r )
    return output_ratings
end

function update(rule::Union{UpdateElo,UpdateEloF},
                ratingA::Real, ratingB::Real, outcome::Real,
                factorA::Union{Missing,Real}, factorB::Union{Missing,Real})
    (expectedA, expectedB, tie) = predict_outcome(rule,  ratingA, ratingB, factorA, factorB)
    outcomeA = (sign( outcome) + 1)/2
    outcomeB = (sign(-outcome) + 1)/2
    ΔrA = rule.K*( outcomeA - expectedA )
    ΔrB = rule.K*( outcomeB - expectedB )
    return ΔrA, ΔrB
end

function predict_outcome(rule::Union{UpdateElo,UpdateEloF},
                         ratingA::Real, ratingB::Real,
                         factorA::Union{Missing,Real}, factorB::Union{Missing,Real})
    rating_diff = ratingA - ratingB
    factorA = coalesce(factorA, 0.0) # replace missing values with 0.0
    factorB = coalesce(factorB, 0.0) # replace missing values with 0.0
    factors = (factorA - factorB)
    expectedA = cdf( rule.dist,  rating_diff + factors )  # do these separately in case asymmetric distribution is chosen
    expectedB = cdf( rule.dist, -rating_diff - factors )  # do these separately in case asymmetric distribution is chosen
    # at present no model for ties
    return (expectedA, expectedB, 0.0 )
end

# function simulate_outcome(rule::Union{UpdateElo,UpdateEloF},
#                           ratingA::Real, ratingB::Real,
#                           factorA::Union{Missing,Real}, factorB::Union{Missing,Real})
#     (expectedA, expectedB, tie ) = predict_outcome(rule, ratingA, ratingB, factorA, factorB)
# end


text:
using Base.Test
using QuantumOptics

alpha = 0.5

fockbasis = FockBasis(2)
spinbasis = SpinBasis(1//2)

a = destroy(fockbasis)
at = create(fockbasis)
n = number(fockbasis)

sx = sigmax(spinbasis)
sy = sigmay(spinbasis)
sz = sigmaz(spinbasis)

Ispin = identity(spinbasis)
Ifock = identity(fockbasis)

# # LazyTensor
# basis = tensor(spinbasis, fockbasis, spinbasis)
# psi_ket = tensor(basis_ket(spinbasis, 2), coherentstate(fockbasis, alpha), basis_ket(spinbasis, 1))
# psi_bra = dagger(psi_ket)

# op1 = LazyTensor(basis, basis, 2, a)
# op2 = LazyTensor(basis, basis, 3, sy)
# op3 = LazyTensor(basis, basis, [2,3], [a,sy])
# op4 = LazyTensor(basis, basis, [1,2,3], [sx,a,sy])

# result_ket = deepcopy(psi_ket)
# result_bra = deepcopy(psi_bra)

# @test Ispin ⊗ a ⊗ Ispin == sparse(op1)
# @test Ispin ⊗ Ifock ⊗ sy == sparse(op2)
# @test Ispin ⊗ a ⊗ sy == sparse(op3)
# @test sx ⊗ a ⊗ sy == sparse(op4)

# @test full(Ispin ⊗ a ⊗ Ispin) == full(op1)
# @test full(Ispin ⊗ Ifock ⊗ sy) == full(op2)
# @test full(Ispin ⊗ a ⊗ sy) == full(op3)
# @test full(sx ⊗ a ⊗ sy) == full(op4)

# for op=[op1, op2, op3, op4]
#     @test typeof(op) == LazyTensor
#     operators.gemv!(Complex(1.), op, psi_ket, Complex(0.), result_ket)
#     @test_approx_eq_eps 0. norm(full(op)*psi_ket - result_ket) 1e-12
#     operators.gemv!(Complex(1.), psi_bra, op, Complex(0.), result_bra)
#     @test_approx_eq_eps 0. norm(psi_bra*full(op) - result_bra) 1e-12
# end

# # LazySum
# psi_ket = coherentstate(fockbasis, alpha)
# psi_bra = dagger(psi_ket)
# op = LazySum(a, at)
# @test typeof(op) == LazySum
# @test (a+at) == sparse(op)
# @test full(a+at) == full(op)

# @test_approx_eq_eps 0. norm(op*psi_ket - (a+at)*psi_ket) 1e-12

# psi_ket2 = Ket(fockbasis)
# operators.gemv!(Complex(1.), LazySum(a, at, at*a), psi_ket, Complex(0.), psi_ket2)
# @test_approx_eq_eps 0. norm((a+at+at*a)*psi_ket - psi_ket2) 1e-12

# psi_bra2 = Bra(fockbasis)
# operators.gemv!(Complex(1.), psi_bra, LazySum(a, at, at*a), Complex(0.), psi_bra2)
# @test_approx_eq_eps 0. norm(psi_bra*(a+at+at*a) - psi_bra2) 1e-12

# # LazyProduct
# n_lazy = LazyProduct(at, a)
# @test_approx_eq_eps 0. norm(n_lazy*psi_ket - n*psi_ket) 1e-12
# @test_approx_eq_eps 0. norm(psi_bra*n_lazy - psi_bra*n) 1e-12

# zero_op = n_lazy - n
# @test typeof(zero_op) == LazySum
# @test (at*a) == sparse(n_lazy)
# @test full(at*a) == full(n_lazy)

# @test_approx_eq_eps 0. norm(zero_op*psi_ket) 1e-12
# @test_approx_eq_eps 0. norm(psi_bra*zero_op) 1e-12

# psi_ket_ = deepcopy(psi_ket)
# psi_ket2 = Ket(fockbasis)
# operators.gemv!(Complex(1.), LazyProduct(a, at, at*a), psi_ket_, Complex(0.), psi_ket2)
# @test_approx_eq_eps 0. norm(a*at*at*a*psi_ket - psi_ket2) 1e-12

# psi_bra_ = deepcopy(psi_bra)
# psi_bra2 = Bra(fockbasis)
# operators.gemv!(Complex(1.), psi_bra_, LazyProduct(a, at, at*a), Complex(0.), psi_bra2)
# @test_approx_eq_eps 0. norm(psi_bra*a*at*at*a - psi_bra2) 1e-12

# LazySum of LazyTensor
basis = tensor(spinbasis, fockbasis)
psi_ket = tensor(basis_ket(spinbasis, 2), coherentstate(fockbasis, alpha))
psi_bra = dagger(psi_ket)

op1 = LazyTensor(basis, basis, 1, sy)
op2 = LazyTensor(basis, basis, 2, a)
op3 = LazyTensor(basis, basis, [1,2], [sx+sy, a])

op = LazySum(op1, op2, op3)

result_ket = deepcopy(psi_ket)
result_ket2 = deepcopy(psi_ket)
result_bra = deepcopy(psi_bra)

@test sy ⊗ Ifock == sparse(op1)
@test Ispin ⊗ a == sparse(op2)
@test (sx+sy) ⊗ a == sparse(op3)
@test sy ⊗ Ifock + Ispin ⊗ a + (sx+sy) ⊗ a == sparse(op)

@test full(sy ⊗ Ifock) == full(sparse(op1))
@test full(Ispin ⊗ a) == full(sparse(op2))
@test full((sx+sy) ⊗ a) == full(sparse(op3))
@test full(sy ⊗ Ifock + Ispin ⊗ a + (sx+sy) ⊗ a) == full(sparse(op))

# op = LazySum(sparse(op1))

# operators.gemv!(Complex(1.), op2, psi_ket, Complex(0.), result_ket)
# operators.gemv!(Complex(1.), sparse(op2), psi_ket, Complex(0.), result_ket2)
# # operators.gemv!(Complex(1.), op1, psi_ket, Complex(1.), result_ket)
# @test_approx_eq_eps 0. norm(result_ket2 - result_ket) 1e-12

@test typeof(op) == LazySum
operators.gemv!(Complex(1.), op, psi_ket, Complex(0.), result_ket)
@test_approx_eq_eps 0. norm(full(op)*psi_ket - result_ket) 1e-12
operators.gemv!(Complex(1.), psi_bra, op, Complex(0.), result_bra)
@test_approx_eq_eps 0. norm(psi_bra*full(op) - result_bra) 1e-12


text:
using AcousticFWI,Seismic,PyPlot

function main()

    ns = 4
    ng = 250
    nz = 150
    nx = 250
    nt = 1024
    nf = 1024
    dz = 0.01
    dx = 0.01
    dt = 0.002
    f0 = 10.
    fmin = 0.5
    fmax = 30.
    ext = 50
    atten_max = 2.5

    vp = 2.1*ones(nz,nx)
    vp[61:100,:] = 3.0
    vp[101:end,:] = 2.7

    ot = zeros(ns)
    isz = 3*ones(Int,ns)
    isx = round(Int,collect(linspace(2,nx-1,ns)))
    igz = 3*ones(Int,ng)
    igx = round(Int,collect(linspace(1,nx,ng)))

    wav = Ricker(f0=f0,dt=dt)
    d = zeros(nt,ng,ns)
    for ishot = 1:ns
        u = Helmholtz(isz[ishot],isx[ishot],ot[ishot],vp,wav,fmin,fmax,nf,nt,dz,dx,dt,ext,atten_max)
        d[:,:,ishot] = u[:,3,igx]
    end

    vp0 = 2.5*ones(vp)

    # Perform 1 iteration of FWI to check if the Sensitivity operator passes the dot-product test

    nz = size(vp0,1) + 2*ext
    nx = size(vp0,2) + 2*ext

    L = Laplacian(nz,nx,dz,dx)
    M = MassMatrix(vp0,nz,nx,ext)
    A = Attenuation(nz,nx,ext,atten_max)
    R = Restriction(nz,nx,ext,igz,igx)

    D = fft([d ; zeros(nf-nt,ng,ns)],1)
    WAV = fft([wav;zeros(nf-length(wav))])
    fs = 1/dt
    df = fs/nf
    faxis = fftshift(-fs/2:df:fs/2-df)
    waxis = 2*pi*faxis

    _,iwmin = findmin(abs(fmin-faxis))
    _,iwmax = findmin(abs(fmax-faxis))

    param = Dict(:R=>R,:A=>A)

    iw = 15
    w = waxis[iw]
    s = zeros(eltype(D),nz*nx,ns)
    for ishot = 1:ns
        s[:,ishot] = Source(isz[ishot],isx[ishot],ot[ishot],WAV,waxis,w,nz,nx,ext)
    end
    H = lufact(L + w^2*M*A)
    U = zeros(eltype(D),nz*nx,ns)
    r = zeros(eltype(D),ng*ns)
    for ishot = 1:ns
        U[:,ishot] = H\s[:,ishot]
        r[(ishot-1)*ng+1:ishot*ng] = R*U[:,ishot] - D[iw,:,ishot]
    end
    param[:H] = H
    param[:U] = U
    param[:w] = w

    m = rand(nz*nx) + im*rand(nz*nx)
    r = rand(ns*ng) + im*rand(ns*ng)

    R = SensitivityMultiShot(m,false;param...)
    M = SensitivityMultiShot(r,true;param...)

    dotprodtest = (m[:]'*M[:])[1] - (r[:]'*R[:])[1]
    println("This number should be close to zero: ",dotprodtest)

end

main()


text:
using FastDEQ
using CUDA
using Flux
using FluxExperimental
using LinearAlgebra
using Random
using Test

@testset "FastDEQ.jl" begin
    mse_loss_function = SupervisedLossContainer(loss_function = Flux.Losses.mse)

    @info "Testing DEQ"
    Random.seed!(0)

    model = gpu(DEQChain(Dense(2, 2),
                         DeepEquilibriumNetwork(Parallel(+, Dense(2, 2; bias=false), Dense(2, 2; bias=false)),
                                                ContinuousDEQSolver(;abstol=0.1f0, reltol=0.1f0))))
    x = gpu(rand(Float32, 2, 1))
    y = gpu(rand(Float32, 2, 1))
    ps = Flux.params(model)
    gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    for _p in ps
        @test all(isfinite.(gs[_p]))
    end

    @info "Testing SkipDEQ"
    Random.seed!(0)

    model = gpu(DEQChain(Dense(2, 2),
                         SkipDeepEquilibriumNetwork(Parallel(+, Dense(2, 2), Dense(2, 2)),
                                                    Dense(2, 2),
                                                    ContinuousDEQSolver(;abstol=0.1f0, reltol=0.1f0);
                                                    sensealg=SteadyStateAdjoint(0.1f0, 0.1f0, 10))))
    x = gpu(rand(Float32, 2, 1))
    y = gpu(rand(Float32, 2, 1))
    ps = Flux.params(model)
    gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    for _p in ps
        @test all(isfinite.(gs[_p]))
    end

    @info "Testing SkipDEQ V2"
    Random.seed!(0)

    model = gpu(DEQChain(Dense(2, 2),
                         SkipDeepEquilibriumNetwork(Parallel(+, Dense(2, 2), Dense(2, 2)),
                                                    ContinuousDEQSolver(;abstol=0.1f0, reltol=0.1f0))))
    x = gpu(rand(Float32, 2, 1))
    y = gpu(rand(Float32, 2, 1))
    ps = Flux.params(model)
    gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    for _p in ps
        @test all(isfinite.(gs[_p]))
    end

    @info "Testing Broyden Solver"
    Random.seed!(0)

    model = gpu(DEQChain(Conv((3, 3), 1 => 1, relu; pad=1, stride=1),
                         SkipDeepEquilibriumNetwork(Parallel(+, Conv((3, 3), 1 => 1, relu; pad=1, stride=1),
                                                                Conv((3, 3), 1 => 1, relu; pad=1, stride=1)),
                                                    Conv((3, 3), 1 => 1, relu; pad=1, stride=1),
                                                    DiscreteDEQSolver(BroydenSolver; abstol=0.001f0,
                                                                      reltol=0.001f0, device=gpu, original_dims=(8 * 8, 1),
                                                                      batch_size=4);
                                                    sensealg=SteadyStateAdjoint(0.1f0, 0.1f0, 10))))
    x = gpu(rand(Float32, 8, 8, 1, 4))
    y = gpu(rand(Float32, 8, 8, 1, 4))
    ps = Flux.params(model)
    gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    for _p in ps
        @test all(isfinite.(gs[_p]))
    end

    @info "Testing L-Broyden Solver"
    Random.seed!(0)

    model = gpu(DEQChain(Conv((3, 3), 1 => 1, relu; pad=1, stride=1),
                         SkipDeepEquilibriumNetwork(Parallel(+, Conv((3, 3), 1 => 1, relu; pad=1, stride=1),
                                                             Conv((3, 3), 1 => 1, relu; pad=1, stride=1)),
                                                    Conv((3, 3), 1 => 1, relu; pad=1, stride=1),
                                                    DiscreteDEQSolver(LimitedMemoryBroydenSolver; abstol=0.001f0,
                                                                      reltol=0.001f0, device=gpu, original_dims=(8 * 8, 1),
                                                                      batch_size=4);
                                                    sensealg=SteadyStateAdjoint(0.1f0, 0.1f0, 10))))
    x = gpu(rand(Float32, 8, 8, 1, 4))
    y = gpu(rand(Float32, 8, 8, 1, 4))
    ps = Flux.params(model)
    gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    for _p in ps
        @test all(isfinite.(gs[_p]))
    end

    @info "Testing MultiScaleDEQ"
    Random.seed!(0)

    model = gpu(MultiScaleDeepEquilibriumNetwork((Parallel(+, Dense(4, 4, tanh_fast), Dense(4, 4, tanh_fast)),
                                                  Dense(3, 3, tanh_fast), Dense(2, 2, tanh_fast),
                                                  Dense(1, 1, tanh_fast)),
                                                 [NoOpLayer() Dense(4, 3, tanh_fast) Dense(4, 2, tanh_fast) Dense(4, 1, tanh_fast);
                                                  Dense(3, 4, tanh_fast) NoOpLayer() Dense(3, 2, tanh_fast) Dense(3, 1, tanh_fast);
                                                  Dense(2, 4, tanh_fast) Dense(2, 3, tanh_fast) NoOpLayer() Dense(2, 1, tanh_fast);
                                                  Dense(1, 4, tanh_fast) Dense(1, 3, tanh_fast) Dense(1, 2, tanh_fast) NoOpLayer()],
                                                  ContinuousDEQSolver(;abstol=0.1f0, reltol=0.1f0)))
    x = gpu(rand(Float32, 4, 2))
    y = tuple([gpu(rand(Float32, i, 2)) for i in 4:-1:1]...)
    sol = model(x)
    ps = Flux.params(model)
    gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    for _p in ps
        @test all(isfinite.(gs[_p]))
    end

    @info "Testing MultiScaleSkipDEQ"
    Random.seed!(0)

    model = gpu(MultiScaleSkipDeepEquilibriumNetwork((Parallel(+, Dense(4, 4, tanh_fast), Dense(4, 4, tanh_fast)),
                                                      Dense(3, 3, tanh_fast), Dense(2, 2, tanh_fast),
                                                      Dense(1, 1, tanh_fast)),
                                                     [NoOpLayer() Dense(4, 3, tanh_fast) Dense(4, 2, tanh_fast) Dense(4, 1, tanh_fast);
                                                      Dense(3, 4, tanh_fast) NoOpLayer() Dense(3, 2, tanh_fast) Dense(3, 1, tanh_fast);
                                                      Dense(2, 4, tanh_fast) Dense(2, 3, tanh_fast) NoOpLayer() Dense(2, 1, tanh_fast);
                                                      Dense(1, 4, tanh_fast) Dense(1, 3, tanh_fast) Dense(1, 2, tanh_fast) NoOpLayer()],
                                                     (Dense(4, 4, tanh_fast), Dense(4, 3, tanh_fast),
                                                      Dense(4, 2, tanh_fast), Dense(4, 1, tanh_fast)),
                                                     ContinuousDEQSolver(;abstol=0.1f0, reltol=0.1f0)))
    x = gpu(rand(Float32, 4, 2))
    y = tuple([gpu(rand(Float32, i, 2)) for i in 4:-1:1]...)
    sol = model(x)
    ps = Flux.params(model)
    gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    for _p in ps
        @test all(isfinite.(gs[_p]))
    end

    # CI gives mutation error though it works locally.
    # @info "Testing MultiScaleSkipDEQV2"
    # Random.seed!(0)

    # model = gpu(MultiScaleSkipDeepEquilibriumNetwork((Parallel(+, Dense(4, 4, tanh_fast), Dense(4, 4, tanh_fast)),
    #                                                   Dense(3, 3, tanh_fast), Dense(2, 2, tanh_fast),
    #                                                   Dense(1, 1, tanh_fast)),
    #                                                  [NoOpLayer() Dense(4, 3, tanh_fast) Dense(4, 2, tanh_fast) Dense(4, 1, tanh_fast);
    #                                                   Dense(3, 4, tanh_fast) NoOpLayer() Dense(3, 2, tanh_fast) Dense(3, 1, tanh_fast);
    #                                                   Dense(2, 4, tanh_fast) Dense(2, 3, tanh_fast) NoOpLayer() Dense(2, 1, tanh_fast);
    #                                                   Dense(1, 4, tanh_fast) Dense(1, 3, tanh_fast) Dense(1, 2, tanh_fast) NoOpLayer()],
    #                                                  ContinuousDEQSolver(;abstol=0.1f0, reltol=0.1f0)))
    # x = gpu(rand(Float32, 4, 2))
    # y = tuple([gpu(rand(Float32, i, 2)) for i in 4:-1:1]...)
    # sol = model(x)
    # ps = Flux.params(model)
    # gs = Flux.gradient(() -> mse_loss_function(model, x, y), ps)
    # for _p in ps
    #     @test all(isfinite.(gs[_p]))
    # end
end


text:
# sample to show world module's graph
module WorldSample
  # include external modules
  include(joinpath(split(@__FILE__, "src")[1], "src/model/world/puddle_world.jl"))

  # main function
  # show world's graph
  function main(delta_time=0.1, end_time=5; is_test=false)
    # save path of output
    path = "src/model/world/world_sample.gif"

    # create instance
    world = PuddleWorld(-5.0, 5.0, -5.0, 5.0,
                        delta_time=delta_time,
                        end_time=end_time,
                        is_test=is_test,
                        save_path=path)

    # draw graph
    draw(world)
  end
end

max_stars_repo_name: ShisatoYano/JuliaAutonomy

max_issues_repo_name: ShisatoYano/JuliaAutonomy

max_forks_repo_name: ShisatoYano/JuliaAutonomy


text:
# this test is designed to test the ability of the package to use a state space that is not an AbstractArray.
# using Revise
using Test, Random, Setfield
using BifurcationKit
const BK = BifurcationKit
####################################################################################################
# test Bordered Arrays
_a = BorderedArray(zeros(2), zeros(2))
_b = BorderedArray(zeros(2), zeros(2))
BK.mul!(_a, 1., _b)
BK.axpby!(1., _a, 1., _b)
BK.getVec(_a)
BK.getVec(_a.u)
BK.getVec(BorderedArray(rand(2),1.))
BK.getP(BorderedArray(rand(2),1.))
####################################################################################################
# We start with a simple Fold problem
using LinearAlgebra
function F0(x::Vector, r)
	out = r .+  x .- x.^3
end

opt_newton0 = NewtonPar(tol = 1e-11, verbose = false)
	out0, hist, flag = newton(F0, [0.8], 1., opt_newton0)

opts_br0 = ContinuationPar(dsmin = 0.001, dsmax = 0.07, ds= -0.02, pMax = 4.1, pMin = -1., newtonOptions = setproperties(opt_newton0; maxIter = 70, tol = 1e-8), detectBifurcation = 0, maxSteps = 150)

BK.isStable(opts_br0, nothing)

br0, u1 = continuation(F0, out0, 1.0, (@lens _), opts_br0, recordFromSolution = (x, p) -> x[1])

outfold, hist, flag = newton(
		F0, (x, r) -> diagm(0 => 1 .- 3 .* x.^2),
		br0, 2;
		Jᵗ = (x, r) -> diagm(0 => 1 .- 3 .* x.^2),
		d2F = (x, r, v1, v2) -> -6 .* x .* v1 .* v2,)
	# flag && printstyled(color=:red, "--> We found a Fold Point at α = ",outfold.p, ", from ", br0.specialpoint[2].param, "\n")

####################################################################################################
# Here is a more involved example

function Fb(x::BorderedArray, p)
	r, s = p
	BorderedArray(r .+  s .* x.u .- (x.u).^3, x.p - 0.0)
end

# there is no finite differences defined, so we need to provide a linearsolve
# we could also have used GMRES. We define a custom Jacobian which will be used for evaluation and jacobian inverse
struct Jacobian
	x
	r
	s
end

# We express the jacobian operator
function (J::Jacobian)(dx)
	BorderedArray( (J.s .- 3 .* ((J.x).u).^2) .* dx.u, dx.p)
end

struct linsolveBd <: BK.AbstractBorderedLinearSolver end

function (l::linsolveBd)(J, dx)
	x = J.x
	r = J.r
	out = BorderedArray(dx.u ./ (J.s .- 3 .* (x.u).^2), dx.p)
	out, true, 1
end

sol0 = BorderedArray([0.8], 0.0)

opt_newton = NewtonPar(tol = 1e-11, verbose = false, linsolver = linsolveBd())
sol, hist, flag = newton(Fb, (x, p) -> Jacobian(x, 1., 1.), sol0, (1., 1.), opt_newton)

opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= -0.01, pMax = 4.1, pMin = -1., newtonOptions = setproperties(opt_newton; maxIter = 70, tol = 1e-8), detectBifurcation = 0, maxSteps = 150, saveSolEveryStep = 1)

	br, u1 = continuation(
		Fb, (x, r) -> Jacobian(x, r[1], r[2]),
		sol,  (1., 1.), (@lens _[1]),
		opts_br; recordFromSolution = (x,p) -> x.u[1])

BK.getSolx(br,1)
BK.getSolp(br,1)

# plot(br);title!("")

outfold, hist, flag = newton(
	Fb, (x, r) -> Jacobian(x, r[1], r[2]),
	br, 1;
	Jᵗ = (x, r) -> Jacobian(x, r[1], r[2]),
	d2F = (x, r, v1, v2) -> BorderedArray(-6 .* x.u .* v1.u .* v2.u, 0.),)
		# flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p, ", from ", br.specialpoint[1].param,"\n")

outfoldco, hist, flag = continuation(
	Fb, (x, r) -> Jacobian(x, r[1], r[2]),
	br, 1, (@lens _[2]), opts_br;
	Jᵗ = (x, r) -> Jacobian(x, r[1], r[2]),
	d2F = ((x, r, v1, v2) -> BorderedArray(-6 .* x.u .* v1.u .* v2.u, 0.)), plot = false)

# try with newtonDeflation
# test with Newton deflation 1
deflationOp = DeflationOperator(2, 1.0, [zero(sol)])
soldef0 = BorderedArray([0.1], 0.0)
soldef1,  = newton(
	Fb, (x, r) -> Jacobian(x, r[1], r[2]),
	soldef0, (0., 1.),
	opt_newton, deflationOp)

push!(deflationOp, soldef1)

Random.seed!(1231)
# test with Newton deflation 2
soldef2, = newton(
	Fb, (x, r) -> Jacobian(x, r[1], r[2]),
	rmul!(soldef0,rand()),  (0., 1.),
	opt_newton, deflationOp)

# test indexing
deflationOp[1]

# test length
length(deflationOp)

# test pop!
pop!(deflationOp)

# test empty
empty!(deflationOp)

####################################################################################################
using KrylovKit

function Fr(x::RecursiveVec, p)
	r, s = p
	out = similar(x)
	for ii=1:length(x)
		out[ii] .= r .+  s .* x[ii] .- x[ii].^3
	end
	out
end

# there is no finite differences defined, so we need to provide a linearsolve
# we could also have used GMRES. We define a custom Jacobian which will be used for evaluation and jacobian inverse
struct JacobianR
	x
	s
end

# We express the jacobian operator
function (J::JacobianR)(dx)
	out = similar(dx)
	for ii=1:length(out)
		out[ii] .= (J.s .- 3 .* (J.x[ii]).^2) .* dx[ii]
	end
	return out
end

struct linsolveBd_r <: BK.AbstractBorderedLinearSolver end

function (l::linsolveBd_r)(J, dx)
	x = J.x
	out = similar(dx)
	for ii=1:length(out)
		out[ii] .= dx[ii] ./ (J.s .- 3 .* (x[ii]).^2)
	end
	out, true, 1
end

opt_newton0 = NewtonPar(tol = 1e-10, maxIter = 50, verbose = false, linsolver = linsolveBd_r())
	out0, hist, flag = newton(
		Fr, (x, p) -> JacobianR(x, p[1]),
		RecursiveVec([1 .+ 0.1*rand(10) for _ = 1:2]), (0., 1.),
		opt_newton0)

Base.:copyto!(dest::RecursiveVec, in::RecursiveVec) = copyto!(dest.vecs, in.vecs)

opts_br0 = ContinuationPar(dsmin = 0.001, dsmax = 0.2, ds= -0.01, pMin = -1.1, pMax = 1.1, newtonOptions = opt_newton0)

br0, u1 = continuation(
	Fr, (x, p) -> JacobianR(x, p[1]),
	out0, (0.9, 1.), (@lens _[1]), opts_br0;
	plot = false,
	recordFromSolution = (x,p) -> x[1][1])

br0, u1 = continuation(
	Fr, (x, p) -> JacobianR(x, p[1]),
	out0, (0.9, 1.), (@lens _[1]), opts_br0;
	tangentAlgo = BorderedPred(),
	plot = false,
	recordFromSolution = (x,p) -> x[1][1])

outfold, hist, flag = newton(
	Fr, (x, p) -> JacobianR(x, p[1]),
	br0, 1; #index of the fold point
	Jᵗ = (x, r) -> JacobianR(x, r[1]),
	d2F = (x, r, v1, v2) -> RecursiveVec([-6 .* x[ii] .* v1[ii] .* v2[ii] for ii=1:length(x)]),)
		# flag && printstyled(color=:red, "--> We found a Fold Point at α = ", outfold.p, ", from ", br0.specialpoint[1].param,"\n")

outfoldco, hist, flag = continuation(
	Fr, (x, p) -> JacobianR(x, p[1]),
	br0, 1,	(@lens _[2]), opts_br0;
	Jᵗ = (x, s) -> JacobianR(x, s[1]),
	d2F = ((x, r, v1, v2) -> RecursiveVec([-6 .* x[ii] .* v1[ii] .* v2[ii] for ii=1:length(x)])),
	tangentAlgo = SecantPred(), plot = false)

outfoldco, hist, flag = continuation(
	Fr, (x, p) -> JacobianR(x, p[1]),
	br0, 1, (@lens _[2]), opts_br0;
	Jᵗ = (x, s) -> JacobianR(x, s[1]),
	d2F = ((x, r, v1, v2) -> RecursiveVec([-6 .* x[ii] .* v1[ii] .* v2[ii] for ii=1:length(x)])),
	tangentAlgo = BorderedPred(), plot = false)

# try with newtonDeflation
# test with Newton deflation 1
deflationOp = DeflationOperator(2, 1.0, [out0])
soldef1, = newton(
	Fr, (x, p) -> JacobianR(x, 0.),
	out0, (0., 1.0),
	opt_newton0, deflationOp)

push!(deflationOp, soldef1)


text:
using PolygonMinkowskiSum
using Test

import Unitful: °
using GeometryTypes: Point
using PolygonMinkowskiSum: SimplePolygon

function make_regular_ngon(n; offset=0°, radius=1)
    points = [radius .* Point(cos(θ+offset), sin(θ+offset))
        for θ = range(0°, 360°, length=(n+1))][1:end-1]
    SimplePolygon(points)
end

using Plots

diamond = make_regular_ngon(4; offset=0°)
square = make_regular_ngon(4; offset=45°)

big_square = minkowski_sum(square, square)
plot(big_square, markershape=:cross)

for offset = 0°:45°:360°
    @show offset
    diamond = make_regular_ngon(4; offset=0°)
    square = make_regular_ngon(4; offset=offset#=+45°=#, radius=0.5)

    octogon = minkowski_sum(square, diamond)
    p = plot(aspect_ratio=1.0)
    plot!(p ,diamond, fillalpha=0.2)
    plot!(p, square, fillalpha=0.2)
    plot!(p, octogon, fillalpha=0.2)
    display(p)
end

Plots.pyplot() # gr doesn't yet support `reuse=false`: https://github.com/JuliaPlots/Plots.jl/issues/1370
ps = DefaultDict(key->plot(reuse=false, title=key, aspect_ratio=1), passkey=true)
minkowski_sum(sweepme, alongme, (ns=Plots, plots=ps))
map(display, values(ps))


text:
using Test
# using Revise
using JDF
using RDatasets

@testset "JDF.jl categorical arrays" begin
    df = DataFrame([collect(1:100) for i =1:3])
    df[!, :x1] = categorical(df[!, :x1])
    df[!, :x2] = categorical(string.(df[!, :x2]))
    savejdf("a3cate.jdf", df)

    df2 = loadjdf("a3cate.jdf", cols=[:x2, :x1])
    @test size(df2, 2) == 2
    @test size(df2, 1) == 100
    @time df2[!, :x1] isa CategoricalVector{Int}
    @time df2[!, :x2] isa CategoricalVector{String}

    rm("a3cate.jdf", force=true, recursive=true)
end

@testset "Guard against Github #27" begin
    iris = dataset("datasets", "iris")
    savejdf(iris, "iris.jdf")
    loadjdf("iris.jdf")

    rm("iris.jdf", force = true, recursive = true)
end


text:
export ft_open

function ft_open(iDevice::Integer)
  ft_handle = Ref{Culong}()
  ft_status = ccall((:FT_Open, d2xx),
                      Cuint,
                      (Cuint,Ref{Culong}),
                      iDevice,ft_handle)
  checkstatus(ft_status)
  return ft_handle[]
end


text:
module Color

using FixedPointNumbers, Compat

typealias Fractional Union(FloatingPoint, FixedPoint)

import Base: ==, convert, eltype, hex, isless, linspace, show, typemin, typemax, writemime

if VERSION < v"0.4.0-dev+3275"
    import Base.Graphics: set_source, set_source_rgb, GraphicsContext
else
    import Graphics: set_source, set_source_rgb, GraphicsContext
end

export ColorValue, color,
       ColourValue,
       AlphaColorValue, AbstractAlphaColorValue, AbstractRGB,
       RGB, HSV, HSL, XYZ, xyY, Lab, LAB, Luv, LUV, LCHab, LCHuv, DIN99, DIN99d, DIN99o, LMS, RGB24,
       RGBA, HSVA, HSLA, XYZA, xyYA, LabA, LuvA, LCHabA, LCHuvA, DIN99A, DIN99dA, DIN99oA, LMSA, ARGB32,
       rgba, hsva, hsla, xyza, xyYa, laba, luva, lchaba, lchuva, din99a, din99da, din99oa, lmsa, argb32,
       colortype, weighted_color_mean, hex,
       protanopic, deuteranopic, tritanopic,
       distinguishable_colors,
       colordiff, DE_2000, DE_94, DE_JPC79, DE_CMC, DE_BFD, DE_AB, DE_DIN99, DE_DIN99d, DE_DIN99o,
       MSC, sequential_palette, diverging_palette, colormap,
       colormatch, CIE1931_CMF, CIE1964_CMF, CIE1931J_CMF, CIE1931JV_CMF

# The core; every other include will need these type definitions
include("colorspaces.jl")

# Early utilities
include("utilities.jl")

# Include other module components
include("conversions.jl")
include("algorithms.jl")
include("parse.jl")
include("differences.jl")
include("colormaps.jl")
include("display.jl")
include("colormatch.jl")

@deprecate RGBA32 ARGB32
@deprecate rgba32 argb32

end # module


text:
module CollisionAvoidanceGame

using JuMP: @NLconstraint, @NLexpression, @objective, @variable
using PartiallyObservedInverseGames.CostUtils: symbol

unique!(push!(LOAD_PATH, @__DIR__))
import TestDynamics

export generate_player_cost_model

function generate_player_cost_model(;
    player_idx,
    control_system::TestDynamics.ProductSystem,
    T,
    goal_position,
    weights = (; state_proximity = 1, state_velocity = 1, control_Δv = 1, control_Δθ = 1),
    cost_prescaling = (;
        state_proximity = 0.1,
        state_velocity = 1,
        control_Δv = 10,
        control_Δθ = 1,
    ),
    fix_costs = (; # encoding soft-constraints rather than preferences
        state_goal = 100,
        state_lane = 1.0,
        state_orientation = 2.0,
    ),
    x_lane_center = nothing,
    y_lane_center = nothing,
    target_speed = 0,
    target_orientation = if !isnothing(x_lane_center)
        0
    elseif !isnothing(y_lane_center)
        pi / 2
    else
        nothing
    end,
    prox_min_regularization = 0.1,
    T_activate_goalcost = T,
)
    state_indices = TestDynamics.state_indices(control_system, player_idx)
    input_indices = TestDynamics.input_indices(control_system, player_idx)

    opponent_position_indices = let
        opponent_indices = filter(!=(player_idx), eachindex(control_system.subsystems))
        map(opponent_indices) do jj
            TestDynamics.state_indices(control_system, jj)[1:2]
        end
    end

    function add_regularized_squared_distance!(opt_model, x_sub_ego, pos_opponent)
        @NLexpression(
            opt_model,
            [t = 1:T],
            (x_sub_ego[1, t] - pos_opponent[1, t])^2 +
            (x_sub_ego[2, t] - pos_opponent[2, t])^2 +
            prox_min_regularization
        )
    end

    function add_objective!(opt_model, x, u; weights)
        T = size(x, 2)
        @views x_sub_ego = x[state_indices, :]
        @views u_sub_ego = u[input_indices, :]
        @views opponent_positions = map(opponent_position_indices) do opp_position_idxs
            x[opp_position_idxs, :]
        end

        J̃ = (;
            state_goal = isnothing(goal_position) ? 0 :
                         sum(el -> el^2, x_sub_ego[1:2, T_activate_goalcost:T] .- goal_position),
            state_lane = let
                # Note the switched indices. This is on purpose. The lane cost for the lane along y acts
                # on the x position and vice versa.
                x_lane_cost =
                    isnothing(x_lane_center) ? 0 :
                    sum(el -> el^2, x_sub_ego[2, :] .- x_lane_center)
                y_lane_cost =
                    isnothing(y_lane_center) ? 0 :
                    sum(el -> el^2, x_sub_ego[1, :] .- y_lane_center)
                x_lane_cost + y_lane_cost
            end,
            state_proximity = let
                prox_cost = sum(opponent_positions) do pos_opponent
                    d_sq = add_regularized_squared_distance!(opt_model, x_sub_ego, pos_opponent)
                    prox_cost = @variable(opt_model, [t = 1:T])
                    @NLconstraint(opt_model, [t = 1:T], prox_cost[t] == -log(d_sq[t]))
                    prox_cost
                end
                sum(prox_cost)
            end,
            state_velocity = sum(el -> el^2, x_sub_ego[3, :] .- target_speed),
            state_orientation = isnothing(target_orientation) ? 0 :
                                sum(el -> el^2, x_sub_ego[4, :] .- target_orientation),
            control_Δv = sum(el -> el^2, u_sub_ego[1, :]),
            control_Δθ = sum(el -> el^2, u_sub_ego[2, :]),
        )
        @objective(
            opt_model,
            Min,
            sum(weights[k] * cost_prescaling[k] * J̃[k] for k in keys(weights)) +
            sum(fix_costs[k] * J̃[k] for k in keys(fix_costs)) * sum(weights) / length(weights)
        )
    end

    function add_objective_gradients!(opt_model, x, u; weights)
        n_states, T = size(x)
        n_controls = size(u, 1)
        @views x_sub_ego = x[state_indices, :]
        @views u_sub_ego = u[input_indices, :]
        @views opponent_positions = map(opponent_position_indices) do opp_position_idxs
            x[opp_position_idxs, :]
        end

        dprox_dxy = sum(opponent_positions) do pos_opponent
            d_sq = add_regularized_squared_distance!(opt_model, x_sub_ego, pos_opponent)
            dproxdx = @variable(opt_model, [t = 1:T])
            @NLconstraint(
                opt_model,
                [t = 1:T],
                dproxdx[t] == -2 * (x_sub_ego[1, t] - pos_opponent[1, t]) / d_sq[t]
            )
            dproxdy = @variable(opt_model, [t = 1:T])
            @NLconstraint(
                opt_model,
                [t = 1:T],
                dproxdy[t] == -2 * (x_sub_ego[2, t] - pos_opponent[2, t]) / d_sq[t]
            )
            [dproxdx'; dproxdy']
        end

        dgoal_dxy =
            isnothing(goal_position) ? zeros(2, T) :
            hcat(
                zeros(2, T_activate_goalcost - 1),
                2 * (x_sub_ego[1:2, T_activate_goalcost:T] .- goal_position),
            )

        # Note: x-lane acts on y gradient and vice versa.
        dlane_dxy = let
            dlane_dx =
                isnothing(y_lane_center) ? zeros(T) : 2 * (x_sub_ego[1, :] .- y_lane_center)
            dlane_dy = isnothing(x_lane_center) ? zeros(T) : 2 * (x_sub_ego[2, :] .- x_lane_center)

            [dlane_dx'; dlane_dy']
        end

        dorientation_dθ =
            isnothing(target_orientation) ? zeros(T)' : 2 * (x_sub_ego[4, :] .- target_orientation)'

        # TODO: Technically this is missing the negative gradient on the opponents state but we
        # can't control that anyway (certainly not in OL Nash). Must be fixed for non-decoupled
        # systems and potentially FB Nash.
        dJdx = let
            dJ̃dx_sub = (;
                state_goal = [dgoal_dxy; zeros(2, T)],
                state_lane = [dlane_dxy; zeros(2, T)],
                state_proximity = [dprox_dxy; zeros(2, T)],
                state_velocity = [zeros(2, T); 2 * (x_sub_ego[3, :] .- target_speed)'; zeros(1, T)],
                state_orientation = [zeros(3, T); dorientation_dθ],
                control_Δv = zeros(size(x_sub_ego)),
                control_Δθ = zeros(size(x_sub_ego)),
            )
            dJdx_sub =
                sum(
                    weights[k] * cost_prescaling[symbol(k)] * dJ̃dx_sub[symbol(k)]
                    for k in keys(weights)
                ) +
                sum(fix_costs[k] * dJ̃dx_sub[k] for k in keys(fix_costs)) * sum(weights) /
                length(weights)
            [
                zeros(first(state_indices) - 1, T)
                dJdx_sub
                zeros(n_states - last(state_indices), T)
            ]
        end

        dJdu = let
            dJdu_sub =
                2 * [weights[:control_Δv], weights[:control_Δθ]] .*
                [cost_prescaling[:control_Δv], cost_prescaling[:control_Δθ]] .* u_sub_ego
            [
                zeros(first(input_indices) - 1, T)
                dJdu_sub
                zeros(n_controls - last(input_indices), T)
            ]
        end

        (; dx = dJdx, du = dJdu)
    end

    (; player_inputs = input_indices, weights, add_objective!, add_objective_gradients!)
end

end