import ..Continuations: step!, new_sweep!
using ..Continuations: AbstractContinuationSolver, ContinuationOptions,
    SweepSetup

abstract type BifurcationCache{PC} <: AbstractContinuationCache{PC} end

struct BifurcationSolver{
        R <: ContinuationSolver,
        P <: BifurcationProblem,
        C <: BifurcationCache,
        S <: BifurcationSolution,
        } <: AbstractContinuationSolver
    super::R
    prob::P
    opts::ContinuationOptions
    cache::C
    sol::S
end

function BifurcationSolver(prob::BifurcationProblem,
                           opts::ContinuationOptions)
    super = ContinuationSolver(prob, opts)
    cache = BifurcationCache(prob, super.cache)
    sol = BifurcationSolution(super.sol, sweeptype(prob, super))
    return BifurcationSolver(super, prob, opts, cache, sol)
end

TimeKind(::Type{<: BifurcationSolver{_, P}}) where {_, P} = TimeKind(P)

function step!(solver::BifurcationSolver)
    step!(solver.super)
    # calling [[../continuations/solver.jl::step!]]

    analyze!(solver.cache, solver.opts)
    record!(solver.sol, solver.cache)
    check_sweep_length(solver.sol.sweeps[end])
    post_step!(solver)  # TODO: remove
end

# This exists only for quick-and-dirty Bogdanov-Takens point handling
# in codim2 solver.  See: [[../codim2/solver.jl::post_step!]]
function post_step!(::Any) end

function new_sweep!(solver::BifurcationSolver, setup::SweepSetup)
    new_sweep!(solver.super, setup)
    # calling [[../continuations/solver.jl::new_sweep!]]

    allocate_sweep!(solver.sol, as(solver, ContinuationSolver))

    for u in setup.past_points
        re_analyze!(solver, u)
    end
    re_analyze!(solver, setup.u0)
    check_sweep_length(solver.sol.sweeps[end])
end
# TODO: `new_sweep!(solver.super, setup)` sets up `solver.super.cache`
# but `re_analyze!(solver, u)` rewrites the cache.  It causes no
# problem at the moment since the last `re_analyze!(solver, setup.u0)`
# set the cache back to the first state.

function allocate_sweep!(sol::BifurcationSolution, solver::ContinuationSolver)
    super = as(sol, ContinuationSolution)
    sweep = BifurcationSweep(super.sweeps[end], solver)
    push!(sol.sweeps, sweep)
end

function push_special_interval!(sweep::BifurcationSweep,
                                cache::BifurcationCache)
    push_special_interval!(sweep,
                           cache.point_type,
                           as(cache, ContinuationCache).J)
end

function push_special_interval!(sweep::BifurcationSweep,
                                point_type,
                                J1)
    super = as(sweep, ContinuationSweep)
    point = SpecialPointInterval(
        timekind(sweep),
        point_type,
        length(sweep) - 1,  # point_index
        super.u[end - 1],   # u0
        super.u[end],       # u1
        J1,
        WeakRef(sweep),
    )
    push!(sweep.special_intervals, point)
end
# `push_special_interval!` is relying on that it is called **after**
# the new point (`u1` in above) is `push!`ed to `super.u`.  With this
# assumption, `point_index` is same as the index used for
# `SimpleBifurcationInterval` in `push_point!` (see:
# [[../continuations/solution.jl::push_point!]]).  Note that `super.u`
# is updated via `push_point!` via `step!(::ContinuationSolver)`
# before `push_special_interval!` is called via `record!`.

"""
    is_special(point_type) :: Bool
"""
is_special(point_type::T) where T = point_type != regular_point(T)

function record!(sol::BifurcationSolution, cache::BifurcationCache)
    super = as(cache, ContinuationCache)
    sweep = sol.sweeps[end]
    push!(sweep.jacobians, copy(super.J))
    push!(sweep.eigvals, copy(cache.eigvals))

    if is_special(cache.point_type)
        push_special_interval!(sweep, cache)
    end
end


# edyna Julia module begins -->
module edyna

# ellipsoids data
ellcount = 0
ellcenter = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
ellradii = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
ellorient = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
             Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
             Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
elldefo = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
           Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
           Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
elldensity = Array{Float32}(undef, ellcount);
ellyoung = Array{Float32}(undef, ellcount);
ellpoisson = Array{Float32}(undef, ellcount);
ellmass = Array{Float32}(undef, ellcount);
elleuler = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
            Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
            Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
ellinveuler = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
               Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
               Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
ellinertia = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
              Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
              Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
ellinvinertia = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
                 Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount),
                 Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
elllinear = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
ellangular = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]

# insert ellipsoid
function ellipsoid(x, a, b, c, R, density, young = 1E9, poisson = 0.25, linear = [0.0 0.0 0.0], angular = [0.0 0.0 0.0])
  push!(ellcenter[1],x[1])
  push!(ellcenter[2],x[2])
  push!(ellcenter[3],x[3])
  push!(ellradii[1],a)
  push!(ellradii[2],b)
  push!(ellradii[3],c)
  push!(ellorient[1],R[1,1])
  push!(ellorient[2],R[2,1])
  push!(ellorient[3],R[3,1])
  push!(ellorient[4],R[1,2])
  push!(ellorient[5],R[2,2])
  push!(ellorient[6],R[3,2])
  push!(ellorient[7],R[1,3])
  push!(ellorient[8],R[2,3])
  push!(ellorient[9],R[3,3])
  push!(elldefo[1],1.0)
  push!(elldefo[2],0.0)
  push!(elldefo[3],0.0)
  push!(elldefo[4],0.0)
  push!(elldefo[5],1.0)
  push!(elldefo[6],0.0)
  push!(elldefo[7],0.0)
  push!(elldefo[8],0.0)
  push!(elldefo[9],1.0)

  (mass, volume, E, J) = ellipsoid_properties(a, b, c, R, density)

  push!(elldensity, density)
  push!(ellyoung, young)
  push!(ellpoisson, poisson)
  push!(ellmass, mass)
  push!(elleuler[1], E[1,1])
  push!(elleuler[2], E[2,1])
  push!(elleuler[3], E[3,1])
  push!(elleuler[4], E[1,2])
  push!(elleuler[5], E[2,2])
  push!(elleuler[6], E[3,2])
  push!(elleuler[7], E[1,3])
  push!(elleuler[8], E[2,3])
  push!(elleuler[9], E[3,3])
  invE = inv(E)
  push!(ellinveuler[1], invE[1,1])
  push!(ellinveuler[2], invE[2,1])
  push!(ellinveuler[3], invE[3,1])
  push!(ellinveuler[4], invE[1,2])
  push!(ellinveuler[5], invE[2,2])
  push!(ellinveuler[6], invE[3,2])
  push!(ellinveuler[7], invE[1,3])
  push!(ellinveuler[8], invE[2,3])
  push!(ellinveuler[9], invE[3,3])
  push!(ellinertia[1], J[1,1])
  push!(ellinertia[2], J[2,1])
  push!(ellinertia[3], J[3,1])
  push!(ellinertia[4], J[1,2])
  push!(ellinertia[5], J[2,2])
  push!(ellinertia[6], J[3,2])
  push!(ellinertia[7], J[1,3])
  push!(ellinertia[8], J[2,3])
  push!(ellinertia[9], J[3,3])
  invJ = inv(J)
  push!(ellinvinertia[1], invJ[1,1])
  push!(ellinvinertia[2], invJ[2,1])
  push!(ellinvinertia[3], invJ[3,1])
  push!(ellinvinertia[4], invJ[1,2])
  push!(ellinvinertia[5], invJ[2,2])
  push!(ellinvinertia[6], invJ[3,2])
  push!(ellinvinertia[7], invJ[1,3])
  push!(ellinvinertia[8], invJ[2,3])
  push!(ellinvinertia[9], invJ[3,3])

  push!(elllinear[1],linear[1])
  push!(elllinear[2],linear[2])
  push!(elllinear[3],linear[3])
  push!(ellangular[1],angular[1])
  push!(ellangular[2],angular[2])
  push!(ellangular[3],angular[3])

  global ellcount
  ellcount += 1
end

# triangles data
tricount = 0
trip = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
triq = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
trir = [Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount), Array{Float32}(undef, ellcount)]
trilinear = Array{Function}(undef, ellcount) # [vx, vy, vz] = linear[i](t)
triangular = Array{Function}(undef, ellcount) # [ox, oy, oz] = angular[i](t)

# insert an obstacle
# triangles format: [[[p1, q1, r1], [p2, q2, r2], [p3, q3, r3]], [[p1, q1, r1], [p2, q2, r2], [p3, q3, r3]], ...]
zerofunc(t) = Float32[0.0; 0.0; 0.0]
function obstacle(triangles, linear = zerofunc, angular = zerofunc)
  for i = 1:size(triangles)[1]
    push!(trip[1], triangles[i][1][1])
    push!(trip[2], triangles[i][1][2])
    push!(trip[3], triangles[i][1][3])
    push!(triq[1], triangles[i][2][1])
    push!(triq[2], triangles[i][2][2])
    push!(triq[3], triangles[i][2][3])
    push!(trir[1], triangles[i][3][1])
    push!(trir[2], triangles[i][3][2])
    push!(trir[3], triangles[i][3][3])
    push!(trilinear, linear)
    push!(triangular, angular)

    global tricount
    tricount += 1
  end
end

# set gravity
grav = [0, 0, 0]
function gravity(g1, g2, g3)
 grav[1] = g1
 grav[2] = g2
 grav[3] = g3
end

# run rigid simulation
function run_rigid(duration, step, outstep, outdir)

  output(0.0, outdir)
end

# run pseudo-rigid simulation
function run_pseudorigid(duration, step, outstep, outdir)
end

export ellipsoid, obstacle, gravity, run_rigid, run_pseudorigid

# --- utility functions ----

# calculate boundingx box
function ellbox()

  amax = maximum(ellradii[1])
  bmax = maximum(ellradii[2])
  cmax = maximum(ellradii[3])

  xlo = minimum(ellcenter[1]) - amax
  ylo = minimum(ellcenter[2]) - bmax
  zlo = minimum(ellcenter[3]) - cmax
  xhi = maximum(ellcenter[1]) + amax
  yhi = maximum(ellcenter[2]) + bmax
  zhi = maximum(ellcenter[3]) + cmax

  return (xlo, ylo, zlo, xhi, yhi, zhi)
end

# output current time step
function output(curtime, outdir)

  mkpath(outdir)

  path = joinpath(outdir,"ellip.dump")

  open(path, "w") do f
    write(f, "ITEM: TIMESTEP\n")
    write(f, curtime)
    write(f, "ITEM: NUMBER OF ATOMS\n")
    write(f, ellcount)
    (x0, y0, z0, x1, y1, z1) = ellbox()
    write(f, "ITEM: BOX BOUNDS\n");
    write(f, "$x0 $x1\n")
    write(f, "$y0 $y1\n")
    write(f, "$z0 $z1\n")
    write(f, "ITEM: ATOMS id x y z radius vx vy vz omegax omegay omegaz\n")
    for i = 1:ellcount
      println(f, i, " ", ellcenter[1][i], " ", ellcenter[2][i], " ", ellcenter[3][i], " ",
                (ellradii[1][i]+ellradii[2][i]+ellradii[3][i])/3.0, " ",
		 elllinear[1][i], " ", elllinear[2][i], " ", elllinear[3][i], " ",
		 ellangular[1][i], " ", ellangular[2][i], " ", ellangular[3][i], " ")
    end
  end
end

# ellipsoid properties
include("ellpro.jl")

end
# <-- edyna Julia module ends


"""
    get_ref_meshes(mtg)

Get all reference meshes from an mtg, usually from an OPF.

# Examples

```julia
using PlantGeom
file = joinpath(dirname(dirname(pathof(PlantGeom))),"test","files","simple_plant.opf")
opf = read_opf(file)
meshes = get_ref_meshes(opf)

using GLMakie
viz(meshes)
```
"""
function get_ref_meshes(mtg)

    if !isroot(mtg)
        @warn "Node is not the root node, using `get_root(mtg)`."
        x = get_root(mtg)
    else
        x = mtg
    end

    return x.attributes[:ref_meshes]
end

function get_ref_mesh_index!(node, ref_meshes = get_ref_meshes(node))

    # If the reference node mesh is unknown, get it:
    if node[:geometry].ref_mesh_index === nothing
        node[:geometry].ref_mesh_index =
            findfirst(x -> x === node[:geometry].ref_mesh, ref_meshes.meshes)
    end

    return node[:geometry].ref_mesh_index
end

function get_ref_mesh_index(node, ref_meshes = get_ref_meshes(node))
    # If the reference node mesh is unknown, get it:
    if node[:geometry].ref_mesh_index === nothing
        return findfirst(x -> x === node[:geometry].ref_mesh, ref_meshes.meshes)
    end

    return node[:geometry].ref_mesh_index
