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

using Markdown
using InteractiveUtils

# ╔═╡ acf67f82-1a85-11eb-2bdf-abc05bfb68eb
begin
	using DataFrames
	using Statistics
	using PyPlot
	using GLM
end

# ╔═╡ 8061b6ba-1a85-11eb-1a73-130f152eb8d8
VERSION

# ╔═╡ ad304f82-1a85-11eb-1b1e-bba3c3d49c45
isfile.(["Project.toml", "Manifest.toml"])

# ╔═╡ acdb9e6a-1a85-11eb-2e56-b91f3f2c3e24
aq = [10.0   8.04  10.0  9.14  10.0   7.46   8.0   6.58
       8.0   6.95   8.0  8.14   8.0   6.77   8.0   5.76
      13.0   7.58  13.0  8.74  13.0  12.74   8.0   7.71
       9.0   8.81   9.0  8.77   9.0   7.11   8.0   8.84
      11.0   8.33  11.0  9.26  11.0   7.81   8.0   8.47
      14.0   9.96  14.0  8.1   14.0   8.84   8.0   7.04
       6.0   7.24   6.0  6.13   6.0   6.08   8.0   5.25
       4.0   4.26   4.0  3.1    4.0   5.39  19.0  12.50
      12.0  10.84  12.0  9.13  12.0   8.15   8.0   5.56
       7.0   4.82   7.0  7.26   7.0   6.42   8.0   7.91
       5.0   5.68   5.0  4.74   5.0   5.73   8.0   6.89];

# ╔═╡ acbde3e8-1a85-11eb-2d9f-b7754fbf0281
df = DataFrame(aq)

# ╔═╡ aca17b68-1a85-11eb-39fd-bf8a9cfe42bb
newnames = vec(string.(["x", "y"], [1 2 3 4]))

# ╔═╡ ac6908d2-1a85-11eb-0733-63fcf56fecd3
rename!(df, newnames)

# ╔═╡ ac4c9b20-1a85-11eb-22d4-cb4290f63002
DataFrame(aq, [:x1, :y1, :x2, :y2, :x3, :y3, :x4, :y4])

# ╔═╡ ac31974e-1a85-11eb-1d08-f1198ac82594
df.x1

# ╔═╡ ac168440-1a85-11eb-3b9f-3595b29077cc
df."y1"

# ╔═╡ abf8d044-1a85-11eb-1df4-3f834e5990b9
df[:, :y1]

# ╔═╡ abdc5e8c-1a85-11eb-3b39-b500d21dd246
df[:, "y1"]

# ╔═╡ aa0e0f54-1ac8-11eb-1085-93fe0cedd3c0
select(df, r"x", :)

# ╔═╡ dcaa992a-1acb-11eb-0015-cdb2cfe62b8d
df

# ╔═╡ a9ee21f8-1ac8-11eb-0555-377010925b9a
select(df, r"x")

# ╔═╡ a9b716d6-1ac8-11eb-0ded-bdad1bea310d
describe(df, :mean => mean, :std => std)

# ╔═╡ a99835e0-1ac8-11eb-2a27-0157951b99fd
describe(df, "mean" => mean, "std" => std)

# ╔═╡ a97963ea-1ac8-11eb-0ee7-7bbf209a59fd
describe(df)

# ╔═╡ a95f9456-1ac8-11eb-14b4-c5aa281000bf
df.id = 1:nrow(df)

# ╔═╡ a91faa3a-1ac8-11eb-0d22-31a5e566a332
nrow(df)

# ╔═╡ a901e5fc-1ac8-11eb-1eb3-a11857b4bdc0
ncol(df)

# ╔═╡ a8e4981e-1ac8-11eb-0c16-07fe13d9de24
select(df, "id", :)

# ╔═╡ a8c5f24c-1ac8-11eb-387a-07909ec37e8a
df

# ╔═╡ a8a989c2-1ac8-11eb-2512-799ec5c564bf
select!(df, "id", :)

# ╔═╡ 51020740-1aca-11eb-0c90-57bad44609d8
xlim = collect(extrema(Matrix(select(df, r"x"))) .+ (-1, 1))

# ╔═╡ a84fad9c-1ac8-11eb-2492-1f456aa4f48b
ylim = collect(extrema(Matrix(select(df, r"y"))) .+ (-1, 1))

# ╔═╡ a8120980-1ac8-11eb-2b9a-69b28da0cc07
collect(10:10:100)

# ╔═╡ 853a316c-1ad2-11eb-30fa-dbc9f5e61c71
begin
	fig, axs = plt.subplots(2, 2)
	fig.tight_layout(pad=4.0)
	for i in 1:4
		x = Symbol("x", i)
		y = Symbol("y", i)
		model = lm(term(y)~term(x), df)
		axs[i].plot(xlim, predict(model, DataFrame(x => xlim)), color="orange")
		axs[i].scatter(df[:, x], df[:, y])
		axs[i].set_xlim(xlim)
		axs[i].set_ylim(ylim)
		axs[i].set_xlabel("x$i")
		axs[i].set_ylabel("y$i")
		a, b = round.(coef(model), digits=2)
		c = round(100 * r2(model), digits=2)
		axs[i].set_title(string("R²=$c%, $y=$a+$b$x"))
	end
	fig
end

# ╔═╡ 85848898-1ad2-11eb-1fa5-713858ae4565
begin
	x = :var1
	y = :var2

	xc = 1:3
	yc = 4:6

	DataFrame(x => xc, y => yc)

end

# ╔═╡ 85a6e21a-1ad2-11eb-00a4-cb349b78caec
DataFrame(var1=xc, var2=yc)

# ╔═╡ 85dc1928-1ad2-11eb-242b-2917e65f2244
df.x1

# ╔═╡ 85f131a0-1ad2-11eb-23ad-df576e2e6b7a

# ╔═╡ 8609a440-1ad2-11eb-0696-09f4c5a8fc0f

# ╔═╡ 8623433e-1ad2-11eb-2589-79e06429f720

# ╔═╡ a7f25360-1ac8-11eb-1c1f-552f8c1a9d8a

# ╔═╡ a7d388c2-1ac8-11eb-1bd7-59a0654fd402

# ╔═╡ a79492ea-1ac8-11eb-098c-f12b9d38f313

# ╔═╡ a76f13ea-1ac8-11eb-3800-b1c27e74a466

# ╔═╡ a76f037a-1ac8-11eb-199c-cb332db6cca1

# ╔═╡ a7282752-1ac8-11eb-0f05-83cace66d3b0

# ╔═╡ ab7c9592-1a85-11eb-342f-3dab4c1d54a6

# ╔═╡ Cell order:
# ╠═8061b6ba-1a85-11eb-1a73-130f152eb8d8
# ╠═ad304f82-1a85-11eb-1b1e-bba3c3d49c45
# ╠═acf67f82-1a85-11eb-2bdf-abc05bfb68eb
# ╠═acdb9e6a-1a85-11eb-2e56-b91f3f2c3e24
# ╟─acbde3e8-1a85-11eb-2d9f-b7754fbf0281
# ╟─aca17b68-1a85-11eb-39fd-bf8a9cfe42bb
# ╠═ac6908d2-1a85-11eb-0733-63fcf56fecd3
# ╠═ac4c9b20-1a85-11eb-22d4-cb4290f63002
# ╠═ac31974e-1a85-11eb-1d08-f1198ac82594
# ╠═ac168440-1a85-11eb-3b9f-3595b29077cc
# ╠═abf8d044-1a85-11eb-1df4-3f834e5990b9
# ╠═abdc5e8c-1a85-11eb-3b39-b500d21dd246
# ╠═aa0e0f54-1ac8-11eb-1085-93fe0cedd3c0
# ╠═dcaa992a-1acb-11eb-0015-cdb2cfe62b8d
# ╠═a9ee21f8-1ac8-11eb-0555-377010925b9a
# ╠═a9b716d6-1ac8-11eb-0ded-bdad1bea310d
# ╠═a99835e0-1ac8-11eb-2a27-0157951b99fd
# ╠═a97963ea-1ac8-11eb-0ee7-7bbf209a59fd
# ╠═a95f9456-1ac8-11eb-14b4-c5aa281000bf
# ╠═a91faa3a-1ac8-11eb-0d22-31a5e566a332
# ╠═a901e5fc-1ac8-11eb-1eb3-a11857b4bdc0
# ╠═a8e4981e-1ac8-11eb-0c16-07fe13d9de24
# ╠═a8c5f24c-1ac8-11eb-387a-07909ec37e8a
# ╠═a8a989c2-1ac8-11eb-2512-799ec5c564bf
# ╠═51020740-1aca-11eb-0c90-57bad44609d8
# ╠═a84fad9c-1ac8-11eb-2492-1f456aa4f48b
# ╠═a8120980-1ac8-11eb-2b9a-69b28da0cc07
# ╠═853a316c-1ad2-11eb-30fa-dbc9f5e61c71
# ╠═85848898-1ad2-11eb-1fa5-713858ae4565
# ╠═85a6e21a-1ad2-11eb-00a4-cb349b78caec
# ╠═85dc1928-1ad2-11eb-242b-2917e65f2244
# ╠═85f131a0-1ad2-11eb-23ad-df576e2e6b7a
# ╠═8609a440-1ad2-11eb-0696-09f4c5a8fc0f
# ╠═8623433e-1ad2-11eb-2589-79e06429f720
# ╠═a7f25360-1ac8-11eb-1c1f-552f8c1a9d8a
# ╠═a7d388c2-1ac8-11eb-1bd7-59a0654fd402
# ╠═a79492ea-1ac8-11eb-098c-f12b9d38f313
# ╠═a76f13ea-1ac8-11eb-3800-b1c27e74a466
# ╠═a76f037a-1ac8-11eb-199c-cb332db6cca1
# ╠═a7282752-1ac8-11eb-0f05-83cace66d3b0
# ╠═ab7c9592-1a85-11eb-342f-3dab4c1d54a6

max_stars_repo_name: martin-fabbri/colab-notebooks

max_issues_repo_name: martin-fabbri/colab-notebooks

max_forks_repo_name: martin-fabbri/colab-notebooks


text:
const BTGdict = GenericIncidenceList{ExVertex,Edge{ExVertex},Array{ExVertex,1},Array{Array{Edge{ExVertex},1},1}}

# BayesTree declarations
"""
$(TYPEDEF)

Data structure for the Bayes (Junction) tree, which is used for inference and constructed from a given `::FactorGraph`.
"""
mutable struct BayesTree
  bt::BTGdict
  btid::Int
  cliques::Dict{Int,TreeClique}
  frontals::Dict{Symbol,Int}
  variableOrder::Vector{Symbol}
  buildTime::Float64
end

function emptyBayesTree()
    bt =   BayesTree(Graphs.inclist(TreeClique,is_directed=true),
                     0,
                     Dict{Int,TreeClique}(),
                     #[],
                     Dict{AbstractString, Int}(),
                     Symbol[],
					 0.0 )
    return bt
end

"""
    $TYPEDEF

Container for upward tree solve / initialization.

TODO
- remove proceed
- more direct clique access (cliq, parent, children), for multi-process solves
"""
mutable struct CliqStateMachineContainer{BTND, T <: AbstractDFG, InMemG <: InMemoryDFGTypes}
  dfg::T
  cliqSubFg::InMemG
  tree::BayesTree
  cliq::TreeClique
  parentCliq::Vector{TreeClique}
  childCliqs::Vector{TreeClique}
  forceproceed::Bool # TODO: bad flag that must be removed by refactoring sm
  incremental::Bool
  drawtree::Bool
  dodownsolve::Bool
  delay::Bool
  opts::SolverParams
  refactoring::Dict{Symbol, String}
  oldcliqdata::BTND
  logger::SimpleLogger
  CliqStateMachineContainer{BTND}() where {BTND} = new{BTND, InMemDFGType, InMemDFGType}() # NOTE JT - GraphsDFG as default?
  CliqStateMachineContainer{BTND}(x1::G,
                                  x2::InMemoryDFGTypes,
                                  x3::BayesTree,
                                  x4::TreeClique,
                                  x5::Vector{TreeClique},
                                  x6::Vector{TreeClique},
                                  x7::Bool,
                                  x8::Bool,
                                  x9::Bool,
                                  x10a::Bool,
                                  x10aa::Bool,
                                  x10aaa::SolverParams,
								  x10b::Dict{Symbol,String}=Dict{Symbol,String}(),
                                  x11::BTND=emptyBTNodeData(),
                                  x13::SimpleLogger=SimpleLogger(Base.stdout) ) where {BTND, G <: AbstractDFG} = new{BTND, G, typeof(x2)}(x1,x2,x3,x4,x5,x6,x7,x8,x9,x10a,x10aa,x10aaa,x10b,x11, x13)
end

function CliqStateMachineContainer(x1::G,
                                   x2::InMemoryDFGTypes,
                                   x3::BayesTree,
                                   x4::TreeClique,
                                   x5::Vector{TreeClique},
                                   x6::Vector{TreeClique},
                                   x7::Bool,
                                   x8::Bool,
                                   x9::Bool,
                                   x10::Bool,
                                   x10aa::Bool,
                                   x10aaa::SolverParams,
                                   x11::BTND=emptyBTNodeData(),
                                   x13::SimpleLogger=SimpleLogger(Base.stdout) ) where {BTND, G <: AbstractDFG}
  #
  CliqStateMachineContainer{BTND}(x1,x2,x3,x4,x5,x6,x7,x8,x9,x10,x10aa,x10aaa,Dict{Symbol,String}(),x11,x13)
end

const CSMHistory = Vector{Tuple{DateTime, Int, Function, CliqStateMachineContainer}}

"""
$(TYPEDEF)

Data structure for each clique in the Bayes (Junction) tree.
"""
mutable struct BayesTreeNodeData
  frontalIDs::Vector{Symbol}
  separatorIDs::Vector{Symbol}
  inmsgIDs::Vector{Symbol} # Int
  potIDs::Vector{Symbol} # Int # this is likely redundant TODO -- remove
  potentials::Vector{Symbol}
  partialpotential::Vector{Bool}

  dwnPotentials::Vector{Symbol}
  dwnPartialPotential::Vector{Bool}

  cliqAssocMat::Array{Bool,2}
  cliqMsgMat::Array{Bool,2}
  directvarIDs::Vector{Symbol} # Int
  directFrtlMsgIDs::Vector{Symbol} # Int
  msgskipIDs::Vector{Symbol} # Int
  itervarIDs::Vector{Symbol} # Int
  directPriorMsgIDs::Vector{Symbol} # Int
  debug
  debugDwn

  # future might concentrate these four fields down to two
  # these should become specialized BeliefMessage type
  upMsg::TempBeliefMsg # Dict{Symbol, BallTreeDensity}
  dwnMsg::TempBeliefMsg # Dict{Symbol, BallTreeDensity}
  upInitMsgs::Dict{Int, TempBeliefMsg}
  downInitMsg::TempBeliefMsg

  allmarginalized::Bool
  initialized::Symbol
  upsolved::Bool
  downsolved::Bool
  initUpChannel::Channel{Symbol}
  initDownChannel::Channel{Symbol}
  solveCondition::Condition
  lockUpStatus::Channel{Int}
  lockDwnStatus::Channel{Int}
  solvableDims::Channel{Dict{Symbol, Float64}}
  statehistory::Vector{Tuple{DateTime, Int, Function, CliqStateMachineContainer}}
  BayesTreeNodeData() = new()
  BayesTreeNodeData(x...) = new(x[1],x[2],x[3],x[4],x[5],x[6],x[7],x[8],x[9],x[10],
                                x[11],x[12],x[13],x[14],x[15],x[16],x[17],x[18],x[19],x[20],
                                x[21], x[22], x[23], x[24], x[25], x[26], x[27], x[28], x[29], x[30], x[31],
                                Vector{Tuple{DateTime, Int, Function, CliqStateMachineContainer}}() )
