text:
using CUDA
using PaddedViews
using TimerOutputs
function PadData(pd, data)
    x = size(data,1)
    y = size(data,2)
    z = size(data,3)
    padx = 1:(x+2pd)
    pady = 1:(y+2pd)
    padz = 1:z
    datax = (pd+1):(x+pd)
    datay = (pd+1):(y+pd)
    dataz = (1:z)
    return PaddedView(0, data, (padx,pady,padz), (datax, datay, dataz))
end

function comm_groups(n)
    gs = []

    for i = 1:n÷2
        id = i*2
        push!(gs, (id-1,id))
    end

    for i = 1:(n-1)÷2
        id = i*2 + 1
        push!(gs, (id-1,id))
    end
    return gs
end

function split_data_to_gpus(n_gpus,data, t_group, radius)
    ## Split data evenly
    len = round(size(data,3)/n_gpus)
    edges = []
    start = 1
    for i = 1:(n_gpus-1)
        push!(edges, [start,start+len-1])
        start += len
    end
    push!(edges, [start,size(data,3)])
    ed = deepcopy(edges)
    ## Determine the data with padding of the data
    pad = t_group*radius
    if n_gpus == 1
        return Int.(edges[1]), ed
    else
        (edges[1])[2] = (edges[1])[2] + pad
        for i = 2:(n_gpus-1)
            (edges[i])[1] = (edges[i])[1] - pad
            (edges[i])[2] = (edges[i])[2] + pad
        end
        (edges[n_gpus])[1] = (edges[n_gpus])[1] - pad
    end
    ## Determine the padding length for transfer
    [[-pad*(i!=1), pad*(i!=n_gpus)]  for i = 1:n_gpus]
    return map(x -> Int.(x) , edges),map(x -> Int.(x) , ed)
end
function k_init(d)
    tx = (blockIdx().x-1)*blockDim().x + threadIdx().x
    ty = (blockIdx().y-1)*blockDim().y + threadIdx().y
    tz = (blockIdx().z-1)*blockDim().z + threadIdx().z
    dx,dy,dz = size(d)
    if tx <= dx && ty <= dy && tz <=dz
        @inbounds d[tx,ty,tz] = 0
    end
    return nothing
end
function init_zero(d, data)
    device!(d)
    threads = (16,8,8)
    blocks = size(data).÷threads
    blocks = blocks .+ 1
    @cuda blocks=blocks threads=threads k_init(data)
    nothing
end
function ApplyMultiGPU(ngpus, st_inst, t_steps, data ;vsq=nothing, t_group=1, dbg=false, mc=false, specific_gpus=false)
    spec_gpus = collect(0:ngpus-1)

    if ngpus == 2
        spec_gpus = [1 2]
    end
    to = TimerOutput()
    #init_gpu_channels(ngpus)
    radius = max(st_inst.front_z_max, st_inst.behind_z_max)
    ## Split Data among gpus

    comm_g = comm_groups(ngpus)
    s_ind, u_ind = split_data_to_gpus(ngpus, data, t_group, radius)
    s_data = []
    padded_data = PadData(t_group*radius, data)
    for i = 1:ngpus
        s = s_ind[i]
        push!(s_data, @view padded_data[:,:,s[1]:s[2]])
    end
    @timeit to "Init" begin
    ## Create Unified for CPU
    buf_host = Mem.alloc(Mem.Unified, prod(size(padded_data))*sizeof(Float32))
    CUDA.Mem.advise(buf_host, Mem.ADVISE_SET_PREFERRED_LOCATION,
            prod(size(padded_data))*sizeof(Float32), device=CUDA.DEVICE_CPU)
    ptr_host = buf_host.ptr
    g_out = unsafe_wrap(Array{Float32, 3},
                        convert(Ptr{Float32}, buf_host),
                        size(padded_data), own=false)
    copyto!(g_out, padded_data)

    ## Create Unified GPU arrays
    gpu_buffers = [Mem.alloc(Mem.Unified, prod(size(s_data[i]))*sizeof(Float32))
                            for i = 1:ngpus]
    # for i = 1:ngpus
    #     CUDA.Mem.advise(gpu_buffers[i], Mem.ADVISE_SET_PREFERRED_LOCATION,
    #             prod(size(s_data[i]))*sizeof(Float32), device=CUDA.devuce_gpu))
    # end
    gpu_arrays_in =  [unsafe_wrap(CuArray{Float32, 3},
                        convert(CuPtr{Float32}, gpu_buffers[i]),
                        (size(padded_data,1),size(padded_data,2),s_ind[i][2]-s_ind[i][1]+1),
                        own=false) for i = 1:ngpus]
    gpu_pointers_in = [convert(CuPtr{Nothing}, gpu_arrays_in[i].ptr) for i = 1:ngpus]

    gpu_buffers_out = [Mem.alloc(Mem.Unified, prod(size(s_data[i]))*sizeof(Float32))
                            for i = 1:ngpus]
    gpu_arrays_out =  [unsafe_wrap(CuArray{Float32, 3},
                        convert(CuPtr{Float32}, gpu_buffers_out[i]),
                        (size(padded_data,1),size(padded_data,2),s_ind[i][2]-s_ind[i][1]+1),
                        own=false) for i = 1:ngpus]
    gpu_pointers_out = [convert(CuPtr{Nothing}, gpu_arrays_out[i].ptr) for i = 1:ngpus]

    # for i = 1:ngpus
    #     CUDA.cuMemsetD32_v2(gpu_pointers_out[i], Float32(0), prod(size(gpu_arrays_out[i])))
    # end
    ## MemCpy Initial Data
    @timeit to "shit" begin
    @sync begin
        for i = 1:ngpus
            @async begin
                device!(spec_gpus[i])
                #(dbg) && println(i, " start kernel")
                j = i#(i+1)%ngpus + 1
                a = call_kernel(st_inst.bdim, i, ngpus,
                    gpu_arrays_in[j], gpu_arrays_out[j], 1, st_inst)
                #(dbg) && println(i, " end kernel")
            end
        end
    end end
    @sync begin
        for i = 1:ngpus
            s = s_ind[i]
            h_ptr = ptr_host + (s[1]-1)*size(padded_data, 1)*size(padded_data,2)*sizeof(Float32)
            bcount = (s[2]-s[1]+1)*size(padded_data, 1)*size(padded_data,2)*sizeof(Float32)
            @async CUDA.cuMemcpy(gpu_pointers_in[i], h_ptr, bcount)
        end
    end
    end

    at_out = [true for i = 1:ngpus]
    t_counter = 0
    flag_break = false
    #@time begin
    @timeit to "loop" begin
    while true
        if t_counter + t_group >= t_steps
            t_group = t_steps - t_counter
            flag_break = true
        end
        ## kernel loop
        #@timeit to "kernels" begin
        @timeit to "kernel" begin
        @sync begin
            for i = 1:ngpus
                @async begin
                    device!(spec_gpus[i])
                    if ngpus == 4
                        if i != 1
                            j = i#(i+1)%ngpus + 1
                            at_out[i] = call_kernel(st_inst.bdim, i, ngpus,
                                gpu_arrays_in[j], gpu_arrays_out[j], t_group, st_inst)
                        end
                    else
                        j = i#(i+1)%ngpus + 1
                        at_out[i] = call_kernel(st_inst.bdim, i, ngpus,
                            gpu_arrays_in[j], gpu_arrays_out[j], t_group, st_inst)
                    end
                    #println(i, "start")
                    #(dbg) && println(i, " start kernel")
                    # j = i#(i+1)%ngpus + 1
                    # at_out[i] = call_kernel(st_inst.bdim, i, ngpus,
                    #     gpu_arrays_in[j], gpu_arrays_out[j], t_group, st_inst)
                    #(dbg) && println(i, " end kernel")
                    #println(i, "end")
                end
            end
        end
        if mc
            @sync begin
                for i = 1:ngpus
                    @async begin
                        device!(spec_gpus[i])
                        CUDA.synchronize()
                    end
                end
            end
        end
        end
        #end
        if !at_out[1]
            gpu_arrays_in,gpu_arrays_out = gpu_arrays_out,gpu_arrays_in
            gpu_pointers_in,gpu_pointers_out = gpu_pointers_out,gpu_pointers_in
        end
        if flag_break
            break
        end
        # ## comm data
        bsize = size(gpu_arrays_in[1], 1)*size(gpu_arrays_in[1], 1)*
                    radius*t_group*sizeof(Float32)
        #@timeit to "comms" begin
        @timeit to "coms" begin
        @sync begin
            for i in comm_g
                @async begin
                    #println(i, " left")
                    dx,dy,dz = size(gpu_arrays_in[i[1]])
                    p1 = gpu_pointers_in[i[1]]
                    p2 = gpu_pointers_in[i[2]]
                    offp1 = p1 + dx*dy*(dz-2radius*t_group)*sizeof(Float32)
                    CUDA.cuMemcpy(p2, offp1, bsize)

                    # offp1 = p1 + dx*dy*(dz-radius*t_group)*sizeof(Float32)
                    # offp2 = p2 + dx*dy*(radius*t_group)*sizeof(Float32)
                    # CUDA.cuMemcpy(offp1, offp2, bsize)
                end
            end
            for i in comm_g
                @async begin
                    #println(i, " right")
                    dx,dy,dz = size(gpu_arrays_in[i[1]])
                    p1 = gpu_pointers_in[i[1]]
                    p2 = gpu_pointers_in[i[2]]
                    # offp1 = p1 + dx*dy*(dz-2radius*t_group)*sizeof(Float32)
                    # CUDA.cuMemcpy(p2, offp1, bsize)

                    offp1 = p1 + dx*dy*(dz-radius*t_group)*sizeof(Float32)
                    offp2 = p2 + dx*dy*(radius*t_group)*sizeof(Float32)
                    CUDA.cuMemcpy(offp1, offp2, bsize)
                end
            end
            if mc
                @sync begin
                    for i = 1:ngpus
                        @async begin
                            device!(spec_gpus[i])
                            CUDA.synchronize()
                        end
                    end
                end
            end
        end
        end
        #end

        t_counter += t_group
    end
    # @sync begin
    # @timeit to "sync" begin
    #     for i = 1:ngpus
    #         @async begin
    #             device!(spec_gpus[i])
    #             println("sync ", i)
    #             CUDA.synchronize()
    #             println("sync comp", i)
    #         end
    #     end
    # end
    # end
    @timeit to "dow" begin
    @sync begin
        for i = 1:ngpus
            @async begin
                s = u_ind[i]
                host_off = size(g_out, 1)*size(g_out,2)*(s[1]-1)*sizeof(Float32)
                dev_off = (i!=1)*radius*t_group*size(g_out, 1)*size(g_out,2)*sizeof(Float32)
                bsize =  size(g_out, 1)*size(g_out,2)*(s[2]-s[1]+1)*sizeof(Float32)
                CUDA.cuMemcpy(ptr_host+host_off, gpu_pointers_in[i] + dev_off,
                            bsize)
            end
        end
    end
    end
    ## Copy data to cpu
    #@timeit to "download" begin
    # @timeit to "dow" begin
    # @sync begin
    #     for i = 1:ngpus
    #         @async begin
    #             s = u_ind[i]
    #             host_off = size(g_out, 1)*size(g_out,2)*(s[1]-1)*sizeof(Float32)
    #             dev_off = (i!=1)*radius*t_group*size(g_out, 1)*size(g_out,2)*sizeof(Float32)
    #             bsize =  size(g_out, 1)*size(g_out,2)*(s[2]-s[1]+1)*sizeof(Float32)
    #             CUDA.cuMemcpy(ptr_host+host_off, gpu_pointers_in[i] + dev_off,
    #                         bsize)
    #         end
    #     end
    # end
    end
    println(to)
    nothing
    #end
    # for i = 1:ngpus
    #     CUDA.Mem.free(gpu_buffers[i])
    #     CUDA.Mem.free(gpu_buffers_out[i])
    # end
    # CUDA.Mem.free(buf_host)
    #println(to)
    # return @view g_out[radius*t_group+1:end-radius*t_group,
    #         radius*t_group+1:end-radius*t_group, :]
end