end

"""
    get_ref_mesh_index!(node, ref_meshes = get_ref_meshes(node))
    get_ref_mesh_index(node, ref_meshes = get_ref_meshes(node))

Get the index of the reference mesh used in the current node.

# Notes

Please use the `ref_meshes` argument preferably as not giving it make the function visit the
root node each time otherwise, and it can become a limitation when traversing a big MTG.
"""
get_ref_mesh_index!, get_ref_mesh_index

"""
    parse_ref_meshes(mtg)

Parse the reference meshes of an OPF into RefMeshes.
"""
function parse_ref_meshes(x)
    meshes = Dict{Int,RefMesh}()
    meshBDD = meshBDD_to_meshes(x[:meshBDD])

    for (id, value) in x[:shapeBDD]
        normals = meshBDD[value["meshIndex"]]["normals"]
        normals = SVector{length(normals) ÷ 3}(Meshes.Point3(normals[[i, i + 1, i + 2]]) for i in 1:3:length(normals))

        if haskey(meshBDD[value["meshIndex"]], "textureCoords")
            text_coord = meshBDD[value["meshIndex"]]["textureCoords"]
            text_coord = SVector{length(text_coord) ÷ 2}(Meshes.Point2(text_coord[[i, i + 1]]) for i in 1:2:length(text_coord))
        else
            text_coord = nothing
        end

        push!(
            meshes,
            id => RefMesh(
                value["name"],
                meshBDD[value["meshIndex"]]["mesh"],
                normals,
                text_coord,
                x[:materialBDD][value["materialIndex"]],
                meshBDD[value["meshIndex"]]["enableScale"]
            )
        )
    end

    # We create RefMeshes just now in case they were not sorted in the opf file
    refmeshes = RefMeshes(RefMesh[])

    for i in sort(collect(keys(meshes)))
        push!(refmeshes.meshes, meshes[i])
    end

    return refmeshes
end

"""
    meshBDD_to_meshes(x)

# Examples

```julia
using MultiScaleTreeGraph
file = joinpath(dirname(dirname(pathof(MultiScaleTreeGraph))),"test","files","simple_plant.opf")
opf = read_opf(file)
meshBDD_to_meshes(opf.attributes[:meshBDD])
```
"""
function meshBDD_to_meshes(x)
    meshes = Dict{Int,Any}()

    for (key, value) in x
        mesh = Dict()
        mesh_points = pop!(value, "points")
        mesh_faces = pop!(value, "faces")

        points3d = Meshes.Point3[mesh_points[p:(p+2)] for p = 1:3:length(mesh_points)]
        faces3d = [Meshes.connect((mesh_faces[p:(p+2)]...,), Meshes.Ngon) for p = 1:3:length(mesh_faces)]

        push!(mesh, "mesh" => Meshes.SimpleMesh(points3d, faces3d))
        merge!(mesh, value)

        push!(meshes, key => mesh)
    end

    return meshes
end

"""
Parse a material in opf format to a [`Phong`](@ref) material.
"""
function materialBDD_to_material(x)
    Phong(
        RGBA(x["emission"]...),
        RGBA(x["ambient"]...),
        RGBA(x["diffuse"]...),
        RGBA(x["specular"]...),
        x["shininess"]
    )
end

"""
    align_ref_meshes(meshes::RefMeshes)

Align all reference meshes along the X axis. Used for visualisation only.
"""
function align_ref_meshes(meshes::RefMeshes)
    meshes_vec = Meshes.SimpleMesh[]
    translation_vec = [0.0, 0.0, 0.0]

    for i in meshes.meshes
        translated_vertices = [i + Meshes.Vec(translation_vec...) for i in Meshes.vertices(i.mesh)]

        push!(meshes_vec, Meshes.SimpleMesh(translated_vertices, Meshes.topology(i.mesh)))

        # Maximum X coordinates of the newly translated mesh:
        xmax_ = maximum([Meshes.coordinates(i)[1] for i in translated_vertices])
        # Update the translation for the next mesh to begin at xmax*1.1 from the last one
        translation_vec[1] = xmax_ * 1.1
    end

    return meshes_vec
end

"""
    get_ref_meshes_color(meshes::RefMeshes)

Get the reference meshes colors (only the diffuse part for now).

# Examples

```julia
using MultiScaleTreeGraph, PlantGeom
file = joinpath(dirname(dirname(pathof(MultiScaleTreeGraph))),"test","files","simple_plant.opf")
opf = read_opf(file)
meshes = get_ref_meshes(opf)
PlantGeom.get_ref_meshes_color(meshes)
```
"""
function get_ref_meshes_color(meshes::RefMeshes)
    [material_single_color(i.material) for i in meshes.meshes]
end

function material_single_color(x::Phong)
    x.diffuse
end

function material_single_color(x::Colorant)
    x
end


importall Base
using Knet,AutoGrad
using Knet: gpuinfo, gpusync
gpu(true)

AutoGrad._dbg{T}(x::KnetArray{T})="K$(join([id2(x),size(x)...,8*sizeof(T)],'_'))"

include(Pkg.dir("AutoGrad/test/profile.jl"))
gpuinfo("before d0kn,w2kn")
d0kn = map(i->map(KnetArray,i),d0)
w2kn = map(KnetArray,w2)
gpuinfo("after d0kn,w2kn")
# w1kn = map(KnetArray,w1)
# timeall(w2kn,d0kn,10)
# loop(fun[9],w2kn,d0kn,10)
# (x1kn,y1kn)=first(d0kn)
# fun[10](w2kn,x1kn,y1kn)

function timeall_kn(w=w2kn,d=d0kn,t=10)
    for i=1:length(fun)
        print(i); printfun(fun[i])
        for j=1:4
            gc_enable(false)
            @time (loop_kn(fun[i],w,d,t); gpusync())
            # @time loop_kn(fun[i],w,d,t)
            gpuinfo("before gc")
            gc_enable(true)
            gc()
            sleep(2)
            gpuinfo("after gc ")
        end
    end
end

function loop_kn(f,w,d,t)
    for i in 1:t
        for (x,y) in d
            f(w,x,y)
        end
    end
end

# gc_enable(false); @time loop_kn(fun[9],w2kn,d0kn,10); gc_enable(true)
# gc_enable(false); @time loop_kn(fun[9],w2kn,d0kn,10); gc_enable(true)
# gc_enable(false); @time loop_kn(fun[9],w2kn,d0kn,10); gc_enable(true)
# gc_enable(false); @time loop_kn(fun[9],w2kn,d0kn,10); gc_enable(true)
# gc_enable(false); @time loop_kn(fun[9],w2kn,d0kn,10); gc_enable(true)
# gc_enable(false); @profile loop_kn(fun[9],w2kn,d0kn,10); gc_enable(true)
# Profile.print()
# Profile.print(format=:flat)

timeall_kn()


"""
    prepare_observables(m::HubbardModelAttractive)

Initializes observables for the attractive Hubbard model and returns a `Dict{String, Observable}`.

See also [`measure_observables!`](@ref) and [`finish_observables!`](@ref).
"""
@inline function prepare_observables(mc::DQMC, m::HubbardModelAttractive)
    obs = Dict{String,Observable}()
    obs["confs"] = Observable(HubbardConf, "Configurations")
    obs["greens"] = Observable(typeof(mc.s.greens), "Equal-times Green's function")
    obs["Eboson"] = Observable(Float64, "Bosonic energy")

    return obs
end

"""
    measure_observables!(m::HubbardModelAttractive, obs::Dict{String,Observable}, conf::HubbardConf, E::Float64)

Measures observables and updates corresponding `Observable` objects in `obs`.

See also [`prepare_observables`](@ref) and [`finish_observables!`](@ref).
"""
@noinline function measure_observables!(mc::DQMC, m::HubbardModelAttractive,
							obs::Dict{String,Observable}, conf::HubbardConf)
    push!(obs["confs"], mc.conf)
    push!(obs["greens"], greens(mc))
    push!(obs["Eboson"], energy_boson(mc, m, conf))
    nothing
end

"""
    measure_observables!(mc::DQMC, m::HubbardModelAttractive, obs::Dict{String,Observable})

Finish measurements of observables.

See also [`prepare_observables`](@ref) and [`measure_observables!`](@ref).
"""
@inline function finish_observables!(mc::DQMC, m::HubbardModelAttractive,
							obs::Dict{String,Observable})
    nothing
end


###
### Particle Filtering and Particle MCMC Samplers.
###

####
#### Generic Sequential Monte Carlo sampler.
####

"""
$(TYPEDEF)

Sequential Monte Carlo sampler.

# Fields

$(TYPEDFIELDS)
"""
struct SMC{space, R} <: ParticleInference
    resampler::R
end

"""
    SMC(space...)
    SMC([resampler = AdvancedPS.ResampleWithESSThreshold(), space = ()])
    SMC([resampler = AdvancedPS.resample_systematic, ]threshold[, space = ()])

Create a sequential Monte Carlo sampler of type [`SMC`](@ref) for the variables in `space`.

If the algorithm for the resampling step is not specified explicitly, systematic resampling
is performed if the estimated effective sample size per particle drops below 0.5.
"""
function SMC(resampler = AdvancedPS.ResampleWithESSThreshold(), space::Tuple = ())
    return SMC{space, typeof(resampler)}(resampler)
end

# Convenient constructors with ESS threshold
function SMC(resampler, threshold::Real, space::Tuple = ())
    return SMC(AdvancedPS.ResampleWithESSThreshold(resampler, threshold), space)
end
SMC(threshold::Real, space::Tuple = ()) = SMC(AdvancedPS.resample_systematic, threshold, space)

# If only the space is defined
SMC(space::Symbol...) = SMC(space)
SMC(space::Tuple) = SMC(AdvancedPS.ResampleWithESSThreshold(), space)

struct SMCTransition{T,F<:AbstractFloat}
    "The parameters for any given sample."
    θ::T
    "The joint log probability of the sample (NOTE: does not work, always set to zero)."
    lp::F
    "The weight of the particle the sample was retrieved from."
    weight::F
end

function SMCTransition(vi::AbstractVarInfo, weight)
    theta = tonamedtuple(vi)

    # This is pretty useless since we reset the log probability continuously in the
    # particle sweep.
    lp = getlogp(vi)

    return SMCTransition(theta, lp, weight)
end

metadata(t::SMCTransition) = (lp = t.lp, weight = t.weight)

DynamicPPL.getlogp(t::SMCTransition) = t.lp

struct SMCState{P,F<:AbstractFloat}
    particles::P
    particleindex::Int
    # The logevidence after aggregating all samples together.
    average_logevidence::F
end

function getlogevidence(samples, sampler::Sampler{<:SMC}, state::SMCState)
    return state.average_logevidence
end

function AbstractMCMC.sample(
    rng::AbstractRNG,
    model::AbstractModel,
    sampler::Sampler{<:SMC},
    N::Integer;
    chain_type=MCMCChains.Chains,
    resume_from=nothing,
    progress=PROGRESS[],
    kwargs...
)
    if resume_from === nothing
        return AbstractMCMC.mcmcsample(rng, model, sampler, N;
                                       chain_type=chain_type,
                                       progress=progress,
                                       nparticles=N,
                                       kwargs...)
    else
        return resume(resume_from, N;
                      chain_type=chain_type, progress=progress, nparticles=N, kwargs...)
    end
end

function DynamicPPL.initialstep(
    rng::AbstractRNG,
    model::AbstractModel,
    spl::Sampler{<:SMC},
    vi::AbstractVarInfo;
    nparticles::Int,
    kwargs...
)
    # Reset the VarInfo.
    reset_num_produce!(vi)
    set_retained_vns_del_by_spl!(vi, spl)
    resetlogp!!(vi)
    empty!!(vi)

    # Create a new set of particles.
    particles = AdvancedPS.ParticleContainer(
        [AdvancedPS.Trace(model, spl, vi) for _ in 1:nparticles],
    )

    # Perform particle sweep.
    logevidence = AdvancedPS.sweep!(rng, particles, spl.alg.resampler)

    # Extract the first particle and its weight.
    particle = particles.vals[1]
    weight = AdvancedPS.getweight(particles, 1)

    # Compute the first transition and the first state.
    transition = SMCTransition(particle.f.varinfo, weight)
    state = SMCState(particles, 2, logevidence)

    return transition, state
end

function AbstractMCMC.step(
    ::AbstractRNG,
    model::AbstractModel,
    spl::Sampler{<:SMC},
    state::SMCState;
    kwargs...
)
    # Extract the index of the current particle.
    index = state.particleindex

    # Extract the current particle and its weight.
    particles = state.particles
    particle = particles.vals[index]
    weight = AdvancedPS.getweight(particles, index)

    # Compute the transition and the next state.
    transition = SMCTransition(particle.f.varinfo, weight)
    nextstate = SMCState(state.particles, index + 1, state.average_logevidence)

    return transition, nextstate