end

# TODO -- this should be a constructor
function emptyBTNodeData()
  BayesTreeNodeData(Symbol[],Symbol[],Symbol[],
                    Symbol[],Symbol[],Bool[], # 6
                    Symbol[],Bool[],
                    Array{Bool}(undef, 0,0),
                    Array{Bool}(undef, 0,0),
                    Int[],Int[],             # 10+2
                    Int[],Int[],Int[],       # 13+2
                    nothing, nothing,        # 15+2
                    Dict{Symbol, BallTreeDensity}(),  # :null => AMP.manikde!(zeros(1,1), [1.0;], (:Euclid,))),
                    Dict{Symbol, BallTreeDensity}(),  # :null => AMP.manikde!(zeros(1,1), [1.0;], (:Euclid,))),
                    Dict{Int, TempBeliefMsg}(),
                    TempBeliefMsg(),         # 19+2
                    false, :null,
                    false, false,            # 23+2
                    Channel{Symbol}(1), Channel{Symbol}(1), Condition(), # 26+2
                    Channel{Int}(1), Channel{Int}(1),
                    Channel{Dict{Symbol,Float64}}(1) )
end

"""
$(TYPEDEF)
"""
mutable struct PotProd
    Xi::Symbol # Int
    prev::Array{Float64,2}
    product::Array{Float64,2}
    potentials::Array{BallTreeDensity,1}
    potentialfac::Vector{Symbol}
end
"""
$(TYPEDEF)
"""
mutable struct CliqGibbsMC
    prods::Array{PotProd,1}
    lbls::Vector{Symbol}
    CliqGibbsMC() = new()
    CliqGibbsMC(a,b) = new(a,b)
end
"""
$(TYPEDEF)
"""
mutable struct DebugCliqMCMC
  mcmc::Union{Nothing, Array{CliqGibbsMC,1}}
  outmsg::NBPMessage
  outmsglbls::Dict{Symbol, Symbol} # Int
  priorprods::Vector{CliqGibbsMC}
  DebugCliqMCMC() = new()
  DebugCliqMCMC(a,b,c,d) = new(a,b,c,d)
end

"""
$(TYPEDEF)
"""
mutable struct UpReturnBPType
  upMsgs::NBPMessage
  dbgUp::DebugCliqMCMC
  IDvals::Dict{Symbol, EasyMessage}
  keepupmsgs::TempBeliefMsg # Dict{Symbol, BallTreeDensity} # TODO Why separate upMsgs?
  totalsolve::Bool
  UpReturnBPType() = new()
  UpReturnBPType(x1,x2,x3,x4,x5) = new(x1,x2,x3,x4,x5)
end

"""
$(TYPEDEF)

TODO refactor msgs into only a single variable
"""
mutable struct DownReturnBPType
  dwnMsg::NBPMessage
  dbgDwn::DebugCliqMCMC
  IDvals::Dict{Symbol,EasyMessage} # Int
  keepdwnmsgs::TempBeliefMsg # Dict{Symbol, BallTreeDensity}
end

"""
$(TYPEDEF)
"""
mutable struct FullExploreTreeType{T, T2, T3 <:InMemoryDFGTypes}
  fg::T3
  bt::T2
  cliq::TreeClique
  prnt::T
  sendmsgs::Vector{NBPMessage}
end

const ExploreTreeType{T} = FullExploreTreeType{T, BayesTree}
const ExploreTreeTypeLight{T} = FullExploreTreeType{T, Nothing}

function ExploreTreeType(fgl::G,
                         btl::BayesTree,
                         vertl::TreeClique,
                         prt::T,
                         msgs::Array{NBPMessage,1} ) where {G <: AbstractDFG, T}
  #
  ExploreTreeType{T}(fgl, btl, vertl, prt, msgs)
end

"""
$(TYPEDEF)
"""
mutable struct MsgPassType
  fg::GraphsDFG
  cliq::TreeClique
  vid::Symbol # Int
  msgs::Array{NBPMessage,1}
  N::Int
end

mutable struct PackedBayesTreeNodeData
  frontalIDs::Vector{Symbol}
  separatorIDs::Vector{Symbol}
  inmsgIDs::Vector{Symbol} # Int
  potIDs::Vector{Symbol} # Int # this is likely redundant TODO -- remove
  potentials::Vector{Symbol}
  partialpotential::Vector{Bool}
  dwnPotentials::Vector{Symbol}
  dwnPartialPotential::Vector{Bool}
  cliqAssocMat::Array{Bool,2}
  cliqMsgMat::Array{Bool,2}
  directvarIDs::Vector{Symbol} # Int
  directFrtlMsgIDs::Vector{Symbol} # Int
  msgskipIDs::Vector{Symbol} # Int
  itervarIDs::Vector{Symbol} # Int
  directPriorMsgIDs::Vector{Symbol} # Int
end

function convert(::Type{PackedBayesTreeNodeData}, btnd::BayesTreeNodeData)
  return PackedBayesTreeNodeData(
    btnd.frontalIDs,
    btnd.separatorIDs,
    btnd.inmsgIDs,
    btnd.potIDs,
    btnd.potentials,
    btnd.partialpotential,
    btnd.dwnPotentials,
    btnd.dwnPartialPotential,
    btnd.cliqAssocMat,
    btnd.cliqMsgMat,
    btnd.directvarIDs,
    btnd.directFrtlMsgIDs,
    btnd.msgskipIDs,
    btnd.itervarIDs,
    btnd.directPriorMsgIDs  )
end

function convert(::Type{BayesTreeNodeData}, pbtnd::PackedBayesTreeNodeData)
  btnd = emptyBTNodeData()
    btnd.frontalIDs = pbtnd.frontalIDs
    btnd.separatorIDs = pbtnd.separatorIDs
    btnd.inmsgIDs = pbtnd.inmsgIDs
    btnd.potIDs = pbtnd.potIDs
    btnd.potentials = pbtnd.potentials
    btnd.partialpotential = pbtnd.partialpotential
    btnd.dwnPotentials = pbtnd.dwnPotentials
    btnd.dwnPartialPotential = pbtnd.dwnPartialPotential
    btnd.cliqAssocMat = pbtnd.cliqAssocMat
    btnd.cliqMsgMat = pbtnd.cliqMsgMat
    btnd.directvarIDs = pbtnd.directvarIDs
    btnd.directFrtlMsgIDs = pbtnd.directFrtlMsgIDs
    btnd.msgskipIDs = pbtnd.msgskipIDs
    btnd.itervarIDs = pbtnd.itervarIDs
    btnd.directPriorMsgIDs = pbtnd.directPriorMsgIDs
  return btnd
end

ext: jl


text:
@doc raw"""
    oneHot(Y; classes = [], numC = 0)

convert array of integer classes into one Hot coding.

# Arguments

- `Y` := a vector of classes as a number

- `classes` := the classes explicity represented (in case not all the classes are present in the labels given)

- `numC` := number of classes as alternative to `classes` variable

# Examples

```julia
Y = rand(0:9, 100); # a 100 item with class of [0-9]
"""
function oneHot(Y; classes = [], numC = 0)
    if numC > 0
        c = numC
        Cs = sort(unique(Y))
    elseif length(classes) > 0
        Cs = sort(classes)
        c = length(Cs)
    else
        Cs = sort(unique(Y))
        c = length(Cs)
    end
    hotY = BitArray{2}(undef, c, length(Y))
    @simd for i=1:length(Y)
        hotY[:,i] .= (Cs .== Y[i])
    end
    return hotY
end #oneHot

export oneHot

@doc raw"""
    resetCount!(outLayer::Layer, cnt::Symbol)

to reset a counter in all layers under `outLayer`.

# Arguments

- `outLayer::Layer` := the layer from start reseting the counter

- `cnt::Symbol` := the counter to be reseted

# Examples

```julia
X_train = rand(128, 100);

X_Input = Input(X_train);
X = FCLayer(50, :relu)(X_Input);
X_out = FCLayer(10, :softmax)(X);

FCache = chainForProp(X_train, X_Input);

# Now to reset the forwCount in all layers

resetCount!(X_out, :forwCount)
```

"""
function resetCount!(outLayer::Layer,
                     cnt::Symbol)
    prevLayer = outLayer.prevLayer
    if outLayer isa Input
        # if outLayer.forwCount != 0
            eval(:($outLayer.$cnt = 0))
        # end #if outLayer.forwCount != 0
    elseif isa(outLayer, MILayer) #if prevLayer == nothing
        for prevLayer in outLayer.prevLayer
            resetCount!(prevLayer, cnt)
        end #for
        eval(:($outLayer.$cnt = 0))
    else #if prevLayer == nothing
        resetCount!(prevLayer, cnt)
        eval(:($outLayer.$cnt = 0))
    end #if prevLayer == nothing
    return nothing
end #function resetForwCount

export resetCount!

### to extend the getindex fun

@doc raw"""
    getindex(it, key; default) = haskey(it, key) ? it[key] : default

# Examples

```julia
D = Dict(:A=>"A", :B=>"B")

A = getindex(D, :A)

## this will return an error
#C = getindex(D: :C)

#instead
C = getindex(D, :C; default="C")
#this will return the `String` C
```
"""
Base.getindex(it, key; default) = haskey(it, key) ? it[key] : default

#### getLayerSlice

export getLayerSlice

function getLayerSlice(cLayer::Layer, nextLayer::ConcatLayer, BCache::Dict{Layer, Dict{Symbol, AbstractArray}})
    N = ndims(BCache[nextLayer][:dA])
    fAx = axes(BCache[nextLayer][:dA])[1:end-2]
    lAx = axes(BCache[nextLayer][:dA])[end]
    LSlice = nextLayer.LSlice[cLayer]
    return BCache[nextLayer][:dA][fAx...,LSlice,lAx]
end #function getLayerSlice(cLayer::Layer, nextLayer::ConcatLayer

"""
    getLayerSlice(cLayer::Layer, nextLayer::Layer, BCache::Dict{Layer, Dict{Symbol, AbstractArray}})

Fall back method for  `Layer`s other than `ConcatLayer`

"""
function getLayerSlice(cLayer::Layer, nextLayer::Layer, BCache::Dict{Layer, Dict{Symbol, AbstractArray}})
    return BCache[nextLayer][:dA]
end #function getLayerSlice(cLayer::Layer, nextLayer::Layer

###

export chParamType!

function chParamType!(cLayer::Layer, T::DataType, cnt::Integer = -1)
    if cnt < 0
        cnt = cLayer.updateCount + 1
    end

    if length(cLayer.nextLayers) == 0
        if cLayer.updateCount < cnt
            try
                cLayer.W = T.(cLayer.W)
                cLayer.B = T.(cLayer.B)
                cLayer.updateCount += 1
            catch
                cLayer.updateCount += 1
            end
        end #if cLayer.forwCount < cnt
        return nothing
    elseif isa(cLayer, MILayer) #if typeof(cLayer)==AddLayer
        if all(
            i -> (i.updateCount == cLayer.prevLayer[1].updateCount),
            cLayer.prevLayer,
        )
            try
                cLayer.W = T.(cLayer.W)
                cLayer.B = T.(cLayer.B)
                cLayer.updateCount += 1
            catch
                cLayer.updateCount += 1
            end
            for nextLayer in cLayer.nextLayers
                chParamType!(nextLayer, T, cnt)
            end
        end #if all

        return nothing
    else #if cLayer.prevLayer==nothing
        if cLayer.updateCount < cnt
            if cLayer isa Input
                try
                    cLayer.W = T.(cLayer.W)
                    cLayer.B = T.(cLayer.B)
                    cLayer.updateCount += 1
                catch
                    cLayer.updateCount += 1
                end
            else
                try
                    cLayer.W = T.(cLayer.W)
                    cLayer.B = T.(cLayer.B)
                    cLayer.updateCount += 1
                catch
                    cLayer.updateCount += 1
                end
            end
            for nextLayer in cLayer.nextLayers
                chParamType!(nextLayer, T, cnt)
            end

        end #if cLayer.forwCount < cnt
        return nothing
    end #if cLayer.prevLayer!=nothing

    return nothing

end #function chParamType(cLayer::Layer, T::DataType)


text:
using DrWatson
@quickactivate
using Test
using HierarchicalAD

include("experiments.jl")

max_stars_repo_name: vitskvara/HierarchicalAD

max_issues_repo_name: vitskvara/HierarchicalAD

max_forks_repo_name: vitskvara/HierarchicalAD


text:
### Processing dataset: INPUT is path_to_csv_data and OUTPUT is datamatrix

function createdatamatrix(path_to_csv_data, number_zones)
    rawdata = CSV.read(path_to_csv_data);
    rawdata = convert(Matrix, rawdata[:,1:5]);
    size_rawdata, placeholder_columns = size(rawdata);
    datamatrix= zeros(number_zones, number_zones, T, 2);
    for i=1:size_rawdata
        if(rawdata[i,1] == 0)
            rawdata[i,1] = number_zones; # in case a zone has an ID of "0", we give it a new zone ID.
        end
        if(rawdata[i,2] == 0)
            rawdata[i,2] = number_zones; # in case a zone has an ID of "0", we give it a new zone ID.
        end
        if(rawdata[i,3] == 0)
            rawdata[i,3] = 24; # Midnight is denoted by a value of zero. We replace it with 24 to get matrix indexes right.
        end
        index1 = convert(Int64, rawdata[i,1]);
        index2 = convert(Int64, rawdata[i,2]);
        index3 = convert(Int64, rawdata[i,3]);
        datamatrix[index1,index2,index3,1] = rawdata[i,4]; # mean_travel_time goes here
        datamatrix[index1,index2,index3,2] = rawdata[i,5]; # standard_deviation_travel_time goes here
    end

    return datamatrix

end

max_stars_repo_name: ArsamAryandoust/CarParkingMaps

max_issues_repo_name: ArsamAryandoust/CarParkingMaps

max_forks_repo_name: ArsamAryandoust/CarParkingMaps


