text:
using MyPkg
using Documenter

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

makedocs(;
    modules=[MyPkg],
    authors="SatoshiTerasaki <terasakisatoshi.math@gmail.com> and contributors",
    repo="https://github.com/terasakisatoshi/jldev_poetry/blob/{commit}{path}#{line}",
    sitename="jldev_poetry",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://terasakisatoshi.github.io/jldev_poetry",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/terasakisatoshi/jldev_poetry",
    devbranch="main",
)


text:
const ALLOWED_FUNCTIONS = Set(union([:+,:-,:/,:*,:^,:\,:%,:!,:(==),:(!=),:<,:(<=),:>,:(>=),:(//)],first.(Calculus.symbolic_derivative_1arg_list)))

### Functions for checking strings before we parse/eval them

# Does it contain a function head?
expr_has_head(s, h) = false
expr_has_head(e::Expr, h::Symbol) = expr_has_head(e, Symbol[h])
function expr_has_head(e::Expr, vh::Vector{Symbol})
    in(e.head, vh) || any(a -> expr_has_head(a, vh), e.args)
end

# Does it contain a function definition?
has_function_def(s::String) = has_function_def(parse(s; raise=false))
function has_function_def(e::Expr)
    expr_has_head(e, Symbol[:(->), :function]) ||
    expr_has_head(e, Symbol[:macro]) ||
    # one line funtion definition:
    (expr_has_head(e, :(=)) && expr_has_head(e.args[], :call))
end

# Does it contain a block definition?
has_block_def(s::String) = has_block_def(parse(s; raise=false))
function has_block_def(e::Expr)
    expr_has_head(e, Symbol[:for]) ||
    expr_has_head(e, Symbol[:while]) ||
    expr_has_head(e, Symbol[:(->), :ccall]) ||
    expr_has_head(e, Symbol[:if]) ||
    expr_has_head(e, Symbol[:begin]) ||
    expr_has_head(e, Symbol[:let]) ||
    expr_has_head(e, Symbol[:do]) ||
    expr_has_head(e, Symbol[:try])
end

# Does it contain a macro using the @-form?
has_macro(s::String) = search(s,'@')!=0

# Does it call any function that's not in a whitelist of basic mathematical functions?
has_bad_call(s) = false
has_bad_call(s::String) = has_bad_call(parse(s; raise=false))
function has_bad_call(e::Expr)
    (in(e.head, Symbol[:call]) && e.args[1] ∉ ALLOWED_FUNCTIONS)|| any(a -> has_bad_call(a), e.args)
end

function sanitize_string(exstr)
  if has_function_def(exstr)
      error("Don't define functions in your system of equations...")
  end
  if has_block_def(exstr)
      error("Don't use Julia control flow blocks in your equations...")
  end
  if has_macro(exstr)
    error("Don't use macros you fowl demon. I'm watching you...")
  end
  if has_bad_call(exstr)
    error("A function which is not allowed was detected. If you're trying to be sneaky, then you're a bad person.")
  end
end

ext: jl

max_stars_repo_name: abhishalya/DiffEqOnlineServer

max_issues_repo_name: abhishalya/DiffEqOnlineServer

max_forks_repo_name: abhishalya/DiffEqOnlineServer


text:
@testset "Independent Bounds" begin
    m = BabyPOMDP()
    rs = MemorizingSource(1,1,MersenneTwister(37))
    sb = ScenarioBelief([1=>true], rs, 0, false)
    b = IndependentBounds(0.0, -1e-5)
    @test bounds(b, m, sb) == (0.0, -1e-5)
    b = IndependentBounds(0.0, -1e-5, consistency_fix_thresh=1e-5)
    @test bounds(b, m, sb) == (0.0, 0.0)
    b = IndependentBounds(0.0, -1e-4, consistency_fix_thresh=1e-5)
    @test bounds(b, m, sb) == (0.0, -1e-4)
end


text:
function AssetPricingTest3bPs(Re1,Re2,f,h,bandwidth=0,q=0,PointEstOnlyQ=false)
#AssetPricingTest3bPs  Testing difference of two linear factor models by testing the
#                      difference of the intercepts, or some other combination of
#                      parameters
#
#
#
#
#  Usage:  fnOutput = AssetPricingTest3bPs(Re1,Re2,f,h[,bandwidth[,q[,PointEstOnlyQ]]])
#
#  Input:   Re1             Txn1, test assets
#           Re2             Txn2, test assets     (in many applications, Re2 = Re1)
#           f               TxK, 1st set of factors
#           h               TxL, 2nd set of factors
#           bandwidth       (optional) scalar of Newey-West band width,  [0]
#           q               (optional) heterogenous array of vectors,
#                           expected values of R*theta: Any[alpha,delta], [0]
#           PointEstOnlyQ   (optional), bool: if true, then only the point estimates are calculated [false]
#
#  Output:  fnOutput      8x1 heterogeneous array with the following content:
#             [1] Joint         3x1, in each row WaldStat for alpha, delta, alpha-delta (if n1=n2)
#             [2] Indiv         3xmax(n1,n2), t-stat of alfa,delta and alpha-delta (if n1=n2)
#             [3] bf            n1x(1+K), [alpha beta] in Re1 = alpha + beta*f + errors
#             [4] bh            n2x(1+L), [delta gamma] in Re2 = delta+ gamma*h + errors
#             [5] EReFit        1x2 cell array, each cell is a nx2 matrix with [actual,fitted ERe]
#             [6] theta         [n1(1+K)+n2(1+L)]x1, estimated parameters, [alpha;vec(beta);delta;vec(gamma)]
#             [7] V             matrix, Cov[sqrt(T)*theta]
#
#
#
#  Notice: (a) the order of parameters in theta are  [alpha;vec(beta);delta;vec(gamma)] from
#          Re1 = alpha + beta*f + errors
#          Re2 = delta + gamma*h + errors,
#          where Re1 is an n1x1 vector, Re2 is n2x1, f is is Kx1 and h is Lx1
#
#          (b) theta2 = reshape(theta,n,1+K+1+L)
#          creates a matrix with [alpha beta delta gamma]
#
#
#
#
#  Uses:    excise4mPs, HDirProdPs, NewEst3Ps
#
#
#
#
#  Paul.Soderlind@unisg.ch, April 2015, To Julia Oct 2015
#------------------------------------------------------------------------------

  n1 = size(Re1,2)
  n2 = size(Re2,2)
  if q == 0                                    #scalar 0
    q = Any[zeros(n1,1),zeros(n2,1)]           #default values for options
  end
  #----------------------------------------------------------

  K  = size(f,2)
  L  = size(h,2)

  (Re1,Re2,f,h,) = excise4mPs(Re1,Re2,f,h)
  T              = size(Re1,1)                      #no. obs

  Ef   = mean(f,1)'
  Eh   = mean(h,1)'
  ERe1 = mean(Re1,1)'
  ERe2 = mean(Re2,1)'

  f1 = [ones(T,1) f]
  K1 = 1 + K
  h1 = [ones(T,1) h]
  L1 = 1 + L
                                         #bf = f1\Re1; epsf = Re1 - bf*f1
  bf         = f1\Re1
  epsf       = Re1 - f1*bf
  bf         = bf'                       #n1xK1
  alfaM      = bf[:,1]
  betaM      = bf[:,2:end]               #n1xK, Re1(t) = alfa * beta*f(t) + e(t)
  gf         = HDirProdPs(f1,epsf)       #moment conditions
  ERe1Hatf   = betaM*Ef                  #n1xK1 * K1x1

  bh         = h1\Re2
  epsh       = Re2 - h1*bh
  bh         = bh'                       #n2xL1
  deltaM     = bh[:,1]
  gammaM     = bh[:,2:end]               #n2xL
  gh         = HDirProdPs(h1,epsh)
  ERe2Hath   = gammaM*Eh                 #n2xL1 * L1x1

  ERe1Fit = [ERe1 ERe1Hatf]              #actual and fitted ERe
  ERe2Fit = [ERe2 ERe2Hath]
  EReFit  = Any[ERe1Fit,ERe2Fit]

  theta = [vec(bf);vec(bh)]              #parameters (combined)

  if PointEstOnlyQ
    fnOutput = Any[NaN,NaN,bf,bh,EReFit,theta,NaN,NaN]
    return fnOutput
  end
  #----------------------------------------------------------

  g     = [gf gh]                        #moment conditions (combined)
  S0    = NewEst3Ps(g,bandwidth)         #Cov[sqrt(T)*gbar]
  g  = nothing                           #can save some memory
  gf = nothing
  gh = nothing

  Sff  = f1'f1/T
  Df_1 = -kron(inv(Sff),eye(n1))         #upper left corner of inverse of Jacobian
  Shh  = h1'h1/T
  Dh_1 = -kron(inv(Shh),eye(n2))         #lower right corner of inverse of Jacobian
  D_1  = cat([1,2],Df_1,Dh_1)            #inverse of Jacobian

  V    = D_1*S0*D_1'                     #Cov[sqrt(T)*theta]
  #----------------------------------------------------------

  RM = Any[[eye(n1)      zeros(n1,n1*K) zeros(n1,n2) zeros(n1,n2*L)],        #alpha
           [zeros(n2,n1) zeros(n2,n1*K) eye(n2)      zeros(n2,n2*L)]]        #delta

  if n1 == n2                     #alpha - delta, if n1=n2
    RM = Any[RM[1],RM[2],RM[1] - RM[2]]
    q  = Any[ q[1], q[2], q[1] -  q[2]]
  end

  Indiv = fill(NaN,(length(RM),max(n1,n2)))     #t-stat
  Joint = fill(NaN,length(RM))                  #WaldStat
  for i = 1:length(RM)                          #loop over (alpha,delta,alpha-delta)
    R          = RM[i]
    q_i        = q[i]
    Rstat      = R*theta - q_i                  #eg. alpha - delta - expected value
    Lambda     = R*V*R'                         #eg. Cov[sqrt(T)*(alpha-delta)]
    stat_std   = sqrt(diag(Lambda/T))
    tstat      = Rstat./stat_std
    Joint[i]   = (T * Rstat'inv(Lambda)*Rstat)[1]     #[1] to make it a scalar
    Indiv[i,1:length(tstat)] = tstat'
  end
  #----------------------------------------------------------

             #    1      2    3 4    5      6   7
  fnOutput = Any[Joint,Indiv,bf,bh,EReFit,theta,V]

  return fnOutput

end
#------------------------------------------------------------------------------

max_stars_repo_name: PaulSoderlind/PaulSoderlindCode

max_issues_repo_name: PaulSoderlind/PaulSoderlindCode

max_forks_repo_name: PaulSoderlind/PaulSoderlindCode


text:
using StaticArrays: SVector

tobig(x) = parse(BigInt, bytes2hex(reverse(x)), base=16)

function int2bytes(x::Integer)
    hex = string(x, base=16)
    if mod(length(hex), 2) != 0
        hex = string("0", hex)
    end

    return reverse(hex2bytes(hex))
end

struct StaticBigInt{N} <: Integer
    x::SVector{N, UInt8}
end

function StaticBigInt(x::Integer)
    bytes = int2bytes(x)
    N = length(bytes)
    sv = SVector{N, UInt8}(bytes)

    return StaticBigInt(sv)
end

Base.BigInt(x::StaticBigInt) = tobig(x.x)

style(x, n) = "\33[1;$(n)m$x\33[0m"

# Some nice printing
intstyle(g::StaticBigInt) = style(" $(bytesize(g)) bytes", 90)

function Base.show(io::IO, a::StaticBigInt)
    show(io, BigInt(a))
    print(io, intstyle(a))
end

Base.display(x::StaticBigInt) = show(x)

abstract type Generator end

untype(G::Type{<:Generator}) = G.parameters[1] # A dirty hack

group(G::Type{<:Generator}) = G.parameters[1]#.parameters[1]
group(::G) where G <: Generator = G.parameters[1]

Base.show(io::IO, 𝓖::Type{Generator}) = print(io, "Generator")
Base.show(io::IO, 𝓖::Type{<:Generator}) = print(io, "Generator[$(group(𝓖))]")

order(G::Type{<:Generator}) = order(G.parameters[1])
order(::G) where G <: Generator = order(G.parameters[1])

Base.convert(::Type{G}, x) where G <: Generator = G(x)
Base.convert(::Type{G}, g::G) where G <: Generator = g # Identity

"""
Returns a value of the function
"""
function value end
Base.isless(x::G, y::G) where G <: Generator = value(x) < value(y)

import Base.==
==(x::Generator, y::Generator) = false
==(x::G, y::G) where G <: Generator = value(x) == value(y)

==(x::Generator, y) = error("Uncomparable: $(typeof(x)) with $(typeof(y))")
==(x, y::Generator) = error("Uncomparable: $(typeof(x)) with $(typeof(y))")

postfix(G::Generator) = postfix(group(G))

import Base.show
Base.show(io::IO, g::Generator) = print(io, "$(value(g)) " * style(postfix(g), 90))

function Base.show(io::IO, g::Vector{G}) where G <: Generator
    show(io, value.(g))
    Group = untype(G)
    str = " " * postfix(Group)
    print(io, style(str, 90))
end

Base.display(x::Vector{G}) where G <: Generator = show(x)

Base.show(io::IO, g::NTuple{N, <:Generator}) where N = print(io, "$(value.(g)) " * style(postfix(g[1]), 90))

# A more generic function can be constructed if necessary
Base.broadcasted(f::Function, x::Generator, y::AbstractVector{<:Integer}) = f.((x for i in 1:length(y)), y)

Base.broadcasted(::typeof(*), x::G, y::Vector{G}) where G <: Generator = (x for i in 1:length(y)) .* y
Base.broadcasted(::typeof(*), x::Vector{G}, y::G) where G <: Generator =  x .* (y for i in 1:length(x))

############################# PRIME GROUP #################################

struct PrimeGroup{N} ### Type parameter is essential to ensure it to be bitstype
    x::StaticBigInt{N}
end

PrimeGroup(x::Integer) = PrimeGroup(StaticBigInt(x))
PrimeGroup(; q) = PrimeGroup(2*q + 1)

Base.getindex(::Type{Generator}, 𝓖::PrimeGroup) = PrimeGenerator{𝓖} # This is somewhat like a trait

const 𝐙 = PrimeGroup
import Base./
/(::Type{𝐙}, p::Integer) = PrimeGroup(p)

modulus(G::PrimeGroup) = BigInt(G.x)
order(G::PrimeGroup) = (modulus(G) - 1) ÷ 2

# Z/(p) notation seems more reasonable for printing this. Seems more appropriate to be left for PrimeGroup
postfix(G::PrimeGroup) = "mod $(modulus(G)) (q = $(order(G)))"

function trimnumber(x::String)
    if length(x) < 30
        return x
    else
        return x[1:10] * "..." * x[end-10:end]
    end
end

trimnumber(x::Integer)= trimnumber(string(x))

groupstr(m) = "𝐙/($(trimnumber(m)))"

Base.show(io::IO, x::PrimeGroup) = print(io, groupstr(modulus(x)))

############################ PRIME GENERATOR ##################################

struct PrimeGenerator{G} <: Generator #{G}
    g::BigInt
end

PrimeGenerator(x::Integer, p::Integer) = PrimeGenerator{PrimeGroup(p)}(x)

modulus(g::PrimeGenerator) = modulus(group(g)) # A method which one could add.
value(g::PrimeGenerator) = g.g

validate(g::PrimeGenerator) = value(g) != 1 && value(g^order(g)) == 1

### NEW METHOD
Base.show(io::IO, 𝓖::Type{PrimeGenerator}) = print(io, "PrimeGenerator")  # Need to remember

Base.convert(::Type{BigInt}, x::PrimeGenerator) = value(x)

import Base.*
*(x::PrimeGenerator{G}, y::PrimeGenerator{G}) where G = PrimeGenerator{G}(mod(value(x) * value(y), modulus(G)))

import Base.^
^(x::PrimeGenerator{G}, n::Integer) where G = PrimeGenerator{G}(powermod(x.g, n, modulus(G)))

modinv(s, q) = gcdx(s, q)[2]

import Base.inv
inv(x::PrimeGenerator{G}) where G = PrimeGenerator{G}(modinv(value(x), modulus(G)))

import Base./
/(x::PrimeGenerator, y::PrimeGenerator) = x * inv(y)

### Need to add also elgamal encrytption primitives as follows:

struct Enc{T<:Generator}
    pk::T
    g::T
end

### Encrytion as we see does eactually
#(enc::Enc{T})(m::T, r::Integer) where T <: Generator = (m*enc.pk^r, enc.g^r) ### Message first?
#(enc::Enc)(r::Integer) = (enc.pk^r, enc.g^r)

(enc::Enc{T})(m::T, r::Integer) where T <: Generator = (enc.g^r, m*enc.pk^r) ### Message first?
(enc::Enc)(r::Integer) = (enc.g^r, enc.pk^r)

a(x::Tuple{T, T}) where T <: Generator = x[1]
b(x::Tuple{T, T}) where T <: Generator = x[2]

*(x::Tuple{G, G}, y::Tuple{G, G}) where G <: Generator = (a(x)*a(y), b(x)*b(y))

(enc::Enc)(e::Tuple{G, G}, r::Integer) where G <: Generator = e * enc(r)

struct ElGamal{G <: Generator} <: AbstractVector{G}
    a::Vector{G}
    b::Vector{G}

    function ElGamal{G}(a::Vector{G}, b::Vector{G}) where {G <: Generator}
        @assert length(a) == length(b)
        return new(a, b)
    end
end

ElGamal(a::Vector{G}, b::Vector{G}) where G <: Generator = ElGamal{G}(a, b)

ElGamal(e::Vector{Tuple{G, G}}) where G <: Generator = ElGamal([a(i) for i in e], [b(i) for i in e])

function ElGamal{G}(a::Vector{T}, b::Vector{T}) where {T, G<:Generator}
    a′ = convert(Vector{G}, a)
    b′ = convert(Vector{G}, b)

    return ElGamal{G}(a′, b′)
end

a(e::ElGamal) = e.a
b(e::ElGamal) = e.b

Base.getindex(e::ElGamal, i::Int) = (a(e)[i], b(e)[i])
Base.length(e::ElGamal) = length(a(e))
Base.size(e::ElGamal) = size(a(e))

### NEW METHOD
function Base.setindex!(e::ElGamal{G}, val::Tuple{G, G}, i) where G <: Generator
    a(e)[i] = a(val)
    b(e)[i] = b(val)
end

function *(x::ElGamal{G}, y::ElGamal{G}) where G

    @assert length(x) == length(y)

    a′ = a(x) .* a(y)
    b′ = b(x) .* b(y)

    return ElGamal(a′, b′)
end

function *(x::ElGamal{G}, y::Tuple{G, G}) where G

    a′ = a(x) .* a(y)
    b′ = b(x) .* b(y)

    return ElGamal(a′, b′)
end

*(x::Tuple{G, G}, y::ElGamal{G}) where G = y * x

(enc::Enc)(e::ElGamal, r::Integer) = enc(r) * e

function (enc::Enc{G})(m::Vector{G}, r::AbstractVector{<:Integer}) where G <: Generator

    a′ = enc.g .^ r
    b′ = m .* (enc.pk .^ r)

    return ElGamal(a′, b′)
end

function (enc::Enc{G})(𝐞::ElGamal{G}, 𝐫::AbstractVector{<:Integer}) where G <: Generator

    ### Need to improve this
    r = ElGamal(enc.(𝐫))
    𝐞′ = 𝐞 * r

    return 𝐞′
end

struct Dec
    sk::Integer
end

#(dec::Dec)(e::Tuple{G, G}) where G = a(e) * b(e)^(-dec.sk) # Only operation for which it is desirable to store encryptions as a tuple.

(dec::Dec)(e::Tuple{G, G}) where G = b(e) * a(e)^(-dec.sk) # Only operation for which it is desirable to store encryptions as a tuple.

(dec::Dec)(e::ElGamal) = [dec(ei) for ei in e]

Base.isless(x::Tuple{G, G}, y::Tuple{G, G}) where G <: Generator = x[1] == y[1] ? x[2] < y[2] : x[1] < y[1]


text:
mutable struct UnivariateOptimizationResults{Tb,Tt,Tf, Tx,M,O<:UnivariateOptimizer} <: OptimizationResults
    method::O
    initial_lower::Tb
    initial_upper::Tb
    minimizer::Tx
    minimum::Tf
    iterations::Int
    iteration_converged::Bool
    converged::Bool
    rel_tol::Tt
    abs_tol::Tt
    trace::OptimizationTrace{M}
    f_calls::Int
end


text:
# precise values were computed with Maple 2018.2.1
@testset "Smallest parameter is equal to 1" begin
  a = Float64[1, 10, 100, 1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_meixner_ortho_poly(i, one(Float64), log(x)) ≈ -log(z[u])
  end
end

@testset "Smallest parameter is equal to 10" begin
  a = Float64[10, 100, 1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F = [-.52628324337533965e-3 -.91738185249562661e-3 -.99103028876012253e-3 -.99905081268556093e-3;
       -2.0317801262571448 -3.6572455938562417 -4.0102491632437109 -4.0501571465136512;
       -3.3504104733399496 -6.1641594451500663 -6.8430557799176588 -6.9224879687084634;
       -6.9013533188221590 -13.505607286760857 -15.749584728918338 -16.058540058411344;
       -8.9931244812245734 -18.252192303278967 -22.280051500153905 -22.945569932182599;
       -13.364668417083967 -28.631101920969909 -40.723125251220806 -45.230948827869905]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_meixner_ortho_poly(i, Float64(10), log(x)) ≈ F[u, v]
  end
end

@testset "Smallest parameter is equal to 100" begin
  a = Float64[100, 1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F = [-.50248118964360310e-2 -.90986888595390629e-2 -.99014706210901467e-2;
       -19.349403266385754 -36.223694926634441 -40.058734895006031;
       -31.843688770301594 -60.988773615097764 -68.343834334910785;
       -65.106469020297578 -133.11309240174212 -157.16644812648946;
       -84.255629726916470 -179.31620948692071 -222.11471731657940;
       -120.73661754514218 -277.84032552047403 -403.55806307301357]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_meixner_ortho_poly(i, Float64(100), log(x)) ≈ F[u, v]
  end
end

@testset "Smallest parameter is equal to 1_000" begin
  a = Float64[1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F = [-.50021886485016083e-1 -.90912435206967968e-1;
       -192.57640014113831 -361.89231624299969;
       -316.86775046147991 -609.24410302537978;
       -647.41688830508753 -1329.2305629085790;
       -837.33067117580264 -1790.0472246195170;
       -1196.8270424842711 -2770.4365897006691]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    @test SimpleBirthDeathProcess.log_meixner_ortho_poly(i, Float64(1_000), log(x)) ≈ F[u, v]
  end
end

@testset "k in [1, 0, -1]: Smallest parameter is equal to 2" begin
  a = Float64[2, 10, 100, 1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    r1, r2, r3 = SimpleBirthDeathProcess.log_meixner_ortho_poly_joint(i, Float64(2), log(x))
    @test r1 ≈ log(i + 2 * z[u] - 1) - 2 * log(z[u]) - log1p(i)
    @test r2 ≈ log(i + z[u] - 1) - log(z[u]) - log(i)
    @test r3 === zero(Float64)
  end
end

@testset "k in [1, 0, -1]: Smallest parameter is equal to 10" begin
  a = Float64[10, 100, 1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F1 = [-.52628324337533965e-3 -.91738185249562661e-3 -.99103028876012253e-3 -.99905081268556093e-3;
        -2.0317801262571448 -3.6572455938562417 -4.0102491632437109 -4.0501571465136512;
        -3.3504104733399496 -6.1641594451500663 -6.8430557799176588 -6.9224879687084634;
        -6.9013533188221590 -13.505607286760857 -15.749584728918338 -16.058540058411344;
        -8.9931244812245734 -18.252192303278967 -22.280051500153905 -22.945569932182599;
        -13.364668417083967 -28.631101920969909 -40.723125251220806 -45.230948827869905]

  F2 = [-.44997154621306883e-3 -.82495557250215841e-3 -.89191929546486659e-3 -.89914565055050351e-3;
        -1.7268423647221549 -3.2876137827946026 -3.6091750283602051 -3.6451409270498669;
        -2.8344139388261806 -5.5392105465265657 -6.1586334908790476 -6.2302379617871075;
        -5.7391285605322099 -12.112750882906229 -14.173760567088424 -14.452676449894550;
        -7.3620201029852153 -16.326937625328312 -20.048881722140631 -20.650973873851784;
        -10.125594848132590 -25.159927670658581 -36.558405822355922 -40.704713967442946]

  F3 = [-.37644616098559351e-3 -.73267076952353524e-3 -.79281005797932914e-3 -.79924050637141279e-3;
        -1.4345092663145545 -2.9187743959822365 -3.2081117219182366 -3.2401248196531340;
        -2.3417773974599485 -4.9159716647216521 -5.4742368237939885 -5.5379882234321887;
        -4.6484173208971883 -10.728016496952154 -12.598124000963975 -12.846814969489424;
        -5.8601107881637621 -14.420485859284776 -17.818384875997870 -18.356386451070455;
        -7.6035926506422622 -21.822978690338372 -32.411557448038426 -36.179147716752838]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    r1, r2, r3 = SimpleBirthDeathProcess.log_meixner_ortho_poly_joint(i, Float64(10), log(x))
    @test r1 ≈ F1[u, v]
    @test r2 ≈ F2[u, v]
    @test r3 ≈ F3[u, v]
  end
end

@testset "k in [1, 0, -1]: Smallest parameter is equal to 100" begin
  a = Float64[100, 1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F1 = [-.50248118964360310e-2 -.90986888595390629e-2 -.99014706210901467e-2;
        -19.349403266385754 -36.223694926634441 -40.058734895006031;
        -31.843688770301594 -60.988773615097764 -68.343834334910785;
        -65.106469020297578 -133.11309240174212 -157.16644812648946;
        -84.255629726916470 -179.31620948692071 -222.11471731657940;
        -120.73661754514218 -277.84032552047403 -403.55806307301357]

  F2 = [-.49496903336308133e-2 -.90068897338010612e-2 -.98024463036859674e-2;
        -19.049868084336226 -35.856894960464313 -39.658088245693884;
        -31.337725781208335 -60.369013437295331 -67.660255610319402;
        -63.975125304549095 -131.73475269960223 -155.59375275504627;
        -82.680492341128495 -177.41449504455335 -219.88985501919318;
        -117.81301911039979 -274.45170026360770 -399.42270183276153]

  F3 = [-.48748212962385064e-2 -.89151055509698546e-2 -.97034221785626762e-2;
        -18.751468651711550 -35.490177672488245 -39.257442779810132;
        -30.833857028198276 -59.749429848752658 -66.976679679969060;
        -62.850044157717020 -130.35722132504727 -154.02107758909523;
        -81.116336230254627 -175.51460195041096 -217.66506396756262;
        -114.93595962960294 -271.07491295298205 -395.28907222661447]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    r1, r2, r3 = SimpleBirthDeathProcess.log_meixner_ortho_poly_joint(i, Float64(100), log(x))
    @test r1 ≈ F1[u, v]
    @test r2 ≈ F2[u, v]
    @test r3 ≈ F3[u, v]
  end
end

@testset "k in [1, 0, -1]: Smallest parameter is equal to 1_000" begin
  a = Float64[1_000, 10_000]
  z = Float64[1.0001, 1.5, 2.0, 5.0, 10.0, 100.0]

  F1 = [-.50021886485016083e-1 -.90912435206967968e-1;
        -192.57640014113831 -361.89231624299969;
        -316.86775046147991 -609.24410302537978;
        -647.41688830508753 -1329.2305629085790;
        -837.33067117580264 -1790.0472246195170;
        -1196.8270424842711 -2770.4365897006691]

  F2 = [-.49946878041294134e-1 -.90820697729011826e-1;
        -192.27737387819252 -361.52579376188207;
        -316.36272606764623 -608.62485012554081;
        -646.28835292904761 -1327.8536325701771;
        -835.76046016206889 -1788.1477816590733;
        -1193.9243639196877 -2767.0554727792707]

  F3 = [-.49871894621978197e-1 -.90728961752889464e-1;
        -191.97846010652304 -361.15927958063496;
        -315.85790900945579 -608.00561493593598;
        -645.16043639731818 -1326.4767830232228;
        -834.19133195599179 -1786.2485202922418;
        -1191.0262000516550 -2763.6755269221370]

  @testset for (v, i) = enumerate(a), (u, x) = enumerate(z)
    r1, r2, r3 = SimpleBirthDeathProcess.log_meixner_ortho_poly_joint(i, Float64(1_000), log(x))
    @test r1 ≈ F1[u, v]
    @test r2 ≈ F2[u, v]
    @test r3 ≈ F3[u, v]
  end
end


text:
"""
	form_group!(G3D::AbstractArray{Float64, 3},
			img::AbstractArray{Float64, 2},
			matchTable::Array{CartesianIndex{2}, 3},
			refIndex::Array{CartesianIndex{2},2},
			patchSize::CartesianIndex{2},
			position::CartesianIndex{2},
			group_transform!::Function)

Forward BM3DDenoise groupings
"""
function form_group!(
	G3D::AbstractArray{Float64,3},
	img::AbstractArray{Float64,2},
	matchTable::Array{CartesianIndex{2},3},
	refIndex::Array{CartesianIndex{2},2},
	patchSize::CartesianIndex{2},
	position::CartesianIndex{2},
	group_transform!::Function,
)

	# Form table of 3D groups
	image_to_group!(img, G3D, matchTable, refIndex, patchSize, position)

	# Apply group transform
	group_transform!(G3D)
end

"""
	invert_group!(img::AbstractArray{Float64, 2},
				G3D::AbstractArray{Float64, 3},
				matchTable::Array{CartesianIndex{2},3},
				refIndex::Array{CartesianIndex{2},2},
				patchSize::CartesianIndex{2},
				position::CartesianIndex{2},
				group_itransform!::Function,
				imgLockPool::Array{ReentrantLock})

Inverse BM3DDenoise groupings
"""
function invert_group!(
	img::AbstractArray{Float64,2},
	G3D::AbstractArray{Float64,3},
	matchTable::Array{CartesianIndex{2},3},
	refIndex::Array{CartesianIndex{2},2},
	patchSize::CartesianIndex{2},
	position::CartesianIndex{2},
	group_itransform!::Function,
	imgLockPool::Array{ReentrantLock},
)

	# Apply inverse group transform
	group_itransform!(G3D)

	group_to_image!(img, G3D, matchTable, refIndex, patchSize, position, imgLockPool)

end

"""
	group_to_image!(img::AbstractArray{Float64, 2},
			G3D::AbstractArray{Float64, 3},
			matchTable::Array{CartesianIndex{2},3},
			refIndex::Array{CartesianIndex{2},2},
			patchSize::CartesianIndex{2},
			position::CartesianIndex{2},
			imgLockPool::Array{ReentrantLock})

group to image
"""
function group_to_image!(
	img::AbstractArray{Float64,2},
	G3D::AbstractArray{Float64,3},
	matchTable::Array{CartesianIndex{2},3},
	refIndex::Array{CartesianIndex{2},2},
	patchSize::CartesianIndex{2},
	position::CartesianIndex{2},
	imgLockPool::Array{ReentrantLock},
)

	nMatch = size(matchTable, 3)

	lock(imgLockPool[position]) do
		@views img[refIndex[position]:(patchSize - CartesianIndex(
			1,
			1,
		) + refIndex[position])] .+= G3D[1, :, :]
	end

	@inbounds @views for k in 1:nMatch
		position2 = matchTable[position, k]
		lock(imgLockPool[position2]) do
			img[refIndex[position2]:(patchSize - CartesianIndex(
				1,
				1,
			) + refIndex[position2])] .+= G3D[k + 1, :, :]
		end
	end
end

"""
	group_to_image!(img::AbstractArray{Float64, 2},
			W::Float64,
			matchTable::Array{CartesianIndex{2},3},
			refIndex::Array{CartesianIndex{2},2},
			patchSize::CartesianIndex{2},
			position::CartesianIndex{2},
			imgLockPool::Array{ReentrantLock})
"""
function group_to_image!(
	img::AbstractArray{Float64,2},
	W::Float64,
	matchTable::Array{CartesianIndex{2},3},
	refIndex::Array{CartesianIndex{2},2},
	patchSize::CartesianIndex{2},
	position::CartesianIndex{2},
	imgLockPool::Array{ReentrantLock},
)

	nMatch = size(matchTable, 3)

	lock(imgLockPool[position]) do
		@views img[refIndex[position]:(patchSize - CartesianIndex(
			1,
			1,
		) + refIndex[position])] .+= W
	end

	@inbounds @views for k in 1:nMatch
		position2 = matchTable[position, k]
		lock(imgLockPool[position2]) do
			img[refIndex[position2]:(patchSize - CartesianIndex(
				1,
				1,
			) + refIndex[position2])] .+= W
		end
	end
end

"""
	image_to_group!(img::AbstractArray{Float64, 2},
			G3D::AbstractArray{Float64, 3},
			matchTable::Array{CartesianIndex{2},3},
			refIndex::Array{CartesianIndex{2},2},
			patchSize::CartesianIndex{2},
			position::CartesianIndex{2})

"""
function image_to_group!(
	img::AbstractArray{Float64,2},
	G3D::AbstractArray{Float64,3},
	matchTable::Array{CartesianIndex{2},3},
	refIndex::Array{CartesianIndex{2},2},
	patchSize::CartesianIndex{2},
	position::CartesianIndex{2},
)

	nMatch = size(matchTable, 3)

	@views G3D[1, :, :] .=
		img[refIndex[position]:(patchSize - CartesianIndex(
			1,
			1,
		) + refIndex[position])]

	@views @inbounds for k in 1:nMatch
		position2 = matchTable[position, k]
		G3D[k + 1, :, :] .=
			img[refIndex[position2]:(patchSize - CartesianIndex(
				1,
				1,
			) + refIndex[position2])]
	end
end


text:
module psfsize

using ..params
using ..psf
using FFTW
export calcsize

function calcsize(p::ParameterSet, obj::Space)
    # we set an arbitrarily high length to test across
    testrange = collect((0:1:(p.sim.subvpix*20)) .* p.sim.subpixelpitch)
    x3max = maximum(abs.(obj.z))

    if p.opt.a0 > zero(Float64)
        steps = 10
        stepz = p.opt.a0/steps
        lineproj = zeros(ComplexF64, length(testrange), obj.zlen)
        linemag = zeros(length(testrange), obj.zlen)
        Hline = fresnelH(lineproj[:,1], p, stepz)
        plan = plan_fft!(lineproj[:,1])
        for i in 1:obj.zlen
            lineproj[:,i] = psfline(testrange, obj.z[i], p)
            plan*lineproj[:,i]
            lineproj[:,i] = lineproj[:,i].*(Hline.^steps)
            plan\lineproj[:,i]
            linemag[:,i] = abs2.(lineproj[:,i]) ./ maximum(abs2.(lineproj[:,i]))
        end
        imgspace = makeimgspace(dropdims(reduce(max, linemag, dims=2); dims=2), p)
    else
        lineproj = psfline(testrange, x3max, p)
        linemag =  abs2.(lineproj) ./ maximum(abs2.(lineproj))
        imgspace = makeimgspace(linemag, p)
    end
    return imgspace
end

function psfline(x₁line::Vector{Float64}, x₃::Float64, p::ParameterSet)

    l²norm² = abs.(x₁line) ./ p.opt.M
    v = l²norm² .* (p.con.k*sin(p.con.alpha))
    u = 4*p.con.k * x₃ * (sin(p.con.alpha/2)^2)
    Kₒ = p.opt.M/(p.opt.fobj*p.opt.lambda)^2 * exp(-im*u/(4(sin(p.con.alpha/2)^2)))

    psfline = Array{ComplexF64,1}(undef,length(v))

    Threads.@threads for i in 1:length(v)
        @inbounds psfline[i] = integratePSF(v[i], u, p.opt.a0, p.con.alpha) * Kₒ
    end

    return psfline
end

function makeimgspace(psfline::Vector{Float64}, p::ParameterSet)
    outArea = psfline .< 0.1
    if sum(outArea) == 0
        error("Estimated PSF size exceeds the limit")
    end
    sizeref = cld(findfirst(outArea)[1],p.sim.subvpix)
    halfwidth = max( p.sim.subvpix*(sizeref + 4), 2*p.sim.subvpix)

    "Create X-Y image space based on halfwidth."
    x =(-halfwidth:1:halfwidth) * p.sim.subpixelpitch
    y =(-halfwidth:1:halfwidth) * p.sim.subpixelpitch
    imgspace = Space(x,y)

    return imgspace
end

end # module


text:
using DataDeps

register(DataDep(
  "insurance-company-benchmark",
  "https://archive.ics.uci.edu/ml/datasets/Insurance+Company+Benchmark+%28COIL+2000%29",
  "https://archive.ics.uci.edu/ml/machine-learning-databases/tic-mld/ticdata2000.txt",
  "5d5bdca1be0b9a6c012b93a8c8104761201bd70f7e08764fa67f2cfb26243723",
  post_fetch_method=(path -> begin
    UCIData.process_dataset(path,
        target_index=86,
        feature_indices=1:85,
        delim='\t',
    )
  end),
))

max_stars_repo_name: JackDunnNZ/uci-data

max_issues_repo_name: JackDunnNZ/uci-data

max_forks_repo_name: JackDunnNZ/uci-data


text:
using AdaMeanShift
using StaticArrays
using LinearAlgebra
using Test

@testset "AdaMeanShift.jl" begin
    vx=-20:0.1:20
    vy=0:0.1:30
    M=[exp(-(x-7)^2/(2*3)-(y-12)^2/(2*2)) for x in vx, y in vy]
    posmax=Tuple(findmax(M)[2])
    P,h,w,δ = atomic_meanshift(M, @SVector([100.0,100.0]) , @SVector([50.0,50.0]), 0.0,0.0)
    for i in 1:500
        P,h,w,δ = atomic_meanshift(M, P , h, 0.0,0.0)
    end
    v=Tuple(P).-posmax
    @test sum(abs.(v))<1e-7
    @test round(Int,atomic_intensity(M, P, h).intensity) == 1263.0

    Pv=[@SVector([100.0,100.0]) for rep=1:200]
    hv=[@SVector([50.0,50.0]) for rep=1:200]
    wv=[0.0 for rep=1:200]
    XX=meanshift!(M, Pv, hv, wv,Inf,smoothing=0.0,isotropy=0.0)
    @test atomic_meanshift(zeros(0,0),[0.0,0.0],[0.0,0.0],0.0,0.0).modeval===NaN
    @test norm(Tuple(Pv[1]).-posmax)<1e-7
    XX=meanshift!(M, Pv, hv, wv,20.0,smoothing=0.0)
    @test norm(hv[1])≈20.0
    @test norm(Tuple(Pv[1]).-posmax)<1e-7
end

@testset "Helper Methods" begin
    vx=-20:0.1:20
    vy=0:0.1:30
    M=[exp(-(x-7)^2/(2*3)-(y-12)^2/(2*2)) for x in vx, y in vy]
    posmax=Tuple(findmax(M)[2])
    P,h,w,δ = atomic_meanshift(M, [100.0,100.0] , [50.0,50.0], 0.0,0.0)
    for i in 1:500
        P,h,w,δ = atomic_meanshift(M, P , h, 0.0,0.0)
    end
    v=Tuple(P).-posmax
    Pjl=[x for x in Tuple(P)]
    hjl=[x for x in Tuple(h)]
    @test sum(abs.(v))<1e-7
    @test round(Int,atomic_intensity(M, Pjl, hjl).intensity) == 1263.0
end

@testset "Non adaptive routines" begin
    vx=-20:0.1:20
    vy=0:0.1:30
    M=[exp(-(x-7)^2/(2*3)-(y-12)^2/(2*2)) for x in vx, y in vy]
    posmax=Tuple(findmax(M)[2])

    Pv=[@SVector([100.0,100.0])]
    hv=[@SVector([50.0,50.0])]
    wv=[0.0]
    XX=meanshift_nonadaptive!(M, Pv, hv, wv,smoothing=0.0)

    @test norm(Tuple(Pv[1]).-posmax)<1e-7
end

@testset "Ratio Calculation" begin
    A=[exp(-(x-4)^2 -(y-3)^2) for x in 0:0.1:8, y in 0:0.1:8]
    A[40,40]=-1.0
    @test prod(&,atomic_ratio(A,A,[40.99999999608069, 30.9999999998036],
            [13.102949390292464, 13.10294939029246]) .≈ (1.0, 0.031935091590899516))
    A.=0.0
    @test prod(&,isnan.(atomic_ratio(A,A,[40.99999999608069, 30.9999999998036],
            [13.102949390292464, 13.10294939029246])))
    A.=eps(0.0)
    @test prod(&,atomic_ratio(A,A,[40.99999999608069, 30.9999999998036],
            [13.102949390292464, 13.10294939029246]).≈(1.0,0.0))
end

@testset "Region Extraction" begin
    T=(ones(8,8),@SVector([3.0,4.0]),@SVector([2.0,2.0]))
    RES=([0.0 0.0 0.0 1.0 0.0 0.0 0.0; 0.0 0.0 1.0 1.0 1.0 0.0 0.0; 0.0 1.0 1.0 1.0 1.0 1.0 0.0; 0.0 0.0 1.0 1.0 1.0 0.0 0.0; 0.0 0.0 0.0 1.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0],[3.0, 4.0], [2.0, 2.0])
    @test prod(&,extract_circ_region(T...).==RES)
end


text:
Distributions.rand(rng::AbstractRNG, d::Dirac, n::Int) = hcat(fill(d.value, n)...)

"""
    BudgetSet{T<:Real, D<:ContinuousMultivariateSampleable} <: CenteredAmbiguitySet{T,D}

Bertsimas's uncertainty set:

```math
\\left\\{ \\mu \\; \\middle| \\begin{array}{ll}
s.t.  \\quad \\mu_i \\leq \\hat{r}_i + z_i \\Delta_i \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad \\mu_i \\geq \\hat{r}_i - z_i \\Delta_i  \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad z_i \\geq 0 \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad z_i \\leq 1 \\quad \\forall i = 1:\\mathcal{N} \\\\
\\quad \\quad \\sum_{i}^{\\mathcal{N}} z_i \\leq \\Gamma \\\\
\\end{array}
\\right\\} \\\\
```

Attributes:
- `d::Sampleable{Multivariate, Continous}`: The parent distribution with an uncertain mean
- `Δ::Array{Float64,1}`: Uncertainty around mean. (default: std(d) / 5)
- `Γ::Float64`: Number of assets in worst case. (default: 0.1 * length(d))

References:
- Bertsimas, D. e Sim, M. (2004). The price of robustness. Operations research, 52(1):35–53.

"""
struct BudgetSet{T<:Real, D<:Sampleable} <: CenteredAmbiguitySet{T,D}
    d::D
    Δ::Vector{T}
    Γ::T
    # Inner constructor for validating arguments
    function BudgetSet{T, D}(d::D, Δ::Vector{T}, Γ::T) where {T<:Real, D<:Sampleable}
        length(d) == length(Δ) || throw(ArgumentError(
            "Distribution ($(length(d))) and Δ ($(length(Δ))) are not the same length"
        ))
        all(>=(0), Δ) || throw(ArgumentError("All uncertainty deltas must be >= 0"))
        Γ >= 0 || throw(ArgumentError("Budget must be >= 0"))
        Γ <= length(d) || @warn "Budget should not exceed the distribution length"
        return new{T, D}(d, Δ, Γ)
    end
end

distribution(s::BudgetSet) = s.d

# Default outer constructor
BudgetSet(d::D, Δ::Vector{T}, Γ::T) where {T<:Real, D<:Sampleable} = BudgetSet{T, D}(d, Δ, Γ)

"""
    calculate_measure!(measure::ExpectedReturn{BudgetSet,WorstCase}, w)

Returns worst case return (WCR) in BudgetSet's uncertainty set ([`BudgetSet`](@ref)).

For further reading:
- BudgetSet' paper: BudgetSet, D. e Sim, M. (2004). The price of robustness. Operations research, 52(1):35–53.
- Original implementation from: https://github.com/andrewrosemberg/PortfolioOpt.jl.

WCR is defined by the following primal problem:

```math
\\begin{aligned}
\\min_{\\mu, z} \\quad & \\mu ' w \\\\
s.t. \\quad & \\mu_i \\leq \\hat{r}_i + z_i \\Delta_i \\quad \\forall i = 1:\\mathcal{N} \\quad &: \\pi^-_i \\\\
& \\mu_i \\geq \\hat{r}_i - z_i \\Delta_i  \\quad \\forall i = 1:\\mathcal{N} \\quad &: \\pi^+_i \\\\
& z_i \\geq 0 \\quad \\forall i = 1:\\mathcal{N} \\\\
& z_i \\leq 1 \\quad \\forall i = 1:\\mathcal{N} \\quad &: \\theta_i \\\\
& \\sum_{i}^{\\mathcal{N}} z_i \\leq \\Gamma \\quad : \\lambda \\\\
\\end{aligned}
```

Which is equivalent to the following dual problem:

```math
\\begin{aligned}
\\max_{\\lambda, \\pi^-, \\pi^+, \\theta} \\quad  \\sum_{i}^{\\mathcal{N}} (\\hat{r}_i (\\pi^+_i - \\pi^-_i) - \\theta_i ) - \\Gamma \\lambda\\\\
s.t. \\quad &  w_i = \\pi^+_i - \\pi^-_i  \\quad \\forall i = 1:\\mathcal{N} \\\\
&  \\Delta_i (\\pi^+_i + \\pi^-_i) - \\theta_i \\leq \\lambda \\quad \\forall i = 1:\\mathcal{N} \\\\
& \\lambda \\geq 0 , \\; \\pi^- \\geq 0 , \\; \\pi^+ \\geq 0 , \\; \\theta \\geq 0 \\\\
\\end{aligned}
```

To avoid solving an optimization problem we enforece the dual constraints in
the upper level problem and return the objective expression (a lower bound of the optimum).

Arguments:
 - `w`: portfolio optimization investment variable ("weights").
 - `s::BudgetSet`: Struct containing atributes of BudgetSet's uncertainty set.
"""
function calculate_measure!(measure::ExpectedReturn{S,WorstCase}, w)  where {S<:BudgetSet}
    model = owner_model(w)
    s = ambiguityset(measure)

    # parameters
    means = Vector(Distributions.mean(s.d))
    n = length(s)
    Δ = s.Δ
    Γ = s.Γ
    # dual variables
    @variable(model, λ >= 0)
    @variable(model, π_neg[i=1:n] >= 0)
    @variable(model, π_pos[i=1:n] >= 0)
    @variable(model, θ[i=1:n] >= 0)
    # constraints: from duality
    @constraints(
        model,
        begin
            constrain_dual1[i=1:n], w[i] == π_pos[i] - π_neg[i]
        end
    )
    @constraints(
        model,
        begin
            constrain_dual2[i=1:n], Δ[i] * (π_pos[i] + π_neg[i]) - θ[i] <= λ
        end
    )

    return return dot(means, w) - sum(θ) - λ * Γ
end

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

"""
    EllipticalSet <: CenteredAmbiguitySet{T,D}

```math
\\left\\{ \\mu \\; \\middle| \\begin{array}{ll}
s.t.  \\quad \\sqrt{(\\hat{r} - \\mu) ' \\Sigma^{-1} (\\hat{r} - \\mu)} \\leq \\delta \\\\
\\end{array}
\\right\\} \\\\
```

Atributes:
- `d::Sampleable{Multivariate, Continous}`: The parent distribution with an uncertain mean
- `Δ::Array{Float64,1}`: Uniform uncertainty around mean. (default: 0.025)

References:
- Ben-Tal, A. e Nemirovski, A. (2000). Robust solutions of linear programming problems contaminated with uncertain data. Mathematical programming, 88(3):411–424.

For more information on how BenTal uncertainty sets are used for RO, please review
the PortfolioOptimization.jl [docs](https://invenia.pages.invenia.ca/PortfolioOptimization.jl/).
"""
struct EllipticalSet{T<:Real, D<:ContinuousMultivariateSampleable} <: CenteredAmbiguitySet{T,D}
    d::D
    Δ::T

    # Inner constructor for validating arguments
    function EllipticalSet{T, D}(d::D, Δ::T) where {T<:Real, D<:ContinuousMultivariateSampleable}
        Δ >= 0 || throw(ArgumentError("Uncertainty delta must be >= 0"))
        return new{T, D}(d, Δ)
    end
end

# Default outer constructor
EllipticalSet(d::D, Δ::T) where {T<:Real, D<:ContinuousMultivariateSampleable} = EllipticalSet{T, D}(d, Δ)

distribution(s::EllipticalSet) = s.d

"""
    calculate_measure!(measure::ExpectedReturn{EllipticalSet,WorstCase}, w)

Returns worst case return (WCR) in EllipticalSet's uncertainty set ([`EllipticalSet`](@ref)).

WCR is defined by the following primal problem:

```math
\\begin{aligned}
\\min_{\\mu} \\quad & \\mu ' w \\\\
s.t. \\quad & ||Σ^{-\\frac{1}{2}} (\\mu - \\hat{r}) || \\leq \\delta \\quad &: \\theta \\\\
\\end{aligned}
```

Which is equivalent to the following dual problem:

```math
\\begin{aligned}
\\max_{\\theta} \\quad &  w ' \\hat{r} - \\theta \\delta \\\\
s.t. \\quad & ||Σ^{\\frac{1}{2}}  w || \\leq \\theta \\\\
\\end{aligned}
```

To avoid solving an optimization problem we enforece the dual constraints in
the upper level problem and return the objective expression (a lower bound of the optimum).

Arguments:
 - `w`: portfolio optimization investment variable ("weights").
 - `s:: EllipticalSet `: Struct containing atributes of EllipticalSet's uncertainty set.
"""
function calculate_measure!(measure::ExpectedReturn{S,WorstCase}, w)   where {S<:EllipticalSet}
    model = owner_model(w)
    s = ambiguityset(measure)

    # parameters
    means = Vector(Distributions.mean(s.d))
    Σ = PDMat(Symmetric(cov(s.d)))
    sqrt_Σ = collect(Σ.chol.U)

    Δ = s.Δ
    # dual variables
    @variable(model, θ)
    # constraints: from duality
    @constraint(model, [θ; sqrt_Σ * w] in JuMP.SecondOrderCone())

    return dot(w, means) - θ * Δ
end

max_stars_repo_name: andrewrosemberg/PortfolioOpt

max_issues_repo_name: andrewrosemberg/PortfolioOpt

max_forks_repo_name: andrewrosemberg/PortfolioOpt


text:
"""
    AbstractRule

A supertype for all decision rules.
"""
abstract type AbstractRule end

"""
    apply(r::AbstractRule, dp)

Determine the probability an agent will switch its strategy based on the rule and the difference in
payoffs of the two strategies.
"""
apply(r::AbstractRule, dp) = error("method must be overloaded")

"""
    AbstractScheme

A supertype for all schemes, mechanisms for agents to choose their next strategy.
"""
abstract type AbstractScheme end

"""
    AbstractArena{Graph, Scheme}

A supertype for all game-playing arenas.
"""
abstract type AbstractArena{Graph <: SimpleGraph, Scheme <: AbstractScheme} end

"""
    game(arena)

The game the agents in the arena play
"""
game(a::AbstractArena) = a.game

"""
    graph(arena)

The graph on which the agents play their games
"""
graph(a::AbstractArena) = a.graph

"""
    scheme(arena)

The scheme the agents use to determine their next strategy
"""
scheme(a::AbstractArena) = a.scheme

"""
    length(arena)

The number of agents in the arena
"""
Base.length(a::AbstractArena) = nv(graph(a))

"""
    edges(arena)

The edges connecting agents in the arena
"""
LightGraphs.edges(a::AbstractArena) = edges(graph(a))

"""
    neighbors(arena, agent)

A vector of the agent's neighbors
"""
LightGraphs.neighbors(a::AbstractArena, i) = neighbors(graph(a), i)

"""
    payoffs(arena, strategies)

Determine the payoff of each agents' possible strategies given the strategies of all of the other
agents.
"""
function payoffs(a::AbstractArena, ss::AbstractVector{Int})
    if length(ss) != length(a)
        throw(ArgumentError("number of strategies should be the same as the number of agents"))
    elseif !all(1 .<= ss .<= 2)
        throw(ArgumentError("all strategies must be either 1 or 2"))
    end
    ps = zeros(Float64, 2, length(ss))
    for i in 1:length(ss), j in neighbors(a, i), s in 1:2
        ps[s, i] += play(game(a), s, ss[j])
    end
    ps
end

"""
    play(arena, strategies)

Play one round of the game in the arena, returning the agents' next strategies.
"""
play(a::AbstractArena, ss::AbstractVector{Int}) = decide(scheme(a), a, ss, payoffs(a, ss))

"""
    play(arena; rounds=1, replicates=1)

Starting from `replicates` many random initial strategies, play the game in the arena `rounds` many
times. The result is a 3-D array of the agents strategies with the first two dimensions of size
`rounds` and `replicates`, respectively.
"""
function play(a::AbstractArena; rounds::Int = 1, replicates::Int = 1)
    if rounds < 1
        throw(DomainError(rounds, "must be at least 1"))
    end
    if replicates < 1
        throw(DomainError(replicates, "must be at least 1"))
    end

    series = Array{Int}(undef, rounds, replicates, length(a))
    @views for i in 1:replicates
        series[1,i,:] = rand(1:2, length(a))
        for j in 2:rounds
            series[j,i,:] = play(a, series[j-1,i,:])
        end
    end
    series
end

"""

    decide(scheme, arena, ss::AbstractVector{Int}, ps::AbstractArray{Float64,2})

Use the scheme to decide what each agent's strategy will be in the next timestep. Here `ss` is a
vector of the agents' current strategies (one for each agent) and `ps` is a matrix of payoffs. Each
column of `ps` should correspond to an agent, and each row is a corresponding strategy.
"""
function decide(scheme::AbstractScheme,
                arena::AbstractArena,
                ss::AbstractVector{Int},
                ps::AbstractArray{Float64,2})
    if !all(1 .<= ss .<= 2)
        throw(ArgumentError("all strategies must be either 1 or 2"))
    end
    if length(ss) != size(ps, 2)
        throw(ArgumentError("length(ss) ≠ size(ps, 2)"))
    elseif size(ps, 1) != 2
        throw(ArgumentError("size(ps, 1) ≠ 2"))
    end
    dps = Vector{Float64}(undef, length(ss))
    for (i, s) in enumerate(ss)
        dps[i] = ps[3 - s, i] - ps[s, i]
    end
    decide(scheme, arena, ss, dps)
end

"""
    decide(scheme, arena, ss::AbstractVector{Int}, dps::AbstractVector{Float64})

Use the scheme to decide what each agent's strategy will be in the next time step. Here `ss` is a
vector of the agents' current strategies (one for each agent) and `dps` is a vector of the
difference between the payoff for the strategy the agent didn't play and the payoff they actually
received.

    decide(scheme, arena, s::Int, dp::Float64)

Use the scheme to decide what an agent's strategy will be in the next time step. Here `s` is the
agent's current strategy and `dp` is the difference between the payoff for the alternative strategy
and the actual payoff they received.
"""
function decide(scheme::AbstractScheme,
                arena::AbstractArena,
                ss::AbstractVector{Int},
                dps::AbstractVector{Float64})
    length(ss) != length(dps) && error("strategies and payoff differences must have the same length")
    map(i -> decide(scheme, arena, ss[i], dps[i]), 1:length(ss))
end


text:
module PartialRejectionSampling

# Imports

using LinearAlgebra
const LA = LinearAlgebra

using Random, Distributions

using LightGraphs
const LG = LightGraphs

using SimpleWeightedGraphs
const SWG = SimpleWeightedGraphs

using Distances

using SpecialFunctions
const SF = SpecialFunctions

# Exports

const PRS = PartialRejectionSampling
export PRS

export AbstractWindow
export AbstractSpatialWindow,
    RectangleWindow,
    SquareWindow,
    BallWindow
export AbstractDiscreteWindow,
    GraphNode

export AbstractPointProcess
export AbstractSpatialPointProcess,
    HomogeneousPoissonPointProcess,
    HardCorePointProcess,
    StraussPointProcess
export AbstractGraphPointProcess,
    Ising,
    HardCoreGraph,
    RootedSpanningForest,
    SinkFreeGraph
export PatternFreeString

export generate_sample,
    generate_sample_prs,
    generate_sample_grid_prs,
    generate_sample_dcftp,         # For Spatial point processes
    generate_sample_gibbs_perfect  # For the Ising model

# Code inclusions

include("common.jl")
include("utils.jl")

include(joinpath("spatial", "window.jl"))
include(joinpath("graph", "window.jl"))

include("grid_prs.jl")

## Spatial point processes
include(joinpath("spatial", "dominated_cftp.jl"))
include(joinpath("spatial", "poisson.jl"))
include(joinpath("spatial", "hard_core.jl"))
include(joinpath("spatial", "strauss.jl"))

## Graph point processes
include(joinpath("graph", "ising.jl"))
include(joinpath("graph", "hard_core.jl"))
include(joinpath("graph", "rooted_spanning_forest.jl"))
include(joinpath("graph", "sink_free_graph.jl"))

# Misc
include(joinpath("misc", "pattern_free_string.jl"))

end


text:
@testset "373.find-k-pairs-with-smallest-sums.jl" begin
    @test k_smallest_pairs([1, 7, 11], [2, 4, 6], 3) == [(1, 2), (1, 4), (1, 6)]
    @test k_smallest_pairs([1,1,2], [1, 2, 3], 2) == [(1, 1), (1, 1)]
    @test k_smallest_pairs([1,2], [3], 2) == [(1, 3), (2, 3)]
end


text:
@testset "StrategyUtils" begin

    @testset "is_in_line_of_sight Tests" begin
        # myPlayer at [3,3]
        # obstacle (opponent) at [6,3] and size 0.3
        # ball at [8,3]
        ws = sb.WorldState([sb.Robot([3.0, 3.0], 0.0)], [sb.Robot([6.0, 3.0], 0.0)], [sb.Ball([8.0, 3.0])])
        @test sb.is_in_line_of_sight(ws, ws.myPlayers[1], ws.obstacles[1]) == false
        ws = sb.WorldState([sb.Robot([3.0, 3.0], 0.0)], [sb.Robot([6.0, 3.0], 0.0)], [sb.Ball([8.0, 3.0])])
        @test sb.is_in_line_of_sight(ws, ws.myPlayers[1], ws.obstacles[1], excludes=[ws.opponents[1]]) == false
        @test sb.is_in_line_of_sight(ws, ws.myPlayers[1], ws.obstacles[1], excludes=[ws.opponents[1], ws.myPlayers[1]]) == true
        ws = sb.WorldState([sb.Robot([3.0, 3.0], 0.0)], [sb.Robot([6.0, 3.15], 0.0)], [sb.Ball([8.0, 3.0])])
        @test sb.is_in_line_of_sight(ws, ws.myPlayers[1], ws.obstacles[1], tol=0.05, excludes=[ws.myPlayers[1]]) == true
        ws = sb.WorldState([sb.Robot([8.0, 6.0], 0.0)], [sb.Robot([6.0, 3.0], 0.0)], [sb.Ball([8.0, 3.0])])
        @test sb.is_in_line_of_sight(ws, ws.myPlayers[1], ws.obstacles[1]) == false
        @test sb.is_in_line_of_sight(ws, ws.myPlayers[1], ws.obstacles[1], excludes=[ws.myPlayers[1]]) == true
    end;

    @testset "iswithin_hfov Tests" begin
        robot = sb.Robot([2.0, 2.0], 0.0)
        @test sb.iswithin_hfov(robot, [4.0, 2.0]) == true
        @test sb.iswithin_hfov(robot, [0.0, 2.0]) == false
        @test sb.iswithin_hfov(robot, [4.0, 9.0]) == false
        @test sb.iswithin_hfov(robot, [3.0, 0.0]) == false
        robot = sb.Robot([0.0, 0.0], 0.0)
        @test sb.iswithin_hfov(robot, [3.0, 0.0]) == true
    end;

    # @testset "what_is_in_line_of_sight Tests" begin
    #     # myPlayer at [3,3]
    #     # obstacle (opponent) at [6,3] and size 0.3
    #     # ball at [8,3]
    #     ws = sb.WorldState([sb.Posture([3.0, 3.0], 0.0)], [sb.Posture([6.0, 3.0], 0.0)], [8.0, 3.0])
    #     @test sb.what_is_in_line_of_sight(ws, ws.myPlayers[1], ws.ball) == [sb.Posture([6.0, 3.0], 0.0)]
    #     ws = sb.WorldState([sb.Posture([3.0, 3.0], 0.0)], [sb.Posture([6.0, 3.0], 0.0)], [8.0, 3.0])
    #     @test sb.what_is_in_line_of_sight(ws, ws.myPlayers[1], ws.ball, excludes=[[6.0, 3.0]]) == []
    #     ws = sb.WorldState([sb.Posture([3.0, 3.0], 0.0)], [sb.Posture([6.0, 3.15], 0.0)], [8.0, 3.0])
    #     @test sb.what_is_in_line_of_sight(ws, ws.myPlayers[1], ws.ball, tol=0.05) == []
    #     ws = sb.WorldState([sb.Posture([8.0, 6.0], 0.0)], [sb.Posture([6.0, 3.0], 0.0)], [8.0, 3.0])
    #     @test sb.what_is_in_line_of_sight(ws, ws.myPlayers[1], ws.ball) == []
    # end;

    @testset "intervall" begin
        @test sb.intervall(1, 4, 4) == [2, 3, 4]
        @test sb.intervall(1, 4, 4, true) == [2, 3, 4]
        @test sb.intervall(3, 1, 4) == [4, 1]
        @test sb.intervall(3, 1, 4, true) == [4, 1]
        @test sb.intervall(5, 3, 8) == [6, 7, 8, 1, 2, 3]
        @test sb.intervall(5, 3, 8, true) == [6, 7, 8, 1, 2, 3]
        @test sb.intervall(2, 2, 4) == []
        @test sb.intervall(2, 2, 4, true) == [3, 4, 1, 2]
        @test sb.intervall(5, 5, 8, true) == [6, 7, 8, 1, 2, 3, 4, 5]
    end;

    @testset "get_edge Tests" begin
        @test sb.get_edge([9.0, 2.32]) == 2
        @test sb.get_edge([9.0, 0.0]) == 1
        @test sb.get_edge([9.0, 6.0]) == 2
        @test sb.get_edge([0.0, 6.0]) == 3
        @test sb.get_edge([0.0, 0.0]) == 1
        @test sb.get_edge([0.0, 2.32]) == 4
        @test sb.get_edge([5.0, 6.0]) == 3
        @test sb.get_edge([1.0, 0.0]) == 1
    end;

    @testset "raster Tests" begin
        @test sb.raster(0.0, 10.0, 9) == 1.0:9.0
        @test sb.raster(0.0, 2.0, 3) == 0.5:0.5:1.5
        @test sb.raster(0.0, 6.0, 59) == 0.1:0.1:5.9
        @test sb.raster(0.0, 9.0, 89) == 0.1:0.1:8.9
    end;

    @testset "get_shadow_polygon Tests" begin
        config = sb.get_config(["--add_myPlayer","\"--position -0.5 1.0\""])
        @test string(sb.get_shadow_polygon([-1.0, 0.7],config.ws.myPlayers[1], config)) == "GeometryTypes.Point{2,Float64}[[-0.458511, 0.855852], [4.5, 2.28302], [1.23115, 3.0], [-0.607665, 1.10444], [-0.458511, 0.855852]]"
        @test string(sb.get_shadow_polygon([-1.5, 0.0],config.ws.myPlayers[1], config)) == "GeometryTypes.Point{2,Float64}[[-0.405782, 0.883282], [2.21643, 3.0], [0.921682, 3.0], [-0.616718, 1.09422], [-0.405782, 0.883282]]"
    end;

    @testset "get_shadow_polygons Tests" begin
        config = sb.get_config(["--add_myPlayer","\"--position -0.5 1.0\"","--add_opponent","\"--position 1.5 1.0\""])
        @test string(sb.get_shadow_polygons(config.ws.myPlayers[1], config)) == "Array{GeometryTypes.Point{2,Float64},1}[[[4.5, 1.3761], [1.4888, 1.1496], [1.4888, 0.8504], [4.5, 0.6239]]]"
        @test string(sb.get_shadow_polygons(config.ws.myPlayers[1], config, true)) == "Array{GeometryTypes.Point{2,Float64},1}[[[2.8974, 3.0], [0.0398, 0.6822], [-0.5, 1.0], [2.8974, 3.0], [-4.5, 3.0], [-0.0894, 0.5774], [-4.5, -3.0], [4.5, -3.0], [4.5, 3.0]]]"
    end;

end


text:
module SingleLayerQG

export
  Problem,
  set_zeta!,
  updatevars!,

  kinetic_energy,
  potential_energy,
  energy_dissipation,
  energy_work,
  energy_drag,
  enstrophy,
  reduced_enstrophy,
  enstrophy_dissipation,
  enstrophy_work,
  enstrophy_drag

using
  CUDA,
  Reexport

@reexport using FourierFlows

using FFTW: rfft
using LinearAlgebra: mul!, ldiv!
using FourierFlows: parsevalsum, parsevalsum2

nothingfunction(args...) = nothing

"""
    Problem(; parameters...)

Construct a BarotropicQG turbulence problem.
"""

function Problem(dev::Device=CPU();
  # Numerical parameters
          nx = 256,
          Lx = 2π,
          ny = nx,
          Ly = Lx,
          dt = 0.01,
  # Physical parameters
           β = 0.0,
        kdef = 0.0,
         eta = nothing,
  # Drag and/or hyper-/hypo-viscosity
           ν = 0.0,
          nν = 1,
           μ = 0.0,
  # Timestepper and equation options
     stepper = "RK4",
       calcF = nothingfunction,
  stochastic = false,
           T = Float64)

  # the grid
  grid = TwoDGrid(dev, nx, Lx, ny, Ly; T=T)
  x, y = gridpoints(grid)

  # topographic PV
  eta === nothing && ( eta = zeros(dev, T, (nx, ny)) )

  params = !(typeof(eta)<:ArrayType(dev)) ? Params(grid, β, kdef, eta, μ, ν, nν, calcF) : Params(β, kdef, eta, rfft(eta), μ, ν, nν, calcF)

  vars = calcF == nothingfunction ? Vars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid))

  equation = Equation(params, grid)

  return FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
end

# ----------
# Parameters
# ----------

"""
    Params(g::TwoDGrid, β, eta, μ, ν, nν, calcF)

Returns the params for an unforced two-dimensional barotropic QG problem.
"""
struct Params{T, Aphys, Atrans} <: AbstractParams
        β :: T            # Planetary vorticity y-gradient
     kdef :: T            # deformation wavenumber
      eta :: Aphys        # Topographic PV
     etah :: Atrans       # FFT of Topographic PV
        μ :: T            # Linear drag
        ν :: T            # Viscosity coefficient
       nν :: Int          # Hyperviscous order (nν=1 is plain old viscosity)
   calcF! :: Function     # Function that calculates the forcing on QGPV q
end

"""
    Params(g::TwoDGrid, β, eta::Function, μ, ν, nν, calcF)

Constructor for Params that accepts a generating function for the topographic PV.
"""
function Params(grid::AbstractGrid{T, A}, β, kdef, eta::Function, μ, ν, nν::Int, calcF) where {T, A}
  etagrid = A([eta(grid.x[i], grid.y[j]) for i=1:grid.nx, j=1:grid.ny])
  etah = rfft(etagrid)
  return Params(β, kdef, etagrid, etah, μ, ν, nν, calcF)
end

# ---------
# Equations
# ---------

"""
    Equation(params, grid)

Returns the equation for two-dimensional barotropic QG problem with `params` and `grid`.
"""
function Equation(params::Params, grid::AbstractGrid)
  L = @. - params.μ - params.ν * grid.Krsq^params.nν + im * params.β * grid.kr * grid.invKrsq
  CUDA.@allowscalar L[1, 1] = 0
  return FourierFlows.Equation(L, calcN!, grid)
end

# ----
# Vars
# ----

abstract type BarotropicQGVars <: AbstractVars end

struct Vars{Aphys, Atrans, F, P} <: BarotropicQGVars
        q :: Aphys
     zeta :: Aphys
      psi :: Aphys
        u :: Aphys
        v :: Aphys
       qh :: Atrans
    zetah :: Atrans
     psih :: Atrans
       uh :: Atrans
       vh :: Atrans
       Fh :: F
  prevsol :: P
end

const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing}
const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}

"""
    Vars(dev, grid)

Returns the vars for unforced two-dimensional barotropic QG problem on device `dev` and with `grid`
"""
function Vars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) q u v psi zeta
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) qh uh vh psih zetah
  Vars(q, zeta, psi, u, v, qh, zetah, psih, uh, vh, nothing, nothing)
end

"""
    ForcedVars(dev, grid)

Returns the vars for forced two-dimensional barotropic QG problem on device dev and with `grid`.
"""
function ForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) q u v psi zeta
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) qh uh vh psih zetah Fh
  return Vars(q, zeta, psi, u, v, qh, zetah, psih, uh, vh, Fh, nothing)
end

"""
    StochasticForcedVars(dev, grid)

Returns the vars for stochastically forced two-dimensional barotropic QG problem on device dev and with `grid`.
"""
function StochasticForcedVars(dev::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) q u v psi zeta
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) qh uh vh psih zetah Fh prevsol
  return Vars(q, zeta, psi, u, v, qh, zetah, psih, uh, vh, Fh, prevsol)
end

# -------
# Solvers
# -------

function calcN_advection!(N, sol, t, clock, vars, params, grid)
  @. vars.zetah = sol
  @. vars.psih  = - vars.zetah / (grid.Krsq + params.kdef^2)
  if params.kdef == 0.0
      CUDA.@allowscalar vars.psih[1, 1] = 0
  end
  @. vars.uh    = -im * grid.l  * vars.psih
  @. vars.vh    =  im * grid.kr * vars.psih

  ldiv!(vars.zeta, grid.rfftplan, vars.zetah)
  ldiv!(vars.u, grid.rfftplan, vars.uh)
  ldiv!(vars.v, grid.rfftplan, vars.vh)

  @. vars.q = vars.zeta + params.eta
  uq = vars.u                                            # use vars.u as scratch variable
  @. uq *= vars.q                                        # u*q
  vq = vars.v                                            # use vars.v as scratch variable
  @. vq *= vars.q                                        # v*q

  uqh = vars.uh                                          # use vars.uh as scratch variable
  mul!(uqh, grid.rfftplan, uq)                           # \hat{u*q}
  vqh = vars.vh                                          # use vars.vh as scratch variable
  mul!(vqh, grid.rfftplan, vq)                           # \hat{v*q}

  @. N = -im * grid.kr * uqh - im * grid.l * vqh         # -∂(u*q)/∂x -∂(v*q)/∂y
  return nothing
end

function calcN!(N, sol, t, clock, vars, params, grid)
  calcN_advection!(N, sol, t, clock, vars, params, grid)
  addforcing!(N, sol, t, clock, vars, params, grid)

  return nothing
end

addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing

function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
  params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  @. N += vars.Fh

  return nothing
end

function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid)
  if t == clock.t # not a substep
    @. vars.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing
    params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  end

  @. N += vars.Fh

  return nothing
end

# ----------------
# Helper functions
# ----------------

"""
    updatevars!(sol, vars, params, grid)

Update the variables in `vars` with the solution in `sol`.
"""
function updatevars!(sol, vars, params, grid)
  @. vars.zetah = sol
  @. vars.psih  = - vars.zetah / (grid.Krsq + params.kdef^2)
  if params.kdef == 0.0
      CUDA.@allowscalar vars.psih[1, 1] = 0
  end
  @. vars.uh    = -im * grid.l  * vars.psih
  @. vars.vh    =  im * grid.kr * vars.psih

  ldiv!(vars.zeta, grid.rfftplan, deepcopy(vars.zetah))
  ldiv!(vars.psi, grid.rfftplan, deepcopy(vars.psih))
  ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh))
  ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh))

  @. vars.q = vars.zeta + params.eta

  return nothing
