text:
using Revise
using AlphaZero
using ProgressLogging
using TensorBoardLogger
using Logging
using ParameterSchedulers: Scheduler, Cos
import Flux
import FileIO
import Random
import CUDA

import LinearAlgebra
LinearAlgebra.BLAS.set_num_threads(1)

include("../mu_game_wrapper.jl")
include("../network.jl")
include("../alphazerolike.jl")
include("../trace.jl")
include("../play.jl")
include("../training.jl")
include("../learning.jl")
include("../benchmark.jl")
include("../probe-games.jl")

device = Flux.cpu
# device = Flux.gpu
CUDA.allowscalar(true)

gspec = Examples.games["connect-four"]
# gspec = ProbeGames.games["a1r1"]
# gspec = ProbeGames.games["a1o2r2"]
# gspec = ProbeGames.games["a1r1s2"] # representation collapses states into the same one
# gspec = ProbeGames.games["a2r2"]
# gspec = ProbeGames.games["twoplayer"]
# gspec = ProbeGames.games["simplertictactoe"]

n=8
self_play = (;
  sim=SimParams(
    num_games=5,
    num_workers=n,
    batch_size=n,
    use_gpu=false,
    # todevice=Flux.cpu, #?
    reset_every=4, #not used, mcts resets everytime
    flip_probability=0.,
    alternate_colors=false),
  mcts = MctsParams(
    num_iters_per_turn=128, #1000 benchmark
    cpuct=1.25,
    temperature=ConstSchedule(1.0),
    dirichlet_noise_ϵ=0.25,
    dirichlet_noise_α=0.1))

# # for mcts rollout generation
# self_play = (;
#   sim=SimParams(
#     num_games=10_000,
#     num_workers=n,
#     batch_size=n,
#     use_gpu=false,
#     reset_every=4, #not used, mcts resets everytime
#     flip_probability=0.,
#     alternate_colors=false),
#   mcts = MctsParams(
#     num_iters_per_turn=500, #1000 benchmark
#     cpuct=2.5,
#     temperature=ConstSchedule(1.0),
#     dirichlet_noise_ϵ=0.25,
#     dirichlet_noise_α=0.1))

arena = (;
  sim=SimParams(
    num_games=50,
    num_workers=n,
    batch_size=n,
    use_gpu=false,
    reset_every=1,
    flip_probability=0., #0.5
    alternate_colors=true),
  mcts = MctsParams(
    self_play.mcts,
    temperature=ConstSchedule(0.3),
    dirichlet_noise_ϵ=0.1),
  update_threshold=0.0,
  device=Flux.cpu)

learning_params = (;
  num_unroll_steps=10, #if =0, g is not learning, muzero-general=20
  td_steps=50, # with max length=9, always go till the end of the game, rootvalues don't count
  discount=0.997,
  #// value_loss_weight = 0.25, #TODO
  l2_regularization=1f-4, #Float32
  #// l2_regularization=0f0, #Float32
  loss_computation_batch_size=512,
  batches_per_checkpoint=10,
  num_checkpoints=1,
  opt=Scheduler(
    Cos(λ0=1e-3, λ1=1e-5, period=10000), # cosine annealing, google 2e4, generat doesn't use any
    Flux.ADAM()
  ),
  device=Flux.gpu
)

benchmark_sim = SimParams(
    num_games=400,
    num_workers=1,
    batch_size=4,
    use_gpu=false,
    reset_every=1,
    flip_probability=0.0, #0.5
    alternate_colors=true)

bench_mcts = MctsParams(
  num_iters_per_turn=500, #1000 benchmark
  cpuct=1.25,
  temperature=ConstSchedule(0.3),
  dirichlet_noise_ϵ=0.25,
  dirichlet_noise_α=0.1)

benchmark = (;
  vanilla_mcts=Benchmark.Duel(
    Mu(arena.mcts),
    Benchmark.MctsRollouts(bench_mcts),
    benchmark_sim),
  # minmax_d5=Benchmark.Duel(
  #   Mu(arena.mcts),
  #   Benchmark.MinMaxTS(depth=5, amplify_rewards=true, τ=1.),
  #   benchmark_sim)
)

  #memory = CircularBuffer{Trace{GI.state_type(gspec)}}(1024)

  # μparams = MuParams(self_play, learning_params, arena, 6, 3000)
  μparams = MuParams(self_play, learning_params, arena, 10_000, 10_000)

  Random.seed!(2137)
  # env = MuEnv(gspec, μparams, αnns, experience=FileIO.load("memory_minmax.jld2", "mem"))
  # env = MuEnv(gspec, μparams, αnns, experience=FileIO.load("memory_minmax_d5_t025.jld2", "mem"))

# muzero-general params
μNetworkHP = MuNetworkHP(
  gspec,
  PredictionHP(hiddenstate_shape=64, width=64, depth_common=0,
    depth_vectorhead=0, depth_scalarhead=0, use_batch_norm=false, batch_norm_momentum=1.),
  DynamicsHP(hiddenstate_shape=64, width=64, depth_common=-1,
   depth_vectorhead=0, depth_scalarhead=0,  use_batch_norm=false, batch_norm_momentum=1.),
  RepresentationHP(hiddenstate_shape=64, width=0, depth=-1))

  # env = MuEnv(gspec, μparams, MuNetwork(μNetworkHP))
  # env = MuEnv(gspec, μparams, deepcopy(env.bestnns))
  env = MuEnv(gspec, μparams, MuNetwork(μNetworkHP), experience=FileIO.load("results/c4_5x4/memory_mctsrollout500.jld2", "mem"))

##
learning_step!(env)
##

  # learning_step
lp = env.params.learning_params
tr = MuTrainer(env.gspec, env.curnns|>lp.device|>Flux.trainmode!, env.memory, env.params.learning_params, env.params.learning_params.opt)
@timed samples = [sample_batch(tr.gspec, tr.memory, tr.hyper) for _ in 1:10] |> Flux.gpu

hyper = tr.hyper
nns = env.curnns|>Flux.gpu
(X, A_mask, As, Ps, Vs, Rs) = samples[1]
##
# function losses(nns, hyper, (X, A_mask, As, Ps, Vs, Rs))
  prediction, dynamics, representation = nns.f, nns.g, nns.h
  creg = hyper.l2_regularization
  Ksteps = hyper.num_unroll_steps

  # initial step, from the real observation
  Hiddenstate = forward(representation, X)
  P̂⁰, V̂⁰ = forward(prediction, Hiddenstate)
  P̂⁰ = normalize_p(P̂⁰, A_mask)
  # R̂⁰ = zero(V̂⁰)

  scale_initial = iszero(Ksteps) ? 1f0 : 0.5f0
  Lp = scale_initial * lossₚ(P̂⁰, Ps[:, 1, :]) # scale=1
  Lv = scale_initial * lossᵥ(V̂⁰, Vs[1:1, :])
  Lr = zero(Lv) # starts at next step (see MuZero paper appendix)

  scale_recurrent = iszero(Ksteps) ? nothing : 0.5f0 / Ksteps #? instead of constant scale, maybe 2^(-i+1)
  # recurrent inference
  # for k in 1:Ksteps
  k=1
    # targets are stored as follows: [A⁰¹ A¹² ...] [P⁰ P¹ ...] [V⁰ V¹ ...] but [R¹ R² ...]
    A = As[k, :]
##
# function forward(nn::DynamicsNetwork, hiddenstate, action)
  nn=dynamics; hiddenstate=Hiddenstate; action=A;
  action_one_hot = onehotbatch(action, GI.actions(gspec)) #! scalar indexing
  hiddenstate_action = vcat(hiddenstate, collect(action_one_hot)) #!CUDA ERROR
  c = nn.common(hiddenstate_action)
  r = nn.scalarhead(c)
  s⁺¹ = nn.vectorhead(c)
  return (r, s⁺¹)
# end

    R̂, Hiddenstate = forward(dynamics, Hiddenstate, A) #
    P̂, V̂ = forward(prediction, Hiddenstate) #? should flip V based on players
    # scale loss so that the overall weighting of the recurrent_inference (g,f nns)
    # is equal to that of the initial_inference (h,f nns)
    Lp += scale_recurrent * lossₚ(P̂, Ps[:, k+1, :]) #? @view
    Lv += scale_recurrent * lossᵥ(V̂, Vs[k+1:k+1, :])
    Lr += scale_recurrent * lossᵣ(R̂, Rs[k:k, :])
  # end
  Lreg = iszero(creg) ? zero(Lv) : creg * sum(sum(w.^2) for w in regularized_params(nns))
  L = Lp + Lv + Lr + Lreg # + Lr
  # L = Lp + Lreg # + Lr
  Zygote.@ignore @info "Loss" loss_total=L loss_policy=Lp loss_value=Lv loss_reward=Lr loss_reg_params=Lreg relative_entropy=Lp-Flux.Losses.crossentropy(Ps, Ps) #? check if compute means inside logger is avaliable
  return (L, Lp, Lv, Lr, Lreg)
# end
##

#TODO check with MCTS, MinMax
#Revise.jl, tensorboard_logger

# # checking Losses
# zero(Ps)
# one.(Ps)
# Flux.Losses.mse(zero(Vs), one.(Vs), agg=x->mean(∇_scale .* x)/sum(∇_scale))
# Flux.Losses.mse(zero(Vs), one.(Vs), agg=x->mean(x .* ∇_scale))
# Flux.Losses.mse(zero(Vs), one.(Vs), agg=x->sum(∇_scale .* x) / sum(∇_scale))
# Flux.Losses.mse(zero(Vs), one.(Vs))
# sum(∇_scale)
# one.(Vs)
# Vs
# one.(Ps) .* ∇_scale

# Flux.Losses.crossentropy(one.(Ps).*ℯ, one.(Ps), agg=x->mean(x,∇_scale))
# log.(one.(Ps) .* ℯ)
# log.(zero(Ps)+eps.(Ps))
# one.(Ps) .*
# one.(Ps) .* log.(one.(Ps).*ℯ)
# mean(.-sum(one.(Ps) .* log.(one.(Ps).*ℯ), dims=1), weights(∇_scale))
# ∇_scale .* 2

"Lenient comparison operator for `struct`, both mutable and immutable (type with \\eqsim)."
@generated function ≂(x, y)
  if !isempty(fieldnames(x)) && x == y
    mapreduce(n -> :(x.$n ≂ y.$n), (a,b)->:($a && $b), fieldnames(x))
  else
    :(x == y)
  end
end

# nns = deepcopy(tr.nns)
# nns == tr.nns
# nns ≂ tr.nns
# env.curnns ≂ env.bestnns
# nns = deepcopy(env.bestnns)
# nns ≂ env.bestnns
# fieldnames(typeof(nns))
# typeof(nns) == typeof(env.bestnns)
# nns.f.common

(X, A_mask, As, Ps, Vs, Rs) = sample_batch(env.gspec, env.memory, env.params.learning_params)
@enter losses(env.curnns, learning_params, (X, A_mask, As, Ps, Vs, Rs))
As
A = As[1:1,:]
Vs

env.curnns = deepcopy(env.bestnns)
@assert env.curnns ≂ env.bestnns    # all their fields are the same
@assert !(env.curnns == env.bestnns) # they have not the same memory adress

tr = MuTrainer(env.gspec, env.curnns, env.memory, env.params.learning_params, Flux.ADAM())
@assert env.curnns ≂ tr.nns
@assert env.curnns == tr.nns # same memory adress (shallow copy)

@timed update_weights!(tr, tr.hyper.batches_per_checkpoint)
@assert env.curnns == tr.nns
@assert !(env.bestnns ≂ tr.nns) # after update! tr.nns and env.curnns changed
@assert !(env.curnns ≂ env.bestnns)

learning_step!(env)
train!(env)

# gnerator with rand
function foo(x)
  y = x+1
  z = x+1
  return (; x,y,z)
end
rands = (rand() for _ in 1:5)
foos = (foo(r) for r in rands)
X = [f.x for f in foos] # Flux.batch(f.x for f in foos)
Y = [f.y for f in foos]
Z = [f.z for f in foos]
X[1], Y[1], Z[1]


