text:
module TestSimplePF
using Test
using JuMP, Ipopt
using DataFrames, CSV
using Printf

include("../src/DistSystem.jl")
import .DistSystem

include("../src/SimplePF.jl")
import .SimplePF
export runtests

struct Case
    name::String
    V_expected::Vector{Float64}
    angle_expected::Vector{Float64}
end

function case3_dist()
    VN = 12500
    Sb = 1e6
    Yb = Sb / (VN^2)
    Vsub = 1.0
    P_limt = 10
    Q_limit = 10
    bus_sub = 1
    nbuses = 3
    VL = 0.95
    VH = 1.05
    sub = DistSystem.Substation(VN, bus_sub, Vsub, P_limt, Q_limit, [1.0, 2, 3])
    data = DataFrame(CSV.File("case3_dist.csv"))
    case = Case(
        "case3_dist",
        [1.0000, 0.9641, 0.9511],
        [0.0, -1.8255, -2.6440]
    )
    sys = DistSystem.factory_system(data, VL, VH, sub)

    return sys, case
end

function case33()
    VN = 12660
    Sb = 1e6
    Yb = Sb / (VN^2)
    Vsub = 1.0
    P_limt = 4
    Q_limit = 2.5
    bus_sub = 1
    nbuses = 33
    VL = 0.90
    VH = 1.05
    sub = DistSystem.Substation(VN, bus_sub, Vsub, P_limt, Q_limit, [1.0, 2, 3])
    data = DataFrame(CSV.File("case33.csv"))
    sys = DistSystem.factory_system(data, VL, VH, sub)
    case = Case(
        "case33",
        [1.0000, 0.9973, 0.9848, 0.9785, 0.9723, 0.9583, 0.9563, 0.9520, 0.9475, 0.9435, 0.9428, 0.9416, 0.9383, 0.9378, 0.9376, 0.9375, 0.9396, 0.9403, 0.9968, 0.9932, 0.9925, 0.9919, 0.9813, 0.9746, 0.9713, 0.9568, 0.9547, 0.9468, 0.9412, 0.9386, 0.9379, 0.9382, 0.9398],
        [0, -0.0150, -0.0934, -0.1490, -0.2104, -0.5805, -0.8319, -0.8952, -1.1107, -1.3188, -1.3396, -1.3817, -1.6797, -1.8333, -1.9543, -2.0837, -2.3423, -2.4561, -0.0258, -0.0927, -0.1121, -0.1324, -0.1242, -0.2126, -0.2561, -0.5839, -0.5881, -0.7285, -0.8235, -0.8304, -1.1211, -1.2100, -1.2906]
    )
    return sys, case
end

function test_complex_multiplication()
    z1 = 1 + -4im
    z2 = -2 + 5im
    z3 = z1 * z2
    z1v = [real(z1), imag(z1)]
    z2v = [real(z2), imag(z2)]

    z3re = SimplePF.mc_re(z1v[1], z1v[2], z2v[1], z2v[2])
    z3im = SimplePF.mc_im(z1v[1], z1v[2], z2v[1], z2v[2])

    @testset "test_complex_multiplication" begin
        @test z3re == real(z3)
        @test z3im == imag(z3)
    end

end

function test_add_variables(sot, sys)
    sot = SimplePF.add_variables(sot, sys)
    @testset "add_variables" begin
        @test sot isa Model
        @test num_variables(sot) == sys.nbuses * 2

        for i = sys.buses
            @test start_value(sot[:V][:Re, i]) == 1.0
            @test start_value(sot[:V][:Im, i]) == 0.0

            @test lower_bound(sot[:V][:Re, i]) == -sys.VH
            @test lower_bound(sot[:V][:Im, i]) == -sys.VH

            @test upper_bound(sot[:V][:Re, i]) == sys.VH
            @test upper_bound(sot[:V][:Im, i]) == sys.VH
        end
    end

    return sot

end

function test_add_voltage_constraints(sot, sys)
    sot = SimplePF.add_voltage_constraints(sot, sys)
    @test sot isa Model

    voltage_constraint = sot[:voltage_constraint]
    V² = sot[:V²]
    V_module = sot[:V_module]
    V = sot[:V]

    for i = sys.buses
        @test V²[i] isa QuadExpr
        @test V²[i] == V[:Re, i]^2 + V[:Im, i]^2

        obj = constraint_object(voltage_constraint[i])
        @test isequal_canonical(obj.func, V²[i])
        @test obj.set == MOI.Interval(sys.VL^2, sys.VH^2)
    end
    return sot
end

function test_add_I_V_relationship(sot, sys)
    sot = SimplePF.add_I_V_relationship(sot, sys)
    @test sot isa Model
    V = sot[:V]
    I = sot[:I]
    G = real(sys.Y)
    B = imag(sys.Y)

    for i = sys.buses
        @test I[:Re, i] == sum(G[i, j] * V[:Re, j] for j in sys.buses) - sum(B[i, j] * V[:Im, j] for j in sys.buses)
        @test I[:Im, i] == sum(B[i, j] * V[:Re, j] for j in sys.buses) + sum(G[i, j] * V[:Im, j] for j in sys.buses)
    end
    return sot
end

function test_add_S_VI_relationship(sot, sys)
    sot = SimplePF.add_S_VI_relationship(sot, sys)
    @test sot isa Model
    V = sot[:V]
    I = sot[:I]
    P = sot[:P]
    Q = sot[:Q]

    for i = sys.buses
        @test P[i] == SimplePF.mc_re(V[:Re, i], V[:Im, i], I[:Re, i], -I[:Im, i])
        @test Q[i] == SimplePF.mc_im(V[:Re, i], V[:Im, i], I[:Re, i], -I[:Im, i])
    end
    return sot
end

function test_add_substation_constraint(sot, sys)
    sot = SimplePF.add_substation_constraint(sot, sys)
    @test sot isa Model

    V = sot[:V]
    I = sot[:I]
    P = sot[:P]
    Q = sot[:Q]
    i = sys.substation.bus
    Plimit = float(sys.substation.P_limit)
    Qlimit = float(sys.substation.Q_limit)

    @test is_fixed(V[:Re, i])
    @test is_fixed(V[:Im, i])
    @test fix_value(V[:Re, i]) == 1.0
    @test fix_value(V[:Im, i]) == 0

    obj = constraint_object(sot[:sub_plimit])
    @test obj.set == MOI.Interval(0.0, Plimit)
    @test isequal_canonical(obj.func, P[i])

    obj = constraint_object(sot[:sub_qlimit])
    @test obj.set == MOI.Interval(0.0, Qlimit)
    @test isequal_canonical(obj.func, Q[i])

    obj = constraint_object(sot[:sub_current])
    @test isequal_canonical(obj.func, I[:Re, i]^2 + I[:Im, i]^2)
    @test obj.set == MOI.GreaterThan(0.0)

    return sot
end

function test_add_power_injection_definition(sot, sys)
    sot = SimplePF.add_power_injection_definition(sot, sys)
    PL = sys.PL
    QL = sys.QL
    subbus = sys.substation.bus
    p = sot[:p]
    q = sot[:q]
    P = sot[:P]
    Q = sot[:Q]

    @test_throws KeyError p[subbus]
    @test_throws KeyError q[subbus]

    for i = [i for i = sys.buses if i != subbus]
        obj = constraint_object(p[i])
        @test isequal_canonical(obj.func, P[i])
        @test obj.set == MOI.EqualTo(-PL[i])

        obj = constraint_object(q[i])
        @test isequal_canonical(obj.func, Q[i])
        @test obj.set == MOI.EqualTo(-QL[i])
    end

    return sot
end

function test_nl_pf(sot, sys, case)

    other_sys, _ = case()
    model = SimplePF.nl_pf(Model(), other_sys)
    @testset "nl_pf" begin
        #Find a better way to test it
        @test sprint(print, sot) == sprint(print, model)
    end
    return sot
end

function util_test_case!(model, sys, case)

    @testset "solve $(case.name)" begin
        set_optimizer(model, Ipopt.Optimizer)
        set_silent(model)
        optimize!(model)

        @testset "Optimization" begin
            @testset "Solved: $(termination_status(model))" begin
                @test termination_status(model) == LOCALLY_SOLVED
            end
            @testset "Feasible: $(primal_status(model))" begin
                @test primal_status(model) == FEASIBLE_POINT
            end
        end

        V = value.(model[:V_module])
        θ = value.(model[:V_angle_deg])
        for i = sys.buses
            @test case.V_expected[i] ≈ V[i] atol = 1e-4
            @test case.angle_expected[i] ≈ θ[i] atol = 1e-4
        end
    end
end

function test_solve_case(case)
    sys, actual_case = case()
    model = SimplePF.nl_pf(Model(), sys)
    util_test_case!(model, sys, actual_case)
end

function runtests()
    @testset "SimplePF" begin
        for case in [case3_dist, case33]
            sys, actual_case = case()
            @testset "$(actual_case.name)" begin
                sot = Model()
                sot = test_add_variables(sot, sys)
                sot = test_add_voltage_constraints(sot, sys)
                sot = test_add_I_V_relationship(sot, sys)
                sot = test_add_S_VI_relationship(sot, sys)
                sot = test_add_substation_constraint(sot, sys)
                sot = test_add_power_injection_definition(sot, sys)
                sot = test_nl_pf(sot, sys, case)
                test_solve_case(case)
            end

        end
        test_complex_multiplication()
    end
end
end

TestSimplePF.runtests()


text:
using Test, BenchmarkTools, StaticArrays, LinearAlgebra, KiteUtils
using KiteModels, KitePodModels

if ! @isdefined SEGMENTS
    const SEGMENTS = se().segments
end
if ! @isdefined kcu
    const kcu = KCU(se())
    const kps = KPS3(kcu)
end

res1 = zeros(SVector{SEGMENTS+1, KiteModels.KVec3})
res2 = deepcopy(res1)
if ! @isdefined res3
    const res3 = reduce(vcat, vcat(res1, res2))
end

msg=""
@testset verbose = true "KPS3 benchmarking....   " begin

function set_defaults()
    KiteModels.clear!(kps)
    kps.set.l_tether = 150.0
    kps.set.elevation = 60.0
    kps.set.area = 20.0
    kps.set.rel_side_area = 50.0
    kps.set.v_wind = 8.0
    kps.set.mass = 11.4
    kps.set.damping =  2 * 473.0
    kps.set.alpha = 1.0/7
    kps.set.c_s = 0.6
end

function init_392()
    KiteModels.clear!(kps)
    kps.set.l_tether = 392.0
    kps.set.elevation = 70.0
    kps.set.area = 10.0
    kps.set.rel_side_area = 50.0
    kps.set.v_wind = 9.1
    kps.set.mass = 6.2
    kps.set.c_s = 0.6
end

set_defaults()
function test_initial_condition(params::Vector)
    my_state = kps
    y0, yd0 = KiteModels.init(my_state, params)
    residual!(res3, yd0, y0, kps, 0.0)
    return norm(res3) # z component of force on all particles but the first
end

# Inputs:
# State vector state_y   = pos1, pos2, ..., posn, vel1, vel2, ..., veln
# Derivative   der_yd    = vel1, vel2, ..., veln, acc1, acc2, ..., accn
# Output:
# Residual     res = res1, res2 = pos1,  ..., vel1, ...
@testset "test_residual!       " begin
    res1 = zeros(SVector{SEGMENTS, KVec3})
    res2 = deepcopy(res1)
    res = reduce(vcat, vcat(res1, res2))
    X = zeros(SimFloat, 2*kps.set.segments)
    y0, yd0 = KiteModels.init(kps, X; delta=1e-6)
    # println(y0)
    # println(yd0)
    p = kps
    t = 0.0
    clear!(kps)
    residual!(res, yd0, y0, p, t)
    res1 = res[1:3*SEGMENTS]
    res2 = res[3*SEGMENTS+1:end]
    @test res1 == zeros(3*(SEGMENTS))
    # TODO: add test for res2
    # println(res2)
end

t = @benchmark residual!(res, yd, y, p, t) setup = (res1 = zeros(SVector{SEGMENTS, KVec3}); res2 = deepcopy(res1);
                                                               res = reduce(vcat, vcat(res1, res2)); pos = deepcopy(res1);
                                                               pos[1] .= [1.0,2,3]; vel = deepcopy(res1); y = reduce(vcat, vcat(pos, vel));
                                                               der_pos = deepcopy(res1); der_vel = deepcopy(res1); yd = reduce(vcat, vcat(der_pos, der_vel));
                                                               p = kps; t = 0.0)

@test t.memory == 0
global msg = "Mean time residual! one point model: $(round(mean(t.times), digits=1)) ns"

end
println(msg)

# julia> include("test/bench.jl")
# Test Summary:         |
# test_initial_residual | No tests
# BenchmarkTools.Trial: 10000 samples with 191 evaluations.
#  Range (min … max):  523.073 ns … 953.518 ns  ┊ GC (min … max): 0.00% … 0.00%
#  Time  (median):     524.859 ns               ┊ GC (median):    0.00%
#  Time  (mean ± σ):   530.646 ns ±  27.683 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

#   █▅  ▂▃                                                        ▁
#   ███▆███▆▆▆▅▁▅▄▄▅▄▅▃▄▃▃▃▃▁▄▇▇▇▆▇▇▇▆▇▆▅▄▃▅▄▃▁▃▁▄▃▄▁▄▄▁▁▃▁▁▃█▇▅▅ █
#   523 ns        Histogram: log(frequency) by time        676 ns <