text:
"""
    conf_enc(input::AbstractString; maxn = 7, spins = 2,
    notation = Dict("s" => 0, "p" => 1, "d" => 2, "f" => 3))

Parse an input string as the configuration of an atom. An electronic
configuration is string like "[He] 2s2 2p1", where the notation [X] indicates
that all subshells associated with the noble gas X are fully occupied.
The function encodes the configuration provided "[He] 2s2 2p1" and returns
`((2.0, 2.0), (1.0,))`, where subtuples `(2.0, 2.0)` and `(1.0,)` are
populations of all the l = 0 and l = 1 shells, i.e. `((1s², 2s²), (2p¹,))`.
Within each subtuple populations of levels with the same azimuthal quantum
number l and subsequent radial quantum numbers nᵣ are listed.

Optional keyword arguments are:

* `maxn`  - maximal permitted principal quantum number
* `spins` - number of spin up/spin down electrons per level (1 or 2)
* `notation` - dictionary encoding the correspondence of quantum numbers l = 0, 1... to letters: s, p, d... See [spectroscopic notation](https://en.wikipedia.org/wiki/Spectroscopic_notation)

# Examples
```jldoctest
julia> conf = conf_enc("[Ar] 3d1 4s2") # Scandium, Z = 21
((2.0, 2.0, 2.0, 2.0), (6.0, 6.0), (1.0,))

julia> sum(collect(Iterators.flatten(conf)))
21.0
```
"""
function conf_enc(input::AbstractString;
    maxn=7,
    spins=2,
    notation=Dict("s" => 0, "p" => 1, "d" => 2, "f" => 3))
    # mnemonic for the s,p,d,f,g,h,i,k notation:
    # Sober Physicists Don't Find Giraffes Hiding In Kitchens

    ls = join(keys(notation))

    p = raw"(^\s*(\[(He|Ne|Ar|Kr|Xe|Rn)\]\s*$)|" *
        raw"^\s*(\[(He|Ne|Ar|Kr|Xe|Rn)\])?(\s+[1-9]+\d*" *
        "[" * ls * "]" *
        raw"{1}\d+\.?\d*)+\s*$|^[1-9]+\d*" *
        "[" * ls * "]" *
        raw"{1}\d+\.?\d*(\s+[1-9]+\d*" *
        "[" * ls * "]" *
        raw"{1}\d+\.?\d*)*\s*$)+"

    configur_pattern = Regex(p, "i")
    subshell_pattern = Regex(raw"(\d+)([" * ls * raw"]{1})(\d+\.?\d*)")

    # check syntax of the input string
    if !occursin(configur_pattern, input)
        throw(DomainError(input,
        "input electron configuration doesn't match common pattern," *
        "such as: \"[He] 2s2 2p1\" ... or \"1s2 2s2 2p1 ...\""))
    end

    # include all fully occupied subshells associated with the noble gas
    input = input |> strip |> lowercase |>
    s -> replace(s, "[rn]" => "[xe] 4f14 5d10 6s2 6p6") |>
    s -> replace(s, "[xe]" => "[kr] 4d10 5s2 5p6") |>
    s -> replace(s, "[kr]" => "[ar] 3d10 4s2 4p6") |>
    s -> replace(s, "[ar]" => "[ne] 3s2 3p6") |>
    s -> replace(s, "[ne]" => "[he] 2s2 2p6") |>
    s -> replace(s, "[he]" => "1s2")

    configuration = Vector{Vector{Float64}}()
    buff = zeros(length(notation), maxn)
    uniq = Set{Tuple{Int,Int}}()
    lmax = 0

    # parse subshells, check for physical sense
    for subshell in split(input)
        _n, _l, _occ = match(subshell_pattern, subshell).captures
        n, l, occ = parse(Int, _n), notation[_l], parse(Float64, _occ)

        if !in(n, 1:maxn)
            throw(DomainError(subshell,
            "Principal quantum number is out of range 1:$maxn"))
            elseif l ≥ n
            throw(DomainError(subshell,
            "Azimuthal quantum number must be l < $n"))
            elseif occ > (2l + 1) * spins
            throw(DomainError(subshell,
            "$_l-subshell cannot be occupied by more then" *
            "$((2l + 1) * spins) electrons"))
            elseif in((n, l), uniq)
            throw(DomainError(input,
            "The $n$_l subshell has already been encountered before."))
        end

        push!(uniq, (n, l))
        if l > lmax
            lmax = l
        end
        buff[l + 1,n] = occ
    end

    for l in 0:lmax
        subshell = []; flag = true
        for n in maxn:-1:(l + 1)
            if buff[l + 1, n] == 0.0 && flag
                continue
            end
            flag = false
            push!(subshell, buff[l + 1, n])
        end
        push!(configuration, reverse(subshell))
    end
    return configuration
end

Nₑ(conf) = sum(Iterators.flatten(conf))

macro c_str(config)
    conf_enc(config)
end

#=
function conf_decode(conf)
    ls = ('s', 'p', 'd', 'f')
    s = []
    for l in 1:length(conf), n in 1:length(conf[l])
            push!(s, tuple(n + l - 1, l, conf[l][n]))
    end
    sort!(s, by = x -> x[1])
    return mapreduce(x -> "$(x[1])$(ls[x[2]])$(x[3]) ", *, s, init="") |>
    s -> replace(s, "1s2.0 2s2.0 2p6.0 3s2.0 3p6.0 3d10.0 "  *
                    "4s2.0 4p6.0 4d10.0 4f14.0 5s2.0 5p6.0 " *
                    "5d10.0 6s2.0 6p6.0"            => "[Rn]") |>
    s -> replace(s, "1s2.0 2s2.0 2p6.0 3s2.0 3p6.0 3d10.0 4s2.0 " *
                    "4p6.0 4d10.0 5s2.0 5p6.0"      => "[Xe]") |>
    s -> replace(s, "1s2.0 2s2.0 2p6.0 3s2.0 3p6.0 3d10.0 4s2.0 " *
                    "4p6.0"                         => "[Kr]") |>
    s -> replace(s, "1s2.0 2s2.0 2p6.0 3s2.0 3p6.0" => "[Ar]") |>
    s -> replace(s, "1s2.0 2s2.0 2p6.0"             => "[Ne]") |>
    s -> replace(s, "1s2.0"                         => "[He]") |>
    strip
end
=#


text:
export execute_chromosome

function evaluate_node(c::Chromosome, node::InputNode, context::Vector{BitString})
    if ! node.active
        node.active = true
        c.number_active_nodes += 1
    end
    return context[node.index]
end

function evaluate_node(c::Chromosome, node::InteriorNode, context::Vector{BitString})
    if ! node.active
        func = node.func
        args = map(node.inputs[1:func.arity]) do position
            (level, index) = position
            evaluate_node(c, c[level, index], context)
        end
        node.active = true
        node.cache = func.func(args...)
        c.number_active_nodes += 1
    end
    return node.cache
end

function evaluate_node(c::Chromosome, node::OutputNode, context::Vector{BitString})
    if ! node.active
        node.active = true
        (level, index) = node.input
        node.cache = evaluate_node(c, c[level, index], context)
        c.number_active_nodes += 1
    end
    return node.cache
end

# TODO: Since we are caching the evaluation results we should no
# longer expose the context since providing two different contexts for
# the same chromosome will produce incorrect results the second time
# around.

function execute_chromosome(c::Chromosome, context::Vector{BitString})
    c.active_set = true
    return BitString[evaluate_node(c, node, context) for node = c.outputs]
end

# Executes chrososome using the standard input context
function execute_chromosome(c::Chromosome)
    params = c.params
    mask = output_mask(params.numinputs)

    ctx = std_input_context(params.numinputs)
    result = execute_chromosome(c, ctx)

    return BitString[x & mask for x = result]
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule JonkerVolgenant_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("JonkerVolgenant")
JLLWrappers.@generate_main_file("JonkerVolgenant", UUID("a885a4bc-65ec-5865-a301-890ce0af48d0"))
end  # module JonkerVolgenant_jll