end

####
#### Particle Gibbs sampler.
####

"""
$(TYPEDEF)

Particle Gibbs sampler.

Note that this method is particle-based, and arrays of variables
must be stored in a [`TArray`](@ref) object.

# Fields

$(TYPEDFIELDS)
"""
struct PG{space,R} <: ParticleInference
    """Number of particles."""
    nparticles::Int
    """Resampling algorithm."""
    resampler::R
end

"""
    PG(n, space...)
    PG(n, [resampler = AdvancedPS.ResampleWithESSThreshold(), space = ()])
    PG(n, [resampler = AdvancedPS.resample_systematic, ]threshold[, space = ()])

Create a Particle Gibbs sampler of type [`PG`](@ref) with `n` particles for the variables
in `space`.

If the algorithm for the resampling step is not specified explicitly, systematic resampling
is performed if the estimated effective sample size per particle drops below 0.5.
"""
function PG(
    nparticles::Int,
    resampler = AdvancedPS.ResampleWithESSThreshold(),
    space::Tuple = (),
)
    return PG{space, typeof(resampler)}(nparticles, resampler)
end

# Convenient constructors with ESS threshold
function PG(nparticles::Int, resampler, threshold::Real, space::Tuple = ())
    return PG(nparticles, AdvancedPS.ResampleWithESSThreshold(resampler, threshold), space)
end
function PG(nparticles::Int, threshold::Real, space::Tuple = ())
    return PG(nparticles, AdvancedPS.resample_systematic, threshold, space)
end

# If only the number of particles and the space is defined
PG(nparticles::Int, space::Symbol...) = PG(nparticles, space)
function PG(nparticles::Int, space::Tuple)
    return PG(nparticles, AdvancedPS.ResampleWithESSThreshold(), space)
end

const CSMC = PG # type alias of PG as Conditional SMC

struct PGTransition{T,F<:AbstractFloat}
    "The parameters for any given sample."
    θ::T
    "The joint log probability of the sample (NOTE: does not work, always set to zero)."
    lp::F
    "The log evidence of the sample."
    logevidence::F
end

function PGTransition(vi::AbstractVarInfo, logevidence)
    theta = tonamedtuple(vi)

    # This is pretty useless since we reset the log probability continuously in the
    # particle sweep.
    lp = getlogp(vi)

    return PGTransition(theta, lp, logevidence)
end

metadata(t::PGTransition) = (lp = t.lp, logevidence = t.logevidence)

DynamicPPL.getlogp(t::PGTransition) = t.lp

function getlogevidence(samples, sampler::Sampler{<:PG}, vi::AbstractVarInfo)
    return mean(x.logevidence for x in samples)
end

function DynamicPPL.initialstep(
    rng::AbstractRNG,
    model::AbstractModel,
    spl::Sampler{<:PG},
    vi::AbstractVarInfo;
    kwargs...
)
    # Reset the VarInfo before new sweep
    reset_num_produce!(vi)
    set_retained_vns_del_by_spl!(vi, spl)
    resetlogp!!(vi)

    # Create a new set of particles
    num_particles = spl.alg.nparticles
    particles = AdvancedPS.ParticleContainer(
        [AdvancedPS.Trace(model, spl, vi) for _ in 1:num_particles],
    )

    # Perform a particle sweep.
    logevidence = AdvancedPS.sweep!(rng, particles, spl.alg.resampler)

    # Pick a particle to be retained.
    Ws = AdvancedPS.getweights(particles)
    indx = AdvancedPS.randcat(rng, Ws)
    reference = particles.vals[indx]

    # Compute the first transition.
    _vi = reference.f.varinfo
    transition = PGTransition(_vi, logevidence)

    return transition, _vi
end

function AbstractMCMC.step(
    rng::AbstractRNG,
    model::AbstractModel,
    spl::Sampler{<:PG},
    vi::AbstractVarInfo;
    kwargs...
)
    # Reset the VarInfo before new sweep.
    reset_num_produce!(vi)
    resetlogp!!(vi)

    # Create reference particle for which the samples will be retained.
    reference = AdvancedPS.forkr(AdvancedPS.Trace(model, spl, vi))

    # For all other particles, do not retain the variables but resample them.
    set_retained_vns_del_by_spl!(vi, spl)

    # Create a new set of particles.
    num_particles = spl.alg.nparticles
    x = map(1:num_particles) do i
        if i != num_particles
            return AdvancedPS.Trace(model, spl, vi)
        else
            return reference
        end
    end
    particles = AdvancedPS.ParticleContainer(x)

    # Perform a particle sweep.
    logevidence = AdvancedPS.sweep!(rng, particles, spl.alg.resampler, reference)

    # Pick a particle to be retained.
    Ws = AdvancedPS.getweights(particles)
    indx = AdvancedPS.randcat(rng, Ws)
    newreference = particles.vals[indx]

    # Compute the transition.
    _vi = newreference.f.varinfo
    transition = PGTransition(_vi, logevidence)

    return transition, _vi
end

DynamicPPL.use_threadsafe_eval(::SamplingContext{<:Sampler{<:Union{PG,SMC}}}, ::AbstractVarInfo) = false

function DynamicPPL.assume(
    rng,
    spl::Sampler{<:Union{PG,SMC}},
    dist::Distribution,
    vn::VarName,
    __vi__::AbstractVarInfo
)
    local vi
    try
        vi = AdvancedPS.current_trace().f.varinfo
    catch e
        # NOTE: this heuristic allows Libtask evaluating a model outside a `Trace`.
        if e == KeyError(:__trace) || current_task().storage isa Nothing
            vi = __vi__
        else
            rethrow(e)
        end
    end
    if inspace(vn, spl)
        if ~haskey(vi, vn)
            r = rand(rng, dist)
            push!!(vi, vn, r, dist, spl)
        elseif is_flagged(vi, vn, "del")
            unset_flag!(vi, vn, "del")
            r = rand(rng, dist)
            vi[vn] = vectorize(dist, r)
            DynamicPPL.setgid!(vi, spl.selector, vn)
            setorder!(vi, vn, get_num_produce(vi))
        else
            DynamicPPL.updategid!(vi, vn, spl)
            r = vi[vn]
        end
    else # vn belongs to other sampler <=> conditionning on vn
        if haskey(vi, vn)
            r = vi[vn]
        else
            r = rand(rng, dist)
            push!!(vi, vn, r, dist, DynamicPPL.Selector(:invalid))
        end
        lp = logpdf_with_trans(dist, r, istrans(vi, vn))
        acclogp!!(vi, lp)
    end
    return r, 0, vi
end

function DynamicPPL.observe(spl::Sampler{<:Union{PG,SMC}}, dist::Distribution, value, vi)
    Libtask.produce(logpdf(dist, value))
    return 0, vi
end

# Convenient constructor
function AdvancedPS.Trace(
    model::Model,
    sampler::Sampler{<:Union{SMC,PG}},
    varinfo::AbstractVarInfo,
)
    newvarinfo = deepcopy(varinfo)
    DynamicPPL.reset_num_produce!(newvarinfo)
    f = Turing.Essential.TracedModel(model, sampler, newvarinfo)
    return AdvancedPS.Trace(f)
end


function append_depthfirst!(g, lookup, node)
    push!(lookup, node)
    pv = nv(g)
    for c in node["children"]
        add_vertex!(g)
        add_edge!(g, pv, nv(g))
        append_depthfirst!(g, lookup, c)
    end
    g, lookup
end

"""
    g, vertexlist = ontology(filename)
    g, vertexlist = ontology("mouse")

Return the structure graph `g` and list `vertexlist` of metadata for nodes.
"""
function ontology(filename)
    str = readstring(dataset(filename, "ontology"))
    hierarchy = JSON.parse(str)
    g = DiGraph()
    add_vertex!(g)  # for the root node
    vertexlist = Dict{String,Any}[]
    append_depthfirst!(g, vertexlist, hierarchy)
    g, vertexlist
end

"""
    matches = findvertices(pat, vertexlist, field="acronym")

Return the list of vertices that Regex-match to `pat` for the
specified field ( `"acronym"` or `"name"`) of each node listed in `vertexlist`.
"""
findvertices(pat, vertexlist, field="acronym") = find(x->ismatch(pat, x[field]), vertexlist)

"""
    ids = structureids(v, vertexlist)

Return the structure_id associated with each vertex in `v`.
"""
structureids(v, vertexlist) = map(i->vertexlist[i]["id"], v)
structureids(pat::Regex, vertexlist, field="acronym") =
    structureids(findvertices(pat, vertexlist, field), vertexlist)

function annotation(resolution=50; species="mouse")
    anno = load(dataset(species, "annotation_"*string(resolution)))
    # The NRRD header is incorrect, we should fix the orientation
    r = resolution*μm
    return AxisArray(anno.data, (:P, :I, :R), (r, r, r))
end

function download_annotation(resolution=50)
    dest = joinpath(mousedir, "annotation_$resolution.nrrd")
    download("http://download.alleninstitute.org/informatics-archive/current-release/mouse_ccf/annotation/ccf_2016/annotation_$resolution.nrrd", dest)
end

"""
    bbleft, bbright = boundingbox(f, annotated)

Return the BoundingBoxes encapsulating structure(s) for which
`f(annotated[i])` returns `true`, separately for the two hemispheres
of the brain. `annotated` is a labeled volume. Units of `bb` are
microns.
"""
function boundingbox(f, annotated)
    scalebb(bx, px) = BoundingBox(map((iv, s) -> minimum(iv)*s .. maximum(iv)*s, bx.intervals, ps))
    local bb
    isfirst = true
    # Search only the left half of the volume
    inds = mapfilter(i->first(i):(first(i)+last(i))÷2, indices(annotated),
                     Axis{:R}, axes(annotated))
    for i in CartesianRange(inds)
        if f(annotated[i])
            if isfirst
                bb = BoundingBox(i, i)
                isfirst = false
            else
                bb |= i
            end
        end
    end
    dim = find(x->x==:R, axisnames(annotated))[1]
    bbr = BoundingBox(ntuple(ndims(annotated)) do d
        iv = bb.intervals[d]
        if d == dim
            id = indices(annotated, d)
            l = first(id) + last(id)
            l-maximum(iv) .. l-minimum(iv)
        else
            iv
        end
    end)
    ps = pixelspacing(annotated)
    scalebb(bb, ps), scalebb(bbr, ps)
end


module RMB

importall NLsolve
import Base.rand
import Distributions: wsample
import Calculus

export solvedp!, u!, pi!
export InventoryModel

include("stochastic-matrices.jl")
include("rust-models.jl")
include("inv_model.jl")

end #module


function main()
    d = setup_data()

end

function train(d::Dataset)
    opt = SGD()
    for epoch = 1:nepochs
        println("epoch: $(epoch)")
        opt.rate = 0.0075 / epoch
        loss = fit(train_x, train_y, model, crossentropy, opt)
        println("loss: $(loss)")

        test_z = map(x -> predict(model,x), test_x)
        acc = accuracy(test_y, test_z)
        println("test acc.: $(acc)")
        println("")
    end
end

include("intdict.jl")
include("io.jl")
include("model.jl")


##########################
# plan_* constructors
##########################

for op in [:nfft, :nfct, :nfst]
planfunc = Symbol("plan_"*"$op")
@eval begin

# The following automatically call the plan_* version for type Array

$(planfunc)(x::AbstractArray, N::Union{Integer,NTuple{D,Int}}, args...; kargs...) where {D} =
    $(planfunc)(Array, x, N, args...; kargs...)

$(planfunc)(x::AbstractArray, y::AbstractArray, args...; kargs...) where {D} =
    $(planfunc)(Array, x, y, args...; kargs...)

# The follow convert 1D parameters into the format required by the plan

$(planfunc)(Q::Type, x::AbstractVector, N::Integer, rest...; kwargs...) where {D}  =
    $(planfunc)(Q, collect(reshape(x,1,length(x))), (N,), rest...; kwargs...)

$(planfunc)(Q::Type, x::AbstractVector, N::NTuple{D,Int}, rest...; kwargs...) where {D} =
    $(planfunc)(Q, collect(reshape(x,1,length(x))), N, rest...; kwargs...)

$(planfunc)(Q::Type, x::AbstractMatrix, N::NTuple{D,Int}, rest...; kwargs...) where {D}  =
    $(planfunc)(Q, collect(x), N, rest...; kwargs...)

end
end

## NNFFT constructor
plan_nnfft(Q::Type, x::AbstractVector, y::AbstractVector, rest...; kwargs...) where {D}  =
    plan_nnfft(Q, collect(reshape(x,1,length(x))), collect(reshape(y,1,length(x))), rest...; kwargs...)

###############################################
# Allocating trafo functions with plan creation
###############################################

for (op,trans) in zip([:nfft, :nfct, :nfst],
                      [:adjoint, :transpose, :transpose])