end

updatevars!(prob) = updatevars!(prob.sol, prob.vars, prob.params, prob.grid)

"""
    set_zeta!(prob, zeta)
    set_zeta!(sol, vars, params, grid)

Set the solution `sol` as the transform of zeta and update variables `vars`
on the `grid`.
"""
function set_zeta!(sol, vars, params, grid, zeta)
  mul!(vars.zetah, grid.rfftplan, zeta)

  @. sol = vars.zetah

  updatevars!(sol, vars, params, grid)

  return nothing
end

set_zeta!(prob, zeta) = set_zeta!(prob.sol, prob.vars, prob.params, prob.grid, zeta)

"""
    kinetic_energy(prob)
    potential_energy(prob)
    kinetic_energy(vars, grid)
    potential_energy(vars, grid, params)

Returns the domain-averaged kinetic energy of solution `sol`: ∫ ½ (u²+v²) dxdy / (Lx Ly) = ∑ ½ k² |ψ̂|² / (Lx Ly).
Returns the domain-averaged potential energy of solution `sol`: ½ kdef² ∫ ψ² dxdy / (Lx Ly) = ½ kdef² ∑ |ψ̂|² / (Lx Ly).

"""
function kinetic_energy(sol, grid, vars, params)
    @. vars.uh = sqrt.(grid.Krsq) * sol /(grid.Krsq + params.kdef^2) ## uh is a dummy variable
    if params.kdef == 0.0
        CUDA.@allowscalar vars.uh[1, 1] = 0
    end
    return parsevalsum2(vars.uh , grid) / (2 * grid.Lx * grid.Ly)