#  Memory estimate: 0 bytes, allocs estimate: 0.

ext: jl


text:
using Pkg
Pkg.activate("submissions/2021feb15")
Pkg.instantiate()

# for std()
using Statistics

# for det()
using LinearAlgebra

# for importing the data from CSV
using DataFrames
using CSV

# For GLOM
import GPLinearODEMaker; GLOM = GPLinearODEMaker

# For this module
using GLOM_RV_Example; GLOM_RV = GLOM_RV_Example

# For units in orbit fitting functions
using UnitfulAstro, Unitful

using Plots

rv_indicators_type = "dcpca0"
kernel_choice = 3
#starid = 101501
#starid = 10700
starid = 26965
#starid = 34411
if starid == 101501
    star_rot_rate = 16.28  # days
elseif starid == 10700
    star_rot_rate = 34.0  # days  # wikipedia, how embarassing
elseif starid == 26965
    star_rot_rate = 37.1  # days  #  Saar & Osten (1997) via https://iopscience.iop.org/article/10.3847/1538-3881/aaa896
elseif starid == 34411
    star_rot_rate = 27.9  # days # https://onlinelibrary.wiley.com/doi/full/10.1002/asna.202013790
end
 init_timescale_1 = star_rot_rate/4
 init_timescale_2 = 1.0/24
 tighten_lengthscale_priors = 3

#=
if starid == 101501
    star_rot_rate = 16.28  # days
elseif starid == 10700
    star_rot_rate = 34.0  # days  # wikipedia, how embarassing
elseif starid == 26965
    star_rot_rate = 37.1  # days  #  Saar & Osten (1997) via https://iopscience.iop.org/article/10.3847/1538-3881/aaa896
elseif starid == 34411
    star_rot_rate = 27.9  # days # https://onlinelibrary.wiley.com/doi/full/10.1002/asna.202013790
end

init_timescale_1 = star_rot_rate/4
init_timescale_2 = 1.0/24
tighten_lengthscale_priors = 3

include("glom_generic.jl")
=#

starid_list = [101501, 26965, 34411, 10700]
rv_indicators_type_list = ["dcpca0", "scalpels0", "dcpca3", "scalpels3", "dcpca4", "scalpels4" ]
kernel_choice_list = [3, 5, 4]

#for kernel_choice in kernel_choice_list
kernel_choice = 3
    for starid in starid_list
#starid = 26965
        if starid == 101501
            star_rot_rate = 16.28  # days
        elseif starid == 10700
            star_rot_rate = 34.0  # days  # wikipedia, how embarassing
        elseif starid == 26965
            star_rot_rate = 37.1  # days  #  Saar & Osten (1997) via https://iopscience.iop.org/article/10.3847/1538-3881/aaa896
        elseif starid == 34411
            star_rot_rate = 27.9  # days # https://onlinelibrary.wiley.com/doi/full/10.1002/asna.202013790
        end
        init_timescale_1 = star_rot_rate/4
        init_timescale_2 = 1.0/24
        if kernel_choice == 4 ||  kernel_choice == 6
            tighten_lengthscale_priors = 3
        else
            tighten_lengthscale_priors = 1
        end

        for rv_indicators_type in rv_indicators_type_list
            #rv_indicators_type = "dcpca0"

            include("glom_generic.jl")
            run_glom_rv(starid, rv_indicators_type, kernel_choice, star_rot_rate, init_timescale_1, init_timescale_2, tighten_lengthscale_priors)
        end
    end
#end

starid = 26965
rv_indicators_type = "dcpca0"
 include("glom_generic.jl")
rv_indicators_type = "scalpels0"
 include("glom_generic.jl")
rv_indicators_type = "scalpels3"
 include("glom_generic.jl")
rv_indicators_type = "scalpels4"
 include("glom_generic.jl")
rv_indicators_type = "dcpca3"
 include("glom_generic.jl")
rv_indicators_type = "dcpca4"
 include("glom_generic.jl")

starid = 34411
 rv_indicators_type = "dcpca0"
  include("glom_generic.jl")
 rv_indicators_type = "scalpels0"
  include("glom_generic.jl")
 rv_indicators_type = "scalpels3"
  include("glom_generic.jl")
 rv_indicators_type = "scalpels4"
  include("glom_generic.jl")
 rv_indicators_type = "dcpca3"
  include("glom_generic.jl")
 rv_indicators_type = "dcpca4"
  include("glom_generic.jl")

starid = 101501
 rv_indicators_type = "dcpca0"
   include("glom_generic.jl")
 rv_indicators_type = "scalpels0"
    include("glom_generic.jl")
 rv_indicators_type = "scalpels3"
    include("glom_generic.jl")
 rv_indicators_type = "scalpels4"
    include("glom_generic.jl")
 rv_indicators_type = "dcpca3"
    include("glom_generic.jl")
 rv_indicators_type = "dcpca4"
    include("glom_generic.jl")

starid = 10700
 rv_indicators_type = "dcpca0"
       include("glom_generic.jl")
 rv_indicators_type = "scalpels0"
        include("glom_generic.jl")
 rv_indicators_type = "scalpels3"
        include("glom_generic.jl")
 rv_indicators_type = "scalpels4"
        include("glom_generic.jl")
 rv_indicators_type = "dcpca3"
        include("glom_generic.jl")
 rv_indicators_type = "dcpca4"
        include("glom_generic.jl")

# NEW KERNEL
kernel_choice = 5
starid = 26965
rv_indicators_type = "dcpca0"
 include("glom_generic.jl")
rv_indicators_type = "scalpels0"
 include("glom_generic.jl")
rv_indicators_type = "scalpels3"
 include("glom_generic.jl")
rv_indicators_type = "scalpels4"
 include("glom_generic.jl")
rv_indicators_type = "dcpca3"
 include("glom_generic.jl")
rv_indicators_type = "dcpca4"
 include("glom_generic.jl")

starid = 34411
 rv_indicators_type = "dcpca0"
  include("glom_generic.jl")
 rv_indicators_type = "scalpels0"
  include("glom_generic.jl")
 rv_indicators_type = "scalpels3"
  include("glom_generic.jl")
 rv_indicators_type = "scalpels4"
  include("glom_generic.jl")
 rv_indicators_type = "dcpca3"
  include("glom_generic.jl")
 rv_indicators_type = "dcpca4"
  include("glom_generic.jl")

starid = 101501
 rv_indicators_type = "dcpca0"
   include("glom_generic.jl")
 rv_indicators_type = "scalpels0"
    include("glom_generic.jl")
 rv_indicators_type = "scalpels3"
    include("glom_generic.jl")
 rv_indicators_type = "scalpels4"
    include("glom_generic.jl")
 rv_indicators_type = "dcpca3"
    include("glom_generic.jl")
 rv_indicators_type = "dcpca4"
    include("glom_generic.jl")

starid = 10700
 rv_indicators_type = "dcpca0"
       include("glom_generic.jl")
 rv_indicators_type = "scalpels0"
        include("glom_generic.jl")
 rv_indicators_type = "scalpels3"
        include("glom_generic.jl")
 rv_indicators_type = "scalpels4"
        include("glom_generic.jl")
 rv_indicators_type = "dcpca3"
        include("glom_generic.jl")
 rv_indicators_type = "dcpca4"
        include("glom_generic.jl")

        # NEW KERNEL
kernel_choice = 4
 starid = 26965
 rv_indicators_type = "dcpca0"
 include("glom_generic.jl")
 rv_indicators_type = "scalpels0"
 include("glom_generic.jl")
 rv_indicators_type = "scalpels3"
 include("glom_generic.jl")
 rv_indicators_type = "scalpels4"
 include("glom_generic.jl")
 rv_indicators_type = "dcpca3"
 include("glom_generic.jl")
 rv_indicators_type = "dcpca4"
 include("glom_generic.jl")

starid = 34411
         rv_indicators_type = "dcpca0"
          include("glom_generic.jl")
         rv_indicators_type = "scalpels0"
          include("glom_generic.jl")
         rv_indicators_type = "scalpels3"
          include("glom_generic.jl")
         rv_indicators_type = "scalpels4"
          include("glom_generic.jl")
         rv_indicators_type = "dcpca3"
          include("glom_generic.jl")
         rv_indicators_type = "dcpca4"
          include("glom_generic.jl")

        starid = 101501
         rv_indicators_type = "dcpca0"
           include("glom_generic.jl")
         rv_indicators_type = "scalpels0"
            include("glom_generic.jl")
         rv_indicators_type = "scalpels3"
            include("glom_generic.jl")
         rv_indicators_type = "scalpels4"
            include("glom_generic.jl")
         rv_indicators_type = "dcpca3"
            include("glom_generic.jl")
         rv_indicators_type = "dcpca4"
            include("glom_generic.jl")

        starid = 10700
         rv_indicators_type = "dcpca0"
               include("glom_generic.jl")
         rv_indicators_type = "scalpels0"
                include("glom_generic.jl")
         rv_indicators_type = "scalpels3"
                include("glom_generic.jl")
         rv_indicators_type = "scalpels4"
                include("glom_generic.jl")
         rv_indicators_type = "dcpca3"
                include("glom_generic.jl")
         rv_indicators_type = "dcpca4"
                include("glom_generic.jl")


text:
### A Pluto.jl notebook ###
# v0.16.4

using Markdown
using InteractiveUtils

# ╔═╡ 9edb35ce-4ea7-11eb-353f-53459c337880
begin
	import Pkg
	Pkg.activate(Base.current_project())
	using ReinforcementLearning
	using Statistics
	using Flux
	using Plots
end

# ╔═╡ 6ce0a200-4ea7-11eb-123c-73e3a260351c
md"""
# Chapter 6.7 Maximization Bias and Double Learning

In example 6.7, authors introduced a MDP problem to compare the different performance of **Q-Learning** and **Double-Q-Learning**. This environment is kind of special compared to the environments we have seen before. In the first step, only `LEFT` and `RIGHT` are allowed. In the second step, if the `LEFT` is chosen previously, then we have 10 valid actions. We call this kind of environment is of `FULL_ACTION_SET`.
"""

# ╔═╡ 172236ac-4ea8-11eb-1393-5bc3683bf6b0
begin
	"""
	states:
	1:A
	2:B
	3:terminal
	actions:
	1: left
	2: right
	"""
	Base.@kwdef mutable struct MaximizationBiasEnv <: AbstractEnv
		position::Int = 1
		reward::Float64 = 0.0
	end
	function (env::MaximizationBiasEnv)(a::Int)
		if env.position == 1
			if a == 1
				env.position = 2
				env.reward = 0.0
			else
				env.position = 3
				env.reward = 0.0
			end
		elseif env.position == 2
			env.position = 3
			env.reward = randn() - 0.1
		end
		nothing
	end
end

# ╔═╡ 19876a70-4ea8-11eb-2b70-9d48875f7ae7
RLBase.state_space(env::MaximizationBiasEnv) = Base.OneTo(3)

# ╔═╡ 247ada66-4ea8-11eb-3c0a-af131b73e9f3
RLBase.action_space(env::MaximizationBiasEnv) = Base.OneTo(10)

# ╔═╡ 26b9f6f2-4ea8-11eb-38f3-3f1822c4acd2
RLBase.ActionStyle(env::MaximizationBiasEnv) = FULL_ACTION_SET

# ╔═╡ 296e0ab6-4ea8-11eb-33c5-71be241708b1
const LEFT = 1

# ╔═╡ 2c06d25a-4ea8-11eb-26e3-01d81d714859
const RIGHT = 2

# ╔═╡ 2e11adfc-4ea8-11eb-02bc-d550994774f8
function RLBase.legal_action_space(env::MaximizationBiasEnv)
    if env.position == 1
        (LEFT, RIGHT)
    else
        Base.OneTo(10)
    end
end

# ╔═╡ 3148b560-4ea8-11eb-0fc1-f77efb282ba4
function RLBase.legal_action_space_mask(env::MaximizationBiasEnv)
    m = fill(false, 10)
    if env.position == 1
        m[LEFT] = true
        m[RIGHT] = true
    else
        m .= true
    end
    m
end

# ╔═╡ 3a674788-4ea8-11eb-1931-b39142fffb47
function RLBase.reset!(env::MaximizationBiasEnv)
    env.position = 1
    env.reward = 0.0
    nothing
end

# ╔═╡ 3ec32b6c-4ea8-11eb-1575-319c528b4f02
RLBase.reward(env::MaximizationBiasEnv) = env.reward

# ╔═╡ 4127acde-4ea8-11eb-0e47-7d82bb5cb310
RLBase.is_terminated(env::MaximizationBiasEnv) = env.position == 3

# ╔═╡ 4309b632-4ea8-11eb-1d46-edd93efd6925
RLBase.state(env::MaximizationBiasEnv) = env.position

# ╔═╡ 45650d82-4ea8-11eb-39c2-71fc325d6c13
md"""
Now the environment is well defined.
"""

# ╔═╡ 56f8ba1c-4ea8-11eb-2ce7-8357675ee6ef
world = MaximizationBiasEnv()

# ╔═╡ 657eca7c-4ea8-11eb-3c42-b1a3085f6129
NS, NA = length(state_space(world)), length(action_space(world))

# ╔═╡ 8d15e872-4ea8-11eb-2fc2-7f3c4a90a85a
md"""
To calculate the percentage of chosing `LEFT` action in the first step, we'll create customized hook here:
"""