planfunc = Symbol("plan_$(op)")
tfunc = Symbol("$(op)_$(trans)")
@eval begin

# TODO fix comments (how?)
"""
nfft(x, f, rest...; kwargs...)

calculates the nfft of the array `f` for the nodes contained in the matrix `x`
The output is a vector of length M=`size(nodes,2)`
"""
function $(op)(x, f::AbstractArray; kargs...)
  p = $(planfunc)(x, size(f); kargs... )
  return p * f
end

"""
nfft_adjoint(x, N, fHat, rest...; kwargs...)

calculates the adjoint nfft of the vector `fHat` for the nodes contained in the matrix `x`.
The output is an array of size `N`
"""
function $(tfunc)(x, N, fHat;  kargs...)
  p = $(planfunc)(x, N;  kargs...)
  return $(trans)(p) * fHat
end

end
end

############################
# Allocating trafo functions
############################

"""
        *(p, f) -> fHat

For a **non**-directional `D` dimensional plan `p` this calculates the NFFT/NNFFT of a `D` dimensional array `f` of size `N`.
`fHat` is a vector of length `M`.
(`M` and `N` are defined in the plan creation)

For a **directional** `D` dimensional plan `p` both `f` and `fHat` are `D`
dimensional arrays, and the dimension specified in the plan creation is
affected.
"""
function Base.:*(p::AbstractComplexFTPlan{T}, f::AbstractArray{Complex{U},D}; kargs...) where {T,U,D}
  fHat = similar(f, Complex{T}, size_out(p))
  mul!(fHat, p, f; kargs...)
  return fHat
end

"""
        *(p::Adjoint{T,<:AbstractFTPlan{T}}, fHat) -> f

For a **non**-directional `D` dimensional plan `p` this calculates the adjoint NFFT/NNFFT of a length `M` vector `fHat`
`f` is a `D` dimensional array of size `N`.
(`M` and `N` are defined in the plan creation)

For a **directional** `D` dimensional plan `p` both `f` and `fHat` are `D`
dimensional arrays, and the dimension specified in the plan creation is
affected.
"""

function Base.:*(p::Adjoint{Complex{T},<:AbstractComplexFTPlan{T}}, fHat::AbstractArray{Complex{U},D}; kargs...) where {T,U,D}
  f = similar(fHat, Complex{T}, size_out(p))
  mul!(f, p, fHat; kargs...)
  return f
end

# The following two methods are redundant but need to be defined because of a method ambiguity with Julia Base
function Base.:*(p::Adjoint{Complex{T},<:AbstractComplexFTPlan{T}}, fHat::AbstractVector{Complex{U}}; kargs...) where {T,U}
  f = similar(fHat, Complex{T}, size_out(p))
  mul!(f, p, fHat; kargs...)
  return f
end
function Base.:*(p::Adjoint{Complex{T},<:AbstractComplexFTPlan{T}}, fHat::AbstractArray{Complex{U},2}; kargs...) where {T,U}
  f = similar(fHat, Complex{T}, size_out(p))
  mul!(f, p, fHat; kargs...)
  return f
end

"""
        *(p, f) -> fHat

For a **non**-directional `D` dimensional plan `p` this calculates the NFCT/NFST of a `D` dimensional array `f` of size `N`.
`fHat` is a vector of length `M`.
(`M` and `N` are defined in the plan creation)

For a **directional** `D` dimensional plan `p` both `f` and `fHat` are `D`
dimensional arrays, and the dimension specified in the plan creation is
affected.
"""
function Base.:*(p::AbstractRealFTPlan{T}, f::AbstractArray{U,D}; kargs...) where {T,U,D}
  fHat = similar(f, T, size_out(p))
  mul!(fHat, p, f; kargs...)
  return fHat
end

"""
        *(p::Transpose{T,AbstractRealFTPlan{T}}, fHat) -> f

For a **non**-directional `D` dimensional plan `p` this calculates the adjoint NFCT/NFST of a length `M` vector `fHat`
`f` is a `D` dimensional array of size `N`.
(`M` and `N` are defined in the plan creation)

For a **directional** `D` dimensional plan `p` both `f` and `fHat` are `D`
dimensional arrays, and the dimension specified in the plan creation is
affected.
"""

function Base.:*(p::Transpose{T,<:AbstractRealFTPlan{T}}, fHat::AbstractArray{U,D}; kargs...) where {T,U,D}
  f = similar(fHat, T, size_out(p))
  mul!(f, p, fHat; kargs...)
  return f
end

# The following two methods are redundant but need to be defined because of a method ambiguity with Julia Base
function Base.:*(p::Transpose{T,<:AbstractRealFTPlan{T}}, fHat::AbstractVector{U}; kargs...) where {T,U}
  f = similar(fHat, T, size_out(p))
  mul!(f, p, fHat; kargs...)
  return f
end
function Base.:*(p::Transpose{T,<:AbstractRealFTPlan{T}}, fHat::AbstractArray{U,2}; kargs...) where {T,U}
  f = similar(fHat, T, size_out(p))
  mul!(f, p, fHat; kargs...)
  return f
end


#################################################################################
#Type Alias for long type signatures
const MinMax = NamedTuple{(:min, :max), NTuple{2, Float64}}
const NamedMinMax = Tuple{String, MinMax}
const UpDown = NamedTuple{(:up, :down), NTuple{2, Float64}}
const InOut = NamedTuple{(:in, :out), NTuple{2, Float64}}

# Type Alias From other Packages
const PM = PowerModels
const PSY = PowerSystems
const PSI = PowerSimulations
const IS = InfrastructureSystems
const MOI = MathOptInterface
const MOIU = MathOptInterface.Utilities
const PJ = ParameterJuMP
const MOPFM = MathOptFormat.MOF.Model
const TS = TimeSeries

#Type Alias for JuMP and PJ containers
const JuMPExpressionMatrix = Matrix{<:JuMP.AbstractJuMPScalar}
const PGAE{V} = PJ.ParametrizedGenericAffExpr{Float64, V} where V<:JuMP.AbstractVariableRef
const GAE{V} = JuMP.GenericAffExpr{Float64, V} where V<:JuMP.AbstractVariableRef
const JuMPAffineExpressionArray = Matrix{GAE{V}} where V<:JuMP.AbstractVariableRef
const JuMPConstraintArray = JuMP.Containers.DenseAxisArray{JuMP.ConstraintRef}
const JuMPParamArray = JuMP.Containers.DenseAxisArray{PJ.ParameterRef}
const DSDA = Dict{Symbol, JuMP.Containers.DenseAxisArray}


export DeepCFR

"""
    DeepCFR(;kwargs...)

Symbols used here follow the paper: [Deep Counterfactual Regret Minimization](https://arxiv.org/abs/1811.00164)

# Keyword arguments

- `K`, number of traverrsal.
- `t`, number of iteration.
- `Π`, the policy network.
- `V`, a dictionary of each player's advantage network.
- `MΠ`, a strategy memory.
- `MV`, a dictionary of each player's advantage memory.
- `reinitialize_freq=1`, the frequency of reinitializing the value networks.
"""
Base.@kwdef mutable struct DeepCFR{TP,TV,TMP,TMV,I,R,P} <: AbstractCFRPolicy
    Π::TP
    V::TV
    MΠ::TMP
    MV::TMV
    K::Int = 20
    t::Int = 1
    reinitialize_freq::Int = 1
    batch_size_V::Int = 32
    batch_size_Π::Int = 32
    n_training_steps_V::Int = 1
    n_training_steps_Π::Int = 1
    rng::R = Random.GLOBAL_RNG
    initializer::I = glorot_normal(rng)
    max_grad_norm::Float32 = 10.0f0
    # for logging
    Π_losses::Vector{Float32} = zeros(Float32, n_training_steps_Π)
    V_losses::Dict{P,Vector{Float32}} =
        Dict(k => zeros(Float32, n_training_steps_V) for (k, _) in MV)
    Π_norms::Vector{Float32} = zeros(Float32, n_training_steps_Π)
    V_norms::Dict{P,Vector{Float32}} =
        Dict(k => zeros(Float32, n_training_steps_V) for (k, _) in MV)
end

function RLBase.prob(π::DeepCFR, env::AbstractEnv)
    I = send_to_device(device(π.Π), state(env))
    m = send_to_device(device(π.Π), ifelse.(legal_action_space_mask(env), 0.0f0, -Inf32))
    logits = π.Π(Flux.unsqueeze(I, ndims(I) + 1)) |> vec
    σ = softmax(logits .+ m)
    send_to_host(σ)
end

(π::DeepCFR)(env::AbstractEnv) =
    sample(π.rng, action_space(env), Weights(prob(π, env), 1.0))

"Run one interation"
function RLBase.update!(π::DeepCFR, env::AbstractEnv)
    for p in players(env)
        if p != chance_player(env)
            for k in 1:π.K
                external_sampling!(π, copy(env), p)
            end
            update_advantage_networks(π, p)
        end
    end
    π.t += 1
end

"Update Π (policy network)"
function RLBase.update!(π::DeepCFR)
    Π = π.Π
    Π_losses = π.Π_losses
    Π_norms = π.Π_norms
    D = device(Π)
    MΠ = π.MΠ
    ps = Flux.params(Π)

    for x in ps
        x .= π.initializer(size(x)...)
    end

    for i in 1:π.n_training_steps_Π
        batch_inds = rand(π.rng, 1:length(MΠ), π.batch_size_Π)
        I = send_to_device(D, Flux.batch([MΠ[:I][i] for i in batch_inds]))
        σ = send_to_device(D, Flux.batch([MΠ[:σ][i] for i in batch_inds]))
        t = send_to_device(D, Flux.batch([MΠ[:t][i] / π.t for i in batch_inds]))
        m = send_to_device(
            D,
            Flux.batch([ifelse.(MΠ[:m][i], 0.0f0, -Inf32) for i in batch_inds]),
        )
        gs = gradient(ps) do
            logits = Π(I) .+ m
            loss = mean(reshape(t, 1, :) .* ((σ .- softmax(logits)) .^ 2))
            ignore() do
                # println(σ, "!!!",m, "===", Π(I))
                Π_losses[i] = loss
            end
            loss
        end
        Π_norms[i] = clip_by_global_norm!(gs, ps, π.max_grad_norm)
        update!(Π, gs)
    end
end

"Update advantage network"
function update_advantage_networks(π, p)
    V = π.V[p]
    V_losses = π.V_losses[p]
    V_norms = π.V_norms[p]
    MV = π.MV[p]
    if π.t % π.reinitialize_freq == 0
        for x in Flux.params(V)
            # TODO: inplace
            x .= π.initializer(size(x)...)
        end
    end
    if length(MV) >= π.batch_size_V
        for i in 1:π.n_training_steps_V
            batch_inds = rand(π.rng, 1:length(MV), π.batch_size_V)
            I = send_to_device(device(V), Flux.batch([MV[:I][i] for i in batch_inds]))
            r̃ = send_to_device(device(V), Flux.batch([MV[:r̃][i] for i in batch_inds]))
            t = send_to_device(device(V), Flux.batch([MV[:t][i] / π.t for i in batch_inds]))
            m = send_to_device(device(V), Flux.batch([MV[:m][i] for i in batch_inds]))
            ps = Flux.params(V)
            gs = gradient(ps) do
                loss = mean(reshape(t, 1, :) .* ((r̃ .- V(I) .* m) .^ 2))
                ignore() do
                    V_losses[i] = loss
                end
                loss
            end
            V_norms[i] = clip_by_global_norm!(gs, ps, π.max_grad_norm)
            update!(V, gs)
        end
    end
end

"CFR Traversal with External Sampling"
function external_sampling!(π::DeepCFR, env::AbstractEnv, p)
    if is_terminated(env)
        reward(env, p)
    elseif current_player(env) == chance_player(env)
        env(rand(π.rng, action_space(env)))
        external_sampling!(π, env, p)
    elseif current_player(env) == p
        V = π.V[p]
        s = state(env)
        I = send_to_device(device(V), Flux.unsqueeze(s, ndims(s) + 1))
        A = action_space(env)
        m = legal_action_space_mask(env)
        σ = masked_regret_matching(V(I) |> send_to_host |> vec, m)
        v = zeros(length(σ))
        v̄ = 0.0
        for i in 1:length(m)
            if m[i]
                v[i] = external_sampling!(π, child(env, A[i]), p)
                v̄ += σ[i] * v[i]
            end
        end
        push!(π.MV[p], I = s, t = π.t, r̃ = (v .- v̄) .* m, m = m)
        v̄
    else
        V = π.V[current_player(env)]
        s = state(env)
        I = send_to_device(device(V), Flux.unsqueeze(s, ndims(s) + 1))
        A = action_space(env)
        m = legal_action_space_mask(env)
        σ = masked_regret_matching(V(I) |> send_to_host |> vec, m)
        push!(π.MΠ, I = s, t = π.t, σ = σ, m = m)
        a = sample(π.rng, A, Weights(σ, 1.0))
        env(a)
        external_sampling!(π, env, p)
    end