function call_kernel(bdim, id, ngpus, dev_in, dev_out, t_steps, st_inst)
    bdimx = bdim
    bdimy = bdim
    radius = st_inst.max_radius
    dx = size(dev_in,1)
    dy = size(dev_in,2)
    bx = Int(floor((dx - 2*radius)/bdim))
    by = Int(floor((dy - 2*radius)/bdim))
    at_out = true
    for t = 0:t_steps-1
        offz_f = (id != 1)*t*radius
        offz_b = (id != ngpus)*t*radius
        args = (dev_in, dev_out, offz_f,offz_b, 0, 0)
        @cuda(blocks=(bx,by,1), threads=(bdimx,bdimy),
                    shmem=((bdimx+2*radius)*(bdimy+2*radius))*sizeof(Float32),
                    st_inst.kernel(args...))
        dev_in,dev_out = dev_out,dev_in
        at_out = !at_out
        #@show "kernel $t $id $ngpus"
    end

    return at_out
end

function ghost_phase(bdim, id, ngpus, dev_in, dev_out, t_steps, st_inst)
    bdimx = bdim
    bdimy = bdim
    radius = st_inst.max_radius
    dx = size(dev_in,1)
    dy = size(dev_in,2)
    bx = Int(floor((dx - 2*radius)/bdim))
    by = Int(floor((dy - 2*radius)/bdim))
    at_out = true
    for t = 0:t_steps-1
        offz_f = (id != 1)*t*radius
        offz_b = (id != ngpus)*t*radius
        args = (dev_in, dev_out, offz_f,offz_b, 0, 0)
        if id != 1
            offz_f =
            @cuda(blocks=(bx,by,1), threads=(bdimx,bdimy),
                        shmem=((bdimx+2*radius)*(bdimy+2*radius))*sizeof(Float32),
                        st_inst.kernel(args...))
        end
        if id != ngpus
            @cuda(blocks=(bx,by,1), threads=(bdimx,bdimy),
                        shmem=((bdimx+2*radius)*(bdimy+2*radius))*sizeof(Float32),
                        st_inst.kernel(args...))
        end
        dev_in,dev_out = dev_out,dev_in
        at_out = !at_out
        #@show "kernel $t $id $ngpus"
    end

    return at_out
end


text:
export solve_mao_cvx!

mutable struct SCPParam_Mao <: SCPParamSpecial
  ρ::Vector         # trust-region acceptance thresholds
  Δ_u0          		# initial trust region size for control
  λ              		# penalty weight
  α               	# trust region shrinkage & growth factor

  r_vec::Vector     # trust region ratios
  Δ_u_vec::Vector   # trust region size for control
end

function SCPParam_Mao(ρ, Δ_u0, λ, α)
	SCPParam_Mao(ρ, Δ_u0, λ, α, [0.], [Δ_u0])
end

function solve_mao_cvx!(SCPS::SCPSolution, SCPP::SCPProblem, solver="Mosek", max_iter=50, force=false; kwarg...)
	# Solves a sequential convex programming problem
	# Inputs:
	#		SCPS - SCP solution structure, including an initial trajectory
	#		SCPP - SCP problem definition structure
	#   max_iter - Set to enforce a maximum number of iterations (0 = unlimited)
	#   force - Set true to force further refinement iterations despite convergence

	if solver == "Mosek"
		set_default_solver(MosekSolver(; kwarg...))
	elseif solver == "Gurobi"
		set_default_solver(GurobiSolver(; kwarg...))
  else
		set_default_solver(SCSSolver(; kwarg...))
	end

	N = SCPP.N
	param = SCPP.param

	param.alg = SCPParam_Mao(SCPP.PD.model)
  α, Δ_u0 = param.alg.α, param.alg.Δ_u0
  r_vec, Δ_u_vec = param.alg.r_vec, param.alg.Δ_u_vec
  ρ0, ρ1, ρ2 = param.alg.ρ

	iter_cap = SCPS.iterations + max_iter
	SCPV = SCPVariables{Convex.Variable,Convex.Variable}(SCPP)
	SCPC = SCPConstraints(SCPP)

	update_model_params!(SCPP, SCPS.traj)
	push!(SCPS.J_true, cost_true(SCPS.traj, SCPS.traj, SCPP))
	param.obstacle_toggle_distance = model.clearance + 1. # TODO(ambyld): Generalize this, workspace-dependent

	first_time = true
  prob = minimize(0.)
	while (SCPS.iterations <= iter_cap)
		time_start = time_ns()

		# Set up, solve problem
		prob.objective = cost_full_convexified_mao(SCPV, SCPS.traj, SCPC, SCPP)
		prob.constraints = add_constraints_mao_cvx(SCPV, SCPS.traj, SCPC, SCPP)
    Convex.solve!(prob,warmstart=!first_time)
    first_time = false

		push!(SCPS.solver_status, prob.status)
    if prob.status != :Optimal
      warn("SCP-Mao failed find optimal solution")
		  push!(SCPS.iter_elapsed_times, (time_ns() - time_start)/10^9)
      return
    end

		# Recover solution
    new_traj = Trajectory(SCPV.X.value, SCPV.U.value, SCPV.Tf.value[1])
		push!(SCPS.convergence_measure, convergence_metric(new_traj, SCPS.traj, SCPP))
    push!(SCPS.J_full, prob.optval)
		SCPS.dual = get_dual_cvx(prob, SCPP, solver)

    # Check trust regions
		push!(r_vec, trust_region_ratio_mao(new_traj, SCPS.traj, SCPP))
    if r_vec[end] < ρ0
			push!(SCPS.accept_solution, false)
      push!(SCPS.J_true, SCPS.J_true[end])
    else
			push!(SCPS.accept_solution, true)
			push!(SCPS.J_true, cost_true(new_traj, SCPS.traj, SCPP))
			copy!(SCPS.traj, new_traj)
    end

    # Adjust trust region radius
    if r_vec[end] < ρ0 || r_vec[end] < ρ1
      # shrink trust region
      push!(Δ_u_vec, 1/α*Δ_u_vec[end])
    elseif r_vec[end] < ρ2
      # maintain trust region
      push!(Δ_u_vec, Δ_u_vec[end])
    else
      # expand trust region
      push!(Δ_u_vec, α*Δ_u_vec[end])
    end

    iter_elapsed_time = toq()
    push!(SCPS.iter_elapsed_times, iter_elapsed_time)
    SCPS.total_time += iter_elapsed_time
    SCPS.iterations += 1

    !SCPS.accept_solution[end] && continue

		if SCPS.convergence_measure[end] <= param.convergence_threshold
			SCPS.converged = true
			force ? continue : break
		end
	end
end

function add_constraints_mao_cvx(SCPV::SCPVariables, traj_prev::Trajectory, SCPC::SCPConstraints, SCPP::SCPProblem)
	constraints = Convex.Constraint[]
	Δ_u = SCPP.param.alg.Δ_u_vec[end]

	update_model_params!(SCPP, traj_prev)

  for (f, k, i) in SCPC.convex_state_eq
		constraints += f(SCPV, traj_prev, SCPP, k, i) == 0.
	end

  for (f, k, i) in (SCPC.convex_state_ineq..., SCPC.convex_control_ineq...)
		constraints += f(SCPV, traj_prev, SCPP, k, i) <= 0.
	end

  for (f, k, i) in SCPC.nonconvex_state_convexified_ineq
    # hack for enforcing SDF constraint when it passes out
    # 0 (i.e. too far from obstacle) rather than Convex expr
    new_constraint = f(SCPV, traj_prev, SCPP, k, i)
    typeof(new_constraint <= 0) == Bool && continue
    constraints += new_constraint <= 0.
	end

	for (f, k, i) in SCPC.control_trust_region_ineq
		constraints += f(SCPV, traj_prev, SCPP, k, i) - Δ_u <= 0.
	end

	return constraints
end

function cost_nonconvex_penalty_convexified_mao(traj, traj_prev::Trajectory, SCPC::SCPConstraints, SCPP::SCPProblem)
	λ = SCPP.param.alg.λ

	J = 0
  for (f, k, i) in SCPC.dynamics
		J += λ*norm(f(traj, traj_prev, SCPP, k, i), Inf)
	end

	return J
end

function cost_full_convexified_mao(traj, traj_prev::Trajectory, SCPC::SCPConstraints, SCPP::SCPProblem)
	cost_true_convexified(traj, traj_prev, SCPP) + cost_nonconvex_penalty_convexified_mao(traj, traj_prev, SCPC, SCPP)
end

ext: jl


text:
using LibSndFile
using FileIO
using LyonPassiveEar
using ClearStacktrace

vowel = load("vowel_a.wav")

mono = Float64.(vowel[:, 1]).data

coch = lyon_passive_ear(mono[1:100], sample_rate = 44100, decimation_factor = 64, differ = false)

lyon_passive_ear([1.0, 0, 0, 0, 0, 0], sample_rate = 400, decimation_factor = 1)


text:
# File holding the base regressor - a single decision tree

check_random_state(seed::Int) = MersenneTwister(seed)
check_random_state(rng::AbstractRNG) = rng

# Binary Tree as most basic form
mutable struct BinaryTree
    children_left::Vector{Int}
    children_right::Vector{Int}
    BinaryTree() = new([], [])
end

# Adding node; default node number is -1; will be changed
function add_node!(tree::BinaryTree)
    push!(tree.children_left, -1)
    push!(tree.children_right, -1)
    return
end

# Add a node to the left side of the parent node
function set_left_child!(tree::BinaryTree, node_id::Int64, child_id::Int)
    tree.children_left[node_id] = child_id
    return
end

# Add a node to the right side of the parent node
function set_right_child!(tree::BinaryTree, node_id::Int64, child_id::Int)
    tree.children_right[node_id] = child_id
    return
end

# Get the children nodes of a particular node
function get_children(tree::BinaryTree, node_id::Int64)
    return tree.children_left[node_id], tree.children_right[node_id]
end

# Check whether node is a terminal node (= leaf)
function is_leaf(tree::BinaryTree, node_id::Int64)
    return tree.children_left[node_id] == tree.children_right[node_id] == -1
end

###############################################################################
abstract type AbstractRegressor end

# Decision Tree Regressor object
mutable struct DTRegressor <: AbstractRegressor

    # Internal parameters
    num_nodes::Int
    binarytree::BinaryTree
    n_samples::Vector{Int} # Number of samples for each node
    split_dimensions::Vector{Union{Int, Nothing}} # Which variable was chosen for splitting
    split_values::Vector{Union{Float64, Nothing}} # Split-point values
    yhat::Vector{Union{Float64, Nothing}} # Predictor in each node
    mse::Vector{Union{Float64, Nothing}} # Reduction in MSE when splitting this node
    n_features::Union{Int, Nothing}
    pl::Vector{Float64} # Probability of falling into left node
    pr::Vector{Float64} # Probability of falling into right node
    depth_list::Vector{Union{Int, Nothing}} # Records depth for each split

    # externally set parameters
    max_depth::Union{Int, Nothing}
    max_features::Union{Int, Nothing}
    min_samples_leaf::Int
    random_state::Union{AbstractRNG, Int}
    α::Float64

    # Initialize object with default parameters
    DTRegressor(;
        max_depth=nothing,          # Not binding by default
        max_features=nothing,       # Set with parameters; = mtry
        min_samples_leaf = 1,       # ν = 1 -> Standard CART
        random_state=Random.GLOBAL_RNG,
        α = 0.01,)  = new(
            0, BinaryTree(), [], [], [], [], [], nothing, [], [], [],
        max_depth,
        max_features,
        min_samples_leaf,
        check_random_state(random_state),
        α)

end

# Fit the function to the data
function fit!(tree::DTRegressor, X::Matrix, Y::Matrix)
    # Sanity checks for dimensions
    @assert size(Y, 1) == size(X, 1) "X and Y have different lengths"
    @assert size(Y, 2) == 1 "Y is not a 1d matrix"

    # Set the internal value
    tree.n_features = size(X, 2)        # Number of features
    # If max_features is not set use all features
    tree.max_features = isnothing(tree.max_features) ? size(X, 2) : tree.max_features

    # Do the fitting - Apply Algorithm 1 starting with whole data set
    split_node!(tree, X, Y, 0)
    return
end

# Function for setting the defaults
function set_defaults!(tree::DTRegressor, Y::Matrix)

    # Default values are usually -2, it's only a placeholder which will be changes later
    # Enables quick checks to see whether everything is alright

    push!(tree.split_dimensions, -2)
    push!(tree.split_values, nothing)
    push!(tree.n_samples, size(Y, 1))   # Sample size of node

    # Default reduction in impurity - corresponds to equation (1)
    push!(tree.mse, 0.0)

    # Here we assign the estimator to be the mean of all Y values in a node
    push!(tree.yhat, mean(Y))

    # As default for node balancedness take -2
    push!(tree.pl, -2)
    push!(tree.pr, -2)
    add_node!(tree.binarytree)
