using Distributed
addprocs(2)
# using Revise
using GraphPlot
using Colors
using DistributedFactorGraphs
using RoME
@everywhere using RoME
using RoMEPlotting
Gadfly.set_default_plot_size(35cm,25cm)

function getCoordCartesianFromKDERange(dfg::AbstractDFG, varsym::Symbol; digits::Int=-2, extend::Float64=0.5)
    ax = round.(getKDERange(getKDE(dfg, varsym), extend=extend), digits=digits)
    if size(ax, 1) == 1
      return Coord.Cartesian(xmin=ax[1,1], xmax=ax[1,2])
    elseif size(ax, 1) == 2
      return Coord.Cartesian(xmin=ax[1,1], xmax=ax[1,2],ymin=ax[2,1], ymax=ax[2,2])
    else
      error("only do 1D or 2D")
    end
end
# import KernelDensityEstimate: getCoordCartesianFromKDERange

function addMaptoPlot!(pl)
    # plot the design coordinates
    map_coords = Array{Array{Float64,1},1}[[[0.0, 0.0], [0.0, 330.0], [40.0, 330.0], [40.0, 310.0], [20.0, 310.0], [20.0, 150.0], [130.0, 150.0], [130.0, 140.0], [20.0, 140.0], [20.0, 20.0], [250.0, 20.0], [250.0, 60.0], [260.0, 60.0], [260.0, 20.0], [290.0, 20.0], [290.0, 60.0], [300.0, 60.0], [300.0, 20.0], [430.0, 20.0], [430.0, 110.0], [300.0, 110.0], [300.0, 100.0], [290.0, 100.0], [290.0, 190.0], [280.0, 190.0], [280.0, 200.0], [320.0, 200.0], [320.0, 190.0], [300.0, 190.0], [300.0, 120.0], [430.0, 120.0], [430.0, 190.0], [410.0, 190.0], [410.0, 200.0], [430.0, 200.0], [430.0, 310.0], [390.0, 310.0], [390.0, 330.0], [450.0, 330.0], [450.0, 0.0], [0.0, 0.0]], [[90.0, 310.0], [90.0, 330.0], [170.0, 330.0], [170.0, 310.0], [90.0, 310.0]], [[230.0, 310.0], [340.0, 310.0], [340.0, 330.0], [210.0, 330.0], [210.0, 310.0], [220.0, 310.0], [220.0, 200.0], [210.0, 200.0], [210.0, 190.0], [240.0, 190.0], [240.0, 200.0], [230.0, 200.0], [230.0, 310.0]]]

    for m_coord in map_coords
        Xpp = map(p->p[1] + 150, m_coord)
        Ypp = map(p->-p[2] +330 + 650, m_coord)

        push!(pl.layers, Gadfly.layer(x=Xpp,y=Ypp,Geom.path(), Theme(line_width=1pt))[1])
    end

    # plot the real map variation
    built_coords = Array{Array{Float64,1},1}[[[230.0, 310.0], [340.0, 310.0], [340.0, 330.0], [210.0, 330.0], [210.0, 310.0], [220.0, 310.0], [220.0, 200.0], [210.0, 200.0], [210.0, 190.0], [240.0, 190.0], [240.0, 200.0], [230.0, 200.0], [230.0, 310.0]]]

    for b_coord in built_coords
        Xpp = map(p->p[1] + 150+30, b_coord)
        Ypp = map(p->-p[2] +330 + 650+20, b_coord)

        push!(pl.layers, Gadfly.layer(x=Xpp,y=Ypp,Geom.path(), Theme(line_width=1pt, default_color=colorant"red"))[1])
    end

    return pl
end

## First add the designed landmark positions
function addLandmarkTrueWithPrior!(dfg::AbstractDFG,
                                   key::Symbol,
                                   vals::Vector{<: Real},
                                   lgt_cov::Array{Float64,2}=0.01*Matrix(LinearAlgebra.I, 2,2)  )
  #
  if !exists(dfg, key)
    addVariable!(dfg, key, Point2)
    pp = PriorPoint2(MvNormal(vals, lgt_cov))
    addFactor!(dfg, [key], pp)
  end
  nothing
end

macro doline(skip, operation)
  :($skip ? $operation : nothing)
end

ground = Dict{Symbol, Vector{Float64}}(
  :x0=>[260.0, 720.0, 0.0],
  :x1=>[260.0, 720.0, -pi/2],
  :x2=>[260.0, 720.0, -pi],
  :x3=>[260.0, 720.0, pi/2],
  :x4=>[260.0, 800.0, 0.0],
  :x5=>[310.0, 800.0, 0.0],
  :x6=>[310.0, 800.0, -pi/2],
  :x7=>[310.0, 800.0, -pi],
  :x8=>[310.0, 800.0, pi/2],
)

## Add all variables as design truth (for use as multihypos later)
# Ground truth for landmarks (designed)
landmarks_real = Dict(
    :l1=>[360+30,780+20],
    :l2=>[370+30,780+20],
    :l3=>[370+30,670+20],
    :l4=>[360+30,670+20],
    :l5=>[360+30,790+20],
    :l8=>[190,670],
    :l6=>[170,670],
    :l7=>[170,830],
    :l9=>[280,830],
    :l10=>[440,790],
    :l11=>[430,790],
    :l12=>[400,920])

# map with Dx=-10 on l1-l4 (ground truth from measured, varies from design)
landmarks_design = Dict(
    :l1=>[360,780],
    :l2=>[370,780],
    :l3=>[370,670],
    :l4=>[360,670],
    :l5=>[360,790],
    :l8=>[190,670],
    :l6=>[170,670],
    :l7=>[170,830],
    :l9=>[280,830],
    :l10=>[440,790],
    :l11=>[430,790],
    :l12=>[400,920])

## parameters
autoinit = true
pp_σx = 0.1
pp_σy = 0.1
pp_σψ = 0.01
br_σb = 0.03
# br_σb = 0.001
br_σρ = 0.3
pMeas, pDesg = 0.7, 0.3

SKL = true

## Start the factor graph construction
fg = LightDFG{SolverParams}(params=SolverParams())