text:
"""
    fit_Cbar(domain, filter::TopHatFilter, u₀, U₀, M, t; λ = 1e-4, method = :ridge)

Fit the filtered convection matrix `C̄` for a `TopHatFilter`, given a list of initial
conditions `u₀` and associated antiderivatives `U₀`.

The unfiltered equation is given by

``\\frac{d u_h}{d t} + C u_h = 0.``

The associated filtered equation is

``\\frac{d \\bar{u}_h}{d t} + \\bar{C} \\bar{u}_h = 0.``

This function fits `C̄` using a regularized least-squares regression based on
the above equation evaluated at different time steps and spatial points.
The regularizatin method is either `:ridge` or `:lasso`, the latter being
able to identify sparsity.

The sparsity pattern of `C̄` is also enforced as a constraint.
"""
function fit_Cbar(domain, filter::TopHatFilter, u₀, U₀, M, t; λ = 1e-4, method = :ridge)
    L = domain.right - domain.left
    x = discretize(domain, M)
    h = filter.width

    # Observation matrices
    U = zeros(M, 0)
    Uₜ = zeros(M, 0)
    for (u₀, U₀) ∈ zip(u₀, U₀)
        ū(x, t) = (U₀(x + h(x) - t) - U₀(x - h(x) - t)) / 2h(x)
        ūₜ(x, t) = -(u₀(x + h(x) - t) - u₀(x - h(x) - t)) / 2h(x)
        U = [U ū.(x, t')]
        Uₜ = [Uₜ ūₜ.(x, t')]
    end

    ℒ = domain isa PeriodicIntervalDomain ? [-L, 0, L] : [0]
    C̄ = spzeros(M, M)
    if method == :ridge
        reg = RidgeRegression(λ; fit_intercept = false)
    elseif method == :lasso
        reg = LassoRegression(λ; fit_intercept = false)
    else
        error("Unsupported method")
    end
    for m = 1:M
        xₘ = x[m]
        nₘ = mapreduce(ℓ -> -h(xₘ) .< (x .+ ℓ .- xₘ) .≤ h(xₘ), .|, ℒ)
        nₘ = mapreduce(i -> circshift(nₘ, i), .|, -1:1)
        Uₘ = U[nₘ, :]
        c̄ = fit(reg, Uₘ', Uₜ[m, :])
        # c̄ = (Uₘ * Uₘ' + λ * I) \ (Uₘ * Uₜ[m, :])
        C̄[m, nₘ] = -c̄
        # c̄ = fit(reg, U', Uₜ[m, :])
        # C̄[m, :] = -c̄
    end

    C̄
end


text:
export AffineMap, square_check, visualize

"""
An `AffineMap` represents a mapping of the plane to itself.
"""
immutable AffineMap
  A::Array{Float64,2}
  b::Array{Float64,1}
end

AffineMap(M::Matrix) = AffineMap(M,[0,0])
AffineMap() = AffineMap(eye(2))

"""
`f(x)` where `f` is an affine map and `x` is a vector
is standard funtion application taking `x` to `Ax+b`.
Equivalent form: `f*x`.
"""
function (f::AffineMap)(x::Vector)
  return f.A * x + f.b
end

(*)(f::AffineMap, x::Vector) = f(x)

function (*)(f::AffineMap, g::AffineMap)
  A = f.A * g.A
  b = f.A*g.b + f.b
  return AffineMap(A,b)
end

"""
`square_check(x)` checks if the vector `x` lies in the
unit square.
"""
square_check(x::Vector) = (0 <= x[1] <= 1) && (0 <= x[2] <= 1)

"""
`square_check(f,x)` checks if `f(x)` is in the unit square.
"""
square_check(f::AffineMap,x::Vector) = square_check(f(x))

"""
`square_check(f)` checks if `f(x)` is in the unit square
for four values of `x` corresponding
to the corners of the unit square.
"""
function square_check(f::AffineMap)
  corners = ([0,0], [0,1], [1,0], [1,1])
  return all(square_check(f,x) for x in corners)
end

function quad_draw(pts, marker::String="-")
  x = [v[1] for v in pts]
  y = [v[2] for v in pts]
  plot(x,y,marker,color="k")
  nothing
end

function visualize(f::AffineMap)
  corners = [ [0,0], [0,1], [1,1], [1,0], [0,0] ]
  quad_draw(corners,":")
  poly = map(f,corners)
  quad_draw(poly)
  axis("equal")
  axis("off")
  axis([-0.25,1.25,-0.25,1.25])
  nothing
end


text:
# include("../utils/env_utils.jl")
mutable struct TabularRoundRobin <: LearningUpdate
end
function update!(self::TabularRoundRobin, reward, action, next_action, state, next_state, args...)
end
Base.size(learner::TabularRoundRobin) = (0,0)

function get_action_probs(self::TabularRoundRobin, state, observation, args...)
    action_probs = zeros(4)
    UP = 1
    RIGHT = 2
    DOWN = 3
    LEFT = 4
    LEFT_JUNCTION = 3
    MIDDLE_JUNCTION = 8
    RIGHT_JUNCTION = 19

    #This relies on being tabular. For continuous TMaze, need different round robin method
    state = if typeof(state) == Int
        state
    else
        #It is a sparse array encoding
        tups = findnz(state)
        tups[1][1]
    end

    if state == LEFT_JUNCTION
        action_probs[DOWN] = 0.5
        action_probs[UP] = 0.5
    elseif state == MIDDLE_JUNCTION
        action_probs[LEFT] = 0.5
        action_probs[RIGHT] = 0.5
    elseif state == RIGHT_JUNCTION
        action_probs[DOWN] = 0.5
        action_probs[UP] = 0.5
    elseif state < 3
        action_probs[UP] = 1
    elseif state < 5
        action_probs[DOWN] = 1
    elseif state < 8
        action_probs[LEFT] = 1
    elseif state < 15
        action_probs[UP] = 1
    elseif state < 17
        action_probs[RIGHT] = 1
    elseif state < RIGHT_JUNCTION
        action_probs[UP] = 1
    else
        action_probs[DOWN] = 1
    end

    return action_probs
end

max_stars_repo_name: MatthewMcLeod/curiosity

max_issues_repo_name: MatthewMcLeod/curiosity

max_forks_repo_name: MatthewMcLeod/curiosity


text:
"""
    BackgroundPlayer(stream::PortAudioStream[, src::AbstractAudioGenerator];
                     block_size=undef)

The object that manages a background audio thread.

Create one instance per audio stream.  Replace or modify the `src` field while
audio is playing with `lock_player` for thread safety.

Functions: `play`, `pause`, `is_playing`, `lock_player`, `finalize`
"""
mutable struct BackgroundPlayer
    _lock::ReentrantLock
    _play_lock::ReentrantLock
    _play_sync::Threads.Condition
    _is_playing::Bool
    stream::PortAudioStream
    src::AbstractAudioGenerator
    block_size::Int
    sample_rate::Float64

    function BackgroundPlayer(stream::PortAudioStream; block_size=undef)
        BackgroundPlayer(
            stream,
            NullAudioGenerator();
            block_size=block_size,
        )
    end
    function BackgroundPlayer(
            stream::PortAudioStream,
            src::AbstractAudioGenerator
            ; block_size=undef)
        # Ensure Julia has multiple threads
        @assert(Threads.nthreads() >= 2,
            "BackgroundPlayer requires its own thread "
            *"(start julia with `\$ JULIA_NUM_THREADS=4 julia)`")
        # Init
        play = src != NullAudioGenerator()  # Auto start playing if not null
        if block_size == undef
            block_size = size(stream.sink.chunkbuf)[2]
        end
        player = new(
            ReentrantLock(),
            ReentrantLock(),
            Threads.Condition(),
            play,
            stream,
            src,
            block_size,
            stream.samplerate,
        )
        # Start background task
        ref = WeakRef(player)
        stop = Threads.Atomic{Bool}(false)
        task = Threads.@spawn _background_play(ref, stop)
        # Register finalizer
        finalizer(player) do p
            # TODO: Check this runs (ensure no memory leaks)
            stop[] = true
            lock(p._play_sync.lock) do
                notify(p._play_sync)
            end
        end
        player
    end
end

function Base.show(io::IO, p::BackgroundPlayer)
    write(io, "BackgroundPlayer "
          *"($(is_playing(p) ? "playing" : "paused"), "
          *"block_size=$(p.block_size))\n"
          *"  Stream: $(p.stream)\n\n  Src: $(p.src)")
end

function _async_println(io, args...)
    @async begin
        println(io, args...)
        flush(io)
    end
end

function _async_print_error(io, msg, e)
    @async begin
        print(io, msg)
        showerror(io, e)
        flush(io)
    end
end

"""Run with Threads.@spawn by BackgroundPlayer."""
function _background_play(ref::WeakRef, stop::Threads.Atomic{Bool})
    try
        while true
            stop[] && break
            quit, play, play_sync = Base.invokelatest(
                                        _background_play_inner, ref)
            quit && break
            play || lock(play_sync.lock) do
                wait(play_sync)  # Wait until allowed to play
            end
        end
    catch e
        _async_print_error(stderr, "Error in BackgroundPlayer thread:\n", e)
    end
end
function _background_play_inner(ref::WeakRef)
    p::Union{Nothing, BackgroundPlayer} = ref.value
    p == nothing && return true, false, nothing
    player = p
    p = nothing
    samples = Ref{Any}(nothing)
    play_sync = player._play_sync
    play = lock(player._play_lock) do
        play = lock(player._lock) do
            play = player._is_playing
            play || return play
            try
                src = player.src
                rate = player.sample_rate
                if !has_samples(src, rate)
                    player._is_playing = false
                    return false
                end
                bs = player.block_size
                samples[] = next_samples!(src, bs, rate)
                @assert(size(samples[])[1] <= bs,
                    "too many samples from source "
                    *"($(size(samples[])[1]) > $bs)")
            catch e
                _async_print_error(
                    stderr, "Error while playing $(player.src):\n", e)
                player._is_playing = false
                player = nothing
                return false
            end
            play
        end
        play && write(player.stream, samples[])
        play
    end
    player = nothing  # Release reference while waiting
    false, play, play_sync
end

"""
    lock_player(player::BackgroundPlayer) do
        # Modify or replace player.src while audio is playing.
        # ...
    end
"""
function lock_player(f::Function, player::BackgroundPlayer)
    lock(f, player._lock)
end

function pause(player::BackgroundPlayer)
    lock(player._play_lock) do
        player._is_playing = false
    end
    nothing
end

function play(player::BackgroundPlayer, src::AbstractAudioGenerator)
    lock_player(player) do
        player.src = src
        play(player)
    end
end
function play(player::BackgroundPlayer)
    lock(player._lock) do
        player._is_playing = true
        lock(player._play_sync.lock) do
            notify(player._play_sync)
        end
    end
    nothing
end

function is_playing(player::BackgroundPlayer)::Bool
    lock(player._lock) do
        player._is_playing
    end
end


text:
path = joinpath(dirname(@__FILE__), "..", "..")
push!(Base.LOAD_PATH, path)
datapath = joinpath(dirname(@__FILE__), "data")

is_julia1 = VERSION >= v"1.0"
if is_julia1
    using Test
    using Dates
    using Random
    using LinearAlgebra
    using DelimitedFiles
    using SparseArrays
else
    using Base.Test
end

use_MOI = false
sets_to_test = Symbol[]
# push!(sets_to_test, :RANDSDP)
# push!(sets_to_test, :SENSORLOC)
push!(sets_to_test, :SDPLIB)
# push!(sets_to_test, :MIMO)

@static if use_MOI
    using ProxSDP, MathOptInterface
    using MosekTools
    if is_julia1
        LinearAlgebra.symmetric_type(::Type{MathOptInterface.VariableIndex}) = MathOptInterface.VariableIndex
        LinearAlgebra.symmetric(v::MathOptInterface.VariableIndex, ::Symbol) = v
        LinearAlgebra.transpose(v::MathOptInterface.VariableIndex) = v
    end
    include("moi_init.jl")
    # optimizer = MOIU.CachingOptimizer(ProxSDPModelData{Float64}(), ProxSDP.Optimizer(log_verbose=true, timer_verbose = true))
    optimizer = ProxSDP.Solver(log_verbose=true, timer_verbose = true, convergence_window=200) #, tol_primal = 1e-3, tol_dual = 1e-3)
else
    solvers = Tuple{String, Function}[]
    using JuMP
    using ProxSDP
    push!(solvers, ("ProxSDP", () -> ProxSDP.Optimizer(log_verbose=true, timer_verbose = true, time_limit = 900.0, log_freq = 1_000, tol_psd=1e-6, tol_primal = 1e-3, tol_dual = 1e-3)))
    # push!(solvers, ("ProxSDP_fullrank", () -> ProxSDP.Optimizer(full_eig_decomp = true, log_verbose=true, timer_verbose = false, time_limit = 900.0, log_freq = 1_000)))
    # using MosekTools
    # push!(solvers, ("MOSEK", () -> Mosek.Optimizer(MSK_DPAR_OPTIMIZER_MAX_TIME = 900.0))) # eps = ???
    # using CSDP
    # push!(solvers, ("CSDP", () -> CSDP.Optimizer(objtol=1e-4, maxiter=100000))) # eps = ???
    # using SDPA
    # push!(solvers, ("SDPA", () -> SDPA.Optimizer())) # eps = ???
    # using SCS
    # push!(solvers, ("SCS", () -> SCS.Optimizer())) # eps = 1e-4
    # using COSMO
    # push!(solvers, ("COSMO", () -> COSMO.Optimizer(time_limit = 900.0, max_iter = 100_000, eps_abs = 1e-3))) # eps = 1e-4
end

# function ProxSDP.get_solution(opt)
#     nothing
# end

NOW = is_julia1 ? replace("$(now())",":"=>"_") : replace("$(now())",":","_")
FILE = open(joinpath(dirname(@__FILE__),"proxsdp_bench_$(NOW).log"),"w")
println(FILE, "class, prob_ref, time, p_obj, d_obj, p_res, d_res")
function println2(FILE, class::String, ref::String, sol::ProxSDP.MOISolution)
    println(FILE, "$class, $ref, $(sol.time), $(sol.objval), $(sol.dual_objval), $(sol.primal_residual), $(sol.dual_residual)")
    flush(FILE)
end
println2(FILE, class::String, ref::String, sol::Nothing) = nothing
function println2(FILE, class::String, ref::String, sol)
    println(FILE, "$class, $ref, $(sol[2]), $(sol[1])")
    flush(FILE)
end
function println2(FILE, solver::String, class::String, ref::String, sol)
    # println(FILE, "$solver, $class, $ref, $(sol[2]), $(sol[1])")
    println(FILE, "$solver, $class, $ref, $(sol[2]), $(sol[1]), $(sol[3])")
    flush(FILE)
end

RANDSDP_TEST_SET = 1:1
SENSORLOC_TEST_SET = [#100:100:300#1000
    100,
    200,
    300,
    # 400, # mosek stops here - maybenot in new form
    # 500, # proxsdp 166s
    # 600,
    # 700,
    # 800,
    # 900,
    # 1000,
]
MIMO_TEST_SET = [
    100,
    500, # mosek stops here
    # 1000, # 700s on scs
    # 1500, # SCS stops here (?)
    # 2000,
    # 2500,
    # 3000,
    # 3500,
    # 4000,
    # 4500,
    # 5000,
    # 6000,
    # 7000,
    # 8000,
    # 9000,
    # 10_000,
    ]
GPP_TEST_SET = [
    "gpp124-1.dat-s",
    "gpp124-1.dat-s",
    "gpp124-2.dat-s",
    "gpp124-3.dat-s",
    "gpp124-4.dat-s",
    "gpp250-1.dat-s",
    "gpp250-2.dat-s",
    "gpp250-3.dat-s",
    "gpp250-4.dat-s",
    "gpp500-1.dat-s", # SCS > 1000 s
    "gpp500-2.dat-s", # SCS > 1000 s
    "gpp500-3.dat-s",
    "gpp500-4.dat-s",
    # "equalG11.dat-s",
    # "equalG51.dat-s",
]
MAXCUT_TEST_SET = [
    "mcp124-1.dat-s",
    "mcp124-1.dat-s",
    "mcp124-2.dat-s",
    "mcp124-3.dat-s",
    "mcp124-4.dat-s",
    "mcp250-1.dat-s",
    "mcp250-2.dat-s",
    "mcp250-3.dat-s",
    "mcp250-4.dat-s",
    "mcp500-1.dat-s",
    "mcp500-2.dat-s",
    "mcp500-3.dat-s",
    "mcp500-4.dat-s",
    # "maxG11.dat-s"  ,
    # "maxG51.dat-s"  ,
    # "maxG32.dat-s"  ,
    # "maxG55.dat-s"  ,
    # "maxG60.dat-s"  ,
]

include("base_randsdp.jl")
include("moi_randsdp.jl")
include("jump_randsdp.jl")

include("base_mimo.jl")
include("moi_mimo.jl")
include("jump_mimo.jl")

include("base_sensorloc.jl")
include("moi_sensorloc.jl")
include("jump_sensorloc.jl")

include("base_sdplib.jl")
include("moi_sdplib.jl")
include("jump_sdplib.jl")

if use_MOI
    _randsdp = moi_randsdp
    _mimo = moi_mimo
    _sensorloc = moi_sensorloc
    _sdplib = moi_sdplib
else
    _randsdp = jump_randsdp
    _mimo = jump_mimo
    _sensorloc = jump_sensorloc
    _sdplib = jump_sdplib
end

for optimizer in solvers
    if :RANDSDP in sets_to_test
        println("RANDSDP")
        _randsdp(optimizer[2], 0, 5, 5)
        for i in RANDSDP_TEST_SET
            @show i
            sol = _randsdp(optimizer[2], i, 5, 5)
            println2(FILE, optimizer[1], "RANDSDP", "$i", sol)
        end
    end
    if :SENSORLOC in sets_to_test
        println("SENSORLOC")
        for n in SENSORLOC_TEST_SET
            @show n
            sol = _sensorloc(optimizer[2], 0, n)
            println2(FILE, optimizer[1], "SENSORLOC", "$n", sol)
        end
    end
    if :SDPLIB in sets_to_test
        println("gpp")
        for name in GPP_TEST_SET
            println(name)
            sol = _sdplib(optimizer[2], joinpath(datapath, name))
            println2(FILE, optimizer[1], "SDPLIB_gp", name, sol)
        end
        println("max_cut")
        for name in MAXCUT_TEST_SET
            println(name)
            sol = _sdplib(optimizer[2], joinpath(datapath, name))
            println2(FILE, optimizer[1], "SDPLIB_mc", name, sol)
        end
    end
    if :MIMO in sets_to_test
        println("MIMO")
        _mimo(optimizer[2], 0, 100)
        for n in MIMO_TEST_SET
            @show n
            sol = _mimo(optimizer[2], 0, n)
            println2(FILE, optimizer[1], "MIMO", "$n", sol)
        end
    end
end

close(FILE)


text:
using Test
using GOR

include("test_gorfile.jl")
#include("test_parquetfile.jl")
include("test_verifyorder.jl")
include("test_merge.jl")
include("test_join.jl")
include("test_map.jl")
include("test_aggregate.jl")
include("test_groupby.jl")


text:
using Distances
using Random

"""
Représente une distance entre deux données identifiées par leur id
"""
mutable struct Distance

    distance::Float64
    ids::Vector{Int}

    function Distance()
        return new()
    end
end

"""
Constructeur d'une distance

Entrées :
- id1 : id de la première donnée
- id2 : id de la seconde donnée
- x   : caractéristique des données d'entraînement
"""
function Distance(id1::Int, id2::Int, x::Matrix{Float64})

    d = Distance()
    d.distance = euclidean(x[id1, :], x[id2, :])
    d.ids = [id1, id2]

    return d

end

##### Norme 1

# function Distance(id1::Int, id2::Int, x::Matrix{Float64})

#     d = Distance()
#     d.distance = sum(abs(x[id1,i] -x[id2,i]) for i in 1:size(x)[2])
#     d.ids = [id1,id2]

#     return d

# end

##### Norme infinie

# function Distance(id1::Int, id2::Int, x::Matrix{Float64})

#     d = Distance()
#     d.distance = max(abs(x[id1,i] -x[id2,i]) for i in 1:size(x)[2])
#     d.ids = [id1,id2]

#     return d

# end

##### Pseudo-norme min

# function Distance(id1::Int, id2::Int, x::Matrix{Float64})

#     d = Distance()
#     d.distance = min(abs(x[id1,i] -x[id2,i]) for i in 1:size(x)[2])
#     d.ids = [id1,id2]

#     return d

# end

##### Norme 2 perturbée

# function Distance(id1::Int, id2::Int, x::Matrix{Float64})

#     d = Distance()
#     acc = euclidean(x[id1, :], x[id2, :])
#     d.distance = acc * (0.6 + 0.4*rand())
#     d.ids = [id1,id2]

#     return d

# end

max_stars_repo_name: boumpteryx/RODM

max_issues_repo_name: boumpteryx/RODM

max_forks_repo_name: boumpteryx/RODM


text:
using Pkg.Artifacts
using Pkg.TOML

# This is the path to the Artifacts.toml we will manipulate
artifact_toml = joinpath(@__DIR__, "Artifacts.toml")
verilog_lib = joinpath(@__DIR__, "verilog-stdlib/")
project_toml = TOML.parsefile(joinpath(@__DIR__, "Project.toml"))
version = project_toml["version"]

# create_artifact() returns the content-hash of the artifact directory once we're finished creating it
dir_hash = create_artifact() do artifact_dir
    # We create the artifact by simply copying over the library
    cp(verilog_lib, artifact_dir; force = true)
end
# compress the artifact dir into a tar archive
tar_hash = archive_artifact(dir_hash, joinpath(@__DIR__, "verilog-stdlib.tar.gz"))

# Now bind that hash within our `Artifacts.toml`
base_url = "https://github.com/UW-PHARM/BitSAD.jl/raw"
dev_url = "$base_url/master/verilog-stdlib.tar.gz"
release_url = "$base_url/v$version/verilog-stdlib.tar.gz"
bind_artifact!(artifact_toml, "verilog-stdlib", dir_hash;
               download_info = [(release_url, tar_hash), (dev_url, tar_hash)],
               force = true,
               lazy = true)


text:
using NCDatasets
using Statistics
using Interpolations
using LinearAlgebra
using Glob
using JLD
"""
generate_mima_params(mima_params::Array{Float64},
                            mima_param_names::Array{String})

Generate a MIMAParameters file, setting the values (mima_params) of
a group of MIMAParameters (mima_param_names).
"""
function generate_mima_params(mima_params::Array{Float64}, mima_param_names::Array{String})
    # Generate version
    version = rand(11111:99999)
    if length(mima_params) == length(mima_param_names)
        open("mima_param_defs_$(version).txt", "w") do io
            for i in 1:length(mima_params)
                write(
                    io,
		    "$(mima_param_names[i]) = $(mima_params[i])\n",
                )
            end
        end
    else
        throw(ArgumentError("Number of parameter names must be equal to number of values provided."))
    end
    return version
end

"""
agg_mima_ekp(n_params::Integer, output_name::String="ekp_mima")

Aggregate all iterations of the parameter ensembles and write to file, given the
number of parameters in each parameter vector (p).
"""
function agg_mima_ekp(n_params::Integer, output_name::String = "ekp_mima")
    # Get versions
    version_files = glob("versions_*.txt")
    # Recover parameters of last iteration
    last_up_versions = readlines(version_files[end])

    ens_all = Array{Float64, 2}[]
    for (it_num, file) in enumerate(version_files)
        versions = readlines(file)
        u = zeros(length(versions), n_params)
        for (ens_index, version_) in enumerate(versions)
            if it_num == length(version_files)
                open("../../../ClimateMachine.jl/test/Atmos/EDMF/$(version_)", "r") do io
                    u[ens_index, :] =
                        [parse(Float64, line) for (index, line) in enumerate(eachline(io)) if index % 3 == 0]
                end
            else
                open("$(version_).output/$(version_)", "r") do io
                    u[ens_index, :] =
                        [parse(Float64, line) for (index, line) in enumerate(eachline(io)) if index % 3 == 0]
                end
            end
        end
        push!(ens_all, u)
    end
    save(string(output_name, ".jld"), "ekp_u", ens_all)
    return
end


text:
# include("../../VSCodeServer/src/repl.jl")

import Sockets, Base64

include("protocol/debug_adapter_protocol.jl")
include("debugger_utils.jl")
include("debugger_core.jl")
include("debugger_requests.jl")

function clean_up_ARGS_in_launch_mode()
    pipename = ARGS[1]
    crashreporting_pipename = ARGS[2]
    deleteat!(ARGS, 1)
    deleteat!(ARGS, 1)

    if ENV["JL_ARGS"] != ""
        cmd_ln_args_encoded = split(ENV["JL_ARGS"], ';')

        delete!(ENV, "JL_ARGS")

        cmd_ln_args_decoded = map(i -> String(Base64.base64decode(i)), cmd_ln_args_encoded)

        for arg in cmd_ln_args_decoded
            push!(ARGS, arg)
        end
    end

    return pipename, crashreporting_pipename
end

function startdebug(socket, error_handler=nothing)
    @debug "Connected to debug adapter."

    try

        endpoint = JSONRPC.JSONRPCEndpoint(socket, socket, error_handler)

        try

            run(endpoint)

            state = DebuggerState()

            msg_dispatcher = JSONRPC.MsgDispatcher()
            msg_dispatcher[disconnect_request_type] = (conn, params) -> disconnect_request(conn, state, params)
            msg_dispatcher[run_notification_type] = (conn, params) -> run_notification(conn, state, params)
            msg_dispatcher[debug_notification_type] = (conn, params) -> debug_notification(conn, state, params)

            msg_dispatcher[exec_notification_type] = (conn, params) -> exec_notification(conn, state, params)
            msg_dispatcher[set_break_points_request_type] = (conn, params) -> set_break_points_request(conn, state, params)
            msg_dispatcher[set_exception_break_points_request_type] = (conn, params) -> set_exception_break_points_request(conn, state, params)
            msg_dispatcher[set_function_exception_break_points_request_type] = (conn, params) -> set_function_break_points_request(conn, state, params)
            msg_dispatcher[stack_trace_request_type] = (conn, params) -> stack_trace_request(conn, state, params)
            msg_dispatcher[scopes_request_type] = (conn, params) -> scopes_request(conn, state, params)
            msg_dispatcher[source_request_type] = (conn, params) -> source_request(conn, state, params)
            msg_dispatcher[variables_request_type] = (conn, params) -> variables_request(conn, state, params)
            msg_dispatcher[continue_request_type] = (conn, params) -> continue_request(conn, state, params)
            msg_dispatcher[next_request_type] = (conn, params) -> next_request(conn, state, params)
            msg_dispatcher[step_in_request_type] = (conn, params) -> setp_in_request(conn, state, params)
            msg_dispatcher[step_in_targets_request_type] = (conn, params) -> step_in_targets_request(conn, state, params)
            msg_dispatcher[step_out_request_type] = (conn, params) -> setp_out_request(conn, state, params)
            msg_dispatcher[evaluate_request_type] = (conn, params) -> evaluate_request(conn, state, params)
            msg_dispatcher[terminate_request_type] = (conn, params) -> terminate_request(conn, state, params)
            msg_dispatcher[exception_info_request_type] = (conn, params) -> exception_info_request(conn, state, params)
            msg_dispatcher[restart_frame_request_type] = (conn, params) -> restart_frame_request(conn, state, params)
            msg_dispatcher[set_variable_request_type] = (conn, params) -> set_variable_request(conn, state, params)
            msg_dispatcher[threads_request_type] = (conn, params) -> threads_request(conn, state, params)
            msg_dispatcher[breakpointslocation_request_type] = (conn, params) -> breakpointlocations_request(conn, state, params)
            msg_dispatcher[set_compiled_items_notification_type] = (conn, params) -> set_compiled_items_request(conn, state, params)
            msg_dispatcher[set_compiled_mode_notification_type] = (conn, params) -> set_compiled_mode_request(conn, state, params)

            @async try
                for msg in endpoint
                    JSONRPC.dispatch_msg(endpoint, msg_dispatcher, msg)
                end
            catch err
                if error_handler === nothing
                    Base.display_error(err, catch_backtrace())
                else
                    error_handler(err, Base.catch_backtrace())
                end
            end

            while true
                msg = take!(state.next_cmd)

                ret = nothing

                if msg.cmd == :run
                    try
                        Main.include(msg.program)
                    catch err
                        Base.display_error(stderr, err, catch_backtrace())
                    end

                    JSONRPC.send(endpoint, finished_notification_type, nothing)
                    break
                elseif msg.cmd == :stop
                    break
                elseif msg.cmd == :set_source_path
                    task_local_storage()[:SOURCE_PATH] = msg.source_path
                else
                    if msg.cmd == :continue
                        ret = our_debug_command(:c, state)
                    elseif msg.cmd == :next
                        ret = our_debug_command(:n, state)
                    elseif msg.cmd == :stepIn
                        if msg.targetId === missing
                            ret = our_debug_command(:s, state)
                        else
                            pc = state.frame.pc
                            itr = 0
                            success = true
                            while state.frame.pc < msg.targetId
                                ret = our_debug_command(:nc, state)
                                if ret isa JuliaInterpreter.BreakpointRef
                                    success = false
                                    break
                                end
                                itr += 1
                                if itr > 100
                                    success = false
                                    @warn "Could not step into sepcified target."
                                    break
                                end
                            end

                            if success
                                ret = our_debug_command(:s, state)
                            end
                        end
                    elseif msg.cmd == :stepOut
                        ret = our_debug_command(:finish, state)
                    end

                    if ret === nothing
                        JSONRPC.send(endpoint, finished_notification_type, nothing)
                        state.debug_mode == :launch && break
                    else
                        send_stopped_msg(endpoint, ret, state)
                    end
                end
            end

            @debug "Finished debugging"
        finally
            close(socket)
        end
    catch err
        if error_handler === nothing
            rethrow(err)
        else
            error_handler(err, Base.catch_backtrace())
        end
    end
end


text:
# This file includes coding methods to map bit streams into symbol streams

export GrayCoding, invmap

abstract type AbstractCoding end

"""
    $TYPEDEF

Gray coding.

# Fields
    $TYPEDFIELDS

# Example
```julia
julia> gc = GrayCoding(2)
GrayCoding(Dict([0, 1] => 2,[1, 1] => 3,[0, 0] => 1,[1, 0] => 4))
```
"""
struct GrayCoding <: AbstractCoding
    pairs::Dict{Vector{Int}, Int}
    function GrayCoding(M::Int)
        ispow2(M) || throw(ArgumentError("Expected `M` to be a power of 2, got $M instead"))
        k = Int(log2(M))
        m = 0 : 1 << k - 1
        codes = m .⊻ (m .>> 1)
        codes = reverse.(digits.(codes, base=2, pad=k))
        new(Dict(zip(codes, m .+ 1)))
    end
end

show(io::IO, coding::GrayCoding) = print(io, "GrayCoding(M:$(length(coding.pairs)))")

# When a GrayCoding object is called, it maps its input bit stream to coressponding symbol stream.
function (coding::GrayCoding)(stream)
    codewords = collect(Iterators.partition(stream, Int(log2(constelsize(coding)))))
    alph = coding.pairs
    map(codewords) do codeword
        alph[codeword]
    end
end

"""
    $SIGNATURES

Returns constellation size of the coding.
"""
constelsize(coding::GrayCoding) = length(coding.pairs)

"""
    $SIGNATURES

Returns the inverse of the mapping. Inverse mapping maps the levels to bit chunks.

# Example
```julia
julia> gray = GrayCoding(3)
GrayCoding(Dict([0, 1, 1] => 3,[0, 0, 1] => 2,[1, 1, 0] => 5,[0, 0, 0] => 1,[1, 0, 0] => 8,[1, 0, 1] => 7,[0, 1, 0] => 4,[1, 1, 1] => 6))

julia> invmap(gray)
Dict{Int64,Array{Int64,1}} with 8 entries:
  7 => [1, 0, 1]
  4 => [0, 1, 0]
  2 => [0, 0, 1]
  3 => [0, 1, 1]
  5 => [1, 1, 0]
  8 => [1, 0, 0]
  6 => [1, 1, 1]
  1 => [0, 0, 0]
```
"""
invmap(coding::AbstractCoding) = Dict([(v, k) for (k, v) in coding.pairs])


text:
"""
Struct for Brownian Motion

		bmProcess=BrownianMotion(σ::num1,μ::num2) where {num1,num2 <: Number}

Where:\n
		σ	=	volatility of the process.
		μ	=	drift of the process.
"""
mutable struct BrownianMotion{num <: Number, num1 <: Number, numtype <: Number} <: AbstractMonteCarloEngine{numtype}
    σ::num
    μ::num1
    function BrownianMotion(σ::num, μ::num1) where {num <: Number, num1 <: Number}
        if σ <= 0
            error("Volatility must be positive")
        else
            zero_typed = zero(num) + zero(num1)
            return new{num, num1, typeof(zero_typed)}(σ, μ)
        end
    end
end

export BrownianMotion;

function simulate!(X, mcProcess::BrownianMotion, mcBaseData::MonteCarloConfiguration{type1, type2, type3, SerialMode, type5}, T::numb) where {numb <: Number, type1 <: Integer, type2 <: Integer, type3 <: StandardMC, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    σ = mcProcess.σ
    μ = mcProcess.μ
    @assert T > 0
    dt = T / Nstep
    mean_bm = μ * dt
    stddev_bm = σ * sqrt(dt)
    isDualZero = mean_bm * stddev_bm * 0
    view(X, :, 1) .= isDualZero
    zero_typed = predict_output_type_zero_(σ, μ)
    Z = Array{typeof(get_rng_type(zero_typed))}(undef, Nsim)
    @inbounds for j = 1:Nstep
        randn!(mcBaseData.rng, Z)
        @. @views X[:, j+1] .= X[:, j] + mean_bm + stddev_bm * Z
    end

    nothing
end

function simulate!(X, mcProcess::BrownianMotion, mcBaseData::MonteCarloConfiguration{type1, type2, type3, SerialMode, type5}, T::numb) where {numb <: Number, type1 <: Integer, type2 <: Integer, type3 <: AntitheticMC, type5 <: Random.AbstractRNG}
    Nsim = mcBaseData.Nsim
    Nstep = mcBaseData.Nstep
    σ = mcProcess.σ
    μ = mcProcess.μ
    @assert T > 0
    dt = T / Nstep
    mean_bm = μ * dt
    stddev_bm = σ * sqrt(dt)
    isDualZero = mean_bm * stddev_bm * 0
    view(X, :, 1) .= isDualZero
    Nsim_2 = div(Nsim, 2)
    zero_typed = predict_output_type_zero_(σ, μ)
    Z = Array{typeof(get_rng_type(zero_typed))}(undef, Nsim_2)
    Xp = @views X[1:Nsim_2, :]
    Xm = @views X[(Nsim_2+1):end, :]
    @inbounds for j = 1:Nstep
        randn!(mcBaseData.rng, Z)
        @. @views Xp[:, j+1] .= Xp[:, j] + mean_bm + stddev_bm * Z
        @. @views Xm[:, j+1] .= Xm[:, j] + mean_bm - stddev_bm * Z
    end

    nothing
end


text:
export cgnr

mutable struct CGNR{vecT,T,Tsparse} <: AbstractLinearSolver
  S
  SHWS
  reg::Regularization
  cl::vecT
  rl::vecT
  zl::vecT
  pl::vecT
  vl::vecT
  xl::vecT
  αl::T
  βl::T
  ζl::T
  weights::vecT
  enforceReal::Bool
  enforcePositive::Bool
  sparseTrafo::Tsparse
  iterations::Int64
  relTol::Float64
  z0::Float64
  normalizeReg::Bool
  regFac::Float64
end

"""
    CGNR(A, x::vecT; λ = 0.0, reg = Regularization("L2", λ), kargs...) where vecT

creates an `CGNR` object for the system matrix `A`.

# Arguments
* `A`                               - system matrix
* `x::vecT`                         - Array with the same type and size as the solution
* (`λ=0.0`)                         - Regularization paramter
* (`reg=Regularization("L2", λ)`)   - Regularization object
* (weights::vecT=eltype(S)[]) - weights for the data term
* (sparseTrafo=nothing)             - sparsifying transform
* (enforceReal::Bool=false)         - constrain the solution to be real
* (enforcePositive::Bool=false)     - constrain the solution to have positive real part
* (iterations::Int64=10)            - number of iterations
* (`relTol::Float64=eps()`)         - rel tolerance for stopping criterion
"""
function CGNR(S, x::vecT=zeros(eltype(S),size(S,2)); λ::Real=0.0, reg::R = Regularization("L2", λ)
              , weights::vecT=similar(x,0)
              , sparseTrafo=nothing
              , enforceReal::Bool=false
              , enforcePositive::Bool=false
              , iterations::Int64=10
              , relTol::Float64=eps()
              , normalizeReg::Bool=false
              , kargs...) where {vecT<:AbstractVector,R<:Union{Regularization, Vector{Regularization}}}

  if typeof(reg)==Vector{Regularization}
    reg = reg[1]
  end

  if (reg.prox!) != (proxL2!)
    @error "CGNR only supports L2 regularizer"
  end

  M,N = size(S)
  T = eltype(S)
  cl = similar(x,N)
  rl = similar(x,M)     #residual vector
  zl = similar(x,N)     #temporary vector
  pl = similar(x,N)     #temporary vector
  vl = similar(x,N)     #temporary vector
  xl = similar(x,M)     #temporary vector
  αl = zero(T)        #temporary scalar
  βl = zero(T)        #temporary scalar
  ζl = zero(T)        #temporary scalar
  SHWS = normalOperator(S, isempty(weights) ? opEye() : WeightingOp(weights))

  return CGNR(S,SHWS,
             reg,cl,rl,zl,pl,vl,xl,αl,βl,ζl,
             weights,enforceReal,enforcePositive,sparseTrafo,iterations,relTol,0.0,normalizeReg,1.0)
end

"""
init!(solver::CGNR{vecT,T,Tsparse}, u::vecT
              ; S::matT=solver.S
              , cl::vecT=similar(u,0)
              , weights::vecT=solver.weights) where {vecT,T,Tsparse,matT}

(re-) initializes the CGNR iterator
"""
function init!(solver::CGNR{vecT,T,Tsparse}, u::vecT
              ; S::matT=solver.S
              , cl::vecT=similar(u,0)
              , weights::vecT=solver.weights) where {vecT,T,Tsparse,matT}

  solver.S = S
  # TODO, the following line is called a second time...
  #solver.SHWS = normalOperator(S, isempty(weights) ? opEye() : WeightingOp(weights))
  if isempty(cl)
    solver.cl[:] .= zero(T)
  else
    solver.cl[:] .= cl
  end
  solver.rl[:] .= u - S*solver.cl
  solver.zl[:] .= zero(T)     #temporary vector
  solver.pl[:] .= zero(T)     #temporary vector
  solver.vl[:] .= zero(T)     #temporary vector
  solver.xl[:] .= zero(T)     #temporary vector
  solver.αl = zero(T)        #temporary scalar
  solver.βl = zero(T)        #temporary scalar
  solver.ζl = zero(T)        #temporary scalar

  #zl = Sᶜ*rl, where ᶜ denotes complex conjugation
  if !isempty(weights)
    solver.xl[:] .= solver.rl .* weights
    ## gemv!('C',one(T), S, xl, zero(T), zl)
    solver.zl[:] .= adjoint(S)*solver.xl
  else
    ## gemv!('C',one(T), S, rl, zero(T), zl)
    solver.zl[:] .= adjoint(S)*solver.rl
  end
  solver.z0 = norm(solver.zl)
  copyto!(solver.pl,solver.zl)

  # normalization of regularization parameters
  if solver.normalizeReg
    solver.regFac = norm(u,1)/length(u)
  else
    solver.regFac = 1.0
  end
end

"""
    solve(solver::CGNR, u::vecT) where vecT

solves Thikhonov-regularized inverse problem using CGNR.

# Arguments
* `solver::CGNR                         - the solver containing both system matrix and regularizer
* `u::vecT`                             - data vector
* (`S::matT=solver.S`)                  - operator for the data-term of the problem
* (`startVector::vecT=similar(u,0)`)    - initial guess for the solution
* (`weights::Vector{T}=solver.weights`) - weights for the data term
* (`solverInfo=nothing`)                - solverInfo for logging

when a `SolverInfo` objects is passed, the residuals `solver.zl` are stored in `solverInfo.convMeas`.
"""
function solve(solver::CGNR{vecT,T,Tsparse}, u::vecT; S::matT=solver.S, startVector::vecT=similar(u,0), weights::vecT=solver.weights, solverInfo=nothing, kargs...) where {vecT,T,Tsparse,matT}
  # initialize solver parameters
  init!(solver, u; S=S, cl=startVector, weights=weights)

  # log solver information
  solverInfo != nothing && storeInfo(solverInfo,solver.cl,norm(solver.zl))

  # perform CGNR iterations
  for (iteration, item) = enumerate(solver)
    solverInfo != nothing && storeInfo(solverInfo,solver.cl,norm(solver.zl))
  end

  return solver.cl
end

"""
  iterate(solver::CGNR{vecT,T,Tsparse}, iteration::Int=0) where {vecT,T,Tsparse}

performs one CGNR iteration.
"""
function iterate(solver::CGNR{vecT,T,Tsparse}, iteration::Int=0) where {vecT,T,Tsparse}
    if done(solver,iteration)
      applyConstraints(solver.cl, solver.sparseTrafo, solver.enforceReal, solver.enforcePositive)
      return nothing
    end

    solver.vl[:] .= solver.SHWS*solver.pl

    # αl = zlᴴ⋅zl/(vlᴴ⋅vl+λ*plᴴ⋅pl)
    solver.ζl= norm(solver.zl)^2
    normvl = dot(solver.pl,solver.vl)

    if solver.reg.λ > 0
      solver.αl = solver.ζl/(normvl+solver.regFac*solver.reg.λ*norm(solver.pl)^2)
    else
      solver.αl = solver.ζl/normvl
    end

    #cl += αl*pl
    BLAS.axpy!(solver.αl,solver.pl,solver.cl)

    #rl += -αl*vl
    BLAS.axpy!(-solver.αl,solver.vl,solver.zl)

    if solver.reg.λ > 0
      BLAS.axpy!(-solver.reg.λ*solver.αl,solver.pl,solver.zl)
    end

    # βl = zl₊₁ᴴ⋅zl₊₁/zlᴴ⋅zl
    solver.βl = dot(solver.zl,solver.zl)/solver.ζl

    #pl = zl + βl*pl
    rmul!(solver.pl,solver.βl)
    BLAS.axpy!(one(eltype(solver.S)),solver.zl,solver.pl)
    return solver.zl, iteration+1
end

#=
function iterate(solver::CGNR{T,Tsparse}, iteration::Int=0) where {T,Tsparse}
    if done(solver,iteration)
      applyConstraints(solver.cl, solver.sparseTrafo, solver.enforceReal, solver.enforcePositive)
      return nothing
    end

    #vl = Sᵗ*pl
    ##gemv!('N',one(T), S, pl, zero(T), vl)
    solver.vl[:] .= solver.S*solver.pl

    # αl = zlᴴ⋅zl/(vlᴴ⋅vl+λ*plᴴ⋅pl)
    solver.ζl= norm(solver.zl)^2
    normvl = isempty(solver.weights) ? dot(solver.vl,solver.vl) : dot(solver.vl,solver.weights.*solver.vl)

    if solver.reg.λ > 0
      solver.αl = solver.ζl/(normvl+solver.reg.λ*norm(solver.pl)^2)
    else
      solver.αl = solver.ζl/normvl
    end

    #cl += αl*pl
    BLAS.axpy!(solver.αl,solver.pl,solver.cl)

    #rl += -αl*vl
    BLAS.axpy!(-solver.αl,solver.vl,solver.rl)

    #zl = Sᶜ*rl-λ*cl
    if !isempty(solver.weights)
      solver.xl[:] .= solver.rl .* solver.weights
      ##gemv!('C',one(T), S, xl, zero(T), zl)
      solver.zl[:] .= adjoint(solver.S)*solver.xl
    else
      ##gemv!('C',one(T), S, rl, zero(T), zl)
      solver.zl[:] .= adjoint(solver.S)*solver.rl
    end
    if solver.reg.λ > 0
      BLAS.axpy!(-solver.reg.λ,solver.cl,solver.zl)
    end

    # βl = zl₊₁ᴴ⋅zl₊₁/zlᴴ⋅zl
    solver.βl = dot(solver.zl,solver.zl)/solver.ζl

    #pl = zl + βl*pl
    rmul!(solver.pl,solver.βl)
    BLAS.axpy!(one(eltype(solver.S)),solver.zl,solver.pl)
    return solver.rl, iteration+1
end
=#

function converged(solver::CGNR)
  return norm(solver.zl)/solver.z0 <= solver.relTol
end

@inline done(solver::CGNR,iteration::Int) = converged(solver) || iteration>=min(solver.iterations, size(solver.S,2))


text:
using ITensors

function site_inds(ny,nx,d=1)
  sh = Matrix{Index}(undef,ny,nx)
  sv = Matrix{Index}(undef,ny,nx)
  for iy = 1:ny, ix = 1:nx
    sh[iy,ix] = Index(d,"site,horiz,x=$ix,y=$iy")
    sv[iy,ix] = Index(d,"site,vert,x=$ix,y=$iy")
  end
  return sh,sv
end

function link_inds(ny,nx,d=1)
  ll = Matrix{Index}(undef,ny,nx)
  lr = Matrix{Index}(undef,ny,nx)
  lu = Matrix{Index}(undef,ny,nx)
  ld = Matrix{Index}(undef,ny,nx)
  for iy = 1:ny, ix = 1:nx
    ll[iy,ix] = Index(d,"link,left,x=$ix,y=$iy")
    lr[iy,ix] = Index(d,"link,right,x=$ix,y=$iy")
    lu[iy,ix] = Index(d,"link,up,x=$ix,y=$iy")
    ld[iy,ix] = Index(d,"link,down,x=$ix,y=$iy")
  end
  return ll,lr,lu,ld
end

per(n,N) = mod(n-1,N)+1

function ctmrg_environment((sh,sv))
  ny,nx = size(sh)
  Clu = Matrix{ITensor}(undef,ny,nx)
  Cru = Matrix{ITensor}(undef,ny,nx)
  Cld = Matrix{ITensor}(undef,ny,nx)
  Crd = Matrix{ITensor}(undef,ny,nx)

  Al = Matrix{ITensor}(undef,ny,nx)
  Ar = Matrix{ITensor}(undef,ny,nx)
  Au = Matrix{ITensor}(undef,ny,nx)
  Ad = Matrix{ITensor}(undef,ny,nx)

  ll,lr,lu,ld = link_inds(ny,nx)
  for iy = 1:ny, ix = 1:nx
    Clu[iy,ix] = randomITensor(ll[iy,ix],lu[iy,ix])
    Cru[iy,ix] = randomITensor(lr[iy,ix],lu[iy,ix])
    Cld[iy,ix] = randomITensor(ll[iy,ix],ld[iy,ix])
    Crd[iy,ix] = randomITensor(lr[iy,ix],ld[iy,ix])
    iyp,ixp = per(iy+1,ny),per(ix+1,nx)
    Al[iy,ix] = randomITensor(sh[iy,ix],ll[iy,ix],ll[iyp,ix])
    Ar[iy,ix] = randomITensor(sh[iy,ix],lr[iy,ix],lr[iyp,ix])
    Au[iy,ix] = randomITensor(sv[iy,ix],lu[iy,ix],lu[iy,ixp])
    Ad[iy,ix] = randomITensor(sv[iy,ix],ld[iy,ix],ld[iy,ixp])
  end
  normalize!((Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))
  return (Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad)
end

function calc_κ(iy,ix,T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir="left")
  ny,nx = size(T)
  iyp,ixp = per(iy+1,ny),per(ix+1,nx)
  normC = scalar(Clu[iy,ix]*Cld[iy,ix]*Cru[iy,ix]*Crd[iy,ix])
  normAlr = scalar(Clu[iy,ix]*Cru[iy,ix]*Al[iy,ix]*Ar[iy,ix]*Cld[iyp,ix]*Crd[iyp,ix])
  normAud = scalar(Clu[iy,ix]*Cld[iy,ix]*Au[iy,ix]*Ad[iy,ix]*Cru[iy,ixp]*Crd[iy,ixp])
  normT = scalar(Clu[iy,ix]*Al[iy,ix]*Cld[iyp,ix]*Au[iy,ix]*T[iy,ix]*Ad[iyp,ix]*Cru[iy,ixp]*Ar[iy,ixp]*Crd[iyp,ixp])
  return normT,normAlr,normAud,normC
end

function calc_κ(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir="left")
  ny,nx = size(T)
  κ = Matrix{Float64}(undef,ny,nx)
  for iy = 1:ny, ix = 1:nx
    normT,normAlr,normAud,normC = calc_κ(iy,ix,T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir=dir)
    κ[iy,ix] = normT*normC/(normAlr*normAud)
  end
  return κ
end

function normalize!((Clu,Cru,Cld,Crd))
  ny,nx = size(Clu)
  for iy ∈ 1:ny, ix ∈ 1:nx
    Clu[iy,ix] /= norm(Clu[iy,ix])
    Cld[iy,ix] /= norm(Cld[iy,ix])
    Cru[iy,ix] /= norm(Cru[iy,ix])
    Crd[iy,ix] /= norm(Crd[iy,ix])
    normC4 = scalar(Clu[iy,ix]*Cld[iy,ix]*Cru[iy,ix]*Crd[iy,ix])
    normC4<0 ? normClu = -abs(normC4)^(1/4) : normClu = normC4^(1/4)
    Clu[iy,ix] /= normClu
    Cld[iy,ix] /= abs(normClu)
    Cru[iy,ix] /= abs(normClu)
    Crd[iy,ix] /= abs(normClu)
  end
end

function normalize!((Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir="left")
  normalize!((Clu,Cru,Cld,Crd))
  ny,nx = size(Clu)
  for iy ∈ 1:ny, ix ∈ 1:nx
    Al[iy,ix] /= norm(Al[iy,ix])
    Ar[iy,ix] /= norm(Ar[iy,ix])
    Au[iy,ix] /= norm(Au[iy,ix])
    Ad[iy,ix] /= norm(Ad[iy,ix])
    iyp,ixp = per(iy+1,ny),per(ix+1,nx)
    normAlr = scalar(Clu[iy,ix]*Cru[iy,ix]*Al[iy,ix]*Ar[iy,ix]*Cld[iyp,ix]*Crd[iyp,ix])
    normAlr<0 ? normAl = -abs(normAlr)^(1/2) : normAl = normAlr^(1/2)
    Al[iy,ix] /= normAl
    Ar[iy,ix] /= abs(normAl)
    normAud = scalar(Clu[iy,ix]*Cld[iy,ix]*Au[iy,ix]*Ad[iy,ix]*Cru[iy,ixp]*Crd[iy,ixp])
    normAud<0 ? normAu = -abs(normAud)^(1/2) : normAu = normAud^(1/2)
    Au[iy,ix] /= normAu
    Ad[iy,ix] /= abs(normAu)
  end
end

function leftright_move!(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir="left",maxdim=5)
  ny,nx = size(T)
  P = Vector{ITensor}(undef,ny)
  P⁻ = Vector{ITensor}(undef,ny)
  if dir=="left" || dir=="up"
    xrange = 1:nx
  elseif dir=="right" || dir=="down"
    xrange = per.(nx-1:-1:0,nx)
  end
  for ix = xrange
    ixm = per(ix-1,nx)
    ixp = per(ix+1,nx)
    ixpp = per(ix+2,nx)
    for iy = 1:ny
      iym = per(iy-1,ny)
      iyp = per(iy+1,ny)

      Cu = Al[iym,ix]*Clu[iym,ix]*Au[iym,ix]*T[iym,ix]*Au[iym,ixp]*T[iym,ixp]*Cru[iym,ixpp]*Ar[iym,ixpp]
      @assert order(Cu)==4
      Cd = Al[iy,ix]*Cld[iyp,ix]*Ad[iyp,ix]*T[iy,ix]*Ad[iyp,ixp]*T[iy,ixp]*Crd[iyp,ixpp]*Ar[iy,ixpp]
      @assert order(Cd)==4
      if dir=="left" || dir=="up"
        li = commonindex(Cru[iy,ixpp],Crd[iy,ixpp])
        si = commonindex(Au[iy,ixp],Ad[iy,ixp])
      elseif dir=="right" || dir=="down"
        li = commonindex(Clu[iy,ix],Cld[iy,ix])
        si = commonindex(Au[iy,ix],Ad[iy,ix])
      end
      Cup = prime(Cu,(li,si))
      ρ = Cd*Cup
      if dir=="left" || dir=="right"
        utags = "$dir,link,x=$ixp,y=$iy"
      elseif dir=="up" || dir=="down"
        utags = "$dir,link,x=$iy,y=$ixp"
      end
      U,S,Vh,spec,u,v = svd(ρ, (li,si); utags=utags, vtags="tmp", maxdim=maxdim, cutoff=0.0)
      V = dag(Vh)
      U *= δ(u,v)
      invsqrtS = S
      for i = 1:dim(u)
        invsqrtS[i,i] = inv(sqrt(S[i,i]))
      end
      P[iy] = Cup*V*invsqrtS
      P⁻[iy] = Cd*dag(U)*invsqrtS
    end
    for iy = 1:ny
      iym = per(iy-1,ny)
      iyp = per(iy+1,ny)
      if dir=="left" || dir=="up"
        Al[iy,ixp] = Al[iy,ix]*P[iy]*T[iy,ix]*P⁻[iyp]
        Clu[iy,ixp] = Clu[iy,ix]*Au[iy,ix]*P⁻[iy]
        Cld[iy,ixp] = Cld[iy,ix]*Ad[iy,ix]*P[iy]
      elseif dir=="right" || dir=="down"
        Ar[iy,ixp] = Ar[iy,ixpp]*P[iy]*T[iy,ixp]*P⁻[iyp]
        Cru[iy,ixp] = Cru[iy,ixpp]*Au[iy,ixp]*P⁻[iy]
        Crd[iy,ixp] = Crd[iy,ixpp]*Ad[iy,ixp]*P[iy]
      end
    end
  end
  normalize!((Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))
end

function swapdiag(M)
  Mp = permutedims(M,[2,1])
  ny,nx = size(Mp)
  for iy = 1:ny, ix = 1:nx
    Mp[iy,ix] = M[ix,iy]
  end
  return Mp
end
function rotate_environment(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))
  return swapdiag(T),swapdiag.((Clu,Cld,Cru,Crd)),swapdiag.((Au,Ad,Al,Ar))
end

printdiv() = println("\n****************************************")
printstepdiv() = println("\n##################################################")

function sweepsdims(stepsizes::Vector{Int},dims::Vector{Int})
  nstep = length(stepsizes)
  maxdims = zeros(Int,stepsizes[end])
  for i = 1:stepsizes[1]
    maxdims[i] = dims[1]
  end
  for j = 2:nstep
    for i = stepsizes[j-1]+1:stepsizes[j]
      maxdims[i] = dims[j]
    end
  end
  return maxdims
end

function check_environment(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))
  ny,nx = size(T)
  for iy = 1:ny
    for ix = 1:nx
      @assert order(Clu[iy,ix])==2
      @assert order(Cru[iy,ix])==2
      @assert order(Cld[iy,ix])==2
      @assert order(Crd[iy,ix])==2
      @assert order(Al[iy,ix])==3
      @assert order(Ar[iy,ix])==3
      @assert order(Au[iy,ix])==3
      @assert order(Ad[iy,ix])==3
      @assert order(T[iy,ix])==4
      @assert order(T[iy,ix])==4
      @assert order(T[iy,ix])==4
      @assert order(T[iy,ix])==4
      @assert length(commoninds(Clu[iy,ix],Cru[iy,ix]))==1
      @assert length(commoninds(Clu[iy,ix],Cld[iy,ix]))==1
      @assert length(commoninds(Cld[iy,ix],Crd[iy,ix]))==1
      @assert length(commoninds(Cru[iy,ix],Crd[iy,ix]))==1
    end
  end
end

function ctmrg(T::Matrix{ITensor},
               (Clu, Cru, Cld, Crd),
               (Al, Ar, Au, Ad);
               verbose = false)
  ny,nx = size(T)

  verbose && println("Original:")
  verbose && @show calc_κ(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))

  nstep = 1000
  maxdim = 10
  dirs = ["left","up","right","down"]
  for ctmrg_step = 1:nstep
    verbose && printstepdiv()
    dir = dirs[per(ctmrg_step,length(dirs))]
    verbose && @show ctmrg_step,dir

    if dir=="left" || dir=="right"
      leftright_move!(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir=dir,maxdim=maxdim)
    elseif dir=="up" || dir=="down"
      T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad) = rotate_environment(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))
      leftright_move!(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir=dir,maxdim=maxdim)
      T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad) = rotate_environment(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))
    end

    check_environment(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad))

    verbose && @show Mκ = calc_κ(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir=dir)
    verbose && @show abs(prod(vec(Mκ)))^(1/(nx*ny))
  end
  Mκ = calc_κ(T,(Clu,Cru,Cld,Crd),(Al,Ar,Au,Ad);dir=dir)
  κave = abs(prod(vec(Mκ)))^(1/(nx*ny))
  return κave
