text:
"""
Module with tests for BLTandPantelides.

* Author: Hilding Elmqvist, Mogram AB
* Date: July-August 2016. Array version: Summer 2018
* License: MIT

"""
module testBLTandPantelides

using Modia.BLTandPantelides
using Modia.BLTandPantelidesUtilities
@static if VERSION < v"0.7.0-DEV.2005"
  using Base.Test
else
  using Test
end

@testset "BLTandPantelides" begin

# testMatch
  println("\nTest match")
  G = Any[
    [2],
    [3,8],
    [4,7],
    [5],
    [3,6],
    [],
     [4,6],
    [1,7],
  ]

  assign = matching(G, 8)

  @show assign
  @testset "testMatch" begin
    @test any(assign .== 6) == false
    @test assign == [8,1,2,7,4,5,3,0]
  end

# testSingular
  println("\nSingular system")
  G = Any[
    [3],
    [3],
    [2,3]
  ]

  assign = matching(G, 4)
  @show assign

  (invAssign, unAssignedVariables) = invertAssign(assign, length(G))
  @show invAssign, unAssignedVariables

  (ass, unAssignedEquations) = invertAssign(invAssign, length(assign))
  @show ass, unAssignedEquations

    @testset "Singular system" begin
    @test assign == [0,3,1,0]
    @test invAssign == [3,0,2]
    @test unAssignedVariables == [1,4]
    @test ass == assign
    @test unAssignedEquations == [2]
  end

# testTarjan
  println("\nTest Tarjans strong connect")
  # Reference: Tarjan, Fig.3., page 158.
  G = Any[
    [2],
    [3,8],
    [4,7],
    [5],
    [3,6],
    [],
     [4,6],
    [1,7],
  ]

  assign = 1:8

  components = BLT(G, assign)
  @show components
  @testset "testTarjan" begin
    @test components == Any[Any[6],Any[7,5,4,3],Any[8,2,1]]
  end

# testPendulum
  println("\nFixed-length pendulum")
  # Reference: Pantelides, page 222-225.
  G = Any[
    [3, 5],
    [4, 6],
    [1, 7, 9],
    [2, 8, 9],
    [1, 2]
  ]

  assign = matching(G, 9)
  @show assign

  # Variable vector is: [x, y, w, z, der(x), der(y), der(w), der(z), T], i.e. vector A is:
  A = [5,6,7,8,0,0,0,0,0]

  # Symbolic output of results

  variables = ["x", "y", "w", "z", "der(x)", "der(y)", "der(w)", "der(z)", "T"]
  equations = [
    "der(x) = w",
    "der(y) = z",
    "der(w) = T*x",
    "der(z) = T*y - g",
    "0 = x^2 + y^2 - L^2"]

  println("\nAssigned original equations:")
  printAssignedEquations(equations, variables, 1:length(equations), assign, A, fill(0, length(equations)))
  printUnassigned(equations, variables, assign, A, fill(0, length(equations)))

  println("\nTest diagnostics for too many equations")
  tooManyEquations = [equations; [
    "x = cos(phi)",
    "y = sin(phi)"]]
  tooFewVariables = [variables; ["phi", "dummy"]]
#=
  Gbig = {G; {
    [1, 10]
    [2, 10]
    }
  }
=#
  Gbig = copy(G)
  push!(Gbig, [1, 10])
  push!(Gbig, [2, 10])
  @show Gbig
  Abig = [A; fill(0, 1) ]
  EGbig = [Gbig; buildExtendedSystem(Abig)]
  EGbig = addDependencies(EGbig, (length(Abig)+1):length(EGbig))
  @show EGbig
  equationsBig = [tooManyEquations; fill("h(., der(.)) = 0", length(EGbig)-length(tooManyEquations))]
  assignBig = matching(EGbig, length(EGbig))
  Abig = [Abig; fill(0, length(EGbig)-length(Abig))]
   printAssignedEquations(equationsBig, tooFewVariables, 1:length(EGbig), assignBig, Abig, fill(0, length(EGbig)))
  printUnassigned(equationsBig, tooFewVariables, assignBig, Abig, fill(0, length(EGbig)))
  componentsBig = BLT(EGbig, assignBig)
  @show componentsBig

  println("\nTest diagnostics for too many variables")
  tooManyVariables = [variables; ["dummy"]]
  Gbig = copy(G)
  Gbig[5] = [1, 10] # Wrong equation involving x and dummy
  Gbig[4] = [10, 8, 9] # Wrong equation involving dummy
  @show Gbig
  Abig = [A; fill(0, 1) ]
  EGbig = [Gbig; buildExtendedSystem(Abig)]
  EGbig = [EGbig; buildFullIncidence(length(Abig)-length(EGbig), length(Abig))]
  @show EGbig
  equationsBig = copy(equations)
  equationsBig[5] = "0 = x^2 + dummy^2 - L^2"
  equationsBig = [equationsBig; fill("h(., der(.)) = 0", length(EGbig)-length(equationsBig))]
  assignBig = matching(EGbig, length(EGbig))
  Abig = [Abig; fill(0, length(EGbig)-length(Abig))]
   printAssignedEquations(equationsBig, tooManyVariables, 1:length(EGbig), assignBig, Abig, fill(0, length(EGbig)))
  printUnassigned(equationsBig, tooManyVariables, assignBig, Abig, fill(0, length(EGbig)))
  componentsBig = BLT(EGbig, assignBig)
  @show componentsBig

  println("\nTest diagnostics for too few equations")
  tooManyVariables = variables
  Gbig = copy(G)
  pop!(Gbig)  # Delete last equation
  @show Gbig
  Abig = A
  EGbig = [Gbig; buildExtendedSystem(Abig)]
  equationsBig = equations[1:4]
  equationsBig = [equationsBig; fill("h(., der(.)) = 0", length(EGbig)-length(equationsBig))]
  EGbig = [EGbig; buildFullIncidence(length(Abig)-length(EGbig), length(Abig))]
  equationsBig = [equationsBig; fill("full(...) = 0", length(EGbig)-length(equationsBig))]
  @show EGbig
  assignBig = matching(EGbig, length(EGbig))
  Abig = [Abig; fill(0, length(EGbig)-length(Abig))]
   printAssignedEquations(equationsBig, tooManyVariables, 1:length(EGbig), assignBig, Abig, fill(0, length(EGbig)))
  printUnassigned(equationsBig, tooManyVariables, assignBig, Abig, fill(0, length(EGbig)))
  componentsBig = BLT(EGbig, assignBig)
  @show componentsBig

  println("\nCheck consistency of equations by matching extended equation set")
  EG = [G; buildExtendedSystem(A)]
  @show EG
  assign = matching(EG, 9)
  @show assign
  @testset "testPendulum 1" begin
    @test all(assign .> 0)
  end

  println("\nPerform index reduction")
  (assign, A, B) = pantelides!(G, 9, A)
  @show G

  @show assign
  @show A
  @show B
  @testset "testPendulum 2" begin
  # According to Pantelides, page 224-225.
    @test assign == [0,0,0,0,1,2,7,4,3,9,8]
    @test A == [5,6,7,8,10,11,0,0,0,0,0]
    @test B == [7,8,0,0,6,9,0,0,0]
  end

  println("------------------------------------------------------")
  println()
  vActive = fill(true, length(A))
  vActive[[1,3]] = false
  @show vActive
  assign = matching(G, length(A), vActive)
  @show assign

  components = BLT(G, assign)
  @show components
  @testset "testPendulum BLT components" begin
    @test assign == [0, 5, 0, 2, 1, 6, 7, 4, 3, 9, 8]
    @test components == Any[Any[1], Any[5], Any[6], Any[2], Any[4, 8, 9, 7, 3]]
  end

  println("------------------------------------------------------")
  println()

  println("\nAll unknowns:")
  printList(variables, 1:length(A), A)
  println("\nAll equations:")
  printList(equations, 1:length(B), B, true)
  println("\nAssigned equations:")
  printAssignedEquations(equations, variables, 1:length(B), assign, A, B)
  println("\nSorted equations:")
  printSortedEquations(equations, variables, components, assign, A, B)
  printUnassigned(equations, variables, assign, A, B)

  println("\nBuild augmented system.")
  AG = [G; buildFullIncidence(length(A)-length(B), length(A))]
  @show AG
  assignAG = matching(AG, length(A))
  @show assignAG
  componentsAG = BLT(AG, assignAG)
  @show componentsAG
  @testset "testPendulum 3" begin
  # See Pantelides, page 222, paragraph 1.
    @test componentsAG == Any[Any[11,3,7,9,8,2,10,4,5,6,1]]
  end

  equationsAG = [equations; fill("", length(B)-length(equations)); fill("full", length(A)-length(B))]
  BAG = [B;fill(0, length(A)-length(B))]
  println("\nAssigned augmented equations:")
  printAssignedEquations(equationsAG, variables, 1:length(BAG), assignAG, A, BAG)
  println("\nSorted augmented equations:")
  printSortedEquations(equationsAG, variables, componentsAG, assignAG, A, BAG)
  printUnassigned(equationsAG, variables, assignAG, A, BAG)

  println("\nSet initial conditions on x and y. Should fail.")
  IG1 = copy(G)
  push!(IG1, [1])
  push!(IG1, [2])
  @show IG1
  assignIG1 = matching(IG1, length(A))
  @show assignIG1
  @testset "testPendulum 4" begin
    @test any(assignIG1 .== 0)
  end
  componentsIG1 = BLT(IG1, assignIG1)
  @show componentsIG1

  equationsIG = [equations; fill("", length(B)-length(equations)); fill("initial", length(A)-length(B))]
  BIG = [B;fill(0, length(A)-length(B))]

  printUnassigned(equationsIG, variables, assignIG1, A, BIG)

  println("\nSet initial conditions on x and w.")
  IG2 = copy(G)
  push!(IG2, [1])
  push!(IG2, [3])
  @show IG2
  assignIG2 = matching(IG2, length(A))
  @show assignIG2
  @testset "testPendulum 5" begin
    @test all(assignIG2 .> 0)
  end
  componentsIG2 = BLT(IG2, assignIG2)
  @show componentsIG2

  println("\nSorted IG2 equations:")
  printSortedEquations(equationsIG, variables, componentsIG2, assignIG2, A, BIG)

  println("\nSet initial conditions on w and z.")
  IG3 = copy(G)
  push!(IG3, [3])
  push!(IG3, [4])
  @show IG3
  assignIG3 = matching(IG3, length(A))
  @show assignIG3
  @testset "testPendulum 6" begin
    @test all(assignIG3 .> 0)
  end
  componentsIG3 = BLT(IG3, assignIG3)
  @show componentsIG3

  println("\nSorted IG3 equations:")
  printSortedEquations(equationsIG, variables, componentsIG3, assignIG3, A, BIG)

# testPendulum1
  println("\nFixed-length pendulum")
  # Reference: Pantelides, page 222-225.
  G = Any[
    [3, 5],
    [4, 6],
    [1, 7, 9],
    [2, 8, 9],
    [1, 2]
  ]

  # Variable vector is: [x, y, w, z, der(x), der(y), der(w), der(z), T], i.e. vector A is:
  A = [5,6,7,8,0,0,0,0,0]

  println("\nPerform index reduction")
  (assign, A, B) = pantelides!(G, 9, A)
  @testset "testPendulum 2" begin
  # According to Pantelides, page 224-225.
    @test assign == [0,0,0,0,1,2,7,4,3,9,8]
    @test A == [5,6,7,8,10,11,0,0,0,0,0]
    @test B == [7,8,0,0,6,9,0,0,0]
  end

  println("\nSet initial conditions on x and w.")
  IG = copy(G)
  push!(IG, [1])
  push!(IG, [3])
  @show IG
  assignIG = matching(IG, length(A))
  @show assignIG
  componentsIG = BLT(IG, assignIG)
  @show componentsIG
  @testset "testPendulum" begin
    @test all(assignIG .> 0)
    @test componentsIG == Any[Any[11], Any[1], Any[10], Any[5], Any[6], Any[2], Any[7,9,8,4,3]]
  end