text:
"""
    drawbasinsR2(f, points [; hasescaped, maxiterations, tolerance])

Return the drawing of the attracting or parabolic basins of given periodic
points of a function \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$,
in a rectangular region in \$\\mathbb{R}^2\$.

The attracting (parabolic) basin of a periodic point \$(x_0,y_0)\$ of period
\$k\$ under \$f\$ is defined as
\$\\mathcal{B}(x_0,y_0)=\\{(x,y)\\in\\mathbb{R}^2\\,|\\,|(f^k)^n(x,y)|\\rightarrow (x_0,y_0)\\,n\\rightarrow\\infty\\}\$

#### Arguments
- `f::Function`: A function \$f:\\mathbb{R}^2\\rightarrow\\mathbb{R}^2\$.
- `points::AbstractArray`: List of points \$(x,y)\$ assumning to be peridic attrating or parabolic.
- `hasescaped::Function`: A boolean function to check if the iterations has escaped (useful for the infinity basin).
- `maxiterations::Integer`: Maximum number of iterations to check.
- `tolerance::Real`: Small number to check closeness to given points.
"""
function drawbasinsR2(f::Function, points::AbstractArray;
    hasescaped::Function=(x,y)->false, maxiterations::Int=100, tolerance::Real=1.0e-6)

    # Veryfying if graphics backend supports functions
    SDDGraphics.supported(:drawpixel)

    # Verifying functions
    @assert typeof(f(1., 1.)) <: Tuple{Real,Real}
    @assert typeof(hasescaped(1., 1.)) <: Bool

    SDDGraphics.newdrawing()

    num_points = length(points)
    num_colors = num_points+1
    SDDGraphics.updatecolorarray(num_colors)

    @sweeprectregion SDDGraphics.xlims() SDDGraphics.ylims() SDDGraphics.canvassize() begin
        xn,yn = x,y
        escapetime = maxiterations
        index = num_colors
        for n in 1:maxiterations
            for k in 1:num_points
                if (points[k][1]-xn)^2 + (points[k][2]-yn)^2 < tolerance
                    index = k
                    break
                end
            end
            if index < num_colors
                escapetime = n
                break
            end
            if hasescaped(xn,yn)
                escapetime = n
                index = num_colors
                break
            end # if hasescaped
            xn,yn = f(xn,yn)
        end # for n maxiterations
        SDDGraphics.colorinterpolationbg(escapetime/maxiterations, index)
        SDDGraphics.drawpixel(i,j)
    end # Implemented algorithm

    SDDGraphics.drawing()
end # function drawtrappedpointsR2

"""
    drawbasinsC(f, points [; hasescaped, maxiterations, tolerance])

Return the drawing of the attracting or parabolic basins of given periodic
points of a function \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$,
in a rectangular region in \$\\mathbb{C}\$.

The attracting (parabolic) basin of a periodic point \$z_0\$ of period
\$k\$ under \$f\$ is defined as
\$\\mathcal{B}(z_0)=\\{z\\in\\mathbb{C}\\,|\\,|(f^k)^n(z)|\\rightarrow z_0\\,n\\rightarrow\\infty\\}\$

#### Arguments
- `f::Function`: A function \$f:\\mathbb{C}\\rightarrow\\mathbb{C}\$.
- `points::AbstractArray`: List of complex numbers \$z\$ assumning to be peridic attrating or parabolic.
- `hasescaped::Function`: A boolean function to check if the iterations has escaped (useful for the infinity basin).
- `maxiterations::Integer`: Maximum number of iterations to check.
- `tolerance::Real`: Small number to check closeness to given points.
"""
function drawbasinsC(f::Function, points::AbstractArray;
    hasescaped::Function=z->false, maxiterations::Int=100, tolerance::Real=1.0e-6)

    # Veryfying if graphics backend supports functions
    SDDGraphics.supported(:drawpixel)

    # Verifying functions
    @assert typeof(f(.1im)) <: Number
    @assert typeof(hasescaped(.1im)) <: Bool

    SDDGraphics.newdrawing()

    num_points = length(points)
    num_colors = num_points+1
    SDDGraphics.updatecolorarray(num_colors)

    @sweeprectregion SDDGraphics.xlims() SDDGraphics.ylims() SDDGraphics.canvassize() begin
        zn = complex(x,y)
        escapetime = maxiterations
        index = num_colors
        for n in 1:maxiterations
            for k in 1:num_points
                if abs2(points[k]-zn) < tolerance
                    index = k
                    break
                end
            end
            if index < num_colors
                escapetime = n
                break
            end
            if hasescaped(zn)
                escapetime = n
                index = num_colors
                break
            end # if hasescaped
            zn = f(zn)
        end # for n maxiterations
        SDDGraphics.colorinterpolationbg(escapetime/maxiterations, index)
        SDDGraphics.drawpixel(i,j)
    end # Implemented algorithm

    SDDGraphics.drawing()
end # function drawtrappedpointsR2

ext: jl


text:
# ------------------------------------------------------------------
# Copyright (c) 2017, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    Ellipsoidal(semiaxes, angles)

A distance defined by an ellipsoid with given `semiaxes` and rotation `angles`.

- For 2D ellipsoids, there are two semiaxes and one rotation angle.
- For 3D ellipsoids, there are three semiaxes and three rotation angles.

## Examples

2D ellipsoid making 45ᵒ with the horizontal axis:

```julia
julia> Ellipsoidal([1.0,0.5], [π/2])
```

3D ellipsoid rotated by 45ᵒ in the xy plane:

```julia
julia> Ellipsoidal([1.0,0.5,0.5], [π/2,0.0,0.0])
```
"""
struct Ellipsoidal{N,T} <: Metric
  dist::Mahalanobis{T}

  function Ellipsoidal{N,T}(semiaxes, angles) where {N,T<:Real}
    @assert length(semiaxes) == N "number of semiaxes must match spatial dimension"
    @assert all(semiaxes .> zero(T)) "semiaxes must be positive"
    @assert N ∈ [2,3] "dimension must be either 2 or 3"

    # scaling matrix
    Λ = Diagonal(one(T)./semiaxes.^2)

    # rotation matrix
    if N == 2
      θ = angles[1]

      cosθ = cos(θ)
      sinθ = sin(θ)

      P = [cosθ -sinθ
           sinθ  cosθ]
    end
    if N == 3
      θxy, θyz, θzx = angles

      cosxy = cos(θxy)
      sinxy = sin(θxy)
      cosyz = cos(θyz)
      sinyz = sin(θyz)
      coszx = cos(θzx)
      sinzx = sin(θzx)

      _1 = one(T)
      _0 = zero(T)

      Rxy = [cosxy -sinxy _0
             sinxy  cosxy _0
             _0     _0 _1]

      Ryz = [_1    _0     _0
             _0 cosyz -sinyz
             _0 sinyz  cosyz]

      Rzx = [ coszx _0 sinzx
             _0 _1    _0
             -sinzx _0 coszx]

      P = Rzx*Ryz*Rxy
    end

    # ellipsoid matrix
    Q = P*Λ*P'

    new(Mahalanobis(Q))
  end
end

Ellipsoidal(semiaxes::AbstractVector{T}, angles::AbstractVector{T}) where {T<:Real} =
  Ellipsoidal{length(semiaxes),T}(semiaxes, angles)

function Distances.evaluate(dist::Ellipsoidal{N,T}, a::AbstractVector, b::AbstractVector) where {N,T<:Real}
  evaluate(dist.dist, a, b)
end


text:
function k3p_cf(f, η1, η2, η3, ξ)

	(ξ^3)*((η1)^2)*(η2)*(

	f((1-ξ, ξ-ξ*η1+ξ*η1*η2), (1-(ξ-ξ*η1*η2*η3), ξ-ξ*η1)) +

	f((1-(ξ-ξ*η1*η2*η3), ξ-ξ*η1), (1-ξ, ξ-ξ*η1+ξ*η1*η2)) +

	f((1-ξ, ξ*η1*(1-η2+η2*η3)), (1-(ξ-ξ*η1*η2), ξ*η1*(1-η2))) +

	f((1-(ξ-ξ*η1*η2), ξ*η1*(1-η2)), (1-ξ, ξ*η1*(1-η2+η2*η3))) +

	f((1-(ξ-ξ*η1*η2*η3), ξ*η1*(1-η2*η3)), (1-ξ, ξ*η1*(1-η2))) +

	f((1-ξ, ξ*η1*(1-η2)), (1-(ξ-ξ*η1*η2*η3), ξ*η1*(1-η2*η3)))

	)
end

function k3p_ce(f, η1, η2, η3, ξ)

	(ξ^3)*((η1)^2) * f((1-ξ, ξ*η1*η3), (1-ξ*(1-η1*η2), ξ*η1*(1-η2))) +

	(ξ^3)*((η1)^2)*(η2)*(

	f((1-ξ, ξ*η1), (1-ξ*(1-η1*η2*η3), ξ*η1*η2*(1-η3))) +

	f((1-ξ*(1-η1*η2), ξ*η1*(1-η2)), (1-ξ, ξ*η1*η2*η3)) +

	f((1-ξ*(1-η1*η2*η3), ξ*η1*η2*(1-η3)), (1-ξ, ξ*η1)) +

	f((1-ξ*(1-η1*η2*η3), ξ*η1*(1-η2*η3)), (1-ξ, ξ*η1*η2))

	)
end

function k3p_cv(f,η1, η2, η3, ξ)

	(ξ^3)*η2*(

	f((1-ξ, ξ*η1), (1-ξ*η2, ξ*η2*η3)) +  f((1-ξ*η2, ξ*η2*η3), (1-ξ, ξ*η1))

	)
end

function k3p_pd(f, η1, η2, η3, ξ)

	u = [η1*(1-η2), η2]		#jacobian of this transformation is (1-η2)
	v = [η3*(1-ξ), ξ]		#jacobian of this transformation is (1-ξ)

	return(f(u,v)*(1-η2)*(1-ξ))
end


text:
using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

@inline function exp_sid_m6_opt_rho2_22(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_sid_m6_opt_rho2_22!(A_copy)
end

@inline function exp_sid_m6_opt_rho2_22!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=9
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    memslots6 = similar(A,T)
    memslots7 = similar(A,T)
    memslots8 = similar(A,T)
    memslots9 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    # Computation order: Bb2 Ba2 B2 Ba3 Bb3 B3 Ba4 Bb4 B4 Ba5 Bb5 B5 Ba6 Bb6 B6 Ba7 Bb7 B7 T2k9
    # Computing Bb2 = x*I+x*A
    coeff1=0.013941257411227636
    coeff2=1.031053526686218
    memslots2 .= coeff2.*memslots1
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Ba2 = x*I+x*A
    coeff1=0.013941257411227636
    coeff2=1.031053526686218
    memslots3 .= coeff2.*memslots1
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B2 with operation: mult
    mul!(memslots4,memslots3,memslots2)
    # Deallocating Ba2 in slot 3
    # Deallocating Bb2 in slot 2
    # Computing Ba3 = x*I+x*A+x*B2
    coeff1=0.0027676995071721643
    coeff2=0.9956823779291472
    coeff3=0.018699401264499645
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb3 = x*I+x*A+x*B2
    coeff1=0.0014050828079002
    coeff2=0.0025755913013428494
    coeff3=0.9956033412300013
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B3 with operation: mult
    mul!(memslots5,memslots2,memslots3)
    # Deallocating Ba3 in slot 2
    # Deallocating Bb3 in slot 3
    # Computing Ba4 = x*I+x*A+x*B2+x*B3
    coeff1=0.00010080993148845571
    coeff2=1.000386931651619
    coeff3=0.014565003723968966
    coeff4=-0.0005245897678761895
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb4 = x*I+x*A+x*B2+x*B3
    coeff1=0.00013960932977845965
    coeff2=0.00027073100561329446
    coeff3=0.00011767015392433715
    coeff4=0.9999937753420692
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B4 with operation: mult
    mul!(memslots6,memslots2,memslots3)
    # Deallocating Ba4 in slot 2
    # Deallocating Bb4 in slot 3
    # Computing Ba5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=1.0034850244473076e-5
    coeff2=6.234107699654157e-5
    coeff3=-0.00010709695216530405
    coeff4=0.0004429662682755111
    coeff5=0.9999999359403395
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb5 = x*I+x*A+x*B2+x*B3+x*B4
    coeff1=-3.420829259855909e-5
    coeff2=1.2724236347474151e-5
    coeff3=7.614113854507912e-7
    coeff4=4.506303286916229e-8
    coeff5=3.3732271485767847e-9
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B5 with operation: mult
    mul!(memslots7,memslots2,memslots3)
    # Deallocating Ba5 in slot 2
    # Deallocating Bb5 in slot 3
    # Computing Ba6 = x*I+x*A+x*B2+x*B3+x*B4+x*B5
    coeff1=0.00024723298310785875
    coeff2=2.9161339295593187
    coeff3=0.19294236215103988
    coeff4=0.012746128809542087
    coeff5=0.001070557957948411
    coeff6=0.91478561889395
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb6 = x*I+x*A+x*B2+x*B3+x*B4+x*B5
    coeff1=0.0061957449875606665
    coeff2=0.022600844232367568
    coeff3=0.035261063716355324
    coeff4=0.00401546253180321
    coeff5=-0.0005271750223542335
    coeff6=1.0628031901019233
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B6 with operation: mult
    mul!(memslots8,memslots2,memslots3)
    # Deallocating Ba6 in slot 2
    # Deallocating Bb6 in slot 3
    # Computing Ba7 = x*I+x*A+x*B2+x*B3+x*B4+x*B5+x*B6
    coeff1=9.183179455916261e-5
    coeff2=8.291209653793697
    coeff3=1.7562193397284078
    coeff4=0.17985908646581833
    coeff5=0.020190024440301577
    coeff6=249.89674278520076
    coeff7=1.0634238665988902
    memslots2 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7 .+ coeff7.*memslots8
    mul!(memslots2,true,I*coeff1,true,true)
    # Computing Bb7 = x*I+x*A+x*B2+x*B3+x*B4+x*B5+x*B6
    coeff1=-0.00015133350603989786
    coeff2=0.028231887322807712
    coeff3=0.012411739707662116
    coeff4=0.0004764037694311965
    coeff5=0.00018495252800435267
    coeff6=1.05795634102314
    coeff7=0.000721682415651672
    memslots3 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7 .+ coeff7.*memslots8
    mul!(memslots3,true,I*coeff1,true,true)
    # Computing B7 with operation: mult
    mul!(memslots9,memslots2,memslots3)
    # Deallocating Ba7 in slot 2
    # Deallocating Bb7 in slot 3
    # Computing T2k9 = x*I+x*A+x*B2+x*B3+x*B4+x*B5+x*B6+x*B7
    coeff1=0.9999550769133566
    coeff2=0.9942510551400959
    coeff3=0.23154125334090794
    coeff4=-0.011469528624315358
    coeff5=-0.0005156243401819641
    coeff6=-0.00015881528950764778
    coeff7=0.022922533875140892
    coeff8=1.0576256651496891
    # Smart lincomb recycle A
    memslots1 .= coeff2.*memslots1 .+ coeff3.*memslots4 .+ coeff4.*memslots5 .+ coeff5.*memslots6 .+ coeff6.*memslots7 .+ coeff7.*memslots8 .+ coeff8.*memslots9
    mul!(memslots1,true,I*coeff1,true,true)
    # Deallocating B2 in slot 4
    # Deallocating B3 in slot 5
    # Deallocating B4 in slot 6
    # Deallocating B5 in slot 7
    # Deallocating B6 in slot 8
    # Deallocating B7 in slot 9
    return memslots1 # Returning T2k9
end

max_stars_repo_name: matrixfunctions/GraphMatFunData

max_issues_repo_name: matrixfunctions/GraphMatFunData

max_forks_repo_name: matrixfunctions/GraphMatFunData


text:
function load_case(case, number_of_peers)
    return load_data(date, h_predict, h_control, h_start, h_end, market_flag, number_of_peers, n_market, case);
end

function load_data(date, h_predict, h_control, h_start, h_end, market_flag, n_peers, n_market, case)
    Input_df = CSV.read("data/200124_datafile_all_details_right_timestamp.csv");
    Flow_df = CSV.read("results/$(date)_results_$(h_predict)_$(h_control)_$(h_start)-$(h_end)"*
                    "_$(market_flag)_$(n_peers)_$(n_market)_$(case).csv");
    # merge on data stamp
    Data_df = leftjoin(Flow_df, Input_df, on=[:month, :day, :hour], makeunique=false,
                        indicator=nothing, validate=(false, false));
    return Data_df
end

function bar_PV(Data_df, date, peer)
    # bars PV
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] groupedbar(
        [:PV_PM :PV_HP :PV_GR :PV_B :PV_DE], color =[:green :firebrick :grey :purple :orange],
        xticks=0:6:24,
        annotations=(2,12, text("$(Dates.month(date))/$(Dates.day(date))", 10)), legendfontsize=6,
        label=["PV_PM" "PV_HP" "PV_GR" "PV_B" "PV_DE"], legend=false, bar_position = :stack, alpha=0.8);

    # bars B
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] groupedbar!(
        (-1) .*[:B_PM :B_HP :B_DE], color =[:green :firebrick :orange],
        xticks=0:6:24,
        annotations=(2,12, text("$(Dates.month(date))/$(Dates.day(date))", 10)), legendfontsize=6,
        label=["B_PM" "B_HP" "B_DE"], legend=false, bar_position = :stack, alpha=0.8);

    # battery cap
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] plot!(
        [:SOC_B], color =[:purple], xticks=0:6:24,
        label=["SOC_b"], legend=false, linewidth= 2.0);

    # PV generation
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] plot!(
        [:PV_generation], xticks=0:6:24, color =[:gold],
        label=["g_e"], legend=false, linewidth= 2.0);

    plot!(0:23, ones(24,1)*13.5, linestyle=:dot, linewidth= 2, color=:purple)
end

function bar_demand(Data_df, date, peer)
    # bars
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] groupedbar(
        [:PM_DE :B_DE :GR_DE :PV_DE ], color =[:green :purple :grey :gold],
        ylim=(0,2), yticks=0:0.5:2, xticks=0:6:24,
        annotations=(2,1.8, text("$(Dates.month(date))/$(Dates.day(date))", 10)), legendfontsize=6,
        label=["PM_DE" "B_DE" "GR_DE" "PV_DE"], legend=false, bar_position = :stack, alpha=0.8);

    # electricity demand
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] plot!(
        [:electkwh], ylim=(0,2), yticks=0:0.5:2, xticks=0:6:24, color =[:orange],
        label=["d_e"], legend=false, linewidth= 2.0);

    plot!(0:23, ones(24,1)*13.5, linestyle=:dash, color=:purple)
end

function bar_heat(Data_df, date, peer)
    # hot water demand
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] plot(
        [:hotwaterkwh]+[:heatingkwh], fillrange=[:heatingkwh], ylim=(0,15), yticks=0:5:15, xticks=0:6:24,
        color =[:steelblue], label=["d_fh"], legend=false, alpha=0.2);
    # heat demand
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] plot!(
        [:heatingkwh], fillrange=zeros(24), ylim=(0,15), yticks=0:5:15, xticks=0:6:24,
        color =[:firebrick], label=["d_fh"], legend=false, alpha=0.2);
    # bars
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] groupedbar!(
        [:PM_HP :B_HP :GR_HP :PV_HP], color =[:green :purple :grey :gold],
        ylim=(0,15), yticks=0:5:15, xticks=0:6:24,
        annotations=(2, 14, text("$(Dates.month(date))/$(Dates.day(date))", 10)), legendfontsize=6,
        label=["PM_DE" "B_DE" "GR_DE" "PV_DE"], legend=false, bar_position = :stack, alpha=0.8);

    plot!(0:23, ones(24,1)*3, linestyle=:dot, linewidth= 2, color=:grey)
end

function bar_comfort(Data_df, date, peer)
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] groupedbar(
        ([:HP_FH :HP_HW].>0.005)*200, color =[:firebrick :steelblue],
        ylim=(19,23), yticks=20:1:22, xticks=0:6:24,
        annotations=(2, 22.5, text("$(Dates.month(date))/$(Dates.day(date))", 10)), legendfontsize=6,
        label=["Mod_fh" "Mod_hw"], legend=false, bar_position = :stack, alpha=0.15);

    # state-of-charge fh
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] plot!(
        [:Temp_FH], ylim=(19,23), yticks=20:1:22, xticks=0:6:24, color =[:firebrick],
        label=["T_fh"], legend=false, linewidth= 2.0);

    plot!(0:23, ones(24,1)*20, linestyle=:dot, linewidth= 2, color=:firebrick)
    plot!(0:23, ones(24,1)*22, linestyle=:dot, linewidth= 2, color=:firebrick)

    plt = twinx()
    # state-of-charge hw
    @df Data_df[(Data_df[:,:day].==Dates.day(date)) .&(Data_df[:,:month].==Dates.month(date)) .&(Data_df[:,:Peer].==peer),:] plot!(
        plt, [:Vol_HW], ylim=(10,190), yticks=20:40:180, xticks=0:6:24, color =[:steelblue],
        label=["V_fh"], legend=false, linewidth= 2.0, linestyle=:dash, grid=false);

    plot!(plt, 0:23, ones(24,1)*20, linestyle=:dot, linewidth= 2, color=:steelblue)
    plot!(plt, 0:23, ones(24,1)*180, linestyle=:dot, linewidth= 2, color=:steelblue)
end

function bar_row(plotfunction, Data_df, date, peer, length)
    plot([plotfunction(Data_df, (date+Day(i-1)), peer)
            for i in 1:length]..., layout=(1,length), size=(300*length,200));
end

function calc_energy_use(Data_df, n_peers)
    return convert.(Int64, round.([ sum(sum(groupby(Data_df, "Peer")[i][!,j] for j in ["PV_DE", "B_DE", "GR_DE", "HP_FH", "HP_HW", "PM_DE"])) for i in 1:n_peers], digits=0))
end

function peer_cons(n_peers)
    multiplier = [100, 100]
    if n_peers >= 2
        for i in 1:(n_peers-2)
            append!(multiplier, 0)
        end
    end
    return multiplier
end

function calc_self_consumption(Data_df, n_peers)
    return convert.(Int64, round.(peer_cons(n_peers) .*[1 - (sum(sum(groupby(Data_df, "Peer")[i][!,j] for j in ["PV_GR", "PV_PM", "B_PM"])) / sum(groupby(Data_df, "Peer")[i][!,"PV_generation"])) for i in 1:n_peers], digits=0))
end

function calc_self_sufficiency(Data_df, n_peers)
    return convert.(Int64, round.(100 .*(1 .- [sum(sum(groupby(Data_df, "Peer")[i][!,j] for j in ["GR_DE", "GR_HP", "PM_DE", "PM_HP"])) for i in 1:n_peers] ./ calc_energy_use(Data_df, n_peers)), digits=0))
end

function calc_comfort_violations(Data_df, n_peers)
    return convert.(Int64, round.([sum(sum(groupby(Data_df, "Peer")[i][!,j] for j in ["V_HW_plus", "V_HW_minus", "T_FH_plus", "T_FH_minus"])) for i in 1:n_peers], digits=0))
end

function calc_profit(Data_df, n_peers)
    return convert.(Int64, round.([sum(groupby(Data_df, "Peer")[i][!,"profits"]) for i in 1:n_peers]./100, digits=0))
end

function calc_levies(Data_df, case, n_peers)
    return hcat([collect(sum([9.87, 12.19,  12.19, 7.71, 2.05, 6.756, 2.667] .*Matrix(CSV.read("data/tariff_matrix_case_$(case).csv")) *Matrix(groupby(Data_df, "Peer")[i][!,13:27])', dims=2)./100) for i in 1:n_peers]...)
end

max_stars_repo_name: lilanger/PEERS

max_issues_repo_name: lilanger/PEERS

max_forks_repo_name: lilanger/PEERS


text:
#!/usr/bin/env julia
#
# Jerzy Wroczyński (nr. indeksu 250075)
# Obliczenia Naukowe, Laboratorium
# Lista 5.
#

include("utilities.jl")
using .utilities

include("blocksys.jl")
using .blocksys

include("generate-a-test.jl")
using .testdata

global l = 4

global n = parse(Int, ARGS[1])

let A, b
end

# whether to run with classic built-in algorithm
classic = "classic" in ARGS

# whether to run using a custom sparse matrix solution
custom = "custom" in ARGS

# whether to run improved algorithm with pivoting
pivot = "pivot" in ARGS

# whether to run using the LU decomposition
lu = "lu" in ARGS

println(stderr, "loading data...")
if n == 16
    # read the already pregenerated testfiles
    A = utilities.Amatrix("test-data/n16/A.txt", custom)
    b = utilities.bmatrix("test-data/n16/b.txt")
else
    pathA, pathb = testdata.generateTestData(n, l)
    A = utilities.Amatrix(pathA, custom)
    b = utilities.bmatrix(pathb)
end

let x̃
end

println(stderr, "finding solution...")
@time if classic
    println(stderr, "classic algorithm")
    x̃ = A\b
else
    println(stderr, "improved algorithm")
    let pivot
    end

    # by default don’t perform pivoting
    pivot_ = [i for i in 1:n]

    if pivot
        # calculate the pivot row permutation vector
        pivot_ = blocksys.partialPivot(A, n, l, true)
    end

    if lu
        LU, _ = blocksys.decomposeIntoLU(A, b, n, l, pivot_)
        x̃ = blocksys.solveFromLU(LU, b, n, l, pivot_)
    else
        # calculate the `x` vector
        x̃ = blocksys.gaussianElimination(A, b, n, l, pivot_)
    end
end

# print relative error
println("relative error = ", utilities.relativeError(x̃))

max_stars_repo_name: jerry-sky/academic-notebook

max_issues_repo_name: jerry-sky/academic-notebook

max_forks_repo_name: jerry-sky/academic-notebook


text:
module Gong

using Statistics, Dates

#export gc_distance, rad2deg, deg2rad, histc, meshgrid, nan, findNaNmin, findNaNmax

function gc_distance(lat1deg::Float64,lon1deg::Float64,lat2deg::Float64,lon2deg::Float64)
    # This code implements Vincenty 1975: https://www.ngs.noaa.gov/PUBS_LIB/inverse.pdf
    # In the future, need to implement: Karney 2013 https://arxiv.org/abs/1109.4448

    a = 6378.1370; # Earth semi major axis in km
    b = 6356.7523142; # Earth semi-minor axis in km
    R1 = 1/3 * (2*a + b);

    f = (a - b) / a; # flattening parameter for an elipsoid
    #f = 1/298.257223563; # defined by WGS84

    lat1 = deg2rad(lat1deg);
    lon1 = deg2rad(lon1deg);
    lat2 = deg2rad(lat2deg);
    lon2 = deg2rad(lon2deg);

    beta1 = atan((1-f)*tan(lat1)); # reduced latitude 1
    beta2 = atan((1-f)*tan(lat2)); # reduced latitude 2
    P = (beta1 + beta2)/2;
    Q = (beta2 - beta1)/2;

    lambda = abs(lon1 - lon2);
    sigma = atan(sqrt((cos(beta2)*sin(lambda))^2 + (cos(beta1)*sin(beta2) - sin(beta1)*cos(beta2)*cos(lambda))^2) / (sin(beta1)*sin(beta2) + cos(beta1)*cos(beta2)*cos(lambda)));

    alpha = asin(cos(beta1) * cos(beta2) * sin(lambda)/sin(sigma));
    sigmam = acos(cos(sigma) - 2*sin(beta1)*sin(beta2) / cos(alpha)^2)/2;

    X = (sigma - sin(sigma)) * (sin(P)^2 * cos(Q)^2) / (cos(sigma/2)^2)
    Y = (sigma + sin(sigma)) * (cos(P)^2 * sin(Q)^2) / (sin(sigma/2)^2)

    dist = a*(sigma - f/2 * (X+Y));
    return dist
end

function rad2deg(rad)
    return rad * 180.0/pi;
end

function deg2rad(deg)
    return deg * pi/180.0;
end

# define a histogram index function similar to matlab histc
function histc(x, xnodes)
    bin = Array{Number}(undef,length(x));
    for i = 1:length(x)
        a = findlast(x[i] .>= xnodes);
        if !isnothing(a)
            bin[i] = a;
        end
    end
    return bin;
end

# emulate the behavior of Matlab's meshgrid
function meshgrid(xgrid::Array{<:AbstractFloat,1},ygrid::Array{<:AbstractFloat,1})
    nx = length(xgrid);
    ny = length(ygrid);
    minx = minimum(xgrid);
    maxx = maximum(xgrid);
    miny = minimum(ygrid);
    maxy = maximum(ygrid);
    dx = (maxx .- minx) ./ (nx-1);
    dy = (maxy .- miny) ./ (ny-1);
    i = [i for j in miny:dy:maxy, i in minx:dx:maxx];
    j = [j for j in miny:dy:maxy, i in minx:dx:maxx];
    return (i,j)
end

function meshgrid(xgrid::Array{<:Integer,1},ygrid::Array{<:Integer,1})
    minx = minimum(xgrid);
    maxx = maximum(xgrid);
    miny = minimum(ygrid);
    maxy = maximum(ygrid);
    i = [i for j in miny:maxy, i in minx:maxx];
    j = [j for j in miny:maxy, i in minx:maxx];
    return (i,j)
end

function meshgrid(xgrid::UnitRange{<:Integer},ygrid::UnitRange{<:Integer})
    #minx = minimum(xgrid);
    #maxx = maximum(xgrid);
    #miny = minimum(ygrid);
    #maxy = maximum(ygrid);
    i = [i for j in ygrid[1]:ygrid[end], i in xgrid[1]:xgrid[end]];
    j = [j for j in ygrid[1]:ygrid[end], i in xgrid[1]:xgrid[end]];
    return (i,j)
end

# emulate the behavior of Matlab's nan function
function nan(m::Unsigned,n::Unsigned)
    nanarray=Array{Float64,2}(undef,m,n);
    nanarray .= NaN;
end

# adopted from https://github.com/mlubin/NaNMath.jl/issues/31
function findNaNmax(x::Array{<:AbstractFloat,1})
	result = convert(eltype(x), NaN)
    indmax = 0
    @inbounds @simd for i in eachindex(x)
    	v = x[i]
        if !isnan(v)
            if (isnan(result) || v > result)
                result = v
                indmax = i
            end
        end
    end
    return (result,indmax) # note the order of result and ind may be different that what's proposed by floswald as of 2019-12-03
end

# adopted from https://github.com/mlubin/NaNMath.jl/issues/31
function findNaNmin(x::Array{<:AbstractFloat,1})
	result = convert(eltype(x), NaN)
    indmin = 0
    @inbounds @simd for i in eachindex(x)
    	v = x[i]
        if !isnan(v)
            if (isnan(result) || v < result)
                result = v
                indmin = i
            end
        end
    end
    return (result,indmin) # note the order of result and ind may be different that what's proposed by floswald as of 2019-12-03
end

function dg_pol2cart(magnitude::AbstractFloat, compassdir::AbstractFloat)
    theta = mod.(360.0 .- compassdir .+ 90.0, 360) .* pi/180.0;
    u = magnitude .* cos.(theta);
    v = magnitude .* sin.(theta);
    return u + v*im
end

# this does not work, left in for reference. can you 'repeat' function with splat or Tuple() to reproduce Matlab's repmat behavior
#function repmat(x, shape)
#    xsize = size(x);
#    xout = reshape(repeat(x, outer=prod(shape)),[collect(size(x));collect(shape)]...);
#    return xout;
#end

# linearly scale between xmin and xmax using the range provided by xscl
function minmaxscaler(x, xscl; zeromin::Bool=false)
    if zeromin == false
        xmin = minimum(skipmissing(x));
    else
        xmin = 0.0;
    end

    xmax = maximum(skipmissing(x));
    if xmin != xmax
        x_std = (x .- xmin) ./ (xmax - xmin);
        x_scl = x_std .* (maximum(xscl) - minimum(xscl)) .+ minimum(xscl);
    else
        x_scl = x;
    end
    return x_scl, xmin, xmax;
end

# linearly scale between xmin and xmax using the range provided in xscl
function minmaxscaler(x, xscl, xminmax)
    xmin = minimum(skipmissing(xminmax));
    xmax = maximum(skipmissing(xminmax));
    if minimum(skipmissing(x)) != maximum(skipmissing(x))
        x_std = (x .- xmin) ./ (xmax - xmin);
        x_scl = x_std .* (maximum(xscl) - minimum(xscl)) .+ minimum(xscl);
    else
        x_scl = x;
    end
    return x_scl, xmin, xmax;
end

# find standard deviation of a variable, then scale to +/- n std. dev.
function varscaler(varin,nstd=2.0)
    varstd = std(skipmissing(varin));
    varmean = mean(skipmissing(varin));
    varscl = (varin .- varmean) ./ (nstd * varstd);
    return varscl, varmean, varstd, nstd;
end

# manually scale the input variable using specified mean and standard deviation
function varscaler_manual(varin, varmean, varstd, nstd=2.0)
    return (varin .- varmean) ./ (nstd * varstd);
end

# use specified mean and standard deviation to 'unscale' the input variable
function varunscaler(varscl, varmean, varstd, nstd)
    return varscl .* (nstd * varstd) .+ varmean;
end

# this function returns the indices for stock within a specified time period
function findtind(tt,trange)
    return findall(minimum(trange) .<= tt .<= maximum(trange));
end

end


text:
using HTTP
using JSON
using Dates

HTTP.listen("0.0.0.0", 8080, reuseaddr=true) do http
   HTTP.setheader(http, "Server" => "Julia-HTTP")
   HTTP.setheader(http, "Date" => Dates.format(Dates.now(), Dates.RFC1123Format) * " GMT")
   if endswith(http.message.target, "/plaintext")
      HTTP.setheader(http, "Content-Type" => "text/plain")
      HTTP.setstatus(http, 200)
      startwrite(http)
      write(http, "Hello, World!")
   elseif endswith(http.message.target, "/json")
      HTTP.setheader(http, "Content-Type" => "application/json")
      HTTP.setstatus(http, 200)
      startwrite(http)
      write(http, JSON.json(Dict(:message => "Hello, World!")))
   else
       HTTP.setstatus(http, 404)
       startwrite(http)
       write(http, "Not Found")
   end
end

max_stars_repo_name: TaurusMindset/FrameworkBenchmarks

max_issues_repo_name: TaurusMindset/FrameworkBenchmarks

max_forks_repo_name: TaurusMindset/FrameworkBenchmarks


text:
EPS = 1e-6

speye(n) = SparseMatrixCSC(Diagonal(ones(n)))
speye(T, n) = SparseMatrixCSC(Diagonal(ones(T, n)))

pos(x::Real) = max(0, x)
neg(x::Real) = max(0, -x)

struct QuadReflectionMatrix end
J = QuadReflectionMatrix()
*(J::QuadReflectionMatrix, x::Vector{T}) where T <: Real = [x[1]; -x[2:end]]
*(x::Vector{T}, J::QuadReflectionMatrix) where T <: Real = J*x


text:
# This is code for generating random trees of various sorts.
# The hope is that one might have low stretch,
# although none is guaranteed to.

using DataStructures

# sample edges with probability proportional to their weight

"""
    tree = randishKruskal(A)