# ╔═╡ 6f798d14-4ea8-11eb-2825-8f6f17ff4f05
begin
	Base.@kwdef mutable struct CountOfLeft <: AbstractHook
		counts::Vector{Bool} = []
	end
	function (f::CountOfLeft)(::PreActStage, agent, env, action)
		if state(env) == 1
			push!(f.counts, action == LEFT)
		end
	end
end

# ╔═╡ d026cbb8-4ea8-11eb-29a2-dd75efe64466
md"""
Next we create two agent factories, Q-Learning and Double-Q-Learning.
"""

# ╔═╡ ea42c3c6-4ea8-11eb-1d82-0566785eeef3
create_double_Q_agent() = Agent(
    policy=QBasedPolicy(
        learner=DoubleLearner(
            L1=TDLearner(
                approximator=TabularQApproximator(
                    n_state=NS,
                    n_action=NA,
                    opt=Descent(0.1),
                ),
                method=:SARS,
                γ=1.,
                n=0,
            ),
            L2=TDLearner(
                approximator=TabularQApproximator(;
                    n_state=NS,
                    n_action=NA,
                    opt=Descent(0.1),
                ),
                method=:SARS,
                γ=1.,
                n=0,
            ),
            ),
        explorer=EpsilonGreedyExplorer(0.1;is_break_tie=true)
    ),
    trajectory=VectorSARTTrajectory()
)

# ╔═╡ eccb09aa-4ea8-11eb-08a0-2f6344c38669
create_Q_agent() = Agent(
    policy=QBasedPolicy(
        learner=TDLearner(
            approximator=TabularQApproximator(;n_state=NS, n_action=NA, opt=Descent(0.1)),
            method=:SARS,
            γ=1.0,
            n=0
            ),
        explorer=EpsilonGreedyExplorer(0.1;is_break_tie=true)
        ),
    trajectory=VectorSARTTrajectory()
)

# ╔═╡ 0e4b6994-4ea9-11eb-111f-7b37d6a716e5
begin
	DQ_stats = []
	for _ in 1:1000
		hook = CountOfLeft()
		run(create_double_Q_agent(), world, StopAfterEpisode(300),hook)
		push!(DQ_stats, hook.counts)
	end
	plot(mean(DQ_stats)*100, legend=:topright, label="double q", xlabel="Episodes", ylabel="% left actions from A")

	Q_stats = []
	for _ in 1:1000
		hook = CountOfLeft()
		run(create_Q_agent(), world, StopAfterEpisode(300),hook)
		push!(Q_stats, hook.counts)
	end
	plot!(mean(Q_stats)*100, legend=:topright, label="q")
	hline!([5], linestyle=:dash, label="optimal")
end

# ╔═╡ Cell order:
# ╟─6ce0a200-4ea7-11eb-123c-73e3a260351c
# ╠═9edb35ce-4ea7-11eb-353f-53459c337880
# ╠═172236ac-4ea8-11eb-1393-5bc3683bf6b0
# ╠═19876a70-4ea8-11eb-2b70-9d48875f7ae7
# ╠═247ada66-4ea8-11eb-3c0a-af131b73e9f3
# ╠═26b9f6f2-4ea8-11eb-38f3-3f1822c4acd2
# ╠═296e0ab6-4ea8-11eb-33c5-71be241708b1
# ╠═2c06d25a-4ea8-11eb-26e3-01d81d714859
# ╠═2e11adfc-4ea8-11eb-02bc-d550994774f8
# ╠═3148b560-4ea8-11eb-0fc1-f77efb282ba4
# ╠═3a674788-4ea8-11eb-1931-b39142fffb47
# ╠═3ec32b6c-4ea8-11eb-1575-319c528b4f02
# ╠═4127acde-4ea8-11eb-0e47-7d82bb5cb310
# ╠═4309b632-4ea8-11eb-1d46-edd93efd6925
# ╟─45650d82-4ea8-11eb-39c2-71fc325d6c13
# ╠═56f8ba1c-4ea8-11eb-2ce7-8357675ee6ef
# ╠═657eca7c-4ea8-11eb-3c42-b1a3085f6129
# ╟─8d15e872-4ea8-11eb-2fc2-7f3c4a90a85a
# ╠═6f798d14-4ea8-11eb-2825-8f6f17ff4f05
# ╟─d026cbb8-4ea8-11eb-29a2-dd75efe64466
# ╠═ea42c3c6-4ea8-11eb-1d82-0566785eeef3
# ╠═eccb09aa-4ea8-11eb-08a0-2f6344c38669
# ╠═0e4b6994-4ea9-11eb-111f-7b37d6a716e5


text:
# Transferring parameters into NamedTuple and SLArray to avoid passing around
# unnecessary information about the parameters.
function parameters_to_named_tuple(pvec::AbstractVector{S}) where {S <: AbstractParameter}
    tuple_names = Tuple(p.key  for p in pvec)
    tuple_vals  = map(p -> p.value, pvec)
    return NamedTuple{tuple_names}(tuple_vals)
end

function parameters_to_named_tuple(pvec::AbstractModel)
    return parameters_to_named_tuple(get_parameters(pvec))
end

"""
```
detexify(s::String)

detexify(s::Symbol)
```

Remove Unicode characters from the string `s`, replacing them with ASCII
equivalents. For example, `detexify(\"π\")` returns `\"pi\"`.

This code is copied directly from DSGE.jl.
"""
function detexify(s::String)
    s = replace(s, "α" => "alpha")
    s = replace(s, "β" => "beta")
    s = replace(s, "γ" => "gamma")
    s = replace(s, "δ" => "delta")
    s = replace(s, "ϵ" => "epsilon")
    s = replace(s, "ε" => "epsilon")
    s = replace(s, "ζ" => "zeta")
    s = replace(s, "η" => "eta")
    s = replace(s, "θ" => "theta")
    s = replace(s, "ι" => "iota")
    s = replace(s, "κ" => "kappa")
    s = replace(s, "λ" => "lambda")
    s = replace(s, "μ" => "mu")
    s = replace(s, "ν" => "nu")
    s = replace(s, "ξ" => "xi")
    s = replace(s, "π" => "pi")
    s = replace(s, "ρ" => "rho")
    s = replace(s, "ϱ" => "rho")
    s = replace(s, "σ" => "sigma")
    s = replace(s, "ς" => "sigma")
    s = replace(s, "τ" => "tau")
    s = replace(s, "υ" => "upsilon")
    s = replace(s, "ϕ" => "phi")
    s = replace(s, "φ" => "phi")
    s = replace(s, "χ" => "chi")
    s = replace(s, "ψ" => "psi")
    s = replace(s, "ω" => "omega")

    s = replace(s, "Α" => "Alpha")
    s = replace(s, "Β" => "Beta")
    s = replace(s, "Γ" => "Gamma")
    s = replace(s, "Δ" => "Delta")
    s = replace(s, "Ε" => "Epsilon")
    s = replace(s, "Ζ" => "Zeta")
    s = replace(s, "Η" => "Eta")
    s = replace(s, "Θ" => "Theta")
    s = replace(s, "Ι" => "Iota")
    s = replace(s, "Κ" => "Kappa")
    s = replace(s, "Λ" => "Lambda")
    s = replace(s, "Μ" => "Mu")
    s = replace(s, "Ν" => "Nu")
    s = replace(s, "Ξ" => "Xi")
    s = replace(s, "Π" => "Pi")
    s = replace(s, "Ρ" => "Rho")
    s = replace(s, "Σ" => "Sigma")
    s = replace(s, "Τ" => "Tau")
    s = replace(s, "Υ" => "Upsilon")
    s = replace(s, "Φ" => "Phi")
    s = replace(s, "Χ" => "Chi")
    s = replace(s, "Ψ" => "Psi")
    s = replace(s, "Ω" => "Omega")

    s = replace(s, "∂" => "d") # This line is not in the original DSGE.detexify though

    s = replace(s, "′" => "'")

    return s
end

function detexify(s::Symbol)
    Symbol(detexify(string(s)))
end

function mat_to_jld2(fn::String, outfn::String)
    mat_data = matread(fn)
    JLD2.jldopen(outfn, true, true, true, IOStream) do file
        for (k, v) in mat_data
            write(file, k, v)
        end
    end
end

"""
```
sparse_accumarray(subs, val, sz=(maximum(subs),))
```

constructs a sparse array using the same algorithm as
`accumarray`. See either `VectorizedRoutines.jl` or
MATLAB's documentation for more details
"""
function sparse_accumarray(subs, val, sz=(maximum(subs),))
    A = spzeros(eltype(val), sz...)
    for i = 1:length(val)
        @inbounds A[subs[i]] += val[i]
    end
    A
end


text:
using PMDPs
using PMDPs.LP
using POMDPs
using BSON, CSV
using DrWatson
using RandomNumbers.Xorshifts
using Random
using MCTS, DiscreteValueIteration
using DataFrames

using POMDPSimulators
using StaticArrays, Distributions # load pp

N = 10

dpw_solver = DPWSolver(;depth=50,
                exploration_constant=40.0, max_time=1.,
                enable_state_pw = false,
                keep_tree=true, show_progress=false, rng=Xorshift128Plus())

mcts_solver = MCTSSolver(;depth=50,
                exploration_constant=40.0, max_time=1.,
                rng=Xorshift128Plus())

"""
 ------ BUS ------------
"""
name = "bus_problem"
out_folder="individual"

display("Evaluating $name")
# sname = string("traces_bus_T=1000_c=55_expected_requests=80_nᵣ=3_objective=revenue_res_budget_μ=5",  ".bson")
sname = string("traces_bus_T=2000_c=55_expected_requests=160_nᵣ=3_objective=revenue_res_budget_μ=5",  ".bson")
# data = load(datadir("traces", sname))
data = PMDPs.load_traces(datadir("traces", sname))

PMDPs.process_data(data, PMDPs.flatrate; folder=out_folder, N=N)
PMDPs.process_data(data, PMDPs.hindsight; folder=out_folder, N=N)
# PMDPs.process_data(data, PMDPs.vi; folder=out_folder, N=N)
# PMDPs.process_data(data, PMDPs.fhvi; folder=out_folder, N=N)
PMDPs.process_data(data, PMDPs.mcts; folder=out_folder, N=N, method_info="dpw", mcts_solver=dpw_solver)
PMDPs.process_data(data, PMDPs.mcts; folder=out_folder, N=N, method_info="vanilla", mcts_solver=mcts_solver)

println("$name Done.")