end

function potential_energy(sol, grid, vars, params)
    @. vars.uh = sol /(grid.Krsq + params.kdef^2) ## uh is a dummy variable
    if params.kdef == 0.0
        CUDA.@allowscalar vars.uh[1, 1] = 0
    end
    return params.kdef^2*parsevalsum2(vars.uh, grid) / (2 * grid.Lx * grid.Ly)
end

kinetic_energy(prob) = kinetic_energy(prob.sol, prob.grid, prob.vars, prob.params)
potential_energy(prob) = potential_energy(prob.sol, prob.grid, prob.vars, prob.params)

"""
    enstrophy(prob)
    enstrophy(sol, grid, vars)
    reduced_enstrophy(prob)
    reduced_enstrophy(sol, grid, vars)

Returns the domain-averaged enstrophy ½ ∫ q² dxdy / (Lx Ly), with q = ζ + η and sol = ζ.
Returns the domain-averaged reduced enstrophy ½ ∫(ζ² + 2ζη) dxdy / (Lx Ly) .

"""
function enstrophy(sol, grid, vars, params)
  @. vars.zetah = sol
  return 0.5*parsevalsum2(vars.zetah + params.etah, grid) / (grid.Lx * grid.Ly)
end
enstrophy(prob) = enstrophy(prob.sol, prob.grid, prob.vars, prob.params)