end

"This is the specific regret matching method used in DeepCFR"
function masked_regret_matching(v, m)
    v⁺ = max.(v .* m, 0.0f0)
    s = sum(v⁺)
    if s > 0
        v⁺ ./= s
    else
        fill!(v⁺, 0.0f0)
        v⁺[findmax(v, m)[2]] = 1.0
    end
    v⁺
end


mutable struct DirectJumpAggregation{T,S,F1,F2,RNG} <: AbstractSSAJumpAggregator
  next_jump::Int
  prev_jump::Int
  next_jump_time::T
  end_time::T
  cur_rates::Vector{T}
  sum_rate::T
  ma_jumps::S
  rates::F1
  affects!::F2
  save_positions::Tuple{Bool,Bool}
  rng::RNG
end
DirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool,Bool}, rng::RNG; kwargs...) where {T,S,F1,F2,RNG} =
  DirectJumpAggregation{T,S,F1,F2,RNG}(nj, nj, njt, et, crs, sr, maj, rs, affs!, sps, rng)

############################# Required Functions #############################

# creating the JumpAggregation structure (tuple-based constant jumps)
function aggregate(aggregator::Direct, u, p, t, end_time, constant_jumps,
                    ma_jumps, save_positions, rng; kwargs...)

  # handle constant jumps using tuples
  rates, affects! = get_jump_info_tuples(constant_jumps)

  build_jump_aggregation(DirectJumpAggregation, u, p, t, end_time, ma_jumps,
                          rates, affects!, save_positions, rng; kwargs...)
end

# creating the JumpAggregation structure (function wrapper-based constant jumps)
function aggregate(aggregator::DirectFW, u, p, t, end_time, constant_jumps,
                    ma_jumps, save_positions, rng; kwargs...)

  # handle constant jumps using function wrappers
  rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps)

  build_jump_aggregation(DirectJumpAggregation, u, p, t, end_time, ma_jumps,
                          rates, affects!, save_positions, rng; kwargs...)
end

# set up a new simulation and calculate the first jump / jump time
function initialize!(p::DirectJumpAggregation, integrator, u, params, t)
  generate_jumps!(p, integrator, u, params, t)
  nothing
end

# execute one jump, changing the system state
@inline function execute_jumps!(p::DirectJumpAggregation, integrator, u, params, t)
  update_state!(p, integrator, u)
  nothing
end

# calculate the next jump / jump time
function generate_jumps!(p::DirectJumpAggregation, integrator, u, params, t)
  p.sum_rate, ttnj = time_to_next_jump(p, u, params, t)
  @fastmath p.next_jump_time = t + ttnj
  @inbounds p.next_jump = searchsortedfirst(p.cur_rates, rand(p.rng) * p.sum_rate)
  nothing
end

######################## SSA specific helper routines ########################

# tuple-based constant jumps
@fastmath function time_to_next_jump(p::DirectJumpAggregation{T,S,F1,F2,RNG}, u, params, t) where {T,S,F1 <: Tuple, F2 <: Tuple, RNG}
  prev_rate = zero(t)
  new_rate  = zero(t)
  cur_rates = p.cur_rates

  # mass action rates
  majumps   = p.ma_jumps
  idx       = get_num_majumps(majumps)
  @inbounds for i in 1:idx
    new_rate     = evalrxrate(u, i, majumps)
    cur_rates[i] = new_rate + prev_rate
    prev_rate    = cur_rates[i]
  end

  # constant jump rates
  rates = p.rates
  if !isempty(rates)
    idx  += 1
    fill_cur_rates(u, params, t, cur_rates, idx, rates...)
    @inbounds for i in idx:length(cur_rates)
      cur_rates[i] = cur_rates[i] + prev_rate
      prev_rate    = cur_rates[i]
    end
  end

  @inbounds sum_rate = cur_rates[end]
  sum_rate, randexp(p.rng) / sum_rate
end

@inline function fill_cur_rates(u, p, t, cur_rates, idx, rate, rates...)
  @inbounds cur_rates[idx] = rate(u, p, t)
  idx += 1
  fill_cur_rates(u, p, t, cur_rates, idx, rates...)
end

@inline function fill_cur_rates(u, p, t, cur_rates, idx, rate)
  @inbounds cur_rates[idx] = rate(u, p, t)
  nothing
end

# function wrapper-based constant jumps
@fastmath function time_to_next_jump(p::DirectJumpAggregation{T,S,F1,F2,RNG}, u, params, t) where {T,S,F1 <: AbstractArray,F2 <: AbstractArray, RNG}
  prev_rate = zero(t)
  new_rate  = zero(t)
  cur_rates = p.cur_rates

  # mass action rates
  majumps   = p.ma_jumps
  idx       = get_num_majumps(majumps)
  @inbounds for i in 1:idx
    new_rate     = evalrxrate(u, i, majumps)
    cur_rates[i] = new_rate + prev_rate
    prev_rate    = cur_rates[i]
  end

  # constant jump rates
  idx  += 1
  rates = p.rates
  @inbounds for i in 1:length(p.rates)
    new_rate       = rates[i](u, params, t)
    cur_rates[idx] = new_rate + prev_rate
    prev_rate      = cur_rates[idx]
    idx           += 1
  end

  @inbounds sum_rate = cur_rates[end]
  sum_rate, randexp(p.rng) / sum_rate
end


# todo: create sampler_instructions type

abstract type EMR_MSM_Model_Estimate{T<:AbstractFloat} end

struct EMR_MSM_Model_PointEstimate{T<:AbstractFloat} <: EMR_MSM_Model_Estimate{T}
    F::Array{T,1}
    Lin_mat::Array{T,2}
    Quad_mats::Array{T,3}
    RCorrs::ResCorrs{T}
    μ::AbstractArray{T,1} # mean of the deepest residual layer
    Σ::AbstractArray{T,2} # covariance of the deepes residual layer
end

# TODO add F, Lin_mat ... views ...
struct EMR_MSM_Model_DistEstimate{T<:AbstractFloat} <: EMR_MSM_Model_Estimate{T}
    estimates::Array{EMR_MSM_Model_PointEstimate{T},1}
end

function EMR_MSM_Model_PointEstimate(F::S, A::S, B::S, L::S,
    μ::S, Σ::S, num_params::Integer, num_layers::Integer) where
    S <: AbstractVector{T} where T <: Real
    L_mats = ResCorrs(L, num_params, num_layers)

    Lin_mat  = reshape(A,(num_params, num_params))
    Quad_mats = reshape(B, (num_params, num_params, num_params))

    Σ = reshape(Σ,(num_params, num_params))

    EMR_MSM_Model_PointEstimate{T}(F, Lin_mat, Quad_mats, L_mats, μ, Σ)
end

function EMR_MSM_Model_PointEstimate(dist::EMR_MSM_Model_DistEstimate{T},
    num_params::Integer, num_layers::Integer,
    aggregate_fun::Function) where T <: Real

    F = mapslices(aggregate_fun,hcat([vec(est.F) for est in dist.estimates]...);dims=2)[:,1]
    A = mapslices(aggregate_fun,hcat([vec(est.Lin_mat) for est in dist.estimates]...);dims=2)[:,1]
    B = mapslices(aggregate_fun,hcat([vec(est.Quad_mats) for est in dist.estimates]...);dims=2)[:,1]
    if num_layers > 0
        L = mapslices(aggregate_fun,hcat([vec(est.RCorrs) for est in dist.estimates]...);dims=2)[:,1]
    else
        L = Array{T}(undef, 0)
    end

    μ = mapslices(aggregate_fun,hcat([vec(est.μ) for est in dist.estimates]...);dims=2)[:,1]
    Σ = mapslices(aggregate_fun,hcat([vec(est.Σ) for est in dist.estimates]...);dims=2)[:,1]

    EMR_MSM_Model_PointEstimate(F, A, B, L,
        μ, Σ, num_params, num_layers)
end