# testReactor
  # Reference: Pantelides, page 225-228.
  println("\nExothermic Reactor Model")
  G = Any[
    [1, 3, 5],
    [2, 4, 5, 6],
    [1, 2, 5],
    [1]
  ]

  matching(G, 6)

  A = [3,4,0,0,0,0]

  (assign, A, B) = pantelides!(G, 6, A)
  @show assign
  @show A
  @show B
  components = BLT(G, assign)
  @show components

  @testset "testReactor" begin
    @test assign == [0,0,1,7,3,2,8,6]
    @test components == Any[Any[3],Any[1],Any[8],Any[6],Any[7],Any[2],Any[4],Any[5]]
  end

  println("\n\n----------------------\n")
end

println("\n----------------------\n")

function bigTest(G)
  n = length(G)
  if n <= 100
    @show G
  end
  assign = matching(G, n)
  if n <= 100
    @show assign
  end
  components = BLT(G, assign)
  if n <= 100
    @show components
  end
end

const n=10000  # Stack overflow for band and n=100000
const nFull=1000

function test()
  println("\nBig tests, n = ", n)
  println("\nBig test: diagonal")
  @static if VERSION < v"0.7.0-DEV.2005"
    G1 = Array{Any}(n)
  else
    G1 = Array{Any}(undef, n)
  end
  for i in 1:n
    G1[i] = [i]
  end
  @time bigTest(G1)

  println("\nBig test: band")
  @static if VERSION < v"0.7.0-DEV.2005"
    G2 = Array{Any}(n)
  else
    G2 = Array{Any}(undef, n)
  end
  for i in 1:n
    G2[i] = [i-1 < 1 ? n : i-1, i, mod(i,n)+1]
  end
  @time bigTest(G2)

  println("\nBig test: full, n=", nFull)

   @static if VERSION < v"0.7.0-DEV.2005"
    G3 = Array{Any}(nFull)
  else
    G3 = Array{Any}(undef, nFull)
  end
  for i in 1:nFull
    G3[i] = [i for i in 1:nFull]
  end
  @time bigTest(G3)

end

test()

end


text:
using DSGE, Test, ModelConstructors

# Initialize model object
m = Model990()
m <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
m <= Setting(:use_population_forecast, true)
m990 = Model990(testing = true)
m990 <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
m990 <= Setting(:use_population_forecast, true)

# Forecast under historical rule
paras = [α.value for α in m.parameters]
df = load_data(m990, verbose = :none)

output_vars = [:histpseudo, :forecastpseudo, :irfpseudo]
sys_hist = compute_system(m)
out_hist = DSGE.forecast_one_draw(m, :mode, :none, output_vars, paras, df, verbose = :none)

# Getting indices
eq_mp    = m.equilibrium_conditions[:eq_mp]
eq_other = setdiff(1:n_equilibrium_conditions(m), eq_mp)

# Assign alternative policy
setup_permanent_altpol!(m, DSGE.taylor93())

# Compare equilibrium conditions under historical and alternative rules
Γ0_hist, Γ1_hist, ~ = eqcond(m)
Γ0_alt93,  Γ1_alt93,  ~ = get_setting(m, :regime_eqcond_info)[2].alternative_policy.eqcond(m)

# Repeat for Taylor99
setup_permanent_altpol!(m, DSGE.taylor99())
Γ0_alt99,  Γ1_alt99,  ~ = get_setting(m, :regime_eqcond_info)[2].alternative_policy.eqcond(m)

## First, check if coefficients in matrices correct
### Taylor 93
endo = m.endogenous_states
@testset "Check Taylor93 rule has the right coefficients." begin
    @test Γ0_alt93[eq_mp, endo[:R_t]] == 1
    @test Γ0_alt93[eq_mp, endo[:π_a_t]] == -1.5/4
    @test Γ0_alt93[eq_mp, endo[:y_t]]  == -0.5/4
    @test Γ0_alt93[eq_mp, endo[:y_f_t]] == 0.5/4
end

### Taylor 99
@testset "Check Taylor99 rule has the right coefficients." begin
    @test Γ0_alt99[eq_mp, endo[:R_t]] == 1
    @test Γ0_alt99[eq_mp, endo[:π_a_t]] == -1.5/4
    @test Γ0_alt99[eq_mp, endo[:y_t]]  == -1/4
    @test Γ0_alt99[eq_mp, endo[:y_f_t]] == 1/4
end

# Testing
@testset "Compare AltPolicy eqconds under hist and alt rules" begin
    @test Γ0_hist[eq_other, :] ≈ Γ0_alt93[eq_other, :]
    @test Γ1_hist[eq_other, :] ≈ Γ1_alt93[eq_other, :]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt93[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt93[eq_mp, :])

    @test Γ0_hist[eq_other, :] ≈ Γ0_alt99[eq_other, :]
    @test Γ1_hist[eq_other, :] ≈ Γ1_alt99[eq_other, :]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt99[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt99[eq_mp, :])
end

# Forecast under Taylor 93 and Taylor 99
setup_permanent_altpol!(m, DSGE.taylor93())
out_alt93 = DSGE.forecast_one_draw(m, :mode, :none, output_vars, paras, df, verbose = :none,
                                   regime_switching = true, n_regimes = get_setting(m, :n_regimes))
setup_permanent_altpol!(m, DSGE.taylor99())
out_alt99 = DSGE.forecast_one_draw(m, :mode, :none, output_vars, paras, df, verbose = :none,
                                   regime_switching = true, n_regimes = get_setting(m, :n_regimes))
m <= Setting(:date_forecast_start, DSGE.iterate_quarters(date_forecast_start(m), -1))
m <= Setting(:date_conditional_end, date_forecast_start(m))
setup_permanent_altpol!(m, DSGE.taylor99(); cond_type = :full)
out_alt99_cond = DSGE.forecast_one_draw(m, :mode, :full, output_vars, paras, df, verbose = :none,
                                        regime_switching = true, n_regimes = get_setting(m, :n_regimes))
m <= Setting(:date_forecast_start, DSGE.iterate_quarters(date_forecast_start(m), 1))
m <= Setting(:date_conditional_end, date_forecast_start(m))
@testset "Check forecast under Taylor 93" begin
    @test out_hist[:histpseudo] ≈ out_alt93[:histpseudo]
    @test !(out_hist[:forecastpseudo] ≈ out_alt93[:forecastpseudo])
end
@testset "Check forecast under Taylor 99" begin
    @test out_hist[:histpseudo] ≈ out_alt99[:histpseudo]
    @test !(out_hist[:forecastpseudo] ≈ out_alt99[:forecastpseudo])
end

####################
# New Policy Rules #
####################

m = Model1002("ss10"; custom_settings = Dict{Symbol, Setting}(:add_altpolicy_pgap => Setting(:add_altpolicy_pgap, true)))
my = Model1002("ss10"; custom_settings = Dict{Symbol, Setting}(:add_altpolicy_pgap => Setting(:add_altpolicy_pgap, true),
                                                               :add_altpolicy_ygap => Setting(:add_altpolicy_ygap, true)))
m <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
m <= Setting(:use_population_forecast, true)
my <= Setting(:date_forecast_start, quartertodate("2015-Q4"))
my <= Setting(:use_population_forecast, true)

# Getting indices
eq_mp    = m.equilibrium_conditions[:eq_mp]
eq_pgap  = m.equilibrium_conditions[:eq_pgap]
eq_other = setdiff(1:n_equilibrium_conditions(m), [eq_mp, eq_pgap])
eq_mpy    = my.equilibrium_conditions[:eq_mp]
eq_pgapy  = my.equilibrium_conditions[:eq_pgap]
eq_ygapy  = my.equilibrium_conditions[:eq_ygap]
eq_othery = setdiff(1:n_equilibrium_conditions(my), [eq_mpy, eq_pgapy, eq_ygapy])
eq_otherymp = 1:n_equilibrium_conditions(my)

# Compare equilibrium conditions under historical and alternative rules
Γ0_hist, Γ1_hist, ~ = eqcond(m)
Γ0_histy, Γ1_histy, ~ = eqcond(my)

# Matrices for alt_inflation
setup_permanent_altpol!(m, DSGE.alt_inflation())
Γ0_alt_inf,  Γ1_alt_inf,  ~ = get_setting(m, :regime_eqcond_info)[2].alternative_policy.eqcond(m)

# Repeat for zero_rate
setup_permanent_altpol!(m, DSGE.zero_rate())
Γ0_alt_zero,  Γ1_alt_zero,  ~ = get_setting(m, :regime_eqcond_info)[2].alternative_policy.eqcond(m)

# Repeat for smooth_ait_gdp
setup_permanent_altpol!(my, DSGE.smooth_ait_gdp())
Γ0_alt_smooth,  Γ1_alt_smooth,  ~ = get_setting(my, :regime_eqcond_info)[2].alternative_policy.eqcond(my)

# Repeat for smooth_ait_gdp_alt
setup_permanent_altpol!(my, DSGE.smooth_ait_gdp_alt())
Γ0_alt_smooth_alt,  Γ1_alt_smooth_alt,  ~ = get_setting(my, :regime_eqcond_info)[2].alternative_policy.eqcond(my)

# Repeat for flexible_ait
setup_permanent_altpol!(my, DSGE.flexible_ait())
Γ0_alt_flex_ait,  Γ1_alt_flex_ait,  ~ = get_setting(my, :regime_eqcond_info)[2].alternative_policy.eqcond(my)

# Repeat for ait
setup_permanent_altpol!(my, DSGE.ait())
my <= Setting(:pgap_type, :ait)
my <= Setting(:pgap_value, 0.)
Γ0_alt_ait,  Γ1_alt_ait,  ~ = get_setting(my, :regime_eqcond_info)[2].alternative_policy.eqcond(my)

# Repeat for ngdp
setup_permanent_altpol!(my, DSGE.ait())
my <= Setting(:pgap_type, :ngdp)
my <= Setting(:pgap_value, 12.)
Γ0_alt_ngdp,  Γ1_alt_ngdp,  ~ = get_setting(my, :regime_eqcond_info)[2].alternative_policy.eqcond(my)

# Repeat for rw
setup_permanent_altpol!(my, DSGE.rw())
Γ0_alt_rw,  Γ1_alt_rw,  ~ = get_setting(my, :regime_eqcond_info)[2].alternative_policy.eqcond(my)

# Repeat for rw_zero_rate
setup_permanent_altpol!(my, DSGE.rw_zero_rate())
Γ0_alt_rw_zero,  Γ1_alt_rw_zero,  ~ = get_setting(my, :regime_eqcond_info)[2].alternative_policy.eqcond(my)