A heuristic for computing low-stretch spanning trees.  Where Kruskal's MST algorithm adds edges in order of weight, this algorithm adds them at random with probability proportional to their weight.
"""
function randishKruskal(mat::SparseMatrixCSC)
    n = size(mat)[1]
    (ai,aj,av) = findnz(triu(mat))
    m = length(av)

    s = Sampler(av)

    comps = IntDisjointSets(n)

    treeinds = zeros(Int64,n-1)
    numintree = 0
    for j in 1:m
        i = pop!(s)
        if !DataStructures.in_same_set(comps,ai[i],aj[i])
            numintree = numintree+1
            treeinds[numintree] = i
            DataStructures.union!(comps,ai[i],aj[i])
        end
    end

    tree = sparse(ai[treeinds],aj[treeinds],av[treeinds],n,n)
    tree = tree + tree'

    return tree

end

"""
    tree = randishPrim(A)

A heuristic for computing low-stretch spanning trees.  Where Prim's MST algorithm grows a cluster by always adding the edge on the boundary of maximum weight, this algorithm adds a boundary edge with probability proportional to its weight.
"""
function randishPrim(mat::SparseMatrixCSC{Tval,Tind}) where {Tval,Tind}

  n = mat.n
  m = nnz(mat)

  (ai, aj, av) = findnz(mat)
  flipInd = flipIndex(mat)

  visited = zeros(Bool,n)

  s = Sampler(m)

  treeEdges = zeros(Tind,n-1)
  numEdges = 0

  # should randomize the initial choice
  v = one(Tind)

  # always add edge by highest vertex number first
  for ind in mat.colptr[v]:(mat.colptr[v+1]-1)
      wt = mat.nzval[ind]
      push!(s, ind, wt)
  end

  visited[v] = true

  while s.nitems > 0

    edge = pop!(s)

    v = aj[edge]
    if visited[v]
        v = ai[edge]
    end

    if !visited[v]

        numEdges += 1
        treeEdges[numEdges] = edge

        for ind in mat.colptr[v]:(mat.colptr[v+1]-1)
            u = mat.rowval[ind]

            wt = mat.nzval[ind]
            useind = v < u ? ind : flipInd[ind]
            if !visited[u]
                push!(s, useind, wt)

            end
        end

        visited[v] = true
    end

  end

  treeEdges = treeEdges[1:numEdges]

  tr = sparse(ai[treeEdges],aj[treeEdges],av[treeEdges],n,n)
  tr = tr + tr';

  return tr

end # randishPrim


text:
"""
The `Operators` module defines the following interfaces for the following operators:

- `is_deterministic(sf::SFunc)::Bool`
- `sample(sf::SFunc{I,O}, i::I)::O where {I,O}`
- `sample_logcpdf(sf::SFunc{I,O}, i::I)::Tuple{O, AbstractFloat} where {I,O}`
- `invert(sf::SFunc{I,O}, o::O)::I where {I,O}`
- `lambda_msg(sf::SFunc{I,O}, i::SFunc{<:__Opt{Tuple{}}, O})::SFunc{<:__Opt{Tuple{}}, I} where {I,O}`
- `marginalize(sf::SFunc{I,O}, i::SFunc{<:__Opt{Tuple{}}, I})::SFunc{<:__Opt{Tuple{}}, O} where {I,O}`
- `logcpdf(sf::SFunc{I,O}, i::I, o::O)::AbstractFloat where {I,O}`
- `cpdf(sf::SFunc{I,O}, i::I, o::O)::AbstractFloat where {I,O}`
- `log_cond_prob_plus_c(sf::SFunc{I,O}, i::I, o::O)::AbstractFloat where {I,O}`
- `f_expectation(sf::SFunc{I,O}, i::I, fn::Function) where {I,O}`
- `expectation(sf::SFunc{I,O}, i::I)::O where {I,O}`
- `variance(sf::SFunc{I,O}, i::I)::O where {I,O}`
- `get_params(sf::SFunc{I,O,P})::P where {I,O,P}`
- `set_params!(sf::SFunc{I,O,P}, p::P)::SFunc{I,O,P} where {I,O,P}`
- `get_score(sf::SFunc{Tuple{I},O}, i::I)::AbstractFloat where {I,O}`
- `get_log_score(sf::SFunc{Tuple{I},O}, i::I)::AbstractFloat where {I,O}`
- ```support(sf::SFunc{I,O},
                        parranges::NTuple{N,Vector},
                        size::Integer,
                        curr::Vector{<:O}) where {I,O,N}```
- `support_quality(sf::SFunc, parranges)`
- ```bounded_probs(sf::SFunc{I,O},
                             range::__OptVec{<:O},
                             parranges::NTuple{N,Vector})::Tuple{Vector{<:AbstractFloat}, Vector{<:AbstractFloat}} where {I,O,N}```
- ```make_factors(sf::SFunc{I,O},
                            range::__OptVec{<:O},
                            parranges::NTuple{N,Vector},
                            id,
                            parids::Tuple)::Tuple{Vector{<:Scruff.Utils.Factor}, Vector{<:Scruff.Utils.Factor}} where {I,O,N}```
- `initial_stats(sf::SFunc)`
- ```expected_stats(sf::SFunc{I,O},
                              range::__OptVec{<:O},
                              parranges::NTuple{N,Vector},
                              pis::NTuple{M,Dist},
                              child_lambda::Score{<:O}) where {I,O,N,M}```
- `accumulate_stats(sf::SFunc, existing_stats, new_stats)`
- `maximize_stats(sf::SFunc, stats)`
- ```compute_bel(sf::SFunc{I,O},
                          range::__OptVec{<:O},
                          pi::Dist{<:O},
                          lambda::Score{<:O})::Dist{<:O} where {I,O}```
- `compute_lambda(sf::SFunc, range::__OptVec, lambda_msgs::Vector{<:Score})::Score`
- ```send_pi(sf::SFunc{I,O},
                       range::__OptVec{O},
                       bel::Dist{O},
                       lambda_msg::Score{O})::Dist{<:O} where {I,O}```
- ```outgoing_pis(sf::SFunc,
                            range::__OptVec,
                            bel::Dist,
                            incoming_lambdas::__OptVec{<:Score})::Vector{<:Dist}```
- ```outgoing_lambdas(sf::SFunc{I,O},
                      lambda::Score{O},
                      range::__OptVec{O},
                      parranges::NTuple{N,Vector},
                      incoming_pis::Tuple)::Vector{<:Score} where {N,I,O}```
- ```compute_pi(sf::SFunc{I,O},
                         range::__OptVec{O},
                         parranges::NTuple{N,Vector},
                         incoming_pis::Tuple)::Dist{<:O} where {N,I,O}```
- ```send_lambda(sf::SFunc{I,O},
                           lambda::Score{O},
                           range::__OptVec{O},
                           parranges::NTuple{N,Vector},
                           incoming_pis::Tuple,
                           parent_idx::Integer)::Score where {N,I,O}```
"""
module Operators

using Scruff

include("operators/op_performance.jl")
include("operators/op_defs.jl")

"""
    module_functions(mod)

Returns the name of all the functions in the given module.
"""
function module_functions(mod)
    list = Symbol[]
    for nm in names(mod; all=true)
        if !startswith(string(nm), r"@|#") && match(r"^(?:eval|include)$", string(nm)) === nothing
            typeof(eval(nm)) <: Function && push!(list,nm)
        end
    end
    return list
end

"""
    export_operators()

A macro that creates an `export` statement for all the functions in the `Scruff.Operators` module.
"""
macro export_operators()
    is = "export " * join(module_functions(Scruff.Operators), ", ")
    quote
        $(Meta.parse(is))
    end
end

"""
    import_operators()

A macro that creates an `import Scruff.Operators:` statement for all the functions in the `Scruff.Operators` module.
"""
macro import_operators()
    is = "import Scruff.Operators: " * join(module_functions(Scruff.Operators), ", ")
    quote
        $(Meta.parse(is))
    end
end

@export_operators()

end


text:
using Compat
# Do not delete this, this will let you load other codes locally
push!(LOAD_PATH, abspath("modules"))
push!(LOAD_PATH, abspath("."))

max_stars_repo_name: Roger-luo/QCBM

max_issues_repo_name: Roger-luo/QCBM

max_forks_repo_name: Roger-luo/QCBM


text:
using .RasterDataSources

using .RasterDataSources: RasterDataSource

const RDS = RasterDataSources

"""
    GeoArray(T::Type{<:RasterDataSource}, [layer]; kw...) => GeoArray

Load a `RasterDataSource` as an `GeoArray`. `T` and `layers` are are passed to
`RasterDataSources.getraster`, while `kw` args are for both `getraster` and `GeoArray`.

# Keywords

- `month`: an `Int` between `1` and `12`, usually for `Climate` datasets
- `date`: a `DateTime` object, usually for `Weather` datasets.
- `res`: a `String` resolution, for datasets with multiple resolutions.

Other `GeoArray` keywords are passed to the `GeoArray` constructor.