end


text:
using RAFF

using Test
using DelimitedFiles
using Distributed
using Random
using SharedArrays
using Logging

include("test_raff.jl")

include("test_solvers.jl")

include("test_utils.jl")

include("test_generator.jl")

include("test_parallel.jl")

include("test_multivariate.jl")

include("test_integration.jl")


text:
using MAT
using MatrixProductStates
BLAS.set_num_threads(1)

#Spin model parameters
const na = 70
const od = 33
const ad = 2.5*60/na #atomic distance in micro m
const cloud_sig = 23.6344/ad
const cloud_coup = exp.(-((1:na)-(na+1)/2).^2/2/cloud_sig^2)
const gam_1d = cloud_coup/sum(cloud_coup)/2*od
const gam_eg = 1.0
const del_p = 0.0
const k_wg = 0.5*pi
const rj = collect(1.0:na) #sort!(rand(na)*na)#

#Pulse shape parameters
const k_in = k_wg
const gam_exp = 2*pi*6.065e6
const f_amp_exp = sqrt.([0.0589,0.0721,0.0995,0.1637,0.3081,0.6395,1.2121,1.8898,
    2.6039,4.2050,7.1291,10.4823,17.0816,27.1592,42.8378,71.7874]/gam_exp/1e-6)
const f_amp = f_amp_exp[1]