const ts_mv_stanmodel = "
functions{
    matrix sym_mat(vector v, int num_params){
        int c;
        matrix[num_params,num_params] sym;

        c = 1;
        sym = rep_matrix(0,num_params,num_params);

        for (i in 1:(num_params-1)){
            for(j in (i+1):num_params){
                sym[i,j] = v[c];
                c = c + 1;
            }
        }

        sym = sym + sym' + diag_matrix(v[((num_params*num_params-num_params)/2+1):
                                    ((num_params*num_params+num_params)/2)]);
        return sym;
    }

    int get_rc_ind(int cur_layer, int num_params2){
        return num_params2*((cur_layer+2)*(cur_layer+1)/2 - 1);
    }

    matrix predict(row_vector start, matrix diff, int num_params, int num_data){
        matrix[num_data, num_params] pred;
        pred[1] = start;
        for(i in 2:num_data){
                pred[i] = pred[i-1] + diff[i-1];
        }
        return pred;
    }
}
data {
  int<lower=0> num_timeseries;
  int<lower=0> num_data[num_timeseries];
  int<lower=0> max_num_data;
  int<lower=0> num_params;
  int<lower=0> num_layers;
  matrix[max_num_data, num_params] x[num_timeseries];
  matrix[max_num_data-1, num_params] dx[num_timeseries];
  matrix[max_num_data, num_timeseries] time_steps;
  real<lower=0> scale_global;
  real<lower=1> nu_global;            // degree of freedom for the half-t prior
  real<lower=1> nu_local;             // degree of freedom for the half-t prior
}
transformed data{
  int num_params2 = num_params*num_params;
  int num_sym_params = (num_params*num_params + num_params) / 2;
  int num_all_sym_params = num_params*(num_params*num_params + num_params) / 2;
  int num_rescor_params =
    num_params*num_params*((num_layers+2)*(num_layers+1)/2 - 1);
}
parameters {
  real<lower=0> r1_global;
  real<lower=0> r2_global;
  real<lower=0> logsigma;
  real<lower=0> logrsigma;

  vector<lower=0>[num_params] r1_local_f;
  vector<lower=0>[num_params] r2_local_f;
  vector<lower=0>[num_params2] r1_local_l;
  vector<lower=0>[num_params2] r2_local_l;
  vector<lower=0>[num_all_sym_params] r1_local_q;
  vector<lower=0>[num_all_sym_params] r2_local_q;
  vector<lower=0>[num_rescor_params] r1_local_rc;
  vector<lower=0>[num_rescor_params] r2_local_rc;

  vector[num_params] f;
  vector[num_params2] l;
  vector[num_all_sym_params] q;
  vector[num_rescor_params] rc;
}
transformed parameters {
  real<lower=0> tau;

  vector<lower=0>[num_params] lambda_f;
  vector<lower=0>[num_params2] lambda_l;
  vector<lower=0>[num_all_sym_params] lambda_q;
  vector<lower=0>[num_rescor_params] lambda_rc;

  real sigma;
  real rsigma;
  vector[num_params] trafo_f;
  vector[num_params2] trafo_l;
  vector[num_all_sym_params] trafo_q;
  vector[num_rescor_params] trafo_rc;
  matrix[num_params,num_params] trafo_qq[num_params];

  matrix[max_num_data-1,(num_layers+1)*num_params] res_vec[num_timeseries];
  matrix[max_num_data-1, num_params] trafo_dx_hat[num_timeseries];
  matrix[max_num_data-1, num_params] r[num_layers+1,num_timeseries];
  matrix[max_num_data-1, num_params] dr[num_layers,num_timeseries];
  matrix[max_num_data-1, num_params] trafo_r_hat[num_layers,num_timeseries];
  matrix[max_num_data-1, num_params] trafo_dr_hat[num_layers,num_timeseries];
  matrix[max_num_data-1, num_params] trafo_r_tilde[num_layers,num_timeseries];
  matrix[max_num_data-1, num_params] trafo_dr_tilde[num_layers,num_timeseries];

  tau = r1_global * sqrt(r2_global);
  lambda_f = r1_local_f .* sqrt(r2_local_f);
  lambda_l = r1_local_l .* sqrt(r2_local_l);
  lambda_q = r1_local_q .* sqrt(r2_local_q);
  lambda_rc = r1_local_rc .* sqrt(r2_local_rc);

  trafo_f = f .* lambda_f * tau;
  trafo_l = l .* lambda_l * tau;
  trafo_q = q .* lambda_q * tau;
  trafo_rc = rc .* lambda_rc * tau;

  sigma = exp(logsigma);
  rsigma = exp(logrsigma);

  for(ts in 1:num_timeseries){
      trafo_dx_hat[ts] = rep_matrix(0, max_num_data-1, num_params);
      res_vec[ts] = rep_matrix(0, max_num_data-1,(num_layers+1)*num_params);
      for(i in 1:num_layers){
          trafo_dr_hat[i, ts] = rep_matrix(0, max_num_data-1, num_params);
          trafo_r_hat[i, ts] = rep_matrix(0, max_num_data-1, num_params);
          trafo_r_tilde[i, ts] = rep_matrix(0, max_num_data-1, num_params);
          dr[i, ts] = rep_matrix(0, max_num_data-1, num_params);
          r[i, ts] = rep_matrix(0, max_num_data-1, num_params);
          trafo_dr_tilde[i, ts] = rep_matrix(0, max_num_data-1, num_params);
      }

      for (i in 1:(num_data[ts]-1)){
        trafo_dx_hat[ts][i] =
            -x[ts][i]*to_matrix(trafo_l, num_params, num_params)' + trafo_f';
        for (j in 1:num_params){
            trafo_qq[j] = sym_mat(
                    trafo_q[(1+(j-1)*num_sym_params):(j*num_sym_params)],
                    num_params);
            trafo_dx_hat[ts][i,j] = trafo_dx_hat[ts][i,j] +
                quad_form(trafo_qq[j],
                    x[ts][i]');
        }
        trafo_dx_hat[ts][i] = time_steps[i+1,ts]*trafo_dx_hat[ts][i];
      }

      r[1, ts] = dx[ts] - trafo_dx_hat[ts];

      if(num_layers > 0){

          res_vec[ts][:,:num_params] = x[ts][:(num_data[ts]-1),:];
          for(i in 1:num_layers){
            dr[i, ts][1:(num_data[ts]-2),:] = r[i ,ts][2:,:] - r[i, ts][:(num_data[ts]-2),:];
            res_vec[ts][:,(i*num_params+1):((i+1)*num_params)] = r[i, ts];
            for(j in 1:(num_data[ts]-1-i)){
                trafo_dr_hat[i, ts][j] = time_steps[j+1, ts] *
                    res_vec[ts][j,:(num_params*(i+1))] * to_matrix(trafo_rc[
                    (get_rc_ind(i-1, num_params2)+1):get_rc_ind(i, num_params2)],
                    num_params, (i+1)*num_params)';
            }

            r[i+1, ts] = dr[i, ts] - trafo_dr_hat[i, ts];
          }

          trafo_dr_tilde[num_layers, ts] = trafo_dr_hat[num_layers ,ts] ;
          for(i in (-num_layers):(-1)){
              trafo_r_tilde[-i, ts] = predict(r[-i, ts][1], trafo_dr_tilde[-i, ts],
                num_params, num_data[ts]-1);
              if(-i > 1){
                    trafo_dr_tilde[-1-i, ts] = trafo_dr_hat[-1-i, ts] +
                        trafo_r_tilde[-i, ts];
              }
          }

          trafo_dx_hat[ts] = trafo_dx_hat[ts] + trafo_r_tilde[1, ts];
      }
  }
}
model {
  r1_global ~ normal(0.0, scale_global*sigma);
  r2_global ~ inv_gamma(0.5*nu_global, 0.5*nu_global);

  r1_local_f ~ normal(0.0, 1);
  r2_local_f ~ inv_gamma(0.5*nu_local, 0.5*nu_local);
  r1_local_l ~ normal(0.0, 1);
  r2_local_l ~ inv_gamma(0.5*nu_local, 0.5*nu_local);
  r1_local_q ~ normal(0.0, 1);
  r2_local_q ~ inv_gamma(0.5*nu_local, 0.5*nu_local);
  r1_local_rc ~ normal(0.0, 1);
  r2_local_rc ~ inv_gamma(0.5*nu_local, 0.5*nu_local);

  f ~ normal(0,1);
  l ~ normal(0,1);
  q ~ normal(0,1);

  for(ts in 1:num_timeseries){
      to_vector(dx[ts][1:(num_data[ts]-1-num_layers)]) ~ normal(
        to_vector(trafo_dx_hat[ts][1:(num_data[ts]-1-num_layers)]), sigma);
  }
}
";

function EMR_MSM_Model_DistEstimate(timeseries::AbstractArray{MSM_Timeseries_Point{T},1},
    num_layers::Integer, num_samples::Integer, num_chains::Integer, tau0::T=one(T);
    uncorr = true, rand_seed::Integer=123) where T <: Real

    num_timeseries = length(timeseries)

    num_obs = length.(timeseries)
    max_num_obs = maximum(num_obs)

    num_params = params(timeseries[1])

    x = values.(timeseries)
    dx = [(x[i][2:num_obs[i],:]) - (x[i][1:num_obs[i]-1, :]) for i in 1:num_timeseries]

    x_stan = permutedims(reshape(hcat(x...),(max_num_obs, num_params, num_timeseries)),
                        [3,1,2])
    dx_stan = permutedims(reshape(hcat(dx...),(max_num_obs-1, num_params, num_timeseries)),
                         [3,1,2])

    # TODO random...
    stanmodel = Stanmodel(name="timeseries_mv" * string(rand(1:10000)),
        model=ts_mv_stanmodel,
        random=CmdStan.Random(rand_seed), nchains=num_chains, num_samples=num_samples)

    ts_data = Dict(
            "num_timeseries" => num_timeseries,
            "num_data" => num_obs,
            "max_num_data" => max_num_obs,
            "num_params" => num_params,
            "num_layers" => num_layers,
            "x" => x_stan,
            "dx" => dx_stan,
            "time_steps" => hcat(timesteps.(timeseries)...),
            "scale_global" => tau0 / sqrt(sum(num_obs)),
            "nu_global" => 1,
            "nu_local" => 1
            );

    rc, chns, cnames = stan(stanmodel, ts_data, "./", CmdStanDir=CMDSTAN_HOME)

    dx_est = [[hcat([[chns.value[:,"trafo_dx_hat." * string(ts) *"." * string(i) *"."* string(p), (div(j-1,num_samples)+1)
                        ][mod(j-1,num_samples)+1]
                        for i in 1:(num_obs[ts]-1)]
                        for p in 1:num_params]...)
                        for j in 1:(num_samples*num_chains)]
                        for ts in 1:num_timeseries]

    res_est = [[[hcat([[chns.value[:,"r."* string(l) *"." * string(ts) *"."  * string(i) *"."* string(p), (div(j-1,num_samples)+1)
                        ][mod(j-1,num_samples)+1]
                        for i in 1:(num_obs[ts]-1)]
                        for p in 1:num_params]...)
                        for l in 1:(num_layers+1)]
                        for j in 1:(num_samples*num_chains)]
                        for ts in 1:num_timeseries]

    dx_est_t = [reshape(hcat(dx_est[ts]...),(num_obs[ts]-1,num_params,num_samples*num_chains))
                for ts in 1:num_timeseries]
    res_est_t = [reshape(hcat(hcat(res_est[ts]...)...),
        (num_obs[ts]-1,num_params,num_layers+1,num_samples*num_chains))
        for ts in 1:num_timeseries]

    x_est = [Array{Float64,3}(undef, num_obs[ts]-1, num_params, num_samples*num_chains)
            for ts in 1:num_timeseries]
    for ts in 1:num_timeseries
        for j in 1:(num_samples*num_chains)
            x_est[ts][1,:,j] = x[ts][1,:]
            for i in 2:(num_obs[ts]-1)
                x_est[ts][i,:,j] = x_est[ts][i-1,:,j] + dx_est_t[ts][i-1,:,j]
            end
        end
    end

    pred_timeseries = [MSM_Timeseries_Dist{T}(x_est[ts], res_est_t[ts],
        timesteps.(timeseries)[ts][1:(num_obs[ts]-1)])
        for ts in 1:num_timeseries]

    all_ll_res = [vcat([res_est_t[ts][:,:,num_layers+1,j] for ts in 1:num_timeseries]...)
                    for j in 1:num_samples*num_chains]

    summary_df = read_summary(stanmodel)
    conv_info = DataFrames.DataFrame()

    for (name,len) in [("trafo_f", num_params), ("trafo_l", num_params^2),
        ("trafo_rc", length(ResCorrs, num_params, num_layers))]
        for i in 1:len
            sdf = summary_df[Symbol(name*"[" * string(i) *"]"),[:r_hat, Symbol("n_eff/s")]]
            sdf[!,Symbol("name")] = [name * "-" * string(i)]

            push!(conv_info, sdf[1,:])
        end
    end

    c = 1
    for i in 1:num_params
        for k in 1:num_params
            for l in 1:num_params
                sdf = summary_df[Symbol(
                "trafo_qq[" * string(i) * "," * string(k) * "," * string(l) * "]"),
                [:r_hat, Symbol("n_eff/s")]]

                sdf[!,Symbol("name")] = ["trafo_qq-" * string(c)]
                push!(conv_info, sdf[1,:])
                c += 1
            end
        end
    end

    return EMR_MSM_Model_DistEstimate{T}([EMR_MSM_Model_PointEstimate(
        [vec(chns.value[:,"trafo_f." * string(i),
        (div(j-1,num_samples)+1)])[mod(j-1,num_samples)+1] for i in 1:num_params],
        [vec(chns.value[:,"trafo_l." * string(i),
        (div(j-1,num_samples)+1)])[mod(j-1,num_samples)+1] for i in 1:num_params^2],
        vcat([vec(hcat([vec([chns.value[:,"trafo_qq." * string(i) *"."* string(k) *"." * string(l),
        (div(j-1,num_samples)+1)][mod(j-1,num_samples)+1]
            for k in 1:num_params]) for l in 1:num_params]...)) for i in 1:num_params]...),
        if num_layers == 0
            Array{T}(undef,length(ResCorrs, num_params, num_layers))
        else
            [vec(chns.value[:,"trafo_rc." * string(i),
            (div(j-1,num_samples)+1)])[mod(j-1,num_samples)+1] for i in 1:length(ResCorrs, num_params, num_layers)]
        end
        ,
        ifelse(uncorr, zeros(T, num_params),mean.([all_ll_res[j][:,i] for i in 1:num_params])),
        ifelse(uncorr, collect(vec(Diagonal(cov(all_ll_res[j])))),
            vec(cov(all_ll_res[j]))),
        num_params,
        num_layers
        ) for j in 1:(num_samples*num_chains)]),
        pred_timeseries, conv_info
end

cmdstan_home!(home::String) = set_cmdstan_home!(home)

Base.size(emr_msm::EMR_MSM_Model_DistEstimate) = size(emr_msm.estimates, 1)

function DataRow(model::EMR_MSM_Model_PointEstimate)
    df = DataFrames.DataFrame()
    fnames = fieldnames(EMR_MSM_Model_PointEstimate)

    for fname in fnames
        item = vec(getfield(model, fname))
        for i in 1:length(item)
            df[!,Symbol(string(fname) * string(i))] = [item[i]]
        end
    end

    df[1,:]
end

function DataFrame(model::EMR_MSM_Model_DistEstimate)
    drs = DataRow.(model.estimates)

    push!(DataFrames.DataFrame(drs[1]),drs[2:length(model.estimates)]...)
end

write(output::String, model::EMR_MSM_Model_PointEstimate) =
    CSV.write(output * ".csv", DataFrames.DataFrame(DataRow(model)))

write(output::String, model::EMR_MSM_Model_DistEstimate) =
    CSV.write(output * ".csv", DataFrame(model))


import DelimitedFiles
import OrderedCollections

function madsmodelrun_adm1(mddata::AbstractDict)
    # function runs the ADM1F model and return selected values from the last result file as predictions

    function get_last_indicator()
        # finds the last indicator file in the current folder
        flist = readdir()
        last_indicator = ""
        for i = 1:length(flist)
            if length(flist[i]) > 9
                if flist[i][1:9] == "indicator"
                    last_indicator = flist[i]
                end
            end
        end
        return last_indicator
    end

    #reads .dat and saves it in mddata
    DelimitedFiles.writedlm("params_tim.dat", Mads.getparamsinit(mddata))

    #removes all previous .out files
    foreach(rm, filter(endswith(".out"), readdir()))

    #runs the ADM1F model with the command line below
    #change the path to the executable below
    command_line=`/Users/elchin/project/ADM1F_WM/build/adm1f -steady -ts_monitor
             -influent_file influent_tim.dat -params_file params_tim.dat
             -ts_type beuler -ts_adapt_type basic -snes_rtol 1.e-5
             -ts_max_snes_failures -1 -Vliq 4981 -t_resx 240`
    Mads.runcmd(command_line; quiet=true, pipe=true)

    #load the outputs from last indicator
    #we are interested only in 7 output values out of 67
    #we return those 7 values as predictions
    last_indicator = get_last_indicator()
    if last_indicator != ""
        o = DelimitedFiles.readdlm(last_indicator)[3:end,:]
        outputs = OrderedCollections.OrderedDict{String, Float64}("o$i"=>o[i] for i=1:67)
        sCOD=0
        for i = 1:7
            sCOD=sCOD+outputs["o$i"]
        end
        Acetate=outputs["o7"]
        Propionate = outputs["o6"]
        Butyrate = outputs["o5"]
        Valerate = outputs["o4"]
        ph = outputs["o27"]
        gaseous = outputs["o41"]*outputs["o44"]*1000
    else
        @error "Output is missing!"
    end
    p=Dict([("o1", sCOD),
            ("o2", Acetate),
            ("o3", Propionate),
            ("o4", Butyrate),
            ("o5", Valerate),
            ("o6", ph),
            ("o7", gaseous)])
    #println("p:",p)
    return p
end


#===============================================================================

    thermal_consistency.jl

    72°C thermal consistency test

===============================================================================#

import DataStructures.OrderedDict
import Memento: debug

#===============================================================================
    constants >>
===============================================================================#

## preset values
# const MIN_FLUORESCENCE_VAL = 8e5 ## not used
const MIN_TM_VAL = 77
const MAX_TM_VAL = 81
const MAX_DELTA_TM_VAL = 2

#===============================================================================
    function >>
===============================================================================#

## called by dispatch()
##
## NB default values supplied to mc_analysis()
## are the same as for melting curve experiments
function act(
    ::Type{Val{thermal_consistency}},
    ## remove MySql dependency
    # db_conn ::MySQL.MySQLHandle,
    # exp_id ::Integer,
    # stage_id ::Integer,
    # calib_info ::Union{Integer,OrderedDict};
    req                 ::Associative;
    dcv                 ::Bool          = DEFAULT_CAL_DCV, ## if true, perform multi-channel deconvolution
    auto_span_smooth    ::Bool          = DEFAULT_MC_AUTO_SPAN_SMOOTH,
    span_smooth_default ::Real          = DEFAULT_MC_SPAN_SMOOTH_DEFAULT,
    span_smooth_factor  ::Real          = DEFAULT_MC_SPAN_SMOOTH_FACTOR,
    max_temperature     ::Real          = DEFAULT_MC_MAX_TEMPERATURE, ## maximum temperature to analyze
    out_format          ::OutputFormat  = pre_json_output,
    reporting           ::Function      = roundoff(JSON_DIGITS) ## reporting function
)
    debug(logger, "at act(::Type{Val{thermal_consistency}})")
    #
    ## required fields
    @get_calibration_data_from_req(thermal_consistency)
    @parse_raw_data_from_req(thermal_consistency)
    #
    ## keyword arguments
    kwargs = MC_FIELD_DEFS |>
        sift(x -> x.key in keys(req)) |>
        mold() do x
			if isa(req[x.key], Int64) && Int_T != Int64
				x.name => Int_T(req[x.key])
			else
				x.name => req[x.key]
			end
        end
    #
    ## create container for data and parameter values
    interface = McInput(
            parsed_raw_data...,
            calibration_data;
            calibration_args = CalibrationParameters(dcv = dcv),
            auto_span_smooth = auto_span_smooth,
            span_smooth_default = span_smooth_default,
            span_smooth_factor = span_smooth_factor,
            max_temperature = max_temperature,
            out_format = full_output,
            reporting = reporting,
            kwargs...)
    #
    ## analyse data analyses melting curve
    mc_w72c =
        mc_analysis(interface)
    #
    ## process the data from only one channel
    ## PROBLEM >> this does not seem appropriate for dual channel analysis
    mc_tm = map(
        field(:peaks_filtered),
        mc_w72c.peak_output[:, CHANNELS[1]]) ## mc_matrix
    min_Tm = max_temperature + 1
    max_Tm = 0
    tm_check_vec = map(mc_tm) do Ta
        if size(Ta, 1) == 0
            TmCheck1w((NaN_T, false), NaN_T)
        else
            top1_Tm = Ta[1,1]
            (top1_Tm < min_Tm) && (min_Tm = top1_Tm)
            (top1_Tm > max_Tm) && (max_Tm = top1_Tm)
            TmCheck1w(
                (top1_Tm, MIN_TM_VAL <= top1_Tm <= MAX_TM_VAL),
                Ta[1,2])
        end ## if size
    end ## do Ta
    #
    ## return values
    delta_Tm_val = max_Tm - min_Tm
    full_out() =
        ThermalConsistencyOutput(
            tm_check_vec,
            (delta_Tm_val, delta_Tm_val .<= MAX_DELTA_TM_VAL),
            true)
    pre_json_out() =
        OrderedDict(
            :tm_check => tm_check_vec,
            :delta_Tm => (reporting(delta_Tm_val),
                            delta_Tm_val .<= MAX_DELTA_TM_VAL),
            :valid    => true)
    (out_format == full_output) && return full_out()
    ## else
    return pre_json_out() |> out(out_format)
end ## act(::Type{Val{thermal_consistency}})


export CuForwardEulerTimeStepper, CuFilteredForwardEulerTimeStepper,
       CuRK4TimeStepper, CuFilteredRK4TimeStepper

struct CuForwardEulerTimeStepper{T,dim} <: AbstractForwardEulerTimeStepper
  dt::Float64
  N::CuArray{Complex{T},dim}    # Explicit linear and nonlinear terms
  CuForwardEulerTimeStepper(dt::T, N::CuArray{Complex{T},dim}) where {T,dim} = new{T,dim}(
    dt, CuArray(zeros(size(N))))
end

CuForwardEulerTimeStepper(dt, N::Array{Complex{T},dim}) where {T,dim} = CuForwardEulerTimeStepper(dt, CuArray(N))

struct CuFilteredForwardEulerTimeStepper{T,dim} <: AbstractFilteredForwardEulerTimeStepper
  dt::Float64
  N::CuArray{Complex{T},dim}    # Explicit linear and nonlinear terms
  filter::CuArray{T,dim}        # Filter for solution
end

function CuFilteredForwardEulerTimeStepper(dt, LC, g; filterkwargs...)
  N = CuArray(zeros(cxeltype(LC), size(LC)))
  filter = makefilter(g, typeof(dt), size(LC); filterkwargs...)
  CuFilteredForwardEulerTimeStepper{ndims(N)}(dt, N, filter)
end

struct CuRK4TimeStepper{T,dim} <: AbstractRK4TimeStepper
  dt::Float64
  sol₁::CuArray{T,dim}
  RHS₁::CuArray{T,dim}
  RHS₂::CuArray{T,dim}
  RHS₃::CuArray{T,dim}
  RHS₄::CuArray{T,dim}
end

function CuRK4TimeStepper(dt, LC)
  @createcuarrays cxeltype(LC) size(LC) sol₁ RHS₁ RHS₂ RHS₃ RHS₄
  CuRK4TimeStepper{cxeltype(LC),ndims(LC)}(dt, sol₁, RHS₁, RHS₂, RHS₃, RHS₄)
end

struct CuFilteredRK4TimeStepper{T,dim} <: AbstractFilteredRK4TimeStepper
  dt::Float64
  sol₁::CuArray{T,dim}
  RHS₁::CuArray{T,dim}
  RHS₂::CuArray{T,dim}
  RHS₃::CuArray{T,dim}
  RHS₄::CuArray{T,dim}
  filter::CuArray{T,dim}    # Filter for solution
end

function CuFilteredRK4TimeStepper(dt, LC, g; filterkwargs...)
  @createcuarrays cxeltype(LC) size(LC) sol₁ RHS₁ RHS₂ RHS₃ RHS₄
  filter = makefilter(g, typeof(dt), size(LC); filterkwargs...)
  CuFilteredRK4TimeStepper{cxeltype(LC),ndims(LC)}(dt, sol₁, RHS₁, RHS₂, RHS₃, RHS₄, filter)
end

ForwardEulerTimeStepper(dt, LC::CuArray{T,dim}) where {T,dim} = CuForwardEulerTimeStepper(dt, LC)
RK4TimeStepper(dt, LC::CuArray{T,dim}) where {T,dim} = CuRK4TimeStepper(dt, LC)

FilteredForwardEulerTimeStepper(
  dt, LC::CuArray{T,dim}, g; filterkwargs...) where {T,dim} = CuFilteredForwardEulerTimeStepper(dt, LC, g;
                                                                                filterkwargs...)
FilteredRK4TimeStepper(
  dt, LC::CuArray{T,dim}, g; filterkwargs...) where {T,dim} = CuFilteredRK4TimeStepper(dt, LC, g; filterkwargs...)


/**
 * This class models a cloud storage API.
 */
class Main {

    static int xTESTOUTPUT = 7;
    static int yTESTOUTPUT = 7;
    static int zTESTOUTPUT = 7;

    public static void main(String[] args) {

	// the following is dangerous data that should be erased sometime.
	condition c = false;

	int dangerous = ({L /c H})args.length;

	// get a key pair
	int pub = (PUBKEY(L /c H){L})42;
	int priv = (PRIVKEY(L /c H){L /c H})42;

	// let's encrypt it, and put the result in the cloud.

	Cloud cld = new Cloud();
	xTESTOUTPUT = CryptoLibrary.encrypt(pub, dangerous); // Should be L

	cld.put(CryptoLibrary.encrypt(pub, dangerous));

	// set the condition.
	c = true;

	yTESTOUTPUT = cld.get(); // should be L, since the ciphertext is still low security.

	// at this point, priv should no longer be accessible at PRIVKEY.
	zTESTOUTPUT = CryptoLibrary.decrypt(priv, cld.get()); // the level should be H, since it can no longer be at L /c H, since c was set.
    }
}


#!/usr/bin/julia
#
# Implementation of simulation variations common across multiple simulation methods.
#
export AdaptiveΔt;

"Adaptive timestep parameters."
struct AdaptiveΔt <: SimVariation
  "Maximum error before reducing Δt."
  ϵmax::Float64;
  "Minimum error before increasing Δt."
  ϵmin::Float64;

  "Factor to reduce Δt by."
  Δtdown::Float64;
  "Factor to increase Δt by."
  Δtup::Float64;
  "Maximum allowable Δt."
  Δtmax::Float64;
  "Minimum allowable Δt."
  Δtmin::Float64;

  "Maximum number of repeated evaluations per iteration."
  maxeval::UInt;

  """
  Maximum number successive iterations where the evaluation limit was reached
  before a different parameter should be altered.

  This might be undefined. As an example, in the Lanczos method, when this
  limit is reached the basis is regenerated.
  """
  maxevallimit::Union{UInt, Nothing};

  # TODO Implement inner constructor to check validity of passed arguments.
end
AdaptiveΔt(;ϵmax::Real=1e-8, ϵmin::Real=1e-12,
            Δtdown::Real=1.2, Δtup::Real=1.5, Δtmax::Real=0.5, Δtmin::Real=1e-5,
            maxeval::Real=5, maxevallimit::Real=3) =
  AdaptiveΔt(ϵmax, ϵmin, Δtdown, Δtup, Δtmax, Δtmin, maxeval, maxevallimit);

"""
    adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number; debug::Bool=true)

Adapt the time step `Δt` using the error in the last propagation step, `ϵ`
according to the settings passed in `par`.
"""
function adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number; debug::Bool=true)
  # Check passed arguments.
  if (ϵ <= 0)
    throw(DomainError(ϵ, "error must be positive."));
  elseif (imag(Δt) == 0 && real(Δt) <= 0)
    throw(DomainError(Δt, "timestep must be positive."));
  elseif (real(Δt) == 0 && imag(Δt) <= 0)
    throw(DomainError(Δt, "timestep must be positive."));
  elseif (real(Δt) > 0 && imag(Δt) > 0)
    throw(DomainError(Δt, "timestep must be either real or imaginary."));
  end

  # Increase timestep if less than the minimum error.
  if (ϵ < par.ϵmin)
    # Check that increasing the timestep is allowed.
    if (abs(Δt) > par.Δtmax)
      println("Reached maximum allowable timestep");
    else
      Δt *= par.Δtup;
      if (debug) println("Increasing timestep to ", Δt); end
    end
  elseif (ϵ > par.ϵmax)
    # Check that decreasing timestep is allowed.
    if (abs(Δt) < par.Δtmin)
      println("Reached minimum allowable timestep");
    # Decrease timestep.
    else
      Δt /= par.Δtdown;
      if (debug) println("Decreasing timestep to ", Δt); end
    end
  end

  return Δt;