end

# This is essentially Algorithm 1; it is applied recursively to each node
function split_node!(tree::DTRegressor, X::Matrix, Y::Matrix, depth::Int)

    tree.num_nodes += 1             # Increase number of total nodes within tree
    node_id = tree.num_nodes        # Reference for node is the node_id
    set_defaults!(tree, Y)          # Set default values
    push!(tree.depth_list, depth)   # Add depth of node to external depth_list

    # In case we have max_features choose randomly that many features to consider
    n_feature_splits = isnothing(tree.max_features) ? tree.n_features : min(tree.n_features, tree.max_features)

    features = randperm(tree.random_state, tree.n_features)[1:n_feature_splits]

    # Default decrease is 0
    max_decrease = 0
    # Loop over all features chosen previously
    @inbounds for j in features
        # Get highest decrease over all features j
        max_decrease = argmax_j(j, tree, X, Y, node_id, max_decrease, depth)
    end

    # If no feature returns a positive decrease, terminate node
    if max_decrease == 0
        return #No split was made
    end

    # Make child nodes if max_depth is not reached yet
    if isnothing(tree.max_depth) || (depth < tree.max_depth)

        x_split = X[:, tree.split_dimensions[node_id]] # Split dimension
        lhs = x_split .<= tree.split_values[node_id] # All obs going into left node
        rhs = x_split .> tree.split_values[node_id] # All obs going into right node

        # Assigning left node as a child node t(L) of the current node t
        set_left_child!(tree.binarytree, node_id, tree.num_nodes + 1)
        # Again calling Algorithm 1 on this child node
        split_node!(tree, X[lhs, :], Y[lhs,:], depth + 1)

        # Assigning right node as a child node t(R) of the current node t
        set_right_child!(tree.binarytree, node_id, tree.num_nodes + 1)
        # Again calling Algorithm 1 on this child node
        split_node!(tree, X[rhs, :], Y[rhs,:], depth + 1)
    end
    return
end

# Function for maximizing decrease in impurity over all features given
function argmax_j(j::Int, tree::DTRegressor, X::Matrix, Y::Matrix, node_id::Int, max_decrease, depth::Int)

    # Only split if min_samples_leaf condition is not binding
    # Times 2 since each of the child nodes needs to contain at least min_samples_leaf observations
    if (size(X, 1) > 2*tree.min_samples_leaf) && (std(Y)!=0)

        # Getting dimension j
        x = X[:, j]
        order = zeros(Int, length(x))           # Allocate space for odered variable
        sortperm!(order, x, alg=InsertionSort)  # Order along dimension j

        x_sorted, y_sorted = x[order], vec(Y[order])
        # Maximizing over split points s;
        tmp_decrease, split_index = argmax_s(y_sorted, tree, depth)

        # If new decrease is bigger than previous value, use this variable instead
        if tmp_decrease > max_decrease
            max_decrease = tmp_decrease
            tree.split_dimensions[node_id] = j                  # Record split dimensions
            tree.split_values[node_id] = x_sorted[split_index]  # Record split point

            # Asigning the number of nodes in left and right nodes for this split
            tree.pl[node_id] = split_index / length(x)      # P(t_L)
            tree.pr[node_id] = 1 - split_index / length(x)  # P(t_R)

        end
        return max_decrease
    else
        return 0
    end
end

# Maximizing over all split points for a given dimension
function argmax_s(arr::Vector{Float64}, tree::DTRegressor, depth::Int)

    # Re-assigining parameters for easier use
    α = tree.α
    min_samples = tree.min_samples_leaf

    # Allocating space
    Δ = 0       # This is Δ(s, t) from equation (2); default is 0
    ind = 0     # Index of split point

    # Length of array then gives probabilities of where element lands
    n = length(arr)
    # Only look within the min_samples interval
    @inbounds for i in min_samples:n-min_samples-1

        # Function for calculating eqaution (2)
        Δ_tmp = get_Δ(arr, i, n, α, depth)
        # Only updating if decrease in variance is higher
        if Δ_tmp > Δ
            Δ = Δ_tmp
            ind = i
        end
    end

    # Retruning decrease and index
    return Δ, ind
end

# This is where the splitting rule comes into play
function get_Δ(arr::Vector{Float64}, ind::Int64, n::Int64, α::Float64, depth::Int)
    # Calculating mean in left and right child node respectively
    Y_left = mean(arr[1:ind])
    Y_rigth = mean(arr[ind+1:end])

    # Decrease in impurity is given by P(t_L)*P(t_R) * (Y_L - Y_R)^2
    # The same as equation (2) but rewritten for efficiency
    Δ = (ind/n)*(1 - ind/n)*(Y_left - Y_rigth)^2

    # This is the weighted splitting rule
    # If α = 0 then we have the standard case

    # Chosen function; α=a compared to paper! Corresponds to equation (18)
    α_k = depth^α

    # Return decrease in impurity from weighted splitting rule
    return (4*(ind/n)*(1 - ind/n))^(α_k) * Δ    # Equation (17)
end

###############################################################################
# Functions used once the tree is built

# Predicting the values of one observation
function predict_arr(tree::DTRegressor, arr::Vector)

    # Start with first node and then drop down through the tree
    next_node = 1
    # While node is not terminal node drop down deeper
    while !is_leaf(tree.binarytree, next_node)
        left, right = get_children(tree.binarytree, next_node)
        next_node = arr[tree.split_dimensions[next_node]] <= tree.split_values[next_node] ? left : right
    end

    # Once terminal node is reached, return yhat from that terminal node
    return tree.yhat[next_node]
end

# Predict values for a whole matrix
function predict(tree::DTRegressor, X::Matrix)
    @assert tree.n_features == size(X, 2) "# of features are not the same"

    # Throw error if tree is not fitted
    if tree.num_nodes == 0
        error("Tree is not fitted")
    end

    # Allocate space for predictions
    prediction = Vector{Float64}(undef, size(X, 1))

    # Transpose matrix for higher efficiency; looping over columns is faster than over rows
    X_prime = transpose(X)

    # Predicting values for each observation of data set
    @inbounds for ind in 1:size(X_prime, 2)
        prediction[ind] = predict_arr(tree, X_prime[:,ind])
    end
    return prediction
end

# Function for getting the selection frequency of strong variables
# Strong variables are all from 1-S; supply S=var_index
# Used for Figures 3 and 5
function strong_selection_freq(tree::DTRegressor, var_index::Int)

    # Allocate space
    res_arr = 0

    # Get splitting dimensions of all nodes which are not terminal nodes
    tmp = tree.split_dimensions[tree.split_dimensions .!= -2]

    # Selection frequency
    res_arr = sum(tmp .<= var_index)/length(tmp)

    return res_arr
end

max_stars_repo_name: christianhilscher/emma

max_issues_repo_name: christianhilscher/emma

max_forks_repo_name: christianhilscher/emma


text:
module TestSetfield

include("test_examples.jl")
include("test_setmacro.jl")
include("test_core.jl")
include("test_functionlenses.jl")
include("test_settable.jl")
include("test_staticarrays.jl")
include("test_kwonly.jl")
include("test_quicktypes.jl")
include("perf.jl")
end  # module


text:
rm( "base/baseConfig2", recursive=true, force=true )
mpSim, isOkay = runFromExcel( joinpath( "base", "baseConfig2" ),
    showInfo=false )
report = nodePopReport( mpSim, 12, "Junior" )
results = Int.( report[:, "Junior"] )
@test all( results .== [189, 187, 123, 124, 100, 151, 164, 138, 145, 115, 160,
    157, 150, 139, 100, 181, 169, 111, 140, 122, 121, 124, 121, 154, 162, 106] )

rm( "base/baseConfig2b", recursive=true, force=true )
mpSim, isOkay = runFromExcel( joinpath( "base", "baseConfig2b" ),
    showInfo=false )
report = nodePopReport( mpSim, 12, "Junior" )
@test all( Int.( report[:, "Junior"] ) .== results )

rm( "base/baseConfig3", recursive=true, force=true )
mpSim, isOkay = runFromExcel( joinpath( "base", "baseConfig3" ),
    showInfo=false )
report = nodePopReport( mpSim, 12, "Junior" )
@test Int.( report[:, "Junior"] ) != results

rm( "base/baseConfig4", recursive=true, force=true )
mpSim, isOkay = runFromExcel( joinpath( "base", "baseConfig4" ),
    showInfo=false )
report = nodePopReport( mpSim, 12, "Junior" )
@test Int.( report[:, "Junior"] ) != results

rm( "base/baseConfig5", recursive=true, force=true )
mpSim, isOkay = runFromExcel( joinpath( "base", "baseConfig5" ),
    showInfo=false )
report = nodePopReport( mpSim, 12, "Junior" )
@test Int.( report[:, "Junior"] ) != results


text:
#GracePlot demonstration 4: simple 2x2 array of subplots
#-------------------------------------------------------------------------------

using GracePlot

#==Input data
===============================================================================#
x=collect(-10:0.1:10)
y = []
for i in 1:4
	push!(y, (x.^i)./(10^i)) #Normalized powers of x
end
titles = ["Linear", "Quadratic", "Cubic", "Quartic"]

#=="Defaults"
===============================================================================#
template = GracePlot.template("plot2x2thick_mono")

#==Generate plot
===============================================================================#
plot = GracePlot.new(fixedcanvas=true, template=template)
#Add 4 subplots:
let g, ds #HIDEWARN_0.7
for gidx in 0:3
	g = add(plot, subtitle = titles[gidx+1])

	#Add 10 datasets:
	for i in 1:10
		ds = add(g, x, i.*y[gidx+1])
	end

	autofit(g)
end
end

#Finalize:
redraw(plot)

#Last line


text:
#####
##### Formulas assuming dry (unsaturated) adiabatic (i.e. isentropic) processes
#####

export DryAdiabaticProcess

export air_pressure_given_θ, air_pressure, air_temperature

"""
    DryAdiabaticProcess

For dispatching to isentropic formulas
"""
struct DryAdiabaticProcess end

"""
    air_pressure_given_θ(param_set, θ::FT, Φ::FT, ::DryAdiabaticProcess)

The air pressure for an isentropic process, where

 - `param_set` an `AbstractParameterSet`, see the [`Thermodynamics`](@ref) for more details
 - `θ` potential temperature
 - `Φ` gravitational potential
"""
function air_pressure_given_θ(
    param_set::APS,
    θ::FT,
    Φ::FT,
    ::DryAdiabaticProcess,
) where {FT <: AbstractFloat}
    _MSLP::FT = MSLP(param_set)
    _R_d::FT = R_d(param_set)
    _cp_d::FT = cp_d(param_set)
    return _MSLP * (1 - Φ / (θ * _cp_d))^(_cp_d / _R_d)
end

"""
    air_pressure(param_set, T::FT, T∞::FT, p∞::FT, ::DryAdiabaticProcess)

The air pressure for an isentropic process, where

 - `param_set` an `AbstractParameterSet`, see the [`Thermodynamics`](@ref) for more details
 - `T` temperature
 - `T∞` ambient temperature
 - `p∞` ambient pressure
"""
function air_pressure(
    param_set::APS,
    T::FT,
    T∞::FT,
    p∞::FT,
    ::DryAdiabaticProcess,
) where {FT <: AbstractFloat}
    _kappa_d::FT = kappa_d(param_set)
    return p∞ * (T / T∞)^(FT(1) / _kappa_d)
end

"""
    air_temperature(param_set, p::FT, θ::FT, Φ::FT, ::DryAdiabaticProcess)

The air temperature for an isentropic process, where

 - `param_set` an `AbstractParameterSet`, see the [`Thermodynamics`](@ref) for more details
 - `p` pressure
 - `θ` potential temperature
"""
function air_temperature(
    param_set::APS,
    p::FT,
    θ::FT,
    ::DryAdiabaticProcess,
) where {FT <: AbstractFloat}
    _R_d::FT = R_d(param_set)
    _cp_d::FT = cp_d(param_set)
    _MSLP::FT = MSLP(param_set)
    return (p / _MSLP)^(_R_d / _cp_d) * θ
end

max_stars_repo_name: mwarusz/CLIMA

max_issues_repo_name: mwarusz/CLIMA

max_forks_repo_name: mwarusz/CLIMA