function reduced_enstrophy(sol, grid, vars, params)
  @. vars.zetah = sol
  return 0.5*parsevalsum(abs2.(vars.zetah) .+ 2* vars.zetah .* params.etah, grid) / (grid.Lx * grid.Ly)
end
reduced_enstrophy(prob) = reduced_enstrophy(prob.sol, prob.grid, prob.vars, prob.params)

"""
    energy_dissipation(prob)
    energy_dissipation(sol, vars, params, grid)

Returns the domain-averaged energy dissipation rate. nν must be >= 1.
"""
@inline function energy_dissipation(sol, vars, params, grid)
  energy_dissipationh = vars.uh # use vars.uh as scratch variable

  @. energy_dissipationh = params.ν * grid.Krsq^(params.nν-1) * abs2(sol)
  CUDA.@allowscalar energy_dissipationh[1, 1] = 0
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dissipationh, grid)
end

@inline energy_dissipation(prob) = energy_dissipation(prob.sol, prob.vars, prob.params, prob.grid)

"""
    enstrophy_dissipation(prob)
    enstrophy_dissipation(sol, vars, params, grid)

Returns the domain-averaged enstrophy dissipation rate. nν must be >= 1.
"""
@inline function enstrophy_dissipation(sol, vars, params, grid)
  enstrophy_dissipationh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_dissipationh = params.ν * grid.Krsq^params.nν * abs2(sol)
  CUDA.@allowscalar enstrophy_dissipationh[1, 1] = 0
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dissipationh, grid)
end