See the docs for
[`RasterDatasources.getraster`](http://docs.ecojulia.org/RasterDataSources.jl/stable/#getraster)
for more specific details about data sources, layers and keyword arguments.
"""
function GeoArray(T::Type{<:RasterDataSource}, layer; crs=_source_crs(T), kw...)
    rds_kw, gd_kw = _filterkw(kw)
    filename = getraster(T, layer; rds_kw...)
    GeoArray(filename; name=RDS.layerkeys(T, layer), crs, gd_kw...)
end

"""
    GeoStack(T::Type{<:RasterDataSource}, [layers::Union{Symbol,AbstractArray,Tuple}]; kw...) => GeoStack

Load a `RasterDataSource` as an `GeoStack`. `T` and `layers` are passed to
`RasterDataSources.getraster`, while `kw` args are for both `getraster` and `GeoStack`.

# Keywords

- `month`: an `Int` between `1` and `12`, usually for `Climate` datasets.
- `date`: a `DateTime` object, usually for `Weather` datasets.
- `res`: a `String` resolution, for datasets with multiple resolutions.

Other `GeoStack` keywords are passed to the `GeoStack` constructor.

See the docs for
[`RasterDatasources.getraster`](http://docs.ecojulia.org/RasterDataSources.jl/stable/#getraster)
for more specific details about data sources, layers and keyword arguments.
"""
GeoStack(T::Type{<:RasterDataSource}; kw...) = GeoStack(T, RDS.layers(T); kw...)
GeoStack(T::Type{<:RasterDataSource}, layer::Symbol; kw...) = GeoStack(T, (layer,); kw...)
function GeoStack(T::Type{<:RasterDataSource}, layers::Tuple; crs=_source_crs(T), kw...)
    rds_kw, gd_kw = _filterkw(kw)
    filenames = map(l -> getraster(T, l; rds_kw...), layers)
    GeoStack(filenames; keys=RDS.layerkeys(T, layers), crs, gd_kw...)
end

"""
    GeoSeries(T::Type{<:RasterDataSource}, [layers::Union{Symbol,AbstractArray,Tuple}]; kw...) => AbstractGeoSeries

Load a `RasterDataSource` as an `AbstractGeoSeries`. `T`, `args` are are passed to
`RasterDataSource.getraster`, while `kw` args are for both `getraster` and `GeoSeries`.

# Keywords

- `month`: a `Vector` or range of `Int` between `1` and `12`, usually for `Climate` datasets.
- `date`: a `Vector` of `DateTime` objects, usually for `Weather` datasets.
- `res`: a `String` resolution, for datasets with multiple resolutions.

Other `GeoSeries` keywords are passed to the `GeoSeries` constructor.

See the docs for
[`RasterDatasources.getraster`](http://docs.ecojulia.org/RasterDataSources.jl/stable/#getraster)
for more specific details about data sources, layers and keyword arguments.
"""
GeoSeries(T::Type{<:RasterDataSource}; kw...) = GeoSeries(T, RDS.layers(T); kw...)
# DateTime time-series
function GeoSeries(T::Type{<:RasterDataSource}, layers;
    resize=_mayberesize(T), crs=_source_crs(T), mappedcrs=nothing, kw...
)
    monthdim = if haskey(kw.data, :month) kw.data[:month] isa AbstractArray
        Dim{:month}(kw.data[:month]; mode=Sampled(; sampling=Intervals(Start())))
    else
        nothing
    end
    datedim = if haskey(kw.data, :date)
        dates = if kw.data[:date] isa Tuple
            dates = RasterDataSources.date_sequence(T, kw.data[:date])
            Ti(dates; mode=Sampled(; sampling=Intervals(Start())))
        elseif kw.data[:date] isa AbstractArray
            Ti(dates; mode=Sampled(; sampling=Intervals(Start())))
        else
            nothing
        end
    else
        nothing
    end
    if isnothing(monthdim) && isnothing(datedim)
        throw(ArgumentError("A GeoSeries can only be constructed from a data source with `date` or `month` keywords that are AbstractArray. For other sources, use GeoStack or GeoArray directly"))
    end

    filenames = getraster(T, layers; kw...)
    can_duplicate = RDS.has_constant_dims(T) && RDS.has_constant_metadata(T)

    if filenames isa AbstractVector{<:AbstractVector}
        series = [GeoSeries(inner_fns, monthdim; resize, crs, mappedcrs, duplicate_first=can_duplicate) for inner_fns in filenames]
        return GeoSeries(series, datedim)
    elseif filenames isa AbstractVector
        dim = isnothing(datedim) ? monthdim : datedim
        return GeoSeries(filenames, dim; resize, crs, mappedcrs, duplicate_first=can_duplicate)
    end
end

_mayberesize(T) = RDS.has_matching_layer_size(T) ? nothing : crop

_source_crs(T) = nothing
_source_crs(T::Type{AWAP}) = crs=EPSG(4326)
_source_crs(T::Type{ALWB}) = crs=EPSG(4326)

function _filterkw(kw)
    rds = []; gd = []
    for p in kw
        dest = first(p) in (:date, :month, :res) ? rds : gd
        push!(dest, p)
    end
    rds, gd
end


text:
models = (
    BayesianLDA,
    BayesianQDA,
    )

fparams = (
    BayesianLDA=(:coef, :intercept, :covariance, :means, :priors, :scalings, :xbar, :classes, :explained_variance_ratio),
    BayesianQDA=(:covariance, :means, :priors, :rotations, :scalings),
    )

@testset "Fit/Predict" begin
    X, y = simple_binaryclf()
    for mod in models
        m = mod()
        acc, f = test_clf(m, X, y)
        @test acc >= 0.75
        @test keys(MB.fitted_params(m, f)) == getproperty(fparams, Symbol(mod))
    end
end

@testset "Docstrings" begin
    for mod in models
        m = mod()
        @test !isempty(MB.docstring(m))
    end
end


text:
# Run model fitting and crossvalidation
# for real or simulated data

# Main runner for model fitting and model comparison
# NOTE: These procedures take a long time to run.
# The settings are different from the ones used in the paper (DEBUG settings).
# Comment/Uncomment the corresponding lines (around 383 and 445)
# to perform the fitting and comparison procedures as in the paper.
function runner_crossval_multipletimes_fit()

    # -------- Fit (once)
    i=""
    result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
    runner_fmri_fit(result_path_base, batch_fit_models_spec)

    # -------- Crossvalidation (multiple runs)
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
        runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec)
    end
end

# Main runner for model and parameter recovery
# NOTE: You need to EDIT which data (real or simulated)
# should be used in
# runner_fmri_fit and runner_fmri_3foldscrossval!!!!
# Right now the real participants' data are used.
# See lines 367 and 428
function runner_multipletimes_fit_recovery()

    # -------- Fit for parameter recovery
    # ---- We run fit multiple times to get some error bars
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_paramrecovery(filestr = string(i))
        runner_fmri_fit(result_path_base, batch_fit_models_spec)

    end

    # -------- Crossvalidation (multiple runs) for model recovery
    nruns = 5
    for i in 1:nruns
        result_path_base, batch_fit_models_spec = model_list_main(filestr = string(i))
        runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec)
    end
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------- Model list -----------------------------------------------------------
# ------------------------------------------------------------------------------
# List and specifications of algorithms for Fig. 3 of paper
function model_list_main(; filestr = "_")

    result_path_base = makehomesavepath("test_main"*filestr)

    stochmax = 20.

    # # Biased Random Walk
    # --------------------
    biasedRandomWalk_prior = [
        RL_Fit.paramBias(0.0, 1.)
    ]
    biasedRandomWalk_prior_std = [.15]
    biasedRandomWalk_spec = RL_Fit.Model_spec_per_subj("Biased Random Walk (0-Model)",
                                RL_Fit.getLL_BiasedRandomWalk_V,
                                biasedRandomWalk_prior, biasedRandomWalk_prior_std,
                                string(result_path_base, "biased_random_walk/"))
    # Sarsa-λ
    # -------
    sarsaL_prior = [
        RL_Fit.paramAlpha(0.0, 1.0),
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramLambda(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.)
        ]
    sarsaL_prior_std = [.15, .15, .15, 1.5]
    sarsaL_spec = RL_Fit.Model_spec_per_subj("Sarsa Lambda",
                    RL_Fit.getLL_SarsaLambda_V,
                    sarsaL_prior, sarsaL_prior_std,
                    string(result_path_base, "sarsa_lambda/"))
    # Prioritized Sweeping with Particle Filtering
    # --------------------------------------------
    priosweep_particle_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNrUpdateCycles(0, 9),
        RL_Fit.paramNParticles(1, 20)
        ]
    priosweep_particle_prior_std = [.15, 1.5, .15, 3., 1.35, 3.]
    priosweep_particle_spec = RL_Fit.Model_spec_per_subj("Prioritized Sweeping Particle Filter",
                                RL_Fit.getLL_Prio_Sweep_TParticleFilterJump_V,
                                priosweep_particle_prior,
                                priosweep_particle_prior_std,
                                string(result_path_base, "priosweep_particle/"))
    # Hybrid-λ-PS-PF
    # --------------
    hybrid_pspf_sarsaL_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNrUpdateCycles(0, 9),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.0),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
        ]
    hybrid_pspf_sarsaL_prior_std = [.15, .15, 3., 1.35, 3., .15, .15, 3., 3.]
    hybrid_pspf_sarsaL_spec = RL_Fit.Model_spec_per_subj("Hybrid: PS with Particle Filtering + Sarsa Lambda",
                                RL_Fit.getLL_Hybrid_Prio_Sweep_TParticleFilterJump_V,
                                hybrid_pspf_sarsaL_prior, hybrid_pspf_sarsaL_prior_std,
                                string(result_path_base, "hybrid_pspf_sarsaL/"))
    # Actor-Critic
    # ------------
    actorCriticTD_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.)
    ]
    actorCriticTD_prior_std = [.15, .15, .15, .15, .15, 1.5]
    actorCriticTD_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda)",
                                RL_Fit.getLL_ActorCritic_TDLambda_V,
                                actorCriticTD_prior, actorCriticTD_prior_std,
                                string(result_path_base, "actor_critic_TD/"))
    # REINFORCE
    # ---------
    REINFORCE_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    REINFORCE_prior_std = [.15, .15, 1.5]
    REINFORCE_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992)",
                            RL_Fit.getLL_REINFORCE_V,
                            REINFORCE_prior, REINFORCE_prior_std,
                            string(result_path_base, "REINFORCE/"))
    # Hybrid Actor-Critic
    # -------------------
    hybrid_actorCritic_FWD_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramEta(0., 1.; name="learning rate fwd"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
    ]
    hybrid_actorCritic_FWD_prior_std = [.15, .15, .15, .15, .15, .15, 3., 3.]
    hybrid_actorCritic_FWD_spec = RL_Fit.Model_spec_per_subj("Hybrid Actor-Critic with FWD",
                                RL_Fit.getLL_Hybrid_AC_FWD_V,
                                hybrid_actorCritic_FWD_prior, hybrid_actorCritic_FWD_prior_std,
                                string(result_path_base, "hybrid_actorcritic_FWD/"))

    # Forward Learner (MB only learner)
    # ---------------------------------
    forwardLearner_prior = [
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    forwardLearner_prior_std = [.15, .15, 1.5]
    forwardLearner_spec = RL_Fit.Model_spec_per_subj("Forward Learner",
                                RL_Fit.getLL_Forward_Learner_V,
                                forwardLearner_prior, forwardLearner_prior_std,
                                string(result_path_base, "forward_learner/"))
    # Hybrid-0 (Glaescher et al. 2010)
    # --------------------------------
    glaescherHybrid_prior = [
        RL_Fit.paramAlpha(0., 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramExpDecayOffset(0., 1.),
        RL_Fit.paramExpDecaySlope(-.1, +.8),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    glaescherHybrid_prior_std = [.15, .15, .15, .15, .15, 1.5]
    glaescherHybrid_spec = RL_Fit.Model_spec_per_subj("Glaescher-Hybrid",
                                RL_Fit.getLL_GlaescherHybrid_V,
                                glaescherHybrid_prior, glaescherHybrid_prior_std,
                                string(result_path_base, "glaescher_hybrid/"))
    # Hybrid-λ (Daw et al. 2011)
    # --------------------------
    hybridSarsaL_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramGamma(0., 1.),
        RL_Fit.paramLambda(0.0, 1.),
        RL_Fit.paramEta(0., 1.),
        RL_Fit.paramExpDecayOffset(0., 1.),
        RL_Fit.paramExpDecaySlope(0., 1.),
        RL_Fit.paramTemp(0.01, 10.)
    ]
    hybridSarsaL_prior_std = [.15, .15, .15, .15, .15, .15, 1.5]
    hybridSarsaL_prior_spec = RL_Fit.Model_spec_per_subj("Hybrid-SarsaLambda",
                                RL_Fit.getLL_HybridSarsaLambda_V,
                                hybridSarsaL_prior, hybridSarsaL_prior_std,
                                string(result_path_base, "hybrid_FWD_SarsaLambda/"))

    # REINFORCE with baseline
    # -----------------------
    REINFORCE_baseline_prior = [
        RL_Fit.paramAlpha(0.0, 1., name="learning rate"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate baseline")
    ]
    REINFORCE_baseline_prior_std = [.15, .15, 1.5, .15]
    REINFORCE_baseline_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992) baseline",
                            RL_Fit.getLL_REINFORCE_baseline_V,
                            REINFORCE_baseline_prior, REINFORCE_baseline_prior_std,
                            string(result_path_base, "REINFORCE_baseline/"))
    # Surprise Sarsa-λ
    # ----------------
    sarsaL_modulated_prior = [
        RL_Fit.paramGamma(0.0, 1.0),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramAlpha(0.0, 1.0),
        RL_Fit.paramLambda(0.0, 1.0)
        ]
    sarsaL_modulated_prior_std = [.15, 1.5, .15, 3., 3., .15, .15]
    sarsaL_modulated_spec = RL_Fit.Model_spec_per_subj("Sarsa Lambda modulated by Particle Filtering",
                                RL_Fit.getLL_sarsaL_modulated_TParticleFilterJump_V,
                                sarsaL_modulated_prior,
                                sarsaL_modulated_prior_std,
                                string(result_path_base, "sarsaL_modulated/"))

    # Surprise Actor-Critic
    # ---------------------
    actorCriticTD_modulatedActor_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    actorCriticTD_modulatedActor_prior_std = [.15, .15, 3., .15, .15, .15, .15, 1.5, 3.]
    actorCriticTD_modulatedActor_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda) modulated actor (1 degree)",
                                RL_Fit.getLL_ActorCritic_TDLambda_modulatedActor_V,
                                actorCriticTD_modulatedActor_prior, actorCriticTD_modulatedActor_prior_std,
                                string(result_path_base, "actor_critic_TD_modulatedActor_1degree/"))
    # Surprise REINFORCE
    # ------------------
    REINFORCE_modulated_prior = [
        RL_Fit.paramAlpha(0.0, 1.),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    REINFORCE_modulated_prior_std = [.15, .15, 3., .15, 1.35, 3.]
    REINFORCE_modulated_spec = RL_Fit.Model_spec_per_subj("REINFORCE (Williams 1992) modulated",
                            RL_Fit.getLL_REINFORCE_modulated_V,
                            REINFORCE_modulated_prior, REINFORCE_modulated_prior_std,
                            string(result_path_base, "REINFORCE_modulated/"))

    # Hybrid Actor-Critic PF
    # ----------------------
    hybrid_actorCritic_FWD_BF_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramNParticles(1, 20),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MB"),
        RL_Fit.paramWeight(0.0, 20.0; name="w_MF")
    ]
    hybrid_actorCritic_FWD_BF_prior_std = [.15, .15, .15, .15, .15, .15, 3., 3., 3., 3.,]
    hybrid_actorCritic_FWD_BF_spec = RL_Fit.Model_spec_per_subj("Hybrid Actor-Critic with FWD - BF",
                                RL_Fit.getLL_Hybrid_AC_FWD_BF_V,
                                hybrid_actorCritic_FWD_BF_prior, hybrid_actorCritic_FWD_BF_prior_std,
                                string(result_path_base, "hybrid_actorcritic_FWD_BF/"))
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --- Run for the paper:
    batch_fit_models_spec = [
                            actorCriticTD_spec,
                            actorCriticTD_modulatedActor_spec,
                            hybrid_actorCritic_FWD_BF_spec,
                            hybrid_actorCritic_FWD_spec,
                            REINFORCE_spec,
                            REINFORCE_modulated_spec,
                            REINFORCE_baseline_spec,
                            sarsaL_spec,
                            sarsaL_modulated_spec,
                            forwardLearner_spec,
                            priosweep_particle_spec,
                            hybridSarsaL_prior_spec,
                            glaescherHybrid_spec,
                            hybrid_pspf_sarsaL_spec,
                            biasedRandomWalk_spec
                          ]
    result_path_base, batch_fit_models_spec
end

# List and specifications of algorithm for parameter recovery
function model_list_paramrecovery(; filestr = "_")

    result_path_base = makehomesavepath("test_main_recovery"*filestr)

    stochmax = 20.

    # Surprise Actor-Critic
    # ---------------------
    actorCriticTD_modulatedActor_prior = [
        RL_Fit.paramAlpha(0.0, 1.; name="learning rate actor"),
        RL_Fit.paramSurpriseProbability(0.0, 1.),
        RL_Fit.paramStochasticity(0.00001, stochmax),
        RL_Fit.paramEta(0.0, 1.; name="learning rate critic"),
        RL_Fit.paramGamma(0.0, 1.),
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace actor"), # actor
        RL_Fit.paramLambda(0.0, 1.; name="elig. trace critic"), # critic
        RL_Fit.paramTemp(0.01, 10.),
        RL_Fit.paramNParticles(1, 20)
    ]
    actorCriticTD_modulatedActor_prior_std = [.15, .15, 3., .15, .15, .15, .15, 1.5, 3.]
    actorCriticTD_modulatedActor_spec = RL_Fit.Model_spec_per_subj("Actor-Critic (TD-Lambda) modulated actor (1 degree)",
                                RL_Fit.getLL_ActorCritic_TDLambda_modulatedActor_V,
                                actorCriticTD_modulatedActor_prior, actorCriticTD_modulatedActor_prior_std,
                                string(result_path_base, "actor_critic_TD_modulatedActor_1degree/"))

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --- Run for parameter recovery:
    batch_fit_models_spec = [
                            actorCriticTD_modulatedActor_spec,
                          ]

    result_path_base, batch_fit_models_spec
end

# ------------------------------------------------------------------------------
# ------------------------------------------------------------------------------
# ------- MCMC settings and helpers --------------------------------------------
# ------------------------------------------------------------------------------

# Settings for data fitting
# Same settings as in runner_fmri_3foldscrossval (see below),
# but here: only fit, not crossvalidation
function runner_fmri_fit(result_path_base,
                                        batch_fit_models_spec
                                        # ; filestr = "_")
                                        )
    rl_module_path = "hi"
    # --------------------------------------------------------------------------
    # ------ Which data? Real or simulated? Edit!

    # Real participants' data:
    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEIC_all_fMRI.csv"

    # Simulated data:
    # data_path = "./mcmc_rl_fit/projects/fmri/data_sim/SARSPEICZVG_all_fMRI_Sim_exporder_SurpAC_run2.csv"

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    description = "Batch fit: Models / subjects"
    # ---- DEBUG settings
    max_nr_episodes = 50  # FIT to the first 10 episodes. For the main analysis, I always used 40
    max_nr_episodes = 2  # FIT to the first 10 episodes. For the main analysis, I always used 40
    nr_repetitions = 2
    top_n_fraction = 0.1
    nr_burnin = 3
    nr_samples = 3
    inter_collect_interval = 4
    # # end debug settings
    # ---------------------------
    # ---- Settings used for paper
    # max_nr_episodes = 50  # FIT to the first n episodes.
    # nr_repetitions = 50 #20 # one repetition generates nr_mcmc_samples. doing several repetitions, starting from different random initial positions helps finding isolated regions (e.g. bimodal distributions).
    # top_n_fraction = 0.001 # not important. I sometimes a the top 1percent of mcmc samples
    # nr_burnin = 1500 # important. mcmc needs to "burn in". The more params (larger space) the longer it takes to walk into the volume of interest.
    # nr_samples = 2000 # how many samples to collect from each independent repetition
    # inter_collect_interval = 10 # samples are correlated. therefore make N steps before collecting a sample. 7 to 10 seems a good value.
    # # end
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------

    fit_spec = RL_Fit.FitSpec_multiple_models(
        batch_fit_models_spec,
        max_nr_episodes,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        # false,  #do_per_subject_fit
        true,  #do_per_subject_fit
        result_path_base
    )
    # --------------------------------------------------------------------------
    # ----------------- RUN -----------------
    RL_Fit.run_fit_multiple_models(description, data_path, fit_spec, rl_module_path)

end

# Settings for crossvalidation
# Same settings as in runner_fmri_fit (see above),
# but here: ONLY crossvalidation, not fit
function runner_fmri_3foldscrossval(result_path_base,
                                        batch_fit_models_spec
                                        )
    rl_module_path = "hi"
    # --------------------------------------------------------------------------
    # ------ Which data? Real or simulated? Edit!

    # Real participants' data:
    data_path = "./mcmc_rl_fit/projects/fmri/data/SARSPEIC_all_fMRI.csv"

    # Simulated data:
    # data_path = "./mcmc_rl_fit/projects/fmri/data_sim/SARSPEICZVG_all_fMRI_Sim_exporder_SurpAC_run2.csv"

    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    description = "Batch fit: Models / subjects"
    # ---- DEBUG settings:
    max_nr_episodes = 50  # FIT to the first 10 episodes. For the main analysis, I always used 40
    max_nr_episodes = 2  # FIT to the first 10 episodes. For the main analysis, I always used 40
    nr_repetitions = 2
    top_n_fraction = 0.1
    nr_burnin = 3
    nr_samples = 3
    inter_collect_interval = 4
    # # end debug settings
    # ---------------------------
    # ---- Settings used for paper
    # max_nr_episodes = 50  # FIT to the first n episodes.
    # nr_repetitions = 50 #20 # one repetition generates nr_mcmc_samples. doing several repetitions, starting from different random initial positions helps finding isolated regions (e.g. bimodal distributions).
    # top_n_fraction = 0.001 # not important. I sometimes a the top 1percent of mcmc samples
    # nr_burnin = 1500 # important. mcmc needs to "burn in". The more params (larger space) the longer it takes to walk into the volume of interest.
    # nr_samples = 2000 # how many samples to collect from each independent repetition
    # inter_collect_interval = 10 # samples are correlated. therefore make N steps before collecting a sample. 7 to 10 seems a good value.
    # # end
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------
    # --------------------------------------------------------------------------

    fit_spec = RL_Fit.FitSpec_multiple_models(
        batch_fit_models_spec,
        max_nr_episodes,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        false,  #do_per_subject_fit
        # true,  #do_per_subject_fit
        result_path_base
    )
    # --------------------------------------------------------------------------
    # # ----------------- CROSSVALIDATE -----------------
    (dataFmriAll, all_data_stats) = RL_Fit.load_SARSPEI(data_path; max_episodes_per_subj = max_nr_episodes)
    # when running just an additional model, reuse the same (initially randomly) sampled
    # (make results comparable with previously fit algorithms)

    # I called kf_Spec = RL_Fit.get_k_folds_all_out(Set(all_data_stats.all_subject_ids), 3) once
    # and then fix the subject ID per fold
    kf_Spec = RL_Fit.get_k_folds( Set(all_data_stats.all_subject_ids),
    [
    Set([7, 2, 19, 5, 8, 12, 20])
    Set([18, 4, 13, 3, 15, 21, 6])
    Set([9, 14, 10, 16, 17, 11, 1])])
    # Otherwise shuffle(collect(1:21))... and split in 3
    CV_multi_model_spec = RL_Fit.CV_multiple_models_spec(
        batch_fit_models_spec,
        kf_Spec,
        nr_burnin,
        nr_samples,
        inter_collect_interval,
        top_n_fraction,
        nr_repetitions,
        result_path_base)
    cv_result = RL_Fit.run_cross_validation(description, dataFmriAll,
                                            CV_multi_model_spec, rl_module_path)
end


text:
include("Graph_class.jl")
include("Tools.jl")
include("Sampling.jl")
include("Algorithm.jl")

using LinearAlgebra

"""
    node_selection_linear_sum(R, k)

    (1) R: a list of tuples, each tuple contains a rr-set, and its corresponding sampled vertex
    (2) n: number of vertices in the graph
    (3) node_weight: a list of node weights
    (4) k: budget for seednodes.

    % Inside the function
    (1) reverse_R_list Vector{Set{Int}}: for each vertex u, stores all the IDs of rr sets that u covers. The values will be updated in the function
    (2) value_list: PriorityQueue{Int, Float64}(Base.Order.Reverse) for each vertex u, stores the sum of weights of the rr sets that u covers. The values will be updated in the function
    (3) weight_list: each rr_set contains a weight, the weight equals to the weight of corresponding vertex if covered.
    (4) size: The length of the rr set
"""
function node_selection_linear_sum(R, n, node_weight, k)
    T = time()

    rrl = [r[1] for r in R]
    nodel = [r[2] for r in R] # the sampled vertex

    #println(R)

    len_rrs = length(rrl)

    rR = [Set{Int}() for _ in 1:n]
    vl = PriorityQueue{Int32, Float64}(Base.Order.Reverse) #init as priorityQueue
    wl = [node_weight[i] for i in nodel]
    tmp_rr = []
    tmp_rr_another = []

    rr3D = falses(len_rrs, n) # Init a bit matrix for quick check
    for j in 1:length(rrl)
        for i in rrl[j]
            rr3D[j, i] = true
        end
    end

    #println("Time prepare", time() - T)

    for i in 1:n
        #tmp_rr = [j for j in 1:len_rrs if i in rrl[j]] # All the rr IDs covered by node i
        tmp_rr = [j for j in 1:len_rrs if rr3D[j, i] == true]
        #println("equal or not", isequal(tmp_rr, tmp_rr_another))
        union!(rR[i], tmp_rr)
        #rR[i] = Set{Int}(tmp_rr)
        #println([weight_list[k] for k in tmp_rr])
        vl[i] = sum([wl[k] for k in rR[i]])
    end

    #println("Time first loop", time() - T)

    SEED = zeros(Int32, k) # In oder to store the seeds
    cR = Set{Int}() # The rr-sets going to be covered
    tmp_value, ris_value = 0, 0
    k_copy = k
    #tmp_size,  cm, fc, ris_value = 0, 0, 0, 0
    # The process of node-selection
    for j in 1:k
        SEED[j] = peek(vl)[1] # Choose the node with the highest priority
        if peek(vl)[2] < 0 #if can not find a seed with positive addition
            k_copy = j - 1
            break
        end
        tmp_value += peek(vl)[2]
        #println("peek value is ", peek(vl)[2])
        #println("seed is ", SEED[j])
        cR = rR[SEED[j]]
        #println("cR are", cR)
        for idx in 1 : n
            #println("wl is ")
            #println(wl)
            #println("intersection is")
            #println("vl[idx] is ", vl[idx])
            vl[idx] -= sum(wl[[v for v in intersect(rR[idx], cR)]]) # Update value_list
            rR[idx] = setdiff(rR[idx], cR) # update reverse_R_list
        end
    end

    #println("Time second loop", time() - T)

    ris_value = tmp_value / len_rrs

    return ris_value, SEED[1:k_copy]
end

"""
    ris_measure_linear(G::Graph, k, node_weight, delta)

    delta: network parameter (ack, spread, reject)
    ℓ: accuracy parameter

    ris for maximize the measures of their linear part.
    Return: ris_value, seed nodes, sample_size

    flag: For some indices that are so small, we set a upper bound on sampling size (aci, ad, ap)

        Tang, Youze, Yanchen Shi, and Xiaokui Xiao.
        "Influence maximization in near-linear time: A martingale approach."
        Proceedings of the 2015 ACM SIGMOD International Conference on Management of Data. 2015.
"""
function ris_measure_linear(G::Graph, k, node_weight, flag = 0, delta = 0.5, ϵ = 0.4, ℓ = 1)
    # Create the "reversed" Link Table
    g = Array{Array{Tuple{Int32, Float64}, 1}, 1}(undef, G.n)
    for i = 1 : G.n
        g[i] = []
    end

    for (_, u, v, w) in G.E
        push!(g[v], (u, w))
    end

    T = time()

    χ = maximum(abs.(node_weight))

    n = G.n
    ϵ′ = 2 * ϵ # let ϵ' large enough
    λ′ = (2 + 4/3*ϵ′)*(logchoose(n, k) + ℓ * log(n) + log(log2(n)))*n / (ϵ′)^2

    λ⁺ = 8*n*χ*(ϵ/3 + 1)*(ℓ * log(n) + log(2) + logchoose(n, k)) / (ϵ)^2
    #LB = χ
    LB = 1 # when LB < 1, then we do not need much accuracy
    R, SEED = [], []
    ris_value, x, θᵢ = 0, 0, 0

    #println("before loop time is ", time() - T)

    for i in 1 : round(log2(n) - 1)
        y = n / (2^i)
        θᵢ = λ′ / y
        if length(R) <= θᵢ
            append!(R, [getRRS(g, n, delta) for _ in 1 : (θᵢ - length(R) + 1)])
        end

        if (flag == 1) && (length(R) >= n*50) # add condition to avoid too much sampling
            break
        end

        #println("loop is ", i)
        #println("Time Before selection is ", time() - T)

        ris_value, _ = node_selection_linear_sum(R, n, node_weight, k)

        #println("Time After selection is ", time() - T)

        if ris_value * n < 1 # avoid too many sampling
            #println("small value")
            break
        else
            if n*ris_value >= (1 + ϵ′)*y*χ
                LB = n * ris_value / (1 + ϵ′)
                break
            end
        end
    end

    θ = λ⁺ / LB
    if (length(R) <= θ)
        if (θ <= n*50) || (flag == 0)
            append!(R, [getRRS(g, n, delta) for _ in 1 : (θ - length(R) + 1)])
        else
            #println("avoid much sampling")
            append!(R, [getRRS(g, n, delta) for _ in 1 : (n*50 - length(R) + 1)])
        end
    end

    #println("Time Before last calculation is ", time() - T)

    ris_value, SEED = node_selection_linear_sum(R, n, node_weight, k)

    #println("Time After last calculation is ", time() - T)

    T = time() - T
    return ris_value * n, SEED, length(R), T
end

## Below for greedy algorithm

function evaluation_ris(R, n, node_weight, pair_weight, seeds)
    rrl = [r[1] for r in R]
    nodel = [r[2] for r in R] # node pair in this case

    len_rrs = length(R)
    nR = div(length(R), 2)

    rR = Set{Int}()
    wl = [n*pair_weight[nodel[i], nodel[i + nR]] for i in 1:nR]
    sl = [node_weight[nodel[i]] for i in 1:nR]

    for i in seeds
        union!(rR, [j for j in 1:len_rrs if i in rrl[j]]) # All the rr IDs covered by node i
    end

    return sum([wl[k] for k in rR if (k + nR) in rR])* n / nR + sum([sl[k] for k in rR if k <= nR])* n / nR # as init, i should cover the pair of sampled rr-id
end

# For general greedy algorithm
function node_selection_pair(R, n, node_weight, pair_weight, k)
    rrl = [r[1] for r in R]
    nodel = [r[2] for r in R] # node pair in this case

    len_rrs = length(R)
    tmp_rr = []
    nR = div(length(R), 2)

    rR = [Set{Int}() for _ in 1:n]
    vl = PriorityQueue{Int32, Float64}(Base.Order.Reverse) #init as priorityQueue
    wl = [n*pair_weight[nodel[i], nodel[i + nR]] for i in 1:nR] # weight for the pairs
    sl = [node_weight[nodel[i]] for i in 1:nR] # weight for single
    cover = [false for _ in 1 : len_rrs] # mark if a node is covered in general
    valid_cover = [false for _ in 1 : nR] # only when the gain is achived

    rr3D = falses(len_rrs, n) # Init a bit matrix for quick check
    for j in 1:length(rrl)
        for i in rrl[j]
            rr3D[j, i] = true
        end
    end

    for i in 1:n
        tmp_rr = [j for j in 1:len_rrs if rr3D[j, i] == true] # All the rr IDs covered by node i
        rR[i] = Set{Int}(tmp_rr)
        vl[i] = sum([wl[k] for k in rR[i] if (k + nR) in rR[i]]) # as init, i should cover the pair of sampled rr-id, only for pairs
        vl[i] += sum([sl[k] for k in rR[i] if k <= nR]) # the single weight
    end

    SEED = zeros(Int32, k)
    cR = Set{Int}() # The rr-sets going to be covered
    tmp_value, ris_value = 0, 0
    k_copy = k
    # The process of node-selection
    for j in 1:k
        SEED[j] = peek(vl)[1] # Choose the node with the highest priorit

        NewValid, NewValid_pair, TmpInter, TmprDiff = Set{Int}(), Set{Int}(), Set{Int}(), Set{Int}() # newly real-covered node, to update valid_cover, only stores the id < n; TmpInter
        if peek(vl)[2] < 0 #if can not find a seed with positive addition
            k_copy = j - 1
            break
        end
        tmp_value += peek(vl)[2]
        cR = rR[SEED[j]] # covered rr-ids by seed[j]

        # check all the covers, can they make pairs?
        for c in cR
            #println("c is ", c)
            if c > nR
                if (cover[c - nR] == true)
                    push!(NewValid, c)
                    valid_cover[c - nR] = true
                end
            else
                if (cover[c + nR] == true)
                    push!(NewValid, c)
                    valid_cover[c] = true
                end
            end
        end

        # check all the pairs in cR, in order to make pairs
        for c in cR
            if c < nR
                if (c + nR) in cR
                    push!(NewValid_pair, c)
                    valid_cover[c] = true
                end
            end
        end

        # update cover
        for c in cR
            cover[c] = true # update cover
        end

        #println("NewValid is ", NewValid)

        for idx in 1 : n
            TmprDiff = setdiff(rR[idx], cR) # update reverse_R_list
            TmpInter = intersect(rR[idx], cR)

            vl[idx] -= sum(sl[[v for v in TmpInter if v <= nR]]) # for single nodes

            for v in TmpInter
                 # Update value_list, subtraction
                if v in NewValid
                    vl[idx] -= wl[(v > nR) ? v - nR : v]
                end
                if (v in NewValid_pair) && ((v + nR) in TmpInter)
                    vl[idx] -= wl[v]
                end
            end

            rR[idx] = TmprDiff # update set

            for v in TmprDiff
                # Update value_list, avoid pair
                if (((v + nR) in cR) && ((v + nR) ∉ TmpInter))  || (((v - nR) in cR) && ((v - nR) ∉ TmpInter))
                    vl[idx] += wl[(v > nR) ? v - nR : v] # Update value_list
                end
            end
        end

    end

    ris_value = tmp_value / nR

    return ris_value, SEED[1:k_copy]
end

function ris_measure_greedy(G::Graph, k, node_weight, pair_weight, delta = 0.5, ϵ = 0.3, ℓ = 1)
    # Create the "reversed" Link Table
    g = Array{Array{Tuple{Int32, Float64}, 1}, 1}(undef, G.n)
    for i = 1 : G.n
        g[i] = []
    end

    for (_, u, v, w) in G.E
        push!(g[v], (u, w))
    end

    T = time()

    χ = maximum(abs.(node_weight)) + maximum(abs.(pair_weight)) * G.n

    n = G.n
    ϵ′ = √2 * ϵ # let ϵ' large enough
    λ′ = (2 + 4/3*ϵ′)*(logchoose(n, k) + ℓ * log(n) + log(log2(n)))*n / (ϵ′)^2

    λ⁺ = 8*n*χ*(ϵ/3 + 1)*(ℓ * log(n) + log(2) + logchoose(n, k)) / (ϵ)^2
    #LB = χ
    LB = 1 # when LB < 1, then we do not need much accuracy
    R, SEED = [], []
    ris_value, θᵢ = 0, 0
    UppSample = n * n * 20  # set upper bound on sampling size

    for i in 1 : round(log2(n))
        y = n^2 / (2^i)
        θᵢ = div(λ′, y)
        if length(R) <= θᵢ
            append!(R, [getRRS(g, n, delta) for _ in 1 :  2*(θᵢ - length(R) + 1)])
        end

        ris_value, _ = node_selection_pair(R, n, node_weight, pair_weight, k)

        if (length(R) >= UppSample) # add condition to avoid too much sampling
            break
        end

        if ris_value * n < 1 # avoid too many sampling
           break
        else
            if n * ris_value >= (1 + ϵ′)*y*χ
                LB = n * ris_value / (1 + ϵ′)
                break
            end
        end
    end

    θ = div(λ⁺, LB)
    if (length(R) <= θ)
        if (θ <= UppSample)
            append!(R, [getRRS(g, n, delta) for _ in 1 : 2*(θ - length(R) + 1)])
        else
            append!(R, [getRRS(g, n, delta) for _ in 1 : 2*(UppSample - length(R) + 1)])
        end
    end
    #append!(R, [getRRSTwo(g, n, delta) for _ in 1 : (θ - length(R) + 1)])

    ris_value, SEED = node_selection_pair(R, n, node_weight, pair_weight, k)

    T = time() - T
    return ris_value * n , SEED, length(R), T, R
end

function assign_pair_weights(G::Graph, Δs)
	sL = getSparseL(G)
    Inv = getW(G)
    sT = Matrix{Float64}(I, G.n, G.n) - ones(G.n,G.n)/G.n

    p = Δs .* (Δs .* (Inv*sT*Inv))' # polarization
    d = Δs .* (Δs .* (Inv*sL*Inv))'
    ci = Δs .* (Δs .* (Inv*sL*sL*Inv))'
    idc = Δs .* (Δs .* (Inv))'

	return ci, d, p, idc
end

"""
    assign_node_weights(G::Graph, s, Δs)

    G: The input graph for computing different measures.
    s: The vector of innate opinions
    Δs: The change of the innate opinion of a vertex once it gets influenced.

    This function is supposed to create different node weights depending on different measures.
    The measures include Internal Conflict Index, Disagreement Index, Polarization Index, and Controversy index + Polarization Index.
"""
function assign_node_weights(G::Graph, s, Δs; eps = 1e-6)
    IpL = getSparseIpL(G)
	sL = getSparseL(G)
	f = approxchol_sddm(IpL, tol=0.1*eps)

	z = f(s)

	# calculate aC_I(G)
	aci = 2 * Δs .* f(sL * (sL * z))

	# calculate aD(G)
	ad = 2 * Δs .* f(sL * z)

	# calculate aP(G)
	ap = 2 * Δs .* (f(z) .- (sum(s) / G.n))

	# calculate aI_dc(G)
	aidc = 2 * Δs .* z

    aidc_upper = 2 * Δs .* z +  Δs .* Δs

	return aci, ad, ap, aidc, aidc_upper
end

"""
    delta_s_marketing(s, epsilon)

    How Δs is calculated under marketing setting
"""
function delta_s_marketing(s, epsilon)
    return [1 > (s[i] + epsilon) ? epsilon : 1 - s[i] for i in 1 : length(s)]
end

"""
    delta_s_backfire(s, epsilon, beta)

    How Δs is calculated under back fire setting
"""
function delta_s_backfire(s, epsilon, beta)
    delta_s = []
    for v in s
        if v >= beta
            if v + epsilon > 1
                push!(delta_s, 1 - v)
            else
                push!(delta_s, epsilon)
            end
        else
            if v - epsilon < 0
                push!(delta_s, -v)
            else
                push!(delta_s, -epsilon)
            end
        end
    end
    return delta_s
end


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.10.3
#   kernelspec:
#     display_name: Julia 1.6.4
#     language: julia
#     name: julia-1.6
# ---

# %%
using Distributions
using Plots
using BlackBoxOptim

# model
Distributions.TDist(μ, ρ, ν) = LocationScale(μ, ρ, TDist(ν))

# test data
X = [-0.01, 0.01, 1.0];

# %%
using Random
Random.seed!(123456789)

# %%
SearchRange = [(-0.1, 0.4), (-15.0, 10.0), (-2.5, 8.0)]
TraceMode = :compact
TraceInterval = 0.001
MaxSteps = 50000
@show SearchRange
o = bboptimize(x -> -loglikelihood(TDist(x[1], 10^x[2], 10^x[3]), X); SearchRange, TraceMode, TraceInterval, MaxSteps);

# %%
SearchRange = [(-0.5, 1.0), (-15.0, 10.0), (-2.5, 8.0)]
TraceInterval = 0.001
MaxSteps = 50000
@show SearchRange
o = bboptimize(x -> -loglikelihood(TDist(x[1], 10^x[2], 10^x[3]), X); SearchRange, TraceMode, TraceInterval, MaxSteps);

# %%
SearchRange = [(-3.0, 4.0), (-15.0, 10.0), (-2.5, 8.0)]
TraceInterval = 0.001
MaxSteps = 50000
@show SearchRange
o = bboptimize(x -> -loglikelihood(TDist(x[1], 10^x[2], 10^x[3]), X); SearchRange, TraceMode, TraceInterval, MaxSteps);

# %%

max_stars_repo_path: 0024/TDist model is not log-convex - BlackBoxOptim.jl

max_stars_repo_name: genkuroki/public

max_issues_repo_name: genkuroki/public

max_forks_repo_name: genkuroki/public