#Rydberg specific parameters
const om = 10.0/6.065/2.0
const k_c = 0.0
const del_s = 0.0
const gam_es = 0.0

# rmax = 20 na = 70
const uu = [524.328 + 89.9671im, 524.328 - 89.9671im, 649.737 + 219.429im, 649.737 - 219.429im, 488.928]
const lam = [-0.228929 + 0.189648im, -0.228929 - 0.189648im, 0.146279 + 0.297638im, 0.146279 - 0.297638im, 0.480032]

# rmax = 30 na = 60
#uu = [1055.32, 1156.39 + 264.599im, 1156.39 - 264.599im, 781.661]
#lam = [-0.220397, 0.0354409 + 0.244048im, 0.0354409 - 0.244048im, 0.396262]
#uu = [15158.7, 11895.1 - 1929.93im, 11895.1 + 1929.93im, 2399.43, 16.4076]
#lam = [-0.103573, 0.015448 + 0.126927im, 0.015448 - 0.126927im, 0.30542, 0.606657]
#uu = [90866.8, 59173.0 - 18223.1im, 59173.0 + 18223.1im, 5257.17, 107.47, 1.26558]
#lam = [-0.0606795, 0.0112804 + 0.0784944im, 0.0112804 - 0.0784944im, 0.246432, 0.48336, 0.720283]