text:
function analyze_run(rs::RunResults,time_trace)
	nL = length(rs.data)
	i=1
	figure()
	for (lname, data) in rs.data
		subplot(1,nL,i)
		# 100 randomly selected neurons
		plot_raster(convert(BitArray{2},data["didSpike"][randperm(100),:]))
		title(lname)
		xticks(time_trace)
		i+=1
	end
	show()
end


text:
module steam_governor
using ReachabilityAnalysis
@taylorize function steam_governor!(du, u, p, t)
    local ϵ = 3.0
    local α = 1.0
    local β = 1.0

    du[1] = u[2]
    du[2] = u[3]^2 * sin(u[1]) * cos(u[1]) - sin(u[1]) - ϵ*u[2]
    du[3] = α*(cos(u[1]) - β)

    return du
end

function model(X0)
    S = @system(x' = steam_governor!(x), dim:3)
    return IVP(S, X0)
end
end # module

max_stars_repo_name: mforets/ReachabilityLibrary

max_issues_repo_name: mforets/ReachabilityLibrary

max_forks_repo_name: mforets/ReachabilityLibrary


text:
export Equiv

"""
    Equiv(nodes::Array{Node,1})

`Equiv` objects `show` a FastHenry .equiv command
"""
struct Equiv <: Element
  nodes :: Array{Node,1}
end


text:
"""
    ComputeDensityFunction(z::Float64, AnalitycalDensity::AnalitycalDensity)

This function returns the source density for a given redshift ``z``.
"""
function ComputeDensity(z::Float64, AnalitycalDensity::AnalitycalDensity)
    return ((z/AnalitycalDensity.Z0)^2)*exp(-(z/AnalitycalDensity.Z0)^(3. / 2.))*
    AnalitycalDensity.Normalization
end

"""
    NormalizeAnalitycalDensity!(AnalitycalDensity::AnalitycalDensity)

This function normalize AnalitycalDensity in order to have the correct
value of the surface density once integrated.
"""
function NormalizeAnalitycalDensity!(AnalitycalDensity::AnalitycalDensity)
    int, err = QuadGK.quadgk(x -> ComputeDensity(x, AnalitycalDensity),
    AnalitycalDensity.ZMin, AnalitycalDensity.ZMax, rtol=1e-12)
    AnalitycalDensity.Normalization *= (AnalitycalDensity.SurfaceDensity/int)
end

"""
    ComputeInstrumentResponse(z::Float64, zp::Float64,
    InstrumentResponse::InstrumentResponse)

This function computes the probability that we actually measure a redshift
``z_p`` if the real redshift is ``z``.
"""
function ComputeInstrumentResponse(z::Float64, zp::Float64,
    InstrumentResponse::InstrumentResponse)
    prob_z =
    (1-InstrumentResponse.fout)/(sqrt(2*pi)*InstrumentResponse.σb*(1+z))*
    exp(-0.5*((z-InstrumentResponse.cb*zp-InstrumentResponse.zb)/
    (InstrumentResponse.σb*(1+z)))^2)+
    InstrumentResponse.fout/(sqrt(2*pi)*InstrumentResponse.σo*(1+z))*
    exp(-0.5*((z-InstrumentResponse.co*zp-InstrumentResponse.zo)/
    (InstrumentResponse.σo*(1+z)))^2)
    return prob_z
end

"""
    ComputeConvolvedDensity(z::Float64, i::Int64, ConvolvedDensity::AbstractConvolvedDensity,
    AnalitycalDensity::AnalitycalDensity, InstrumentResponse::InstrumentResponse)

This function computes the Convolved density function for a single bin at a
given redshift ``z``.
"""
function ComputeConvolvedDensity(z::Float64, i::Int64, ConvolvedDensity::AbstractConvolvedDensity,
    AnalitycalDensity::AnalitycalDensity, InstrumentResponse::InstrumentResponse)
    int, err = QuadGK.quadgk(x -> ComputeInstrumentResponse(z, x,
    InstrumentResponse),
    ConvolvedDensity.ZBinArray[i],
    ConvolvedDensity.ZBinArray[i+1], rtol=1e-12)
    return int*ComputeDensity(z,
    AnalitycalDensity) * ConvolvedDensity.DensityNormalizationArray[i]
end

"""
    NormalizeConvolvedDensity!(ConvolvedDensity::AbstractConvolvedDensity,
    AnalitycalDensity::AnalitycalDensity,InstrumentResponse::InstrumentResponse,
    CosmologicalGrid::CosmologicalGrid)

This function normalizes ConvolvedDensity such that the integrals of the
convolved densities are normalized to 1.
"""
function NormalizeConvolvedDensity!(ConvolvedDensity::AbstractConvolvedDensity,
    AnalitycalDensity::AnalitycalDensity, InstrumentResponse::InstrumentResponse,
    CosmologicalGrid::CosmologicalGrid)
    for idx in 1:length(ConvolvedDensity.DensityNormalizationArray)
        int, err = QuadGK.quadgk(x -> ComputeConvolvedDensity(x, idx,
        ConvolvedDensity, AnalitycalDensity, InstrumentResponse),
        first(CosmologicalGrid.ZArray),
        last(CosmologicalGrid.ZArray), rtol=1e-12)
        ConvolvedDensity.DensityNormalizationArray[idx] /= int
    end
end

function ComputeSurfaceDensityBins!(ConvolvedDensity::AbstractConvolvedDensity,
    AnalitycalDensity::AnalitycalDensity)
    for idx in 1:length(ConvolvedDensity.SurfaceDensityArray)
        int, err = QuadGK.quadgk(x -> ComputeDensity(x, AnalitycalDensity),
        ConvolvedDensity.ZBinArray[idx], ConvolvedDensity.ZBinArray[idx+1], rtol=1e-12)
        ConvolvedDensity.SurfaceDensityArray[idx] = int
    end
end

"""
    ComputeConvolvedDensityGrid!(CosmologicalGrid::CosmologicalGrid,
    ConvolvedDensity::AbstractConvolvedDensity, AnalitycalDensity::AnalitycalDensity,
    InstrumentResponse::InstrumentResponse)

This function computes the convolved density function for all tomographic bins
on the ``z``-grid provided by CosmologicalGrid.
"""
function ComputeConvolvedDensityGrid!(CosmologicalGrid::CosmologicalGrid,
    ConvolvedDensity::AbstractConvolvedDensity, AnalitycalDensity::AnalitycalDensity,
    InstrumentResponse::InstrumentResponse)
    ConvolvedDensity.DensityGridArray = zeros(Float64,length(ConvolvedDensity.ZBinArray)-1,
    length(CosmologicalGrid.ZArray))
    for idx_ZBinArray in 1:length(ConvolvedDensity.ZBinArray)-1
        for idx_ZArray in 1:length(CosmologicalGrid.ZArray)
            ConvolvedDensity.DensityGridArray[idx_ZBinArray, idx_ZArray] =
            ComputeConvolvedDensity(CosmologicalGrid.ZArray[idx_ZArray],
            idx_ZBinArray, ConvolvedDensity, AnalitycalDensity,
            InstrumentResponse)
        end
    end
end

function ShiftConvolvedDensityGrid!(CosmologicalGrid::CosmologicalGrid,
    ConvolvedDensity::AbstractConvolvedDensity)
    for idx_ZBinArray in 1:length(ConvolvedDensity.ZBinArray)-1
        InterpConvDens = Dierckx.Spline1D(CosmologicalGrid.ZArray,
        ConvolvedDensity.DensityGridArray[idx_ZBinArray,:], k=4,
        bc="extrapolate")
        for idx_ZArray in 1:length(CosmologicalGrid.ZArray)
            ConvolvedDensity.DensityGridArray[idx_ZBinArray,idx_ZArray] =
            InterpConvDens(CosmologicalGrid.ZArray[idx_ZArray]+
            ConvolvedDensity.ShiftArray[idx_ZBinArray])
        end
    end
end


text:
#=
    The File contains Sparsik tests
=#

#------------------------------------------------------------------------------

@testset "W. square, det" begin
    # here the correctness is checked
    # Ay = b

    A = from_dense([1 0 0;
                    0 2 0;
                    0 0 3;],
                    Nemo.QQ)
    b = from_dense([3, 2, 1], Nemo.QQ)

    y = square_nonsingular_deterministic_wiedemann(A, b)

    @test y == from_dense([Nemo.QQ(3), Nemo.QQ(1), Nemo.QQ(1, 3)], Nemo.QQ)

    # -----

    Z13 = GF(13)
    ZZ =  GF(2^31 - 1)

    # a series of tests
    sizes = [10, 10, 10, 15, 30, 50, 50, 75, 100]
    for n in sizes
        A, b = 0, 0
        while true
            A = random_sparsik((n, n), ZZ, density=0.3)
            b = random_sparsik((n, ), ZZ)
            try
                y = square_nonsingular_deterministic_wiedemann(A, b)
            catch e
                if isa(e, SingularException)
                    continue
                else
                    rethrow(e)
                end
            end

            break
        end

        X = [convert(Int, x) for x in to_dense(A)]
        #if iszero(det(X))
            #error("W. did't notice singularity")
        #end

        @test A * y == b
    end

    A = from_dense([1 1; 1 1;], Nemo.QQ)
    b = from_dense([1, 2], Nemo.QQ)
    @test_throws SingularException square_nonsingular_deterministic_wiedemann(A, b)
    @test_throws SingularException square_nonsingular_randomized_wiedemann(A, b)

end

@testset "W. matrix minpolynomial" begin

end

@info "OK"


text:
# module NAssets
using Agents: vertices
using MetaGraphs: add_vertex!, has_edge
using Agents, AgentsPlots, Plots, LightGraphs, MetaGraphs, GraphPlot, GraphRecipes, NetworkLayout
using Tables, DataFrames
using CSV, JSON, Serialization, DelimitedFiles
using BritishNationalGrid
using ZipFile, Shapefile
using Random
using Match
using LinearAlgebra
using StatsBase
using Distributions
using StatsPlots
using SparseArrays
using Laplacians
using DataStructures
using RollingFunctions
using BenchmarkTools
using Statistics
using Logging,LoggingExtras, LoggingFacilities
using Dates,TimeZones

#Core structures
include("core/core_structs.jl")
#graph-related
include("core/graph_functions.jl")
#events
include("eve/artificial_events.jl")
#various util functions
include("utils/util_functions.jl")
#logging
include("utils/logging_functions.jl")
#Plotting functions
include("utils/plotting_functions.jl")
# time-to-event functions
include("utils/tte_functions.jl")
#running functions
include("utils/running_functions.jl")

#maintenance model
include("phy/maintenance_model.jl")

include("ntw/of_switch.jl")
include("ntw/of_control.jl")
include("ctl/agent_control.jl")

include("phy/physical_model.jl")

include("ntw/network_model.jl")
include("phy/geo_model.jl")

#Main Functions
include("model/netManFunctions.jl")

#Agents.jl function implementation for this model
include("model/netManModel.jl")

#Basic queries
#include("model/queries_basic.jl")
#Multiple queries
include("ctl/queries_multiple.jl")

export load_run_configs, single_run_with_logging

configs = []
exp_params = deserialize("data/services.bin")

for exp_param in exp_params
    for cg=6:6
        push!(configs, ( seed = first(exp_param)
                    ,ctl_model= GraphModel(cg) #NAssets.GraphModel(6)
                    ,ntw_topo = GraphModel(6) #NAssets.GraphModel(6)
                    ,size=100
                    ,n_steps=200
                    ,deterioration = 0. # physical deterioration parameter. 0: switched off.
                    # random asset drops
                    ,drop_proportion=0.2 #proportion of assets to randomly fail. 0: switched off.
                    ,drop_stabilisation=50 # time (from simulation start) to allow assets to stabilise before trigger random drops.
                    ,prob_random_walks = 1. #BA(6): 0.5#star(5): 1. #complete(3): 0.3 #ring(2): 1. # grid(4): 0.8 #0.5
                    ,benchmark = false
                    ,animation = false
                    ,k=2
                    ,B = 0.5
                    ,ctl_k=2
                    ,ctl_B = 0.5
                    ,custom_topo = nothing
                    ,ctl_custom_topo = nothing
                    ,mnt_policy = 1
                    ,mnt_wc_duration = 20
                    ,mnt_bc_duration = 10
                    ,mnt_wc_cost = 150.
                    ,mnt_bc_cost = 100.
                    ,traffic_dist_params = [1,0.05]
                    ,traffic_proportion = 0.01
                    ,clear_cache_graph_freq = 50#50,#25, # How often the ntw graph is cleared to initial state, 0: no cache. A value of 10, is not enough in a 16 mesh network to find paths when queries are not repeated, prob_eq_query.
                    ,interval_tpt = 10 #interval used to calculate tpt
                    ,pkt_size = 1#0.065, # (in MB) pkt size  IP between 21 is 65536 bytes Ref: Internet Core Protocols: The Definitive Guide by Eric Hall
                    ,pkt_per_tick = 2000#0, # How many packets are processsed per tick. #TODO which number is reasonable?
                    # I am setting this to 2000 as the expectations is nes
                    # are able to process. Check references e.g. Nokia SR 7750.
                    ,max_queue_ne = 300#700 #This indicates how many pkts/msgs can be stored in tick to be processed the next tick
                    ,ofmsg_reattempt=10#4,# greater number to avoid duplicated install flows
                    ,max_cache_paths = 2
                    ,data_dir = "data/exp1/"
                    ,plots_dir = "plots/exp1/"
                    # ,ntw_services = [(1,7),(4,1),(5,14),(12,8)]
                    ,ntw_services = last(exp_param)
                ))
    end
end

BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

Threads.@threads for config in configs
    single_run_with_file_logging(config)
end
# end # module


text:
import LinearAlgebra: mul!, axpy!
import SparseArrays as SArrays

export FspMatrixSparse, get_parameters, get_states, get_rowcount, get_colcount, get_propensities

"""
The transition rate matrix of the finite Markov chain approximation to the CME (including sink states). This type is based on the sparse representation of the truncated state space.
"""
Base.@kwdef mutable struct FspMatrixSparse{NS,NR,IntT<:Integer,RealT<:AbstractFloat} <: FspMatrix
    parameters::Vector{Any}
    states::Vector{MVector{NS,IntT}}

    rowcount::IntT
    colcount::IntT

    propensities::Vector{<:Propensity}

    timeinvariant_propensity_ids::Vector{IntT}
    timeinvariant_matrix::SparseMatrixCSC

    t_cache::RealT
    jointtv_propensity_ids::Vector{Int}
    jointtv_matrices::Vector{SparseMatrixCSC}

    separabletv_propensity_ids::Vector{IntT}
    separabletv_factormatrices::Vector{SparseMatrixCSC}
end
# Accessors
get_parameters(A::FspMatrixSparse) = A.parameters
get_states(A::FspMatrixSparse) = A.states
get_rowcount(A::FspMatrixSparse) = A.rowcount
get_colcount(A::FspMatrixSparse) = A.colcount
get_propensities(A::FspMatrixSparse) = A.propensities
get_time(A::FspMatrixSparse) = A.t_cache
get_jointtv_propensity_ids(A::FspMatrixSparse) = A.jointtv_propensity_ids
get_separabletv_propensity_ids(A::FspMatrixSparse) = A.separabletv_propensity_ids
get_timeinvariant_propensity_ids(A::FspMatrixSparse) = A.timeinvariant_propensity_ids

"""
FspMatrixSparse{RealT}(space::AbstractStateSpaceSparse{NS,NR,IntT}, propensity_functions::Vector{<:Propensity}; θ = [])

Return an instance of `FspMatrixSparse` type with nonzero entries of type `RealT`, based on a sparse state space of type `AbstractStateSpaceSparse` and a vector `propensity_functions` of propensity functions.

# See also
`Propensity`
"""
function FspMatrixSparse{RealT}(space::AbstractStateSpaceSparse{NS,NR,IntT,SizeT}, propensity_functions::Vector{<:Propensity}; parameters = []) where {NS,NR,IntT<:Integer,RealT<:AbstractFloat,SizeT<:Integer}
    state_count = get_state_count(space)
    sink_count = get_sink_count(space)
    rowcount = state_count + sink_count
    colcount = rowcount

    timeinvariant_propensity_ids = Vector{IntT}()
    jointtv_propensity_ids = Vector{IntT}()
    separabletv_propensity_ids = Vector{IntT}()

    for (ir, propensity) in enumerate(propensity_functions)
        if !istimevarying(propensity)
            push!(timeinvariant_propensity_ids, ir)
        else
            push!(istimeseparable(propensity) ? separabletv_propensity_ids : jointtv_propensity_ids, ir)
        end
    end

    rowindices = Vector{IntT}()
    colindices = Vector{IntT}()
    vals = Vector{RealT}()
    for reactionidx in timeinvariant_propensity_ids
        ridxs, cidxs, vs = _generate_sparsematrix_entries(space, propensity_functions[reactionidx].f, reactionidx, parameters, valtype=RealT)
        rowindices = vcat(rowindices, ridxs)
        colindices = vcat(colindices, cidxs)
        vals = vcat(vals, vs)
    end
    timeinvariant_matrix = !isempty(timeinvariant_propensity_ids) ? sparse(rowindices, colindices, vals, rowcount, rowcount) : nothing

    separabletv_factormatrices = Vector{SparseMatrixCSC}()
    for reactionidx in separabletv_propensity_ids
        rowindices, colindices, vals = _generate_sparsematrix_entries(space, propensity_functions[reactionidx].statefactor, reactionidx, parameters, valtype=RealT)
        push!(
            separabletv_factormatrices,
            sparse(rowindices, colindices, vals, rowcount, rowcount)
        )
    end

    jointtv_matrices = Vector{SparseMatrixCSC}()
    for reactionidx in jointtv_propensity_ids
        rowindices, colindices, vals = _generate_sparsematrix_entries(space, nothing, reactionidx, parameters, valtype=RealT)
        push!(
            jointtv_matrices,
            sparse(rowindices, colindices, vals, rowcount, rowcount)
        )
    end

    return FspMatrixSparse{NS,NR,IntT,RealT}(
        parameters = parameters,
        t_cache = -Inf,
        states = deepcopy(space.states),
        rowcount = rowcount,
        colcount = colcount,
        propensities = propensity_functions,
        timeinvariant_propensity_ids = timeinvariant_propensity_ids,
        timeinvariant_matrix = timeinvariant_matrix,
        jointtv_propensity_ids = jointtv_propensity_ids,
        jointtv_matrices = jointtv_matrices,
        separabletv_propensity_ids = separabletv_propensity_ids,
        separabletv_factormatrices = separabletv_factormatrices
    )
end

"""
FspMatrixSparse(space::AbstractStateSpaceSparse, propensity_functions::Vector; θ = [])

Return an instance of `FspMatrixSparse` type with double-precision(i.e.`Float64`) nonzero entries, based on a sparse state space of type `AbstractStateSpaceSparse` and a vector `propensity_functions` of propensity functions.
"""
function FspMatrixSparse(space::AbstractStateSpaceSparse{NS,NR,IntT,SizeT}, propensity_functions::Vector{<:Propensity}; parameters = []) where {NS,NR,IntT<:Integer,SizeT<:Integer}
    return FspMatrixSparse{Float64}(space, propensity_functions; parameters)
end

function _generate_sparsematrix_entries(space::AbstractStateSpaceSparse{NS,NR,IntT,SizeT}, statefactor::Union{Nothing,Function}, reactionidx::Integer, θ::Vector = []; valtype::Type{RealT}=Float64) where {NS,NR,IntT<:Integer,RealT<:AbstractFloat,SizeT<:Integer}
    state_count = get_state_count(space)
    rowindices = zeros(IntT, 2 * state_count)
    colindices = zeros(IntT, 2 * state_count)
    vals = zeros(valtype, 2 * state_count)

    # Not optimal, should evaluate propensity function only once per CME state
    @simd for idx = 1:state_count
        # Populate the diagonal entries
        dval = (statefactor === nothing) ? 0.0 : statefactor(space.states[idx], θ)
        @inbounds rowindices[idx] = idx
        @inbounds colindices[idx] = idx
        @inbounds vals[idx] = -1.0 * dval

        # Populate the row corresponding to a reachable sink state (if any)
        sink_idx = space.sink_connectivity[idx][reactionidx]
        if sink_idx ≠ 0
            @inbounds colindices[state_count+idx] = idx
            @inbounds rowindices[state_count+idx] = state_count + sink_idx
            @inbounds vals[state_count+idx] = dval
        end
    end
    @simd for idx = 1:state_count
        # If nonzero, this will be the column index of the nonzero entry on the idx-th row
        cidx = space.state_connectivity[idx][reactionidx]
        if cidx ≠ 0
            @inbounds rowindices[cidx+state_count] = idx
            @inbounds colindices[cidx+state_count] = cidx
            @inbounds vals[cidx+state_count] = -1.0 * vals[cidx]
        end
    end
    return rowindices[colindices.≠0], colindices[colindices.≠0], vals[colindices.≠0]
end

function _update_sparsematrix!(matrix::SparseMatrixCSC, states::Vector, propensity::Any, t::AbstractFloat, θ::Vector = [])
    m = matrix.m
    n = length(states)
    rowvals = SArrays.getrowval(matrix)
    nzvals = SArrays.getnzval(matrix)
    @simd for icol in 1:n
        val = propensity.f(t, states[icol], θ)
        for i in nzrange(matrix, icol)
            @inbounds nzvals[i] = (rowvals[i] == icol) ? -1.0 * val : val
        end
    end
    nothing
end

# Functions to query the dimensions of the matrix
import Base: size

"""
Return the size of the FSP matrix.
"""
function size(A::FspMatrixSparse)
    return (A.rowcount, A.colcount)
end

"""
Return number of rows (if `dim=1`) or columns (if `dim=2`) of the FSP matrix.
"""
function size(A::FspMatrixSparse, dim::Integer)
    if !(1 ≤ dim ≤ 2)
        throw(ArgumentError("Second argument must be either 1 or 2."))
    end
    return (dim == 1) ? A.rowcount : A.colcount
end

# Matrix-vector multiplications
"""
matvec!(out::Vector{RealT}, t::AbstractFloat, A::FspMatrixSparse, v::Vector{RealT})

Perform matrix-vector multiplication for the FSP matrix `A` at time `t` and vector `v`. The result is written to the vector `out`.
"""
function matvec!(out, t, A::FspMatrixSparse, v)
    if A.timeinvariant_matrix isa Nothing
        out .= 0.0
    else
        # out .= A.timeinvariant_matrix*v
        SArrays.mul!(out, A.timeinvariant_matrix, v)
    end
    θ = A.parameters
    for (i, r) in enumerate(A.separabletv_propensity_ids)
        # out .+= A.propensities[r].tfactor(t, θ...)*A.separabletv_factormatrices[i]*v
        SArrays.mul!(out, A.separabletv_factormatrices[i], v, A.propensities[r].tfactor(t, θ), 1)
    end
    needupdate = false
    if t ≠ A.t_cache
        needupdate = true
        A.t_cache = t
    end
    for (i, r) in enumerate(A.jointtv_propensity_ids)
        (needupdate) && _update_sparsematrix!(A.jointtv_matrices[i], A.states, A.propensities[r],
            t, θ)
        # out .+= A.jointtv_matrices[i]*v
        SArrays.mul!(out, A.jointtv_matrices[i], v, 1.0, 1.0)
    end
    return nothing
end

export matvecadd!

"""
    matvecadd!(out, t, A::FspMatrixSparse, v)

Perform the in-place update `out = out + A(t)*v`.
"""
function matvecadd!(out, t, A::FspMatrixSparse, v)
    if !(A.timeinvariant_matrix isa Nothing    )
        SArrays.mul!(out, A.timeinvariant_matrix, v, 1, 1)
    end
    θ = A.parameters
    for (i, r) in enumerate(A.separabletv_propensity_ids)
        # out .+= A.propensities[r].tfactor(t, θ...)*A.separabletv_factormatrices[i]*v
        SArrays.mul!(out, A.separabletv_factormatrices[i], v, A.propensities[r].tfactor(t, θ), 1)
    end
    needupdate = false
    if t ≠ A.t_cache
        needupdate = true
        A.t_cache = t
    end
    for (i, r) in enumerate(A.jointtv_propensity_ids)
        (needupdate) && _update_sparsematrix!(A.jointtv_matrices[i], A.states, A.propensities[r],
            t, θ)
        # out .+= A.jointtv_matrices[i]*v
        SArrays.mul!(out, A.jointtv_matrices[i], v, 1.0, 1.0)
    end
    return nothing
end

"""
w = matvec(out::Vector{RealT}, t::AbstractFloat, A::FspMatrixSparse, v::Vector{RealT})

Perform matrix-vector multiplication for the FSP matrix `A` at time `t` and vector `v` and return `w` as the result.
"""
function matvec(t, A::FspMatrixSparse, v)
    w = similar(v)
    matvec!(w, t, A, v)
    return w
end

function (A::FspMatrixSparse)(t::AbstractFloat)
    return (t, A)
end

import Base: *
function *(A_at_t::Tuple{AbstractFloat,FspMatrixSparse}, v::Vector{RealT}) where {RealT<:AbstractFloat}
    t = A_at_t[1]
    A = A_at_t[2]
    return matvec(t, A, v)
end

function *(A::FspMatrixSparse, v::Vector{RealT}) where {RealT<:AbstractFloat}
    return matvec(0.0, A, v)
end


text:
mutable struct Decoder
    chardict::Dict
    tagset
    nn
end

function Decoder(vocabfile::String, trainfile::String, testfile::String, nepochs::Int, learnrate::Float64, batchsize::Int)
    chars = open(readlines, vocabfile)
    chardict = Dict("UNKNOWN"=>1, "[EOL]"=>2)

    tagset = BIOES()
    traindata = readdata(trainfile, chardict, tagset)
    testdata = readdata(testfile, chardict, tagset)
    testdata = batch(testdata, 100)
    charembeds = embeddings(Float32, length(chardict), 20, init_w=Normal(0,0.05))
    nn = NN(length(chardict), length(tagset.tag2id))

    @info "#Training examples:\t$(length(traindata))"
    @info "#Testing examples:\t$(length(testdata))"
    @info "#Chars:\t$(length(chardict))"
    @info "#Wordtags:\t$(length(tagset.tag2id))"

    opt = SGD()
    for epoch = 1:nepochs
        println("Epoch:\t$epoch")
        opt.rate = learnrate * batchsize / sqrt(batchsize) / (1 + 0.05*(epoch-1))
        println("Learning rate: $(opt.rate)")

        shuffle!(traindata)
        batches = batch(traindata, batchsize)
        prog = Progress(length(batches))
        loss = 0.0
        for i in 1:length(batches)
            c, t1, t2 = batches[i]
            y = nn(c, t1, t2)
            loss += sum(y.data)
            params = gradient!(y)
            foreach(opt, params)
            ProgressMeter.next!(prog)
        end
        loss /= length(batches)
        println("Loss:\t$loss")
        continue

        # test
        println("Testing...")
        preds = Int[]
        golds = Int[]
        for (w,c,t) in testdata
            y = nn(w, c)
            append!(preds, y)
            append!(golds, t.data)
        end
        length(preds) == length(golds) || throw("Length mismatch: $(length(preds)), $(length(golds))")

        preds = decode(tagset, preds)
        golds = decode(tagset, golds)
        fscore(golds, preds)
        println()
    end
    Decoder(chardict, tagset, nn)
end

function encode(chardict::Dict, chars::Vector{String})
    unk = chardict["UNKNOWN"]
    ids = map(c -> get(chardict,c,unk), chars)
    Var(ids)
end

function readdata(path::String, chardict::Dict, tagset)
    data = NTuple{2,Var}[]
    chars, wordtags, senttags = String[], String[], Int[]
    lines = open(readlines, path)
    push!(lines, "")
    for line in lines
        if isempty(line)
            isempty(chars) && continue
            c = encode(chardict, chars)
            t1 = Var(encode(tagset,tags))
            push!(data, (w,c,t))
            empty!(chars)
            empty!(wordtags)
            empty!(senttags)
        else
            items = split(line, "\t")
            push!(chars, String(items[1]))
            push!(wordtags, String(items[2]))
            push!(senttags, parse(Int,items[3]))
        end
    end
    data
end

function train(dec::Decoder, trainfile::String, testfile::String)

end


text:
module YaoFlux

using Zygote, LuxurySparse, Yao, YaoBase, SparseArrays, BitBasis
import Zygote: Context

Zygote.@adjoint (::Type{T})(perms, vals) where T <: PermMatrix = T(perms, vals), Δ -> nothing
Zygote.@adjoint (::Type{T})() where T <: IMatrix = T(), Δ -> nothing
Zygote.@adjoint Base.:(*)(A::Number, B::PermMatrix) = A * B, Δ->(sum(Δ .* B), A * Δ)
Zygote.@adjoint Base.:(*)(A::PermMatrix, B::Number) = A * B, Δ->(A * Δ, sum(Δ .* B))

Zygote.@adjoint SparseArrays.SparseMatrixCSC(A::PermMatrix) = SparseMatrixCSC(A), Δ->(Δ, )

Zygote.@adjoint BitBasis.onehot(::Type{T}, nbits::Int, x::Integer, nbatch::Int) where T = onehot(T, nbits, x, nbatch), Δ -> nothing

# YaoArrayRegister

# struct GradReg{Reg, G}
#     r::Reg # this is the register to back prop
#     grad::G # gradient of the register
# end

# no gradient if meets a constructor
# Zygote.@adjoint (::Type{T})(r) where T <: ArrayReg = T(r), Δ->nothing

# Zygote.@adjoint function YaoArrayRegister.state(r::ArrayReg)
#     state(r), Δ -> (GradReg(r, Δ), )
# end

Zygote.@adjoint function Base.getfield(r::ArrayReg, n::Symbol)
    getfield(r, n), Δ -> (ArrayReg(Δ), nothing)
end

# YaoBlocks

# Zygote.@adjoint (::Type{T})(P::AbstractBlock, θ) where T <: RotationGate = T(P, θ), Δ->(nothing, real(Δ),)

# U * r
# delta * r', U' * delta
# in most cases block do not have parameters
# Zygote.@adjoint function YaoBlocks.apply!(r::ArrayReg, x::ConstantGate)
#     apply!(r, x), function (Δ)
#         apply!(copy(Δ), x'), nothing # no grad for constant gates
#     end
# end

# Zygote.@adjoint function YaoBlocks.apply!(r::ArrayReg{B, T}, x::RotationGate) where {B, T}
#     r_ = copy(r)
#     apply!(r, x), function (Δ)
#         # dU = reshape(Δ * state(r)', size(U))
#         # dr = reshape(U' * Δ, size(U))
#         # apply!(Δ, x')
#         dispatch!(-, x, π)
#         apply!(r_, -x/2)
#         dispatch!(+, x, π)

#         return (apply!(copy(Δ), x'), sum(state(Δ) * state(r_)'))
#     end
# end

# upstreams
Zygote.@adjoint Base.:(-)(a, b) = a-b, Δ -> (Δ, -Δ)
Zygote.@adjoint Base.:(+)(a, b) = a+b, Δ -> (Δ, Δ)

# require mutate
Zygote.@adjoint! function copyto!(xs::AbstractVector, ys::Tuple)
    xs_ = copy(xs)
    copyto!(xs, ys), function (dxs)
        copyto!(xs_, xs)
        return (nothing, Tuple(dxs))
    end
end

end # module


text:
using DataFrames,DataArrays,Distributions,StatsBase,MixedModels

"""
Create a DataFrame with crossed grouping factors, `S` and `I`
with `M` and `N` levels, respectively.  The condition, `C`,
should be a vector of length 2.  The effect of all the `rep` calls
is similar to `expand.grid` in R.
"""
function mkdata(C::Vector=[zero(Int8),one(Int8)],N=10,M=100)
    l = length(C)
    DataFrame(y = zeros(l*N*M),
              C = rep(C,N*M),
              I = rep(rep(compact(pool(collect(1:N))),fill(l,N)),M),
              S = rep(compact(pool(collect(1:M))),fill(l*N,M)))
end

"""
Simulate a single response vector for a linear mixed model with
fixed-effects parameter β, and independent random effects for intercept
and condition, `C`, by subject and by item.  The function modifies the
first argument, `dat`, replacing dat[:y].
"""
function sim1!(dat::DataFrame,β::Vector,σ=600.,σRS=300.,σRI=30.)
    y,C,I,S = dat[:y],dat[:C],dat[:I],dat[:S]
    length(β) == length(unique(C)) == 2 || throw(DimensionMismatch())
    # vectors of random effects (subject-int,subject-C,item-int,item-C)
    M,N = length(S.pool),length(I.pool)
    re = Vector{Float64}[σRS*randn(M),σRS*randn(M)*0.5,
                         σRI*randn(N),σRI*randn(N)*0.5]
    I,S = I.refs,S.refs
    for ii in eachindex(y)
        yi = β[1] + re[1][S[ii]] + re[3][I[ii]] + σ*randn()
        if (cc = C[ii]) ≠ 0
            yi += cc*(β[2]+re[2][S[ii]]+re[4][I[ii]])
        end
        y[ii] = yi
    end
    dat
end

const dat = mkdata();
const Xs = hcat(ones(2000),dat[:C]);
const Xst = Xs'
srand(1234321);

sim1!(dat,[2000.,0.]);
ms1 = fit!(lmm(y ~ 1+C + (1+C|S) + (1+C|I),dat));
#ms2 = fit!(lmm(y ~ 1+C + (1+C|S) + (1|I),dat))

max_stars_repo_name: werbac/jlib

max_issues_repo_name: werbac/jlib

max_forks_repo_name: werbac/jlib


text:
using Quandl, DataFrames, Redis, Winston

qapi_key = "aaaaaaaaaaaaaaaaaaaa"	# Your Quandl key
set_auth_token(qapi_key)

qdf1 = quandl("WIKI/AAPL", format="DataFrame", order="asc");
aapl = convert(Array, qdf1[:Close]);
scf1 = 1.0/aapl[1];  aapl = scf1 .* aapl;

qdf2 = quandl("WIKI/MSFT", format="DataFrame", order="asc");
msft = convert(Array, qdf2[:Close]);
scf2 = 1.0/msft[1];  msft = scf2 .* msft;

n = [1:length(aapl)];
conn = RedisConnection()
for i = 1:n
  rpush(conn,'APPL~Close',aapl[i])
  rpush(conn,'MSFT~Close',msft[i])
end

t = [1:n];		# Just plot the data not the dates

aapl-data = float32(lrange(conn,'AAPL~Close',0,-1);
msft-data = float32(lrange(conn,'MSFT~Close',0,-1);
plot(t, aapl-data,"b--", t, msft-data, "r.");

max_stars_repo_path: Module 3/Chapter08/redis-quandl.jl

max_issues_repo_path: Module 3/Chapter08/redis-quandl.jl

max_forks_repo_path: Module 3/Chapter08/redis-quandl.jl


text:
using FluxRC, KitBase, Plots, OrdinaryDiffEq, LinearAlgebra
using ProgressMeter: @showprogress

cd(@__DIR__)
ps = UnstructPSpace("../assets/square.msh")

N = deg = 2
Np = (N + 1) * (N + 2) ÷ 2
ncell = size(ps.cellid, 1)
nface = size(ps.faceType, 1)

J = rs_jacobi(ps.cellid, ps.points)

spg = global_sp(ps.points, ps.cellid, N)
fpg = global_fp(ps.points, ps.cellid, N)

pl, wl = tri_quadrature(N)

V = vandermonde_matrix(N, pl[:, 1], pl[:, 2])
Vr, Vs = ∂vandermonde_matrix(N, pl[:, 1], pl[:, 2])
∂l = ∂lagrange(V, Vr, Vs)

ϕ = correction_field(N, V)

pf, wf = triface_quadrature(N)
ψf = zeros(3, N + 1, Np)
for i = 1:3
    ψf[i, :, :] .= vandermonde_matrix(N, pf[i, :, 1], pf[i, :, 2])
end

lf = zeros(3, N + 1, Np)
for i = 1:3, j = 1:N+1
    lf[i, j, :] .= V' \ ψf[i, j, :]
end

γ = 5 / 3
u0 = zeros(size(ps.cellid, 1), Np, 4)
for i in axes(u0, 1), j in axes(u0, 2)
    prim = [
        max(exp(-300 * ((spg[i, j, 1] - 0.5)^2 + (spg[i, j, 2] - 0.5)^2)), 1e-3),
        0.0,
        0.0,
        1.0,
    ]
    #prim = [1., 0., 0., 1.]
    u0[i, j, :] .= prim_conserve(prim, γ)
end

cell_normal = zeros(ncell, 3, 2)
for i = 1:ncell
    pids = ps.cellid[i, :]

    cell_normal[i, 1, :] .= unit_normal(ps.points[pids[1], :], ps.points[pids[2], :])
    cell_normal[i, 2, :] .= unit_normal(ps.points[pids[2], :], ps.points[pids[3], :])
    cell_normal[i, 3, :] .= unit_normal(ps.points[pids[3], :], ps.points[pids[1], :])

    p =
        [
            ps.points[pids[1], :] .+ ps.points[pids[2], :],
            ps.points[pids[2], :] .+ ps.points[pids[3], :],
            ps.points[pids[3], :] .+ ps.points[pids[1], :],
        ] / 2

    for j = 1:3
        if dot(cell_normal[i, j, :], p[j][1:2] - ps.cellCenter[i, 1:2]) < 0
            cell_normal[i, j, :] .= -cell_normal[i, j, :]
        end
    end
end

function dudt!(du, u, p, t)
    du .= 0.0

    deg, γ = p
    ncell = size(u, 1)
    nsp = size(u, 2)

    f = zeros(ncell, nsp, 4, 2)
    for i in axes(f, 1)
        for j in axes(f, 2)
            fg, gg = euler_flux(u[i, j, :], γ)
            for k in axes(f, 3)
                f[i, j, k, :] .= inv(J[i]) * [fg[k], gg[k]]
            end
        end
    end

    u_face = zeros(ncell, 3, deg + 1, 4)
    f_face = zeros(ncell, 3, deg + 1, 4, 2)
    for i = 1:ncell, j = 1:3, k = 1:deg+1, l = 1:4
        u_face[i, j, k, l] = sum(u[i, :, l] .* lf[j, k, :])
        f_face[i, j, k, l, 1] = sum(f[i, :, l, 1] .* lf[j, k, :])
        f_face[i, j, k, l, 2] = sum(f[i, :, l, 2] .* lf[j, k, :])
    end

    n = [[0.0, -1.0], [1 / √2, 1 / √2], [-1.0, 0.0]]

    fn_face = zeros(ncell, 3, deg + 1, 4)
    for i = 1:ncell, j = 1:3, k = 1:deg+1, l = 1:4
        fn_face[i, j, k, l] = dot(f_face[i, j, k, l, :], n[j])
    end

    fn_interaction = zeros(ncell, 3, deg + 1, 4)
    for i = 1:ncell, j = 1:3, k = 1:deg+1
        uL = local_frame(u_face[i, j, k, :], cell_normal[i, j, 1], cell_normal[i, j, 2])

        ni, nj, nk = neighbor_fpidx([i, j, k], ps, fpg)

        fwn_local = zeros(4)
        if ni > 0
            uR = local_frame(
                u_face[ni, nj, nk, :],
                cell_normal[i, j, 1],
                cell_normal[i, j, 2],
            )

            flux_hll!(fwn_local, uL, uR, γ, 1.0)
        end
        fwn_global = global_frame(fwn_local, cell_normal[i, j, 1], cell_normal[i, j, 2])

        fwn_xy = zeros(4, 2)
        for idx = 1:4
            fwn_xy[idx, :] .= fwn_global[idx] .* cell_normal[i, j, :]
        end

        fws_rs = zeros(4, 2)
        for idx = 1:4
            fws_rs[idx, :] .= inv(J[i]) * fwn_xy[idx, :]
        end

        fwns = [sum(fws_rs[idx, :] .* n[j]) for idx = 1:4]
        #fwns = [sum(fws_rs[idx, :]) for idx in 1:4]

        fn_interaction[i, j, k, :] .= fwns
    end
    #=
        for i = 1:ncell, j = 1:3, k = 1:deg+1
            uL = local_frame(u_face[i, j, k, :], 1., 0.)
            uD = local_frame(u_face[i, j, k, :], 0., 1.)

            ni, nj, nk = neighbor_fpidx([i, j, k], ps, fpg)

            fw1_local = zeros(4)
            fw2_local = zeros(4)
            if ni > 0
                uR = local_frame(u_face[ni, nj, nk, :], 1., 0.)
                uU = local_frame(u_face[ni, nj, nk, :], 0., 1.)

                flux_hll!(fw1_local, uL, uR, γ, 1.0)
                flux_hll!(fw2_local, uD, uU, γ, 1.0)
            end
            fw1_global = global_frame(fw1_local, 1., 0.)
            fw2_global = global_frame(fw2_local, 0., 1.)

            fw_xy = hcat(fw1_global, fw2_global)

            fw_rs = zeros(4, 2)
            for idx in 1:4
                fw_rs[idx, :] .= inv(J[i]) * fw_xy[idx, :]
            end

            fwns = [sum(fw_rs[idx, :] .* n[j]) for idx in 1:4]
            #fwns = [sum(fws_rs[idx, :]) for idx in 1:4]

            fn_interaction[i, j, k, :] .= fwns
        end=#

    #=
        for i = 1:ncell, j = 1:3, k = 1:deg+1
            fw1 = zeros(4)
            fw2 = zeros(4)

            ni, nj, nk = neighbor_fpidx([i, j, k], ps, fpg)

            if ni > 0
                flux_roe!(fw1, u_face[i, j, k, :], u_face[ni, nj, nk, :], γ, 1., [1., 0.])
                flux_roe!(fw2, u_face[i, j, k, :], u_face[ni, nj, nk, :], γ, 1., [0., 1.])
            end

            fi = zeros(4, 2)
            for id in 1:4
                fi[id, :] .= inv(J[i]) * [fw1[id], fw2[id]]
            end

            for l = 1:4
                fn_interaction[i, j, k, l] = sum(fi[l, :] .* n[j])
            end
        end=#

    #=
        for i = 1:ncell, j = 1:3, k = 1:deg+1
            ni, nj, nk = neighbor_fpidx([i, j, k], ps, fpg)

            if ni > 0
                fR = f_face[ni, nj, nk, :, :]
                fRg = zeros(4, 2)
                for id = 1:4
                    fRg[id, :] .= J[ni] * f_face[ni, nj, nk, id, :]
                end
                fRl = zeros(4, 2)
                for id = 1:4
                    fRl[id, :] .= inv(J[i]) * fRg[id, :]
                end

                _f0 = (f_face[i, j, k, :, :] + fRl) ./ 2
                _f1 = -det(J[i]) * 0.5 .* (u_face[i, j, k, :] + u_face[ni, nj, nk, :]) .* (ps.cellArea[i]^0.5 / 2) ./ dt
                _f = hcat(_f0[:, 1] + _f1, _f0[:, 2] + _f1)

                fn_interaction[i, j, k, :] .= [sum(_f[id, :] .* n[j]) for id = 1:4]
            end
        end=#

    rhs1 = zeros(ncell, nsp, 4)
    for i in axes(rhs1, 1), j in axes(rhs1, 2), k = 1:4
        if ps.cellType[i] == 0
            rhs1[i, j, k] =
                -sum(f[i, :, k, 1] .* ∂l[j, :, 1]) - sum(f[i, :, k, 2] .* ∂l[j, :, 2])
        end
    end

    rhs2 = zero(rhs1)
    for i = 1:ncell
        if ps.cellType[i] == 0
            for j = 1:nsp, k = 1:4
                rhs2[i, j, k] =
                    -sum((fn_interaction[i, :, :, k] .- fn_face[i, :, :, k]) .* ϕ[:, :, j])
                #rhs2[i, j, k] = - sum((fn_interaction[i, :, :, k] .- fn_face[i, :, :, k])) / 3
            end
        end
    end

    for i = 1:ncell
        if ps.cellType[i] == 0
            du[i, :, :] .= rhs1[i, :, :] .+ rhs2[i, :, :]
            #du[i, :, :] .= rhs2[i, :, :]
        end
    end

    return nothing
end

tspan = (0.0, 0.1)
p = (N, 5 / 3)
prob = ODEProblem(dudt!, u0, tspan, p)
dt = 0.002
itg = init(prob, Euler(), save_everystep = false, adaptive = false, dt = dt)

@showprogress for iter = 1:20
    step!(itg)
end

begin
    prim = zero(itg.u)
    for i in axes(prim, 1), j in axes(prim, 2)
        prim[i, j, :] .= conserve_prim(itg.u[i, j, :], γ)
        prim[i, j, 4] = 1 / prim[i, j, 4]
    end
    write_vtk(ps.points, ps.cellid, prim[:, 2, :])
end

du = zero(u0)
f1, f2 = dudt!(du, u0, p, 1.0)

idx = 1211

f1[idx, 2, 2, :]

f2[idx, 2, 2, :]

fw = zeros(4)
flux_roe!(fw, [1.0, 0.0, 0.0, 1.0], []fw::X, wL::Y, wR::Y, γ, dt, n = [1.0, 0.0])


text:
"""
    FeasibilityEvaluator{T} <: AbstractNLPEvaluator

TODO

"""
mutable struct FeasibilityEvaluator{Evaluator<:AbstractNLPEvaluator, T, VT} <: AbstractNLPEvaluator
    inner::Evaluator
    x_min::VT
    x_max::VT
    cons::VT
end
function FeasibilityEvaluator(nlp::AbstractNLPEvaluator)
    if !is_constrained(nlp)
        error("Input problem must have inequality constraints")
    end
    x_min, x_max = bounds(nlp, Variables())
    cx = similar(x_min, n_constraints(nlp))
    return FeasibilityEvaluator{typeof(nlp), eltype(x_min), typeof(x_min)}(nlp, x_min, x_max, cx)
end
function FeasibilityEvaluator(datafile::String)
    nlp = SlackEvaluator(datafile)
    return FeasibilityEvaluator(nlp)
end

n_variables(nlp::FeasibilityEvaluator) = n_variables(nlp.inner)
n_constraints(nlp::FeasibilityEvaluator) = 0

constraints_type(::FeasibilityEvaluator) = :bound

has_hessian(nlp::FeasibilityEvaluator) = has_hessian(nlp.inner)
has_hessian_lagrangian(nlp::FeasibilityEvaluator) = has_hessian(nlp)

# Getters
get(nlp::FeasibilityEvaluator, attr::AbstractNLPAttribute) = get(nlp.inner, attr)
get(nlp::FeasibilityEvaluator, attr::AbstractVariable) = get(nlp.inner, attr)
get(nlp::FeasibilityEvaluator, attr::PS.AbstractNetworkAttribute) = get(nlp.inner, attr)

# Setters
function setvalues!(nlp::FeasibilityEvaluator, attr::PS.AbstractNetworkValues, values)
    setvalues!(nlp.inner, attr, values)
end

# Bounds
bounds(nlp::FeasibilityEvaluator, ::Variables) = bounds(nlp.inner, Variables())
bounds(nlp::FeasibilityEvaluator, ::Constraints) = (Float64[], Float64[])

initial(nlp::FeasibilityEvaluator) = initial(nlp.inner)

function update!(nlp::FeasibilityEvaluator, u)
    conv = update!(nlp.inner, u)
    constraint!(nlp.inner, nlp.cons, u)
    return conv
end

# f(x) = 0.5 * || c(x) ||²
function objective(nlp::FeasibilityEvaluator, u)
    return 0.5 * dot(nlp.cons, nlp.cons)
end

function constraint!(nlp::FeasibilityEvaluator, cons, u)
    @assert length(cons) == 0
    return
end

# Gradient
# ∇f = J' * c(x)
function gradient!(nlp::FeasibilityEvaluator, grad, u)
    σ = 0.0
    ojtprod!(nlp.inner, grad, u, σ, nlp.cons)
    return
end

jacobian_structure(ag::FeasibilityEvaluator) = (Int[], Int[])
function jacobian!(ag::FeasibilityEvaluator, jac, u)
    @assert length(jac) == 0
    return
end

# H = ∇²c(x) + J'*J
function hessprod!(nlp::FeasibilityEvaluator, hessvec, u, v)
    σ = 0.0
    # Need to update the first-order adjoint λ first
    hessian_lagrangian_penalty_prod!(nlp.inner, hessvec, u, nlp.cons, σ, v, 0.0)
    # J' * J * v
    jv = similar(nlp.cons)
    jtv = similar(u)
    jprod!(nlp.inner, jv, u, v)
    jtprod!(nlp.inner, jtv, u, jv)
    hessvec .+= jtv
    return
end

function hessian_structure(nlp::FeasibilityEvaluator)
    n = n_variables(nlp)
    rows = Int[r for r in 1:n for c in 1:r]
    cols = Int[c for r in 1:n for c in 1:r]
    return rows, cols
end

function reset!(nlp::FeasibilityEvaluator)
    reset!(nlp.inner)
end


text:
# Random array with two phases
rand_array = rand(Float32, (50, 50, 50))
rand_array = map(x -> (x<0.3) ? 0 : 1, rand_array)

@testcase "Check s²(a, x) = l²(a, x) = 0 for all x where a is an array without phase 2." begin
    for p in (false, true)
        for func in (Directional.s2, Directional.l2)
            corr = mean(func(rand_array, 2; periodic = p, directions = known_directions))
            @test all(x -> x == 0, corr)
        end
    end
end

@testcase "Check that corr(a, len1, phase) = corr(a, len2, phase)[1:len1] for len2>len1" begin
    for phase in (0, 1)
        for func in (Directional.s2,
                     Directional.l2,
                     Directional.c2,
                     Directional.surfsurf,
                     Directional.surfvoid)
            corr1 = mean(func(rand_array, phase; len = 30))
            corr2 = mean(func(rand_array, phase; len = 35))
            @test corr1 == corr2[1:30]
        end
    end
end

# Probability of randomly choosing 0 or 1
p = count(x -> x == 0, rand_array) / length(rand_array) # P{random dot ∈ phase 0}
q = 1 - p                                               # P{random dot ∈ phase 1}
prob = [p, q]

@testcase "sᶦ(a,0) = lᶦ(a,0) = P{randomly choosing i}" begin
    for p in (false, true)
        for phase in 0:1
            for func in (Directional.s2, Directional.l2)
                f = mean(func(rand_array, phase; len = 1, directions = known_directions))[1]
                @test relerr(f, prob[phase+1]) < 0.02
            end
        end
    end
end

@testcase "sᶦ(a,1) = lᶦ(a,1)" begin
    for p in (false, true)
        for phase in 0:1
            s = mean(Directional.s2(rand_array, phase;
                                    len = 2, directions = known_directions))[2]
            l = mean(Directional.l2(rand_array, phase;
                                    len = 2, directions = known_directions))[2]
            @test s ≈ l
        end
    end
end

@testcase "Check that l2 is non-increasing" begin
    for p in (false, true)
        for phase in 0:1
            l = mean(Directional.l2(rand_array, phase; directions = known_directions))
            @test minimum(map(-, l, l[2:end])) >= 0
        end
    end
end

@testcase "Check that sⁱ(x) = P{randomly independently choosing i twice} (x > 0)" begin
    for p in (false, true)
        for phase in 0:1
            s = mean(Directional.s2(rand_array, phase;
                                    directions = known_directions, periodic = p))
            @test relerr(maximum(s[2:end]), prob[phase+1]^2) < 0.04
        end
    end
end

@testcase "Check pore size and chord length sum" begin
    for phase in 0:1
        @test sum(pore_size(rand_array, phase).weights) ≈ 1
        @test sum(Directional.chord_length(rand_array, phase)[1].weights) ≈ 1
    end
end


text:
# --------------- Subsetting IO table -----------------------------------------------------------------------------------------------------------------------------

"""
    index_subset(N::Integer, S::Integer, ctry_all::Vector{String}, ctry_subset::Vector{String}, ind_all::Vector{String}, ind_subset::Vector{String})

The function yields the row/column indices for the IO table according to desired countries and industries. The subsetting should be done for Z and F,
    which then needs a recalculation of Y and VA!

# Arguments
- `N::Integer`: number of origin/destination countries.
- `S::Integer`: number of origin/destination industries.
- `ctry_all::Vector{String}`: ordered vector of all country names in IO table.
- `ctry_subset::Vector{String}`: vector of country names to subset.
- `ind_all::Vector{String}`: ordered vector of all industry codes (2 digits) in IO table.
- `ind_subset::Vector{String}`: vector of industry codes (2 digits) to subset.

# Output
- `index_subset::Vector{Integer}`: indices to subset Z and F matrices.

# Examples
```julia-repl
julia> index_subset = subset(N, S, countries, ctry_EU, industries, ind_sub)
```
"""

function index_subset(N::Integer, S::Integer, ctry_all::Vector{String}, ctry_subset::Vector{String}, ind_all::Vector{String}, ind_subset::Vector{String})

    # subsetting rows
    N_subset = length(ctry_subset)
    S_subset = length(ind_subset)

    all = repeat(ctry_all, inner=S) .* "_" .* repeat(ind_all, outer=N) # all row country-industry pairs
    subset = repeat(ctry_subset, inner=S_subset) .* "_" .* repeat(ind_subset, outer=N_subset) # subset of country-industry pairs
    index_subset = findall(in(subset), all)

    return index_subset
end

# ---------------

"""
    ctry_aggregation(Z::Matrix, F::Matrix, N::Integer, S::Integer,
        ctry_all::Vector{String}, ctry_subset::Vector{String}, ind_all::Vector{String}, ind_subset::Vector{String})

The function creates new intermediate and final demand matrices by aggregating the countries in the vector "ctry_subset". The aggregation only works for
    countries, therefore, the two industry vectors need to be the same.

# Arguments
- `Z::Matrix`: NS×NS, origin country-industry destination country-industry intermediate demand matrix Z.
- `F::Matrix`: NS×N, origin country-industry destination country final demand matrix F.
- `N::Integer`: number of origin/destination countries.
- `S::Integer`: number of origin/destination industries.
- `ctry_all::Vector{String}`: ordered vector of all country names in IO table.
- `ctry_subset::Vector{String}`: vector of country names to subset.
- `ind_all::Vector{String}`: ordered vector of all industry codes (2 digits) in IO table.
- `ind_subset::Vector{String}`: needs to be identical to the vector "ind_all".

# Output
- `Z::Matrix`: N_new*S×N_new*S, new origin country-industry destination country-industry intermediate demand matrix Z.
- `F::Matrix`: N_new*S×N_new, new origin country-industry destination country final demand matrix F.

# Examples
```julia-repl
julia> Z_new, F_new = aggregation(Z, F, N, S, countries, ctry_EU, industries, industries)
```
"""

function ctry_aggregation(Z::Matrix, F::Matrix, N::Integer, S::Integer,
    ctry_all::Vector{String}, ctry_subset::Vector{String}, ind_all::Vector{String}, ind_subset::Vector{String})

    # obtain indices for subsetting Z and F
    index_row_subset = index_subset(N, S, ctry_all, ctry_subset, ind_all, ind_subset)
    index_col_subset = index_subset(N, S, ctry_all, ctry_subset, ind_all, ind_subset)
    index_col_subset_F = findall(in(ctry_subset), ctry_all) # only have countries in columns

    # subset columns, construct aggregate and rebuild matrix
    N_col_subset = length(ctry_subset)
    S_col_subset = length(ind_subset)

    Z_col_subset = Z[:, index_col_subset]
    Z_col_subset = [sum(Z_col_subset[i,j:S_col_subset:(N_col_subset-1)*S_col_subset+j]) for i in 1:N*S, j in 1:S_col_subset] # sum over industries to create aggregate
    Z_col_other = Z[:, Not(index_col_subset)] # other countries
    Z_new = hcat(Z_col_other, Z_col_subset) # rebuild Z with new aggregate as last columns

    F_col_subset = F[:, index_col_subset_F]
    F_col_subset = [sum(F_col_subset[i,:]) for i in 1:N*S]
    F_col_other = F[:, Not(index_col_subset_F)] # other countries
    F_new = hcat(F_col_other, F_col_subset) # rebuild F with new aggregate as last column

    # subset rows from new matrix, construct aggregate and rebuild matrix
    N_row_subset = length(ctry_subset)
    S_row_subset = length(ind_subset)
    N_new = size(F_new, 2) # new number of countries

    Z_row_subset = Z_new[index_row_subset, :]
    Z_row_subset = [sum(Z_row_subset[i:S_row_subset:(N_row_subset-1)*S_row_subset+i,j]) for i in 1:S_row_subset, j in 1:N_new*S_row_subset]
    Z_row_other = Z_new[Not(index_row_subset), :]
    Z_new = vcat(Z_row_other, Z_row_subset)

    F_row_subset = F_new[index_row_subset, :]
    F_row_subset = [sum(F_row_subset[i:S_row_subset:(N_row_subset-1)*S_row_subset+i,j]) for i in 1:S_row_subset, j in 1:N_new]
    F_row_other = F_new[Not(index_row_subset), :]
    F_new = vcat(F_row_other, F_row_subset)

    return Z_new, F_new
end

# ---------------- aggregation in one function --------------------------------------------------------------------------------------------------------------------------

"""
    ctry_aggregation(Z::Matrix, F::Matrix, N::Integer, S::Integer, ctry_all::Vector{String}, ctry_subset::Vector{String})

The function creates new intermediate and final demand matrices by aggregating the countries in the vector "ctry_subset". The aggregation only works for
    countries not for industries.

# Arguments
- `Z::Matrix`: NS×NS, origin country-industry destination country-industry intermediate demand matrix Z.
- `F::Matrix`: NS×N, origin country-industry destination country final demand matrix F.
- `N::Integer`: number of origin/destination countries.
- `S::Integer`: number of origin/destination industries.
- `ctry_all::Vector{String}`: ordered vector of all country names in IO table.
- `ctry_subset::Vector{String}`: vector of country names to subset.

# Output
- `Z::Matrix`: N_new*S×N_new*S, new origin country-industry destination country-industry intermediate demand matrix Z.
- `F::Matrix`: N_new*S×N_new, new origin country-industry destination country final demand matrix F.

# Examples
```julia-repl
julia> Z_new, F_new = aggregation(Z, F, N, S, countries, ctry_EU)
```
"""

function ctry_aggregation(Z::Matrix, F::Matrix, N::Integer, S::Integer, ctry_all::Vector{String}, ctry_subset::Vector{String})

    # obtain indices for subsetting Z and F
    N_subset = length(ctry_subset)
    industries = lpad.(1:S, 2, '0') # vector of industries

    all = repeat(ctry_all, inner=S) .* "_" .* repeat(industries, outer=N) # all row country-industry pairs
    subset = repeat(ctry_subset, inner=S) .* "_" .* repeat(industries, outer=N_subset) # subset of country-industry pairs

    index_subset = findall(in(subset), all) # can be used for both rows and columns since IO table is a square matrix
    index_subset_F = findall(in(ctry_subset), ctry_all) # only have countries in columns

    # subset columns, construct aggregate and rebuild matrix
    Z_col_subset = Z[:, index_subset]
    Z_col_subset = [sum(Z_col_subset[i, j:S:(N_subset-1)*S+j]) for i in 1:N*S, j in 1:S] # sum over industries to create aggregate
    Z_col_other = Z[:, Not(index_subset)] # other countries
    Z_new = hcat(Z_col_other, Z_col_subset) # rebuild Z with new aggregate as last columns

    F_col_subset = F[:, index_subset_F]
    F_col_subset = [sum(F_col_subset[i,:]) for i in 1:N*S]
    F_col_other = F[:, Not(index_subset_F)] # other countries
    F_new = hcat(F_col_other, F_col_subset) # rebuild F with new aggregate as last column

    # subset rows from new matrix, construct aggregate and rebuild matrix
    N_new = size(F_new, 2) # new number of countries

    Z_row_subset = Z_new[index_subset, :]
    Z_row_subset = [sum(Z_row_subset[i:S:(N_subset-1)*S+i, j]) for i in 1:S, j in 1:N_new*S]
    Z_row_other = Z_new[Not(index_subset), :]
    Z_new = vcat(Z_row_other, Z_row_subset)

    F_row_subset = F_new[index_subset, :]
    F_row_subset = [sum(F_row_subset[i:S:(N_subset-1)*S+i, j]) for i in 1:S, j in 1:N_new]
    F_row_other = F_new[Not(index_subset), :]
    F_new = vcat(F_row_other, F_row_subset)

    return Z_new, F_new
end


text:
using StatsBase

PROJECT_IMG_FOLDER = joinpath("data", "book-covers")
DATASETS = ["test", "train"]

if ~isdir(PROJECT_IMG_FOLDER) mkdir(PROJECT_IMG_FOLDER) end

for dataset_name in DATASETS

    categories = Int[]

    current_folder = joinpath(PROJECT_IMG_FOLDER, dataset_name)
    file_name = joinpath(PROJECT_IMG_FOLDER, "bookcover30-labels-$dataset_name.txt")
    f = open(file_name); lines = readlines(f); close(f);

    for line in lines
        value = parse(Int, split(line, " ")[2])
        push!(categories, value)
    end

    println(countmap(categories))
end