@inline enstrophy_dissipation(prob) = enstrophy_dissipation(prob.sol, prob.vars, prob.params, prob.grid)

"""
    energy_work(prob)
    energy_work(sol, vars, grid)

Returns the domain-averaged rate of work of energy by the forcing `Fh`.
"""
@inline function energy_work(sol, vars::ForcedVars, grid)
  energy_workh = vars.uh # use vars.uh as scratch variable

  @. energy_workh = grid.invKrsq * sol * conj(vars.Fh)
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(energy_workh, grid)
end

@inline function energy_work(sol, vars::StochasticForcedVars, grid)
  energy_workh = vars.uh # use vars.uh as scratch variable

  @. energy_workh = grid.invKrsq * (vars.prevsol + sol)/2 * conj(vars.Fh) # Stratonovich
  # @. energy_workh = grid.invKrsq * vars.prevsol * conj(vars.Fh)             # Ito
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(vars.uh, grid)
end

@inline energy_work(prob) = energy_work(prob.sol, prob.vars, prob.grid)

"""
    enstrophy_work(prob)
    enstrophy_work(sol, vars, grid)

Returns the domain-averaged rate of work of enstrophy by the forcing `Fh`.
"""
@inline function enstrophy_work(sol, vars::ForcedVars, grid)
  enstrophy_workh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_workh = sol * conj(vars.Fh)
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid)
end