N=10
"""
 ------ LP ------------
"""
pps = [
    # Dict(pairs((nᵣ=3, c=3, T=10, expected_res=3., res_budget_μ=5.))),
    # # Dict(pairs((nᵣ=6, c=5, T=100, expected_res=60., res_budget_μ=5.))),
    # Dict(pairs((nᵣ=10, c=5, T=100, expected_res=100., res_budget_μ=5.))),
    # Dict(pairs((nᵣ=10, c=40, T=1000, expected_res=Float64(800), res_budget_μ=5.))),
    # Dict(pairs((nᵣ=50, c=40, T=1000, expected_res=Float64(4000), res_budget_μ=5.))),
    Dict(pairs((nᵣ=3, c=3, T=10, expected_res=3., res_budget_μ=5., objective=:utilization))),
    # Dict(pairs((nᵣ=6, c=5, T=100, expected_res=60., res_budget_μ=5., objective=:utilization))),
    Dict(pairs((nᵣ=10, c=5, T=100, expected_res=100., res_budget_μ=5., objective=:utilization))),
    Dict(pairs((nᵣ=10, c=40, T=1000, expected_res=Float64(800), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((nᵣ=50, c=40, T=1000, expected_res=Float64(4000), res_budget_μ=5., objective=:utilization)))
]
name = "linear_problem"
out_folder="individual"

for pp_params in pps
    display("Evaluating $name with $pp_params")
    sname = savename("traces_lp", pp_params,  "bson")
    # data = load(datadir("traces", sname))
    data = PMDPs.load_traces(datadir("traces", sname))

    # PMDPs.process_data(data, PMDPs.flatrate; folder=out_folder, N=N)
    PMDPs.process_data(data, PMDPs.hindsight; folder=out_folder, N=N)
    # PMDPs.process_data(data, PMDPs.vi; folder=out_folder, N=N)
    # PMDPs.process_data(data, PMDPs.fhvi; folder=out_folder, N=N)

    # PMDPs.process_data(data, PMDPs.mcts; folder=out_folder, N=N, method_info="dpw", mcts_solver=dpw_solver)
    # PMDPs.process_data(data, PMDPs.mcts; folder=out_folder, N=N, method_info="vanilla", mcts_solver=mcts_solver)
end
println("LP Done.")

"""
 ------ Graph ------------
"""
seed = 12
gpps = [
    # Dict(pairs((NV=5, NE=8, seed=1, NP=20, c=5, T=100, expected_res=Float64(80), res_budget_μ=5.))),
    # Dict(pairs((NV=8, NE=20, seed=1, NP=50, c=10, T=1000, expected_res=Float64(400), res_budget_μ=5.))),
    # Dict(pairs((NV=15, NE=30, seed=1, NP=100, c=10, T=1000, expected_res=Float64(600), res_budget_μ=5.))),
    # Dict(pairs((NV=30, NE=45, seed=1, NP=100, c=10, T=1000, expected_res=Float64(900), res_budget_μ=5.)))
    Dict(pairs((NV=5, NE=8, seed=1, NP=20, c=5, T=100, expected_res=Float64(80), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((NV=8, NE=20, seed=1, NP=50, c=10, T=1000, expected_res=Float64(400), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((NV=15, NE=30, seed=1, NP=100, c=10, T=1000, expected_res=Float64(600), res_budget_μ=5., objective=:utilization))),
    Dict(pairs((NV=30, NE=45, seed=1, NP=100, c=10, T=1000, expected_res=Float64(900), res_budget_μ=5., objective=:utilization)))
]
name = "graph_problem"
out_folder="individual"
for pp_params in gpps
    name = "graph_problem"
    display("Evaluating $name with $pp_params,")
    sname = savename("traces_gp", pp_params,  "bson")
    data = load(datadir("traces", sname))
    data = PMDPs.load_traces(datadir("traces", sname))

    PMDPs.process_data(data, PMDPs.flatrate; folder=out_folder, N=N)
    PMDPs.process_data(data, PMDPs.hindsight; folder=out_folder, N=N)
    # PMDPs.process_data(data, PMDPs.vi; N=N)
    # PMDPs.process_data(data, PMDPs.fhvi; folder=out_folder, N=N)
    # PMDPs.process_data(data, PMDPs.mcts; N=N)
    PMDPs.process_data(data, PMDPs.mcts; folder=out_folder, N=N, method_info="dpw", mcts_solver=dpw_solver)
    PMDPs.process_data(data, PMDPs.mcts; folder=out_folder, N=N, method_info="vanilla", mcts_solver=mcts_solver)
end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/FEMBasis.jl/blob/master/LICENSE

using Calculus

function get_reference_element_coordinates end
function eval_basis! end
function eval_dbasis! end

function get_dim(p::Expr)
    pargs = p.args[2:end]
    :w in pargs && return 3
    :v in pargs && return 2
    :u in pargs && return 1
end

"""
    function calculate_basis_coefficients(polynomial::String, coordinates::Vararg{Tuple})

Calculate "interpolate coefficient matrix" for some polynomial p.

# Examples
That is, if we have polynomial p = 1 + u + v and coordinates (0,0), (1,0), (0,1),
we find A*p such that first row is the first coordinate, second row is second
coordinate and so on:
```julia
julia> p = "1 + u + w"
julia> X = ((0.0,0.0), (1.0,0.0), (0.0,1.0))
julia> calculate_basis_coefficient(p, X)
[1.0 0.0 0.0 # <-- p(0.0,0.0) = 1.0     = [1.0 0.0 0.0] * [1.0, u, v]
 1.0 1.0 0.0 # <-- p(1.0,0.0) = 1.0 + u = [1.0 1.0 0.0] * [1.0, u, v]
 1.0 0.0 1.0] # <- p(0.0,1.0) = 1.0 + v = [1.0 0.0 1.0] * [1.0, u, v]
```
"""
function calculate_basis_coefficients(p::Expr, X::Tuple)
    pargs = p.args[2:end]
    dim = get_dim(p)
    nbasis = length(X)
    sandbox = Module(:__SANDBOX__)
    vars = (:u, :v, :w)

    # calculate coefficient matrix A and inversion
    A = zeros(nbasis, nbasis)
    for i=1:nbasis
        for j=1:nbasis
            for (k,c) in enumerate(X[j])
                code = :( $(vars[k]) = $c )
                eval(sandbox, code)
            end
            A[j,i] = eval(sandbox, pargs[i])
        end
    end
    return A
end

"""
"""
function calculate_interpolation_polynomials(p::Expr, A::Matrix)
    pargs = p.args[2:end]
    dim = get_dim(p)
    nbasis = size(A, 1)
    invA = inv(A)
    equations = Expr[]
    for j=1:nbasis
        equation = Expr(:call, :+)
        for i=1:nbasis
            isapprox(invA[i,j], 0.0) && continue
            if typeof(pargs[i]) <: Int && pargs[i] == 1
                push!(equation.args, invA[i,j])
            else
                if isapprox(invA[i,j], 1.0)
                    push!(equation.args, :($(pargs[i])))
                elseif isapprox(invA[i,j], -1.0)
                    push!(equation.args, :(-$(pargs[i])))
                else
                    push!(equation.args, :($(invA[i,j]) * $(pargs[i])))
                end
            end
        end
        push!(equations, equation)
    end
    return equations
end

"""
Calculate derivatives of basis functions with respect to parameters u, v, w.
"""
function calculate_interpolation_polynomial_derivatives(basis, dim)
    equations = Vector[]
    vars = [:u, :v, :w]
    nbasis = length(basis)
    for j=1:nbasis
        deq = []
        for k=1:dim
            #println("∂($(basis[j])) / ∂$(vars[k])")
            der = differentiate(basis[j], vars[k])
            push!(deq, der)
        end
        push!(equations, deq)
    end
    return equations
end

function create_basis{nbasis,dim}(name::Symbol, description::String, X::NTuple{nbasis, NTuple{dim, Float64}}, basis, dbasis)

    Q = Expr(:block)
    for i=1:nbasis
        push!(Q.args, :(N[$i] = $(basis[i])))
    end

    V = Expr(:block)
    for i=1:nbasis
        for k=1:dim
            push!(V.args, :(dN[$k,$i] = $(dbasis[i][k])))
        end
    end

    # FIXME: this can be done better somehow
    if dim == 1
        unpack = :((u,) = xi)
    elseif dim == 2
        unpack = :((u, v) = xi)
    else
        unpack = :((u, v, w) = xi)
    end

    code = quote
        type $name <: FEMBasis.AbstractBasis
        end

        function Base.size(::Type{$name})
            return ($dim, $nbasis)
        end

        function Base.size(::Type{$name}, j::Int)
            j == 1 && return $dim
            j == 2 && return $nbasis
        end

        function Base.length(::Type{$name})
            return $nbasis
        end

        function FEMBasis.get_reference_element_coordinates(::Type{$name})
            return $X
        end

        function FEMBasis.eval_basis!{T}(::Type{$name}, N::Matrix{T}, xi)
            $unpack
            $Q
            return N
        end

        function FEMBasis.eval_dbasis!{T}(::Type{$name}, dN::Matrix{T}, xi)
            $unpack
            $V
            return dN
        end
    end

    return code
end

function create_basis{nbasis,dim}(name::Symbol, description::String, X::NTuple{nbasis, NTuple{dim, Float64}}, p::Expr)
    A = calculate_basis_coefficients(p, X)
    basis = calculate_interpolation_polynomials(p, A)
    dbasis = calculate_interpolation_polynomial_derivatives(basis, dim)
    return create_basis(name, description, X, basis, dbasis)
end

function create_basis{nbasis,dim}(name::Symbol, description::String, X::NTuple{nbasis, NTuple{dim, Float64}}, p::String)
    create_basis(name, description, X, parse(p))
end

function create_basis{nbasis,dim}(name::Symbol, description::String, X::NTuple{nbasis, NTuple{dim, Float64}}, basis_::NTuple{nbasis, String})
    basis = [parse(b) for b in basis_]
    dbasis = calculate_interpolation_polynomial_derivatives(basis, dim)
    return create_basis(name, description, X, basis, dbasis)
end

function create_basis{nbasis,dim}(name::Symbol, description::String, X::NTuple{nbasis, NTuple{dim, Float64}}, basis::NTuple{nbasis, Expr})
    dbasis = calculate_interpolation_polynomial_derivatives(basis, dim)
    return create_basis(name, description, X, basis, dbasis)
end


text:
"""
Two different temporal plots, one with memory and one with some percolation rate.
While varying the percolation rate, we compare the ratio of the depths of the
graphs
"""

using QuNet
using LightGraphs, SimpleWeightedGraphs, GraphPlot, MetaGraphs, Plots, Colors, Statistics
using LaTeXStrings
using JLD
using Parameters

datafile = "data/bandwidth_wtih_memory"

# Params
num_trials = 500::Int64
perc_range = (0.00, 0.05, 1.00)::Tuple{Float64, Float64, Float64}
grid_size = 10::Int64
time_depth = 20::Int64
num_pairs = 50::Int64
asynchronus_weight = 100::Int64
maxpaths=4

generate_new_data = true
if generate_new_data == true

    # Generate ixi graph and extend it in time
    G = GridNetwork(grid_size, grid_size)

    max_depth_data = [[],[],[],[]]
    max_depth_err = [[],[],[],[]]
    max_depth_mem_data = [[],[],[],[]]
    max_depth_mem_err = [[],[],[],[]]

    for maxpaths in 1:4
        println("## Collecting for maxpath $maxpaths")
        for i in perc_range[1]:perc_range[2]:perc_range[3]
            println("Collecting for memory percolation rate : $i")
            raw_max_depth = []
            raw_max_depth_mem = []
            for j in 1:num_trials
                # Extend in time without memory
                T = QuNet.TemporalGraph(G, time_depth, memory_prob=0.0)
                # Extend in time with memory links:
                T_mem = QuNet.TemporalGraph(G, time_depth, memory_prob=i, memory_costs = zero_costvector())

                # Get i random userpairs. Ensure src nodes are fixed on T=1, dst nodes are asynchronus.
                user_pairs = make_user_pairs(T, num_pairs, src_layer=-1, dst_layer=-1)
                mem_user_pairs = make_user_pairs(T, num_pairs, src_layer=-1, dst_layer=-1)

                # Add async nodes
                QuNet.add_async_nodes!(T_mem, mem_user_pairs, ϵ=asynchronus_weight)
                QuNet.add_async_nodes!(T, user_pairs, ϵ=asynchronus_weight)

                # Get pathset data
                pathset_mem, dum1, dum2 = QuNet.greedy_multi_path!(T_mem, QuNet.purify, mem_user_pairs, maxpaths)
                pathset, dum1, dum2 = QuNet.greedy_multi_path!(T, QuNet.purify, user_pairs, maxpaths)
                # Pathset is an array of vectors containing edges describing paths between end-user pairs
                # Objective: find the largest timedepth used in the pathsets

                max_depth = QuNet.max_timedepth(pathset, T)
                max_depth_mem = QuNet.max_timedepth(pathset_mem, T)

                # Get the ratio of these two quantities. Add it to data array
                push!(raw_max_depth, max_depth)
                push!(raw_max_depth_mem, max_depth_mem)
            end

            # Average the raw data, add it to plot data:
            push!(max_depth_data[maxpaths], mean(raw_max_depth))
            push!(max_depth_err[maxpaths], std(raw_max_depth)/sqrt(num_trials-1))
            push!(max_depth_mem_data[maxpaths], mean(raw_max_depth_mem))
            push!(max_depth_mem_err[maxpaths], std(raw_max_depth)/sqrt(num_trials-1))
        end
    end

    # Save data
    d = Dict{Symbol, Any}()
    @pack! d = num_trials, perc_range, grid_size, time_depth, num_pairs, asynchronus_weight, plot_data, error_data,
    max_depth_data, max_depth_err, max_depth_mem_data, max_depth_mem_err
    save("$datafile.jld", "data", d)

else
    # Load data
    if !isfile("$datafile.jld")
        error("$datafile.jld not found in working directory")
    end
    d = load("$datafile.jld")["data"]
    @unpack num_trials, perc_range, grid_size, time_depth, num_pairs, asynchronus_weight, plot_data, error_data, max_depth_data, max_depth_err, max_depth_mem_data, max_depth_mem_err = d
end

# Plot
x = collect(perc_range[1]:perc_range[2]:perc_range[3])
# plot(x, plot_data, yerr = error_data, legend = false)

plot(x, max_depth_data[1], seriestype=:scatter, yerr = max_depth_err[1], label="max depth 1", legend=:topleft,
guidefontsize=14, tickfontsize=12, legendfontsize=8, fontfamily="computer modern",
color_palette = palette(:Spectral_4, 4), markersize=5)
plot!(x, max_depth_data[2], seriestype=:scatter, yerr = max_depth_mem_err[2], label="max depth 2", markersize=5)
plot!(x, max_depth_data[3], seriestype=:scatter, yerr = max_depth_mem_err[3], label="max depth 3", markersize=5)
plot!(x, max_depth_data[4], seriestype=:scatter, yerr = max_depth_mem_err[4], label="max depth 4", markersize=5)
plot!(x, max_depth_mem_data[1], seriestype=:scatter, yerr = max_depth_mem_err[1], label="max depth with memory 1", markershape=:utriangle, markersize=5)
plot!(x, max_depth_mem_data[2], seriestype=:scatter, yerr = max_depth_mem_err[2], label="max depth with memory 2", markershape=:utriangle, markersize=5)
plot!(x, max_depth_mem_data[3], seriestype=:scatter, yerr = max_depth_mem_err[3], label="max depth with memory 3", markershape=:utriangle, markersize=5)
plot!(x, max_depth_mem_data[4], seriestype=:scatter, yerr = max_depth_mem_err[4], label="max depth with memory 4", markershape=:utriangle, markersize=5)
xaxis!(L"$\textrm{Proportion of Nodes with Quantum Memory}$")
yaxis!("Time depth")

savefig("plots/bandwidth_with_memory_rate.png")
savefig("plots/bandwidth_with_memory_rate.pdf")

max_stars_repo_name: FalafelGood/QuNet

max_issues_repo_name: FalafelGood/QuNet

max_forks_repo_name: FalafelGood/QuNet


text:
using Pkg
Pkg.activate(".")
Pkg.instantiate()

using Remark

paths_ignore = [".git", ".github", "javascript", "src", "stylesheets"]
note_paths = setdiff(filter(x -> isdir(x), readdir(".")), paths_ignore)

for path in note_paths
    Remark.slideshow("overview", options = Dict("ratio" => "16:9", "highlightStyle" => "github", "highlightLanguage" => "julia"))
end

max_stars_repo_name: vsrikrish/systems-lecture-notes

max_issues_repo_name: vsrikrish/systems-lecture-notes

max_forks_repo_name: vsrikrish/systems-lecture-notes


text:
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

module BAData

export BAInput, BASparseMatrix, BAOutput, insert_reproj_err_block!, insert_w_err_block!, empty_ba_output, load_ba_input,
    N_CAM_PARAMS, ROT_IDX, C_IDX, F_IDX, X0_IDX, RAD_IDX

const N_CAM_PARAMS = 11
const ROT_IDX = 1
const C_IDX = 4
const F_IDX = 7
const X0_IDX = 8
const RAD_IDX = 10

struct BAInput
    n::Int
    m::Int
    p::Int
    cams::Matrix{Float64}
    X::Matrix{Float64}
    w::Vector{Float64}
    feats::Matrix{Float64}
    obs::Matrix{Int}
end

struct BASparseMatrix
    "Number of cams"
    n::Int
    "Number of points"
    m::Int
    "Number of observations"
    p::Int
    nrows::Int
    ncols::Int
    """
    Int[nrows + 1]. Defined recursively as follows:
    rows[0] = 0
    rows[i] = rows[i-1] + the number of nonzero elements on the i-1 row of the matrix
    """
    rows::Vector{Int}
    "Column index in the matrix of each element of vals. Has the same size"
    cols::Vector{Int}
    "All the nonzero entries of the matrix in the left-to-right top-to-bottom order"
    vals::Vector{Float64}
    BASparseMatrix(n::Int, m::Int, p::Int) = new(n, m, p, 2 * p + p, N_CAM_PARAMS * n + 3 * m + p, [0], [], [])
end

mutable struct BAOutput
    reproj_err::Matrix{Float64}
    w_err::Vector{Float64}
    jacobian::BASparseMatrix
end

empty_ba_output() = BAOutput(Array{Float64}(undef, 0, 0), [], BASparseMatrix(0, 0, 0))

function insert_reproj_err_block!(matrix::BASparseMatrix, obsIdx::Int, camIdx::Int, ptIdx::Int, J::AbstractMatrix{Float64})
    # We use zero-based indexing for storage, but Julia uses 1-bsed indexing
    # Hence, the conversion
    obsIdxZeroBased = obsIdx - 1
    camIdxZeroBased = camIdx - 1
    ptIdxZeroBased = ptIdx - 1
    n_new_cols = N_CAM_PARAMS + 3 + 1
    lastrow = matrix.rows[end]
    push!(matrix.rows, lastrow + n_new_cols, lastrow + n_new_cols + n_new_cols)
    for i_row ∈ 1:2
        for i ∈ 1:N_CAM_PARAMS
            push!(matrix.cols, N_CAM_PARAMS * camIdxZeroBased + (i - 1))
            push!(matrix.vals, J[i_row, i])
        end
        col_offset = N_CAM_PARAMS * matrix.n
        for i ∈ 1:3
            push!(matrix.cols, col_offset + 3 * ptIdxZeroBased + (i - 1))
            push!(matrix.vals, J[i_row, N_CAM_PARAMS + i])
        end
        col_offset += 3 * matrix.m
        val_offset = N_CAM_PARAMS + 3
        push!(matrix.cols, col_offset + obsIdxZeroBased);
        push!(matrix.vals, J[i_row, val_offset + 1]);
    end
end

function insert_w_err_block!(matrix::BASparseMatrix, wIdx::Int, w_d::Float64)
    # We use zero-based indexing for storage, but Julia uses 1-bsed indexing
    # Hence, the conversion
    wIdxZeroBased = wIdx - 1
    push!(matrix.rows, matrix.rows[end] + 1)
    push!(matrix.cols, N_CAM_PARAMS * matrix.n + 3 * matrix.m + wIdxZeroBased)
    push!(matrix.vals, w_d)
end

function load_ba_input(fn::AbstractString)::BAInput
    fid = open(fn)
    lines = readlines(fid)
    close(fid)
    line=split(lines[1]," ")
    n = parse(Int,line[1])
    m = parse(Int,line[2])
    p = parse(Int,line[3])
    off = 2

    one_cam = zeros(Float64,N_CAM_PARAMS,1)
    line=split(lines[off]," ")
    for i in 1:N_CAM_PARAMS
        one_cam[i] = parse(Float64,line[i])
    end
    cams = repeat(one_cam,1,n)
    off += 1

    one_X = zeros(Float64,3,1)
    line=split(lines[off]," ")
    for i in 1:3
        one_X[i] = parse(Float64,line[i])
    end
    X = repeat(one_X,1,m)
    off += 1

    one_w = parse(Float64,lines[off])
    w = repeat([one_w],p)
    off += 1

    one_feat = zeros(Float64,2,1)
    line=split(lines[off]," ")
    for i in 1:2
        one_feat[i] = parse(Float64,line[i])
    end
    feats = repeat(one_feat,1,p)

    camIdx = 1
    ptIdx = 1
    obs = zeros(Int,2,p)
    for i in 1:p
        obs[1,i] = camIdx
        obs[2,i] = ptIdx
        camIdx = (camIdx%n) + 1
        ptIdx = (ptIdx%m) + 1
    end

    BAInput(n, m, p, cams, X, w, feats, obs)
end

end

max_stars_repo_name: dsyme/ADBench

max_issues_repo_name: dsyme/ADBench

max_forks_repo_name: dsyme/ADBench


text:
#=
Functions:
  contract_edge(e,x,k)
  contract_edge_k_1(e,x)
  contract(A,x,m)
  contract_k_1(A,x)
  contract_k_1!(A(type needs to be determined),x,k,y)
  contract_multi(A,vs)
  contract_edge_k_1!(indices,val,x,res)
#------------------------------------------------------------------------------
						           Contraction
------------------------------------------------------------------------------=#

"""-----------------------------------------------------------------------------
    contract_edge(e,x,k)

This function takes in an edge of a super symmetric tensor and computes the
resulting edges which result from contracting the the edge along k modes with
the vector x.

Input:
------
* e -(Tuple(Array{Int,1},Float)):

    A list of indices paired with an edge value. Note that the list of indices
    corresponds to multiple sets of indices because we consider all.
    permutations.
* x -(Array{Float,1})

    The vector to contract with.
* k - (Int):

    A positive integer which corresponds to the number of modes to contract
    along, must be greater than 0, and less than or equal to the cardinality
    of the edge.

Output:
-------
* condensed_dict - (Dict{Array{Int,1},AbstractFloat}

    The hyper edges in the lower order tensor which are formed by contracting
    the vector along the hyperedge e.
-----------------------------------------------------------------------------"""
function contract_edge(e::Tuple{Array{Int,1},M},x::Array{N,1},k::Int) where {N <: AbstractFloat, M<:AbstractFloat}
    order = length(e)

    (indices,val) = e
    condensed_dict = Dict{Array{Int,1},N}()
    visited_sub_indices = Dict{Array{Int,1},Dict{Array{Int,1},N}}()

    for i in 1:length(indices)
        sub_edge = deleteat!(copy(indices),i)
        if !haskey(visited_sub_indices,sub_edge)
            if k == 1
                condensed_dict[sub_edge] = val*x[indices[i]]
            else
                visited_sub_indices[sub_edge] = contract_edge((sub_edge,val*x[indices[i]]),x,k-1)
            end
        end
    end

    if k != 1
        for (_,sub_dict) in visited_sub_indices
            reduce_dictionaries!(condensed_dict,sub_dict)
        end
    end

    return condensed_dict
end

"""-----------------------------------------------------------------------------
    contract_edge_k_1(e,x)

  This function takes in an edge of a super symmetric tensor and computes the
resulting edges which result from contracting the the edge along k-1 modes with
the vector x, where k is the order of the hyper edge.

Input:
------
* e -(Tuple(Array{Int,1},AbstractFloat)):

    a list of sorted indices paired with an edge value. Note that the list of
    indices corresponds to multiple sets of indices because we consider all
    permutations.
* x -(Array{AbstractFloat,1})

    The vector of the same dimenionality of the tensor, to contract with.

Output:
-------
* contraction_vals - (Array{Tuple{Array{Int,1},AbstractFloat}})

    The hyper edges in the lower order tensor which are formed by contracting
    the vector along the hyperedge e.
-----------------------------------------------------------------------------"""
function contract_edge_k_1(e::Tuple{Array{Int,1},N},x::Array{N,1}) where N <: AbstractFloat
    (indices,val) = e
    order = length(indices)

    visited_sub_indices = Set{Array{Int,1}}()
    contraction_vals = Array{Tuple{Array{Int,1},N}}(undef,1)

    for i in 1:order
        sub_edge = deleteat!(copy(indices),i)
        if !in(sub_edge,visited_sub_indices)#haskey(scaling_factors,sub_edge)
            scaling = multiplicity_factor(sub_edge)
            push!(visited_sub_indices,sub_edge)
            push!(contraction_vals,([indices[i]],scaling*val*prod(x[sub_edge])))
        end
    end
    return contraction_vals
end

"""-----------------------------------------------------------------------------
    contract(A,x,m)

  This function contracts the tensor along m modes. Note that when the tensor is
dense this function uses Base.Cartesian, and thus in order to generate the loops
with a variable used a trick which instantiates empty arrays of length 0, and
passes them to another function which can pull the orders out to generate the
loops.

Input:
------
* A -(SSSTensor or Array{AbstractFloat,k}):

    The tensor to contract.
* x - (Array{AbstractFloat,1}):

    A vector of numbers to contract with.
* m - (Int)

    The number of modes to contract A with x along.

Output:
-------
* y - (SSSTensor or CSC Matrix or Array{Float64,k-m}):

    The output vector of Ax^m. THe output will be sparse if the input tensor is
    sparse, and dense otherwise. When the output is second order, and A is
    sparse, then the output will be a sparse matrix.
-----------------------------------------------------------------------------"""
function contract(A::SSSTensor, x::Array{N,1},m::Int) where {N <: AbstractFloat}
    @assert length(x) == A.cubical_dimension
    k = order(A)
    @assert 0 < m <= k

    new_edges = Dict{Array{Int,1},N}()
    #compute contractions
    for edge in A.edges
        new_e = contract_edge(Tuple(edge),x,m)
        reduce_dictionaries!(new_edges,new_e)
    end

    if k == m
        for (_,v) in new_edges
            return v
        end
    elseif k - m == 1
        y = zeros(length(x))
        for (e,v) in new_edges
            y[e[1]] = v
        end
        return y
    elseif k - m == 2
	  index = 0
	  nnzs = length(new_edges)
	  I = zeros(Int64,2*nnzs)
	  J = zeros(Int64,2*nnzs)
	  V = zeros(N,2*nnzs)

	  for (e,val) in new_edges
	     i,j = e
		 if i == j
		   index += 1
		   I[index] = i
		   J[index] = j
		   V[index] = val
	     else
		   index += 2
 	       I[index-1] = i
		   J[index-1] = j
		   I[index] = j
		   J[index] = i
		   V[index-1] = val
		   V[index] = val
		 end
	  end
	  return sparse(I[1:index],J[1:index],V[1:index],A.cubical_dimension,A.cubical_dimension)
	else
        return SSSTensor(new_edges,A.cubical_dimension)
    end
end

#Dense Case
function contract(A::Array{N,k}, x::Array{M,1},m::Int64) where {M <: AbstractFloat,N <: AbstractFloat,k}

    return dense_contract(A,x,zeros(Int,repeat([0],m)...),
	                      zeros(Int,repeat([0],k-m)...))
end

@generated function dense_contract(A::Array{N,k}, x::Array{M,1},
                                   B::Array{Int,m}, C::Array{Int,p}) where
								   {M<:AbstractFloat,N<:AbstractFloat,k,m,p}
    quote
        n = size(A)[1]
        @assert n == length(x)
        @assert $k >= m

        y = zeros(N,repeat([n],$k - $m)...)

        @nloops $k i A begin
            xs = prod(x[collect(@ntuple $m j-> i_{j+$p})])
            (@nref $p y i) += xs*(@nref $k A i)
        end
		if $k == m
	  		for val in y
	    		return val
	  		end
		else
          return y
        end
      end
end

"""-----------------------------------------------------------------------------
    contract_k_1(A,x)

This function contracts the tensor along k-1 modes to produce a vector. This
will produce the same result as contract(A,x,k-1), but runs in a much faster
time.

Inputs
------
* A -(SSSTensor):

    The tensor to contract.
* x - (Array{AbstractFloat,1}):

    A vector of numbers to contract with.
Outputs
-------
* y - (Array{AbstractFloat,1}):

    The output vector of Ax^{k-1}.
-----------------------------------------------------------------------------"""
function contract_k_1(A::Ten, x::Array{N,1}) where {N <: AbstractFloat,Ten <: AbstractSSTen}
    @assert length(x) == A.cubical_dimension

    new_edges = Array{Tuple{Array{Int,1},N}}(undef,0)
    y = zeros(A.cubical_dimension)

    #compute contractions
    for edge in A.edges
        contracted_edges = contract_edge_k_1(Tuple(edge),x)
        push!(new_edges,contracted_edges...)
    end
    #reduce edges and copy into new vector
    edge_dict = reduce_edges(new_edges)

    for (i,v) in edge_dict
        y[i[1]] = v
    end
    return y
end

function contract_k_1!(A::Ten, x::Array{N,1},y::Array{N,1}) where {N <: AbstractFloat,Ten <: AbstractSSTen}
    @assert length(x) == A.cubical_dimension

    new_edges = Array{Tuple{Array{Int,1},N}}(undef,0)
	@inbounds for i in 1:A.cubical_dimension
		y[i] = zero(N)
	end

    #compute contractions
    for edge in A.edges
        contracted_edges = contract_edge_k_1(Tuple(edge),x)
        push!(new_edges,contracted_edges...)
    end
    #reduce edges and copy into new vector
    edge_dict = reduce_edges(new_edges)

    for (i,v) in edge_dict
        y[i[1]] = v
    end
    return y
end

"""-----------------------------------------------------------------------------
    contract_multi(A,vs)

  This function computes the result of contracting the tensor A by the columns
of the array Vs.

-----------------------------------------------------------------------------"""
function contract_multi(A::SSSTensor, Vs::Array{N,2}) where N <: AbstractFloat
  k = order(A)
  n,m = size(Vs)
  @assert m <= k
  @assert n == A.cubical_dimension

  i = 1
  while true
    if k - i >= 2
      global A_sub = contract(A,Vs[:,i],1)
    elseif k - i == 1
	  global A_sub = A_sub*Vs[:,i]
    elseif k - i == 0
	  global A_sub = dot(A_sub,Vs[:,i])
    end
	i += 1

	if i > m #no more vectors
	  return A_sub
    end
  end
end

#TODO: make this take in an arbitrary number of vectors.
function contract(A::SSSTensor,v::Array{N,1},u::Array{N,1}) where N <: AbstractFloat
  return contract_multi(A,hcat(v,u))
end

"""-----------------------------------------------------------------------------
    contract_edge_k_1!(indices,val,x,res)

  This function computes the value of contracting one hyper edge with an input
vector x to all but 1 mode, and saves the result in the appropriate locations
in res.

Inputs:
-------
* indices - (Array{Int,2}):

    2D array with all the indices of the hyper edges in the tensor. Note that
    the number of rows is not necesarily equal to the non-zeros in the tensors.
    The indices of each row is assumed to be sorted in increasing order.
* nnz_val -  (N <: AbstractFloat):

    The non-zero associated with the hyper edges.
* x - (Array{N <: AbstractFloat, 1}):

    The vector to contract the hyper edge with.
* res_val - (Array{N <:AbstractFloat, 1}):

    The location to update the solution to. res is assumed to be initialized and
    so the values are simply added to the entries of y, this is done because
    this function is considered a helper function to the contract function o
    verloaded for the dense array representation of symmetric tensors.
-----------------------------------------------------------------------------"""
@inline function contract_edge_k_1!(edge::Tuple{Array{Int,1},N},
                                    x::Array{N,1},res::Array{N,1},
									sub_edge::Array{Int,1},
									multiplicities::Dict{Int,Int},
									final_counts::Array{Int,1},
									ord::Int)where {M,T, N <: AbstractFloat}
	(indices, nnz_val) = edge
    #ord = length(indices)
	prev_index = -1
	#sub_edge = Array{Int,1}(undef,ord - 1)

	for j = 1:ord
		#only compute contraction once per index
		if prev_index != indices[j]

			#form sub_edge
			i = 1
			for k in Base.Iterators.flatten((1:j-1,j+1:ord))
				sub_edge[i] = indices[k]
				i += 1
			end

			val = nnz_val * multiplicity_factor(sub_edge,multiplicities,final_counts)
			for k in sub_edge
				val *= x[k]
			end
			res[indices[j]] += val
		end
		prev_index = indices[j]
	end

end

"""-----------------------------------------------------------------------------
    contract_k_1(A,x,k)

  Contracts the vector x with the tensor to all but one mode

-----------------------------------------------------------------------------"""
function contract_k_1(A::COOTen,x::Array{N,1}) where {N <: AbstractFloat}
	@assert A.cubical_dimension == length(x)

	y = zeros(A.cubical_dimension) #resulting vector

	if A.order == 3
		#need a guarantee of unique indices in the hyper edges.
		tri_contract!(A,x,y)
	else
		#preallocate helper functions
		sub_edge = Array{Int,1}(undef,A.order - 1)
		multiplicities = Dict{Int,Int}()
		final_counts = zeros(Int,A.order)

		for edge in A
			contract_edge_k_1!(edge,x,y,sub_edge,multiplicities,final_counts,A.order)
		end

	end
	return y
end

"""-----------------------------------------------------------------------------
    contract_k_1(A,x,k)

  Contracts the vector x with the tensor to all but one mode

-----------------------------------------------------------------------------"""
function contract_k_1!(A::COOTen,x::Array{N,1},y::Array{N,1}) where {N <: AbstractFloat}
	@assert A.cubical_dimension == length(x)

	for i in 1:A.cubical_dimension
		y[i] = zero(N)
	end

	if A.order == 3
		tri_contract!(A,x,y)
	else

		#preallocate helper functions
		sub_edge = Array{Int,1}(undef,A.order - 1)
		multiplicities = Dict{Int,Int}()
		final_counts = zeros(Int,A.order)

		for edge in A
			contract_edge_k_1!(edge,x,y,sub_edge,multiplicities,final_counts,A.order)
		end
	end

end

"""-----------------------------------------------------------------------------
    tri_contract!(A,x,y)

    Special case of contract_k_1!, when the input tensor is one of unique
  triangles then when we write out the contraction operation and loop over the
  rows of COOTen.indices.

  Inputs:
  -------
  A - (COOTen sparse tensor):
    The tensor to contract.

  x - (Array{AbstractFloat,1}):
    The vector to contract the tensor with

  y - (Array{AbstractFloat,1}):
    The place to store the output results, assumed to be zeroed.

-----------------------------------------------------------------------------"""
function tri_contract!(A::COOTen,x::Array{N,1},y::Array{N,1}) where {N <: AbstractFloat}

	@assert A.order == 3
	#y = zeros(Float64,A.cubical_dimension)

	for i in 1:A.unique_nnz
		ci = A.indices[i,1]
        cj = A.indices[i,2]
        ck = A.indices[i,3]
		y[ci] += 2*x[cj]*x[ck]*A.vals[i]
        y[cj] += 2*x[ci]*x[ck]*A.vals[i]
        y[ck] += 2*x[cj]*x[ci]*A.vals[i]
	end
end

"""-----------------------------------------------------------------------------
    inner_product(A::COOTen,B::COOTen)

    Computes the inner product between two COOTensors. currently assuming all
    indices in each row are unique

-----------------------------------------------------------------------------"""
function inner_product(A::COOTen,B::COOTen)

	@assert A.cubical_dimension == B.cubical_dimension
	@assert A.order == B.order

	multiplicity_factor = factorial(A.order )
	inner_product_val = 0.0

	A_head = 1
	B_head = 1
	while A_head <= A.unique_nnz && B_head <= B.unique_nnz
		if A.indices[A_head,:] == B.indices[B_head,:]
			println(A_head, A.vals[A_head])
			inner_product_val += A.vals[A_head]*B.vals[B_head]*multiplicity_factor
			A_head += 1
			B_head += 1
		elseif A.indices[A_head,:] < B.indices[B_head,:]
			A_head += 1
		else
			B_head += 1
		end
	end
	return inner_product_val
end


text:
module Pathways

export assembly, thurber

include("assembly.jl")
include("thurber.jl")

end


text:
using RandomWalks, Distributions
using Test

include("mortal_test.jl")
include("orw_test.jl")
include("points_test.jl")
include("actors_test.jl")


text:
module CalibrationBase

Base.Experimental.@optlevel 3

export
    calibrate_image,
    calibrate_image!,
    calibrate_tomogram,
    calibrate_tomogram!

using SimpleTraits
using ..Monads

function calibration_data end
function calibrate_image! end
function calibrate_image end
function calibrate_tomogram! end
function calibrate_tomogram end

for nm ∈ (:calibrate_image, :calibrate_image!, :calibrate_tomogram, :calibrate_tomogram!)
    @eval begin
        $nm(; kwargs...) = x -> $nm(x; kwargs...)
        @traitfn $nm(m::M; kwargs...) where {M; Monad{M}} =
            mmap($nm(; kwargs...), m)
        @traitfn $nm(m::M, x; kwargs...) where {M; Monad{M}} =
            mmap($nm(x; kwargs...), m)
    end
end

end # module


text:
# Julia wrapper for header: /usr/include/scip/pub_bandit.h
# Automatically generated using Clang.jl wrap_c

function SCIPbanditSelect(bandit, action)
    ccall((:SCIPbanditSelect, libscip), SCIP_RETCODE, (Ptr{SCIP_BANDIT}, Ptr{Cint}), bandit, action)
end

function SCIPbanditUpdate(bandit, action, score)
    ccall((:SCIPbanditUpdate, libscip), SCIP_RETCODE, (Ptr{SCIP_BANDIT}, Cint, Cdouble), bandit, action, score)
end

function SCIPbanditvtableGetName(banditvtable)
    ccall((:SCIPbanditvtableGetName, libscip), Cstring, (Ptr{SCIP_BANDITVTABLE},), banditvtable)
end

function SCIPbanditGetRandnumgen()
    ccall((:SCIPbanditGetRandnumgen, libscip), Ptr{Cint}, ())
end

function SCIPbanditGetNActions(bandit)
    ccall((:SCIPbanditGetNActions, libscip), Cint, (Ptr{SCIP_BANDIT},), bandit)
end


text:
cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

#This script simulates the Fisher-KPP equation and fits
#a neural PDE to the data with the growth (aka reaction) term replaced
#by a feed-forward neural network and the diffusion term with a CNN

using PyPlot, Printf
using LinearAlgebra
using Flux, DiffEqFlux, Optim, DiffEqSensitivity
using BSON: @save, @load
using Flux: @epochs
using OrdinaryDiffEq

#parameter
D = 0.01; #diffusion
r = 1.0; #reaction rate

#domain
X = 1.0; T = 5;
dx = 0.04; dt = T/10;
x = collect(0:dx:X);
t = collect(0:dt:T);
Nx = Int64(X/dx+1);
Nt = Int64(T/dt+1);

#initial conditions
Amp = 1.0;
Delta = 0.2
#IC-1
rho0 = Amp*(tanh.((x .- (0.5 - Delta/2))/(Delta/10)) - tanh.((x .- (0.5 + Delta/2))/(Delta/10)))/2
#IC-2
#rho0 = Amp*(1 .- tanh.((x .- 0.2)/(Delta/6)))/2.

save_folder = "data"

if isdir(save_folder)
    rm(save_folder, recursive=true)
end
mkdir(save_folder)

close("all")
figure()
plot(x, rho0)
title("Initial Condition")
gcf()

########################
# Generate training data
########################
reaction(u) = r * u .* (1 .- u)
lap = diagm(0 => -2.0 * ones(Nx), 1=> ones(Nx-1), -1 => ones(Nx-1)) ./ dx^2
#Periodic BC
lap[1,end] = 1.0/dx^2
lap[end,1] = 1.0/dx^2
#Neumann BC
#lap[1,2] = 2.0/dx^2
#lap[end,end-1] = 2.0/dx^2

function rc_ode(rho, p, t)
    #finite difference
    D * lap * rho + reaction.(rho)
end

prob = ODEProblem(rc_ode, rho0, (0.0, T), saveat=dt)
sol = solve(prob, Tsit5());
ode_data = Array(sol);

figure(figsize=(8,3))

subplot(121)
pcolor(x,t,ode_data')
xlabel("x"); ylabel("t");
colorbar()

subplot(122)
for i in 1:2:Nt
    plot(x, ode_data[:,i], label="t=$(sol.t[i])")
end
xlabel("x"); ylabel(L"$\rho$")
legend(frameon=false, fontsize=7, bbox_to_anchor=(1, 1), loc="upper left", ncol=1)
tight_layout()
savefig(@sprintf("%s/training_data.pdf", save_folder))
gcf()

########################
# Define the neural PDE
########################
n_weights = 3

#for the reaction term
rx_nn = Chain(Dense(1, n_weights, tanh),
                Dense(n_weights, 1),
                x -> x[1])

#conv with bias with initial values as 1/dx^2
w_err = 0.0
init_w = reshape([1.1 -2.5 1.0], (3, 1, 1, 1))
diff_cnn_ = Conv(init_w, [0.], pad=(0,0,0,0))

#initialize D0 close to D/dx^2
D0 = [6.5]

p1,re1 = Flux.destructure(rx_nn)
p2,re2 = Flux.destructure(diff_cnn_)
p = [p1;p2;D0]
full_restructure(p) = re1(p[1:length(p1)]), re2(p[(length(p1)+1):end-1]), p[end]

function nn_ode(u,p,t)
    rx_nn = re1(p[1:length(p1)])

    u_cnn_1   = [p[end-4] * u[end] + p[end-3] * u[1] + p[end-2] * u[2]]
    u_cnn     = [p[end-4] * u[i-1] + p[end-3] * u[i] + p[end-2] * u[i+1] for i in 2:Nx-1]
    u_cnn_end = [p[end-4] * u[end-1] + p[end-3] * u[end] + p[end-2] * u[1]]

    # Equivalent using Flux, but slower!
    #CNN term with periodic BC
    #diff_cnn_ = Conv(reshape(p[(end-4):(end-2)],(3,1,1,1)), [0.0], pad=(0,0,0,0))
    #u_cnn = reshape(diff_cnn_(reshape(u, (Nx, 1, 1, 1))), (Nx-2,))
    #u_cnn_1 = reshape(diff_cnn_(reshape(vcat(u[end:end], u[1:1], u[2:2]), (3, 1, 1, 1))), (1,))
    #u_cnn_end = reshape(diff_cnn_(reshape(vcat(u[end-1:end-1], u[end:end], u[1:1]), (3, 1, 1, 1))), (1,))

    [rx_nn([u[i]])[1] for i in 1:Nx] + p[end] * vcat(u_cnn_1, u_cnn, u_cnn_end)
end

########################
# Soving the neural PDE and setting up loss function
########################
prob_nn = ODEProblem(nn_ode, rho0, (0.0, T), p)
sol_nn = concrete_solve(prob_nn,Tsit5(), rho0, p)

function predict_rd(θ)
  # No ReverseDiff if using Flux
  Array(concrete_solve(prob_nn,Tsit5(),rho0,θ,saveat=dt,sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP())))
end

#match data and force the weights of the CNN to add up to zero
function loss_rd(p)
    pred = predict_rd(p)
    sum(abs2, ode_data .- pred) + 10^2 * abs(sum(p[end-4 : end-2])), pred
end

########################
# Training
########################

#Optimizer
opt = ADAM(0.001)

global count = 0
global save_count = 0
save_freq = 50

train_arr = Float64[]
diff_arr = Float64[]
w1_arr = Float64[]
w2_arr = Float64[]
w3_arr = Float64[]

#callback function to observe training
cb = function (p,l,pred)
    rx_nn, diff_cnn_, D0 = full_restructure(p)
    push!(train_arr, l)
    push!(diff_arr, p[end])

    weight = diff_cnn_.weight[:]
    push!(w1_arr, weight[1])
    push!(w2_arr, weight[2])
    push!(w3_arr, weight[3])

    println(@sprintf("Loss: %0.4f\tD0: %0.4f Weights:(%0.4f,\t %0.4f, \t%0.4f) \t Sum: %0.4f"
            ,l, D0[1], weight[1], weight[2], weight[3], sum(weight)))

    global count

    if count==0
        fig = figure(figsize=(8,2.5));
        ttl = fig.suptitle(@sprintf("Epoch = %d", count), y=1.05)
        global ttl
        subplot(131)
        pcolormesh(x,t,ode_data')
        xlabel(L"$x$"); ylabel(L"$t$"); title("Data")
        colorbar()

        subplot(132)
        img = pcolormesh(x,t,pred')
        global img
        xlabel(L"$x$"); ylabel(L"$t$"); title("Prediction")
        colorbar(); clim([0, 1]);

        ax = subplot(133); global ax
        u = collect(0:0.01:1)
        rx_line = plot(u, rx_nn.([[elem] for elem in u]), label="NN")[1];
        global rx_line
        plot(u, reaction.(u), label="True")
        title("Reaction Term")
        legend(loc="upper right", frameon=false, fontsize=8);
        ylim([0, r*0.25+0.2])

        subplots_adjust(top=0.8)
        tight_layout()
    end

    if count>0
        println("updating figure")
        img.set_array(pred[1:end-1, 1:end-1][:])
        ttl.set_text(@sprintf("Epoch = %d", count))

        u = collect(0:0.01:1)
        rx_pred = rx_nn.([[elem] for elem in u])
        rx_line.set_ydata(rx_pred)
        u = collect(0:0.01:1)

        min_lim = min(minimum(rx_pred), minimum(reaction.(u)))-0.1
        max_lim = max(maximum(rx_pred), maximum(reaction.(u)))+0.1

        ax.set_ylim([min_lim, max_lim])
    end

    global save_count
    if count%save_freq == 0
        println("saved figure")
        savefig(@sprintf("%s/pred_%05d.png", save_folder, save_count), dpi=200, bbox_inches="tight")
        save_count += 1
    end

    display(gcf())
    count += 1

    l < 0.01 # Exit when fit to 2 decimal places
end

#train
@time begin
    res1 = DiffEqFlux.sciml_train(loss_rd, p, ADAM(0.001), cb=cb, maxiters = 100)
    res2 = DiffEqFlux.sciml_train(loss_rd, res1.minimizer, ADAM(0.001), cb=cb, maxiters = 300)
    res3 = DiffEqFlux.sciml_train(loss_rd, res2.minimizer, BFGS(), cb=cb, maxiters = 1000, allow_f_increases=true)
end

pstar = res3.minimizer

## Save trained model
@save @sprintf("%s/model.bson", save_folder) pstar

########################
# Plot for paper
########################
@load @sprintf("%s/model.bson", save_folder) pstar
#re-defintions for newly loaded data

diff_cnn_ = Conv(reshape(pstar[(end-4):(end-2)],(3,1,1,1)), [0.0], pad=(0,0,0,0))
diff_cnn(x) = diff_cnn_(x) .- diff_cnn_.bias
D0 = res3.minimizer[end]

fig = figure(figsize=(4,4))

rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams")
rcParams["font.size"] = 10
rcParams["text.usetex"] = true
rcParams["font.family"] = "serif"
rcParams["font.sans-serif"] = "Helvetica"
rcParams["axes.titlesize"] = 10

subplot(221)
pcolormesh(x,t,ode_data', rasterized=true)
xlabel(L"$x$"); ylabel(L"$t$"); title("Data")
yticks([0, 1, 2, 3, 4, 5])

ax = subplot(222)
cur_pred = predict_rd(pstar)[1]
img = pcolormesh(x,t,cur_pred', rasterized=true)
global img
xlabel(L"$x$"); ylabel(L"$t$"); title("Prediction")
yticks([0, 1, 2, 3, 4, 5])
cax = fig.add_axes([.48,.62,.02,.29])
colb = fig.colorbar(img, cax=cax)
colb.ax.set_title(L"$\rho$")
clim([0, 1]);
colb.set_ticks([0, 1])

subplot(223)
plot(Flux.data(w1_arr ./ w3_arr) .- 1, label=L"$w_1/w_3 - 1$")
plot(Flux.data(w1_arr .+ w2_arr .+ w3_arr), label=L"$w_1 + w_2 + w_3$")
axhline(0.0, linestyle="--", color="k")
xlabel("Epochs"); title("CNN Weights")
xticks([0, 1500, 3000]); yticks([-0.4, -0.3,-0.2, -0.1, 0.0, 0.1])
legend(loc="lower right", frameon=false, fontsize=6)

subplot(224)
u = collect(0:0.01:1)
plot(u, rx_nn.([[elem] for elem in u]), label="UPDE")[1];
plot(u, reaction.(u), linestyle="--", label="True")
xlabel(L"$\rho$")
title("Reaction Term")
legend(loc="lower center", frameon=false, fontsize=6);
ylim([0, 0.3])

tight_layout(h_pad=1)
gcf()
savefig(@sprintf("%s/fisher_kpp.pdf", save_folder))

#plot loss vs epochs and save
figure(figsize=(6,3))
plot(log.(train_arr), "k.", markersize=1)
xlabel("Epochs"); ylabel("Log(loss)")
tight_layout()
savefig(@sprintf("%s/loss_vs_epoch.pdf", save_folder))
gcf()

#=
# Success rate

# 2 decimal places 15 parameters

Loss: 0.0062    D0: 5.9522 Weights:(1.0385,      -2.0765,
    1.0380)          Sum: -0.0000                               1.0380)
updating figure
3430.385576 seconds (4.28 G allocations: 248.824 GiB, 1.31%
gc time)

Loss: 0.0095    D0: 6.1954 Weights:(0.9715,      -1.9432,
        0.9716)          Sum: -0.0000
updating figure
2824.449183 seconds (3.35 G allocations: 194.702 GiB, 1.32%
gc time)

Loss: 0.0094    D0: 5.7375 Weights:(1.0357,      -2.0714,
        1.0358)          Sum: -0.0000
updating figure
1174.592376 seconds (1.41 G allocations: 83.035 GiB, 1.22% gc time)

Loss: 0.0084    D0: 5.9525 Weights:(1.0049,      -2.0096,
        1.0047)          Sum: 0.0000
updating figure
1334.078451 seconds (1.61 G allocations: 94.637 GiB, 1.23% gc time)

Loss: 0.0075    D0: 7.4841 Weights:(0.8076,      -1.6145,       0.8069)          Sum: -0.0000
updating figure
saved figure
1053.729274 seconds (1.16 G allocations: 68.289 GiB, 1.12% gc time)

# 2 decimal places 7 parameters

Loss: 0.0095    D0: 6.1389 Weights:(0.9875,      -1.9749,
        0.9873)          Sum: 0.0000
updating figure
1415.089830 seconds (1.43 G allocations: 82.891 GiB, 1.37% gc time)

Loss: 0.0095    D0: 6.1381 Weights:(0.9447,      -1.8887,       0.9441)          Sum: -0.0000
updating figure
3293.038574 seconds (4.09 G allocations: 234.305 GiB, 1.10% gc time)

Loss: 0.0095    D0: 5.9216 Weights:(0.9869,      -1.9738,
        0.9869)          Sum: 0.0000
updating figure
3233.307375 seconds (4.09 G allocations: 234.477 GiB, 0.98%
gc time

Loss: 0.0095    D0: 5.9216 Weights:(0.9869,      -1.9738,
        0.9869)          Sum: 0.0000
updating figure
3265.894690 seconds (4.09 G allocations: 234.477 GiB, 0.98%
gc time)

Loss: 0.0093    D0: 6.4483 Weights:(0.9128,      -1.8256,       0.9128)         Sum: 0.0000
updating figure
1332.252349 seconds (1.44 G allocations: 83.367 GiB, 1.09% gc time)

# 2 decimal places 4 parameters

Loss: 0.4370    D0: 0.0457 Weights:(138.9983,    -277.9936,     138.9953)       Sum: 0.0000
updating figure
2210.916386 seconds (2.53 G allocations: 146.694 GiB, 1.07% gc time)

Loss: 0.3760    D0: 103.9210 Weights:(0.0533,    -0.1073,       0.0533)         Sum: -0.0007
updating figure
2296.853421 seconds (2.67 G allocations: 154.814 GiB, 1.08% gc time)

Loss: 0.3894    D0: 160.0180 Weights:(0.0367,    -0.0737,       0.0367)         Sum: -0.0003
updating figure
2262.799906 seconds (2.61 G allocations: 151.395 GiB, 1.08% gc time)

Loss: 0.3894    D0: 160.0180 Weights:(0.0367,    -0.0737,       0.0367)         Sum: -0.0003
updating figure
2346.022370 seconds (2.61 G allocations: 151.395 GiB, 1.09% gc time)

Loss: 0.2225    D0: 2739.4200 Weights:(0.0020,   -0.0040,       0.0020)         Sum: -0.0001
updating figure
5764.320007 seconds (6.49 G allocations: 372.446 GiB, 1.22% gc time)
=#

x = [1415.089830,3293.038574,3233.307375,3265.894690,1332.252349]
mean(x)
using Statistics
std(x)


text:
"""
Construct a `BaseURL` object given the base URL.
"""
BaseURL(base_url::AbstractString) = BaseURL(HTTP.URI(base_url))


text:
# Device Enumeration API

function enumDevices(enumfilter::ENUMFILTER = ENUMFILTERALL)
  x = Ref{Cint}(0)
  succ = ccall((:FDwfEnum,libdwf),Cint,(Cint,Ptr{Cint}),
    enumfilter,x)
  succ == 0 && error("Error calling FDwfEnum.")
  return x[]
end

function enumDeviceType(idxDevice::Int32)
  deviceId = Ref{Cint}(0)
  deviceRevision = Ref{Cint}(0)
  succ = ccall((:FDwfEnumDeviceType,libdwf),Cint,(Cint,Ptr{Cint},Ptr{Cint}),
    idxDevice,deviceId,deviceRevision)
  succ == 0 && error("Error calling FDwfEnumDeviceType.")
  return (DEVID(deviceId[]),DEVVER(deviceRevision[]))
end

function enumDeviceIsOpened(idxDevice::Int32)
  x = Ref{Cint}(0)
  succ = ccall((:FDwfEnumDeviceIsOpened,libdwf),Cint,(Cint,Ptr{Cint}),
    idxDevice,x)
  succ == 0 && error("Error calling FDwfEnumDeviceIsOpened.")
  return Bool(x[])
end

const dwfStringSize = 32

function enumUserName(idxDevice::Int32)
  userName = Vector{UInt8}(dwfStringSize)
  succ = ccall((:FDwfEnumUserName,libdwf),Cint,(Cint,Ptr{UInt8}),
    idxDevice,userName)
  succ == 0 && error("Error calling FDwfEnumUserName.")
  i = 0; found = false
  @inbounds while !found
    i += 1
    i > dwfStringSize && break
    found = (userName[i] == UInt8(0))
  end
  return String(copy(userName[1:i-1]))
end

function enumDeviceName(idxDevice::Int32)
  deviceName = Vector{UInt8}(dwfStringSize)
  succ = ccall((:FDwfEnumDeviceName,libdwf),Cint,(Cint,Ptr{UInt8}),
    idxDevice,deviceName)
  succ == 0 && error("Error calling FDwfEnumDeviceName.")
  i = 0; found = false
  @inbounds while !found
    i += 1
    i > dwfStringSize && break
    found = (deviceName[i] == UInt8(0))
  end
  return String(copy(deviceName[1:i-1]))
end

function enumSN(idxDevice::Int32)
  sn = Vector{UInt8}(dwfStringSize)
  succ = ccall((:FDwfEnumSN,libdwf),Cint,(Cint,Ptr{UInt8}),
    idxDevice,sn)
  succ == 0 && error("Error calling FDwfEnumSN.")
  i = 0; found = false
  @inbounds while !found
    i += 1
    i > dwfStringSize && break
    found = (sn[i] == UInt8(0))
  end
  return String(copy(sn[1:i-1]))
end

function enumConfig(idxDevice::Int32)
  x = Ref{Cint}(0)
  succ = ccall((:FDwfEnumConfig,libdwf),Cint,(Cint,Ptr{Cint}),
    idxDevice,x)
  succ == 0 && error("Error calling FDwfEnumConfig.")
  return x[]
end

function enumConfigInfo(idxConfig::Int32,infoType::DWFENUMCONFIGINFO)
  x = Ref{Cint}(0)
  succ = ccall((:FDwfEnumConfigInfo,libdwf),Cint,(Cint,Cint,Ptr{Cint}),
    idxConfig,infoType,x)
  succ == 0 && error("Error calling FDwfEnumConfigInfo.")
  return x[]
end


text:
@testset "Solver Methods" begin

    # Test Sover Methods
    # Test solver on a linear unconstrained problem with one player
    # Should be solved in one iteration
    T = Float64
    N = 20
    dt = 0.1
    p = 1
    model = DoubleIntegratorGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 1.0, 0.0, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 1
    opts.inner_iter = 1
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    newton_solve!(prob)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 1e-6
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

    # Test solver on a non linear unconstrained problem with one player
    T = Float64
    N = 20
    dt = 0.1
    p = 1
    model = UnicycleGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 1.0, 0.0, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 7
    opts.inner_iter = 20
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    newton_solve!(prob)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 1e-6
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

    # Test solver on a linear unconstrained problem with 2 players
    T = Float64
    N = 20
    dt = 0.1
    p = 2
    model = DoubleIntegratorGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 2.0, 1.0, 2.0, 0.0, 0.0, 0.9, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 1
    opts.inner_iter = 1
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    newton_solve!(prob)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 1e-6
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

    # Test solver on a non linear unconstrained problem with 2 players
    T = Float64
    N = 20
    dt = 0.1
    p = 2
    model = UnicycleGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 2.0, 1.0, 2.0, 0.0, 0.0, 0.9, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 7
    opts.inner_iter = 20
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    newton_solve!(prob)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 1e-6
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

    # Test solver on a non linear constrained problem with 2 players
    T = Float64
    N = 20
    dt = 0.1
    p = 2
    model = UnicycleGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 2.0, 1.1, 2.0, 0.0, 0.0, 0.9, 0.9])

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    # Define the constraints that each player must respect
    game_con = GameConstraintValues(probsize)
    # Add collision avoidance
    radius = 0.05
    add_collision_avoidance!(game_con, radius)
    # Add control bounds
    u_max =  1*ones(SVector{model.m,T})
    u_min = -1*ones(SVector{model.m,T})
    add_control_bound!(game_con, u_max, u_min)
    # Add circle constraint
    xc = [1.50, 0.2, 0.3]
    yc = [1.25, 0.2, 0.3]
    radius = [0.2, 0.2, 0.3]
    add_circle_constraint!(game_con, xc, yc, radius)

    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts = Options()
    opts.inner_print = false
    opts.outer_print = false
    opts.ls_iter = 15
    opts.outer_iter = 20
    opts.inner_iter = 20
    opts.reg_0 = 1e-8
    opts.α_dual = 1.0
    opts.λ_max = 1e7
    opts.ϵ_dyn = 1e-3
    opts.ϵ_sta = 1e-3
    opts.ϵ_con = 1e-3
    opts.ϵ_opt = 1e-3
    newton_solve!(prob)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 1e-3
    @test prob.stats.dyn_vio[end].max < 1e-3
    @test prob.stats.sta_vio[end].max < 1e-3
    @test prob.stats.con_vio[end].max < 1e-3
    @test prob.stats.opt_vio[end].max < 1e-3

end

@testset "IBR Solver Methods" begin

    # Test Sover Methods
    # Test solver on a linear unconstrained problem with one player
    # Should be solved in one iteration
    T = Float64
    N = 20
    dt = 0.1
    p = 1
    model = DoubleIntegratorGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 1.0, 0.0, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 1
    opts.inner_iter = 1
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    ibr_newton_solve!(prob, 1)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 1e-6
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

    # Test solver on a non linear unconstrained problem with one player
    T = Float64
    N = 20
    dt = 0.1
    p = 1
    model = UnicycleGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 1.0, 0.0, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 7
    opts.inner_iter = 20
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    ibr_newton_solve!(prob, 1)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 1e-6
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

    # Test solver on a linear unconstrained problem with 2 players
    T = Float64
    N = 20
    dt = 0.1
    p = 2
    model = DoubleIntegratorGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 2.0, 1.0, 2.0, 0.0, 0.0, 0.9, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 1
    opts.inner_iter = 1
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    ibr_newton_solve!(prob)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 5e-2
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

    # Test solver on a non linear unconstrained problem with 2 players
    T = Float64
    N = 20
    dt = 0.1
    p = 2
    model = UnicycleGame(p=p)
    probsize = ProblemSize(N,model)
    x0 = SVector{model.n,T}([1.0, 2.0, 1.0, 2.0, 0.0, 0.0, 0.9, 0.9])
    opts = Options()
    opts.inner_print = false
    opts.outer_print = false

    Q = [Diagonal(1*ones(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(0.5*ones(SVector{model.mi[i],T})) for i=1:p]
    xf = [0*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [-1*ones(SVector{model.mi[i],T}) for i=1:p]
    game_obj = GameObjective(Q,R,xf,uf,N,model)

    game_con = GameConstraintValues(probsize)
    prob = GameProblem(N, dt, x0, model, opts, game_obj, game_con)

    opts.outer_iter = 7
    opts.inner_iter = 20
    opts.ls_iter = 25
    opts.reg_0 = 1e-7
    opts.ϵ_dyn = 1e-10
    opts.ϵ_opt = 1e-10
    ibr_newton_solve!(prob)
    @test norm(prob.core.res, 1)/length(prob.core.res) < 5e-2
    @test dynamics_violation(model, prob.pdtraj).max < 1e-6

end


text:
using EPT
using Turing
using Test
using Random

import AnnealedIS

# rng = MersenneTwister(42)
Random.seed!(42)

@testset "EPT.jl" begin
    @testset "Expectation Macro" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x^2
        end

        yval = 1
        expct_conditioned = expct(yval)

        xval = 2
        vi = Turing.VarInfo(expct_conditioned.gamma1_pos)
        vi[@varname(x)] = [xval;]
        # vi[@varname(y)] = [yval;]

        # Check that the three different models return the right score for a given trace.
        gamma2_lp = logpdf(Normal(0, 1), xval) + logpdf(Normal(xval, 1), yval)
        @test expct_conditioned.gamma2(vi) == xval^2
        @test Turing.getlogp(vi) == gamma2_lp

        gamma1_pos_lp = gamma2_lp + log(max(xval^2, 0))
        @test expct_conditioned.gamma1_pos(vi) == xval^2
        @test Turing.getlogp(vi) == gamma1_pos_lp

        gamma1_neg_lp = gamma2_lp + log(-min(xval^2, 0))
        @test expct_conditioned.gamma1_neg(vi) == xval^2 # TODO: Should it return 0 instead?
        @test Turing.getlogp(vi) == gamma1_neg_lp
    end

    @testset "Correct scoping" begin
        # Checks that we can access functions that are in the current scope
        # inside the model body (here f).
        f(x) = x^2
        @expectation function expct()
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return f(x)
        end

        fx = expct.gamma1_pos()()
        @test isa(fx, Float64)
    end

    @testset "Expectation Estimation" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x^2
        end

        yval = 3
        expct_conditioned = expct(yval)

        num_annealing_dists = 10
        num_samples = 10

        tabi = TABI(
            AIS(num_annealing_dists, num_samples, SimpleRejection())
        )

        expct_estimate, diagnostics = estimate_expectation(expct_conditioned, tabi)
        @test !isnan(expct_estimate)
    end

    # Comment this out because it takes a while.
    # @testset "Convergence test" begin
    #    @expectation function expct(y)
    #        x ~ Normal(0, 1)
    #        y ~ Normal(x, 1)
    #        return x
    #    end
    #
    #    yval = 3
    #    expct_conditioned = expct(yval)

    #    num_annealing_dists = 100
    #    num_samples = 1000

    #    tabi = TABI(
    #        AIS(num_samples, num_annealing_dists)
    #    )

    #    expct_estimate, diagnostics = estimate_expectation(expct_conditioned, tabi)
    #    @test_broken isapprox(expct_estimate, 1.5, atol=1e-2)
    # end

    @testset "Diagnostics" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x
        end

        yval = 3
        expct_conditioned = expct(yval)

        num_annealing_dists = 10
        num_samples = 10

        tabi = TABI(
            AIS(num_samples, num_annealing_dists, SimpleRejection())
        )

        expct_estimate, diagnostics = estimate_expectation(
            expct_conditioned,
            tabi;
            store_intermediate_samples=true
        )

        keys = [:Z2_info, :Z1_negative_info, :Z1_positive_info]

        for k in keys
            @test haskey(diagnostics, k)

            @test typeof(diagnostics[k][:ess]) == Float64
            @test typeof(diagnostics[k][:Z_estimate]) == Float64
            @test size(diagnostics[k][:samples]) == (num_samples,)
            @test haskey(diagnostics[k], :intermediate_samples)
        end
    end

    @testset "Rejection Samplers" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x
        end

        yval = 3
        expct_conditioned = expct(yval)

        num_annealing_dists = 10
        num_samples = 10

        tabi_no_rejection = TABI(
            AIS(num_samples, num_annealing_dists, SimpleRejection())
        )
        _, _ =  estimate_expectation(
            expct_conditioned,
            tabi_no_rejection;
            store_intermediate_samples=true
        )

        tabi_rejection = TABI(
            AIS(num_samples, num_annealing_dists, RejectionResample())
        )
        _, _ =  estimate_expectation(
            expct_conditioned,
            tabi_rejection;
            store_intermediate_samples=true
        )
    end

    @testset "Disable Z1_pos or Z1_neg" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x^2
        end

        yval = 3
        expct_conditioned = expct(yval)

        num_annealing_dists = 10
        num_samples = 2

        tabi_no_Z1_neg = TABI(
            AIS(num_samples, num_annealing_dists, SimpleRejection()),
            AIS(0, num_annealing_dists, SimpleRejection()),
            AIS(num_samples, num_annealing_dists, SimpleRejection())
        )

        _, d =  estimate_expectation(
            expct_conditioned,
            tabi_no_Z1_neg;
            store_intermediate_samples=true
        )

        full_tabi = TABI(AIS(num_samples, num_annealing_dists, SimpleRejection()))
        _, d_full =  estimate_expectation(
            expct_conditioned,
            tabi_no_Z1_neg;
            store_intermediate_samples=true
        )

        # Check that estimate_expectation is type-stable.
        @test typeof(d_full) == typeof(d)
    end

    @testset "Turing Importance Sampling" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x
        end

        yval = 3
        expct_conditioned = expct(yval)

        num_samples = 10

        tabi = TABI(
            TuringAlgorithm(IS(), num_samples)
        )

        expct_estimate, diag = estimate_expectation(
            expct_conditioned,
            tabi;
            progress=false
        )

        @test typeof(expct_estimate) == Float64
        for key in [:Z1_positive_info, :Z1_negative_info, :Z2_info]
            @test typeof(diag[key]) <: MCMCChains.Chains
        end
    end

    @testset "Prior extraction" begin
        @expectation function expct(y)
        x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x
        end

        yval = 3
            expct_conditioned = expct(yval)

        log_prior = AnnealedIS.make_log_prior_density(
            expct_conditioned.gamma1_pos
        )

        xval = 0.0
        true_prior = logpdf(Normal(0, 1), xval)
        @test log_prior((x = xval,)) == true_prior
    end

    @testset "Turing AnIS" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x
        end

        yval = 3
        expct_conditioned = expct(yval)

        num_samples = 10
        num_annealing_dists = 10

        tabi = TABI(
            TuringAlgorithm(AnnealedIS.AnIS(num_annealing_dists), num_samples)
        )

        expct_estimate, diag = estimate_expectation(expct_conditioned, tabi)

        @test typeof(expct_estimate) == Float64
        for key in [:Z1_positive_info, :Z1_negative_info, :Z2_info]
            @test typeof(diag[key]) <: MCMCChains.Chains
        end
    end

    @testset "Multiple Expectations" begin
        @expectation function expct(y)
            x ~ Normal(0, 1)
            y ~ Normal(x, 1)
            return x, x^2, x^3
        end

        @test isa(expct, Array{EPT.Expectation})
        @test length(expct) == 3
    end
end