# Testing
@testset "Compare AltPolicy eqconds under hist and alt rules" begin
    @test Γ0_hist[eq_other, :] ≈ Γ0_alt_inf[eq_other, :]
    @test Γ1_hist[eq_other, :] ≈ Γ1_alt_inf[eq_other, :]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt_inf[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt_inf[eq_mp, :])
    @test Γ0_hist[eq_other, eq_other] ≈ Γ0_alt_zero[eq_other, eq_other]
    @test Γ1_hist[eq_other, eq_other] ≈ Γ1_alt_zero[eq_other, eq_other]
    @test !(Γ0_hist[eq_mp, :]  ≈ Γ0_alt_zero[eq_mp, :])
    @test !(Γ1_hist[eq_mp, :]  ≈ Γ1_alt_zero[eq_mp, :])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_smooth[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_smooth[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_smooth[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_smooth[eq_mp, eq_othery])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_smooth_alt[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_smooth_alt[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_smooth_alt[eq_mp, eq_other])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_smooth_alt[eq_mp, eq_other])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_flex_ait[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_flex_ait[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_flex_ait[eq_mp, eq_other])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_flex_ait[eq_mp, eq_other])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_ait[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_ait[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_ait[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_ait[eq_mp, eq_othery])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_ngdp[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_ngdp[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_ngdp[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_ngdp[eq_mp, eq_othery])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_rw[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_rw[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_rw[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_rw[eq_mp, eq_othery])

    @test Γ0_histy[eq_othery, eq_otherymp] ≈ Γ0_alt_rw_zero[eq_othery, eq_otherymp]
    @test Γ1_histy[eq_othery, eq_otherymp] ≈ Γ1_alt_rw_zero[eq_othery, eq_otherymp]
    @test !(Γ0_histy[eq_mp, eq_other]  ≈ Γ0_alt_rw_zero[eq_mp, eq_othery])
    @test !(Γ1_histy[eq_mp, eq_other]  ≈ Γ1_alt_rw_zero[eq_mp, eq_othery])
end

# Checking if setting variables works
## These values are random and not actually good choices per se
m <= Setting(:ait_Thalf, 14.2)
m <= Setting(:gdp_Thalf, 5.6)
m <= Setting(:smooth_ait_gdp_ρ_smooth, 0.23)
m <= Setting(:smooth_ait_gdp_φ_π, 0.36)
m <= Setting(:smooth_ait_gdp_φ_y, 0.56)
m <= Setting(:smooth_ait_gdp_alt_ρ_smooth, 0.76)
m <= Setting(:smooth_ait_gdp_alt_φ_π, 0.11)
m <= Setting(:smooth_ait_gdp_alt_φ_y, 0.89)
m <= Setting(:rw_ρ_smooth, 0.45)
m <= Setting(:rw_φ_π, 0.83)
m <= Setting(:rw_φ_y, 0.94)
m <= Setting(:ρ_rw, 0.63)

@testset "Checking Settings assignments" begin
    @test get_setting(m, :ait_Thalf) == 14.2
    @test get_setting(m, :gdp_Thalf) == 5.6
    @test get_setting(m, :smooth_ait_gdp_ρ_smooth) == 0.23
    @test get_setting(m, :smooth_ait_gdp_φ_π) == 0.36
    @test get_setting(m, :smooth_ait_gdp_φ_y) == 0.56
    @test get_setting(m, :smooth_ait_gdp_alt_ρ_smooth) == 0.76
    @test get_setting(m, :smooth_ait_gdp_alt_φ_π) == 0.11
    @test get_setting(m, :smooth_ait_gdp_alt_φ_y) == 0.89
    @test get_setting(m, :rw_ρ_smooth) == 0.45
    @test get_setting(m, :rw_φ_π) == 0.83
    @test get_setting(m, :rw_φ_y) == 0.94
    @test get_setting(m, :ρ_rw) == 0.63
end

ext: jl


text:
using ITensors,
      Test

@testset "MPO Basics" begin
  N = 6
  sites = [Index(QN(-1)=>1, QN(1)=>1; tags="Site,n=$n") for n=1:N]
  links = [Index(QN()=>1; tags="Links,l=$n") for n=1:N-1]
  @test length(MPO()) == 0
  #O = MPO(sites)
  O = MPO(N)
  for i in 1:length(O)
    O[i] = randomITensor(QN(), sites[i], sites[i]')
  end
  @test length(O) == N

  O[1] = emptyITensor(sites[1], prime(sites[1]))
  @test hasind(O[1],sites[1])
  @test hasind(O[1],prime(sites[1]))
  P = copy(O)
  @test hasind(P[1],sites[1])
  @test hasind(P[1],prime(sites[1]))
  # test constructor from Vector{ITensor}

  K = MPO(N)
  K[1] = randomITensor(QN(), dag(sites[1]), sites[1]', links[1])
  for i in 2:N-1
    K[i] = randomITensor(QN(), dag(sites[i]), sites[i]', dag(links[i-1]), links[i])
  end
  K[N] = randomITensor(QN(), dag(sites[N]), sites[N]', dag(links[N-1]))

  J = MPO(N)
  J[1] = randomITensor(QN(), dag(sites[1]), sites[1]', links[1])
  for i in 2:N-1
    J[i] = randomITensor(QN(), dag(sites[i]), sites[i]', dag(links[i-1]), links[i])
  end
  J[N] = randomITensor(QN(), dag(sites[N]), sites[N]', dag(links[N-1]))

  L = MPO(N)
  L[1] = randomITensor(QN(), dag(sites[1]), sites[1]', links[1])
  for i in 2:N-1
    L[i] = randomITensor(QN(), dag(sites[i]), sites[i]', dag(links[i-1]), links[i])
  end
  L[N] = randomITensor(QN(), dag(sites[N]), sites[N]', dag(links[N-1]))

  @test length(K) == N
  @test ITensors.data(MPO(copy(ITensors.data(K)))) == ITensors.data(K)

  phi = MPS(N)
  phi[1] = randomITensor(QN(-1), sites[1], links[1])
  for i in 2:N-1
    phi[i] = randomITensor(QN(-1), sites[i], dag(links[i-1]), links[i])
  end
  phi[N] = randomITensor(QN(-1), sites[N], dag(links[N-1]))

  psi = MPS(N)
  psi[1] = randomITensor(QN(-1), sites[1], links[1])
  for i in 2:N-1
    psi[i] = randomITensor(QN(-1), sites[i], dag(links[i-1]), links[i])
  end
  psi[N] = randomITensor(QN(-1), sites[N], dag(links[N-1]))

  @testset "orthogonalize!" begin
    orthogonalize!(phi, 1)
    orthogonalize!(K, 1)
    orig_inner = ⋅(phi, K, phi)
    orthogonalize!(phi, div(N, 2))
    orthogonalize!(K, div(N, 2))
    @test ⋅(phi, K, phi) ≈ orig_inner
  end

  @testset "inner <y|A|x>" begin
    @test maxlinkdim(K) == 1
    phidag = dag(phi)
    prime!(phidag)
    phiKpsi = phidag[1]*K[1]*psi[1]
    for j = 2:N
      phiKpsi *= phidag[j]*K[j]*psi[j]
    end
    @test phiKpsi[] ≈ inner(phi, K, psi)
  end

  @testset "inner <By|A|x>" begin
    phidag = dag(phi)
    prime!(phidag, 2)
    Jdag = dag(J)
    prime!(Jdag)
    for j ∈ eachindex(Jdag)
      swapprime!(Jdag[j],2,3)
      swapprime!(Jdag[j],1,2)
      swapprime!(Jdag[j],3,1)
    end

    phiJdagKpsi = phidag[1]*Jdag[1]*K[1]*psi[1]
    for j ∈ eachindex(psi)[2:end]
      phiJdagKpsi = phiJdagKpsi*phidag[j]*Jdag[j]*K[j]*psi[j]
    end

    @test phiJdagKpsi[] ≈ inner(J,phi,K,psi)

    badsites = [Index(2,"Site") for n=1:N+1]
    badpsi = randomMPS(badsites)
    @test_throws DimensionMismatch inner(J,phi,K,badpsi)
  end

  @testset "error_contract" begin
    dist = sqrt(abs(1 + (inner(phi,phi) - 2*real(inner(phi,K,psi)))
                        /inner(K,psi,K,psi)))
    @test dist ≈ error_contract(phi,K,psi)
  end

  @testset "contract" begin
    @test maxlinkdim(K) == 1
    psi_out = contract(K, psi, maxdim=1)
    @test inner(phi,psi_out) ≈ inner(phi,K,psi)
    @test_throws ArgumentError contract(K, psi, method="fakemethod")
  end

  # TODO: implement add for QN MPOs and add this test back
  #@testset "add(::MPO, ::MPO)" begin
  #  shsites = siteinds("S=1/2", N)
  #  M = add(K, L)
  #  @test length(M) == N
  #  k_psi = contract(K, psi, maxdim=1)
  #  l_psi = contract(L, psi, maxdim=1)
  #  @test inner(psi, k_psi + l_psi) ≈ ⋅(psi, M, psi) atol=5e-3
  #  @test inner(psi, sum([k_psi, l_psi])) ≈ dot(psi, M, psi) atol=5e-3
  #  for dim in 2:4
  #    shsites = siteinds("S=1/2",N)
  #    K = basicRandomMPO(N, shsites; dim=dim)
  #    L = basicRandomMPO(N, shsites; dim=dim)
  #    M = K + L
  #    @test length(M) == N
  #    psi = randomMPS(shsites)
  #    k_psi = contract(K, psi)
  #    l_psi = contract(L, psi)
  #    @test inner(psi, k_psi + l_psi) ≈ dot(psi, M, psi) atol=5e-3
  #    @test inner(psi, sum([k_psi, l_psi])) ≈ inner(psi, M, psi) atol=5e-3
  #    psi = randomMPS(shsites)
  #    M = add(K, L; cutoff=1E-9)
  #    k_psi = contract(K, psi)
  #    l_psi = contract(L, psi)
  #    @test inner(psi, k_psi + l_psi) ≈ inner(psi, M, psi) atol=5e-3
  #  end
  #end

  @testset "contract(::MPO, ::MPO)" begin
    @test maxlinkdim(K) == 1
    @test maxlinkdim(L) == 1
    KL = contract(prime(K), L; maxdim = 1)
    Lpsi = contract(L, psi; maxdim = 1)
    psi_kl_out = contract(prime(K), Lpsi; maxdim = 1)
    @test inner(psi,KL,psi) ≈ inner(psi, psi_kl_out) atol=5e-3
  end

  @testset "*(::MPO, ::MPO)" begin
    @test maxlinkdim(K) == 1
    @test maxlinkdim(L) == 1
    KL = *(prime(K), L; maxdim = 1)
    psi_kl_out = *(prime(K), *(L, psi; maxdim = 1); maxdim = 1)
    @test ⋅(psi, KL, psi) ≈ dot(psi, psi_kl_out) atol=5e-3
  end

  sites = siteinds("S=1/2",N)
  O = MPO(sites,"Sz")
  @test length(O) == N # just make sure this works

  @test_throws ArgumentError randomMPO(sites, 2)
  @test isnothing(linkind(MPO(fill(ITensor(), N), 0, N + 1), 1))
end

@testset "splitblocks" begin
  N = 4
  sites = siteinds("S=1", N; conserve_qns = true)
  ampo = AutoMPO()
  for j=1:N-1
    ampo .+= 0.5, "S+", j, "S-", j+1
    ampo .+= 0.5, "S-", j, "S+", j+1
    ampo .+=      "Sz", j, "Sz", j+1
  end
  H = MPO(ampo, sites)

  # Split the tensors to make them more sparse
  # Drops zero blocks by default
  H̃ = splitblocks(linkinds, H)

  @test prod(H) ≈ prod(H̃)

  @test nnz(H[1]) == 13
  @test nnz(H[2]) == 51
  @test nnz(H[3]) == 51
  @test nnz(H[4]) == 13

  @test nnzblocks(H[1]) == 7
  @test nnzblocks(H[2]) == 11
  @test nnzblocks(H[3]) == 11
  @test nnzblocks(H[4]) == 7

  @test nnz(H̃[1]) == nnzblocks(H̃[1]) == count(≠(0), H[1]) == count(≠(0), H̃[1]) == 9
  @test nnz(H̃[2]) == nnzblocks(H̃[2]) == count(≠(0), H[2]) == count(≠(0), H̃[2]) == 18
  @test nnz(H̃[3]) == nnzblocks(H̃[3]) == count(≠(0), H[3]) == count(≠(0), H̃[3]) == 18
  @test nnz(H̃[4]) == nnzblocks(H̃[4]) == count(≠(0), H[4]) == count(≠(0), H̃[4]) == 9
end

nothing


text:
function lempel_ziv_complexity(sequence)
    sub_strings = Set()
    n = length(sequence)

    ind = 1
    inc = 1
    while true
        if ind + inc > n
            break
        end
        sub_str = sequence[ind : ind + inc]
        if sub_str in sub_strings
            inc += 1
        else
            push!(sub_strings, sub_str)
            ind += inc
            inc = 1
        end
    end
    return length(sub_strings)
end

function lempel_ziv(ca::AbstractCA)
    ca_size = size(ca.evolution, 1)
    lz_tot = 0

    for i=1:ca_size
        lz_tot += lempel_ziv_complexity(ca.evolution[i,:])
    end
    lz_tot/ca_size
end


text:
try
    using EventTrackerTests
    true
catch
    false
end || begin
    push!(LOAD_PATH, joinpath(@__DIR__, "EventTrackerTests"))
    using EventTrackerTests
end


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    AbstractFilter

A method to filter spatial objects.
"""
abstract type AbstractFilter end

"""
    filter(object, filt)

Filter spatial `object` with filtering method `filt`.
"""
Base.filter(object::AbstractSpatialObject, filt::AbstractFilter) =
  @error "not implemented"

#------------------
# IMPLEMENTATIONS
#------------------
include("filtering/unique_coords.jl")


text:
"""
This example uses the ProfileData struct and GP.
"""

include("GP1.jl")

log_γs = -3.0:0.1:3.0

V_name = Dict("T" =>"Temperature [°C]", "wT"=>"Temperature flux [°C⋅m/s]")

v_str = "T"
# v_str = "wT"

N = 1
# N = 2

function animate_profile_trained_on_any(filename_predict, filenames_train, k::Int64, γ, d, D)

    𝒟predict = construct_profile_data(filename_predict, v_str, D; N=N)
    kernel = get_kernel(k, γ, 0.0)

    𝒟train = construct_profile_data_multiple(filenames_train, v_str, D; N=N)
    𝒢 = construct_gpr(𝒟train.x_train, 𝒟train.y_train, kernel; distance_fn=d, z=𝒟train.zavg);
    println(length(𝒟predict.x_train))

    println(length(𝒟train.x_train))
    gpr_prediction = get_gpr_pred(𝒢, 𝒟predict)

    animation_set = 1:10:(𝒟predict.Nt-2)
    anim = @animate for i in animation_set
        exact = 𝒟predict.v[:,i]
        day_string = string(floor(Int, 𝒟predict.t[i]/86400))
        p1 = scatter(gpr_prediction[i], 𝒟predict.zavg, label = "GP")
        # xlims=(minimum(data.v[:,1]),maximum(data.v[:,1]))
        xlims=(18,20)

        if i<size(𝒟train.v)[2]
            exact16 = 𝒟train.v[:,i]
        else
            exact16 = 𝒟train.v[:,end]
        end

        plot!(exact, 𝒟predict.z, legend = :topleft, label = "LES", xlabel = "$(V_name["T"])", ylabel = "Depth [m]", title = "day $(day_string)", xlims=xlims)
        # plot!(exact16, 𝒟train.z, legend = :topleft, label = "LES gs 16", xlabel = "$(V_name["T"])", ylabel = "depth", xlims=xlims)
        display(p1)
    end

    return anim
end

# function l2norm_strat_penalty(a,b,z) # d(x,x') = || x - x' ||
#     α_proxy(x) = x[2] - x[1]
#     # println("hello")
#     # println("$(abs(α_proxy(a)-α_proxy(b)))")
#     if abs(α_proxy(a)-α_proxy(b))>0.05
#         return l2_norm(a,b) + 0.0001
#     end
#     return l2_norm(a,b)
# end

filename="general_strat_32_profiles"
filenames = [
             "general_strat_4_profiles.jld2",
             # "general_strat_8_profiles.jld2",
             # "general_strat_12_profiles.jld2",
             # "general_strat_16_profiles.jld2",
             # "general_strat_20_profiles.jld2",
             # "general_strat_24_profiles.jld2",
             # "general_strat_28_profiles.jld2",
             # "general_strat_32_profiles.jld2"
             ]
filenames_descript = "8"
# filenames = ["general_strat_8_profiles.jld2","general_strat_32_profiles.jld2"]

γ=5.5
k=2
anim = animate_profile_trained_on_any("$(filename).jld2", filenames, k, γ, l2norm_strat_penalty, 16)
gif(anim, pwd() * "/../les/data_sandreza/$(filename)/gp_γ$(γ)_k$(k)_l2norm_strat_penalty_trainedOnGs$(filenames_descript)_scaledwithtime_N$(N).gif", fps = 15)

max_stars_repo_name: adelinehillier/LearnConvection

max_issues_repo_name: adelinehillier/LearnConvection

max_forks_repo_name: adelinehillier/LearnConvection


text:
# Undergraduate Student: Arturo Burgos
# Professor: João Rodrigo Andrade
# Federal University of Uberlândia - UFU, Fluid Mechanics Laboratory - MFLab, Block 5P, Uberlândia, MG, Brazil

# Third exercise: Fibonacci sequence - by a common loop

using Printf
using Statistics
using CPUTime

println("Choose a indice value for n: ")
N = chomp(readline())
n = parse(Int32,N)
println("\n")

Fbn = 0

##############################
# Variable recursive fibonacci  --> Worse performance, why?
##############################

function recurfibonacci(k)

    if k <= 2
     return 1
    else
       return S = recurfibonacci(k-1) + recurfibonacci(k-2)
    end

end

Fbn = recurfibonacci(n)

println("The correspond indices number is: ",Fbn)

@time @CPUtime recurfibonacci(n)

max_stars_repo_name: arturofburgos/Comparing-Languages


text:
function update_weights(data, e, closest_same, closest_other, weights, weights_mult, m, k, max_f_vals, min_f_vals)

	# Distance weights for penalty term

	# Compute weights for each nearest hit
	d_vals_closest_same = 1 ./ (sum(abs.(e .- closest_same)./((max_f_vals .- min_f_vals) .+ eps(Float64)), dims=2) .+ eps(Float64))
	dist_weights_penalty = d_vals_closest_same ./ sum(d_vals_closest_same)

	# Distance weights for reward term
	d_vals_closest_other = 1 ./ (sum(abs.(e .- closest_other)./((max_f_vals .- min_f_vals) .+ eps(Float64)), dims=2) .+ eps(Float64))
	dist_weights_reward = d_vals_closest_other ./ sum(d_vals_closest_other)

	# Penalty term
	penalty = sum(dist_weights_penalty .* (abs.(e .- closest_same)./((max_f_vals .- min_f_vals) .+ eps(Float64))), dims=1)

	# Reward term
	reward = sum(weights_mult .* (dist_weights_reward .* (abs.(e .- closest_other)./((max_f_vals .- min_f_vals) .+ eps(Float64)))), dims=1)

	# Weights update
	weights = weights .- penalty./m .+ reward./m

	# Return updated weights.
	return vec(weights)
end

max_stars_repo_name: Architecton/diplomska-naloga-koda

max_issues_repo_name: Architecton/diplomska-naloga-koda

max_forks_repo_name: Architecton/diplomska-naloga-koda


text:
__precompile__()

module PlyIO

# Types for the ply data model
export Ply, PlyElement, PlyComment, ArrayProperty, ListProperty
export plyname  # Is there something in base we could overload for this?

# High level file IO
# (TODO: FileIO?)
export load_ply, save_ply

include("types.jl")
include("io.jl")

end


text:
# Autogenerated wrapper script for pigz_jll for i686-w64-mingw32
export pigz, unpigz

using Zlib_jll
JLLWrappers.@generate_wrapper_header("pigz")
JLLWrappers.@declare_executable_product(pigz)
JLLWrappers.@declare_executable_product(unpigz)
function __init__()
    JLLWrappers.@generate_init_header(Zlib_jll)
    JLLWrappers.@init_executable_product(
        pigz,
        "bin\\pigz.exe",
    )

    JLLWrappers.@init_executable_product(
        unpigz,
        "bin\\unpigz.exe",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
#####
##### Memory Buffer:
##### datastructure to collect self-play experience
#####

"""
    TrainingSample{Board}

Type of a training sample. A sample features the following fields:
- `b::Board` is the board position (by convention, white is to play)
- `π::Vector{Float64}` is the recorded MCTS policy for this position
- `z::Float64` is the reward collected at the end of the game
- `t::Float64` is the number of moves remaining before the end of the game
- `n::Int` is the number of times the board position `b` was recorded

As revealed by the last field `n`, several samples that correspond to the
same board position can be merged, in which case the `π`, `z` and `t`
fields are averaged together.
"""
struct TrainingSample{Board}
  b :: Board
  π :: Vector{Float64}
  z :: Float64
  t :: Float64 # average time before game end
  n :: Int # number of times the board position `b` has been seen
end

"""
    MemoryBuffer{Board}

A circular buffer to hold memory samples.

# How to use

- Use `new_batch!(mem)` to start a new batch, typically once per iteration
  before self-play.
- Use `push_sample!(mem, board, policy, white_playing, turn)` to record
  a sample during a game, where `turn` is the number of actions that have
  been played by both players since the start of the game.
- Use `push_game!(mem, white_reward, game_length)` when a game terminates
  for which samples have been collected.
"""
mutable struct MemoryBuffer{Board}
  # State-policy pairs accumulated during the current game.
  # The z component is 1 if white was playing and -1 otherwise
  cur :: Vector{TrainingSample{Board}}
  # Long-term memory
  mem :: CircularBuffer{TrainingSample{Board}}
  last_batch_size :: Int
  function MemoryBuffer{B}(size, experience=[]) where B
    mem = CircularBuffer{TrainingSample{B}}(size)
    append!(mem, experience)
    new{B}([], mem, 0)
  end
end

function Base.empty!(b::MemoryBuffer)
  empty!(b.cur)
  empty!(b.mem)
  b.last_batch_size = 0
end

Base.length(b::MemoryBuffer) = length(b.mem)

function merge_samples(es::Vector{TrainingSample{B}}) where B
  b = es[1].b
  π = mean(e.π for e in es)
  z = mean(e.z for e in es)
  n = sum(e.n for e in es)
  t = mean(e.t for e in es)
  return TrainingSample{B}(b, π, z, t, n)
end

# Merge samples that correspond to identical boards
function merge_by_board(es::AbstractVector{TrainingSample{B}}) where B
  dict = Dict{B, Vector{TrainingSample{B}}}()
  sizehint!(dict, length(es))
  for e in es
    if haskey(dict, e.b)
      push!(dict[e.b], e)
    else
      dict[e.b] = [e]
    end
  end
  return [merge_samples(es) for es in values(dict)]
end

function apply_symmetry(::Type{Game}, sample, (symboard, aperm)) where Game
  mask = GI.actions_mask(Game(sample.b))
  symmask = GI.actions_mask(Game(symboard))
  π = zeros(eltype(sample.π), length(mask))
  π[mask] = sample.π
  π = π[aperm]
  @assert iszero(π[.~symmask])
  π = π[symmask]
  return TrainingSample{typeof(sample.b)}(
    symboard, π, sample.z, sample.t, sample.n)
end

function augment_with_symmetries(Game, samples)
  symsamples = [apply_symmetry(Game, s, sym)
    for s in samples for sym in GI.symmetries(Game, s.b)]
  return [samples ; symsamples]
end

get_experience(buf::MemoryBuffer) = buf.mem[:]

last_batch(buf::MemoryBuffer) = buf.mem[end-last_batch_size(buf)+1:end]

last_batch_size(buf::MemoryBuffer) = min(buf.last_batch_size, length(buf))

function push_sample!(buf::MemoryBuffer, board, policy, white_playing, turn)
  player_code = white_playing ? 1. : -1.
  ex = TrainingSample(board, policy, player_code, float(turn), 1)
  push!(buf.cur, ex)
  buf.last_batch_size += 1
end

function push_game!(buf::MemoryBuffer, white_reward, game_length)
  for ex in buf.cur
    r = ex.z * white_reward
    t = game_length - ex.t
    push!(buf.mem, TrainingSample(ex.b, ex.π, r, t, ex.n))
  end
  empty!(buf.cur)
end

function new_batch!(buf::MemoryBuffer)
  buf.last_batch_size = 0
end


text:
using LinearAlgebra, ImageView, Images, ColorVectorSpace, Colors, Dates, SignedDistanceFields

include("raymarching.jl")
include("camera.jl");
include("scene.jl")
include("shaders.jl")
@time include("render.jl")

imshow(myPic)

#begin stamp = Dates.format(Dates.DateTime(Dates.now()), "dd-u-yyyy-HH:MM:SS")
#save(pwd()*"/myPic-"*stamp*".png",map(clamp01nan,myPic)) end


text:
using Roots
using Optim
using LineSearches
using Parameters, Setfield

using BifurcationKit
const BK = BifurcationKit

################################
# mode-1

v2d_m1(x) = vcat([[y[1];y[2]] for y in x]...)
d2v_m1(x) = [[x[2*i-1];x[2*i];0.0] for i in 1:(length(x)÷2)]
Iv2Id_m1(II) = vcat([[2*x-1;2*x] for x in II]...)

function hess_v2d_m1(at,HH)
    II = Int64[]
    J = Int64[]
    Z = Float64[]

    III, JJJ, VVV = findnz(HH)

    for iii in 1:length(III)
        aa = Iv2Id_m1(III[iii])
        bb = Iv2Id_m1(JJJ[iii])

        for ii in 1:2
            for jj in 1:2
                append!(II,aa[ii])
                append!(J,bb[jj])
                append!(Z,VVV[iii][ii,jj])
            end
        end
    end
    return sparse(II,J,Z,2*length(at.X),2*length(at.X))
end

#######################
# (M1-2) K,α,U format for Mode I
# Here U is in dofs format of size 2*N1

function f_en_m1(at,Ubar,K,α)
    U_CLE = UHAT(at,α)
    Ubar_v = d2v_m1(vcat(Ubar,zeros(2*(length(at.X)-at.Ifree))))
    return energy_pp(at,K*U_CLE .+ Ubar_v)
end

function g_en_m1(at,Ubar,K,α)
    U_CLE = UHAT(at,α)
    Ubar_v = d2v_m1(vcat(Ubar,zeros(2*(length(at.X)-at.Ifree))))
    return v2d_m1(grad_pp(at,K*U_CLE.+ Ubar_v))
end

function g_en_m1_alpha(at,Ubar,K,α)
    N2 = at.Iclamp - 1
    VHAT_d_m1 = v2d_m1(d1UHAT(at,α))
#    xx = vcat(Ubar,zeros(length(at.X)-at.Ifree))
    g = g_en_m1(at,Ubar,K,α)[1:(2*N2)]
    return dot(g,K*VHAT_d_m1[1:(2*N2)])
end

function g_en_m1_tot(at,Ubar,K,α)
    N2 = at.Iclamp - 1
    N1 = at.Ifree
    return [g_en_m1(at,Ubar,K,α)[1:(2*N1)]; g_en_m1_alpha(at,Ubar,K,α)]
end

function h_en_m1(at,Ubar,K,α)
    U_CLE = UHAT(at,α)
    Ubar_v = d2v_m1(vcat(Ubar,zeros(2*(length(at.X)-at.Ifree))))
    return hess_v2d_m1(at,hess_pp(at,K*U_CLE.+ Ubar_v))
end

function h_en_m1_tot(at,Ubar,K,α)
    N2 = at.Iclamp - 1
    N1 = at.Ifree
    hess = spzeros(2*N1+1,2*N1+1)
    hess_base = h_en_m1(at,Ubar,K,α)
    g_base = g_en_m1(at,Ubar,K,α)
    hess[1:(2*N1),1:(2*N1)] = hess_base[1:(2*N1),1:(2*N1)]
    VHAT = v2d_m1(d1UHAT(at,α))
    VVHAT = v2d_m1(d1d1UHAT(at,α))
#    VHAT = f_d1UHAT(lambda)
#    VVHAT = f_d1d1UHAT(lambda)
#    balpha = (hess_base[1:N2,1:N2]*VHAT[1:N2])
    balpha = (hess_base*(K*VHAT))
    Calpha = dot(balpha[1:(2*N2)],K*VHAT[1:(2*N2)]) + dot(g_base[1:(2*N2)],K*VVHAT[1:(2*N2)])
    hess[2*N1+1,1:(2*N1)] = balpha[1:(2*N1)]
    hess[1:(2*N1),2*N1+1] = balpha[1:(2*N1)]
    hess[2*N1+1,2*N1+1] = Calpha
    return hess
end

mutable struct NCF
    at
    f
    g
    h
    Ks
    αs
    Us
    data::Dict{Any,Any}
end

function NCF(; R=5.0, R_star = 1.0, lt=tri(), mode1 = true)
    at = AtmModel(R=R, R_star = R_star, lt=lt, mode1 = mode1)
    if at.mode1
        f = (Ubar,K,α) -> f_en_m1(at,Ubar,K,α)
       # f = f_en_m1
        g = (Ubar,K,α) -> g_en_m1_tot(at,Ubar,K,α)
#        g = g_en_m1_tot
        h = (Ubar,K,α) -> h_en_m1_tot(at,Ubar,K,α)
#        h = h_en_m1_tot
    end
    Ks = nothing
    αs = nothing
    Us = nothing
    data = Dict{Any,Any}()
    return NCF(at,f,g,h,Ks,αs,Us,data)
end

function Newton_static(ncf,xx,K,α; show = false)
    at = ncf.at
    N1 = at.Ifree
    x = copy(xx)
    If = 1:2*N1
    show && @printf("------Newton Iteration------\n")
    nit = 0
    for nit = 0:15
        ∇E = ncf.g(x,K,α)[If]
        show && @printf("%d : %4.2e \n", nit, norm(∇E[If], Inf))
        norm(∇E[If], Inf) < 1e-8 && break
        x[If] = x[1:2*N1] - ncf.h(x,K,α)[If, If] \ ∇E[If]
    end
    nit == 12 && warn("the Newton iteration did not converge")
    show && @printf("----------------------------\n")
    return x
end

function Newton_flex(ncf,xx,K;show=false)
    at = ncf.at
    N1 = at.Ifree
    x = copy(xx)
    If = 1:2*N1
    show && @printf("------Newton Iteration------\n")
    nit = 0
    for nit = 0:15
        ∇E = ncf.g(x[If],K,x[end])
        show && @printf("%d : %4.2e \n", nit, norm(∇E, Inf))
        norm(∇E, Inf) < 1e-8 && break
        x = x - ncf.h(x[If],K,x[end]) \ ∇E
    end
    nit == 12 && warn("the Newton iteration did not converge")
    show && @printf("----------------------------\n")
    return x
end

function alpha_to_K(ncf,α;K = 1.0, show=false)
    at = ncf.at
    N1 = at.Ifree
    Ubar0 = zeros(2*N1)
    fk = K -> ncf.g(Ubar0,K,α)[2*N1+1]
    gk(K) = ForwardDiff.derivative(fk,K)
    show && @printf("------Newton Iteration------\n")
    nit = 0
    for nit = 0:15
        #∇E = g_en_m1_tot(x[If],K,x[end])
        show && @printf("%d : %4.2e \n", nit, norm(fk(K)))
        norm(fk(K)) < 1e-8 && break
        K = K - fk(K)/gk(K)
    end
    nit == 12 && warn("the Newton iteration did not converge")
    show && @printf("----------------------------\n")
    return K
end

function finding_k!(ncf,xbar,K,α_0)
    N1 = ncf.at.Ifree
    xbar2 = Newton_static(ncf,xbar,K,α_0,show=false)
    xbar[1:2*N1] = xbar2[1:2*N1]
    return ncf.g(xbar,K,α_0)[end]
end

function preparation(ncf)
    if ncf.Us == nothing
        N1 = ncf.at.Ifree
        α_0 = -0.5*norm(ncf.at.X[1].-ncf.at.X[2])

        #step1: estimate (K_-,K_+) using CLE displacements only
        ks = []
        αs = range(-0.4,0.0,length=10)

        for α in αs
            if length(ks) == 0
                push!(ks,NCFlex.alpha_to_K(ncf,α, K = 0.4, show=false))
            else
                push!(ks,NCFlex.alpha_to_K(ncf,α, K = ks[end], show=false))
            end
        end
        println("(1/3) CLE-only estimation done")

        #finding K for which solution computed with the static scheme is also the solution to the flex scheme
        ff0(x) = ncf.f(x,minimum(ks),α_0)
        gg0! = (gg, x) -> copyto!(gg, ncf.g(x,minimum(ks),α_0)[1:2*N1])
        xbar0 = Optim.minimizer(optimize(ff0, gg0!, zeros(2*N1),ConjugateGradient(linesearch = BackTracking(order=2,maxstep=Inf)),Optim.Options(show_trace=false,g_tol=1e-3, f_tol=1e-32)))
        println("(2/3) First static equilibrium found")
        xbar1 = NCFlex.Newton_static(ncf,xbar0,minimum(ks),α_0,show=false);

        fff = K-> NCFlex.finding_k!(ncf,xbar1,K,α_0)

        K_best = find_zero(fff, (minimum(ks),maximum(ks)),FalsePosition(),atol = 1e-8,verbose=false)

        x0 = NCFlex.Newton_flex(ncf,[xbar1;α_0],K_best,show=false)

        Us = []
        push!(Us,x0[1:end-1])
        ncf.Us = Us
        ncf.Ks = [K_best]
        ncf.αs = [x0[end]]
        ncf.data["K_est"] = extrema(ks)
        println("(3/3) First flex equilibrium found")
    else
        println("Nothing to prepare")
    end
end

function simple_continuation(ncf ; dsmin = 0.000001, dsmax = 0.0003, ds= 0.00004,
            pMax = 4.1, maxSteps=100,theta=0.5,
            tangentAlgo = BorderedPred(),dotPALC = (x, y) -> dot(x, y)/length(x),
            linsolver = DefaultLS(),linearAlgo = BorderingBLS())
    N1 = ncf.at.Ifree
    function FF_m1(x, p)
        @unpack K0 = p
        return ncf.g(x[1:2*N1],K0,x[end])
    end
    function JJ_m1(x,p)
        @unpack K0 = p
    return ncf.h(x[1:2*N1],K0,x[end])
    end

    optcont1 = ContinuationPar(dsmin = dsmin, dsmax = dsmax, ds= ds, pMax = pMax, maxSteps=maxSteps,theta=theta,
        newtonOptions = NewtonPar(tol = 1e-8,linsolver = linsolver),saveSolEveryStep = 0,doArcLengthScaling = false);

    x0 = [ncf.Us[1];ncf.αs[1]]
    par = (K0 = ncf.Ks[1],)

    iter1 = BK.ContIterable(FF_m1, JJ_m1, x0, par, (@lens _.K0),
        optcont1; plot = false, verbosity = 0,
        tangentAlgo = tangentAlgo,dotPALC = dotPALC, linearAlgo = linearAlgo)
        #tangentAlgo = SecantPred(),dotPALC = (x, y) -> dot(x, y))

    Kdot_count = 0
    Sts = []
    dαs = []
    dKs = []

    cstate1 = nothing
    cstate2 = nothing

    next = nothing

    if ~isnothing(ncf.Us) && length(ncf.Us) > 1
        x0 = [ncf.Us[end-1];ncf.αs[end-1]]
        x1 = [ncf.Us[end];ncf.αs[end]]
        k0 = ncf.Ks[end-1]
        k1 = ncf.Ks[end]
        state = BK.iterateFromTwoPoints(iter1, x0, k0, x1, k1);
        next = iterate(iter1,state[1])
    else
        next = iterate(iter1)
    end

    while next !== nothing
        (i,state) = next

        ########iteration
      #  println(state.stepsizecontrol)
        print(state.step)
        print(" ")
        if state.step == 0 && state.tau.u[end] < 0.0
      #      println(state.z_old)
            state.tau.u = -state.tau.u
            state.tau.p = -state.tau.p
    #        state.stepsizecontrol = false
        end
        if ~isnothing(cstate1) && cstate1.tau.p*state.tau.p < 0.0
            if abs(state.tau.p) > 3e-5
                state.step = cstate1.step
                state.z_pred  = cstate1.z_pred
                state.tau = cstate1.tau
                state.z_old = cstate1.z_old
                state.stepsizecontrol = false
                state.ds = cstate1.ds/2
                #println("HA")
            else
                state.stepsizecontrol=true
                state.ds = iter1.contParams.ds
                Kdot_count+= 1
                println(" ")
                println("Fold point found")
                push!(ncf.αs, getx(state)[end])
                push!(ncf.Ks, getp(state))
                push!(Sts,state.step)
                push!(dαs,state.tau.u[end])
                push!(dKs,state.tau.p)
                push!(ncf.Us,getx(state)[1:(2*N1)])
            end
        else
            push!(ncf.αs, getx(state)[end])
            push!(ncf.Ks, getp(state))
            push!(Sts,state.step)
            push!(dαs,state.tau.u[end])
            push!(dKs,state.tau.p)
            push!(ncf.Us,getx(state)[1:(2*N1)])
        end
         cstate1 = copy(state)
        if Kdot_count > 0
            state.stopcontinuation = true
        end
        #################
        next = iterate(iter1,state)
    end

    ncf.data["Sts"] = Sts
    ncf.data["dαs"] = dαs
    ncf.data["dKs"] = dKs
    println(" ")
    println("End")
end

function Du(Us,at)
    Uv = NCFlex.d2v_m1(vcat(Us,zeros(2*(length(at.X)-at.Ifree))))
    return vcat([ [Uv[j] .- Uv[i] for i in at.Ns[j]] for j in 1:length(at.X)]...)
end

function doth1(x,y)
    Dx = Du(x[1:end-1],test.at)
    Dy = Du(y[1:end-1],test.at)
    return dot(Dx,Dy) + x[end]*y[end]
end


text:
using Test

include("game-logic.jl")

@testset "fast attack" begin
    @test !can_do_fast_attack(true)
    @test  can_do_fast_attack(false)
end

@testset "spying" begin
    character_state_combinations = [
        (false, false, false, false),
        (false, false, true, true),
        (false, true, false, true),
        (false, true, true, true),
        (true, false, false, true),
        (true, false, true, true),
        (true, true, true, true),
    ]

    for state in character_state_combinations
        @test can_spy(state[1:3]...) == state[4]
    end
end

@testset "signaling prisoner" begin
    character_state_combinations = [
        (false, false, false),
        (false, true, true),
        (true, false, false),
        (true, true, false),
    ]

    for state in character_state_combinations
        @test can_signal_prisoner(state[1:2]...) == state[3]
    end
end

@testset "freeing prisoner" begin
    character_state_combinations = [
        (false, false, false, false, false),
        (false, false, false, true, true),
        (false, false, true, false, true),
        (false, false, true, true, true),
        (false, true, false, false, false),
        (false, true, false, true, false),
        (false, true, true, false, false),
        (false, true, true, true, false),
        (true, false, false, false, false),
        (true, false, false, true, true),
        (true, false, true, false, false),
        (true, false, true, true, true),
        (true, true, false, false, false),
        (true, true, false, true, false),
        (true, true, true, false, false),
        (true, true, true, true, false),
    ]

    for state in character_state_combinations
        @test can_free_prisoner(state[1:4]...) == state[5]
    end
end


text:
# TODO: Fix this function
function generate_linear_truth(params, system_matrix; single_mode_verification=false)
    logfile = initialize_log(params)
    timing_info = Dict()
    total_runtime = 0.
    # TODO: Move this
    @info "Determining the post-images of the regions under the linear map."
    region_time = @elapsed begin
        region_dict, extents_dict = create_region_data_new(params.domain, params.discretization_step)
        region_post_dict = Dict()

        for i=1:region_dict.count
            region = region_dict[i]

            # Lazy representation saves time
            # TODO: Add the process noise term here for verification
            region_post = LinearMap(0.9999*system_matrix, region)
            region_post_dict[i] = region_post
        end

        region_data = Dict()
        region_data[:extents] = extents_dict
        region_data[:posts] = region_post_dict
    end
    timing_info["region_bound_time_s"] = region_time
    total_runtime += region_time
    @info "Region generation time: " region_time
    save_region_data(params, region_data)
    @info "Generating the transition bounds..."
    bound_time = @elapsed begin
        @info "Calculating transition probability bounds between regions..."
        results_df = DataFrame(Set1 = Int[], Set2 = Int[], MinPr = Float64[], MaxPr = Float64[], MinPrPoint = Array[], MaxPrPoint = Array[])
        for region_pair in Base.product(1:region_dict.count, 1:region_dict.count)
            r1 = region_dict[region_pair[1]]
            r2 = region_dict[region_pair[2]]
            r1_post = region_post_dict[region_pair[1]]

            # Check lazy intersection
            cap = Intersection(r1_post, r2)
            if isempty(cap)
                prange = [0., 0.]
            else
                if isequivalent(r1_post, cap)
                    prange = [1., 1.]
                else
                    prange = [0., 1.]
                end
            end

            df_row = [region_pair[1], region_pair[2], prange[1], prange[2], [-1.], [-1.]]
            push!(results_df, df_row)
        end

        # TODO: Move this constructor to its own fcn
        num_states = region_dict.count
        minPr_mat = spzeros(num_states, num_states)
        maxPr_mat = spzeros(num_states, num_states)
        for i = 1:1:num_states
            sub_row = results_df[results_df.Set1 .== i, :]
            if i == num_states
                minPr_mat[i, i] = 1.
                maxPr_mat[i, i] = 1.
            else
                for j = 1:1:num_states
                    if j == num_states
                        subsub_row = sub_row[sub_row.Set2 .== j, :]
                        minPr_mat[i, j] = 1. - subsub_row.MaxPr[1]
                        maxPr_mat[i, j] = 1. - subsub_row.MinPr[1]
                    else
                        subsub_row = sub_row[sub_row.Set2 .== j, :]
                        minPr_mat[i, j] = subsub_row.MinPr[1]
                        maxPr_mat[i, j] = subsub_row.MaxPr[1]
                    end
                end
            end
        end

        result_mats = Dict("minPr" => minPr_mat, "maxPr" => maxPr_mat)
    end
    total_runtime += bound_time
    timing_info["transition_bound_time_s"] = bound_time
    @info "Bound generation time: " bound_time
    save_transition_matrices(params, result_mats)

    verification_result_mat = nothing
    exp_dir = params.experiment_directory
    if single_mode_verification
        @info "Performing safety verification on single mode for 1 step..."
        verification_time = @elapsed begin
            imdp = create_simple_imdp(result_mats["minPr"], result_mats["maxPr"])
            horizon=1
            verification_result_mat = Globally(imdp, "safe", horizon, "$exp_dir/imdp.txt")
            save_legacy_mats(verification_result_mat, exp_dir, horizon)
        end
        timing_info["single_verification_time_s"] = verification_time
        @info "Verification time: " verification_time
    end

    @info "Total runtime: " total_runtime
    timing_info["total_runtime_s"] = total_runtime
    save_time_info(params, timing_info)

    flush(logfile)
    close(logfile)

    return timing_info, verification_result_mat, result_mats
end


text:
#--------------------------------------------------------------------
# Spacetime Discretization methods in Julia
# Soham 01-2018
#--------------------------------------------------------------------

function restrictmodes!(coeffs::Array{Float64,1}, M::Int)::Array{Float64,1}
    if M+1 > size(coeffs)[1]
        error("Cannot restrict to larger number of modes")
    else
        for index in CartesianRange(size(coeffs))
            index[1] > M+1 ? coeffs[index] = 0 : coeffs[index] = coeffs[index]
        end
    end
    return coeffs
end

function restrictmodes!(coeffs::Array{Float64,2}, M::Int, N::Int)::Array{Float64,2}
    if M+1 > size(coeffs)[1] || N+1 > size(coeffs)[2]
        error("Cannot restrict to larger number of modes")
    else
        for index in CartesianRange(size(coeffs))
            index[1] > M + 1 || index[2] > N + 1 ? coeffs[index] = 0 : coeffs[index] = coeffs[index]
        end
    end
    return coeffs
end

function prolongatemodes(coeffs::Array{Float64,1}, M::Int)::Array{Float64,1}
    if M+1 < size(coeffs)[1]
        error("Cannot prolongate to smaller  number of modes")
    else
        ncoeffs = zeros(M+1)
        oldM = size(coeffs)[1] - 1
        for index in CartesianRange(size(ncoeffs))
            index[1] < oldM+1 ? ncoeffs[index] = coeffs[index] : ncoeffs[index] = ncoeffs[index]
        end
    end
    return ncoeffs
end

function prolongatemodes(coeffs::Array{Float64,2}, M::Int, N::Int)::Array{Float64,2}
    if M+1 < size(coeffs)[1] || N+1 < size(coeffs)[2]
        error("Cannot prolongate to smaller number of modes")
    else
        ncoeffs = zeros(M+1,N+1)
        oldM = size(coeffs)[1] - 1
        oldN = size(coeffs)[2] - 1
        for index in CartesianRange(size(ncoeffs))
            index[1] < oldM+1 && index[2] < oldN+1 ? ncoeffs[index] = coeffs[index] : ncoeffs[index] = ncoeffs[index]
        end
    end
    return ncoeffs
end

function prolongateOP(Nx::Int, M::Int)::Array{Float64,2}
    invndmx = inv(vandermonde(Nx))
    Px = zeros((Nx+1)*M, Nx+1)
    for k in 1:M
        pvndmx = pseudovandermonde(Nx, chebgrid(Nx, M, k))
        Px[1+(k-1)*(Nx+1):k*(Nx+1), :] = pvndmx*invndmx
    end
    return Px
end

function restrictOP(Nx::Int, M::Int)::Array{Float64,2}
    return pinv(prolongateOP(Nx, M))
end

function prolongatePatch(patch::Patch, M::Int)::Dict{Array{Int, 1}, Patch}
    Nx = size(patch.value)[1] - 1
    Ny = size(patch.value)[2] - 1
    Px = prolongateOP(Nx, M)
    Py = prolongateOP(Ny, M)
    mPatch = Px*patch.value*Py'
    return array2dict(mPatch, Nx, Ny, M)
end

function restrictPatch(dbase::Dict{Array{Int, 1}, Patch})::Patch
    Nx = size(dbase[[1,1]].value)[1] - 1
    Ny = size(dbase[[1,1]].value)[2] - 1
    M  = convert(Int, sqrt(length(dbase)))
    Rx = restrictOP(Nx, M)
    Ry = restrictOP(Ny, M)
    sPatch = dict2array(dbase)
    return Patch([1,1], Rx*sPatch*Ry')
end


text:
using Test
using QuantumOptics
using CollectiveSpins

@testset "mpc-equations" begin

edipole = [0, 0, 1]
γ = 1
T = [0:0.1:1;]

spinbasis = SpinBasis(1//2)
sx = sigmax(spinbasis)
sy = sigmay(spinbasis)
sz = sigmaz(spinbasis)

N = 2
systemgeometry = CollectiveSpins.geometry.chain(0.3, N)
system = CollectiveSpins.SpinCollection(systemgeometry, edipole; gammas=γ)
basis = CollectiveSpins.quantum.basis(system)
I = identityoperator(spinbasis)

H = CollectiveSpins.quantum.Hamiltonian(system)
Γ, J = CollectiveSpins.quantum.JumpOperators(system)
Jdagger = [dagger(j) for j=J]
Ω = CollectiveSpins.interaction.OmegaMatrix(system)

Ψ₀ = CollectiveSpins.quantum.blochstate(0., pi/2., N)
ρ₀ = Ψ₀ ⊗ dagger(Ψ₀)

function test_mpc_2particles(t, ρ)
    dρ_h = -1im*(H*ρ - ρ*H)
    dρ_l = Operator(basis)
    for m=1:length(J), n=1:length(J)
       dρ_l += Γ[m,n]*(J[m]*ρ*Jdagger[n] - Jdagger[n]*(J[m]*ρ)/Complex(2) - ρ*Jdagger[n]*J[m]/Complex(2))
    end
    # ========== Test Hamiltonian part ==========
    dsxx_master = expect(sx⊗sx, dρ_h)
    dsyy_master = expect(sy⊗sy, dρ_h)
    dszz_master = expect(sz⊗sz, dρ_h)

    # <sxx>, <syy>, <szz>
    dsxx_mpc = 0.
    dsyy_mpc = 0.
    dszz_mpc = 0.

    @test abs(dsxx_master - dsxx_mpc) < 1e-12
    @test abs(dsyy_master - dsyy_mpc) < 1e-12
    @test abs(dszz_master - dszz_mpc) < 1e-12

    # ========== Test Lindblad part ==========
    dsxx_master = expect(sx⊗sx, dρ_l)
    dsyy_master = expect(sy⊗sy, dρ_l)
    dszz_master = expect(sz⊗sz, dρ_l)

    # <sx>, <sy>, <sz>
    dsxx_mpc = -expect(sx⊗sx, ρ) + Γ[1,2]*(expect(sz⊗sz, ρ) + 0.5*expect(sz⊗I, ρ) + 0.5*expect(I⊗sz, ρ))
    dsyy_mpc = -expect(sy⊗sy, ρ) + Γ[1,2]*(expect(sz⊗sz, ρ) + 0.5*expect(sz⊗I, ρ) + 0.5*expect(I⊗sz, ρ))
    dszz_mpc = -2.0*expect(sz⊗sz, ρ) - 0.5*expect(sz⊗I, ρ) - 0.5*expect(I⊗sz, ρ) +  Γ[1,2]*(expect(sx⊗sx, ρ) + expect(sy⊗sy, ρ))
    println("$(dsxx_master) <-> $(dsxx_mpc)")
    #@test abs(dsxx_master - dsxx_mpc) < 1e-12
    #@test abs(dsyy_master - dsyy_mpc) < 1e-12
    #@test abs(dszz_master - dszz_mpc) < 1e-12
end

# CollectiveSpins.quantum.timeevolution(T, system, ρ₀; fout=test_mpc_2particles)

N = 3
systemgeometry = CollectiveSpins.geometry.chain(0.3, N)
system = CollectiveSpins.SpinCollection(systemgeometry, edipole; gammas=γ)
basis = CollectiveSpins.quantum.basis(system)
I = identityoperator(spinbasis)

H = CollectiveSpins.quantum.Hamiltonian(system)
Γ, J = CollectiveSpins.quantum.JumpOperators(system)
Jdagger = [dagger(j) for j=J]
Ω = CollectiveSpins.interaction.OmegaMatrix(system)

Ψ₀ = CollectiveSpins.quantum.blochstate(0., pi/2., N)
ρ₀ = Ψ₀ ⊗ dagger(Ψ₀)

function test_mpc_3particles(t, ρ::DenseOpType)
    dρ_h = -1im*(H*ρ - ρ*H)
    dρ_l = DenseOperator(basis)
    for m=1:length(J), n=1:length(J)
       dρ_l += Γ[m,n]*(J[m]*ρ*Jdagger[n] - Jdagger[n]*(J[m]*ρ)/Complex(2) - ρ*Jdagger[n]*J[m]/Complex(2))
    end
    # ========== Test Hamiltonian part ==========
    dsxx_master = expect(sx⊗sx⊗I, dρ_h)
    dsyy_master = expect(sy⊗sy⊗I, dρ_h)
    dszz_master = expect(sz⊗sz⊗I, dρ_h)
    dsxy_master = expect(sx⊗sy⊗I, dρ_h)
    dsxz_master = expect(sx⊗sz⊗I, dρ_h)
    dsyz_master = expect(sy⊗sz⊗I, dρ_h)

    # <sxx>, <syy>, <szz>
    dsxx_mpc = Ω[1,3]*expect(sz⊗sx⊗sy, ρ) + Ω[2,3]*expect(sx⊗sz⊗sy, ρ)
    dsyy_mpc = -Ω[1,3]*expect(sz⊗sy⊗sx, ρ) - Ω[2,3]*expect(sy⊗sz⊗sx, ρ)
    dszz_mpc = Ω[1,3]*(expect(sy⊗sz⊗sx, ρ) - expect(sx⊗sz⊗sy, ρ)) + Ω[2,3]*(expect(sz⊗sy⊗sx, ρ) - expect(sz⊗sx⊗sy, ρ))
    dsxy_mpc = Ω[1,2]*(expect(sz⊗I⊗I, ρ) - expect(I⊗sz⊗I, ρ)) + Ω[1,3]*expect(sz⊗sy⊗sy, ρ) - Ω[2,3]*expect(sx⊗sz⊗sx, ρ)
    dsxz_mpc = Ω[1,2]*expect(I⊗sy⊗I, ρ) + Ω[1,3]*expect(sz⊗sz⊗sy, ρ) + Ω[2,3]*(expect(sx⊗sy⊗sx, ρ) - expect(sx⊗sx⊗sy, ρ))
    dsyz_mpc = -Ω[1,2]*expect(I⊗sx⊗I, ρ) - Ω[1,3]*expect(sz⊗sz⊗sx, ρ) + Ω[2,3]*(expect(sy⊗sy⊗sx, ρ) - expect(sy⊗sx⊗sy, ρ))

    @test abs(dsxx_master - dsxx_mpc) < 1e-12
    @test abs(dsyy_master - dsyy_mpc) < 1e-12
    @test abs(dszz_master - dszz_mpc) < 1e-12
    @test abs(dsxy_master - dsxy_mpc) < 1e-12
    @test abs(dsxz_master - dsxz_mpc) < 1e-12
    @test abs(dsyz_master - dsyz_mpc) < 1e-12

    # # ========== Test Lindblad part ==========
    dsxx_master = expect(sx⊗sx⊗I, dρ_l)
    dsyy_master = expect(sy⊗sy⊗I, dρ_l)
    dszz_master = expect(sz⊗sz⊗I, dρ_l)
    dsxy_master = expect(sx⊗sy⊗I, dρ_l)
    dsxz_master = expect(sx⊗sz⊗I, dρ_l)
    dsyz_master = expect(sy⊗sz⊗I, dρ_l)

    # <sx>, <sy>, <sz>
    dsxx_mpc = -expect(sx⊗sx⊗I, ρ) + Γ[1,2]*(expect(sz⊗sz⊗I, ρ) + 0.5*expect(sz⊗I⊗I, ρ) + 0.5*expect(I⊗sz⊗I, ρ)) + 0.5*Γ[1,3]*expect(sz⊗sx⊗sx, ρ) + 0.5*Γ[2,3]*expect(sx⊗sz⊗sx, ρ)
    dsyy_mpc = -expect(sy⊗sy⊗I, ρ) + Γ[1,2]*(expect(sz⊗sz⊗I, ρ) + 0.5*expect(sz⊗I⊗I, ρ) + 0.5*expect(I⊗sz⊗I, ρ)) + 0.5*Γ[1,3]*expect(sz⊗sy⊗sy, ρ) + 0.5*Γ[2,3]*expect(sy⊗sz⊗sy, ρ)
    dszz_mpc = -2.0*expect(sz⊗sz⊗I, ρ) - expect(sz⊗I⊗I, ρ) - expect(I⊗sz⊗I, ρ) +  Γ[1,2]*(expect(sx⊗sx⊗I, ρ) + expect(sy⊗sy⊗I, ρ)) - 0.5*Γ[1,3]*(expect(sx⊗sz⊗sx, ρ) + expect(sy⊗sz⊗sy, ρ)) - 0.5*Γ[2,3]*(expect(sz⊗sx⊗sx, ρ) + expect(sz⊗sy⊗sy, ρ))
    dsxy_mpc = -expect(sx⊗sy⊗I, ρ) + 0.5*Γ[1,3]*expect(sz⊗sy⊗sx, ρ) + 0.5*Γ[2,3]*expect(sx⊗sz⊗sy, ρ)
    dsxz_mpc = -1.5*expect(sx⊗sz⊗I, ρ) - expect(sx⊗I⊗I, ρ) - Γ[1,2]*(expect(sz⊗sx⊗I, ρ) + 0.5*expect(I⊗sx⊗I, ρ)) + 0.5*Γ[1,3]*expect(sz⊗sz⊗sx, ρ) - 0.5*Γ[2,3]*(expect(sx⊗sx⊗sx, ρ) + expect(sx⊗sy⊗sy, ρ))
    dsyz_mpc = -1.5*expect(sy⊗sz⊗I, ρ) - expect(sy⊗I⊗I, ρ) - Γ[1,2]*(expect(sz⊗sy⊗I, ρ) + 0.5*expect(I⊗sy⊗I, ρ)) + 0.5*Γ[1,3]*expect(sz⊗sz⊗sy, ρ) - 0.5*Γ[2,3]*(expect(sy⊗sx⊗sx, ρ) + expect(sy⊗sy⊗sy, ρ))

    @test abs(dsxx_master - dsxx_mpc) < 1e-12
    @test abs(dsyy_master - dsyy_mpc) < 1e-12
    @test abs(dszz_master - dszz_mpc) < 1e-12
    @test abs(dsxy_master - dsxy_mpc) < 1e-12
    @test abs(dsxz_master - dsxz_mpc) < 1e-12
    @test abs(dsyz_master - dsyz_mpc) < 1e-12
    #println("$(dsxx_master) <-> $(dsxx_mpc)")
end

CollectiveSpins.quantum.timeevolution(T, system, ρ₀; fout=test_mpc_3particles)

end # testset


text:
using Base.Test

d = DateTime("2016-02-19T12:34:56")
datestr = map(string,range(d,Dates.Millisecond(123),5_000_000))

using DateParsing

println("Testing DateParsing.parse")
@time a1 = DateParsing.parse(datestr)
@time a1 = DateParsing.parse(datestr)
@time a1 = DateParsing.parse(datestr)

println("Testing DateParsing.fastparse")
@time a2 = DateParsing.fastparse(datestr)
@time a2 = DateParsing.fastparse(datestr)
@time a2 = DateParsing.fastparse(datestr)

@test a1 == a2

println("Testing Base.DateTime on a 1/100th")
subset = 1:100:length(datestr)
datestr = datestr[subset]

function testDateTime{S<:AbstractString}(ads::Vector{S})
    l = length(ads)
    a = Array(DateTime, l)
    for i in 1:l
        a[i] = DateTime(ads[i])
    end
    a
end

@time a3 = testDateTime(datestr)
@time a3 = testDateTime(datestr)
@time a3 = testDateTime(datestr)

@test a1[subset] == a3


text:
struct Model{model} end
Model(model::Symbol) = Model{model}()
Model(model::String) = Model{Symbol(model)}()

macro Model_str(s)
  return :(Model{$(Expr(:quote, Symbol(s)))})
end

struct Observable{obs} end
Observable(obs::Symbol) = Observable{obs}()
Observable(obs::String) = Observable{Symbol(obs)}()

macro Observable_str(s)
  return :(Observable{$(Expr(:quote, Symbol(s)))})
end

∫(f, a, b) = quadgk(f, a, b)[1]

# Create an infinite sum of Hamiltonian terms
function InfiniteITensorSum(model::Model, s::Vector; kwargs...)
  return InfiniteITensorSum(model, infsiteinds(s); kwargs...)
end

function InfiniteITensorSum(model::Model, s::CelledVector; kwargs...)
  N = length(s)
  H = InfiniteITensorSum(N)
  tensors = [ITensor(model, s[n], s[n + 1]; kwargs...) for n in 1:N]
  return InfiniteITensorSum(tensors)
end

# Version accepting IndexSet
function ITensors.ITensor(model::Model, s1, s2; kwargs...)
  return ITensor(model, Index(s1), Index(s2); kwargs...)
end

function ITensors.ITensor(model::Model, s1::Index, s2::Index; kwargs...)
  n1, n2 = 1, 2
  opsum = OpSum(model, n1, n2; kwargs...)
  return prod(MPO(opsum, [s1, s2]))
end


text:
using DrWatson
@quickactivate "StatReth"

# %%
using Plots
using Distributions

# %%
len = 100
p = range(0, 1, length = len)

# %%
prior = ones(len)
prior ./= sum(prior)

likelihood(w, l) = @. pdf(Binomial(w + l, p), w)

function posterior(prior, likelih)
    post = prior .* likelih
    return post ./ sum(post)
end

posterior_a = posterior(prior, likelihood(3, 0))  # www
posterior_b = posterior(prior, likelihood(3, 1))  # wwwl
posterior_c = posterior(prior, likelihood(5, 2))  # lwwlwww

# %%
plot(xlabel = "p", ylabel = "probability", dpi = 200, legend = :topleft)
plot!(p, posterior_a, label = "a")
plot!(p, posterior_b, label = "b")
plot!(p, posterior_c, label = "c")

# %%
prior = ones(len)
prior[p .< 0.5] .= 0.0
prior ./= sum(prior)

posterior_a = posterior(prior, likelihood(3, 0))  # www
posterior_b = posterior(prior, likelihood(3, 1))  # wwwl
posterior_c = posterior(prior, likelihood(5, 2))  # lwwlwww

# %%
plot(xlabel = "p", ylabel = "probability", dpi = 200, legend = :topleft)
plot!(p, posterior_a, label = "a")
plot!(p, posterior_b, label = "b")
plot!(p, posterior_c, label = "c")


text:
# HierarchicalMatrix of GreensFuns on TensorSpace of HierarchicalSpaces

converttoPiecewiseSpace(H::Space) = H
converttoPiecewiseSpace(H::HierarchicalSpace) = PiecewiseSpace(H)

function partition(ss::ApproxFun.TensorSpace{Tuple{HS1,HS2}}) where {HS1<:HierarchicalSpace,HS2<:HierarchicalSpace}
    ss11,ss12 = partition(factor(ss.space,1))
    ss21,ss22 = partition(factor(ss.space,2))
    (ss11⊗ss21,ss12⊗ss22),(converttoPiecewiseSpace(ss12)⊗converttoPiecewiseSpace(ss21),converttoPiecewiseSpace(ss11)⊗converttoPiecewiseSpace(ss22))
end

function partition(ss::CauchyWeight{O,Tuple{HS1,HS2}}) where {O,HS1<:HierarchicalSpace,HS2<:HierarchicalSpace}
    ss11,ss12 = partition(factor(ss.space,1))
    ss21,ss22 = partition(factor(ss.space,2))
    (CauchyWeight(ss11⊗ss21,O),CauchyWeight(ss12⊗ss22,O)),(CauchyWeight(converttoPiecewiseSpace(ss12)⊗converttoPiecewiseSpace(ss21),O),CauchyWeight(converttoPiecewiseSpace(ss11)⊗converttoPiecewiseSpace(ss22),O))
end

function GreensFun(f::DFunction,ss::AbstractProductSpace{Tuple{HS1,HS2}};method::Symbol=:lowrank,kwds...) where {HS1<:HierarchicalSpace,HS2<:HierarchicalSpace}
    (ss11,ss22),(ss21,ss12) = partition(ss)
    meth1 = method == :Cholesky || method == :lowrank ? :standard : method
    G11 = GreensFun(f,ss11;method=method,kwds...)
    G22 = GreensFun(f,ss22;method=method,kwds...)
    G21 = GreensFun(LowRankFun(f,ss21;method=meth1,kwds...))
    G12 = method == :Cholesky ? transpose(G21) : GreensFun(LowRankFun(f,ss12;method=meth1,kwds...))
    return HierarchicalMatrix((G11,G22),(G21,G12))
end

blocksize(H::HierarchicalMatrix{F,G}) where {F<:GreensFun,G<:GreensFun} = map(ncomponents,domain(H).domains)

function domain(H::HierarchicalMatrix{F,G}) where {F<:GreensFun,G<:GreensFun}
    H11,H22 = diagonaldata(H)
    H21,H12 = offdiagonaldata(H)
    m1,n2 = factor(domain(H12),1),factor(domain(H12),2)
    m2,n1 = factor(domain(H21),1),factor(domain(H21),2)
    @assert (m1,n1) == (factor(domain(H11),1),factor(domain(H11),2))
    @assert (m2,n2) == (factor(domain(H22),1),factor(domain(H22),2))
    (m1∪m2)*(n1∪n2)
end

for TYP in (:(ApproxFun.DefiniteLineIntegralWrapper),:DefiniteLineIntegral)
    @eval function Base.getindex(⨍::$TYP,H::HierarchicalMatrix{G,GreensFun{L,T}}) where {G<:GreensFun,L<:LowRankFun,T}
        H11,H22 = diagonaldata(H)
        wsp = components(domainspace(⨍))
        if ncomponents(factor(domain(H11),2)) ≥ 2
            ⨍1 = DefiniteLineIntegral(PiecewiseSpace(wsp[1:ncomponents(factor(domain(H11),2))]))
        else
            ⨍1 = DefiniteLineIntegral(wsp[1])
        end
        if ncomponents(factor(domain(H22),2)) ≥ 2
            ⨍2 = DefiniteLineIntegral(PiecewiseSpace(wsp[end-ncomponents(factor(domain(H22),2))+1:end]))
        else
            ⨍2 = DefiniteLineIntegral(wsp[end])
        end
        HierarchicalOperator((qr(⨍1[H11]),qr(⨍2[H22])),
                             map(LowRankIntegralOperator,offdiagonaldata(H)))
    end
end

qr(H::HierarchicalOperator) = H # trivial no-op for now.


text:
macro debug_info(expression)
    return :($expression)
end


text:
function vtk_save(vtk::DatasetFile)
    if isopen(vtk)
        if vtk.appended
            save_with_appended_data(vtk)
        else
            save_file(vtk.xdoc, vtk.path)
        end
        @assert isopen(vtk)
        close(vtk)
    end
    return [vtk.path] :: Vector{String}
end

"""
Write VTK XML file containing appended binary data to disk.

In this case, the XML file is written manually instead of using the `save_file`
function of `LightXML`, which doesn't allow to write raw binary data.
"""
function save_with_appended_data(vtk::DatasetFile)
    @assert vtk.appended
    @assert isopen(vtk.buf)

    # Convert XML document to a string, and split the last two lines.
    lines = rsplit(string(vtk.xdoc), '\n', limit=3, keepempty=true)

    # Verify that the last two lines are what they're supposed to be.
    @assert lines[2] == "</VTKFile>"
    @assert lines[3] == ""

    open(vtk.path, "w") do io
        # Write everything but the last two lines.
        write(io, lines[1])
        write(io, "\n")

        # Write raw data (contents of buffer vtk.buf).
        # An underscore "_" is needed before writing appended data.
        write(io, "  <AppendedData encoding=\"raw\">")
        write(io, "\n_")
        write(io, take!(vtk.buf))
        write(io, "\n  </AppendedData>")
        write(io, "\n</VTKFile>")

        close(vtk.buf)
    end

    nothing
end


text:
using AdvancedTopics

include("util.jl")

function perform_tests(training_samples::T₁,
                       training_targets::T₂,
                       testing_samples::T₃,
                       testing_targets::T₄,
                       n_neurons::Int;
                       boundary_offset::Int = 100) where {T₁ <: AbstractMatrix,
                                                          T₂ <: AbstractVector,
                                                          T₃ <: AbstractMatrix,
                                                          T₄ <: AbstractVector}
      X₀ = training_samples
      t₀ = training_targets
      X₁ = testing_samples
      t₁ = testing_targets
      L = n_neurons

      W, fs = train_kde_comparators(Float64, X₀, t₀, L, boundary_offset = boundary_offset)

      elm_sig = ELM{Float64}(X₀, t₀, W, [Sigmoid() for f in 1:length(fs)])
      elm_kde = ELM{Float64}(X₀, t₀, W, fs)

      y₀_sig = predict(elm_sig, X₀)
      y₀_kde = predict(elm_kde, X₀)
      y₁_sig = predict(elm_sig, X₁)
      y₁_kde = predict(elm_kde, X₁)

      plt_sig = plot_decision_boundary(elm_sig, X₀, t₀)
      plt_kde = plot_decision_boundary(elm_kde, X₀, t₀)

      training_error_sig = calculate_error(y₀_sig, t₀)
      training_error_kde = calculate_error(y₀_kde, t₀)
      testing_error_sig = calculate_error(y₁_sig, t₁)
      testing_error_kde = calculate_error(y₁_kde, t₁)

      (plt_sig, training_error_sig, testing_error_sig,
       plt_kde, training_error_kde, testing_error_kde)
end

for n_neurons in 1:2:5
    N = 200
    radius₁ = 4
    radius₂ = 11
    σ = 1.2

    X₀, t₀ = synthetic_circular_data(N, radius₁, radius₂, σ)
    X₁, t₁ = synthetic_circular_data(N, radius₁, radius₂, σ)

    (plt_sig, training_error_sig, testing_error_sig,
     plt_kde, training_error_kde, testing_error_kde) = perform_tests(X₀, t₀, X₁, t₁, n_neurons)

    savefig(plt_sig, "figures/db_circular_sig_$(n_neurons)_neurons.png")
    savefig(plt_kde, "figures/db_circular_kde_$(n_neurons)_neurons.png")
end


text:
import Distributions

function laplace_mechanism(result::Number; sensitivity::Number, epsilon::Number)
    @assert sensitivity > 0
    @assert epsilon > 0
    lap = Distributions.Laplace(0.0, sensitivity/epsilon) # mean 0, scale=sensitivity/epsilon
    result + rand(lap)
end

function laplace_mechanism(result::Array; sensitivity::Number, epsilon::Number)
    [laplace_mechanism(x, sensitivity=sensitivity, epsilon=epsilon) for x in result]
end

function geometric_mechanism(result::Int; sensitivity::Int, epsilon::Number)
    @assert sensitivity > 0
    @assert epsilon > 0
    p = 1-exp(-epsilon/sensitivity)
    geo = Distributions.Geometric(p)
    result + rand(geo) - rand(geo)
end

function noisy_max(answers::Vector{S}; sensitivity::Number, epsilon::Number) where S <: Number
    noisy_answers = [laplace_mechanism(x, sensitivity=2*sensitivity, epsilon=epsilon) for x in answers]
    return argmax(noisy_answers)
end

function sparse_vector(;T::Number, N::Int, answers::Vector{S}, sensitivity::Number, epsilon::Number) where S <: Number
    noisy_T = laplace_mechanism(T; sensitivity=sensitivity, epsilon=epsilon/2)
    answered = 0
    i=1
    output = Bool[] # empty array of booleans
    while(answered < N && i <= length(answers))
        noisy_query = laplace_mechanism(answers[i]; sensitivity=2*N, epsilon=epsilon/2)
	    i = i + 1
	    if noisy_query >= noisy_T
	       append!(output, true)
	       answered = answered + 1
	    else
           append!(output, false)
	    end
    end
    output
end

max_stars_repo_name: dkifer/dpcourse

max_issues_repo_name: dkifer/dpcourse

max_forks_repo_name: dkifer/dpcourse