end

"""
    adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number, evals::Real; debug::Bool=true)

Return adapted time step and decision on if previous evaluation should be
repeated using the number of evaluations already completed.

This returns a `Tuple` with the first value being the new time step to use and
the second a `Boolean` indicating if re-evaluation is required.
"""
function adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number, evals::Real;
                  debug::Bool=true)
  newΔt = adapt_Δt(par, ϵ, Δt, debug=debug);

  # Repeat evaluation if time step was reduced and maximum evaluations not
  # reached.
  if (newΔt < Δt && evals < par.maxeval) return (newΔt, true);
  else return (newΔt, false); end
end
"""
    adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number, evals::Vector{T};
             debug::Bool=true) where T<:Union{Int,UInt}

Return adapted time step, re-evaluation decision and if the `maxevallimit` was
reached.

The `evals` vector contains the number of evaluations for all completed
propagation steps (including the current one) with zero values for steps not yet
completed.
"""
function adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number, evals::Vector{T};
                  debug::Bool=true) where T<:Union{Int,UInt}
  # Make copy of `evals` to avoid changing its values.
  evalscopy = copy(evals);
  # Remove evaluations which haven't occurred yet (evaluation counter is zero).
  filter!(v->(v!=0), evalscopy);

  (Δt, reeval) = adapt_Δt(par, ϵ, Δt, evalscopy[end], debug=debug);

  # Forward result if re-evaluation was not required.
  if !reeval return (Δt, false, false); end

  maxevallimitreached = true;
  # Convert to Int, otherwise result will be cast to UInt and will overflow if result is negative.
  min = max(1, length(evalscopy)-convert(Int,par.maxevallimit));
  for i ∈ evalscopy[min]:evalscopy[end]
    if i < par.maxeval
      maxevallimitreached = false;
      break;
    end
  end

  return (Δt, reeval, maxevallimitreached);