# rmax = 20 na = 60
#const uu = [0.704528 + 135.691im, 0.704528 - 135.691im, 225.917]
#const lam = [0.226209 + 0.32374im, 0.226209 - 0.32374im, 0.475997]

#uu = [ 2255.07, 686.082 + 91.2384im, 686.082 - 91.2384im, 477.07]
#lam = [-0.0998312, 0.0646318 + 0.25377im, 0.0646318-0.25377im, 0.424938]
#uu = [-508.38 - 1018.61im, -508.38 + 1018.61im, -1546.1 - 1942.74im, -1546.1 + 1942.74im,
#        28.5203]
#lam = [ -0.0374226 + 0.193987im, -0.0374226 - 0.193987im, 0.250214 + 0.0887921im,
#       0.250214 - 0.0887921im, 0.58104]
#uu = [-658.963 + 91.3732im, -658.963 - 91.3732im, 309.325 + 1085.49im, 309.325 - 1085.49im,
#       709.204, 4.33617]
#lam = [-0.0921073 + 0.171392im, -0.0921073 - 0.171392im, 0.124872 + 0.194013im,
#       0.124872 - 0.194013im, 0.382177, 0.670268]

# rmax = 20 na = 50
#uu = [1740.29, 2277.16 + 454.662im, 2277.16 - 454.662im, 945.959, 6.06471]
#lam = [-0.126487, 0.0183872 + 0.140738im, 0.0183872 - 0.140738im, 0.296793, 0.601841]
#uu = [ 57.1125, 170.91 + 211.667im, 170.91 - 211.667im, 336.281]
#lam = [ -0.298989, 0.0548962 + 0.260436im, 0.0548962 - 0.260436im, 0.382366]

# rmax = 20 na = 40
#uu = [19014.9 - 3038.05im, 19014.9 + 3038.05im, 2635.02, 47.4863, 0.531659]
#lam = [-0.0167917 + 0.0377784im, -0.0167917 - 0.0377784im, 0.155726, 0.386972, 0.653397]

# rmax = 20 na = 30
#uu = [-2652.81-1864.93im, -2652.81 + 1864.93im, 274.344, 5.91125, 0.0704647]
#lam = [0.0313486 + 0.0360902im, 0.0313486 - 0.0360902im, 0.179561, 0.407673, 0.666632]

# rmax = 20 na = 20
#uu = [1350.25, 707.234, 15.7581, 0.426333, 0.0053846]
#lam = [-0.0138335, 0.0500349, 0.19712, 0.418572, 0.672633]

# pure dephasing
const gam_gg = 0.0
const gam_ss = 40e3/6.065e6
const gam_ee = 0.0

# simulation parameters
const setrand = 1 # random number generator seed
const dt = 0.01
const t_fin = 10000.0
const d = 3
const d_max = 10
const measure_int = 5
const path_data = "/home/jdouglas/data/" # save directory
const base_filename = string(path_data, "Ryd_Jumps_N", na, "_D", d_max,
    "_Tf", t_fin, "_f", round(f_amp, 3), "_dt", dt, "_traj", setrand)

# envelope of input pulse
function f(t)

   f_amp

end

# local spin operators
const hgg = [1 0 0; 0 0 0; 0 0 0]
const hee = [0 0 0; 0 1 0; 0 0 0]
const hss = [0 0 0; 0 0 0; 0 0 1]
const hge = [0 1 0; 0 0 0; 0 0 0]
const heg = hge'
const hes = [0 0 0; 0 0 1; 0 0 0]
const hse = hes'
const id = eye(3)

# types used in mps representation
const TN = Complex{Float64}
const TA = Array