## develop the robotic factor graph elements
# x0
x0_prior = [ground[:x0],[0.1,0.1,0.01]] #fixed
# x0_prior = [ground[:x0],[100.,100,1]] #fixed
addVariable!(fg, :x0, Pose2)
addFactor!(fg, [:x0],  PriorPose2( MvNormal(x0_prior[1], Matrix(Diagonal(x0_prior[2].^2))) ), autoinit=autoinit )

    br_calc = calcPosePointBearingRange(ground[:x0], landmarks_real[:l1])
    # x0l1_br = [[deg2rad(27.65), 0.02], [118.53, 0.2]]#(tagψ, σtagψ), Normal(tagrange, σtagrange)
    x0l1_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    @doline SKL addVariable!(fg, :l1, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l1_0, landmarks_design[:l1])
    @doline SKL addFactor!(fg, [:x0, :l1, :l1_0], Pose2Point2BearingRange(Normal(x0l1_br[1]...), Normal(x0l1_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

    br_calc = calcPosePointBearingRange(ground[:x0], landmarks_real[:l2])
    x0l2_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x0l2_br = [[deg2rad(25.56), 0.02], [127.48, 0.2]]
    @doline SKL addVariable!(fg, :l2, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l2_0, landmarks_design[:l2])
    @doline SKL addFactor!(fg, [:x0, :l2, :l2_0], Pose2Point2BearingRange(Normal(x0l2_br[1]...), Normal(x0l2_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

    br_calc = calcPosePointBearingRange(ground[:x0], landmarks_real[:l3])
    x0l3_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x0l3_br = [[deg2rad(-25.56), 0.02], [127.48, 0.2]]
    @doline SKL addVariable!(fg, :l3, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l3_0, landmarks_design[:l3])
    @doline SKL addFactor!(fg, [:x0, :l3, :l3_0], Pose2Point2BearingRange(Normal(x0l3_br[1]...), Normal(x0l3_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

    br_calc = calcPosePointBearingRange(ground[:x0], landmarks_real[:l4])
    x0l4_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x0l4_br = [[deg2rad(-27.65), 0.02], [118.53, 0.2]]
    @doline SKL addVariable!(fg, :l4, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l4_0, landmarks_design[:l4])
    @doline SKL addFactor!(fg, [:x0, :l4, :l4_0], Pose2Point2BearingRange(Normal(x0l4_br[1]...), Normal(x0l4_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

## test for multimodality
# dfgplot(fg)
# drawGraph(fg)
# # ensureAllInitialized!(fg)
getSolverParams(fg).drawtree=true
getSolverParams(fg).showtree=true
# getSolverParams(fg).dbg=true
# getSolverParams(fg).async=false
tree = solveTree!(fg);
pl = drawPosesLandms(fg, drawhist=true, contour=false)
addMaptoPlot!(pl)
pl.coord = Coord.Cartesian(xmin=100, xmax=450, ymin=600, ymax=850)
pl
# pl |> PDF(ENV["HOME"]*"/Downloads/test.pdf")

# plot l1
pl1 = drawLandms(fg, to=1, drawhist=true, contour=false)
addMaptoPlot!(pl1)
pl1.coord = getCoordCartesianFromKDERange(fg, :l1, digits=0)
pl1

# plot l2
pl2 = drawLandms(fg, from=2,to=2, drawhist=true, contour=false)
addMaptoPlot!(pl2)
pl2.coord = getCoordCartesianFromKDERange(fg, :l2, digits=0)
pl2

# plot l3
pl3 = drawLandms(fg, from=3,to=3, drawhist=true, contour=false)
addMaptoPlot!(pl3)
pl3.coord = getCoordCartesianFromKDERange(fg, :l3, digits=0)
pl3

# plot l4
pl4 = drawLandms(fg, from=4,to=4, drawhist=true, contour=false)
addMaptoPlot!(pl4)
pl4.coord = getCoordCartesianFromKDERange(fg, :l4, digits=0)
pl4

pl = vstack(hstack(pl1,pl2), hstack(pl3,pl4))
# pl |> PDF(ENV["HOME"]*"/Downloads/stack.pdf")

# x1
x0x1_pp = [[0,0,deg2rad(-90)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x1, Pose2)
addFactor!(fg, [:x0, :x1], Pose2Pose2(MvNormal(x0x1_pp[1], Matrix(Diagonal(x0x1_pp[2].^2)))), autoinit=autoinit)
# ensureAllInitialized!(fg)
# x2
x1x2_pp = [[0,0,deg2rad(-90)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x2, Pose2)
addFactor!(fg, [:x1, :x2], Pose2Pose2(MvNormal(x1x2_pp[1], Matrix(Diagonal(x1x2_pp[2].^2)))), autoinit=autoinit)

    br_calc = calcPosePointBearingRange(ground[:x2], landmarks_real[:l5])
    x2l5_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x2l5_br = [[deg2rad(-143.75+180), 0.1], [92.65, 1]]
    @doline SKL addVariable!(fg, :l5, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l5_0, landmarks_design[:l5])
    @doline SKL addFactor!(fg, [:x2, :l5, :l5_0], Pose2Point2BearingRange(Normal(x2l5_br[1]...), Normal(x2l5_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

    br_calc = calcPosePointBearingRange(ground[:x2], landmarks_real[:l6])
    x2l6_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x2l6_br = [[deg2rad(-149.93+180), 0.1], [109.77, 1]]
    @doline SKL addVariable!(fg, :l6, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l6_0, landmarks_design[:l6])
    @doline SKL addFactor!(fg, [:x2, :l6, :l6_0], Pose2Point2BearingRange(Normal(x2l6_br[1]...), Normal(x2l6_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

# ensureAllInitialized!(fg)
# x3:
x2x3_pp = [[0,0,deg2rad(-90)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x3, Pose2)
addFactor!(fg, [:x2, :x3], Pose2Pose2(MvNormal(x2x3_pp[1], Matrix(Diagonal(x2x3_pp[2].^2)))), autoinit=autoinit)

    br_calc = calcPosePointBearingRange(ground[:x3], landmarks_real[:l7])
    x3l7_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x3l7_br = [[deg2rad(132.4-90), 0.1], [141.32, 1]]
    @doline SKL addVariable!(fg, :l7, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l7_0, landmarks_design[:l7])
    @doline SKL addFactor!(fg, [:x3, :l7, :l7_0], Pose2Point2BearingRange(Normal(x3l7_br[1]...), Normal(x3l7_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

# ensureAllInitialized!(fg)
# x4 drive + 80
x3x4_pp = [[80,0,deg2rad(-90)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x4, Pose2)
addFactor!(fg, [:x3, :x4], Pose2Pose2(MvNormal(x3x4_pp[1], Matrix(Diagonal(x3x4_pp[2].^2)))), autoinit=autoinit)

# ensureAllInitialized!(fg)
# x5 turn -90 -> drive + 50
x4x5_pp = [[50,0,deg2rad(0)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x5, Pose2)
addFactor!(fg, [:x4, :x5], Pose2Pose2(MvNormal(x4x5_pp[1], Matrix(Diagonal(x4x5_pp[2].^2)))), autoinit=autoinit)

    br_calc = calcPosePointBearingRange(ground[:x5], landmarks_real[:l8])
    x5l8_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x5l8_br = [[deg2rad(-15.26), 0.1], [57.01, 1]]
    @doline SKL addVariable!(fg, :l8, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l8_0, landmarks_design[:l8])
    @doline SKL addFactor!(fg, [:x5, :l8, :l8_0], Pose2Point2BearingRange(Normal(x5l8_br[1]...), Normal(x5l8_br[2]...)), autoinit=autoinit, multihypo=[1.0;0.5;0.5])

    br_calc = calcPosePointBearingRange(ground[:x5], landmarks_real[:l10])
    x5l10_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x5l10_br = [[deg2rad(-6.84), 0.1], [105.84, 1]]
    @doline SKL addVariable!(fg, :l10, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l10_0, landmarks_design[:l10])
    @doline SKL addFactor!(fg, [:x5, :l10, :l10_0], Pose2Point2BearingRange(Normal(x5l10_br[1]...), Normal(x5l10_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

    br_calc = calcPosePointBearingRange(ground[:x5], landmarks_real[:l11])
    x5l11_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x5l11_br = [[deg2rad(-7.43), 0.1], [113.59, 1]]
    @doline SKL addVariable!(fg, :l11, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l11_0, landmarks_design[:l11])
    @doline SKL addFactor!(fg, [:x5, :l11, :l11_0], Pose2Point2BearingRange(Normal(x5l11_br[1]...), Normal(x5l11_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

#
drawGraph(fg)
# ensureAllInitialized!(fg)
# drawPosesLandms(fg)
# getSolverParams(fg).drawtree=true
# getSolverParams(fg).showtree=true
tree = solveTree!(fg);

pl = drawPosesLandms(fg, drawhist=true, contour=false)
addMaptoPlot!(pl)
pl.coord = Coord.Cartesian(xmin=100, xmax=500, ymin=600, ymax=900)
pl
# pl |> PDF(ENV["HOME"]*"/Downloads/explore1.pdf")

pl5 = drawLandms(fg, from=5,to=5, drawhist=true, contour=false)
addMaptoPlot!(pl5)
pl5.coord = getCoordCartesianFromKDERange(fg, :l5, digits=0)
pl5

# reportFactors(fg, Pose2Pose2)
#

# drawLandms(fg, to=1, drawhist=true, contour=false)

# plot the bimodal wall
plotVariable2D(fg, :l1, refs=[landmarks_design, landmarks_real])
plotVariable2D(fg, :l2, refs=[landmarks_design, landmarks_real])
plotVariable2D(fg, :l3, refs=[landmarks_design, landmarks_real])
plotVariable2D(fg, :l4, refs=[landmarks_design, landmarks_real])
plotVariable2D(fg, :l5, refs=[landmarks_design, landmarks_real])
# plot unimodals
plotVariable2D(fg, :l6, refs=[landmarks_design, landmarks_real])
plotVariable2D(fg, :l7, refs=[landmarks_design, landmarks_real])
plotVariable2D(fg, :l8, refs=[landmarks_design, landmarks_real])

reportFactors(fg, Pose2Pose2)

# plotLocalProduct(fg, :x4, levels=5)

# ensureAllInitialized!(fg)
# x6
x5x6_pp = [[0,0,deg2rad(-90)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x6, Pose2)
addFactor!(fg, [:x5, :x6], Pose2Pose2(MvNormal(x5x6_pp[1], Matrix(Diagonal(x5x6_pp[2].^2)))), autoinit=autoinit)

    br_calc = calcPosePointBearingRange(ground[:x6], landmarks_real[:l3])
    x6l3_br  = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x6l3_br = [[deg2rad(-64.29+90), 0.1], [145.07, 1]]
    # addVariable!(fg, :l3, Point2)
    @doline SKL addFactor!(fg, [:x6, :l3, :l3_0], Pose2Point2BearingRange(Normal(x6l3_br[1]...), Normal(x6l3_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

    br_calc = calcPosePointBearingRange(ground[:x6], landmarks_real[:l4])
    x6l4_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x6l4_br = [[deg2rad(-67.83+90), 0.1], [141.01, 1]]
    # addVariable!(fg, :l4, Point2)
    @doline SKL addFactor!(fg, [:x6, :l4, :l4_0], Pose2Point2BearingRange(Normal(x6l4_br[1]...), Normal(x6l4_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

## debug check solve here
# drawGraph(fg)

# ensureAllInitialized!(fg)
# drawPosesLandms(fg)
# getSolverParams(fg).drawtree=true
# getSolverParams(fg).showtree=true
tree = solveTree!(fg);
# tree = solveTree!(fg, tree);
pl = drawPosesLandms(fg, drawhist=true, contour=false)
addMaptoPlot!(pl)
# pl |> PDF(ENV["HOME"]*"/Downloads/explore2.pdf")

# x7
x6x7_pp = [[0,0,deg2rad(-90)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x7, Pose2)
addFactor!(fg, [:x6, :x7], Pose2Pose2(MvNormal(x6x7_pp[1], Matrix(Diagonal(x6x7_pp[2].^2)))), autoinit=autoinit)

    br_calc = calcPosePointBearingRange(ground[:x7], landmarks_real[:l7])
    x7l7_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x7l7_br = [[deg2rad(137-180), 0.1], [147.89, 1]]
    pp_x7l7_br = Pose2Point2BearingRange(Normal(x7l7_br[1]...), Normal(x7l7_br[2]...))
    @doline SKL addFactor!(fg, [:x7; :l7; :l7_0], pp_x7l7_br, autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

    br_calc = calcPosePointBearingRange(ground[:x7], landmarks_real[:l9])
    x7l9_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x7l9_br = [[deg2rad(144.46-180), 0.1], [42.04, 1]]
    @doline SKL addVariable!(fg, :l9, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l9_0, landmarks_design[:l9])
    @doline SKL addFactor!(fg, [:x7, :l9, :l9_0], Pose2Point2BearingRange(Normal(x7l9_br[1]...), Normal(x7l9_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

# drawGraph(fg)
# pts = approxConv(fg, :x7l7l7bf1, :l7)
# pts = approxConv(fg, :x7l7l7bf1, :l7b)

# ensureAllInitialized!(fg)
# x8
x7x8_pp = [[0,0,deg2rad(-90)], [pp_σx, pp_σy, pp_σψ]]
addVariable!(fg, :x8, Pose2)
addFactor!(fg, [:x7, :x8], Pose2Pose2(MvNormal(x7x8_pp[1], Matrix(Diagonal(x7x8_pp[2].^2)))), autoinit=autoinit)

    br_calc = calcPosePointBearingRange(ground[:x8], landmarks_real[:l12])
    x8l12_br = [[br_calc[1], br_σb], [br_calc[2], br_σρ]]
    # x8l12_br = [[deg2rad(53.53-90), 0.1], [133.32, 1]]
    @doline SKL addVariable!(fg, :l12, Point2)
    @doline SKL addLandmarkTrueWithPrior!(fg, :l12_0, landmarks_design[:l12])
    @doline SKL addFactor!(fg, [:x8, :l12, :l12_0], Pose2Point2BearingRange(Normal(x8l12_br[1]...), Normal(x8l12_br[2]...)), autoinit=autoinit, multihypo=[1.0;pMeas;pDesg])

# ensureAllInitialized!(fg)

# getSolverParams(fg).async = false
# tree = solveTree!(fg, tree);
tree = solveTree!(fg);

pl = drawPosesLandms(fg, drawhist=true, contour=false)
addMaptoPlot!(pl)
# pl |> PDF(ENV["HOME"]*"/Downloads/mapped.pdf")

pl = drawLandms(fg, drawhist=true, to=5, contour=false)
addMaptoPlot!(pl)
pl.coord = Coord.Cartesian(xmin=100, xmax=450, ymin=600, ymax=850)
pl
# pl |> PDF(ENV["HOME"]*"/Downloads/badwall.pdf")

drawGraph(fg)

# reportFactors(fg, Pose2Pose2);

## =============================================================================
## Deeper look at various other aspects.

# pl = plotVariable2D(fg, :l10)
#
# pl = drawLandms(fg, from=10, to=11, contour=false, drawhist=true)
# addMaptoPlot!(pl)
#
# pl.coord = getCoordCartesianFromKDERange(fg, :l10, digits=0, extend=10.0)
# pl

# ## quick check
#
# pl = drawLandms(fg, from=4, to=4, drawhist=true, contour=false)
# plotKDE(fg, :l4, levels=20)
# plotKDE(fg, :l4_0, levels=20)
# plotKDE(fg, [:l4;:l4_0])

##

#
# ## 3D VIS ==================
#
# using Amphitheatre
#
# fgmax = deepcopy(fg)
# fgmax.robotId = "Max"
#
# fgvis = BasicFactorGraphPose(fg, meanmax=:mean)
# fgvismax = BasicFactorGraphPose(fgmax, meanmax=:max)
#
# visdatasets = AbstractAmphitheatre[fgvis]
# push!(visdatasets, fgvismax)
# vis, vistask = visualize(visdatasets, start_browser=true)
#
# stopAmphiVis!()
# vistask = Amphitheatre.restartVisualize(vis, visdatasets)
#
# ## END VIS 3D

# plotVariable2D(fg, :l9, refs=[landmarks_design, landmarks_real])
# plotVariable2D(fg, :l10, refs=[landmarks_design, landmarks_real])
# plotVariable2D(fg, :l11, refs=[landmarks_design, landmarks_real])
# plotVariable2D(fg, :l12, refs=[landmarks_design, landmarks_real])
#
# plotLocalProduct(fg, :l3, levels=10)
# plotLocalProduct(fg, :l4, levels=10)
#
# # plotLocalProduct(fg, :x6, levels=10)
#
#
#
# syms = sortDFG(ls(fg, r"x"))
# ppes = Base.map((x)->calcVariablePPE(fg,x), syms)
# trivial_poses = Dict(syms .=> ppes)
#
# # syms = sortDFG(ls(fg, r"l\d*_"))
# # ppes = Base.map((x)->calcVariablePPE(fg,x), syms)
# # design_landms = Dict(syms .=> ppes)
# #
# # syms = sortDFG(setdiff(ls(fg, r"l\d*"),ls(fg, r"l\d*_")) )
# # ppes = Base.map((x)->calcVariablePPE(fg,x), syms)
# # design_mapped = Dict(syms .=> ppes)
# #
#
#
#
#
# ##
# p = DFGPlotProps(  (var=colorant"seagreen", fac=colorant"cyan3"),
#                             (var=1.0, fac=0.3),
#                             (var=:box, fac=:elipse),
#                             spring_layout,#spectral_layout
#                             true)
# dfgplot(fg,p)
#
#
#
# ##
#
# # using RoMEPlotting
#
# plotPose(fg, [:x0])
# plotPose(fg, [:x1])
# plotPose(fg, [:x2])
# plotPose(fg, [:x3])
# plotPose(fg, [:x0,:x1,:x2,:x3], levels=1)
#
#
# ##
#
# plotLocalProduct(fg, :x0)
# plotLocalProduct(fg, :x1)
# plotLocalProduct(fg, :x2)
#
# plotLocalProduct(fg, :l7, levels=3)
#
#
# plotLocalProduct(fg, :l1_0)
# plotLocalProduct(fg, :l2_0)
# plotLocalProduct(fg, :l3_0)
# plotLocalProduct(fg, :l4_0)
#
# plotLocalProduct(fg, :l4)
#
# ls(fg, :l7)
#
#
# drawPosesLandms(fg)

# ensureAllInitialized!(fg)
# cu, pr = predictVariableByFactor(fg, :l7, pp_x7l7_br, [:x7;:l7])
# if minkld(cu, pr) < 3
    # addFactor!(fg, [:x7, :l7], pp_x7l7_br, autoinit=autoinit)
# else
#     addVariable!(fg, :l7_1, Point2)
# end

##


"""
    SerialExecution

Functor object for using serial execution in a lshaped algorithm. Create by supplying a [`Serial`](@ref) object through `execution` in the `LShapedSolver` factory function and then pass to a `StochasticPrograms.jl` model.

"""
struct SerialExecution{H <: AbstractFeasibilityHandler,
                       T <: AbstractFloat,
                       A <: AbstractVector} <: AbstractLShapedExecution
    subproblems::Vector{SubProblem{H,T}}
    decisions::Decisions
    subobjectives::A
    model_objectives::A

    function SerialExecution(structure::VerticalStructure{2, 1, <:Tuple{ScenarioProblems}},
                             ::Type{F}, ::Type{T}, ::Type{A}) where {F <: AbstractFeasibility,
                                                                     T <: AbstractFloat,
                                                                     A <: AbstractVector}
        H = HandlerType(F)
        return new{H,T,A}(Vector{SubProblem{H,T}}(), structure.decisions[2], A(), A())
    end
end

function num_thetas(lshaped::AbstractLShaped, ::SerialExecution)
    return num_thetas(num_subproblems(lshaped), lshaped.aggregation)
end

function initialize_subproblems!(execution::SerialExecution{H,T},
                                 scenarioproblems::ScenarioProblems) where {H <: AbstractFeasibilityHandler,
                                                                            T <: AbstractFloat}
    for i in 1:num_subproblems(scenarioproblems)
        push!(execution.subproblems, SubProblem(
            subproblem(scenarioproblems, i),
            i,
            T(probability(scenario(scenarioproblems, i))),
            execution.decisions.knowns,
            H))
    end
    return nothing
end

function finish_initilization!(lshaped::AbstractLShaped, execution::SerialExecution)
    append!(execution.subobjectives, fill(1e10, num_thetas(lshaped)))
    append!(execution.model_objectives, fill(-1e10, num_thetas(lshaped)))
    return nothing
end

function restore_subproblems!(lshaped::AbstractLShaped, execution::SerialExecution)
    for subproblem in execution.subproblems
        restore_subproblem!(subproblem)
    end
    return nothing
end

function resolve_subproblems!(lshaped::AbstractLShaped, execution::SerialExecution{H,T}) where {H <: AbstractFeasibilityHandler, T <: AbstractFloat}
    # Update subproblems
    update_known_decisions!(execution.decisions, lshaped.x)
    change = KnownValuesChange()
    # Assume no cuts are added
    added = false
    # Update and solve subproblems
    for subproblem in execution.subproblems
        update_subproblem!(subproblem, change)
        cut::SparseHyperPlane{T} = subproblem(lshaped.x)
        added |= aggregate_cut!(lshaped, lshaped.aggregation, cut)
    end
    added |= flush!(lshaped, lshaped.aggregation)
    # Return current objective value and cut_added flag
    return current_objective_value(lshaped), added
end

# API
# ------------------------------------------------------------
function (execution::Serial)(structure::VerticalStructure{2, 1, <:Tuple{ScenarioProblems}},
                             ::Type{F}, ::Type{T}, ::Type{A}) where {F <: AbstractFeasibility,
                                                                     T <: AbstractFloat,
                                                                     A <: AbstractVector}
    return SerialExecution(structure, F, T, A)
end

function str(::Serial)
    return ""
end


const NOpsType = Int#Union{Int,Vector{Int}}

struct UpperBoundedInteger{N, T<: Base.BitInteger} <: Integer
  i::T
  @inline UpperBoundedInteger{N}(i::T) where {N,T<:Base.BitInteger} = new{N,T}(i)
end

@inline UpperBoundedInteger(i::T, ::StaticInt{N}) where {N,T<:Base.BitInteger} = UpperBoundedInteger{N}(i)
@inline UpperBoundedInteger(::StaticInt{M}, ::StaticInt{N}) where {N,M} = StaticInt{M}()
@inline UpperBoundedInteger{N}(::StaticInt{M}) where {N,M} = StaticInt{M}()
@inline Base.:(%)(a::UpperBoundedInteger, ::Type{T}) where {T<:Base.BitInteger} = a.i % T
Base.promote_rule(::Type{T}, ::Type{UpperBoundedInteger{N,S}}) where {N,T<:Base.BitInteger,S} = promote_rule(T,S)
Base.promote_rule(::Type{UpperBoundedInteger{N,S}}, ::Type{T}) where {N,T<:Base.BitInteger,S} = promote_rule(S,T)
Base.convert(::Type{T}, i::UpperBoundedInteger) where {T<:Number} = convert(T, i.i)
Base.convert(::Type{UpperBoundedInteger{N,T}}, i::UpperBoundedInteger{N,T}) where {N,T<:Base.BitInteger} = i
Base.convert(::Type{Any}, i::UpperBoundedInteger) = i
upper_bound(_) = typemax(Int)
upper_bound(::Type{CloseOpen{T,UpperBoundedInteger{N,S}}}) where {T,N,S} = N - 1

@inline Base.last(r::CloseOpen{<:Integer,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i) - One()
@inline ArrayInterface.static_last(r::CloseOpen{<:Integer,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i) - One()
@inline Base.length(r::CloseOpen{<:Integer,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i) - getfield(r,:start)
@inline Base.length(r::CloseOpen{Zero,<:UpperBoundedInteger}) = getfield(getfield(r,:upper),:i)

function Loop(ls::LoopSet, ex::Expr, sym::Symbol, f, s, l, ub::Int)
    if (f !== nothing) && (s !== nothing) && (l !== nothing)
        return static_loop(sym, f, s, l)
    end
    ssym = String(sym)
    rangesym = gensym(ssym * "_loop");
    lensym = gensym(ssym * "_looplen")
    pushpreamble!(ls, Expr(:(=), rangesym, ex))
    pushpreamble!(ls, Expr(:(=), lensym, Expr(:call, GlobalRef(ArrayInterface,:static_length), rangesym)))
    F = if f === nothing
        start = gensym(ssym*"_loopstart")
        pushpreamble!(ls, Expr(:(=), start, Expr(:call, %, Expr(:call, lv(:first), rangesym), Int)))
        MaybeKnown(start, 1)
    else
        MaybeKnown(f)
    end
    S = if s === nothing
        step = gensym(ssym*"_loopstep")
        pushpreamble!(ls, Expr(:(=), step, Expr(:call, %, Expr(:call, lv(:step), rangesym), Int)))
        MaybeKnown(step, 1)
    else
        MaybeKnown(s)
    end
    L = if l === nothing
        stop = gensym(ssym*"_loopstop")
        pushpreamble!(ls, Expr(:(=), stop, Expr(:call, %, Expr(:call, lv(:last), rangesym), Int)))
        MaybeKnown(stop, min(ub, 1024))
    else
        MaybeKnown(l)
    end
    loopiteratesatleastonce!(ls, Loop(sym, F, L, S, rangesym, lensym))
end
function Loop(ls::LoopSet, ex::Expr, sym::Symbol, ::Type{R}) where {R<:AbstractRange}
  f = ArrayInterface.known_first(R)
  s = ArrayInterface.known_step(R)
  l = ArrayInterface.known_last(R)
  ub = upper_bound(R)
  Loop(ls, ex, sym, f, s, l, ub)
end

function static_loop(sym::Symbol, L::Int, S::Int, U::Int)
    Loop(sym, MaybeKnown(L,0), MaybeKnown(U,0), MaybeKnown(S,0), Symbol(""), Symbol(""))
end
function Loop(::LoopSet, ::Expr, sym::Symbol, ::Type{OptionallyStaticUnitRange{Static{L}, Static{U}}}) where {L,U}
    static_loop(sym, L, 1, U)
end
function Loop(::LoopSet, ::Expr, sym::Symbol, ::Type{ArrayInterface.OptionallyStaticStepRange{StaticInt{L}, StaticInt{S}, StaticInt{U}}}) where {L,S,U}
    static_loop(sym, L, S, U)
end
function Loop(::LoopSet, ::Expr, sym::Symbol, ::Type{CloseOpen{Static{L}, Static{U}}}) where {L,U}
    static_loop(sym, L, 1, U - 1)
end

extract_loop(l) = Expr(:call, GlobalRef(Core,:getfield), Symbol("#loop#bounds#"), l, false)

function add_loops!(ls::LoopSet, LPSYM, LB)
    n = max(length(LPSYM), length(LB))
    for i = 1:n
        sym, l = LPSYM[i], LB[i]
        if l<:CartesianIndices
            add_loops!(ls, i, sym, l)
        else
            add_loop!(ls, Loop(ls, extract_loop(i), sym, l)::Loop)
            push!(ls.loopsymbol_offsets, ls.loopsymbol_offsets[end]+1)
        end
    end
end
function add_loops!(ls::LoopSet, i::Int, sym::Symbol, @nospecialize(l::Type{<:CartesianIndices}))
    N, T = l.parameters
    ssym = String(sym)
    for k = N:-1:1
        axisexpr = :(getfield(getfield(getfield(var"#loop#bounds#", $i, false), :indices), $k, false))
        add_loop!(ls, Loop(ls, axisexpr, Symbol(ssym*'#'*string(k)*'#'), T.parameters[k])::Loop)
    end
    push!(ls.loopsymbol_offsets, ls.loopsymbol_offsets[end]+N)
end

function ArrayReferenceMeta(
    ls::LoopSet, @nospecialize(ar::ArrayRefStruct), arraysymbolinds::Vector{Symbol},
    opsymbols::Vector{Symbol}, nopsv::Vector{NOpsType}, expandedv::Vector{Bool}
)
    # unpack the `ArrayRefStruct`
    # we don't want to specialize on it, as it is typed on symbols.
    index_types = (ar.index_types)::UInt128
    indices = (ar.indices)::UInt128
    offsets = (ar.offsets)::UInt128
    strides = (ar.strides)::UInt128
    arrayar, ptrar = array_and_ptr(ar)::Tuple{Symbol,Symbol}
    # ptrar = ptr(ar)::Symbol
    # TODO, check if this matters at all. The compiler still knows it is an `::ArrayRefStruct`, just not `arrayar` or `ptrar`?
    ArrayReferenceMeta(
        ls, index_types, indices, offsets, strides,
        arrayar, ptrar, arraysymbolinds, opsymbols, nopsv, expandedv
    )
end
function ArrayReferenceMeta(
    ls::LoopSet, index_types::UInt128, indices::UInt128, offsets::UInt128, strides::UInt128,
    arrayar::Symbol, ptrar::Symbol, arraysymbolinds::Vector{Symbol},
    opsymbols::Vector{Symbol}, nopsv::Vector{NOpsType}, expandedv::Vector{Bool}
)
    ni = filled_8byte_chunks(index_types)
    index_vec = Symbol[]
    offset_vec = Int8[]
    stride_vec = Int8[]
    loopedindex = Bool[]
    while index_types != zero(UInt128)
        ind = indices % UInt8
        offsetᵢ = offsets % Int8
        strideᵢ = strides % Int8
        if index_types == LoopIndex
            if ind == zero(Int8) # CONSTANTZEROINDEX
                pushfirst!(index_vec, CONSTANTZEROINDEX)
                pushfirst!(offset_vec, offsetᵢ)
                pushfirst!(stride_vec, strideᵢ)
                pushfirst!(loopedindex, true)
            else
                for inda in ls.loopsymbol_offsets[ind]+1:ls.loopsymbol_offsets[ind+1]
                    pushfirst!(index_vec, ls.loopsymbols[inda])
                    pushfirst!(offset_vec, offsetᵢ)
                    pushfirst!(stride_vec, strideᵢ)
                    pushfirst!(loopedindex, true)
                end
            end
        else#if index_types == ComputedIndex
            @assert index_types == ComputedIndex
            opsym = opsymbols[ind]
            if expandedv[ind]
                nops = nopsv[ind]
                for j ∈ 0:nops-1
                    pushfirst!(index_vec, expandedopname(opsym, j))
                    pushfirst!(offset_vec, offsetᵢ)
                    pushfirst!(stride_vec, strideᵢ)
                    pushfirst!(loopedindex, false)
                end
            else
                pushfirst!(index_vec, opsym)
                pushfirst!(offset_vec, offsetᵢ)
                pushfirst!(stride_vec, strideᵢ)
                pushfirst!(loopedindex, false)
            end
        # else
            # @assert index_types == SymbolicIndex
            # pushfirst!(index_vec, arraysymbolinds[ind])
            # pushfirst!(offset_vec, offset)
            # pushfirst!(loopedindex, false)
        end
        index_types >>>= 8
        indices >>>= 8
        offsets >>>= 8
        strides >>>= 8
        ni -= 1
    end
    ArrayReferenceMeta(
        ArrayReference(arrayar, index_vec, offset_vec, stride_vec),
        loopedindex, ptrar
    )
end

extract_varg(i) = :(getfield(var"#vargs#", $i, false))
# _extract(::Type{Static{N}}) where {N} = N
extract_gsp!(sptrs::Expr, name::Symbol) = (push!(sptrs.args, name); nothing)
tupleranks(R::NTuple{8,Int}) = ntuple(n -> sum(R[n] .≥ R), Val{8}())
function rank_to_sortperm((R,N)::Tuple{NTuple{8,Int},Int})::Vector{Int}
  r = tupleranks(R)
  sp = Vector{Int}(undef, N)
  @inbounds for n = 1:N
    sp[r[n]] = n
  end
  sp
end
function loop_indexes_bit!(ls::LoopSet, ar::ArrayReferenceMeta)
    li = ar.loopedindex;
    ind = first(getindices(ar))
    ind === DISCONTIGUOUS && return
    first(li) || throw(LoopError("The contiguous index of a `BitArray` shouldn't be a complex function.", )ind)
    ls.loopindexesbit[getloopid(ls,ind)] = true
    nothing
end
function add_mref!(
    sptrs::Expr, ls::LoopSet, ar::ArrayReferenceMeta, @nospecialize(_::Type{Ptr{T}}),
    C::Int, B::Int, sp::Vector{Int}, name::Symbol
) where {T}
  @assert B ≤ 0 "Batched arrays not supported yet."
  _add_mref!(sptrs, ls, ar, typetosym(T), C, B, sp, name)
  sizeof(T)
end
typetosym(::Type{T}) where {T<:NativeTypes} = (VectorizationBase.JULIA_TYPES[T])::Symbol
typetosym(T) = T
function _add_mref!(sptrs::Expr, ls::LoopSet, ar::ArrayReferenceMeta, T_sym, C::Int, B::Int, sp::Vector{Int}, name::Symbol)
    # maybe no change needed? -- optimize common case
  li = ar.loopedindex;
  if sp == eachindex(sp) || isone(length(li))
    # don't set `bit` to true if our vector width is ≥ 8
    ((T_sym === :Bit) && (ls.vector_width < 8)) && loop_indexes_bit!(ls, ar)
    return extract_gsp!(sptrs, name)
  end
  permute_mref!(ar, C, sp)
  # must now sort array's inds, and stack pointer's
  tmpsp = gensym(name)
  extract_gsp!(sptrs, tmpsp)
  strd_tup = Expr(:tuple)
  offsets_tup = Expr(:tuple)
  gf = GlobalRef(Core,:getfield)
  offsets = gensym(:offsets); strides = gensym(:strides)
  pushpreamble!(ls, Expr(:(=), offsets, Expr(:call, gf, tmpsp, QuoteNode(:offsets))))
  pushpreamble!(ls, Expr(:(=), strides, Expr(:call, gf, tmpsp, QuoteNode(:strd))))
  for (i, p) ∈ enumerate(sp)
    push!(strd_tup.args, Expr(:call, gf, strides, p, false))
    push!(offsets_tup.args, Expr(:call, gf, offsets, p, false))
  end
  #TODO: fix for `T === Bit`.
  column_major = Expr(:tuple)
  for n ∈ eachindex(sp)
    push!(column_major.args, n)
  end
  sptype = Expr(:curly, lv(:StridedPointer), T_sym, length(sp), (C == -1 ? -1 : 1), B, column_major)
  sptr = Expr(:call, sptype, Expr(:call, :pointer, tmpsp), strd_tup, offsets_tup)
  pushpreamble!(ls, Expr(:(=), name, sptr))
  nothing
end
function permute_mref!(ar::ArrayReferenceMeta, C::Int, sp::Vector{Int})
  sp == eachindex(sp) && return nothing
  li = ar.loopedindex; lic = copy(li);
  inds = getindices(ar); indsc = copy(inds);
  offsets = ar.ref.offsets; offsetsc = copy(offsets);
  strides = ar.ref.strides; stridesc = copy(strides);
  for (i, p) ∈ enumerate(sp)
    li[i] = lic[p]
    inds[i] = indsc[p]
    offsets[i] = offsetsc[p]
    strides[i] = stridesc[p]
  end
  C == -1 && makediscontiguous!(getindices(ar))
  return nothing
end
function add_mref!(
    sptrs::Expr, ::LoopSet, ::ArrayReferenceMeta, @nospecialize(_::Type{VectorizationBase.FastRange{T,F,S,O}}),
    ::Int, ::Int, sp::Vector{Int}, name::Symbol
) where {T,F,S,O}
  extract_gsp!(sptrs, name)
  sizeof(T)
end
function create_mrefs!(
    ls::LoopSet, arf::Vector{ArrayRefStruct}, as::Vector{Symbol}, os::Vector{Symbol},
    nopsv::Vector{NOpsType}, expanded::Vector{Bool}, ::Type{Tuple{}}
)
    length(arf) == 0 || throw(ArgumentError("Length of array ref vector should be 0 if there are no stridedpointers."))
    Vector{ArrayReferenceMeta}(undef, length(arf)), Int[]
end
function stabilize_grouped_stridedpointer_type(C, B, R)
  N = (length(C))::Int
  Cv = Vector{Int}(undef, N)
  Bv = Vector{Int}(undef, N)
  Rv = Vector{Tuple{NTuple{8,Int},Int}}(undef, N)
  for n ∈ 1:N
    Cv[n] = C[n]
    Bv[n] = B[n]
    Rₙ = R[n]
    let L::Int = length(Rₙ)
      Rv[n] = (ntuple(i -> i > L ? typemax(Int) : (Rₙ[i])::Int, Val(8))::NTuple{8,Int}, L)
    end
  end
  Cv, Bv, Rv
end
function create_mrefs!(
  ls::LoopSet, arf::Vector{ArrayRefStruct}, as::Vector{Symbol}, os::Vector{Symbol},
  nopsv::Vector{NOpsType}, expanded::Vector{Bool}, @nospecialize(_::Type{VectorizationBase.GroupedStridedPointers{P,C,B,R,I,X,O}})
) where {P,C,B,R,I,X,O}

  Cv,Bv,Rv = stabilize_grouped_stridedpointer_type(C, B, R)
  _create_mrefs!(ls, arf, as, os, nopsv, expanded, P.parameters, Cv, Bv, Rv)
end
function _create_mrefs!(
  ls::LoopSet, arf::Vector{ArrayRefStruct}, as::Vector{Symbol}, os::Vector{Symbol},
  nopsv::Vector{NOpsType}, expanded::Vector{Bool}, P::Core.SimpleVector, C::Vector{Int}, B::Vector{Int}, R::Vector{Tuple{NTuple{8,Int},Int}}
)
    mrefs::Vector{ArrayReferenceMeta} = Vector{ArrayReferenceMeta}(undef, length(arf))
    elementbytes::Vector{Int} = Vector{Int}(undef, length(arf))
    sptrs = Expr(:tuple)
    # pushpreamble!(ls, Expr(:(=), sptrs, :(VectorizationBase.stridedpointers(getfield(vargs, 1, false)))))
    pushpreamble!(ls, Expr(:(=), sptrs, :(VectorizationBase.stridedpointers(getfield(var"#vargs#", 1, false)))))
    j = 0
    rank_to_sps = Vector{Tuple{Int,Vector{Int}}}(undef, length(arf))
    for i ∈ eachindex(arf)
        ar = ArrayReferenceMeta(ls, arf[i], as, os, nopsv, expanded)
        duplicate = false
        vptrar = vptr(ar)
        for k ∈ 1:i-1
            if vptr(mrefs[k]) === vptrar
                duplicate = true
                # if isassigned(rank_to_sps, k)
                Cₖ, sp = rank_to_sps[k]
                permute_mref!(ar, Cₖ, sp)
                elementbytes[i] = elementbytes[k]
                # end
                break
            end
        end
        if !duplicate
            j += 1
            sp = rank_to_sortperm(R[j])::Vector{Int}
            rank_to_sps[i] = (C[j],sp)
            elementbytes[i] = add_mref!(sptrs, ls, ar, P[j], C[j], B[j], sp, vptr(ar))
        end
        mrefs[i] = ar
    end
    mrefs, elementbytes
end

function num_parameters(AM)
    num_param::Int = AM[1]
    # num_param += length(AM[2].parameters)
    num_param + length(AM[3])
end
function gen_array_syminds(AM)
    Symbol[Symbol("##arraysymbolind##"*i*'#') for i ∈ 1:(AM[1])::Int]
end
function process_metadata!(ls::LoopSet, AM, extractind::Int)
    opoffsets = ls.operation_offsets
    expandbyoffset!(ls.outer_reductions, AM[2], opoffsets)
    for (i,si) ∈ enumerate(AM[3])
        sii = si::Int
        opid = opoffsets[sii]+1
        if instruction(operations(ls)[opid]) ≠ DROPPEDCONSTANT
            s = gensym(:symlicm)
            push!(ls.preamble_symsym, (opid, s))
            pushpreamble!(ls, Expr(:(=), s, extract_varg((extractind += 1))))
        end
    end
    expandbyoffset!(ls.preamble_symint, AM[4], opoffsets)
    expandbyoffset!(ls.preamble_symfloat, AM[5], opoffsets)
    expandbyoffset!(ls.preamble_zeros, AM[6], opoffsets)
    expandbyoffset!(ls.preamble_funcofeltypes, AM[7], opoffsets)
    return extractind
end
function expandbyoffset!(indexpand::Vector{T}, inds, offsets::Vector{Int}, expand::Bool = true) where {T <: Union{Int,Tuple{Int,<:Any}}}
    for _ind ∈ inds
        ind = T === Int ? _ind : first(_ind)
        base = offsets[ind] + 1
        for inda ∈ base:(expand ? offsets[ind+1] : base)
            if T === Int
                push!(indexpand, inda)
            elseif T === Tuple{Int,Tuple{Int,Int32,Bool}}
                li = last(_ind)
                push!(indexpand, (inda,(li[1],li[2],li[3])))
            else
                push!(indexpand, (inda,last(_ind)))
            end
        end
    end
    indexpand
end
expandbyoffset(inds::Vector{Int}, offsets::Vector{Int}, expand::Bool) = expandbyoffset!(Int[], inds, offsets, expand)
function loopindex!(idxs::Vector{Int}, ls::LoopSet, u::Unsigned, shift::Unsigned)
  mask = (one(shift) << shift) - one(shift) # mask to zero out all but shift-bits
  while u != zero(u)
    pushfirst!(idxs, ( u % typeof(shift) ) & mask)
    u >>= shift
  end
  idxs
end
loopindex(ls::LoopSet, u::Unsigned, shift::Unsigned) = reverse!(loopindex!(Int[], ls, u, shift))
function loopindexoffset(ls::LoopSet, u::Unsigned, li::Bool, expand::Bool = false)
    if li
        shift = 0x04
        offsets = ls.loopsymbol_offsets
    else
        shift = 0x08
        offsets = ls.operation_offsets
    end
    idxs = loopindex(ls, u, shift)
    expandbyoffset(idxs, offsets, expand)
end
function parents_symvec(ls::LoopSet, u::Unsigned, expand, offset)
    idxs = loopindexoffset(ls, u, true, expand)   # TODO DRY  (undesirable that this gets hard-coded in multiple places)
    return Symbol[getloopsym(ls, i + offset) for i ∈ idxs]
end
loopdependencies(ls::LoopSet, os::OperationStruct, expand = false, offset = 0) = parents_symvec(ls, os.loopdeps, expand, offset)
reduceddependencies(ls::LoopSet, os::OperationStruct, expand = false, offset = 0) = parents_symvec(ls, os.reduceddeps, expand, offset)
childdependencies(ls::LoopSet, os::OperationStruct, expand = false, offset = 0) = parents_symvec(ls, os.childdeps, expand, offset)

# parents(ls::LoopSet, u::UInt128) = loopindexoffset(ls, u, false)
function parents(ls::LoopSet, u₀::UInt128, u₁::UInt128)
  idxs = Int[]
  u₁ == zero(u₁) || loopindex!(idxs, ls, u₁, 0x0010)
  loopindex!(idxs, ls, u₀, 0x0010)
  reverse!(idxs)
end
parents(ls::LoopSet, os::OperationStruct) = parents(ls, os.parents₀, os.parents₁)

expandedopname(opsymbol::Symbol, offset::Integer) = Symbol(String(opsymbol)*'#'*string(offset+1)*'#')
function calcnops(ls::LoopSet, os::OperationStruct)
    optyp = optype(os)
    if (optyp != loopvalue) && (optyp != compute)
        return 1
    end
    offsets = ls.loopsymbol_offsets
    idxs = loopindex(ls, os.loopdeps, 0x04)  # TODO DRY
    iszero(length(idxs)) && return 1
    return maximum(i->offsets[i+1]-offsets[i], idxs)
end
function isexpanded(ls::LoopSet, ops::Vector{OperationStruct}, nopsv::Vector{NOpsType}, i::Int)
    nops = nopsv[i]
    # nops isa Vector{Int} only if accesses_memory(os), which means isexpanded must be false
    (nops === 1 || isa(nops, Vector{Int})) && return false
    os = ops[i]
    optyp = optype(os)
    if optyp == compute
        any(j -> isexpanded(ls, ops, nopsv, j), parents(ls, os))
    elseif optyp == loopvalue
        true
    else
        false
    end
end
function mref_elbytes(os::OperationStruct, mrefs::Vector{ArrayReferenceMeta}, elementbytes::Vector{Int})
  if isload(os) | isstore(os)
    mrefs[os.array], elementbytes[os.array]
  else
    NOTAREFERENCE, 4
  end
end
function add_op!(
    ls::LoopSet, instr::Instruction, ops::Vector{OperationStruct}, nopsv::Vector{NOpsType}, expandedv::Vector{Bool}, i::Int,
    mrefs::Vector{ArrayReferenceMeta}, opsymbol, elementbytes::Vector{Int}
)
    os = ops[i]
    mref, elbytes = mref_elbytes(os, mrefs, elementbytes)
    # opsymbol = (isconstant(os) && instr != LOOPCONSTANT) ? instr.instr : opsymbol
    # If it's a CartesianIndex add or subtract, we may have to add multiple operations
    expanded = expandedv[i]# isexpanded(ls, ops, nopsv, i)
    opoffsets = ls.operation_offsets
    # offsets = ls.loopsymbol_offsets
    optyp = optype(os)
    if !expanded
        op = Operation(
            length(operations(ls)), opsymbol, elbytes, instr,
            optyp, loopdependencies(ls, os, true), reduceddependencies(ls, os, true),
            Operation[], mref, childdependencies(ls, os, true)
        )
        push!(ls.operations, op)
        push!(opoffsets, opoffsets[end] + 1)
        return
    end
    nops = (nopsv[i])::Int # if it were a vector, it would have to have been expanded
    # if expanded, optyp must be either loopvalue, or compute (with loopvalues in its ancestry, not cutoff by loads)
    for offset = 0:nops-1
        sym = nops === 1 ? opsymbol : expandedopname(opsymbol, offset)
        op = Operation(
            length(operations(ls)), sym, elbytes, instr, optyp,
            loopdependencies(ls, os, false, offset), reduceddependencies(ls, os, false, offset),
            Operation[], mref, childdependencies(ls, os, false, offset)
        )
        push!(ls.operations, op)
    end
    push!(opoffsets, opoffsets[end] + nops)
    nothing
end
function add_parents_to_op!(ls::LoopSet, op::Operation, up₀::UInt128, up₁::UInt128, k::Int, Δ::Int)
    vparents = parents(op)
    ops = operations(ls)
    offsets = ls.operation_offsets
    if isone(Δ) # not expanded
        @assert isone(k)
        for i ∈ parents(ls, up₀, up₁)
            # FIXME; children also filled in cacheunrolled
            for j ∈ offsets[i]+1:offsets[i+1] # if parents are expanded, add them all
                opp = ops[j]
                pushfirst!(vparents, opp)
                push!(children(opp), op)
            end
        end
    else#if isexpanded
        # Do we want to require that all Δidxs are equal?
        # Because `CartesianIndex((2,3)) - 1` results in a methoderorr, I think this is reasonable for now
        # FIXME; children also filled in cacheunrolled
        for i ∈ parents(ls, up₀, up₁)
            opp = ops[offsets[i]+k]
            pushfirst!(vparents, opp)
            push!(children(opp), op)
        end
    end
end
function add_parents_to_ops!(ls::LoopSet, ops::Vector{OperationStruct}, constoffset)
    offsets = ls.operation_offsets
    for i in 1:length(offsets)-1
        pos = offsets[i]
        Δ = offsets[i+1]-pos
        for k ∈ 1:Δ
            op = ls.operations[pos+k]
            if isconstant(op)
                instr = instruction(op)
                if !skip_constant(instr)
                    constoffset += 1
                    pushpreamble!(ls, Expr(:(=), instr.instr, extract_varg(constoffset)))
                end
            elseif !isloopvalue(op)
                add_parents_to_op!(ls, op, ops[i].parents₀, ops[i].parents₁, k, Δ)
            end
        end
    end
    constoffset
end
function add_ops!(
    ls::LoopSet, instr::Vector{Instruction}, ops::Vector{OperationStruct}, mrefs::Vector{ArrayReferenceMeta}, elementbytes::Vector{Int},
    opsymbols::Vector{Symbol}, constoffset::Int, nopsv::Vector{NOpsType}, expandedv::Vector{Bool}
)
    # @show ls.loopsymbols ls.loopsymbol_offsets
    for i ∈ eachindex(ops)
        os = ops[i]
        opsymbol = opsymbols[os.symid]
        add_op!(ls, instr[i], ops, nopsv, expandedv, i, mrefs, opsymbol, elementbytes)
    end
    add_parents_to_ops!(ls, ops, constoffset)
    # for op ∈ operations(ls)
    #     if isstore(op) && isreduction(op) && iszero(length(loopdependencies(op)))
    #         addreduct_to_outer_reductions!(ls, op)
    #     end
    # end
    # for op in operations(ls)
        # @show op
    # end
end

# elbytes(::VectorizationBase.AbstractPointer{T}) where {T} = sizeof(T)::Int
# typeeltype(::Type{P}) where {T,P<:VectorizationBase.AbstractStridedPointer{T}} = T
typeeltype(::Type{Ptr{T}}) where {T} = T
# typeeltype(::Type{Core.LLVMPtr{T,0}}) where {T} = T
typeeltype(::Type{VectorizationBase.FastRange{T,F,S,O}}) where {T,F,S,O} = T
typeeltype(::Type{T}) where {T<:Real} = T
# typeeltype(::Any) = Int8

function add_array_symbols!(ls::LoopSet, arraysymbolinds::Vector{Symbol}, offset::Int)
    for as ∈ arraysymbolinds
        pushpreamble!(ls, Expr(:(=), as, extract_varg((offset+=1))))
    end
    return offset
end
function extract_external_functions!(ls::LoopSet, offset::Int, vargs)
    for op ∈ operations(ls)
        if iscompute(op)
            instr = instruction(op)
            if instr.mod != :LoopVectorization
                offset += 1
                instr_new = get(FUNCTIONSYMBOLS, vargs[offset], instr)
                if instr_new === instr
                    extractf = Expr(:call, GlobalRef(Core,:getfield), Symbol("#vargs#"), offset, false)
                    pushpreamble!(ls, Expr(:(=), instr.instr, extractf))
                else
                    op.instruction = instr_new
                end
            end
        end
    end
    offset
end
outer_reduct_init_typename(op::Operation) = Symbol(mangledvar(op), "#or#init#type#")
function extract_outerreduct_types!(ls::LoopSet, offset::Int, vargs)
  # for op
  for or ∈ ls.outer_reductions
    extractt = Expr(:call, GlobalRef(Core,:getfield), Symbol("#vargs#"), (offset+=1), false)
    pushpreamble!(ls, Expr(:(=), outer_reduct_init_typename(operations(ls)[or]), extractt))
  end
  offset
end
function sizeofeltypes(v)::Int
    num_arrays = length(v)::Int
    if num_arrays == 0
        return 8
    end
    T = typeeltype(v[1])
    sz = if (VectorizationBase.simd_integer_register_size() != VectorizationBase.register_size()) && T <: Integer # hack
        (VectorizationBase.register_size() ÷ VectorizationBase.simd_integer_register_size()) * sizeof(T)
    else
        sz = sizeof(T)
    end
    for i ∈ 2:num_arrays
        Ttemp = typeeltype(v[i])
        szᵢ = if (VectorizationBase.simd_integer_register_size() != VectorizationBase.register_size()) && T <: Integer # hack
            (VectorizationBase.register_size() ÷ VectorizationBase.simd_integer_register_size()) * sizeof(T)
        else
            sizeof(Ttemp)
        end
        # if !VectorizationBase.SIMD_NATIVE_INTEGERS && Ttemp <: Integer # hack
        #     return VectorizationBase.register_size()
        # end
        # T = promote_type(T, Ttemp)
        sz = max(szᵢ, sz)
    end
    sz
    # sizeof(T)
end

function avx_loopset!(
    ls::LoopSet, instr::Vector{Instruction}, ops::Vector{OperationStruct}, arf::Vector{ArrayRefStruct},
    AM::Vector{Any}, LPSYM::Vector{Any}, LB::Core.SimpleVector, vargs::Core.SimpleVector
)
    pushpreamble!(ls, :((var"#loop#bounds#", var"#vargs#") = var"#lv#tuple#args#"))
    add_loops!(ls, LPSYM, LB)
    resize!(ls.loop_order, ls.loopsymbol_offsets[end])
    arraysymbolinds = gen_array_syminds(AM)
    opsymbols = [gensym("op") for _ ∈ eachindex(ops)]
    nopsv = NOpsType[calcnops(ls, op) for op in ops]
    expandedv = [isexpanded(ls, ops, nopsv, i) for i ∈ eachindex(ops)]

    resize!(ls.loopindexesbit, length(ls.loops)); fill!(ls.loopindexesbit, false);
    mrefs, elementbytes = create_mrefs!(ls, arf, arraysymbolinds, opsymbols, nopsv, expandedv, vargs[1])
    for mref ∈ mrefs
        push!(ls.includedactualarrays, vptr(mref))
    end
    # extra args extraction
    extractind = add_ops!(ls, instr, ops, mrefs, elementbytes, opsymbols, 1, nopsv, expandedv)
    extractind = process_metadata!(ls, AM, extractind)
    extractind = add_array_symbols!(ls, arraysymbolinds, extractind)
    extractind = extract_external_functions!(ls, extractind, vargs)
    extractind = extract_outerreduct_types!(ls, extractind, vargs)
    ls
end
function avx_body(ls::LoopSet, UNROLL::Tuple{Bool,Int8,Int8,Int8,Bool,Int,Int,Int,Int,Int,Int,Int,UInt})
    inline, u₁, u₂, v, isbroadcast, W, rs, rc, cls, l1, l2, l3, nt = UNROLL
    q = (iszero(u₁) & iszero(v)) ? lower_and_split_loops(ls, inline % Int) : lower(ls, u₁ % Int, u₂ % Int, v % Int, inline % Int)
    ls.isbroadcast = isbroadcast
    iszero(length(ls.outer_reductions)) ? push!(q.args, nothing) : push!(q.args, loopset_return_value(ls, Val(true)))
    q
end

function _turbo_loopset_debug(::Val{UNROLL}, ::Val{OPS}, ::Val{ARF}, ::Val{AM}, ::Val{LPSYM}, _vargs::Tuple{LB,V}) where {UNROLL, OPS, ARF, AM, LPSYM, LB, V}
    # @show OPS ARF AM LPSYM _vargs
    _turbo_loopset(OPS, ARF, AM, LPSYM, _vargs[1].parameters, V.parameters, UNROLL)
end
function tovector(@nospecialize(t))
    v = Vector{Any}(undef, length(t))
    for i ∈ eachindex(v)
        tᵢ = t[i]
        if tᵢ isa Tuple # reduce specialization?
            v[i] = tovector(tᵢ)
        else
            v[i] = tᵢ
        end
    end
    v
end
function _turbo_loopset(
  @nospecialize(OPSsv), @nospecialize(ARFsv), @nospecialize(AMsv), @nospecialize(LPSYMsv), LBsv::Core.SimpleVector, vargs::Core.SimpleVector,
  UNROLL::Tuple{Bool,Int8,Int8,Int8,Bool,Int,Int,Int,Int,Int,Int,Int,UInt}
)
    nops = length(OPSsv) ÷ 3
    instr = Instruction[Instruction(OPSsv[3i+1], OPSsv[3i+2]) for i ∈ 0:nops-1]
    ops = OperationStruct[ OPSsv[3i] for i ∈ 1:nops ]
    ls = LoopSet(:LoopVectorization)
    inline, u₁, u₂, v, isbroadcast, W, rs, rc, cls, l1, l2, l3, nt = UNROLL
    set_hw!(ls, rs, rc, cls, l1, l2, l3); ls.vector_width = W; ls.isbroadcast = isbroadcast
    arsv = Vector{ArrayRefStruct}(undef, length(ARFsv))
    for i ∈ eachindex(arsv)
        arsv[i] = ARFsv[i]
    end
    avx_loopset!(ls, instr, ops, arsv, tovector(AMsv), tovector(LPSYMsv), LBsv, vargs)
end

@static if VERSION ≥ v"1.7.0-DEV.421"
  using Base: @aggressive_constprop
else
  macro aggressive_constprop(ex); esc(ex); end
end

"""
    _turbo_!(unroll, ops, arf, am, lpsym, lb, vargs...)

Execute an `@turbo` block. The block's code is represented via the arguments:
- `unroll` is `Val((u₁,u₂))` and specifies the loop unrolling factor(s).
  These values may be supplied manually via the `unroll` keyword
  of [`@turbo`](@ref).
- `ops` is `Tuple{mod1, sym1, op1, mod2, sym2, op2...}` encoding the operations of the loop.
  `mod` and `sym` encode the module and symbol of the called function; `op` is an [`OperationStruct`](@ref)
  encoding the details of the operation.
- `arf` is `Tuple{arf1, arf2...}`, where each `arfi` is an [`ArrayRefStruct`](@ref) encoding
  an array reference.
- `am` contains miscellaneous data about the LoopSet (see `process_metadata!`)
- `lpsym` is `Tuple{:i,:j,...}`, a Tuple of the "loop symbols", i.e. the item variable `i` in `for i ∈ iter`
- `lb` is `Tuple{RngTypei,RngTypej,...}`, a Tuple encoding syntactically-knowable information about
  the iterators corresponding to `lpsym`. For example, in `for i ∈ 1:n`, the `1:n` would be encoded with
  `StaticLowerUnitRange(1)` because the lower bound of the iterator can be determined to be 1.
- `vargs...` holds the encoded pointers of all the arrays (see `VectorizationBase`'s various pointer types).
"""
@aggressive_constprop @generated function _turbo_!(
    ::Val{var"#UNROLL#"}, ::Val{var"#OPS#"}, ::Val{var"#ARF#"}, ::Val{var"#AM#"}, ::Val{var"#LPSYM#"}, ::Val{Tuple{var"#LB#",var"#V#"}}, var"#flattened#var#arguments#"::Vararg{Any,var"#num#vargs#"}
) where {var"#UNROLL#", var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#", var"#V#", var"#num#vargs#"}
  # 1 + 1 # Irrelevant line you can comment out/in to force recompilation...
  ls = _turbo_loopset(var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#".parameters, var"#V#".parameters, var"#UNROLL#")
  pushfirst!(ls.preamble.args, :(var"#lv#tuple#args#" = reassemble_tuple(Tuple{var"#LB#",var"#V#"}, var"#flattened#var#arguments#")))
  # return @show avx_body(ls, var"#UNROLL#")
  if last(var"#UNROLL#") > 1
    inline, u₁, u₂, v, isbroadcast, W, rs, rc, cls, l1, l2, l3, nt = var"#UNROLL#"
    # wrap in `var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#"` in `Expr` to homogenize types
    avx_threads_expr(
      ls, (inline, u₁, u₂, v, isbroadcast, W, rs, rc, cls, l1, l2, l3, one(UInt)), nt,
      :(Val{$(var"#OPS#")}()), :(Val{$(var"#ARF#")}()), :(Val{$(var"#AM#")}()), :(Val{$(var"#LPSYM#")}())
    )
  else
    # Main.BODY[] = avx_body(ls, var"#UNROLL#")
    # @show avx_body(ls, var"#UNROLL#")
    avx_body(ls, var"#UNROLL#")
  end
  # @show var"#UNROLL#", var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#"
end


push!(LOAD_PATH, abspath("src"))
using Yao
using QMPS
using DelimitedFiles, JLD2, FileIO, Pkg

# CUDA switch
const USE_CUDA = haskey(Pkg.installed(), "CuYao")
USE_CUDA && println("Hint: Using CUDA since `CuYao` is detected. Edit file `applications.jl` to modify CUDA settings, like switching computing devices.")
USE_CUDA && include("CuChem.jl")
USE_CUDA && device!(CuDevice(0))

"""Heisenberg model without frustration, with open boundary condition."""
simple_model_heis(size...) = Heisenberg(size...; periodic=false)

"""Heisenberg model with frustration strenth J2 = 0.5, with open boundary condition."""
simple_model_j1j2(size...) = J1J2(size...; periodic=false, J2=0.5)

"""
    simple_ansatz(nbit::Int, symmetry::Symbol, depth::Int; load_params::Bool=false)

Load a predefined MPS inspired ansatz with 4 virtual qubits, batch size 4096.
If `load_params` is `true`, load parameters in training step 500.
"""
function simple_ansatz(nbit::Int, symmetry::Symbol, depth::Int; load_params::Bool=false)
    V = 4     # number of virtual qubits
    batch_size = 4096
    load_step = 500

    ansatz = model(Val(symmetry); nbit=nbit, nlayer=depth, V=V, B=batch_size, pairs=pair_ring(V+1))
    USE_CUDA && (ansatz = ansatz |> cu)

    if load_params
        filename = "data/chem_$(symmetry)_d$(depth)_N$(nbit)_V$(V)_S$(load_step).jld2"
        params = load_training(filename)[3]
        println("loading file $filename")
        println("Number of parameters is ", params |> length)
        dispatch!(ansatz.circuit, params)
    end

    ansatz
end

"""
    save_training(filename, qopt::Adam, loss, params, fidelity)

Save training status.
"""
function save_training(filename, qopt::Adam, loss, params, fidelity)
    save(filename, "qopt", qopt, "loss", loss, "params", params, "fidelity", fidelity)
end

"""
    load_training(filename) -> Tuple

Load training status (qopt, loss, params, fidelity).
"""
function load_training(filename)
    res = load(filename)
    res["qopt"], res["loss"], res["params"], res["fidelity"]
end

"""
    run_train(ansatz, model; SAVE_ID, niter=500, start_point=0, save_step=0)

Run a variational quantum eigensolver to solve a `model` Hamiltonian with MPS inspired `ansatz`

* SAVE_ID, token used for saving data,
* niter, number of training steps,
* start_point, load a save point, default is 0 (random parameters).
* save_step, save the training every `save_step` steps.
"""
function run_train(ansatz, model; SAVE_ID, niter::Int=500, start_point::Int=0, save_step::Int=10)
    nbit = nbit_simulated(ansatz)
    V = ansatz.nbit_virtual
    filename(k::Int) = "data/chem_$(SAVE_ID)_N$(nbit)_V$(V)_S$(k).jld2"
    if start_point==0
        optimizer = Adam(lr=0.1)
        history = Float64[]
        fidelities = Float64[]
    else
        optimizer, history, _params, fidelities = load_training(filename(start_point))
        dispatch!(ansatz.circuit, _params)
    end
    qo = QMPSOptimizer(ansatz, model, optimizer)
    gs = ground_state(model)

    for (k_, p) in enumerate(qo)
        k = k_ + start_point
        curr_loss = energy(ansatz, model)/nbit
        fid = fidelity_exact(ansatz, gs)[]
        push!(history, curr_loss)
        push!(fidelities, fid)
        println("step = $k, energy/site = $curr_loss, fidelity = $(fid)")
        flush(stdout)
        if k%save_step == 0 || k == niter
            save_training(filename(k), qo.optimizer, history, parameters(ansatz.circuit), fidelities)
            k >= niter && break
        end
    end
end

"""
    correlation_matrix(ansatz, op; SAVE_ID)

Calculate and save correlation matrix <Z_i Z_j>.
"""
function correlation_matrix(ansatz; SAVE_ID)
    nbit = nbit_simulated(ansatz)
    V = ansatz.nbit_virtual

    om = zeros(Float64, nbit, nbit)
    for i =1:nbit, j=1:nbit
        if i!=j
            om[i,j] = measure_corr(ansatz, i=>Z, j=>Z) |> real
            println("<σz($i)σz($j)> = $(om[i,j])")
        end
    end

    for (token, var) in [
                         ("om", om),
                        ]
        writedlm("data/chem_$(SAVE_ID)_$(token)$(tag)_N$(nbit)_V$V.dat", var)
    end
end


using .Threads, BenchmarkTools

function machin_series(n::Int; handover=false)
    qpi = 0.0
    for i in 1:n
        qpi += (-1)^(i+1)/(2i-1)
        handover && yield()  # we yield here !!!!
    end
    qpi*4
end

function startonthread(id::Int, f::F) where {F<:Function}
    t = Task(nothing)
    @threads for i in 1:nthreads()
        if i == id
            t = @async f()
        end
    end
    fetch(t)
end

println("Benchmark results on thread 1-4 without yielding:")
@btime startonthread(1, ()->machin_series(10_000))
@btime startonthread(2, ()->machin_series(10_000))
@btime startonthread(3, ()->machin_series(10_000))
@btime startonthread(4, ()->machin_series(10_000))

println("Benchmark results on thread 1-4 with yielding:")
@btime startonthread(1, ()->machin_series(10_000, handover=true))
@btime startonthread(2, ()->machin_series(10_000, handover=true))
@btime startonthread(3, ()->machin_series(10_000, handover=true))
@btime startonthread(4, ()->machin_series(10_000, handover=true))


# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   ModiaMath.Variables (ModiaMath/Variables/_module.jl)
#

# Define SimulationModel struct and functions

"""
    sm = ModiaMath.SimulationModel(model::ModiaMath.AbstractComponentWithVariables;
                                   startTime = 0.0, stopTime = 1.0, tolerance = 1e-4,
                                   interval = (stopTime-startTime)/500.0)

Return a simulationModel `sm` struct that can be simulated with [`ModiaMath.simulate!`](@ref).
The given `model` is typically constructed with the [`@component`](@ref) macro.
As keyword arguments, default values for `startTime`, `stopTime`, `tolerance`, `interval` can be given.
"""
mutable struct SimulationModel <: ModiaMath.AbstractSimulationModel
    modelName::String
    simulationState::ModiaMath.SimulationState
    var::ModelVariables
    model::ModiaMath.AbstractComponentWithVariables

    function SimulationModel(model::ModiaMath.AbstractComponentWithVariables;
                             startTime = 0.0,
                             stopTime  = 1.0,
                             tolerance = 1e-4,
                             interval  = (stopTime-startTime)/500.0,
                             structureOfDAE = ModiaMath.DAE_LinearDerivativesAndConstraints,
							 has_constraintDerivatives = false,
                             hev = 1e-8,
                             scaleConstraintsAtEvents::Bool = true)
        modelName = ModiaMath.componentName(model)
        var = ModelVariables(model)
        x = zeros(var.nx)
        x_fixed   = fill(false,var.nx)
        x_nominal = fill(1.0, var.nx)
        ModiaMath.copy_start_to_x!(var,x,x_fixed,x_nominal)

        # Last nfc equations are the constraint equations
        is_constraint = fill(false, var.nx)
        for i = (var.nx-var.nfc+1):var.nx
            is_constraint[i] = true
        end

        simulationState = ModiaMath.SimulationState(modelName, getModelResidues!, x, getVariableName;
                                                    structureOfDAE = structureOfDAE,
                                                    is_constraint = is_constraint,
                                                    has_constraintDerivatives = has_constraintDerivatives,
                                                    x_fixed = x_fixed, x_nominal=x_nominal,
                                                    getResultNames = getResultNames, storeResult! = storeVariables!,
                                                    getResult        = getResult,
                                                    defaultStartTime = startTime,
                                                    defaultStopTime  = stopTime,
                                                    defaultTolerance = tolerance,
                                                    defaultInterval  = interval,
                                                    hev = hev,
                                                    scaleConstraintsAtEvents = scaleConstraintsAtEvents)

        # model._internal.simulationState = simulationState
        new(String(modelName), simulationState, var, model)
    end
end

print_ModelVariables(simulationModel::SimulationModel) = print_ModelVariables(simulationModel.var)

getVariableName(model,vcat,vindex) = ModiaMath.DAE.getVariableName(model,vcat,vindex;
                                                                   xNames = model.var.x_names)

getResultNames(model::Any) = model.var.result_names

function storeVariables!(res::ModiaMath.RawResult, model, t::Float64, x::Vector{Float64}, derx::Vector{Float64}, w::Vector{Float64})
    res.nt += 1
    i = res.nt
    ntMax = size(res.data,1)
    if i > ntMax
        # Result storage is full, double the current storage
        res.data = [res.data; zeros(ntMax,size(res.data,2))]
    end

    # Copy data in raw result storage
    for v in model.var.result_var
        if ModiaMath.isScalar(v)
            res.data[i, v.iresult] = v.value
        else
            vv = v.value
            for j in 1:length(vv)
                res.data[i, v.iresult+j-1 ] = vv[j]
            end
        end
    end
    return nothing
end

function getResult(model::ModiaMath.AbstractSimulationModel, res::ModiaMath.RawResult)
    nt         = res.nt
    data       = res.data

    seriesDict = Dict{AbstractString,Any}()
    varDict    = Dict{Symbol,ModiaMath.AbstractVariable}()
    for v in model.var.result_var
        name    = instanceName(v)
        len     = length(v.value)
        iresult = v.iresult
        if len == 1
            seriesDict[string(name)] = @view data[1:nt, iresult]
        else
            seriesDict[string(name)] = @view data[1:nt, iresult:iresult+len-1]
        end
        varDict[name] = v
   end
   result = ModiaMath.ResultWithVariables(seriesDict, varDict, model.modelName)
   return result, nt
end

computeVariables!(model::ModiaMath.AbstractComponentWithVariables,
                  simulationState::ModiaMath.SimulationState) =
          error("Function ModiaMath.computeVariables!(model) not defined for model ", typeof(model))

function getModelResidues!(m::SimulationModel, t::Float64, _x::Vector{Float64}, _derx::Vector{Float64}, _r::Vector{Float64}, _w::Vector{Float64})
    # Copy _x and _derx values to variables
    ModiaMath.copy_x_and_derx_to_variables!(t, _x, _derx, m.var)

    # Compute model variables
    m.simulationState.time = t
    computeVariables!(m.model, m.simulationState)

    # Copy variables to residues
    ModiaMath.copy_variables_to_residue!(m.var,_x,_derx,_r)
    return nothing
end

# Return a table of actual variable and residue values from nonlinear solver in case of error
function ModiaMath.getVariableAndResidueValues(m::SimulationModel)
   var = m.var
   v_table = get_variableValueTable(var)
   r_table = get_residueValueTable(var, m.simulationState.derxev)
   return (v_table, r_table)
end


using Flux
using Transformers.Basic

struct Encoder
    layers::Array{Transformer}
end

Flux.@functor Encoder
Flux.trainable(encoder::Encoder) = (encoder.layers,)

function (encoder::Encoder)(embedding)
    for transformer ∈ encoder.layers
        embedding = transformer(embedding)
    end
    return embedding
end

function Encoder(
    size::Integer,
    head::Integer,
    hs::Integer,
    ps::Integer,
    layer::Integer;
    act=relu,
    pdrop=0.1
)::Encoder
    return Encoder([
        Transformer(size, head, hs, ps, act=act, pdrop=pdrop)
        for i ∈ 1:layer
    ])
end


# Filters using Nullable types to represent missing measurements

function update{T}(kf::KalmanFilter,y::Nullable{Vector{T}})
    if isnull(y.y)
        return kf
    else
        return update(kf,Observation(get(y)))
    end
end

function update!{T}(kf::KalmanFilter,y::Nullable{Vector{T}})
    if isnull(y)
        return kf
    else
        update!(kf,Observation(get(y)))
        return kf
    end
end

function predictupdate{T}(kf::BasicKalmanFilter,y::Nullable{Vector{T}})
    update(predict(kf),y)
end

function predictupdate!{T}(kf::BasicKalmanFilter,y::Nullable{Vector{T}})
    update!(predict!(kf),y)
end


module PointsOfInterest

using AutoHashEquals, Base.Dates, IntervalSets

export POI

import Base: ∈, push!, delete!, empty!

@auto_hash_equals struct Video
    name::String
    time::DateTime
end

struct Metadata
    pois::Vector{String}
    files::Vector{Video}
end

const sides = (:left, :right)

@auto_hash_equals struct Point
    file::Int
    time::Millisecond

    function Point(md::Metadata, file::String, time::Millisecond)
        i = findfirst(x -> x.name == file, md.files)
        @assert i ≠ 0 "file not found in metadata"
        @assert time ≥ Millisecond(0) "negative times not allowed"
        new(i, time)
    end
end

Base.isless(p1::Point, p2::Point) = p1.file == p2.file ? p1.time < p2.time : p1.file < p2.file

@auto_hash_equals mutable struct POI
    name::Int
    interval::ClosedInterval{Point}
    comment::String

    function POI(md::Metadata, name::String, interval::ClosedInterval{Point}, comment::String)
        i = findfirst(md.pois, name)
        @assert i ≠ 0 "POI not found in metadata"
        @assert !isempty(interval) "interval is non-existent"
        new(i, interval, comment)
    end
end

POI(md::Metadata, name::String, start::Tuple{String, Millisecond}, stop::Tuple{String, Millisecond}, comment::String) = POI(md, name, Point(md, start...), Point(md, stop...), comment)

# POI(md::Metadata, name::String, label::String, start_file::String, start_time::Millisecond, stop_file::String, stop_time::Millisecond, comment::String) = POI(name, label, Point(md, start_file, start_time), Point(md, stop_file, stop_time), comment)

struct Log
    md::Metadata

    pois::Vector{POI}
end

Log(md::Metadata) = Log(md, POI[])

# findfirst

# findfirst(a::Log, x::Rep) = findfirst(a.reps, x)

# in

∈(x::POI, a::Log) = x ∈ a.pois

∈(x::Video, a::Log) = x ∈ a.files

# pushes

function push!(a::Log, x::Video)
    x ∈ a && return a
    i = findfirst(v -> v.time > x.time, a.md.files)
    insert!(a.md.files, i, x)
    for p in a.pois
        left = p.interval.left
        if left.file ≥ i
            right = p.interval.right
            p.interval = Point(left.file + 1, left.time)..Point(right.file + 1, right.time)
        end
    end
    return a
end

function push!(a::Log, x::POI)
    x ∉ a && push!(a.pois, x)
    return a
end

# deletes

function delete!(a::Log, x::POI)
    i = findfirst(a.pois, x)
    i ≠ 0 && deleteat!(a.pois, i)
    return a
end

function delete!(a::Log, x::Video)
    i = findfirst(a.md.files, x)
    i == 0 && return a
    deleteat!(a.md.files, i)
    filter!(p -> p.interval.left.file == i || p.interval.right.file == i, a.pois)
    for p in a.pois
        left = p.interval.left
        if left.file > i
            right = p.interval.right
            p.interval = Point(left.file - 1, left.time)..Point(right.file - 1, right.time)
        end
    end
    return a
end

# replace

function replace!(a::Log, o::POI, n::POI)
    o == n && return a
    n ∈ a && delete!(a, o)
    i = findfirst(a.pois, o)
    @assert i ≠ 0 "old POI not found"
    a.pois[i] = n
    return a
end

# empty

function empty!(a::Log, ::Type{POI})
    empty!(a.pois)
    return a
end

# pop

pop(md::Metadata, x::Point) = (md.files[x.file], x.time)

pop(md::Metadata, x::ClosedInterval{Point}) = (pop(md, x.left), pop(md, x.right))

pop(md::Metadata, x::POI) = (md.pois[x.name], pop(md, x.interval), x.comment)

# combine

function combine_mds(logs::Log...)
    pois = String[]
    videos = Video[]
    for log in logs
        union!(pois, log.md.pois)
        union!(videos, log.md.files)
    end
    sort!(videos, by = v -> v.time)
    return Metadata(pois, videos)
end

function combine(logs::Log...)
    md = combine_mds(logs)
    a = Log(md)
    for log in logs, p in log.pois
        push!(a, POI(md, pop(log.md, p)))
    end
    return a
end

end # module


# this is the default stack page
# ultimately this should be a property of either the CPU or Chipset
const STACK_PAGE = 0x0100

# TODO WTF is with the break (B) flag?
mutable struct FlagsRegister
    C::Bool
    Z::Bool
    I::Bool
    D::Bool
    B::Bool
    Ω::Bool
    V::Bool
    N::Bool

    FlagsRegister() = new(false, false, false, false, false, false, false, false)
end

mutable struct CPU
    A::UInt8  # accumulator register
    X::UInt8  # index registers
    Y::UInt8
    SP::UInt8  # stack pointer register, I'm tempted to make this 16-bit
    PC::UInt16  # program counter register
    flags::FlagsRegister

    function CPU(A::UInt8=0x00, X::UInt8=0x00, Y::UInt8=0x00,
                 SP::UInt8=0xff, PC::UInt16=0x0600, flags::FlagsRegister=FlagsRegister())
        new(A, X, Y, SP, PC, flags)
    end
end

@inline Base.fetch(c::CPU, reg::Symbol) = getfield(c, reg)
@inline store!(c::CPU, reg::Symbol, val::Integer) = setfield!(c, reg, val)

counter!(c::CPU, ℓ::Unsigned) = (c.PC += ℓ)

function describe(c::CPU)
    st = status_string(c)
    A = hexstring(c.A); X = hexstring(c.X); Y = hexstring(c.Y)
    SP = hexstring(c.SP); PC = hexstring(c.PC)
    """
    MOS6502 CPU
        Registers:
            A=$A  X=$X  Y=$Y
            SP=$SP  PC=$PC
        Status: NV-BDIZC
                $st
    """
end

Base.show(io::IO, c::CPU) = print(describe(c))

#==============================================================================================
    <status register access>
==============================================================================================#
status(fr::FlagsRegister, flag::Symbol)::Bool = getfield(fr, flag)
status(c::CPU, flag::Symbol)::Bool = getfield(fr, flag)

# these are toggles
status!(fr::FlagsRegister, flag::Symbol) = setfield!(fr, ~getfield(fr, flag))
status!(c::CPU, flag::Symbol) = status!(c.flags, flag)

status!(fr::FlagsRegister, flag::Symbol, val::Bool) = setfield!(fr, flag, val)
status!(c::CPU, flag::Symbol, val::Bool) = status!(c.flags, flag, val)

function status_string(fr::FlagsRegister)
    string((Int(getfield(fr, f)) for f ∈ reverse(fieldnames(typeof(fr))))...)
end
status_string(c::CPU) = status_string(c.flags)

#==============================================================================================
    </status register access>
==============================================================================================#

#==============================================================================================
    <special pointers>
==============================================================================================#
stackpointer(c::CPU) = Π(STACK_PAGE + c.SP)
#==============================================================================================
    </special pointers>
==============================================================================================#


module InteractionNetworkModels

using Distributed

# Alises for referring to Paths/Interaction Sequences/Samples of Interaction Sequences (purely for readability)
include("aliases.jl")

# Types
include("Types/PathDistributions.jl")
include("Types/MarkovModels.jl")
include("Types/distributions.jl")
include("Types/hollywood_model.jl")

# Data Processing
# include("Data Processing/EdgeCounts.jl")
include("Data Processing/VertexCounts.jl")
include("Data Processing/PathSequences.jl")
include("Data Processing/Multigraphs.jl")
include("Data Processing/string_to_int.jl")
include("Data Processing/high_order_heuristic.jl")
include("Data Processing/remove_repeats.jl")
include("Data Processing/aggregate_graph.jl")

# Metric Models
# ===============

# Model types and utils
# ---------------------
include("Metric Models/SPF.jl")
include("Metric Models/SIS.jl")
include("Metric Models/SIM.jl")
include("Metric Models/utils.jl")

# Model sampling
# ---------------
# Types
include("Metric Models/Model Sampling/Types/initalisers.jl")
include("Metric Models/Model Sampling/Types/SIM_samplers.jl")
include("Metric Models/Model Sampling/Types/SIS_samplers.jl")
include("Metric Models/Model Sampling/Types/SPF_samplers.jl")
include("Metric Models/Model Sampling/Types/outputs.jl")
# SIS
include("Metric Models/Model Sampling/SIS/SIS_sampler.jl")
include("Metric Models/Model Sampling/SIS/SIS_centered_simple.jl")
include("Metric Models/Model Sampling/SIS/SIS_gibbs.jl")
include("Metric Models/Model Sampling/SIS/split_merge/SIS_split_merge_helpers.jl")
include("Metric Models/Model Sampling/SIS/split_merge/SIS_split_merge.jl")
# SIM
include("Metric Models/Model Sampling/SIM/SIM_sampler.jl")
include("Metric Models/Model Sampling/SIM/SIM_subpath.jl")
include("Metric Models/Model Sampling/SIM/SIM_gibbs.jl")
include("Metric Models/Model Sampling/SIM/SIM_length_centered.jl")
include("Metric Models/Model Sampling/SIM/SIM_proportional.jl")
# SPF
include("Metric Models/Model Sampling/SPF/SPF_sampler.jl")
include("Metric Models/Model Sampling/SPF/SPF_dc_sampler.jl")
# Summaries
include("Metric Models/Model Sampling/Summaries/plot_recipes.jl")
include("Metric Models/Model Sampling/Summaries/misc_summaries.jl")

# Posterior Sampling
# ------------------
# Types
include("Metric Models/Posterior Sampling/Types/SIM_samplers.jl")
include("Metric Models/Posterior Sampling/Types/SIS_samplers.jl")
include("Metric Models/Posterior Sampling/Types/SPF_samplers.jl")
include("Metric Models/Posterior Sampling/Types/outputs.jl")
include("Metric Models/Posterior Sampling/Types/predictives.jl")
# Sampler files
include("Metric Models/Posterior Sampling/auxiliary_terms_eval.jl")
include("Metric Models/Posterior Sampling/cooccurrence_matrices.jl")
include("Metric Models/Posterior Sampling/informed_insertion_distribution.jl")
include("Metric Models/Posterior Sampling/iex_SIM.jl")
include("Metric Models/Posterior Sampling/iex_SIM_split_merge.jl")
include("Metric Models/Posterior Sampling/iex_SIM_proportional.jl")
include("Metric Models/Posterior Sampling/iex_SIM_dependent.jl")
include("Metric Models/Posterior Sampling/iex_SIM_with_kick.jl")
include("Metric Models/Posterior Sampling/iex_SIS.jl")
include("Metric Models/Posterior Sampling/iex_SPF.jl")
# Summaries
include("Metric Models/Posterior Sampling/Summaries/plot_recipes.jl")
include("Metric Models/Posterior Sampling/Summaries/misc_summaries.jl")
include("Metric Models/Posterior Sampling/Summaries/predictive_summaries.jl")

# iMCMC Moves (new structure)
include("Metric Models/iMCMC Moves/imcmc_move_type.jl")
include("Metric Models/iMCMC Moves/edit_allocation_move.jl")
include("Metric Models/iMCMC Moves/permutation_move.jl")
include("Metric Models/iMCMC Moves/insert_delete_move.jl")
include("Metric Models/iMCMC Moves/split_merge_move.jl")
include("Metric Models/iMCMC Moves/mixture_move.jl")
include("Metric Models/Samplers/model_sampler.jl")
include("Metric Models/Samplers/posterior_sampler.jl")

# Graphs Models
# =============
include("Graph Models/utils.jl")
include("Graph Models/CER/CER.jl")
include("Graph Models/CER/CER_posterior.jl")

include("Graph Models/SNF/SNF.jl")
include("Graph Models/SNF/MCMC Moves/move_type.jl")
include("Graph Models/SNF/MCMC Moves/gibbs_moves.jl")
# include("Graph Models/SNF/SNF_samplers.jl")
# include("Graph Models/SNF/SNF_multigraph.jl")
# include("Graph Models/SNF/SNF_multigraph_gibbs.jl")
include("Graph Models/SNF/Samplers/SNF_model_sampler.jl")
include("Graph Models/SNF/Samplers/SNF_posterior_sampler.jl")
# include("Graph Models/SNF/SNF_multigraph_gibbs_posterior.jl")

end


import Bridge: kernelr3!, R3!, target, auxiliary, constdiff, llikelihood, _b!, B!, σ!, b!

"""
    (t, obsinfo, mT, P, ITER, subsamples, xinit, pars, priorθ, priormom, updatescheme, outdir, printskip)

This is the main function call for doing either MCMC or SGD for landmark models
Backward ODEs used are in terms of the LMμ-parametrisation

## Arguments
- `t`:      time grid
- `obsinfo`: struct of type `ObsInfo`
- `mT`: vector of momenta at time T used for constructing guiding term
- `P`: target process
- `ITER`: number of iterations
- `subsamples`: vector of indices of iterations that are to be saved
- `xinit`: initial value on state at time 0 (`x0`)
- `pars`: tuning pars for updates (see `?Pars_ms()` and `?Pars_ahs()` for documentation)
- `priorθ`: prior on θ = (a,c,γ)
- `priormom`: prior on initial momenta
- `updatescheme`: vector specifying mcmc-updates
- `outdir` output directory for animation
- `printskip` skip every `printskip` iterations in writing output to console

## Returns:
- `Xsave`: saved iterations of all states at all times in tt_
- `parsave`: saved iterations of all parameter updates
- `initendstates_save`: saved iterations of initial and endstate
- `accinfo`: acceptance indicators of mcmc-updates
- `δ`: value of `(δpos, δmom)` af the final iteration of the algorithm (these are stepsize parameters for updating initial positions and momenta respectively)
- `ρ`: value of `ρinit`  at the final iteration of the algorithm
- `δa`: value of `δa` at the final iteration of the algorithm (step size for updating Hamiltonian kernel parameter `a`)
"""
function lm_mcmc(t, obsinfo, mT, P, ITER, subsamples, xinit, pars, priorθ, priormom, updatescheme, outdir, printskip)
    lt = length(t)
    StateW = PointF
    dwiener = dimwiener(P)
    nshapes = obsinfo.nshapes

    if P.n < 25
        chunksize = d*P.n
    elseif 25 <= P.n < 100
        chunksize = P.n
    else
        chunksize = P.n÷2
    end

    guidrec = [GuidRecursions(t,obsinfo)  for _ ∈ 1:nshapes]  # initialise guiding terms
    gramT_container = [gram_matrix(obsinfo.xobsT[k], P) for k ∈ 1:nshapes]  # gram matrices at observations at time T

    Paux = [auxiliary(P, State(obsinfo.xobsT[k],mT), gramT_container[k]) for k ∈ 1:nshapes] # auxiliary process for each shape
    Q = GuidedProposal(P,Paux,t,obsinfo.xobs0,obsinfo.xobsT,guidrec,nshapes,[mT for _ ∈ 1:nshapes])
    Q = update_guidrec!(Q, obsinfo)   # compute backwards recursion

    X = [initSamplePath(t, xinit) for _ ∈ 1:nshapes]
    W = [initSamplePath(t,  zeros(StateW, dwiener)) for _ ∈ 1:nshapes]
    for k ∈ 1:nshapes   sample!(W[k], Wiener{Vector{StateW}}())  end

    x = deepvec(xinit)
    q, p = split_state(xinit)
    qᵒ = deepcopy(q); pᵒ = deepcopy(q); ∇ = deepcopy(q); ∇ᵒ = deepcopy(q)

    # memory allocations, actual state at each iteration is (X, W, Q, x, ∇x) (x, ∇x are initial state and its gradient)
    Xᵒ = deepcopy(X)
    Wᵒ = initSamplePath(t,  zeros(StateW, dwiener))
    Wnew = initSamplePath(t,  zeros(StateW, dwiener))
    # sample guided proposal and compute loglikelihood (write into X)
    At = zeros(UncF, 2P.n, 2P.n)  # container for to compute a=σσ' for auxiliary process.
    ll, X = gp!(LeftRule(), X, xinit, W, Q, At; skip=sk)

    # setup containers for saving objects
    Xsave = typeof(zeros(length(t) * P.n * 2 * d * nshapes))[]
    push!(Xsave, convert_samplepath(X))
    parsave = Vector{Float64}[]
    push!(parsave, getpars(Q))
    initendstates_save = [extract_initial_and_endstate(0,X[1]) ]
    accinfo = DataFrame(fill([], length(updatescheme)+1), [updatescheme...,:iteration]) # columns are all update types + 1 column for iteration number
    acc = zeros(ncol(accinfo))

    δ = [pars.δpos, pars.δmom, pars.δsgd_mom]
    δa, δγ = pars.δa, pars.δγ
    ρ = pars.ρlowerbound
    adapt(i) = (i > 1.5*pars.adaptskip) & (mod(i,pars.adaptskip)==0)

    x0 = X[1].yy[1]
    dK = gramkernel(x0.q, P)
    if :rmrw_mom ∈ updatescheme ||  :rmmala_mom  ∈ updatescheme
        if :parameter in updatescheme
            @error "Cannot use :rmrw_mom or :rmmala_mom in case parameterupdating is also done (because this requries inv_dK te be recomputed at each iteration)"
        else
        inv_dK = inv(dK)
        end
    else
        inv_dK = 0.0 # not used then
    end

    for i ∈ 1:ITER
        k = 1
        for update ∈ updatescheme
            if update == :innov
                accinfo_, X, W, ll = update_path!(X, W, ll, Xᵒ, Wᵒ, Wnew, Q, ρ, At)
                if adapt(i)
                    ρ = adaptpcnstep(ρ, i, accinfo[!,update], pars.η, pars.adaptskip)
                end
            elseif update ∈ [:mala_mom, :tmala_mom, :rmmala_mom, :rmrw_mom]
                accinfo_, X, x, ∇, ll = update_initialstate!(X,Xᵒ,W,ll, x, qᵒ, pᵒ,∇, ∇ᵒ, Q, δ, update, priormom, (dK, inv_dK), At, chunksize)
                if adapt(i)
                    δ[2] = adapt_pospar_step(δ[2], i, accinfo[!,update], pars.η, pars.adaptskip)
                end
            elseif update ∈ [:mala_pos, :tmala_pos, :rmmala_pos]
                accinfo_ , X, x, ∇, ll= update_initialstate!(X,Xᵒ,W,ll, x, qᵒ, pᵒ,∇, ∇ᵒ, Q, δ, update, priormom, (dK, inv_dK), At, chunksize)
                if adapt(i)
                    δ[1] = adapt_pospar_step(δ[1], i, accinfo[!,update], pars.η, pars.adaptskip)
                end
            elseif update == :parameter
                Q, accinfo_, X, ll = update_pars!(X, ll, Xᵒ,W, Q, priorθ, obsinfo, δa, δγ, gramT_container, At)
                if adapt(i)
                    δa = adapt_pospar_step(δa, i, accinfo[!,update], pars.η, pars.adaptskip)
                end
            elseif update == :matching
                obsinfo, accinfo_, Q, X, ll = update_cyclicmatching(X, ll, obsinfo, Xᵒ, W, Q, At)
            elseif update == :sgd_mom
                accinfo_, X, x, ∇, ll = update_initialstate!(X,Xᵒ,W,ll, x, qᵒ, pᵒ,∇, ∇ᵒ, Q, δ, update, priormom, (dK, inv_dK), At)
            else @error "This update scheme has not been implemented"
            end
            acc[k] = accinfo_
            k += 1
        end
        acc[end] = i
        push!(accinfo, acc)

        # if mod(i,pars.adaptskip)==0 && i < 100  # adjust mT (the momenta at time T used in the construction of the guided proposal)
        #     mTvec = [X[k][lt][2].p  for k in 1:nshapes]     # extract momenta at time T for each shape
        #    Q = update_mT!(Q, mTvec, obsinfo)
        # end

        # save some of the results
        if i ∈ subsamples
            push!(Xsave, convert_samplepath(X))
        end
        push!(parsave, getpars(Q))
        push!(initendstates_save, extract_initial_and_endstate(i,X[1]))

        if mod(i,printskip) == 0
            #println();  println("iteration $i, ρ = $ρ,  δ = $δ,  δa = $δa")
            #print("Average acceptance rates over last ",  printskip, " iterations: ")
            if i >= pars.adaptskip
                ac = accinfo[end-pars.adaptskip+1:end, 1:end-1]
                println(round.([mean(x) for x in eachcol(ac)];digits=2))
            end
            #println("parameter a and γ equal: ", getpars(Q))
        end
    end
    Xsave, parsave, initendstates_save, accinfo, δ, ρ, δa
end

"""
    update_path!(X,Xᵒ,W,Wᵒ,Wnew,ll,x, Q, ρ)

Update bridges for all shapes using Crank-Nicholsen scheme with parameter `ρ` (only in case the method is mcmc).
Newly accepted bridges are written into `(X,W)`, loglikelihood on each segment is written into vector `ll`
All Wiener increments are always updated.

## Write into
- `X`: diffusion paths
- `W`: innovations
- `ll`: vector of loglikehoods (k-th element is for the k-th shape)
"""
function update_path!(X, W, ll, Xᵒ,Wᵒ, Wnew, Q, ρlowerbound, At)
    acc = Int64[]
    x0 = X[1].yy[1]
    ρ = rand(Uniform(ρlowerbound, 1.0))
    ρ_ = sqrt(1.0-ρ^2)
    # From current state (x,W) with loglikelihood ll, update to (x, Wᵒ)
    for k ∈ 1:Q.nshapes
        sample!(Wnew, Wiener{Vector{PointF}}())
        for i ∈ eachindex(W[1].yy)
            for j ∈ eachindex(W[1].yy[1])
                Wᵒ.yy[i][j] = ρ * W[k].yy[i][j] + ρ_ * Wnew.yy[i][j]
            end
        end
        llᵒ, Xᵒ[k] = gp!(LeftRule(), Xᵒ[k], x0, Wᵒ, Q, k, At; skip=sk)
        diff_ll = llᵒ - ll[k]
        if log(rand()) <= diff_ll
            for i ∈ eachindex(W[1].yy)
                X[k].yy[i] .= Xᵒ[k].yy[i]
                W[k].yy[i] .= Wᵒ.yy[i]
            end
            ll[k] = llᵒ
            push!(acc,1)
        else
            push!(acc,0)
        end
    end
    mean(acc), X, W, ll
end

"""
    slogρ!(q, p, Q, W, X, priormom,llout)

Main use of this function is to get gradient information of the loglikelihood with respect ot the initial state.

## Arguments
- `q` and `p` make up the initial state (which could be obtained from `x0 = merge_state(q,p)`)
- `Q`: GuidePropsoal!
- `W`: vector of innovations
- `X`: sample path
- `priormom`: prior on the initial momenta
- `llout`: vector where loglikelihoods for each shape are written into

## Writes into
- The guided proposal is written into `X`.
- The computed loglikelihood for each shape is written into `llout`

## Returns
- loglikelihood (summed over all shapes) + the logpriordensity of the initial momenta
"""
function slogρ!(q,p, Q, W, X, priormom, llout, At)
    lltemp, X = gp!(LeftRule(), X, q, p, W, Q, At; skip=sk)   # writes into X
    llout .= ForwardDiff.value.(lltemp)
    sum(lltemp) + logpdf(priormom, p)
end

"""
    slogρ_pos!(p, Q, W, X,priormom, llout) = (q) -> slogρ!(q, p , Q, W, X, priormom, llout)
"""
slogρ_pos!(p, Q, W, X,priormom, llout, At) = (q) -> slogρ!(q, p , Q, W, X, priormom, llout, At)

"""
    slogρ_mom!(q, Q, W, X,priormom, llout) = (p) -> slogρ!(q, p , Q, W, X, priormom, llout)
"""
slogρ_mom!(q, Q, W, X,priormom, llout, At) = (p) -> slogρ!(q, p , Q, W, X, priormom, llout, At)

"""
    update_initialstate!(X,Xᵒ,W,ll,x,xᵒ,∇, ∇ᵒ, Q::GuidedProposal, δ, update,priormom, (dK, inv_dK), At)

## Arguments
- `X`:  current iterate of vector of sample paths
- `Xᵒ`: vector of sample paths to write proposal into
- `W`:  current vector of Wiener increments
- `ll`: current value of loglikelihood
- `x`, `xᵒ`, `∇x`, `∇xᵒ`: allocated vectors for initial state and its gradient (both actual and proposed values)
- `sampler`: either sgd (not checked yet) or mcmc
- `Q`: GuidedProposal
- `δ`: vector with MALA stepsize for initial state positions (δ[1]) and initial state momenta (δ[2])
- `update`:  can be `:mala_pos`, `:mala_mom`, :rmmala_pos`, `:rmmala_mom`, ``:rmrw_mom`
- `priormom`: prior on initial momenta
- (dK, inv_dK): gramkernel at x0.q and its inverse

## Writes into / modifies
- `X`: landmarks paths
- `x`: initial state of the landmark paths
- ``∇`: gradient of loglikelihood with respect to initial state `x`
- `ll`: a vector with as k-th element the loglikelihood for the k-th shape

## Returns
-  0/1 indicator (reject/accept)
"""
function update_initialstate!(X,Xᵒ,W,ll, x, qᵒ, pᵒ,∇, ∇ᵒ,
                 Q::GuidedProposal, δ, update, priormom, (dK, inv_dK), At, chunksize)

    P = Q.target;   n = P.n
    x0 = deepvec2state(x)
    q, p = split_state(x0)
    dn = d * n

    if update ==:sgd_mom
        # @warn "Option :sgd_mom has not been checked/tested so far; don't use as yet."
        for k ∈ 1:Q.nshapes
            sample!(W[k], Wiener{Vector{PointF}}())
        end
        u = slogρ_mom!(q, Q, W, X, priormom,ll, At)
        cfg = ForwardDiff.GradientConfig(u, p, ForwardDiff.Chunk{chunksize}()) # d*P.n is maximal
        ForwardDiff.gradient!(∇, u, p, cfg) # X gets overwritten but does not change
        pᵒ = p .- δ[3] * ∇
        x0ᵒ = merge_state(q, pᵒ)
        ll, X = gp!(LeftRule(), X, x0ᵒ, W, Q, At)
        x .= deepvec(x0ᵒ)
        accepted = 1.0
    else
        accinit = 0.0
        llᵒ = copy(ll)
        if update in [:mala_pos, :tmala_pos, :rmmala_pos]
            u = slogρ_pos!(p, Q, W, X, priormom,ll, At)
            uᵒ = slogρ_pos!(p, Q, W, Xᵒ, priormom,llᵒ, At)
            cfg = ForwardDiff.GradientConfig(u, q, ForwardDiff.Chunk{chunksize}()) # d*P.n is maximal
            ForwardDiff.gradient!(∇, u, q, cfg) # X and ll get overwritten but do not change
            stepsize = sample(δ[1])
        end
        if update in [:mala_mom, :tmala_mom, :rmmala_mom]
            u = slogρ_mom!(q, Q, W, X, priormom,ll, At)
            uᵒ = slogρ_mom!(q, Q, W, Xᵒ, priormom,llᵒ, At)
            cfg = ForwardDiff.GradientConfig(u, p, ForwardDiff.Chunk{chunksize}()) # d*P.n is maximal
            ForwardDiff.gradient!(∇, u, p, cfg) # X and ll get overwritten but do not change
            #stepsize = δ[2]
            stepsize = sample(δ[2])
        end
        if update == :rmrw_mom
            u = slogρ_mom!(q, Q, W, X, priormom,ll, At)
            uᵒ = slogρ_mom!(q, Q, W, Xᵒ, priormom,llᵒ, At)
            #stepsize = δ[2]
            stepsize = sample(δ[2])
        end

        ## compute proposal and log(acceptance ratio)
        if update == :mala_pos
            ndistr = MvNormal(dn,sqrt(stepsize))
            qᵒ .= q .+ .5 * stepsize *  ∇ .+ rand(ndistr)
            cfgᵒ = ForwardDiff.GradientConfig(uᵒ, qᵒ, ForwardDiff.Chunk{chunksize}())
            ForwardDiff.gradient!(∇ᵒ, uᵒ, qᵒ, cfgᵒ)
            accinit = sum(llᵒ) - sum(ll) -
                      logpdf(ndistr, qᵒ - q - .5*stepsize * ∇) +
                      logpdf(ndistr, q - qᵒ - .5*stepsize * ∇ᵒ)
        elseif update == :tmala_pos
              ndistr = MvNormal(dn,sqrt(stepsize))
              qᵒ .= q .+ .5 * tame(stepsize *  ∇) .+ rand(ndistr)
              cfgᵒ = ForwardDiff.GradientConfig(uᵒ, qᵒ, ForwardDiff.Chunk{chunksize}())
              ForwardDiff.gradient!(∇ᵒ, uᵒ, qᵒ, cfgᵒ)
              accinit = sum(llᵒ) - sum(ll) -
                        logpdf(ndistr, qᵒ - q - .5 * tame(stepsize * ∇)) +
                        logpdf(ndistr, q - qᵒ - .5 * tame(stepsize * ∇ᵒ))
        elseif update == :mala_mom
            ndistr = MvNormal(dn, sqrt(stepsize))
            pᵒ .= p .+ .5 * stepsize * ∇ .+ rand(ndistr)
            cfgᵒ = ForwardDiff.GradientConfig(uᵒ, pᵒ, ForwardDiff.Chunk{chunksize}()) # d*P.n is maximal
            ForwardDiff.gradient!(∇ᵒ, uᵒ, pᵒ, cfgᵒ)
            accinit = sum(llᵒ) - sum(ll) -
                      logpdf(ndistr, pᵒ - p - .5 * stepsize * ∇) +
                      logpdf(ndistr, p - pᵒ - .5 * stepsize * ∇ᵒ)
      elseif update == :tmala_mom
          ndistr = MvNormal(dn, sqrt(stepsize))
          pᵒ .= p .+ .5 * tame(stepsize * ∇) .+ rand(ndistr)
          cfgᵒ = ForwardDiff.GradientConfig(uᵒ, pᵒ, ForwardDiff.Chunk{chunksize}())
          ForwardDiff.gradient!(∇ᵒ, uᵒ, pᵒ, cfgᵒ)
          accinit = sum(llᵒ) - sum(ll) -
                    logpdf(ndistr, pᵒ - p - .5 * tame(stepsize * ∇)) +
                    logpdf(ndistr, p - pᵒ - .5 * tame(stepsize * ∇ᵒ))
        elseif update == :rmmala_pos
            #dK = gramkernel(x0.q, P)
            dK = gramkernel(reinterpret(PointF,q), P)
            ndistr = MvNormal(zeros(d*n),stepsize*dK)
            qᵒ .= q .+ .5 * dK * tame(stepsize * ∇) .+ rand(ndistr)
            cfgᵒ = ForwardDiff.GradientConfig(uᵒ, qᵒ, ForwardDiff.Chunk{chunksize}())
            ForwardDiff.gradient!(∇ᵒ, uᵒ, qᵒ, cfgᵒ)
            dKᵒ = gramkernel(reinterpret(PointF,qᵒ), P)
            ndistrᵒ = MvNormal(zeros(d*n),stepsize*dKᵒ)
            accinit = sum(llᵒ) - sum(ll) -
                     logpdf(ndistr, qᵒ - q - .5 * dK * tame(stepsize * ∇)) +
                     logpdf(ndistrᵒ, q - qᵒ - .5 * dKᵒ * tame(stepsize * ∇ᵒ))
        elseif update == :rmmala_mom
             ndistr = MvNormal(stepsize*inv_dK)
             pᵒ .= p .+ .5 * stepsize * inv_dK * ∇ .+  rand(ndistr)
             cfgᵒ = ForwardDiff.GradientConfig(uᵒ, pᵒ, ForwardDiff.Chunk{chunksize}())
             ForwardDiff.gradient!(∇ᵒ, uᵒ, pᵒ, cfgᵒ) # Xᵒ gets overwritten but does not change
             accinit = sum(llᵒ) - sum(ll) -
                       logpdf(ndistr, pᵒ - p - .5*stepsize * inv_dK * ∇) +
                       logpdf(ndistr, p - pᵒ - .5*stepsize * inv_dK * ∇ᵒ)
       elseif update == :rmrw_mom
            ndistr = MvNormal(zeros(d*n),stepsize*inv_dK)
            pᵒ .= p  .+ rand(ndistr)
            uᵒ(pᵒ)  # writes into llᵒ, don't need gradient info here
            accinit = sum(llᵒ) - sum(ll)   # proposal is symmetric
        end

        # MH acceptance decision
        if log(rand()) <= accinit
            X = copypaths!(X,Xᵒ)
            ll .= llᵒ
            if update in [:mala_pos, :rmmala_pos]
                x .= deepvec(merge_state(qᵒ, p))
            end
            if update in [:mala_mom, :rmmala_mom, :rmrw_mom]
                x .= deepvec(merge_state(q, pᵒ))
            end
            accepted = 1.0
        else
            accepted = 0.0
        end
    end
    accepted, X, x, ∇, ll
end

"""
    trun(x,h) = x/max(1.0,0.5*h*norm(x))

useful function for truncated version of mala
"""
trun(x,h) = x/max(1.0,0.5*h*norm(x))

"""
    tame(x) = x/(I + norm(x))

Useful for tamed version of Langevin algorithm
"""
tame(x) = x/(I + norm(x))

"""
    update_pars!(X, Xᵒ,W, Q, x, ll, priorθ, covθprop, obsinfo, At)

For fixed Wiener increments and initial state, update parameters by random-walk-MH.

## Writes into / modifies
`X`: landmark paths under landmarks process with parameter θ
`ll`: vector of loglikelihoods (one for each shape)

## Returns
`Q, accept`: where `accept` is 1/0 according to accept/reject in the MH-step.
"""
function update_pars!(X, ll, Xᵒ, W, Q , priorθ,  obsinfo, δa, δγ, gramT_container, At)
    θ = getpars(Q)
    a = θ[1];  γ = θ[2]
    aᵒ = a * exp(δa*randn())
    γᵒ = γ * exp(δγ*randn())
    θᵒ = [aᵒ, γᵒ]
    Qᵒ = adjust_to_newpars(Q, θᵒ, obsinfo, gramT_container)
    llᵒ, Xᵒ = gp!(LeftRule(), Xᵒ, X[1].yy[1], W, Qᵒ, At; skip=sk)
    A = sum(llᵒ) - sum(ll) +
        logpdf(priorθ, θᵒ) - logpdf(priorθ, θ) +
        log(γᵒ) - log(γ) + log(aᵒ) - log(a)
    if log(rand()) <= A
        ll .= llᵒ
        X = copypaths!(X,Xᵒ)
        Q = Qᵒ
        #Q = deepcopy(Qᵒ)
        accept = 1.0
    else
        accept = 0.0
    end
    Q, accept, X, ll
end

"""
    function copypaths!(X,Xᵒ)

Write Xᵒ into X
"""
function copypaths!(X,Xᵒ)
    for k ∈ eachindex(X), i ∈ eachindex(X[1].yy)
        X[k].yy[i] .= Xᵒ[k].yy[i]
    end
    X
end

######### adaptation of mcmc tuning parameters #################

"""
    adapt_pospar_step!(δ, n, accinfo, η, update; adaptskip = 15, targetaccept=0.5)

Adapt the step size for updates, where tuning par has to be positive.
The adaptation is multiplication/divsion by exp(η(n)).
"""
function adapt_pospar_step(δ, n, accinfo, η, adaptskip; targetaccept=0.5)
    recent_mean = mean(accinfo[end-adaptskip+1:end])
    if recent_mean > targetaccept
        return δ * exp(η(n))
    else
        return δ * exp(-η(n))
    end
end

"""
    adaptpcnstep(n, accpcn, ρ, nshapes, η; adaptskip = 15, targetaccept=0.5)

Adjust pcN-parameter `ρ` adaptive every `adaptskip` steps.
For that `ρ` is first tranformed to (-∞, ∞), then updated, then mapped back to (0,1)
"""
function adaptpcnstep(ρ, n, accinfo, η, adaptskip; targetaccept=0.5)
    recentmean = mean(accinfo[end-adaptskip+1:end])
    if recentmean > targetaccept
        out =  logistic(logit(ρ) - η(n))
    else
        out =  logistic(logit(ρ) + η(n))
    end
    out
end

"""
    show_updates()

Convenience function to show the available implemted updates.
"""
function show_updates()
    println("Available implemented update steps are:")
    println(":innov")
    println(":sgd_mom")
    println(":mala_mom, :tmala_mom, :rmmala_mom, :rmrw_mom")
    println(":mala_pos, :tmala_pos, :rmmala_pos")
    println(":parameter, :matching")
end


function visualize{T, D}(style::Style, img::Texture{T, D, 2}, data::Dict{Symbol, Any})
  kernel = data[:kernel]
  w, h  = img.dims
  texparams = [
     (GL_TEXTURE_MIN_FILTER, GL_NEAREST),
    (GL_TEXTURE_MAG_FILTER, GL_NEAREST),
    (GL_TEXTURE_WRAP_S,  GL_CLAMP_TO_EDGE),
    (GL_TEXTURE_WRAP_T,  GL_CLAMP_TO_EDGE)
  ]

  v, uv, indexes = genquad(0f0, 0f0, w, h)
  if typeof(kernel) <: Real
    filterkernel = float32(kernel)
  elseif eltype(kernel) <: Union(AbstractArray, Real)
    filterkernel = Texture(kernel, parameters=texparams)
  end

  data = @compat Dict(
    :vertex           => GLBuffer(v, 2),
    :index            => indexbuffer(indexes),
    :uv               => GLBuffer(uv, 2),
    :image            => img,
    :normrange        => data[:normrange],
    :filterkernel     => filterkernel,
    :projectionview   => data[:screen].orthographiccam.projectionview,
    :model            => data[:model]
  )

  fragdatalocation = [(0, "fragment_color"),(1, "fragment_groupid")]
  textureshader    = TemplateProgram(joinpath(shaderdir, "uv_vert.vert"), joinpath(shaderdir, "texture.frag"), attributes=data, fragdatalocation=fragdatalocation)

  obj = RenderObject(data, textureshader)

  prerender!(obj, glDisable, GL_DEPTH_TEST, enabletransparency, glDisable, GL_CULL_FACE)
  postrender!(obj, render, obj.vertexarray)
  obj
end