@inline function enstrophy_work(sol, vars::StochasticForcedVars, grid)
  enstrophy_workh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_workh = (vars.prevsol + sol) / 2 * conj(vars.Fh) # Stratonovich
  # @. enstrophy_workh = grid.invKrsq * vars.prevsol * conj(vars.Fh)           # Ito
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_workh, grid)
end

@inline enstrophy_work(prob) = enstrophy_work(prob.sol, prob.vars, prob.grid)

"""
    energy_drag(prob)

Returns the extraction of domain-averaged energy by drag μ.
"""
@inline function energy_drag(prob)
  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
  energy_dragh = vars.uh # use vars.uh as scratch variable

  @. energy_dragh = params.μ * grid.invKrsq * abs2(sol)
  CUDA.@allowscalar energy_dragh[1, 1] = 0
  return  1 / (grid.Lx * grid.Ly) * parsevalsum(energy_dragh, grid)
end

"""
    enstrophy_drag(prob)

Returns the extraction of domain-averaged enstrophy by drag/hypodrag μ.
"""
@inline function enstrophy_drag(prob)
  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
  enstrophy_dragh = vars.uh # use vars.uh as scratch variable

  @. enstrophy_dragh = params.μ * abs2(sol)
  CUDA.@allowscalar enstrophy_dragh[1, 1] = 0
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(enstrophy_dragh, grid)
end

end # module


text:
import PyCall: PyObject

type Point
    x::Float64
    y::Float64
end
Base.zero(::Point) = Point(0.0, 0.0)

function Base.convert(::Type{Point}, t::Union{Tuple, AbstractArray})
    if length(t) != 2
        throw(ArgumentError("Only length 2 objects can be converted to Point"))
    end
    Point(t[1], t[2])
end
PyObject(t::Point) = PyObject((t.x, t.y))

function Base.convert{T <: Real}(::Type{Array{Point}}, a::Array{T, 2})
    (na, ma) = size(a)
    #TODO: need to think about a 2x2 array, since it is ambigous, currently I assume
    # column major
    if na == 2
        squeeze(mapslices(x->Point(x), a, 1), 1)
    elseif ma == 2
        squeeze(mapslices(x->Point(x), a, 2), 2)
    else
        throw(DimensionError())
    end
end

##TODO: think about the best name for this ...
abstract Grob

max_stars_repo_name: gabrielgellner/OdinSon

max_issues_repo_name: gabrielgellner/OdinSon

max_forks_repo_name: gabrielgellner/OdinSon


text:
rng = ParisiRapuano()
n=1000
@test all(0 .<= rand(rng, n) .< 1.)
# @test all(-5 .<= randn(rng, 1000) .< 5.)
n=100000
# FAILING. Some proble with Close1Open2?
#@test abs(mean(randn(rng, n))) < 5 / √n
for n=[1_000, 10_000]
    @test abs(mean(rand(rng, n))-0.5) < 20abs(mean(rand(n))-0.5)
    @test abs(var(rand(rng, n)) - 1/12) < 20abs(var(rand(n)) - 1/12)
end

max_stars_repo_name: JuliaTagBot/LittleScienceTools

max_issues_repo_name: JuliaTagBot/LittleScienceTools

max_forks_repo_name: JuliaTagBot/LittleScienceTools


text:
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin

function save_solution_file(u, time, dt, timestep,
                            mesh::Union{SerialTreeMesh, StructuredMesh, UnstructuredMesh2D, P4estMesh},
                            equations, dg::DG, cache,
                            solution_callback, element_variables=Dict{Symbol,Any}();
                            system="")
  @unpack output_directory, solution_variables = solution_callback

  # Filename without extension based on current time step
  if isempty(system)
    filename = joinpath(output_directory, @sprintf("solution_%06d.h5", timestep))
  else
    filename = joinpath(output_directory, @sprintf("solution_%s_%06d.h5", system, timestep))
  end

  # Convert to different set of variables if requested
  if solution_variables === cons2cons
    data = u
    n_vars = nvariables(equations)
  else
    # Reinterpret the solution array as an array of conservative variables,
    # compute the solution variables via broadcasting, and reinterpret the
    # result as a plain array of floating point numbers
    data = Array(reinterpret(eltype(u),
           solution_variables.(reinterpret(SVector{nvariables(equations),eltype(u)}, u),
                      Ref(equations))))

    # Find out variable count by looking at output from `solution_variables` function
    n_vars = size(data, 1)
  end

  # Open file (clobber existing content)
  h5open(filename, "w") do file
    # Add context information as attributes
    attributes(file)["ndims"] = ndims(mesh)
    attributes(file)["equations"] = get_name(equations)
    attributes(file)["polydeg"] = polydeg(dg)
    attributes(file)["n_vars"] = n_vars
    attributes(file)["n_elements"] = nelements(dg, cache)
    attributes(file)["mesh_type"] = get_name(mesh)
    attributes(file)["mesh_file"] = splitdir(mesh.current_filename)[2]
    attributes(file)["time"] = convert(Float64, time) # Ensure that `time` is written as a double precision scalar
    attributes(file)["dt"] = convert(Float64, dt) # Ensure that `dt` is written as a double precision scalar
    attributes(file)["timestep"] = timestep

    # Store each variable of the solution data
    for v in 1:n_vars
      # Convert to 1D array
      file["variables_$v"] = vec(data[v, .., :])

      # Add variable name as attribute
      var = file["variables_$v"]
      attributes(var)["name"] = varnames(solution_variables, equations)[v]
    end

    # Store element variables
    for (v, (key, element_variable)) in enumerate(element_variables)
      # Add to file
      file["element_variables_$v"] = element_variable

      # Add variable name as attribute
      var = file["element_variables_$v"]
      attributes(var)["name"] = string(key)
    end
  end

  return filename
end