function time_evolve()

    # create jump and measurement operators
    jumpleft = makempo(TN, TA, [jj->id jj->im*sqrt(gam_1d[jj]/2)*exp(im*k_wg*rj[jj])*hge;
                                jj->0 jj->id], na, d)
    jumpright = makempo(TN, TA, [jj->id jj->im*sqrt(gam_1d[jj]/2)*exp(-im*k_wg*rj[jj])*hge;
                                 jj->0 jj->id], na, d)
    jumpright[1][1, :, 2, :] = f(0.0)*id + im*sqrt(gam_1d[1]/2)*exp(-im*k_wg*rj[1])*hge
    ir_mpo = applyMPOtoMPO(jumpright, conj_mpo(jumpright))
    ir2_mpo = applyMPOtoMPO(applyMPOtoMPO(jumpright, ir_mpo), conj_mpo(jumpright))

    # step times and measurement times
    t = 0.0:dt:t_fin
    t_m = t[1:measure_int:end]
    tstep = length(t) - 1

    # initialize output measures
    p_coh = zeros(tstep)
    p_sum = zeros(tstep)
    p_acc = zeros(tstep, 2na + 4)
    e_pop = zeros(na, length(t_m))
    s_pop = zeros(na, length(t_m))
    epop = zeros(na)
    spop = zeros(na)
    I_r = zeros(TN, length(t_m))
    I2_r = zeros(TN, length(t_m))
    times = zeros(tstep, 4)
    t_r = Float64[]
    t_l = Float64[]
    t_eg = Float64[]
    pos_eg = Float64[]
    t_ss = Float64[]
    pos_ss = Float64[]
    cont = zeros(Int64, 4)

    #Construct left-canonical MPS for the initial state and temp arrays for time evolution
    A, dims =  mpsgroundstate(TN, TA, na, d_max, d)
    A_coh = similar(A)
    A_r = similar(A)
    A_l = similar(A)
    A1 = similar(A)
    A2 = similar(A)
    A3 = similar(A)
    A_coh .= copy.(A)
    A_r .= copy.(A)
    A_l .= copy.(A)
    A1 .= copy.(A)
    A2 .= copy.(A)
    A3 .= copy.(A)
    mpo_size = 4 + length(lam)
    envop = build_env(TN, TA, dims, dims, ones(na + 1)*mpo_size)
    envop_jump = build_env(TN, TA, dims, dims, ones(na + 1)*2)
    env1 = build_env(TN, TA, dims, dims)
    env2 = build_env(TN, TA, dims, dims)
    env3 = build_env(TN, TA, dims, dims)

    # initial measurements
    I_r[1] = scal_op_prod(A, ir_mpo, A)
    I2_r[1] = scal_op_prod(A, ir2_mpo, A)
    measure_excitations!((@view e_pop[:, 1]), (@view s_pop[:, 1]), A)

    # random jump variables
    srand(setrand)
    r1 = rand(tstep)
    r2 = rand(tstep)

    # create evolution operators for Runge-Kutta algorithm
    H1 = rydberg_hamiltonian(TN, TA, 1, dt/2, t[1])
    H2 = rydberg_hamiltonian(TN, TA, 0, 1, t[1] + dt/2)
    H3 = rydberg_hamiltonian(TN, TA, 1, dt/2, t[2])

    H = [H1, H2, H3]

    println("Initialisation complete")

    #EVOLUTION
    for i = 1:tstep
        time_in = time()

        # update time evolution operators
        update_rydberg_hamiltonian!(H1, 1, dt/2, t[i])
        update_rydberg_hamiltonian!(H2, 0, 1, t[i] + dt/2)
        update_rydberg_hamiltonian!(H3, 1, dt/2, t[i+1])

        # hoherent time evolution probability
        RK4stp_apply_H_half!(dt, H, A_coh, A, A1, A2, A3, env1, env2, env3, envop)
        p_coh[i] = normalize_lo!(A_coh)
        times[i, 1] = time() - time_in
        if r1[i] < p_coh[i]
            A .= copy.(A_coh)
            times[i,2] = 0
        else
            #Update waveguide jump operator
            jumpright[1][1,:,2,:] = f(t[i])*id + im*sqrt(gam_1d[1]/2)*exp(-im*k_wg*rj[1])*hge

            #Jump probabilities: decay into the waveguide
            compress_var_apply_H!(A_r, envop_jump, A, A, jumpright, 1)
            p_r = normalize_lo!(A_r)
            compress_var_apply_H!(A_l, envop_jump, A, A, jumpleft, 1)
            p_l = normalize_lo!(A_l)

            #Jump probabilities: decay outside the waveguide
            measure_excitations!(epop, spop, A)
            p_acc[i, 1] = 0.0
            for j = 1:na
                p_acc[i, j + 1] = p_acc[i, j] + dt*gam_eg*epop[j]
            end
            for j = 1:na
                p_acc[i, na + j + 1] = p_acc[i, na + j] + dt*gam_ss*spop[j]
            end
            p_acc[i, 2na + 2] = p_acc[i, 2na + 1] + dt*p_r
            p_acc[i, 2na + 3] = p_acc[i, 2na + 2] + dt*p_l
            p_sum[i] = p_acc[i, 2na + 3]
            println(p_sum[i] + p_coh[i])
            println(t[i])
            p_acc[i, :] = p_acc[i, :]/p_acc[i, 2na + 3]

            #Do jumps according to above probabilities
            if r2[i] < p_acc[i, na + 1]
                for j = 1:na
                    if r2[i] < p_acc[i, j + 1]
                        apply_site_operator!(A[j], hge)
                        leftorthonormalizationQR!(A)
                        push!(t_eg, t[i])
                        push!(pos_eg, j)
                        cont[1] += 1
                        break
                    end
                end
            elseif r2[i] < p_acc[i, 2na + 1]
                for j = 1:na
                    if r2[i] < p_acc[i, na + j + 1]
                        apply_site_operator!(A[j], hss)
                        leftorthonormalizationQR!(A)
                        push!(t_ss, t[i])
                        cont[2] += 1
                        break
                    end
                end
            elseif r2[i] < p_acc[i, 2na + 2]
                A .= copy.(A_r)
                push!(t_r, t[i])
                cont[3] += 1
            else
                A .= copy.(A_l)
                push!(t_l, t[i])
                cont[4] += 1
            end
            times[i, 2] = time() - times[i, 1] - time_in

        end

        # measurements (every measure_int time steps)
        if rem(i, measure_int) == 0

            mind = div(i, measure_int)+1
            # atomic state populations
            measure_excitations!((@view e_pop[:, mind]), (@view s_pop[:, mind]), A)
            times[i, 3] = time() - times[i, 2]  - times[i, 1] - time_in

            # output right field update
            jumpright[1][1, :, 2, :] = f(t[i + 1])*id +
                    im*sqrt(gam_1d[1]/2)*exp(-im*k_wg*rj[1])*hge

            ir_mpo[1] =  apply_site_MPOtoMPO(jumpright[1], conj_site_mpo(jumpright[1]))
            ir2_mpo[1] =  apply_site_MPOtoMPO(
                apply_site_MPOtoMPO(jumpright[1], ir_mpo[1]),
                conj_site_mpo(jumpright[1]))

            I_r[mind] = scal_op_prod(A, ir_mpo, A)
            I2_r[mind] = scal_op_prod(A, ir2_mpo, A)

            times[i,4] = time() - times[i, 3] - times[i, 2] - times[i, 1] - time_in

        end

        # saving temp file
        if rem(i, measure_int*2000) == 0

            write_data_file(string(base_filename,"_temp.mat"), t, t_m, e_pop,
                            s_pop, p_coh, p_acc, p_sum, I_r, I2_r, cont, t_r,
                            t_l, t_eg, t_ss, pos_eg, A, times)

        end

    end

    # saving output file
    write_data_file(string(base_filename,".mat"), t, t_m, e_pop,
                    s_pop, p_coh, p_acc, p_sum, I_r, I2_r, cont, t_r,
                    t_l, t_eg, t_ss, pos_eg, A, times)

end

function write_data_file(filename, t, t_m, e_pop, s_pop, p_coh, p_acc,
                         p_sum, I_r, I2_r, cont, t_r, t_l, t_eg, t_ss,
                         pos_eg, A, times)

    file = matopen(string(filename),"w")
    write(file, "TN", string(TN))
    write(file, "TA", string(TA))
    write(file, "na", na)
    write(file, "gam_1d", gam_1d)
    write(file, "gam_eg", gam_eg)
    write(file, "del_p", del_p)
    write(file, "k_wg", k_wg)
    write(file, "rj", rj)
    write(file, "k_in", k_in)
    write(file, "f_amp", f_amp)
    write(file, "f", f.(t_m))
    write(file, "om", om)
    write(file, "k_c", k_c)
    write(file, "del_s", del_s)
    write(file, "gam_es", gam_es)
    write(file, "uu", uu)
    write(file, "lam", lam)
    write(file, "gam_gg", gam_gg)
    write(file, "gam_ss", gam_ss)
    write(file, "gam_ee", gam_ee)
    write(file, "dt", dt)
    write(file, "t_fin", t_fin)
    write(file, "d_max", d_max)
    write(file, "measure_int", measure_int)

    write(file, "t", collect(t))
    write(file, "t_m", collect(t_m))
    write(file, "e_pop", e_pop)
    write(file, "s_pop", s_pop)
    write(file, "p_coh", p_coh)
#    write(file, "p_acc", p_acc)
    write(file, "p_sum", p_sum)
    write(file, "I_r", I_r)
    write(file, "I2_r", I2_r)
    write(file, "times", sum(times,1))
    write(file, "cont", cont)
    write(file, "t_r", t_r)
    write(file, "t_l", t_l)
    write(file, "t_eg", t_eg)
    write(file, "t_ss", t_ss)
    write(file, "pos_eg", pos_eg)
    write(file, "A", collect.(A))

    close(file)

end

function rydberg_hamiltonian(::Type{TN}, ::Type{TA}, con, delt, time) where {TN, TA}

    H = Array{TA{TN, 4}, 1}(na)

    no_exp = length(lam)
    max_ind = 4 + no_exp
    drj = diff(rj)
    ph = exp.(im*k_wg*drj)
    cp = sqrt.(delt*gam_1d/2)

    H[1] = zeros(TN, 1, 3, max_ind, 3)
    H[1][1, :, 1, :] = id
    H[1][1, :, 2, :] = -cp[1]*ph[1]*heg
    H[1][1, :, 3, :] = -cp[1]*ph[1]*hge
    for exp_ind = 1:no_exp
        H[1][1, :, 3 + exp_ind, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*hss
    end
    H[1][1, :, max_ind, :] = local_h(1, con, delt, time)

    for jj = 2:(na - 1)
        H[jj] = zeros(TN, max_ind, 3, max_ind, 3)
        H[jj][1, :, 1, :] = id
        H[jj][1, :, 2, :] = -cp[jj]*ph[jj]*heg
        H[jj][1, :, 3, :] = -cp[jj]*ph[jj]*hge
        H[jj][1, :, max_ind, :] = local_h(jj, con, delt, time)
        H[jj][2, :, max_ind, :] = cp[jj]*hge
        H[jj][3, :, max_ind, :] = cp[jj]*heg
        H[jj][2, :, 2, :] = ph[jj]*id
        H[jj][3, :, 3, :] = ph[jj]*id
        for exp_ind = 1:no_exp
            H[jj][1, :, 3 + exp_ind, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*hss
            H[jj][3 + exp_ind, :, max_ind, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*hss
            H[jj][3 + exp_ind, :, 3 + exp_ind, :] = lam[exp_ind]*id
        end
        H[jj][max_ind, :, max_ind, :] = id
    end

    H[na] = zeros(TN, max_ind, 3, 1, 3)
    H[na][1, :, 1, :] = local_h(na, con, delt, time)
    H[na][2, :, 1, :] = cp[na]*hge
    H[na][3, :, 1, :] = cp[na]*heg
    for exp_ind = 1:no_exp
        H[na][3 + exp_ind, :, 1, :] = sqrt(-im*delt*uu[exp_ind]*lam[exp_ind])*hss
    end
    H[na][max_ind, :, 1, :] = id

    return H

end

function local_h(jj, con, delt, time)

    con*id/na + delt*((im*del_p -
        (gam_eg + gam_ee + gam_1d[jj])/2)*hee +
        im*f(time)*sqrt(gam_1d[jj]/2)*exp(im*k_in*rj[jj])*heg -
        1/(2na)*abs2(f(time))*id + (im*del_s - gam_ss/2)*hss +
        im*om*(exp(im*k_c*rj[jj])*hes + exp(-im*k_c*rj[jj])*hse))

end

function update_rydberg_hamiltonian!(H, con, delt, time)

    no_exp = length(lam)
    max_ind = 4 + no_exp

    H[1][1, :, max_ind, :] = local_h(1, con, delt, time)

    for jj = 2:(na - 1)
        H[jj][1, :, max_ind, :] = local_h(jj, con, delt, time)
    end

    H[na][1, :, 1, :] = local_h(na, con, delt, time)

end

# calculate local expectation values, assuming mps is already left normalized
function measure_excitations!(e_pop, s_pop, A)

    # initialize right environment
    env = ones(1, 1)

    for n = na:(-1):1
        Aenv = prod_LR(A[n], env)
        e_pop[n] = real(local_exp(Aenv, A[n], hee))
        s_pop[n] = real(local_exp(Aenv, A[n], hss))
        env = update_renv(A[n], Aenv)
    end

end

time_evolve()


text:
#=
Benchmarking routines for characters and strings

Copyright 2017-2020 Gandalf Software, Inc., Scott P. Jones
Licensed under MIT License, see LICENSE.md

Functions to load up sample books
=#

const inppath = "textsamples"
const gutpath = "gutenberg"
const smppath = "samples"

const gutenbergbooks =
    (("2600/2600-0",   "English"), # War & Peace, some other languages in quotes
     ("1400/1400-0",   "English"), # Great Expectations, uses Unicode quotes
     ("42286/42286-0", "Hungarian"),
     ("31536/31536-0", "Polish"),
     ("32941/32941-0", "Japanese"),
     ("24264/24264-0", "Chinese"),
     ("40687/40687-0", "Telugu"), # Third most spoken in India, official
     ("50513/50513-0", "French"),
     ("43007/43007-0", "Arabic"),
     ("38496/38496-8", "Portuguese"), # Latin1
     ("2000/2000-0",   "Spanish"), # Don Quijote
     ("48750/48750-8", "Swedish"), # Latin1
#     ("48322/48322-8", "German"),  # Latin1
     )

const downloadedbooks = (
#                         ("LYSAIa GORA DIeVICh'Ia - SIeRGIeI GOLOVAChIoV.txt", "Russian"),
                         )

getdefdir(dir)::String = dir === nothing ? homedir() : dir

function find_beg(lines)
    CSE = "Character set encoding: "
    csebeg = sizeof(CSE) + 1
    last = length(lines)
    cse = ""
    ln = 0
    while ln < last
        l = lines[ln += 1]
        if starts_with(l, CSE)
            cse = l[csebeg:end]
        elseif sizeof(l) > 41 && starts_with(l, "***") && ends_with(l, "***") &&
            occurs_in(" PROJECT GUTENBERG EBOOK", l) &&
            occurs_in("START OF TH", l)
            # Skip over empty lines
            while ln < last && (l = lines[ln += 1]) == "" ; end
            if starts_with(l, "Produced by ")
                # Skip over non-empty lines after "Produced by"
                while ln < last && lines[ln += 1] != "" ; end
                # Skip over empty lines after "Produced by"
                while ln < last && lines[ln += 1] == "" ; end
            end
            return (ln, cse)
        end
    end
    (1, cse)
end

function find_end(lines, beg)
    ln = last = length(lines)
    while ln > beg
        l = lines[ln]
        if sizeof(l) > 41 && starts_with(l, "***") && ends_with(l, "***") &&
            occurs_in(" PROJECT GUTENBERG EBOOK", l) &&
            occurs_in("END OF TH", l)
            #print("Found END OF at line $ln")
            while ln > beg && (l = lines[ln -= 1]) == "" ; end
            #println(" => $ln")
            if starts_with(l, "End of ") && occurs_in("Project Gutenberg", l)
                #print("Found End of at line $ln")
                while ln > beg && lines[ln -= 1] == "" ; end
                #println(" => $ln")
            end
            return ln
        end
        ln -= 1
    end
    last
end

"""
Load books from Project Gutenberg site, removing lines added at beginning and end that
are not part of the book, as much as possible
"""
function load_gutenberg!(books, list, dict, gutenbergdir)
    mkpath(gutenbergdir)
    for (nam, lang) in list
        cnt = get(dict, lang, 0)
        dict[lang] = cnt + 1
        outnam = cnt == 0 ? "$lang.txt" : "$lang-$cnt.txt"
        lname = joinpath(gutenbergdir, outnam)
        download(string("http://www.gutenberg.org/files/", nam, ".txt"), lname)
        println("Saved to: ", lname)
        lines = readlines(lname)
        (beg, cse) = find_beg(lines)
        lst = find_end(lines, beg)
        if cse == "ISO-8859-1"
            filt = [convert(String, Str(LatinCSE, lines[ln])) for ln in beg:lst]
        elseif cse == "UTF-8" || cse == ""
            filt = lines[beg:lst]
        else
            error("Unknown character set encoding: $cse")
        end
        push!(books, (outnam, filt))
    end
    books
end

function remove_empty(lines)
    # Eliminate empty lines
    len = length(lines)
    out = Vector{String}()
    sizehint!(out, len)
    for l in lines
        is_empty(l) || push!(out, l)
    end
    out
end

"""
load_books(; dir=nothing)

Loads a set of books from a local directory, and downloads a set of books from Project Gutenberg
Returns them as a dictionary with names -> vectors of strings
"""
function load_books(; dir::Any=nothing)
    defdir = getdefdir(dir)
    inputdir = joinpath(defdir, inppath)
    dict = Dict{String,Int}()
    books = Vector{Tuple{String, Vector{String}}}()
    for (nam, lang) in downloadedbooks
        cnt = get(dict, lang, 0)
        dict[lang] = cnt + 1
        outnam = cnt == 0 ? "$lang.txt" : "$lang-$cnt.txt"
        push!(books, (outnam, readlines(joinpath(inputdir, nam))))
    end
    load_gutenberg!(books, gutenbergbooks, dict, joinpath(defdir, gutpath))
end

"""
save_books(books; dir=nothing)

Saves the collection of downloaded books into the given directory, in a "samples" subdirectory.
If the directory is not set, it will default to the user's home directory
"""
function save_books(books; dir::Any=nothing)
    sampledir = joinpath(getdefdir(dir), smppath)
    mkpath(sampledir)
    for (nam, book) in books
        outnam = joinpath(sampledir, nam)
        open(outnam, "w") do io
            for lin in book
                println(io, lin)
            end
            println("Saved $nam")
        end
    end
end