end
"""
    adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number, curreval::Real,
             evals::Vector{T}; debug::Bool=true) where T<:Union{Int,UInt}

Here `curreval` contains the index storing the number of evaluations for the
current propagation step within `evals`.
"""
function adapt_Δt(par::AdaptiveΔt, ϵ::Real, Δt::Number, curreval::Real,
                  evals::Vector{T}; debug::Bool=true) where T<:Union{Int,UInt}
  if par.maxevallimit === nothing
    error("maxium evalution limit was not initalised.");
  end

  (Δt, reeval) = adapt_Δt(par, ϵ, Δt, evals[curreval], debug=debug);

  # Forward result if re-evaluation was not required.
  if !reeval return (Δt, false, false); end

  # TODO: Check UInt overflow characteristics here.
  maxevallimitreached = true;
  for i ∈ max(curreval-par.maxevallimit, 1):curreval
    if evals[i] < par.maxeval
      maxevallimitreached = false;
      break;
    end
  end

  return (Δt, reeval, maxevallimitreached);
end


include("abstract_bond.jl")

"""
Analysis data of determinant quantum Monte Carlo (DQMC) simulation
"""
@with_kw mutable struct DQMC_bondAnalysis
    acc_rate::Float64 = 0.
    prop_local::Int = 0
    acc_local::Int = 0
    acc_rate_global::Float64 = 0.
    prop_global::Int = 0
    acc_global::Int = 0
end

"""
Parameters of determinant quantum Monte Carlo (DQMC)
"""
@with_kw struct DQMC_bondParameters
    global_moves::Bool = false
    global_rate::Int = 5
    thermalization::Int = 100 # number of thermalization sweeps
    num_ch::Int = 4
    sweeps::Int = num_ch*25 # number of sweeps (after thermalization)
    all_checks::Bool = true # e.g. check if propagation is stable/instable
    safe_mult::Int = num_ch*1

    delta_tau::Float64 = 0.01
    beta::Float64

    time_slices::Int = beta / delta_tau
    slices::Int  = num_ch*time_slices
    @assert isinteger(beta / delta_tau) string("beta/delta_tau", "
        (= number of imaginary time slices) must be an integer but is",
        beta / delta_tau, ".")

    measure_every_nth::Int = 10
end

"""
Determinant quantum Monte Carlo (DQMC_bond) simulation
"""
#mutable struct DQMC_bond{M<:Model, CB<:Checkerboard, ConfType<:Any,
mutable struct DQMC_bond{M<:Model, ConfType<:Any,
            Stack<:AbstractDQMC_bondStack} <: MonteCarloFlavor
    model::M
    conf::ConfType
    #hopping_mat::Array{Float64,5}
    hopping_term_up:: Array{Float64,2}
    hopping_term_down:: Array{Float64,2}
    diag_terms::Array{Float64,1} #diagonal terms including chemical potential and onsite disorder
    diag_terms_inv::Array{Float64,1} #inverse diagonal terms including chemical potential and onsite disorder
    s::Stack

    p::DQMC_bondParameters
    a::DQMC_bondAnalysis
    obs::Dict{String, Observable}

    K_xy::Float64
    K_tau::Float64
    DQMC_bond{M, ConfType, Stack}() where {M<:Model, ConfType<:Any, Stack<:AbstractDQMC_bondStack} = new()
end

include("stack_bond.jl")
include("slice_matrices_bond.jl")

"""
    DQMC_bond(m::M; kwargs...) where M<:Model

Create a determinant quantum Monte Carlo simulation for model `m` with
keyword parameters `kwargs`.
"""
function DQMC_bond(m::M; seed::Int=-1, kwargs...) where M<:Model

    p = DQMC_bondParameters(; kwargs...)
    geltype = greenseltype(DQMC_bond, m)
    conf = ones(Int8,2*m.L^m.dims, p.time_slices)#rand(DQMC_bond, m, p.time_slices)#
    mc = DQMC_bond{M, typeof(conf), DQMC_bondStack{geltype,Float64}}()
    mc.conf = conf
    mc.model = m
    mc.p = p
    mc.s = DQMC_bondStack{geltype,Float64}()
    mc.K_xy = m.J*p.delta_tau
    mc.K_tau = 0.5*log(coth(m.h*p.delta_tau))
    init!(mc)# , seed=seed, conf=conf)#why is the conf initialized twice??
    return mc
end

"""
    DQMC(m::M, params::Dict)
    DQMC(m::M, params::NamedTuple)

Create a determinant quantum Monte Carlo simulation for model `m` with
(keyword) parameters as specified in the dictionary/named tuple `params`.
"""
DQMC_bond(m::Model, params::Dict{Symbol, T}) where T = DQMC_bond(m; params...)
DQMC_bond(m::Model, params::NamedTuple) = DQMC_bond(m; params...)

# convenience
@inline beta(mc::DQMC_bond) = mc.p.beta
@inline nslices(mc::DQMC_bond) = mc.p.slices
@inline model(mc::DQMC_bond) = mc.model
@inline conf(mc::DQMC_bond) = mc.conf
@inline current_slice(mc::DQMC_bond) = mc.s.current_slice

# cosmetics
import Base.summary
import Base.show
Base.summary(mc::DQMC_bond) = "DQMC_bond simulation of $(summary(mc.model))"
function Base.show(io::IO, mc::DQMC_bond)
    print(io, "Determinant quantum Monte Carlo simulation\n")
    print(io, "Model: ", mc.model, "\n")
    print(io, "Beta: ", mc.p.beta, " (T ≈ $(round(1/mc.p.beta, sigdigits=3)))")
end
Base.show(io::IO, m::MIME"text/plain", mc::DQMC_bond) = print(io, mc)

"""
    init!(mc::DQMC[; seed::Real=-1])

Initialize the determinant quantum Monte Carlo simulation `mc`.
If `seed !=- 1` the random generator will be initialized with `Random.seed!(seed)`.
"""
function init!(mc::DQMC_bond)#; seed::Real=-1, conf=rand(DQMC_bond,model(mc),nslices(mc)))
    #seed == -1 || Random.seed!(seed)

    #mc.conf = conf
    #mc.hopping_mat = init_hopping_matrices(mc, mc.model)
    init_hopping_matrices(mc, mc.model)
    init_diag_terms(mc, mc.model)
    initialize_stack(mc)
    mc.obs = prepare_observables(mc, mc.model)
    mc.a = DQMC_bondAnalysis()
    nothing
end

"""
    run!(mc::DQMC[; verbose::Bool=true, sweeps::Int, thermalization::Int])

Runs the given Monte Carlo simulation `mc`.
Progress will be printed to `stdout` if `verbose=true` (default).
"""
function run!(mc::DQMC_bond; verbose::Bool=true, sweeps::Int=mc.p.sweeps,
    thermalization=mc.p.thermalization)
    num_ch=mc.p.num_ch #number of checkerboards
    total_sweeps = num_ch*(sweeps + thermalization)

    start_time = now()
    verbose && println("Started: ", Dates.format(start_time, "d.u yyyy HH:MM"))

    # fresh stack
    verbose && println("Preparing Green's function stack")
    initialize_stack(mc) # redundant ?!
    build_stack(mc)

    propagate(mc)

    _time = time()
    verbose && println("\n\nThermalization stage - ", thermalization)
    for i in 1:total_sweeps
        verbose && (i == thermalization + 1) &&
            println("\n\nMeasurement stage - ", sweeps)
        for u in 1:2 * nslices(mc)
            update(mc, i)

            if i > thermalization && current_slice(mc) == nslices(mc) &&
                    mc.s.direction == -1 && (i-1)%mc.p.measure_every_nth == 0
                measure_observables!(mc, mc.model, mc.obs, mc.conf)
            end
        end

        if mod(i, 10) == 0
            mc.a.acc_rate = mc.a.acc_rate / (10 * 2 * nslices(mc))
            mc.a.acc_rate_global = mc.a.acc_rate_global / (10 / mc.p.global_rate)
            sweep_dur = (time() - _time)/10
            if verbose
                println("\t", i)
                @printf("\t\tsweep dur: %.3fs\n", sweep_dur)
                @printf("\t\tacc rate (local) : %.1f%%\n", mc.a.acc_rate*100)
                if mc.p.global_moves
                  @printf("\t\tacc rate (global): %.1f%%\n", mc.a.acc_rate_global*100)
                  @printf("\t\tacc rate (global, overall): %.1f%%\n",
                    mc.a.acc_global/mc.a.prop_global*100)
                end
            end

            mc.a.acc_rate = 0.0
            mc.a.acc_rate_global = 0.0
            flush(stdout)
            _time = time()
        end
    end
    finish_observables!(mc, mc.model, mc.obs)

    mc.a.acc_rate = mc.a.acc_local / mc.a.prop_local
    mc.a.acc_rate_global = mc.a.acc_global / mc.a.prop_global

    end_time = now()
    verbose && println("Ended: ", Dates.format(end_time, "d.u yyyy HH:MM"))
    verbose && @printf("Duration: %.2f minutes", (end_time - start_time).value/1000. /60.)

    nothing
end

"""
    update(mc::DQMC, i::Int)

Propagates the Green's function and performs local and global updates at
current imaginary time slice.
"""
function update(mc::DQMC_bond, i::Int)
    propagate(mc)
    sweep_spatial(mc)

    nothing
end

"""
    sweep_spatial(mc::DQMC_bond,cb::Int)

Performs a sweep of local moves along spatial dimension at current
imaginary time slice.
"""
function sweep_spatial(mc::DQMC_bond)

    m = model(mc)
    N = div(nsites(m),2) # number of bonds in a checkerboard
    bond_checkerboard = m.bond_checkerboard
    slice = current_slice(mc)
    num_ch=mc.p.num_ch
    time_slice = cld(slice,num_ch) # imagenary time
    cb = mod1(slice,num_ch) # number of checkerboard in time slice
    @inbounds for b in 1:N
        i = bond_checkerboard[1,b,cb] # index of bond i in checkerboard ch
        detratio, ΔE_boson, Δ = propose_local(mc, m, i, slice, conf(mc))
        mc.a.prop_local += 1
        if abs(imag(detratio)) > 1e-6
            println("Did you expect a sign problem? imag. detratio: ",
                abs(imag(detratio)))
            @printf "%.10e" abs(imag(detratio))
        end
        p = real(exp(- ΔE_boson) * detratio)

        # Metropolis
        if p > 1 || rand() < p
            accept_local!(mc, m, i, cb, time_slice, conf(mc), Δ, detratio,
                ΔE_boson)
            mc.a.acc_rate += 1/N
            mc.a.acc_local += 1
        end
    end
    nothing
end
#=
"""
    greens(mc::DQMC)

Obtai the current equal-time Green's function.

Internally, `mc.s.greens` is an effective Green's function. This method transforms
this effective one to the actual Green's function by multiplying hopping matrix
exponentials from left and right.
"""
function greens(mc::DQMC_CBFalse)
    eThalfminus = mc.s.hopping_matrix_exp
    eThalfplus = mc.s.hopping_matrix_exp_inv

    greens = copy(mc.s.greens)
    greens .= greens * eThalfminus
    greens .= eThalfplus * greens
    return greens
end
function greens(mc::DQMC_CBTrue)
    chkr_hop_half_minus = mc.s.chkr_hop_half
    chkr_hop_half_plus = mc.s.chkr_hop_half_inv

    greens = copy(mc.s.greens)

    @inbounds @views begin
        for i in reverse(1:mc.s.n_groups)
          greens .= greens * chkr_hop_half_minus[i]
        end
        for i in reverse(1:mc.s.n_groups)
          greens .= chkr_hop_half_plus[i] * greens
        end
    end
    return greens
end
=#
include("DQMC_bond_mandatory.jl")
include("DQMC_bond_optional.jl")