function save_solution_file(u, time, dt, timestep,
                            mesh::ParallelTreeMesh, equations, dg::DG, cache,
                            solution_callback, element_variables=Dict{Symbol,Any}();
                            system="")
  @unpack output_directory, solution_variables = solution_callback

  # Filename without extension based on current time step
  if isempty(system)
    filename = joinpath(output_directory, @sprintf("solution_%06d.h5", timestep))
  else
    filename = joinpath(output_directory, @sprintf("solution_%s_%06d.h5", system, timestep))
  end

  # Convert to different set of variables if requested
  if solution_variables === cons2cons
    data = u
    n_vars = nvariables(equations)
  else
    # Reinterpret the solution array as an array of conservative variables,
    # compute the solution variables via broadcasting, and reinterpret the
    # result as a plain array of floating point numbers
    data = Array(reinterpret(eltype(u),
           solution_variables.(reinterpret(SVector{nvariables(equations),eltype(u)}, u),
                      Ref(equations))))

    # Find out variable count by looking at output from `solution_variables` function
    n_vars = size(data, 1)
  end

  # Calculate element and node counts by MPI rank
  element_size = nnodes(dg)^ndims(mesh)
  element_counts = convert(Vector{Cint}, collect(cache.mpi_cache.n_elements_by_rank))
  node_counts = element_counts * Cint(element_size)

  # non-root ranks only send data
  if !mpi_isroot()
    # Send nodal data to root
    for v in 1:n_vars
      MPI.Gatherv!(vec(data[v, .., :]), nothing, mpi_root(), mpi_comm())
    end

    # Send element data to root
    for (key, element_variable) in element_variables
      MPI.Gatherv!(element_variable, nothing, mpi_root(), mpi_comm())
    end

    return filename
  end

  # Open file (clobber existing content)
  h5open(filename, "w") do file
    # Add context information as attributes
    attributes(file)["ndims"] = ndims(mesh)
    attributes(file)["equations"] = get_name(equations)
    attributes(file)["polydeg"] = polydeg(dg)
    attributes(file)["n_vars"] = n_vars
    attributes(file)["n_elements"] = nelementsglobal(dg, cache)
    attributes(file)["mesh_type"] = get_name(mesh)
    attributes(file)["mesh_file"] = splitdir(mesh.current_filename)[2]
    attributes(file)["time"] = convert(Float64, time) # Ensure that `time` is written as a double precision scalar
    attributes(file)["dt"] = convert(Float64, dt) # Ensure that `dt` is written as a double precision scalar
    attributes(file)["timestep"] = timestep

    # Store each variable of the solution data
    for v in 1:n_vars
      # Convert to 1D array
      recv = Vector{eltype(data)}(undef, sum(node_counts))
      MPI.Gatherv!(vec(data[v, .., :]), MPI.VBuffer(recv, node_counts), mpi_root(), mpi_comm())
      file["variables_$v"] = recv

      # Add variable name as attribute
      var = file["variables_$v"]
      attributes(var)["name"] = varnames(solution_variables, equations)[v]
    end

    # Store element variables
    for (v, (key, element_variable)) in enumerate(element_variables)
      # Add to file
      recv = Vector{eltype(data)}(undef, sum(element_counts))
      MPI.Gatherv!(element_variable, MPI.VBuffer(recv, element_counts), mpi_root(), mpi_comm())
      file["element_variables_$v"] = recv

      # Add variable name as attribute
      var = file["element_variables_$v"]
      attributes(var)["name"] = string(key)
    end
  end

  return filename
end

end # @muladd


text:
module TensorNetworkEvolve

using Yao
using Yao.ConstGate: P1
using Yao: Add
using Graphs
using Graphs: SimpleEdge
using Random
using LinearAlgebra
using OMEinsumContractionOrders
using OMEinsumContractionOrders: CodeOptimizer, CodeSimplifier, SlicedEinsum
using OMEinsum
using OMEinsum: DynamicEinCode, NestedEinsum
using KrylovKit
#using Zygote: gradient, @non_differentiable

export PEPS, VidalPEPS, SimplePEPS, zero_vidalpeps, zero_simplepeps, rand_simplepeps, rand_vidalpeps
export state, statevec, getvlabel, getphysicallabel, newlabel, findbondtensor, virtualbonds
export apply_onbond!, apply_onsite!, inner_product, norm, normalize!
export variables, load_variables!, load_variables

# load AD submodule
include("TensorAD/TensorAD.jl")
using .TensorAD: gradient, DiffTensor

include("peps.jl")
include("tebd.jl")
include("timeevolve.jl")
#include("cracker.jl")

end


text:
###############################################################################
#
# Quadratic function solver
#
###############################################################################
"""

    lower_quadratic(a::FT, b::FT, c::FT) where {FT<:AbstractFloat}

Return the lower quadratic solution or NaN, given
- `a` Parameter in `a*x^2 + b*x + c = 0`
- `b` Parameter in `a*x^2 + b*x + c = 0`
- `c` Parameter in `a*x^2 + b*x + c = 0`

---
Example
```julia
low_x = lower_quadratic(1.0, 2.0, 0.0);
low_x = lower_quadratic(-1.0, 2.0, 0.0);
low_x = lower_quadratic(1.0, 2.0, 2.0);
```
"""
function lower_quadratic(a::FT, b::FT, c::FT) where {FT<:AbstractFloat}
    discr = b^2 - 4*a*c;

    if (discr >= 0) && (a > 0)
        return (-b - sqrt(discr)) / (2 * a)
    elseif (discr >= 0) && (a < 0)
        return (-b + sqrt(discr)) / (2 * a)
    else
        return FT(NaN)
    end
end

"""

    upper_quadratic(a::FT, b::FT, c::FT) where {FT<:AbstractFloat}

Return the upper quadratic solution or NaN, given
- `a` Parameter in `a*x^2 + b*x + c = 0`
- `b` Parameter in `a*x^2 + b*x + c = 0`
- `c` Parameter in `a*x^2 + b*x + c = 0`

---
Example
```julia
up_x = upper_quadratic(1.0, 2.0, 0.0);
up_x = upper_quadratic(-1.0, 2.0, 0.0);
up_x = upper_quadratic(1.0, 2.0, 2.0);
```
"""
function upper_quadratic(a::FT, b::FT, c::FT) where {FT<:AbstractFloat}
    discr = b^2 - 4*a*c;

    if (discr >= 0) && (a > 0)
        return (-b + sqrt(discr)) / (2 * a)
    elseif (discr >= 0) && (a < 0)
        return (-b - sqrt(discr)) / (2 * a)
    else
        return FT(NaN)
    end
end


text:
function usegpu(gpu::Bool)
global GPU=gpu
global CPU=~gpu

    f=open(ADDIR*"/src/initfile.jl","w")
    if gpu==true
        write(f,"GPU=true\n")
    else
        write(f,"GPU=false\n")
    end
close(f)
end

max_stars_repo_name: davidbarber/AutoDiff

max_issues_repo_name: davidbarber/AutoDiff

max_forks_repo_name: davidbarber/AutoDiff


text:
using FirstOrderOptimization

## The state of this code:
# * seems like the algo works.
# 	* obj converges to a nice low number
# 	* i haven't checked it against an sdp solver yet
# 	* or tested for recovery
# * it does a super weird thing where
# 	the rank and obj shoot up in the 14th iteration. why?

srand(1)

type PrevRank
	r::Integer
end

### max norm regularized matrix completion problem

# minimize f(W) + g(W) + h(W)

# where W = [A X; X' B] and

# (data fitting)   f(W) = vecnorm(P.*(X - M))^2
# (regularization) g(W) = lambda*max(diag(X))
# (psd)            h(W) = (W \succeq 0) ? 0 : Inf

function generate_maxnorm_problem(m,n,lambda,k)
	# generate data
	X0 = randn(m,k)
	Y0 = randn(k,n)
	A = X0*Y0 + .1*randn(m,n)
	P = float(randn(m,n) .>= .5) # observed entries

	X(W) = W[1:m, m+1:end]

	function grad_f(W)
		gX = P.*(X(W)-A)
		return [zeros(m,m) gX; gX' zeros(n,n)]
	end

	function prox_g(W, alpha)
		Z = copy(W)
		oldmax = maximum(diag(W))
		newmax = oldmax - lambda*alpha/2
		for i=1:size(Z,1)
			if Z[i,i] > newmax
				Z[i,i] = newmax
			end
		end
		Z
	end

	# we're going to use a closure over prevrank
	# to remember what the rank of prox_h(W) was the last time we computed it
	# in order to avoid calculating too many eigentuples of W
	prevrank = PrevRank(k)

	function prox_h(W, alpha=0; TOL=1e-10)
		# debugging: @show prevrank.r
		while prevrank.r < size(W,1)
			l,v = eigs(Symmetric(W), nev = prevrank.r+1, which=:LR) # v0 = [v zeros(size(W,1), prevrank.r+1 - size(v,2))]
			if l[end] <= TOL
				prevrank.r = sum(l.>=TOL)
				return v*diagm(max(l,0))*v'
			else
				prevrank.r = 2*prevrank.r # double the rank and try again
			end
		end
		# else give up on computational cleverness
		l,v = eig(Symmetric(W))
		prevrank.r = sum(l.>=TOL)
		return v*diagm(max(l,0))*v'
	end

	# we're not going to bother checking whether W is psd or not
	# when evaluating the objective; in the course of the prisma
	# algo this makes no difference
	obj(W) = sum((P.*(X(W) - A)).^2) + lambda*maximum(diag(W))

	return grad_f, prox_g, prox_h, obj
end

function test_prisma(m,n,lambda,k)

	grad_f, prox_g, prox_h, obj = generate_maxnorm_problem(m,n,lambda,k)

	# initialize
	W = zeros(m+n,m+n)
	# lipshitz constant for f
	L_f = 2
	# orabona starts stepsize at
	# beta = lambda/sqrt((m+n)^2*mean(A.^2))
	beta   = lambda/sqrt(obj(W))
	ssr    = PrismaStepsize(beta)
	params = PrismaParams(ssr, 10, 1)

	# recover
	W = PRISMA(W, L_f,
		   grad_f,
		   prox_g,
		   prox_h,
		   obj,
		   params)
end

test_prisma(5,5,10,3)

max_stars_repo_name: madeleineudell/FirstOrderOptimization


text:
module ImageMorphology

using ImageCore
using Base.Cartesian # TODO: delete this

include("convexhull.jl")
include("connected.jl")

include("dilation_and_erosion.jl")
include("thinning.jl")
include("imfill.jl")

export
    dilate,
    dilate!,
    erode,
    erode!,
    opening,
    closing,
    tophat,
    bothat,
    morphogradient,
    morpholaplace,

    # connected.jl
    label_components,
    label_components!,
    component_boxes,
    component_lengths,
    component_indices,
    component_subscripts,
    component_centroids,

    # convexhull.jl
    convexhull,

    # thinning.jl
    thinning,
    GuoAlgo,
    imfill

end # module


text:
export NN

"""
# a tiny coarse-grained neural network framework

The `nn` module only offers a protocal to calculate
neural networks on any array type of Julia that overloads
the interface of `AbtractArray`.

You can use native `Array`, `AbstractSparseArray`, and also
`StaticArray` for small calculations. Moreover, you can also
use GPU arrays, like `ArrayFire`, `CuArray` and etc.

NOTE: GPU part is under development.

each subtype of `AbstractModel` has the following interface

    forward(operator, input)

For some block/operators/models, the forward method will record the input
and output value for backward use, thus, be careful, make sure you provide
a new array with differnet reference to input for the output of forward, or
the content stored for backward may change.

    backward(operator, grad)

backward method will calculate gradient, it will store parameter gradient into
the operator and return the gradient of input.
"""
module NN

import QMTK: forward, backward
using Compat, StaticArrays, QMTK

include("linear/Linear.jl")

# activations
include("activation/Activation.jl")

end


text:
#
# Contraction sequence
#

# Tools for contracting a network with a sequence
function _contract(label1::String, label2::String)
  return string("(", label1, "*", label2, ")")
end

function _contract(tensor1::ITensor, tensor2::ITensor)
  return ITensor(noncommoninds(tensor1, tensor2))
end

sequence_traversal(sequence) = reverse(collect(StatelessBFS(sequence)))

function contract_dict(tensors, sequence, traversal=sequence_traversal(sequence))
  net_tensors = Dict()
  traversal = reverse(collect(StatelessBFS(sequence)))
  for net in traversal
    if net isa Int
      net_tensors[net] = tensors[net]
    else # net isa Vector
      net_tensors[net] = _contract(net_tensors[net[1]], net_tensors[net[2]])
    end
  end
  return net_tensors
end

# Return all of the contractions involved in the sequence.
function contraction_sequence(
  tensors,
  sequence,
  traversal=sequence_traversal(sequence),
  contract_dict=contract_dict(tensors, sequence, traversal),
)
  all_tensors = Any[]
  tensors_1 = Vector{Union{Nothing,eltype(tensors)}}(tensors)
  net_position = Dict()
  N = length(tensors)
  n = N + 1
  for net in traversal
    if net isa Int
      net_position[net] = net
    else
      net_position[net] = n
      n += 1
    end
    if !isa(net, Int)
      for n in net
        tensors_1[net_position[n]] = nothing
      end
      push!(tensors_1, contract_dict[net])
      push!(all_tensors, copy(tensors_1))
    end
  end
  return convert.(Vector{eltype(tensors)}, filter.(!isnothing, all_tensors))
end

#
# Convert a tree to a graph
#

tree_to_graph(tr) = tree_to_graph(Tree(tr))

function tree_to_graph(tr::Tree)
  g = SimpleDiGraph()
  labels = Any[]
  walk_tree!(g, labels, tr)
  return (g, labels)
end

function walk_tree!(g, labels, tr::Tree)
  add_vertex!(g)
  top_vertex = vertices(g)[end]
  push!(labels, tr.x)
  for i in 1:length(tr.x)
    if isa(tr[i], Vector)
      child = walk_tree!(g, labels, Tree(tr[i]))
      add_edge!(g, child, top_vertex)
    else
      add_vertex!(g)
      n = vertices(g)[end]
      add_edge!(g, n, top_vertex)
      push!(labels, tr[i])
    end
  end
  return top_vertex
end

# Visualization function interface. Ultimately calls a beckend.

function visualize(g::AbstractGraph, sequence=nothing; backend=get_backend(), kwargs...)
  # TODO: do something with the sequence (show sequence, add labels indicating sequence, etc.)
  return visualize(Backend(backend), g; kwargs...)
end

function visualize(tn::Vector{ITensor}, sequence=nothing; kwargs...)
  return visualize(MetaDiGraph(tn), sequence; kwargs...)
end

function visualize(tn::Tuple{Vector{ITensor}}, args...; kwargs...)
  return visualize(only(tn), args...; kwargs...)
end
visualize(ψ::MPS, args...; kwargs...) = visualize(data(ψ), args...; kwargs...)
function visualize(tn::Tuple{Vararg{ITensor}}, args...; kwargs...)
  return visualize(collect(tn), args...; kwargs...)
end
visualize(tn::ITensor...; kwargs...) = visualize(collect(tn); kwargs...)

function visualize!(fig, g::AbstractGraph; backend=get_backend(), kwargs...)
  return visualize!(Backend(backend), fig, g; kwargs...)
end

function visualize!(fig, tn::Vector{ITensor}, sequence=nothing; kwargs...)
  return visualize!(fig, MetaDiGraph(tn); kwargs...)
end
visualize!(fig, ψ::MPS, sequence=nothing; kwargs...) = visualize!(fig, data(ψ); kwargs...)
function visualize!(fig, tn::Tuple{Vararg{ITensor}}, sequence=nothing; kwargs...)
  return visualize!(fig, collect(tn); kwargs...)
end
visualize!(fig, tn::ITensor...; kwargs...) = visualize!(fig, collect(tn); kwargs...)

function visualize!(fig, tn::Tuple{Vector{ITensor}}, sequence=nothing; kwargs...)
  return visualize!(fig, tn[1], sequence; kwargs...)
end
function visualize!(
  fig, f::Function, tn::Tuple{Vararg{ITensor}}, sequence=nothing; kwargs...
)
  return visualize!(fig, tn, sequence; kwargs...)
end

function visualize(f::Union{Function,Type}, As...; kwargs...)
  # TODO: specialize on the function type. Also accept a general collection.
  return visualize(As...; kwargs...)
end

function visualize!(fig, f::Union{Function,Type}, As...; kwargs...)
  # TODO: specialize of the function type. Also accept a general collection.
  return visualize!(fig, As...; kwargs...)
end

function _visualize_sequence!(fig, tn, sequence, n; kwargs...)
  return error("Not implemented")
end

function sequence_labels(sequence, all_sequences, vertex_labels)
  traversal = sequence_traversal(sequence)
  labels_dict = contract_dict(vertex_labels, sequence, traversal)
  all_labels = [labels_dict[s] for s in all_sequences]
  return all_labels
end

function _graphplot(backend::Backend, graph; all_labels)
  return error("Not implemented for backend $backend.")
end

function visualize_sequence(sequence, vertex_labels)
  graph, all_sequences = tree_to_graph(sequence)
  all_labels = sequence_labels(sequence, all_sequences, vertex_labels)
  fig = _graphplot(Backend"Makie"(), graph; all_labels=all_labels)
  return fig
end

function default_sequence(tn::Vector{ITensor})
  N = length(tn)
  return foldl((x, y) -> [x, y], 1:N)
end

function visualize_sequence(
  f::Union{Function,Type}, tn::Vector{ITensor}, sequence::Nothing; kwargs...
)
  return visualize_sequence(f, tn, default_sequence(tn); kwargs...)
end

function visualize_sequence(
  f::Union{Function,Type}, tn::Vector{ITensor}, sequence=default_sequence(tn); kwargs...
)
  N = length(tn)

  # TODO: clean this up a bit
  vertex_labels_prefix = get(
    kwargs,
    :vertex_labels_prefix,
    default_vertex_labels_prefix(Backend("Makie"), MetaDiGraph(tn)),
  )
  vertex_labels = get(
    kwargs,
    :vertex_labels,
    default_vertex_labels(Backend(""), MetaDiGraph(tn), vertex_labels_prefix),
  )

  fig = visualize_sequence(sequence, vertex_labels)

  visualize!(fig[1, 2], tn; vertex_labels=vertex_labels, kwargs...)

  traversal = sequence_traversal(sequence)
  labels_sequence = contraction_sequence(vertex_labels, sequence, traversal)
  tn_sequence = contraction_sequence(tn, sequence, traversal)

  for n in 1:(length(tn_sequence) - 1)
    visualize!(fig[1, n + 2], tn_sequence[n]; vertex_labels=labels_sequence[n], kwargs...)
  end

  return fig
end

function visualize_sequence(
  f::Union{Function,Type}, tn::Tuple{Vector{ITensor}}, sequence; kwargs...
)
  return visualize_sequence(f, tn[1], sequence; kwargs...)
end

## is_kwarg(arg_or_kwarg::Symbol) = false
## is_kwarg(arg_or_kwarg::Expr) = (arg_or_kwarg.head == :parameters)
##
## function has_kwargs(args_kwargs::Vector)
##   isempty(args_kwargs) && return false
##   return is_kwarg(first(args_kwargs))
## end
##
## function get_kwargs(args_kwargs::Vector)
##   @assert has_kwargs(args_kwargs)
##   return first(args_kwargs)
## end
##
## function get_kwarg(kwargs::Expr, key::Symbol)
##   n = findfirst(kw -> kw.args[1] == :sequence, kwargs.args)
##   if !isnothing(n)
##     @assert kwargs.args[n].head == :kw
##     return esc(kwargs.args[n].args[2])
##   end
##   return nothing
## end
##
## function args_kwargs(ex::Vector)
##   kwargs = has_kwargs(ex) ? get_kwargs(ex) : :()
##   args = has_kwargs(ex) ? ex[2:end] : ex
##   return args, kwargs
## end
##
## function function_args_kwargs(ex::Symbol)
##   func = :identity
##   args = [ex]
##   kwargs = :()
##   iscollection = true
##   return func, args, kwargs, iscollection
## end
##
## function function_args_kwargs(ex::Expr)
##   if ex.head == :call
##     func = first(ex.args)
##     args, kwargs = args_kwargs(ex.args[2:end])
##     iscollection = true
##   elseif ex.head == :ref
##     #func, args, kwargs, iscollection = function_args_kwargs(Symbol(ex.args))
##     func = :identity
##     args = [ex]
##     kwargs = :()
##     iscollection = false
##   else
##     dump(ex)
##     error("Visualizing expression $ex not supported right now.")
##   end
##   return func, args, kwargs, iscollection
## end
##
## expr_to_string(s::Symbol) = String(s)
## expr_to_string(ex::Expr) = String(repr(ex))[3:(end - 1)]
##
## # Take the symbols of the arguments and output
## # the labels if there are multiple inputs or
## # the prefix for the labels if there is only
## # one input.
## function vertex_labels_kwargs(args, iscollection)
##   if iscollection && isone(length(args))
##     vertex_labels_kw = :vertex_labels_prefix
##     vertex_labels_arg = string(only(args))
##   else
##     vertex_labels_kw = :vertex_labels
##     vertex_labels_arg = string.(args)
##   end
##   return vertex_labels_kw, vertex_labels_arg
## end
##
## function func_args_sequence_kwargs(ex, vis_kwargs...)
##   func, args, kwargs, iscollection = function_args_kwargs(ex)
##   sequence = get_kwarg(kwargs, :sequence)
##   vertex_labels_kw, vertex_labels_arg = vertex_labels_kwargs(args, iscollection)
##   # Merge labels kwarg with kwargs
##   vis_kwargs_dict = Dict([
##     vis_kwarg.args[1] => vis_kwarg.args[2] for vis_kwarg in vis_kwargs
##   ])
##   vertex_labels_kwarg_dict = Dict(vertex_labels_kw => vertex_labels_arg)
##   merged_kwargs_dict = merge(vertex_labels_kwarg_dict, vis_kwargs_dict)
##   merged_kwargs_expr = [:($k = $v) for (k, v) in pairs(merged_kwargs_dict)]
##   return func, esc.(args), sequence, esc.(merged_kwargs_expr)
## end
##
## function visualize_expr(vis_func, ex::Union{Symbol,Expr}, vis_kwargs::Expr...)
##   func, args, sequence, kwargs = func_args_sequence_kwargs(ex, vis_kwargs...)
##   e = quote
##     $(vis_func)($(func), ($(args...),), $(sequence); $(kwargs...))
##   end
##   return e
## end
##
## function visualize_expr!(fig, vis_func!, ex::Union{Symbol,Expr}, vis_kwargs::Expr...)
##   func, args, sequence, kwargs = func_args_sequence_kwargs(ex, vis_kwargs...)
##   e = quote
##     $(vis_func!)($(esc(fig)), $(func), ($(args...),), $(sequence); $(kwargs...))
##   end
##   return e
## end
##
## """
##     @visualize
##
## Visualize a contraction of ITensors, returning the result of the contraction.
##
## The contraction should be written in terms of a series of ITensors contracted with `*`.
##
## # Examples
## ```julia
## using ITensors
## using ITensorVisualization
##
## i = Index(2, "index_i")
## j = Index(10, "index_j")
## k = Index(40, "index_k")
## l = Index(40, "index_l")
## m = Index(40, "index_m")
## A = randomITensor(i, j, k)
## B = randomITensor(i, j, l, m)
## C = randomITensor(k, l)
##
## # Contract the tensors over the common indices
## # and visualize the results
## ABC = @visualize A * B * C
##
## AB = @visualize A * B
## # Use readline() to pause between plots
## readline()
## ABC = @visualize AB * C vertex_labels = ["A*B", "C"]
## readline()
##
## # Save the results to figures for viewing later
## AB = @visualize fig1 A * B
## ABC = @visualize fig2 AB * C vertex_labels = ["A*B", "C"]
##
## display(fig1)
## readline()
## display(fig2)
## readline()
## ```
##
## # Keyword arguments:
## - `vertex_labels`: Custom tensor labels to display on the vertices of the digram. If not specified, they are determined automatically from the input to the macro.
## - `edge_labels=IndexLabels()`: A list of the edge labels or an `AbstractEdgeLabels` object specifying how they should be made.
## - `arrow_show`: Whether or not to show arrows on the edges.
## """
## macro visualize(fig::Symbol, ex::Symbol, kwargs::Expr...)
##   e = quote
##     $(esc(fig)) = $(visualize_expr(visualize, ex, kwargs...))
##     $(esc(ex))
##   end
##   return e
## end
##
## macro visualize!(fig, ex::Symbol, kwargs::Expr...)
##   e = quote
##     $(visualize_expr!(fig, visualize!, ex, kwargs...))
##     $(esc(ex))
##   end
##   return e
## end
##
## macro visualize(ex::Symbol)
##   e = quote
##     display($(visualize_expr(visualize, ex)))
##     $(esc(ex))
##   end
##   return e
## end
##
## macro visualize(ex_or_fig::Symbol, ex_or_kwarg::Expr, last_kwargs::Expr...)
##   if ex_or_kwarg.head == :(=)
##     # The second input is a keyword argument which means that the
##     # first input is the collection to visualize (no figure output binding specified)
##     ex = ex_or_fig
##     kwargs = (ex_or_kwarg, last_kwargs...)
##     e = quote
##       display($(visualize_expr(visualize, ex, kwargs...)))
##       $(esc(ex))
##     end
##   else
##     # The second input is not a keyword argument which means that the
##     # first input is the binding for the figure output, the second is the expression
##     # to visualize
##     fig = ex_or_fig
##     ex = ex_or_kwarg
##     kwargs = last_kwargs
##     e = quote
##       $(esc(fig)) = $(visualize_expr(visualize, ex, kwargs...))
##       $(esc(ex))
##     end
##   end
##   return e
## end
##
## macro visualize!(fig, ex::Expr, kwargs::Expr...)
##   e = quote
##     $(visualize_expr!(fig, visualize!, ex, kwargs...))
##     $(esc(ex))
##   end
##   return e
## end
##
## macro visualize(ex::Expr, kwargs::Expr...)
##   e = quote
##     display($(visualize_expr(visualize, ex, kwargs...)))
##     $(esc(ex))
##   end
##   return e
## end
##
## macro visualize_noeval(ex::Symbol, kwargs::Expr...)
##   e = quote
##     $(visualize_expr(visualize, ex, kwargs...))
##   end
##   return e
## end
##
## macro visualize_noeval(ex::Expr, kwargs::Expr...)
##   e = quote
##     $(visualize_expr(visualize, ex, kwargs...))
##   end
##   return e
## end
##
## macro visualize_noeval!(fig, ex::Symbol, kwargs::Expr...)
##   e = quote
##     $(visualize_expr!(fig, visualize!, ex, kwargs...))
##   end
##   return e
## end
##
## macro visualize_noeval!(fig, ex::Expr, kwargs::Expr...)
##   e = quote
##     $(visualize_expr!(fig, visualize!, ex, kwargs...))
##   end
##   return e
## end
##
## macro visualize_sequence(fig::Symbol, ex::Expr, kwargs::Expr...)
##   e = quote
##     $(esc(fig)) = $(visualize_expr(visualize_sequence, ex, kwargs...))
##     $(esc(ex))
##   end
##   return e
## end
##
## macro visualize_sequence(ex::Expr, kwargs::Expr...)
##   e = quote
##     display($(visualize_expr(visualize_sequence, ex, kwargs...)))
##     $(esc(ex))
##   end
##   return e
## end
##
## macro visualize_sequence_noeval(ex::Expr, kwargs::Expr...)
##   e = quote
##     $(visualize_expr(visualize_sequence, ex, kwargs...))
##   end
##   return e
## end


text:
module TestPlugin

using Test
using Comonicon.AST
using Comonicon.JuliaExpr
using Comonicon.JuliaExpr: emit, emit_body, emit_norm_body, emit_dash_body

function called()
    @test false
end

cmd = Entry(;
    version = v"1.2.0",
    root = NodeCommand(;
        name = "node",
        subcmds = Dict(
            "cmd1" => LeafCommand(; fn = called, name = "cmd1"),
            "cmd2" => LeafCommand(; fn = called, name = "cmd2"),
        ),
    ),
)

eval(emit(cmd, JuliaExpr.Configs(plugin = true)))
end

using Test
using Comonicon
withenv("PATH" => "$(pkgdir(Comonicon, "test", "codegen", "julia")):$(ENV["PATH"])") do
    Sys.which("node-cmd3")
    @test TestPlugin.command_main(["cmd3"]) == 0
end


text:
together_rules = [
    # TODO: This rule needs to be defined in three different ways to account for
    # the possible orderings of the numerators and denominators (the middle rule
    # covers two cases because it is associative). It would be nice to replace
    # them all with a single rule.
    @acrule(~a * (~b)^-1 + ~c * (~d)^-1 => (~a*~d + ~b*~c) / (~b*~d))
    @acrule((~b)^-1 * ~a + ~c * (~d)^-1 => (~a*~d + ~b*~c) / (~b*~d))
    @acrule((~b)^-1 * ~a + (~d)^-1 * ~c => (~a*~d + ~b*~c) / (~b*~d))
]
together_rewriter =
    SymbolicUtils.Postwalk(
        SymbolicUtils.PassThrough(
            SymbolicUtils.Chain(together_rules)
        )
    )

together(ex; kwargs...) = simplify(ex; rewriter=together_rewriter, kwargs...)