text:
using Test
using NCTiles, NCDatasets, MeshArrays
include("testutils.jl")

println("Testing...")

testvars = maketestdata()

ncfilarray2d,tmp = mktemp()
ncfilarray3d,tmp = mktemp()
ncfilbin2d,tmp = mktemp()
ncfilbin3d,tmp = mktemp()
ncfilnc2d,tmp = mktemp()
ncfilnc3d,tmp = mktemp()
ncfiltile2d,tmp = mktemp()

tempfiles = vcat(testvars["fnames2d"], testvars["fnames3d"],
                    testvars["tile_ex"]["fnamestile2d"],
                    ncfilarray2d, ncfilarray3d,
                    ncfilbin2d, ncfilbin3d,
                    ncfilnc2d, ncfilnc3d)

@testset "NCtiles Tests" begin

    # Test in memory write
    @testset "In Memory Data" begin
        fld2d = NCvar("data2d",testvars["units"],[testvars["dims"][1:2]; testvars["dims"][4]],testvars["data2d"],Dict("long_name" =>testvars["longname2d"]),NCDatasets)
        fld3d = NCvar("data3d",testvars["units"],testvars["dims"],testvars["data3d"],Dict("long_name" =>testvars["longname3d"]),NCDatasets)

        writetestfile(ncfilarray2d,fld2d,NCDatasets)
        writetestfile(ncfilarray3d,fld3d,NCDatasets)

        @test testfile(ncfilarray2d,fld2d)
        @test testfile(ncfilarray3d,fld3d)
    end

    # Test BinData write
    @testset "Binary File Data" begin
        bdatafld2d = BinData(testvars["fnames2d"],Float32,(testvars["dims"][1].dims,testvars["dims"][2].dims))
        bdatafld3d = BinData(testvars["fnames3d"],Float32,(testvars["dims"][1].dims,testvars["dims"][2].dims,testvars["dims"][3].dims))

        bfld2d = NCvar("data2d",testvars["units"],[testvars["dims"][1:2]; testvars["dims"][4]],bdatafld2d,Dict("long_name" =>testvars["longname2d"]),NCDatasets)
        bfld3d = NCvar("data3d",testvars["units"],testvars["dims"],bdatafld3d,Dict("long_name" =>testvars["longname3d"]),NCDatasets)

        writetestfile(ncfilbin2d,bfld2d,NCDatasets)
        writetestfile(ncfilbin3d,bfld3d,NCDatasets)

        @test testfile(ncfilbin2d,bfld2d)
        @test testfile(ncfilbin3d,bfld3d)
    end

    #Test NCData write
    @testset "NetCDF File Data" begin
        ncdatafld2d = NCData(ncfilarray2d,"data2d",NCDatasets,Float32)
        ncdatafld3d = NCData(ncfilarray3d,"data3d",NCDatasets,Float32)

        ncfld2d = NCvar("data2d",testvars["units"],[testvars["dims"][1:2]; testvars["dims"][4]],ncdatafld2d,Dict("long_name" =>testvars["longname2d"]),NCDatasets)
        ncfld3d = NCvar("data3d",testvars["units"],testvars["dims"],ncdatafld3d,Dict("long_name" =>testvars["longname3d"]),NCDatasets)

        writetestfile(ncfilnc2d,ncfld2d,NCDatasets)
        writetestfile(ncfilnc3d,ncfld3d,NCDatasets)

        @test testfile(ncfilnc2d,ncfld2d)
        @test testfile(ncfilnc3d,ncfld3d)
    end

    @testset "Tile Data" begin
        tilesize = testvars["tile_ex"]["tilesize"]
        dims = testvars["tile_ex"]["dims"]
        grid = testvars["tile_ex"]["grid"]
        gridvars = GridLoad(grid)
        land = gridvars["hFacC"]
        for f in land.fIndex
            for d in 1:size(land,2)
                land[f,d][land[f,d].==0] .= NaN
                land[f,d][land[f,d].>0] .= 1
            end
        end
        tiledatafld2d = BinData(testvars["tile_ex"]["fnamestile2d"],Float32,Tuple(grid.ioSize))
        tilfld2d = TileData(tiledatafld2d,testvars["tile_ex"]["tilesize"],grid)
        tillat = TileData(gridvars["YC"],tilfld2d.tileinfo,tilfld2d.tilesize,tilfld2d.precision,tilfld2d.numtiles)
        tillon = TileData(gridvars["XC"],tilfld2d.tileinfo,tilfld2d.tilesize,tilfld2d.precision,tilfld2d.numtiles)
        tilarea = TileData(gridvars["RAC"],tilfld2d.tileinfo,tilfld2d.tilesize,tilfld2d.precision,tilfld2d.numtiles)
        tilland = TileData(land,tilfld2d.tileinfo,tilfld2d.tilesize,tilfld2d.precision,tilfld2d.numtiles)
        thic = gridvars["RC"][:,1]
        flds = Dict(["data2d" => NCvar("data2d",testvars["units"],[dims[1:2]; dims[4]],tilfld2d,Dict(),NCDatasets),
                    "lon" => NCvar("lon","degrees_east",dims[1:2],tillon,Dict("long_name" => "longitude"),NCDatasets),
                    "lat" => NCvar("lat","degrees_north",dims[1:2],tillat,Dict("long_name" => "latitude"),NCDatasets),
                    "area" => NCvar("area","m^2",dims[1:2],tilarea,Dict(["long_name" => "grid cell area", "standard_name" => "cell_area"]),NCDatasets),
                    "land" => NCvar("land","1",dims[1:3],tilland,Dict(["long_name" => "land mask", "standard_name" => "land_binary_mask"]),NCDatasets),
                    "thic" => NCvar("thic","m",dims[3],thic,Dict("standard_name" => "cell_thickness"),NCDatasets)
        ])
        writetestfile(ncfiltile2d,flds,NCDatasets)
        savenames = joinpath.(ncfiltile2d*".".*lpad.(string.(1:tilfld2d.numtiles),4,"0").*".nc")

        @test all([testfile(fname,flds[fld]) for fname in savenames for fld in keys(flds)])
    end

end

# Delete temporary files
for fil in tempfiles
    rm(fil)
end


text:
"""
    NSVerify(optimizer, m::Float64)

NSVerify finds counter examples using mixed integer linear programming.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: hyperrectangle or hpolytope
3. Output: PolytopeComplement

# Return
`CounterExampleResult`

# Method
MILP encoding (using `m`). No presolve.
Default `optimizer` is `GLPKSolverMIP()`. Default `m` is `1000.0` (should be large enough to avoid approximation error).

# Property
Sound and complete.

# Reference
[A. Lomuscio and L. Maganti,
"An Approach to Reachability Analysis for Feed-Forward Relu Neural Networks,"
*ArXiv Preprint ArXiv:1706.07351*, 2017.](https://arxiv.org/abs/1706.07351)
"""
@with_kw struct NSVerify <: Solver
    optimizer = GLPK.Optimizer
    m = nothing
end

function solve(solver::NSVerify, problem::Problem)
    # Set up and solve the problem
    model = Model(solver)
    z = init_vars(model, problem.network, :z, with_input=true)
    δ = init_vars(model, problem.network, :δ, binary=true)
    # Set M automatically if not already set.
    if isnothing(solver.m)
        model[:M] = set_automatic_M(problem)
    else
        @warn "M should be chosen carefully. An M which is too small will cause
        the problem to be solved incorrectly. Not setting the `m` keyword, or setting
        it equal to `nothing` will cause a safe value to be calculated automatically.
        E.g. NSVerify()." maxlog = 1
        model[:M] = solver.m
    end

    add_set_constraint!(model, problem.input, first(z))
    add_complementary_set_constraint!(model, problem.output, last(z))
    encode_network!(model, problem.network, MixedIntegerLP())
    feasibility_problem!(model)
    optimize!(model)

    if termination_status(model) == OPTIMAL
        return CounterExampleResult(:violated, value(first(z)))
    elseif termination_status(model) == INFEASIBLE
        return CounterExampleResult(:holds)
    else
        return CounterExampleResult(:unknown)
    end
end

function set_automatic_M(problem)
    # Compute the largest pre-activation bound absolute value.
    # M must be larger than any value a variable in the problem
    # can take.
    bounds = get_bounds(problem, false)
    M = maximum(abs, Iterators.flatten(abs.(hr.center) + hr.radius for hr in bounds))
end

ext: jl


text:
"""
    emrf(flow_graph, source, target, capacity_matrix, flow_algorithm, routes=0)

Compute the maximum multiroute flow (for any number of `route`s)
between `source` and `target` in `flow_graph` via flow algorithm `flow_algorithm`.

If a number of routes is given, return the value of the multiroute flow as
well as the final flow matrix, along with a multiroute cut if the
Boykov-Kolmogorov max-flow algorithm is used as a subroutine.
Otherwise, return the vector of breaking points of the parametric
multiroute flow function.

### References
- [Extended Multiroute Flow algorithm](http://dx.doi.org/10.1016/j.disopt.2016.05.002)
"""
# EMRF (Extended Multiroute Flow) algorithms
function emrf(
    flow_graph::AbstractGraph,                   # the input graph
    source::Integer,                       # the source vertex
    target::Integer,                       # the target vertex
    capacity_matrix::AbstractMatrix,  # edge flow capacities
    flow_algorithm::AbstractFlowAlgorithm, # keyword argument for algorithm
    routes::Real = 0
    )
    breakingpoints = breakingPoints(flow_graph, source, target, capacity_matrix)
    if routes > 0
        x, f = intersection(breakingpoints, routes)
        return maximum_flow(flow_graph, source, target, capacity_matrix,
        algorithm = flow_algorithm, restriction = x)
    end
    return breakingpoints
end

@doc_str """
    auxiliaryPoints(flow_graph, source, target, capacity_matrix)

Output a set of (point, slope) that compose the restricted max-flow function
of `flow_graph` from `source to `target` using capacities in `capacity_matrix`.

### Performance
One point by possible slope is enough (hence \mathcal{O}(λ×max_flow) complexity).
"""
function auxiliaryPoints end
@traitfn function auxiliaryPoints(
    flow_graph::::IsDirected,                   # the input graph
    source::Integer,                           # the source vertex
    target::Integer,                           # the target vertex
    capacity_matrix::AbstractMatrix   # edge flow capacities
    )
    # Problem descriptors
    λ = maximum_flow(flow_graph, source, target)[1] # Connectivity
    n = nv(flow_graph) # number of vertices
    r1, r2 = minmaxCapacity(capacity_matrix) # restriction left (1) and right (2)
    auxpoints = fill((0., 0.), λ + 1)

    # Initialisation of left side (1)
    f1, F1, cut1 = maximum_flow(flow_graph, source, target, capacity_matrix,
    algorithm = BoykovKolmogorovAlgorithm(), restriction = r1)
    s1 = slope(flow_graph, capacity_matrix, cut1, r1) # left slope
    auxpoints[λ + 1 - s1] = (r1, f1) # Add left initial auxiliary point

    # Initialisation of right side (2)
    f2, F2, cut2 = maximum_flow(flow_graph, source, target, capacity_matrix,
    algorithm = BoykovKolmogorovAlgorithm(), restriction = r2)
    s2 = slope(flow_graph, capacity_matrix, cut2, r2) # right slope
    auxpoints[λ + 1 - s2] = (r2, f2) # Add right initial auxiliary point

    # Loop if the slopes are distinct by at least 2
    if s1 > s2 + 1
        queue = [((f1, s1, r1), (f2, s2, r2))]

        while !isempty(queue)
            # Computes an intersection (middle) with a new associated slope
            (f1, s1, r1), (f2, s2, r2) = pop!(queue)
            r, expectedflow = intersection(r1, f1, s1, r2, f2, s2)
            f, F, cut = maximum_flow(flow_graph, source, target, capacity_matrix,
            algorithm = BoykovKolmogorovAlgorithm(), restriction = r)
            s = slope(flow_graph, capacity_matrix, max.(cut, 1), r) # current slope
            auxpoints[λ + 1 - s] = (r, f)
            # If the flow at the intersection (middle) is as expected
            if expectedflow ≉ f # approximatively not equal (enforced by floating precision)
                # if the slope difference between (middle) and left is at least 2
                # push (left),(middle)
                if s1 > s + 1 && !approximately_equal((r2, f2), (r, f))
                    q = (f1, s1, r1), (f, s, r)
                    push!(queue, q)
                end
                # if the slope difference between (middle) and right is at least 2
                # push (middle),(right)
                if s > s2 + 1 && !approximately_equal((r1, f1), (r, f))
                    q = (f, s, r), (f2, s2, r2)
                    push!(queue, q)
                end
            end
        end
    end
    return auxpoints
end

"""
    breakingPoints(flow_graph::::IsDirected, source, target, capacity_matrix)

Calculates the breaking of the restricted max-flow from a set of auxiliary points
for `flow_graph` from `source to `target` using capacities in `capacity_matrix`.
"""
function breakingPoints end
@traitfn function breakingPoints(
    flow_graph::::IsDirected,                   # the input graph
    source::Integer,                           # the source vertex
    target::Integer,                           # the target vertex
    capacity_matrix::AbstractMatrix   # edge flow capacities
    )
    auxpoints = auxiliaryPoints(flow_graph, source, target, capacity_matrix)
    λ = length(auxpoints) - 1
    left_index = 1
    T = eltype(capacity_matrix)
    breakingpoints = Vector{Tuple{T,T,Int}}()

    for (id, point) in enumerate(auxpoints)
        if id == 1
            push!(breakingpoints, (0., 0., λ))
        else
            pleft = breakingpoints[left_index]
            if point[1] != 0
                x, y = intersection(pleft[1], pleft[2], pleft[3],
                point[1], point[2], λ + 1 - id)
                push!(breakingpoints, (x, y, λ + 1 - id))
                left_index += 1
            end
        end
    end
    return breakingpoints
end

"""
    minmaxCapacity(capacity_matrix)

Return the nonzero min and max function of `capacity_matrix`.

Note: this is more efficient than maximum() / minimum() / extrema()
since we have to ignore zero values.
"""

# Function to get the nonzero min and max function of a Matrix
# note: this is more efficient than maximum() / minimum() / extrema()
#       since we have to ignore zero values.
function minmaxCapacity(
    capacity_matrix::AbstractMatrix    # edge flow capacities
    )
    T = eltype(capacity_matrix)
    cmin, cmax = typemax(T), typemin(T)
    for c in capacity_matrix
        if c > zero(T)
            cmin = min(cmin, c)
        end
        cmax = max(cmax, c)
    end
    return cmin, cmax
end

"""
    slope(flow_graph, capacity_matrix, cut, restriction)

Return the slope of `flow_graph` using capacities in `capacity_matrix` and
a cut vector `cut`. The slope is initialized at 0 and is incremented for
each edge whose capacity does not exceed `restriction`.
"""
function slope end
# Function to get the slope of the restricted flow
@traitfn function slope(
    flow_graph::::IsDirected,                   # the input graph
    capacity_matrix::AbstractMatrix,  # edge flow capacities
    cut::Vector,                      # cut information for vertices
    restriction::Number                 # value of the restriction
    )
    slope = 0
    for e in edges(flow_graph)
        ## Chain comparison to wether an edge cross the cut from the source side of
        # the cut to the target side of the cut. Then the edge is selected iff the
        # capacity of the edge is larger then the restriction argument.
        # cut[dst(e)] == 2 > cut[src(e)] is equivalent to
        # cut[dst(e)] == 2 && 2 > cut[src(e)]
        # Description of chain comparisons can be found at https://goo.gl/IJpCqe
        if cut[dst(e)] == 2 > cut[src(e)] &&
            capacity_matrix[src(e), dst(e)] > restriction
            slope += 1
        end
    end
    return slope
end

"""
        intersection(x1, y1, a1, x2, y2, a2)

Return the intersection of two lines defined by `x` and `y` with slopes `a`.
2) A set of segments and a linear function of slope k passing by the origin.
Requires argument:
1) - x1, y1, a1, x2, y2, a2::T<:AbstractFloat # Coordinates/slopes
2) - points::Vector{Tuple{T, T, Int}}         # vector of points with T<:AbstractFloat
- k::R<:Real                             # number of routes (slope of the line)
"""
function intersection(
    x1::T,          # x coordinate of point 1
    y1::T,          # y coordinate of point 1
    a1::Integer,        # slope passing by point 1
    x2::T,          # x coordinate of point 2
    y2::T,          # y coordinate of point 2
    a2::R           # slope passing by point 2
    ) where T<:AbstractFloat where R<:Real

    (a1 == a2) && return -1., -1. # result will be ignored in other intersection method
    b1 = y1 - a1 * x1
    b2 = y2 - a2 * x2
    x = (b2 - b1) / (a1 - a2)
    y = a1 * x + b1
    return x, y
end

"""
    intersection(points, k)

Return the intersection of a set of line segments and a line of slope `k`
passing by the origin. Segments are defined as a triple (x, y, slope).
"""
function intersection(
    points::Vector{Tuple{T,T,I}},  # vector of breaking points
    k::R                               # number of routes (slope of the line)
    ) where T<:AbstractFloat where I<:Integer where R<:Real
    λ = points[1][1] # Connectivity

    # Loop over the segments (pair of breaking points)
    for (id, p) in enumerate(points[1:(end - 1)])
        if id == 1
            k ≈ λ && return points[2]
        else
            x, y = intersection(p[1], p[2], p[3], 0., 0., k)
            (p[1] ≤ x ≤ points[id + 1][1]) && return x, y
        end
    end
    p = points[end]
    return intersection(p[1], p[2], p[3], 0., 0., k)
end

"""
    approximately_equal(a, b)

Return true if each element in the tuple is approximately equal to its counterpart.

### Implementation Notes:
This is a separate function because we don't want to hijack isapprox for tuples.
"""
approximately_equal(a::Tuple{T,T}, b::Tuple{T,T}) where T <: AbstractFloat =
    a[1] ≈ b[1] && a[2] ≈ b[2]


text:
using Gadfly

set_default_plot_size(6inch, 6inch)

plot(x=rand(100), y=rand(100), color=rand(1:8, 100),
     Geom.point)


text:
# This is the adding problem from: Le, Q. V., Jaitly, N., & Hinton,
# G. E. (2015). A Simple Way to Initialize Recurrent Networks of
# Rectified Linear Units. arXiv preprint arXiv:1504.00941.
# Usage: julia adding.jl [opts], use --help for a full list of opts.

module Adding
using Knet, AutoGrad, ArgParse, JLD

function main(args=ARGS)
    info("Adding problem from Le et al. 2015.")
    isa(args, AbstractString) && (args=split(args))
    opts = parse_commandline(args) # TODO: get option processing back in here
    println(opts)
    opts["seed"] > 0 && setseed(opts["seed"])
    global data = Data(opts["length"], opts["batchsize"], opts["epochsize"])
    if opts["load"] != nothing
        global net = load(opts["load"],"net")
    else
        # global net = S2C(p1, p2; rnn=eval(parse(opts["nettype"])), hidden=opts["hidden"], winit=Gaussian(0,opts["winit"]), fbias=opts["fbias"])
        global net = compile(:adding; rnn=symbol(opts["nettype"]), hidden=opts["hidden"], winit=Gaussian(0,opts["winit"]), fbias=opts["fbias"])
    end
    # TODO: can pass opts in directly to compile as hash if keys were symbols: parse_args has an as_symbols option.
    # TODO: winit should probably be Xavier, at least give the user the option
    setp(net; lr=opts["lrate"])
    if opts["nosharing"]
        setp(net, :forwoverwrite, false)
        setp(net, :backoverwrite, false)
    end
    mse = 0; l=[0f0,0f0]; m=[0f0,0f0]
    for epoch=1:opts["epochs"]
        train(net, data, quadloss; gclip=opts["gclip"], losscnt=fill!(l,0), maxnorm=fill!(m,0))
        mse = 2*l[1]/l[2]
        println(tuple(epoch*data.epochsize,mse,m...))
        opts["gcheck"] > 0 && gradcheck(net,f->gradloss(f,data,quadloss;grad=true),f->gradloss(f,data,quadloss);gcheck=opts["gcheck"])
        flush(STDOUT)
    end
    return (mse, m...)
end

@knet function adding(x; rnn=nothing, hidden=0, o...)
    h = rnn(x; o..., out=hidden)
    if predict
        return wb(h; o..., out=1)
    end
end

function train(f, data, loss; gclip=0, losscnt=nothing, maxnorm=nothing)
    reset!(f)
    for (x,ygold) in data
        if ygold == nothing
            sforw(f, x; predict=false)
        else
            ypred = sforw(f, x; predict=true)
            losscnt[1] += loss(ypred, ygold); losscnt[2] += 1
            sback(f, ygold, loss)
            while !stack_isempty(f); sback(f); end
            g = gnorm(f); g > maxnorm[2] && (maxnorm[2]=g)
            gscale = (g > gclip > 0 ? gclip/g : 1)
            update!(f; gscale=gscale)
            w = wnorm(f); w > maxnorm[1] && (maxnorm[1]=w) # TODO: take this maxnorm calculation elsewhere, have a statistics optional function? just be smart about scaling?  batch-normalization, msra etc.
            reset!(f)
        end
    end
end

function gradloss(f, data, loss; grad=false, seed=42)
    data_rng = data.rng
    data.rng = MersenneTwister(seed)
    reset!(f)
    myforw = grad ? sforw : forw
    loss1 = 0
    for (x,ygold) in data
        if ygold == nothing
            myforw(f, x; predict=false)
        else
            ypred = myforw(f, x; predict=true)
            loss1 = loss(ypred, ygold)
            if grad
                sback(f, ygold, loss)
                while !stack_isempty(f); sback(f); end
            end
            break
        end
    end
    data.rng = data_rng
    return loss1
end

# Data generator:

import Base: start, next, done

type Data; len; batchsize; epochsize; batch; sum; cnt; rng;
    Data(len, batchsize, epochsize; rng=Base.GLOBAL_RNG) =
    new(len, batchsize, epochsize, zeros(Float32,2,batchsize),
        zeros(Float32,1,batchsize), zeros(Int,1,batchsize), rng)
end

# state keeps track of instance number, time step
start(a::Data)=(0,0)

done(a::Data,s)=(s[1] >= a.epochsize)

function next(a::Data, s)
    (n,t) = s
    t == 0 && (fill!(a.sum, 0); fill!(a.cnt, 0))
    rand!(a.rng, sub(a.batch,1,:))
    fill!(sub(a.batch,2,:),0)
    togo = a.len - t
    for b=1:a.batchsize
        if (a.cnt[b]==0 ? rand(a.rng) <= 2/togo :
            a.cnt[b]==1 ? rand(a.rng) <= 1/togo : false)
            a.batch[2,b] = 1
            a.cnt[b] += 1
            a.sum[b] += a.batch[1,b]
        end
    end
    if t+1 < a.len
        return ((a.batch,nothing), (n,t+1))
    else
        return ((a.batch,a.sum), (n+a.batchsize,0))
    end
end

function parse_commandline(args)
    s = ArgParseSettings()
    @add_arg_table s begin
        "--epochs"
        help = "Number of epochs to train"
        arg_type = Int
        default = 20 # 100
        "--epochsize"
        help = "number of training examples per epoch"
        arg_type = Int
        default = 2000 # 10000
        "--batchsize"
        help = "minibatch size"
        arg_type = Int
        default = 16
        "--length"
        help = "length of the input sequence"
        arg_type = Int
        default = 10 # 100
        "--hidden"
        help = "number of hidden units"
        arg_type = Int
        default = 5 # 100
        "--lrate"
        help = "learning rate"
        arg_type = Float64
        default = 0.05 # 0.01
        "--gclip"
        help = "gradient clip"
        arg_type = Float64
        default = 0.0 # 1.0
        "--gcheck"
        help = "gradient check"
        arg_type = Int
        default = 0
        "--nettype"
        help = "type of network"
        default = "irnn" # "lstm"
        "--load"
        help = "load network from file"
        "--fbias"
        help = "forget gate bias (for lstm)"
        arg_type = Float64
        default = 1.0
        "--winit"
        help = "stdev for weight initialization (for irnn)"
        arg_type = Float64
        default =  0.01 # 0.001
        "--nosharing"
        help = "Do not share register arrays (this is for debugging, should not change results)"
        action = :store_true
        "--seed"
        help = "Random seed"
        arg_type = Int
        default = 1003
    end
    parse_args(args,s)
end

!isinteractive() && !isdefined(Core.Main,:load_only) && main(ARGS)

end #module

### SAMPLE RUNS: Mon Sep 28 21:22:01 PDT 2015

# [dy_052@hpc3013 examples]$ julia adding.jl
# Dict{AbstractString,Any}("hidden"=>5,"batchsize"=>16,"lrate"=>0.05,"length"=>10,"gclip"=>0.0,"std"=>0.01,"gcheck"=>10,"fbias"=>1.0,"epochs"=>20,"seed"=>1003,"epochsize"=>2000,"type"=>"irnn")
# (:gc,3,5,0.006502123f0,0.007450580596923828)
# (2000,0.22126968f0,3.401929f0,5.9513392f0)
# (4000,0.15322119f0,3.428556f0,4.774887f0)
# (6000,0.14666896f0,3.5303187f0,5.8686156f0)
# (8000,0.14103648f0,3.6803625f0,3.3791587f0)
# (10000,0.14484541f0,3.730668f0,3.6132479f0)
# (:gc,3,1,-0.008531229f0,-0.007376074790954634)
# (12000,0.13800602f0,3.8211298f0,3.4911914f0)
# (14000,0.13632888f0,3.8948088f0,3.0174415f0)
# (:gc,3,5,0.08060618f0,0.07353723049163863)
# (16000,0.14532445f0,4.0158386f0,4.1214657f0)
# (:gc,3,1,-0.07339329f0,-0.07551163434982083)
# (18000,0.13509625f0,4.1566772f0,3.2232072f0)
# (20000,0.12674719f0,4.329744f0,3.4179578f0)
# (:gc,1,5,0.072262116f0,0.07338821887969761)
# (22000,0.10952723f0,4.54435f0,6.0535564f0)
# (24000,0.106625326f0,4.769523f0,3.0753422f0)
# (26000,0.09111189f0,4.863088f0,3.6007364f0)
# (28000,0.08520666f0,4.994016f0,5.3230906f0)
# (30000,0.07867393f0,5.157688f0,6.0482206f0)
# (:gc,3,5,-0.09625465f0,-0.10021030902863652)
# (32000,0.07547985f0,5.239566f0,4.534141f0)
# (:gc,3,2,0.12882976f0,0.12787058949471927)
# (34000,0.06994344f0,5.2877007f0,5.1183004f0)
# (:gc,3,2,0.4618018f0,0.46066939830785103)
# (:gc,3,5,0.10286887f0,0.09141862392426543)
# (36000,0.05994114f0,5.384079f0,4.1663117f0)
# (38000,0.06648293f0,5.5462112f0,5.588353f0)
# (40000,0.048857126f0,5.6036315f0,3.805253f0)
#  16.701477 seconds (22.66 M allocations: 1.070 GB, 2.03% gc time)

# [dy_052@hpc3013 examples]$ julia adding.jl --type lstm
# Dict{AbstractString,Any}("hidden"=>5,"batchsize"=>16,"lrate"=>0.05,"length"=>10,"gclip"=>0.0,"std"=>0.01,"gcheck"=>10,"fbias"=>1.0,"epochs"=>20,"seed"=>1003,"epochsize"=>2000,"type"=>"lstm")
# (:gc,10,4,-0.0032901722f0,-0.002384185791015619)
# (:gc,10,6,-0.0020495693f0,-0.0011920928955078197)
# (:gc,10,7,0.0009976054f0,-0.0)
# (:gc,10,10,0.002146461f0,0.0011920928955078095)
# (:gc,12,5,0.008184996f0,0.0095367431640625)
# (:gc,13,4,-0.0014661846f0,-0.0005960464477539098)
# (2000,0.25688365f0,3.5912254f0,1.1217515f0)
# (4000,0.16381074f0,3.5942483f0,0.2719825f0)
# (6000,0.16333067f0,3.6215296f0,0.4068982f0)
# (8000,0.1633528f0,3.6199408f0,0.30277357f0)
# (10000,0.16167283f0,3.6377945f0,0.3591256f0)
# (12000,0.16095969f0,3.6270704f0,0.26166877f0)
# (14000,0.15849201f0,3.6007977f0,0.3542457f0)
# (16000,0.17288947f0,3.6203938f0,0.28349283f0)
# (18000,0.16657862f0,3.6240203f0,0.3288658f0)
# (20000,0.16593297f0,3.6399243f0,0.31581247f0)
# (22000,0.15971881f0,3.6533332f0,0.2592598f0)
# (24000,0.17853853f0,3.6475596f0,0.36688146f0)
# (26000,0.16586855f0,3.6920733f0,0.2798261f0)
# (28000,0.16403076f0,3.6743417f0,0.37097052f0)
# (30000,0.16771112f0,3.7095623f0,0.39885145f0)
# (32000,0.16923295f0,3.7148843f0,0.25710136f0)
# (34000,0.16505243f0,3.7452762f0,0.30730715f0)
# (36000,0.17146099f0,3.7212203f0,0.33892983f0)
# (38000,0.1655073f0,3.7557068f0,0.32211268f0)
# (40000,0.16552721f0,3.7773042f0,0.34974277f0)
#  55.710653 seconds (87.87 M allocations: 4.044 GB, 2.16% gc time)

### DEAD CODE:

# nb = args["batch"]
# xb = Array(Float32, nx, nb, nt)
# xx = [ KUdense(CudaArray, Float32, nx, nb) for t=1:nt ]
# for b=1:nb:ntrn
#     e=min(ntrn, b + nb - 1)
#     permutedims!(xb, sub(xtrn,:,:,b:e), [1,3,2])
#     for t=1:nt
#         copy!(xx, 1, xb, 1+(t-1)*stride(xb,3), length(xx))
#     end
# end

# for b=1:nb:args["test"]

# end

    # x = zeros(Float32, nx, nt, ni)
    # y = zeros(Float32, ny, ni)
    # for i=1:ni
    #     t1 = rand(1:nt)
    #     t2 = rand(1:nt)
    #     while t1==t2; t2 = rand(1:nt); end
    #     x[2,t1,i] = x[2,t2,i] = 1
    #     y[i] = 0
    #     for t=1:nt
    #         x[1,t,i] = rand()
    #         x[2,t,i] == 1 && (y[i] += x[1,t,i])
    #     end
    # end
    # return (x,y)

        # if true
        #     nt = length(xi)
        #     y = convert(Array,r.out[nops(r)])
        #     a = convert(Array,yi[nt])
        #     e = sum((y-a).^2)/ccount(y)
        #     @show (e,2*err)
        # end

# Organization of the training set and training batches:
# D:dimensionality, I:instance, B:batch-instance, T:time
# In FFNN size(xtrain)=(D,I), size(xbatch)=(D,B) where D fastest
# In Net size(xtrain)=(D,T,I) or xtrain=[(D,T1),(D,T2),...]
# i.e. we want each instance to be contiguous.
# In Net size(xbatch)=(D,B,T) or xbatch=[(D,B1),(D,B2),...]
# i.e. we want each time-step to be contiguous.
# train->batch will need to do some shuffling

# CANCEL: splitting net into two parts may be more efficient
# this way forw calculation by the last three layers are wasted
# but this way we are testing the general input/output

# import Knet: forw, back, ninputs, param, similar!, gpu, initforw, initback, setparam!, update, loss, axpy! # push, pop, get1
# import Knet: backprop, train, predict, nz
# include("../src/rnn.jl")

# setparam!(net1; nesterov=0.01)
# net1.dbg = true

# net0.dbg = net1.dbg = net2.dbg = true
# setseed(42); e1 = forwback1(net0, xtst[1], ytst[1][end])
# setseed(42); e2 = forwback2(net1, net2, xtst[1], ytst[1][end])
# @test @show e1 == e2
# n0 = nops(net0); n1 = nops(net1); n2 = nops(net2)
# for n=n0:-1:1
#     @show n
#     p0 = param(net0.op[n])
#     p1 = (n<=n1 ? param(net1.op[n]) : param(net2.op[n-n1]))
#     p0 == nothing && p1 == nothing && continue
#     @test convert(Array, p1.arr)==convert(Array, p0.arr)
#     @test convert(Array, p1.diff)==convert(Array, p0.diff)
# end

# trnerr = trnerr0 = 0
# maxg = maxg0 = 0

    # c = gradcheck(net0, xtrn[1], ytrn[1]; ncheck=10)
    # c == nothing || println(tuple(:gradcheck, c...))
    # (trnerr0,maxg0) = train2(net0, xtrn, ytrn)
    # @test isapprox(trnerr, trnerr0; atol=1e-3)
    # @test isapprox(maxg, maxg0; atol=1e-2)
    # tsterr = test(net1, xtst, ytst)
    # println(tuple(epoch,sqrt(2trnerr),sqrt(2tsterr),maxgnorm))

# DONE:
# + minibatching (16)
# + dif = nothing, dif0 = array
# + profiling
# x separate into two networks
# + xfer train/predict to knet (tforw? figure out right interface)
# + gradient clipping (1/10/100)
# x adadelta (not this paper)

# setparam!(net0; lr=args["lr"], gc=args["gc"])
# setparam!(net1; lr=args["lr"], gc=args["gc"])
# setparam!(net2; lr=args["lr"], gc=args["gc"])
# args["type"] == "lstm" && setparam!(net0.op[9]; init=fill!, initp=args["fb"])

# function gendata(ni, nt)
#     x = cell(ni)
#     y = cell(ni)
#     for i=1:ni
#         x[i] = cell(nt)
#         y[i] = cell(nt)
#         for t=1:nt
#             x[i][t] = Float32[rand(), 0]
#             y[i][t] = nothing
#         end
#         t1 = rand(1:nt)
#         t2 = rand(1:nt)
#         while t1==t2; t2 = rand(1:nt); end
#         x[i][t1][2] = 1
#         x[i][t2][2] = 1
#         y[i][nt] = Float32[x[i][t1][1]+x[i][t2][1]]
#     end
#     return (x,y)
# end

# # DONE: fix this so y[i] is not a sequence
# function batch(x, y, nb)
#     isempty(x) && return (x,y)
#     xx = Any[]
#     yy = Any[]
#     ni = length(x)    # assume x and y are same length
#     nt = length(x[1]) # assume all x[i] are same length
#     for i1=1:nb:ni
#         i2=min(ni,i1+nb-1)
#         xi = Any[]
#         yi = Any[]
#         for t=1:nt
#             xit = x[i1][t]
#             xt = similar(xit, tuple(size(xit)..., i2-i1+1))
#             for i=i1:i2; xt[:,i-i1+1] = x[i][t]; end
#             push!(xi, xt)
#             yit = y[i1][t]
#             if yit == nothing
#                 yt = nothing # assumes yit=nothing for all i
#             else
#                 yt = similar(yit, tuple(size(yit)..., i2-i1+1))
#                 for i=i1:i2; yt[:,i-i1+1] = y[i][t]; end
#             end
#             push!(yi, yt)
#         end
#         push!(xx, xi)
#         push!(yy, yi)
#     end
#     return(xx, yy)
# end
# net0 = (args["type"] == "irnn" ? Net(irnn(nh),quadlosslayer(ny)) :
#         args["type"] == "lstm" ? Net(lstm(nh),quadlosslayer(ny)) :
#         error("Unknown network type "*args["type"]))
# setparam!(net; lr=args["lr"], gc=args["gc"])  # do a global gclip instead of per parameter
        # "--test"
        # help = "number of testing examples"
        # arg_type = Int
        # default = 2000
    # trnmse = tstmse = maxg = maxw = 0
    # for i=1:length(xtrn)
    #     (l,w,g) = train(net, xtrn[i], ytrn[i][end]; getloss=true, getnorm=true, gclip=args["gc"])
    #     trnmse += l
    #     w > maxw && (maxw = w)
    #     g > maxg && (maxg = g)
    # end
    # for i=1:length(xtst)
    #     tstmse += test(net, xtst[i], ytst[i][end])
    # end
    # trnmse = 2*trnmse/length(xtrn)
    # tstmse = 2*tstmse/length(xtst)
# ntrn = args["train"]
# ntst = args["test"]
# nt = args["length"]
# (xtst1,ytst1) = gendata(ntst, nt)
# (xtst,ytst) = batch(xtst1, ytst1, args["batch"])
# (xtrn1,ytrn1) = gendata(ntrn, nt)
# (xtrn,ytrn) = batch(xtrn1, ytrn1, args["batch"])

# gradcheck(net, xtrn[1], ytrn[1][end]; ncheck=100, rtol=.01, atol=.01)
    # gradcheck(deepcopy(net), xtrn[1], ytrn[1][end]; ncheck=10, rtol=.01, atol=.01)
    # gradcheck(deepcopy(net), xtrn[1], ytrn[1][end]; ncheck=typemax(Int), rtol=.01, atol=0.001)
# DONE: move batch somewhere else

# Sample run for debugging:
# julia> include("adding.jl")
# Dict{AbstractString,Any}("hidden"=>5,"batchsize"=>16,"lrate"=>0.05,"length"=>10,"gclip"=>0.0,"std"=>0.01,"gcheck"=>10,"fbias"=>1.0,"epochs"=>20,"seed"=>1003,"epochsize"=>2000,"type"=>"irnn")
# (2000,0.22126971078821514,3.401929f0,5.9513392f0)
# (4000,0.15322119775290655,3.428556f0,4.774887f0)
# (6000,0.14666901898846646,3.5303187f0,5.8686156f0)
# (8000,0.1410364930989558,3.6803625f0,3.3791587f0)
# (10000,0.14484541128223494,3.730668f0,3.6132479f0)
# (:gc,3,1,-0.008531229f0,-0.0073469345807098295)
# (12000,0.13800603687082819,3.8211298f0,3.4911914f0)
# (14000,0.13632891371942893,3.8948088f0,3.0174415f0)
# (:gc,3,5,0.08060618f0,0.07355656634899832)
# (16000,0.1453244475960048,4.0158386f0,4.1214657f0)
# (:gc,3,1,-0.07339329f0,-0.07550548616563685)
# (18000,0.13509626036176464,4.1566772f0,3.2232072f0)
# (20000,0.1267472031369284,4.329744f0,3.4179578f0)
# (:gc,1,5,0.072262116f0,0.07338319164773356)
# (22000,0.10952723081752598,4.54435f0,6.0535564f0)
# (24000,0.10662533053077042,4.769523f0,3.0753422f0)
# (26000,0.09111188449872992,4.863088f0,3.6007364f0)
# (28000,0.08520665419544599,4.994016f0,5.3230906f0)
# (30000,0.07867394052687623,5.157688f0,6.0482206f0)
# (:gc,3,5,-0.09625465f0,-0.10007152013714596)
# (32000,0.07547986407922849,5.239566f0,4.534141f0)
# (:gc,3,2,0.12882976f0,0.1279206117033168)
# (34000,0.06994342621574566,5.2877007f0,5.1183004f0)
# (:gc,3,2,0.4618018f0,0.460615519841718)
# (:gc,3,5,0.10286887f0,0.09136574590230435)
# (36000,0.059941143911885625,5.384079f0,4.1663117f0)
# (38000,0.06648293199506522,5.5462112f0,5.588353f0)
# (40000,0.048857121442293316,5.6036315f0,3.805253f0)
#  13.332866 seconds (23.97 M allocations: 1017.019 MB, 2.08% gc time)
# :ok

# len	hidden	lr	gc	mse<0.1
# 2	1	0.3		4000
# 3	2	0.2		4000
# 5	2	0.1		8000
# 10	5	0.05		26000
# 20	10	0.03		80000
# 40	30	0.01		220000
# 60	40	0.01	10	440000
# 100	100	0.01	1	870000
# 100	100	0.01	10	960000
# 150	100	0.01	1	1390000
# 200	100	0.01	1	2080000
# 300	100	0.01	1	3860000
# 400	100	0.01	1	7400000 speed=400k/h, unstable

# DONE: lstm does not work with 10x5 find out why: much larger --lr=1.0, also large --fb >= 1 helps.
# The following settings solve 10x5 in 12000 iterations:
# "hidden"=>5,"lr"=>0.7,"batchsize"=>16,"length"=>10,"gclip"=>1.0,"fb"=>100.0,"gcheck"=>10,"epochs"=>10,"seed"=>1003,"epochsize"=>2000,"type"=>"lstm")

    # args = parse_commandline()
    # args = parse_commandline(split("--epochsize 2000 --length 10 --hidden 5 --lrate 0.05 --gc 0 --epochs 20 --seed 1003"))
    # args = parse_commandline(split("--epochsize 10000 --test 2000 --length 100 --hidden 100 --lrate 0.01 --gc 1.0 --epochs 100"))

# Sample run for LSTM:
# julia> main(split("--type lstm"))
# Dict{AbstractString,Any}("hidden"=>5,"batchsize"=>16,"lrate"=>0.05,"length"=>10,"gclip"=>0.0,"std"=>0.01,"gcheck"=>10,"fbias"=>1.0,"epochs"=>20,"seed"=>1003,"epochsize"=>2000,"type"=>"lstm")
# (:gc,10,4,-0.0032901722f0,-0.002384185791015619)
# (:gc,10,6,-0.0020495693f0,-0.0011920928955078197)
# (:gc,10,7,0.0009976054f0,-0.0)
# (:gc,10,10,0.002146461f0,0.0011920928955078095)
# (:gc,12,5,0.008184996f0,0.0095367431640625)
# (:gc,13,4,-0.0014661846f0,-0.0005960464477539098)
# (2000,0.25688365f0,3.5912254f0,1.1217515f0)
# (4000,0.16381074f0,3.5942483f0,0.2719825f0)
# (6000,0.16333067f0,3.6215296f0,0.4068982f0)
# (8000,0.1633528f0,3.6199408f0,0.30277357f0)
# (10000,0.16167283f0,3.6377945f0,0.3591256f0)
# (12000,0.16095969f0,3.6270704f0,0.26166877f0)
# (14000,0.15849201f0,3.6007977f0,0.3542457f0)
# (16000,0.17288947f0,3.6203938f0,0.28349283f0)
# (18000,0.16657862f0,3.6240203f0,0.3288658f0)
# (20000,0.16593297f0,3.6399243f0,0.31581247f0)
# (22000,0.15971881f0,3.6533332f0,0.2592598f0)
# (24000,0.17853853f0,3.6475596f0,0.36688146f0)
# (26000,0.16586855f0,3.6920733f0,0.2798261f0)
# (28000,0.16403076f0,3.6743417f0,0.37097052f0)
# (30000,0.16771112f0,3.7095623f0,0.39885145f0)
# (32000,0.16923295f0,3.7148843f0,0.25710136f0)
# (34000,0.16505243f0,3.7452762f0,0.30730715f0)
# (36000,0.17146099f0,3.7212203f0,0.33892983f0)
# (38000,0.1655073f0,3.7557068f0,0.32211268f0)
# (40000,0.16552721f0,3.7773042f0,0.34974277f0)
#  55.892459 seconds (93.09 M allocations: 4.137 GB, 2.54% gc time)

# OLD GENERATOR FOR COMPARISON:

# type Adding1; len; batchsize; epochsize; b; x; y;
#     Adding1(len, batchsize, epochsize; o...)=new(len,batchsize,epochsize,Adding0(len,batchsize,epochsize; o...))
# end

# start(a::Adding1)=(0,0)
# done(a::Adding1,s)=(s[1] >= a.b.epochsize)

# function next(a::Adding1, s)
#     (n,t) = s
#     t == 0 && ((xy,n1) = next(a.b,n); (a.x,a.y)=xy)
#     t += 1
#     if t < a.b.len
#         return ((a.x[t], nothing), (n,t))
#     else
#         return ((a.x[t], a.y), (n+a.b.batchsize, 0))
#     end
# end

# type Adding0; len; batchsize; epochsize; rng;
#     Adding0(len, batchsize, epochsize; rng=MersenneTwister(0))=new(len, batchsize, epochsize, rng)
# end

# start(a::Adding0)=0

# done(a::Adding0,n)=(n >= a.epochsize)

# function next(a::Adding0, n)
#     nb = min(a.batchsize, a.epochsize-n)
#     x = [ vcat(rand(a.rng,Float32,1,nb),zeros(Float32,1,nb)) for t=1:a.len ]
#     y = Array(Float32,1,nb)
#     t1 = rand(a.rng,1:a.len,nb)
#     t2 = rand(a.rng,1:a.len,nb)
#     for b=1:nb
#         while t2[b]==t1[b]
#             t2[b]=rand(a.rng,1:a.len)
#         end
#         x[t1[b]][2,b]=1
#         x[t2[b]][2,b]=1
#         y[b] = x[t1[b]][1,b] + x[t2[b]][1,b]
#     end
#     return ((x,y), n+nb)
# end

    # p1 = (opts["nettype"] == "irnn" ? Net(irnn; out=opts["hidden"], winit=Gaussian(0,opts["winit"])) :
    #       opts["nettype"] == "lstm" ? Net(lstm; out=opts["hidden"], fbias=opts["fbias"]) :
    #       error("Unknown network type "*opts["nettype"]))
    # p2 = Net(wb; out=1, winit=Gaussian(0,opts["winit"]))

# @knet function p1(x; rnn=nothing, hidden=0, o...)
#     y = rnn(x; o..., out=hidden)
# end

# @knet function p2(x; o...)
#     y = wb(x; o..., out=1)
# end


text:
mutable struct UnequalTimeStack{GT<:Number, GMT<:AbstractArray{GT}} <: AbstractDQMCStack
    # B_{n*safe_mult} ... B_1
    forward_idx::Int64
    forward_u_stack::Vector{GMT}
    forward_d_stack::Vector{Vector{Float64}}
    forward_t_stack::Vector{GMT}

    # B_{n*safe_mult+1}^† ... B_M^†
    backward_idx::Int64
    backward_u_stack::Vector{GMT}
    backward_d_stack::Vector{Vector{Float64}}
    backward_t_stack::Vector{GMT}

    # B_{n*safe_mult+1}^-1 .. B_{(n+1)*safe_mult}^-1
    inv_done::Vector{Bool}
    inv_u_stack::Vector{GMT}
    inv_d_stack::Vector{Vector{Float64}}
    inv_t_stack::Vector{GMT}

    # Greens construction
    greens::GMT
    tmp::GMT
    U::GMT
    D::Vector{Float64}
    T::GMT

    # To avoid recalculating
    last_update::Int64
    last_slices::Tuple{Int64, Int64}

    function UnequalTimeStack{GT, GMT}() where {GT<:Number, GMT<:AbstractArray{GT}}
        @assert isconcretetype(GT);
        @assert isconcretetype(GMT);
        new{GT, GMT}()
    end
end

# TODO
# function Base.sizeof(::Type{UnequalTimeStack}, mc::DQMC)
#     GreensEltype =  MonteCarlo.geltype(mc)
#     N = length(mc.model.l)
#     flv = mc.model.flv
#     blocks = mc.s.n_elements

#     sizeof(GreensEltype) * flv*N * flv*N * blocks * 4 +         # forward+backward U, T
#     sizeof(GreensEltype) * flv*N * blocks * 2 +                 # forward+backward D
#     sizeof(GreensEltype) * flv*N * flv*N * (blocks-1) * 2 +     # inv U, T
#     sizeof(GreensEltype) * flv*N * (blocks-1) +                 # inv Dr
#     5*8 + sizeof(Bool) * (blocks-1) +                           # book-keeping/skipping
#     sizeof(GreensEltype) * flv*N * flv*N * 3 +                  # greens, U, Tr
#     sizeof(GreensEltype) * flv*N                                # D
# end

function initialize_stack(mc::DQMC, s::UnequalTimeStack)
    GreensElType  = geltype(mc)
    GreensMatType = gmattype(mc)
    N = length(lattice(mc))
    flv = nflavors(mc.model)
    M = convert(Int, mc.p.slices / mc.p.safe_mult) + 1

    # B_{n*safe_mult} ... B_1
    s.forward_idx = 1
    s.forward_u_stack = [GreensMatType(undef, flv*N, flv*N) for _ in 1:M]
    s.forward_d_stack = [zeros(Float64, flv*N) for _ in 1:M]
    s.forward_t_stack = [GreensMatType(undef, flv*N, flv*N) for _ in 1:M]

    # B_{n*safe_mult+1}^† ... B_M^†
    s.backward_idx = M-1
    s.backward_u_stack = [GreensMatType(undef, flv*N, flv*N) for _ in 1:M]
    s.backward_d_stack = [zeros(Float64, flv*N) for _ in 1:M]
    s.backward_t_stack = [GreensMatType(undef, flv*N, flv*N) for _ in 1:M]

    # B_{n*safe_mult+1}^-1 .. B_{(n+1)*safe_mult}^-1
    s.inv_done = fill(false, M-1)
    s.inv_u_stack = [GreensMatType(undef, flv*N, flv*N) for _ in 1:M-1]
    s.inv_d_stack = [zeros(Float64, flv*N) for _ in 1:M-1]
    s.inv_t_stack = [GreensMatType(undef, flv*N, flv*N) for _ in 1:M-1]

    # Greens construction
    s.greens = GreensMatType(undef, flv*N, flv*N)
    s.tmp = GreensMatType(undef, flv*N, flv*N)
    s.U = GreensMatType(undef, flv*N, flv*N)
    s.D = zeros(Float64, flv*N)
    s.T = GreensMatType(undef, flv*N, flv*N)

    # To avoid recalculating
    last_update = -1
    last_slices = (-1, -1)

    # maybe skip identities?
    copyto!(s.forward_u_stack[1], I)
    s.forward_d_stack[1] .= 1
    copyto!(s.forward_t_stack[1], I)

    copyto!(s.backward_u_stack[end], I)
    s.backward_d_stack[end] .= 1
    copyto!(s.backward_t_stack[end], I)

    s
end

########################################
# Stack building (full & lazy)
########################################

@bm function build_stack(mc::DQMC, s::UnequalTimeStack)
    # stack = [0, Δτ, 2Δτ, ..., β]
    if mc.last_sweep == s.last_update && all(s.inv_done) &&
        (s.forward_idx == length(mc.s.ranges)+1) && (s.backward_idx == 1)
        return nothing
    end

    # forward
    @bm "forward build" begin
        @inbounds for idx in 1:length(mc.s.ranges)
            copyto!(mc.s.curr_U, s.forward_u_stack[idx])
            for slice in mc.s.ranges[idx]
                multiply_slice_matrix_left!(mc, mc.model, slice, mc.s.curr_U)
            end
            vmul!(mc.s.tmp1, mc.s.curr_U, Diagonal(s.forward_d_stack[idx]))
            udt_AVX_pivot!(
                s.forward_u_stack[idx+1], s.forward_d_stack[idx+1],
                mc.s.tmp1, mc.s.pivot, mc.s.tempv
            )
            vmul!(s.forward_t_stack[idx+1], mc.s.tmp1, s.forward_t_stack[idx])
        end
    end

    # backward
    @bm "backward build" begin
        @inbounds for idx in length(mc.s.ranges):-1:1
            copyto!(mc.s.curr_U, s.backward_u_stack[idx + 1])
            for slice in reverse(mc.s.ranges[idx])
                multiply_daggered_slice_matrix_left!(mc, mc.model, slice, mc.s.curr_U)
            end
            vmul!(mc.s.tmp1, mc.s.curr_U, Diagonal(s.backward_d_stack[idx + 1]))
            udt_AVX_pivot!(
                s.backward_u_stack[idx], s.backward_d_stack[idx],
                mc.s.tmp1, mc.s.pivot, mc.s.tempv
            )
            vmul!(s.backward_t_stack[idx], mc.s.tmp1, s.backward_t_stack[idx+1])
        end
    end

    # inverse
    # TODO should this multiply to U's instead?
    @bm "inverse build" begin
        @inbounds for idx in 1:length(mc.s.ranges)
            copyto!(s.inv_t_stack[idx], I)
            for slice in reverse(mc.s.ranges[idx])
                multiply_slice_matrix_inv_left!(mc, mc.model, slice, s.inv_t_stack[idx])
            end
            udt_AVX_pivot!(
                s.inv_u_stack[idx], s.inv_d_stack[idx], s.inv_t_stack[idx],
                mc.s.pivot, mc.s.tempv
            )
        end
    end

    s.inv_done .= true
    s.forward_idx = length(mc.s.ranges)+1
    s.backward_idx = 1
    s.last_update = mc.last_sweep

    nothing
end

@bm function lazy_build_forward!(mc::DQMC, s::UnequalTimeStack, upto)
    # stack = [0, Δτ, 2Δτ, ..., β]
    if s.last_update != mc.last_sweep
        s.last_update = mc.last_sweep
        s.inv_done .= false
        s.forward_idx = 1
        s.backward_idx = length(mc.s.ranges)+1
    end

    # forward
    @bm "forward build" begin
        @inbounds for idx in s.forward_idx:upto-1
            copyto!(mc.s.curr_U, s.forward_u_stack[idx])
            for slice in mc.s.ranges[idx]
                multiply_slice_matrix_left!(mc, mc.model, slice, mc.s.curr_U)
            end
            vmul!(mc.s.tmp1, mc.s.curr_U, Diagonal(s.forward_d_stack[idx]))
            udt_AVX_pivot!(
                s.forward_u_stack[idx+1], s.forward_d_stack[idx+1],
                mc.s.tmp1, mc.s.pivot, mc.s.tempv
            )
            vmul!(s.forward_t_stack[idx+1], mc.s.tmp1, s.forward_t_stack[idx])
        end
    end

    s.forward_idx = max(upto, s.forward_idx)
    nothing
end

@bm function lazy_build_backward!(mc::DQMC, s::UnequalTimeStack, downto)
    if s.last_update != mc.last_sweep
        s.last_update = mc.last_sweep
        s.inv_done .= false
        s.forward_idx = 1
        s.backward_idx = length(mc.s.ranges)+1
    end

    # backward
    @bm "backward build" begin
        @inbounds for idx in s.backward_idx-1:-1:downto
            copyto!(mc.s.curr_U, s.backward_u_stack[idx + 1])
            for slice in reverse(mc.s.ranges[idx])
                multiply_daggered_slice_matrix_left!(mc, mc.model, slice, mc.s.curr_U)
            end
            vmul!(mc.s.tmp1, mc.s.curr_U, Diagonal(s.backward_d_stack[idx + 1]))
            udt_AVX_pivot!(
                s.backward_u_stack[idx], s.backward_d_stack[idx],
                mc.s.tmp1, mc.s.pivot, mc.s.tempv
            )
            vmul!(s.backward_t_stack[idx], mc.s.tmp1, s.backward_t_stack[idx+1])
        end
    end

    s.backward_idx = min(downto, s.backward_idx)
end

@bm function lazy_build_inv!(mc::DQMC, s::UnequalTimeStack, from, to)
    if s.last_update != mc.last_sweep
        s.last_update = mc.last_sweep
        s.inv_done .= false
        s.forward_idx = 1
        s.backward_idx = length(mc.s.ranges)+1
    end

    # inverse
    @bm "inverse build" begin
        @inbounds for idx in from:to
            s.inv_done[idx] && continue
            s.inv_done[idx] = true
            copyto!(s.inv_t_stack[idx], I)
            for slice in reverse(mc.s.ranges[idx])
                multiply_slice_matrix_inv_left!(mc, mc.model, slice, s.inv_t_stack[idx])
            end
            udt_AVX_pivot!(
                s.inv_u_stack[idx], s.inv_d_stack[idx], s.inv_t_stack[idx],
                mc.s.pivot, mc.s.tempv
            )
        end
    end

    nothing
end

########################################
# Calculating greens
########################################

"""
    greens(mc::DQMC, k, l)

Calculates the unequal-time Greens function at slice `k` and `l`, i.e.
`G(k <- l) = G(kΔτ <- lΔτ) = ⟨cᵢ(k⋅Δτ)cⱼ(l⋅Δτ)^†⟩` where `nslices(mc) ≥ k > l ≥ 0`.

Note that `G(0, 0) = G(nslices, nslices) = G(β, β) = G(0, 0)`.

This method requires the `UnequalTimeStack` to be available. The stack
variables `Ul`, `Dl`, `Tl`, `Ur`, `Dr`, `Tr`, `curr_U` and `tmp1` will
be overwritten (but are avaible as `output` and `temp`).
"""
@bm greens(mc::DQMC, slice1::Int64, slice2::Int64) = copy(_greens!(mc, slice1, slice2))
"""
    greens!(mc::DQMC, k, l[; output=mc.s.greens_temp, temp=mc.s.tmp1])
"""
@bm function greens!(
        mc::DQMC, slice1::Int64, slice2::Int64;
        output = mc.s.greens_temp, temp = mc.s.tmp1
    )
    _greens!(mc, slice1, slice2, output, temp)
end
"""
    _greens!(mc::DQMC, k, l, output=mc.s.greens_temp, temp=mc.s.tmp1)
"""
function _greens!(
        mc::DQMC, slice1::Int64, slice2::Int64,
        output = mc.s.greens_temp, temp = mc.s.tmp1
    )
    calculate_greens(mc, slice1, slice2)
    return _greens!(mc, output, mc.ut_stack.greens, temp)
end
@bm function calculate_greens(mc::DQMC, slice1::Int64, slice2::Int64)
    @assert 0 ≤ slice1 ≤ mc.p.slices
    @assert 0 ≤ slice2 ≤ mc.p.slices
    s = mc.ut_stack
    if (s.last_slices != (slice1, slice2)) || (s.last_update != mc.last_sweep)
        s.last_slices = (slice1, slice2)
        if slice1 ≥ slice2
            calculate_greens_full1!(mc, s, slice1, slice2)
        else
            calculate_greens_full2!(mc, s, slice1, slice2)
        end
    end
    s.greens
end

# low = slice2, high = slice1
"""
    compute_inverse_udt_block!(mc, s::UnequalTimeStack, low, high[, ...])

Computes a inverse UDT block `U D T = B_{low+1}^-1 B_{low+2}^-1 ⋯ B_high^-1`.

Default extra args are: (in order)
* `U = s.U`
* `D = s.D`
* `T = s.T`
* `tmp1 = mc.s.tmp1`
* `tmp2 = mc.s.tmp2`
"""
function compute_inverse_udt_block!(
        mc, s, low, high,
        U = s.U, D = s.D, T = s.T, tmp1 = mc.s.tmp1, tmp2 = mc.s.tmp2
    )
    # @assert low ≤ high
    # B_{low+1}^-1 B_{low+2}^-1 ⋯ B_{high-1}^-1 B_high^-1

    # Combine pre-computed blocks
    @bm "inverse pre-computed" begin
        lower = div(low+1 + mc.p.safe_mult - 2, mc.p.safe_mult) + 1
        upper = div(high, mc.p.safe_mult)
        lazy_build_inv!(mc, s, lower, upper)
        copyto!(U, I)
        D .= 1
        copyto!(T, I)
        # UDT combining style
        for idx in lower:upper
            # () is operation order, [] marks UDT decomposition
            # U [(D (T stack_U)) stack_D] stack_T
            # (U [U') D (T'] stack_T)
            # U D T
            vmul!(tmp1, T, s.inv_u_stack[idx])
            vmul!(tmp2, Diagonal(D), tmp1)
            vmul!(tmp1, tmp2, Diagonal(s.inv_d_stack[idx]))
            udt_AVX_pivot!(tmp2, D, tmp1, mc.s.pivot, mc.s.tempv)
            vmul!(T, tmp1, s.inv_t_stack[idx])
            vmul!(tmp1, U, tmp2)
            copyto!(U, tmp1)
        end
    end

    # remaining multiplications to reach specified bounds
    @bm "inverse fine tuning" begin
        lower_slice = (lower-1) * mc.p.safe_mult + 1
        upper_slice = upper * mc.p.safe_mult

        for slice in min(lower_slice-1, high) : -1 : low+1
            multiply_slice_matrix_inv_left!(mc, mc.model, slice, U)
        end
        if min(lower_slice-1, high) ≥ low+1
            vmul!(tmp1, U, Diagonal(D))
            udt_AVX_pivot!(U, D, tmp1, mc.s.pivot, mc.s.tempv)
            vmul!(tmp2, tmp1, T)
            copyto!(T, tmp2)
        end
        for slice in max(upper_slice+1, min(lower_slice-1, high)+1) : high
            multiply_slice_matrix_inv_right!(mc, mc.model, slice, T)
        end
        # U D T = B_{l+1}^-1 B_{l+2}^-1 ⋯ B_{k-1}^-1 B_k^-1
    end
    return nothing
end

"""
    compute_forward_udt_block!(mc, s::UnequalTimeStack, slice[, ...])

Computes a UDT block `U D T = B_slice B_{slice-1} ⋯ B_1`.

Default extra args are: (in order)
* `U = mc.s.Ul`
* `D = mc.s.Dl`
* `T = mc.s.Tl`
* `tmp = mc.s.tmp1`
"""
function compute_forward_udt_block!(
        mc, s, slice, apply_pivot = Val(true),
        U = mc.s.Ul, D = mc.s.Dl, T = mc.s.Tl, tmp = mc.s.tmp1,
    )
    # B(slice, 1) = B_slice B_{slice-1} ⋯ B_1 = Ul Dl Tl
    @bm "forward B" begin
        idx = div(slice-1, mc.p.safe_mult)
        lazy_build_forward!(mc, s, idx+1)
        copyto!(T, s.forward_u_stack[idx+1])
        for l in mc.p.safe_mult * idx + 1 : slice
            multiply_slice_matrix_left!(mc, mc.model, l, T)
        end
        vmul!(tmp, T, Diagonal(s.forward_d_stack[idx+1]))
        udt_AVX_pivot!(U, D, tmp, mc.s.pivot, mc.s.tempv, apply_pivot)
        vmul!(T, tmp, s.forward_t_stack[idx+1])
    end
    return nothing
end

"""
    compute_backward_udt_block!(mc, s::UnequalTimeStack, slice[, ...])

Computes a UDT block `(U D T)^† = T^† D U^† = B_N B_{N-1} ⋯ B_slice`.

Default extra args are: (in order)
* `U = mc.s.Ur`
* `D = mc.s.Dr`
* `T = mc.s.Tr`
* `tmp = mc.s.tmp1`
"""
function compute_backward_udt_block!(
        mc, s, slice,
        U = mc.s.Ur, D = mc.s.Dr, T = mc.s.Tr, tmp = mc.s.tmp1
    )
    # B(N, slice) = B_N B_{N-1} ⋯ B_{slice+1} = (Ur Dr Tr)^† = Tr^† Dr^† Ur^†
    @bm "backward B" begin
        idx = div.(slice + mc.p.safe_mult - 1, mc.p.safe_mult)
        lazy_build_backward!(mc, s, idx+1)
        copyto!(U, s.backward_u_stack[idx+1])
        for l in mc.p.safe_mult * idx : -1 : slice+1
            multiply_daggered_slice_matrix_left!(mc, mc.model, l, U)
        end
        vmul!(tmp, U, Diagonal(s.backward_d_stack[idx+1]))
        udt_AVX_pivot!(U, D, tmp, mc.s.pivot, mc.s.tempv)
        vmul!(T, tmp, s.backward_t_stack[idx + 1])
    end
    return nothing
end

@bm function calculate_greens_full1!(mc, s, slice1, slice2)
    # stack = [0, Δτ, 2Δτ, ..., β] = [0, safe_mult, 2safe_mult, ... N]
    # @assert slice1 ≥ slice2

    # k ≥ l or slice1 ≥ slice2
    # B_{l+1}^-1 B_{l+2}^-1 ⋯ B_{k-1}^-1 B_k^-1
    compute_inverse_udt_block!(mc, s, slice2, slice1) # low high

    # B(slice2, 1) = Ul Dl Tl
    compute_forward_udt_block!(mc, s, slice2)

    # B(N, slice1) = (Ur Dr Tr)^† = Tr^† Dr^† Ur^†
    compute_backward_udt_block!(mc, s, slice1)

    # U D T remains valid here
    @bm "compute G" begin
        # [B_{l+1}^-1 B_{l+2}^-1 ⋯ B_k^-1 + B_l ⋯ B_1 B_N ⋯ B_{k+1}]^-1
        # [U D T + Ul (Dl Tl Tr^† Dr) Ur^†]^-1
        @bm "B1" begin
            vmul!(s.greens, mc.s.Tl, adjoint(mc.s.Tr))
            vmul!(mc.s.tmp1, s.greens, Diagonal(mc.s.Dr))
            vmul!(s.greens, Diagonal(mc.s.Dl), mc.s.tmp1)
        end
        # [U D T + Ul (G) Ur^†]^-1
        @bm "udt" begin
            udt_AVX_pivot!(mc.s.Tr, mc.s.Dr, s.greens, mc.s.pivot, mc.s.tempv, Val(false))
        end
        # [U D T + (Ul Tr) Dr (G Ur^†)]^-1
        @bm "B2" begin
            vmul!(mc.s.Tl, mc.s.Ul, mc.s.Tr)
            # (G Ur^†) = (Ur / G)^-1
            # Ur := Ur / G
            rdivp!(mc.s.Ur, s.greens, mc.s.Ul, mc.s.pivot)
        end
        # [U D T + Tl Dr Ur^-1]^-1
        # [U (D T Ur + U^† Tl Dr) Ur^-1]^-1
        # [U D_max (D_min T Ur 1/Dr_max + 1/D_max U^† Tl Dr_min) Dr_max Ur^-1]^-1
        @bm "B3" begin
            # 1/D_max U^† Tl Dr_min
            vmul!(mc.s.Tr, adjoint(s.U), mc.s.Tl)
            vmaxinv!(mc.s.Dl, s.D) # mc.s.Dl .= 1.0 ./ max.(1.0, s.D)
            vmul!(mc.s.tmp1, Diagonal(mc.s.Dl), mc.s.Tr)
            vmin!(mc.s.Dl, mc.s.Dr) # mc.s.Dl .= min.(1.0, mc.s.Dr)
            vmul!(mc.s.Tr, mc.s.tmp1, Diagonal(mc.s.Dl))
        end
        # [U D_max (D_min T Ur 1/Dr_max + Tr) Dr_max Ur^-1]^-1
        @bm "B4" begin
            # D_min T Ur 1/Dr_max
            vmul!(mc.s.Tl, s.T, mc.s.Ur)
            vmin!(mc.s.Dl, s.D) # mc.s.Dl .= min.(1.0, s.D)
            vmul!(mc.s.tmp1, Diagonal(mc.s.Dl), mc.s.Tl)
            vmaxinv!(mc.s.Dl, mc.s.Dr) # mc.s.Dl .= 1.0 ./ max.(1.0, mc.s.Dr)
            vmul!(mc.s.Tl, mc.s.tmp1, Diagonal(mc.s.Dl))
        end
        # [U D_max (Tl + Tr) Dr_max Ur^-1]^-1
        @bm "sum, UDT" begin
            rvadd!(mc.s.Tl, mc.s.Tr)
            udt_AVX_pivot!(mc.s.Tr, mc.s.Dl, mc.s.Tl, mc.s.pivot, mc.s.tempv, Val(false))
        end
        # [U D_max (Tr Dl Tl) Dr_max Ur^-1]^-1
        # Ur 1/Dr_max Tl^-1 1/Dl Tr^† D_max U^†
        @bm "B5" begin
            # [[((1/Dr_max) / Tl) 1/Dl] Tr^†] D_max
            vmaxinv!(mc.s.Dr, mc.s.Dr) # mc.s.Dr .= 1.0 ./ max.(1.0, mc.s.Dr)
            copyto!(mc.s.Ul, Diagonal(mc.s.Dr))
            rdivp!(mc.s.Ul, mc.s.Tl, mc.s.tmp1, mc.s.pivot)
            vinv!(mc.s.Dl) # mc.s.Dl[i] .= 1.0 ./ mc.s.Dl[i]
            vmul!(mc.s.tmp1, mc.s.Ul, Diagonal(mc.s.Dl))
            vmul!(mc.s.Ul, mc.s.tmp1, adjoint(mc.s.Tr))
            vmaxinv!(mc.s.Dl, s.D) # mc.s.Dl .= 1.0 ./ max.(1.0, s.D)
            vmul!(s.greens, mc.s.Ul, Diagonal(mc.s.Dl))
        end
        # Ur G U^†
        @bm "B6" begin
            vmul!(mc.s.Tr, s.greens, adjoint(s.U))
            vmul!(s.greens, mc.s.Ur, mc.s.Tr)
        end
    end

    s.greens
end

@bm function calculate_greens_full2!(mc, s, slice1, slice2)
    # stack = [0, Δτ, 2Δτ, ..., β] = [0, safe_mult, 2safe_mult, ... N]
    # @assert slice1 ≤ slice2

    # k ≤ l or slice1 ≤ slice2
    # B_{k+1}^-1 B_{k+2}^-1 ⋯ B_{l-1}^-1 B_l^-1
    compute_inverse_udt_block!(mc, s, slice1, slice2) # low high

    # B(slice1, 1) = B_k ⋯ B_1 = Ul Dl Tl
    compute_forward_udt_block!(mc, s, slice1)

    # B(M, slice2) = B_M ⋯ B_slice2+1 = (Ur Dr Tr)^† = Tr^† Dr^† Ur^†
    compute_backward_udt_block!(mc, s, slice2)

    # U D T remains valid here
    @bm "compute G" begin
        # [B_{l} B_{l-1} ⋯ B_{k+1} + (B_k ⋯ B_1 B_M ⋯ B_{l+1})^-1]^-1
        # [T^-1 D^-1 U^† + (Ul Dl Tl Tr^† Dr Ur^†)^-1]^-1
        # [T^-1 D^-1 U^† + Ur (Dl Tl Tr^† Dr)^-1 Ul^†]^-1
        @bm "B1" begin
            vmul!(s.greens, mc.s.Tl, adjoint(mc.s.Tr))
            vmul!(mc.s.tmp1, Diagonal(mc.s.Dl), s.greens)
            vmul!(s.greens, mc.s.tmp1, Diagonal(mc.s.Dr))
        end
        # [T^-1 D^-1 U^† + Ur G^-1 Ul^†]^-1
        @bm "udt" begin
            udt_AVX_pivot!(mc.s.Tr, mc.s.Dr, s.greens, mc.s.pivot, mc.s.tempv, Val(false))
        end
        # [T^-1 D^-1 U^† + Ur G^-1 Dr^-1 Tr^† Ul^†]^-1
        # [T^-1 D_min^-1 (D_max^-1 U^† (Ul Tr) Dr_min + D_min T Ur G^-1 Dr_max^-1) Dr_min^-1 (Ul Tr)^†]^-1
        @bm "B2" begin
            # D_max^-1 U^† (Ul Tr) Dr_min
            vmul!(mc.s.Tl, mc.s.Ul, mc.s.Tr) # keep me alive
            vmul!(mc.s.Ul, adjoint(s.U), mc.s.Tl)
            vmaxinv!(mc.s.Dl, s.D)
            vmul!(s.U, Diagonal(mc.s.Dl), mc.s.Ul)
            vmin!(mc.s.Dl, mc.s.Dr)
            vmul!(mc.s.Ul, s.U, Diagonal(mc.s.Dl))
        end
        # [T^-1 D_min^-1 (Ul + D_min T Ur G^-1 Dr_max^-1) Dr_min^-1 Tl^†]^-1
        @bm "B3" begin
            # D_min T Ur G^-1 Dr_max^-1
            vmul!(s.U, s.T, mc.s.Ur)
            rdivp!(s.U, s.greens, mc.s.Ur, mc.s.pivot)
            vmin!(mc.s.Dl, s.D)
            vmul!(mc.s.Ur, Diagonal(mc.s.Dl), s.U)
            vmaxinv!(mc.s.Dl, mc.s.Dr)
            vmul!(mc.s.Tr, mc.s.Ur, Diagonal(mc.s.Dl))
        end
        # [T^-1 D_min^-1 (Ul + Tr) Dr_min^-1 Tl^†]^-1
        @bm "sum, udt" begin
            rvadd!(mc.s.Tr, mc.s.Ul)
            udt_AVX_pivot!(mc.s.Ul, mc.s.Dl, mc.s.Tr, mc.s.pivot, mc.s.tempv, Val(false))
        end
        # [T^-1 D_min^-1 Ul Dl Tr Dr_min^-1 Tl^-1]^-1
        # Tl ({[(Dr_min / Tr) Dl^-1] Ul^†} D_min) T
        @bm "B4" begin
            vmin!(mc.s.Dr, mc.s.Dr)
            copyto!(s.U, Diagonal(mc.s.Dr))
            rdivp!(s.U, mc.s.Tr, mc.s.Ur, mc.s.pivot)
            vinv!(mc.s.Dl)
            vmul!(mc.s.Ur, s.U, Diagonal(mc.s.Dl))
            vmul!(s.U, mc.s.Ur, adjoint(mc.s.Ul))
            vmin!(s.D, s.D)
            vmul!(mc.s.Ur, s.U, Diagonal(s.D))
        end
        # Tl Ur T
        @bm "B6" begin
            vmul!(mc.s.Tr, mc.s.Ur, s.T)
            vmul!(s.greens, mc.s.Tl, mc.s.Tr)
            rmul!(s.greens, -1.0)
        end
    end

    s.greens
end

################################################################################
### Iterators
################################################################################

# The method above provides the unequal-time Greens function in a stable,
# general but expensive way. The idea of iterators is to provide what is needed
# more quickly, without the UnequalTimeStack and without having to worry about
# managing stack variables.

# Currently we have:
# GreensIterator
#   - attempts to be general (but isn't yet)
#   - requires UnequalTimeStack
#   - should have high accuracy, but still faster than manual
# CombinedGreensIterator
#   - does just l = 0, k = 0:M-1, but returns G_kl and G_ll
#   - does not require UnequalTimeStack
#   - same accuracy as GreensIterator

abstract type AbstractGreensIterator end
abstract type AbstractUnequalTimeGreensIterator <: AbstractGreensIterator end

init!(::AbstractUnequalTimeGreensIterator) = nothing
"""
    verify(iterator::AbstractUnequalTimeGreensIterator[, maxerror=1e-6])

Returns true if the given `iterator` is accurate up to the given `maxerror`.

See also: (@ref)[`accuracy`]
"""
verify(it::AbstractUnequalTimeGreensIterator, maxerror=1e-6) = maximum(accuracy(it)) < maxerror

# Maybe split into multiple types?
struct GreensIterator{slice1, slice2, T <: DQMC} <: AbstractUnequalTimeGreensIterator
    mc::T
    recalculate::Int64
end

"""
    GreensIterator(mc::DQMC[, ks=Colon(), ls=0, recalculate=4mc.p.safe_mult])

Creates an Iterator which calculates all `[G[k <- l] for l in ls for k in ks]`
efficiently. Currently only allows `ks = :, ls::Integer`.

For stability it is necessary to fully recompute `G[k <- l]` every so often. one
can adjust the frequency of recalculations via `recalculate`. To estimate the
resulting accuracy, one may use `accuracy(GreensIterator(...))`.

This iterator requires the `UnequalTimeStack`. Iteration overwrites the
`DQMCStack` variables `curr_U`, `Ul`, `Dl`, `Tl`, `Ur`, `Dr`, `Tr`, `tmp1` and
`tmp2`, with `curr_U` acting as the output. For the iteration to remain valid,
the UnequalTimeStack must not be overwritten. As such:
- `greens!(mc)` can be called and remains valid
- `greens!(mc, slice)` can be called and remains valid
- `greens!(mc, k, l)` will break iteration but remains valid (call before iterating)
"""
function GreensIterator(
        mc::T, slice1 = Colon(), slice2 = 0, recalculate = 4mc.p.safe_mult
    ) where {T <: DQMC}
    GreensIterator{slice1, slice2, T}(mc, recalculate)
end
init!(it::GreensIterator) = initialize_stack(it.mc, it.mc.ut_stack)
Base.length(it::GreensIterator{:, i}) where {i} = it.mc.p.slices + 1 - i

# Slower, versatile version:
function Base.iterate(it::GreensIterator{:, i}) where {i}
    s = it.mc.ut_stack
    calculate_greens_full1!(it.mc, s, i, i)
    copyto!(s.T, s.greens)
    udt_AVX_pivot!(s.U, s.D, s.T, it.mc.s.pivot, it.mc.s.tempv)
    G = _greens!(it.mc, it.mc.s.curr_U, s.greens, it.mc.s.Ur)
    return (G, (i+1, i))
end
function Base.iterate(it::GreensIterator{:}, state)
    s = it.mc.ut_stack
    k, l = state
    if k > it.mc.p.slices
        return nothing
    elseif k % it.recalculate == 0
        # Recalculate
        calculate_greens_full1!(it.mc, s, k, l) # writes s.greens
        G = _greens!(it.mc, it.mc.s.curr_U, s.greens, it.mc.s.Tl)
        copyto!(s.T, s.greens)
        udt_AVX_pivot!(s.U, s.D, s.T, it.mc.s.pivot, it.mc.s.tempv)
        return (G, (k+1, l))
    elseif k % it.mc.p.safe_mult == 0
        # Stabilization
        multiply_slice_matrix_left!(it.mc, it.mc.model, k, s.U)
        vmul!(it.mc.s.curr_U, s.U, Diagonal(s.D))
        vmul!(it.mc.s.tmp1, it.mc.s.curr_U, s.T)
        udt_AVX_pivot!(s.U, s.D, it.mc.s.curr_U, it.mc.s.pivot, it.mc.s.tempv)
        vmul!(it.mc.s.tmp2, it.mc.s.curr_U, s.T)
        copyto!(s.T, it.mc.s.tmp2)
        G = _greens!(it.mc, it.mc.s.curr_U, it.mc.s.tmp1, it.mc.s.tmp2)
        return (G, (k+1, l))
    else
        # Quick advance
        multiply_slice_matrix_left!(it.mc, it.mc.model, k, s.U)
        vmul!(it.mc.s.curr_U, s.U, Diagonal(s.D))
        vmul!(it.mc.s.tmp1, it.mc.s.curr_U, s.T)
        G = _greens!(it.mc, it.mc.s.curr_U, it.mc.s.tmp1, it.mc.s.Ur)
        s.last_slices = (-1, -1) # for safety
        return (G, (k+1, l))
    end
end
"""
    accuracy(iterator::AbstractUnequalTimeGreensIterator)

Compares values from the given iterator to more verbose computations, returning
the maximum differences for each. This can be used to check numerical stability.
"""
function accuracy(iter::GreensIterator{:, l}) where {l}
    mc = iter.mc
    Gk0s = [deepcopy(greens(mc, k, l)) for k in l:nslices(mc)]
    [maximum(abs.(Gk0s[i] .- G)) for (i, G) in enumerate(iter)]
end

"""
    CombinedGreensIterator(mc::DQMC, recalculate=4mc.p.safe_mult)

Returns an iterator which iterates `[(G[k, 0], G[k, k]) for k in 0:nslices-1]`.
Does a full recalculation of `G[k, 0]` and `G[k, k]` if `k % recalculate == 0`.

This iterator uses `s.tmp1` and `s.tmp2` from the normal `s::DQMCSTack` and
`uts.greens` from the `uts::UnequalTimeStack` as outputs. For correct iteration
it requires  `s.Ul`, `s.Dl`, `s.Tl`, `s.Ur`, `s.Dr`, `s.Tr`, `uts.U`, `uts.D`
and `uts.T` to remain unchanged. Further `s.curr_U` and `uts.tmp` are
overwritten. As such
- `greens!(mc)` will break iteration but remains valid (call before iterating)
- `greens!(mc, slice)` will break iteration but remains valid
- `greens!(mc, k, l)` will break iteration but remains valid
"""
struct CombinedGreensIterator{T <: DQMC} <: AbstractUnequalTimeGreensIterator
    mc::T
    recalculate::Int64
end

function CombinedGreensIterator(mc::T, model::Model, recalculate::Int64=4mc.p.safe_mult) where T
    CombinedGreensIterator{T}(mc, recalculate)
end
function CombinedGreensIterator(mc::T, recalculate::Int64=4mc.p.safe_mult) where T
    CombinedGreensIterator{T}(mc, recalculate)
end
function init!(it::CombinedGreensIterator)
    if it.recalculate < nslices(it.mc)
        initialize_stack(it.mc, it.mcm.ut_stack)
    end
end

Base.length(it::CombinedGreensIterator) = nslices(it.mc)

# Fast specialized version
function Base.iterate(it::CombinedGreensIterator)
    # Measurements take place at current_slice = 1 <> l = τ = 0
    @assert current_slice(it.mc) == 1
    s = it.mc.s
    uts = it.mc.ut_stack
    copyto!(s.Tl, s.greens)
    # Need full built stack for curr_U to be save
    build_stack(it.mc, uts)

    udt_AVX_pivot!(s.Ul, s.Dl, s.Tl, it.mc.s.pivot, it.mc.s.tempv)
    copyto!(uts.U, s.Ul)
    copyto!(uts.D, s.Dl)
    copyto!(uts.T, s.Tl)

    # G01 = (G00-I)B_1^-1; then G0l+1 = G0l B_{l+1}^-1
    vsub!(s.Tr, s.greens, I)
    udt_AVX_pivot!(s.Ur, s.Dr, s.Tr, it.mc.s.pivot, it.mc.s.tempv)
    return iterate(it, 1)
end
# probably need extra temp variables
function Base.iterate(it::CombinedGreensIterator, l)
    # l is 1-based
    s = it.mc.s
    uts = it.mc.ut_stack
    if l > it.mc.p.slices
        return nothing
    elseif l % it.recalculate == 0
        # Recalculation will overwrite
        # Ul, Dl, Tl, Ur, Dr, Tr, U, D, T, tmp1, tmp2, greens
        # curr_U is used only if the stack is rebuilt
        # this leaves uts.tmp (and curr_U)

        calculate_greens_full1!(it.mc, it.mc.ut_stack, l, 0)
        copyto!(s.curr_U, uts.greens)
        calculate_greens_full2!(it.mc, it.mc.ut_stack, 0, l)
        copyto!(uts.tmp, uts.greens)
        calculate_greens_full1!(it.mc, it.mc.ut_stack, l, l)

        # input and output can be the same here
        copyto!(uts.T, uts.greens)
        Gll = _greens!(it.mc, uts.greens, uts.T, s.tmp2)
        udt_AVX_pivot!(uts.U, uts.D, uts.T, s.pivot, s.tempv)

        Gl0 = _greens!(it.mc, s.tmp1, s.curr_U, s.tmp2)
        copyto!(s.Tl, s.curr_U)
        udt_AVX_pivot!(s.Ul, s.Dl, s.Tl, s.pivot, s.tempv)

        G0l = _greens!(it.mc, s.tmp2, uts.tmp, s.curr_U)
        copyto!(s.Tr, uts.tmp)
        udt_AVX_pivot!(s.Ur, s.Dr, s.Tr, s.pivot, s.tempv)

        return ((G0l, Gl0, Gll), l+1)

    elseif l % it.mc.p.safe_mult == 0
        # Stabilization
        # Reminder: These overwrite s.tmp1 and s.tmp2
        multiply_slice_matrix_left!(it.mc, it.mc.model, l, s.Ul)
        multiply_slice_matrix_inv_right!(it.mc, it.mc.model, l, s.Tr)
        multiply_slice_matrix_left!(it.mc, it.mc.model, l, uts.U)
        multiply_slice_matrix_inv_right!(it.mc, it.mc.model, l, uts.T)

        # Gl0
        vmul!(s.tmp1, s.Ul, Diagonal(s.Dl))
        vmul!(s.tmp2, s.tmp1, s.Tl)
        udt_AVX_pivot!(s.Ul, s.Dl, s.tmp1, s.pivot, s.tempv)
        vmul!(s.curr_U, s.tmp1, s.Tl)
        copyto!(s.Tl, s.curr_U)
        Gl0 = _greens!(it.mc, s.tmp1, s.tmp2, s.curr_U)
        # tmp1 must not change anymore

        # G0l
        vmul!(s.curr_U, Diagonal(s.Dr), s.Tr)
        vmul!(uts.greens, s.Ur, s.curr_U)
        copyto!(s.Tr, s.curr_U)
        udt_AVX_pivot!(s.tmp2, s.Dr, s.Tr, s.pivot, s.tempv)
        vmul!(s.curr_U, s.Ur, s.tmp2)
        copyto!(s.Ur, s.curr_U)
        G0l = _greens!(it.mc, s.tmp2, uts.greens, s.curr_U)
        # tmp1, tmp2 must not change anymore

        # Gll
        vmul!(uts.tmp, uts.U, Diagonal(uts.D))
        vmul!(uts.greens, uts.tmp, uts.T)
        udt_AVX_pivot!(s.curr_U, uts.D, uts.tmp, it.mc.s.pivot, it.mc.s.tempv)
        vmul!(uts.U, uts.tmp, uts.T)
        vmul!(uts.T, Diagonal(uts.D), uts.U)
        udt_AVX_pivot!(uts.tmp, uts.D, uts.T, it.mc.s.pivot, it.mc.s.tempv)
        vmul!(uts.U, s.curr_U, uts.tmp)
        Gll = _greens!(it.mc, uts.greens, uts.greens, s.curr_U)

        return ((G0l, Gl0, Gll), l+1)
    else
        # Quick advance
        multiply_slice_matrix_left!(it.mc, it.mc.model, l, s.Ul)
        # multiply_slice_matrix_inv_left!(it.mc, it.mc.model, nslices(it.mc)-l+1, s.Ur)
        multiply_slice_matrix_inv_right!(it.mc, it.mc.model, l, s.Tr)
        multiply_slice_matrix_left!(it.mc, it.mc.model, l, uts.U)
        multiply_slice_matrix_inv_right!(it.mc, it.mc.model, l, uts.T)

        # Gl0
        vmul!(s.curr_U, s.Ul, Diagonal(s.Dl))
        vmul!(s.tmp2, s.curr_U, s.Tl)
        Gl0 = _greens!(it.mc, s.tmp1, s.tmp2, s.curr_U)

        # G0l
        vmul!(s.curr_U, s.Ur, Diagonal(s.Dr))
        vmul!(uts.greens, s.curr_U, s.Tr)
        G0l = _greens!(it.mc, s.tmp2, uts.greens, s.curr_U)

        # Gll
        vmul!(s.curr_U, uts.U, Diagonal(uts.D))
        vmul!(uts.greens, s.curr_U, uts.T)
        Gll = _greens!(it.mc, uts.greens, uts.greens, s.curr_U)

        return ((G0l, Gl0, Gll), l+1)
    end
end

function accuracy(iter::CombinedGreensIterator)
    mc = iter.mc
    Gk0s = [deepcopy(greens(mc, k, 0)) for k in 0:nslices(mc)-1]
    [maximum(abs.(Gk0s[i] .- G)) for (i, G) in enumerate(iter)]
end


text:
using PyPlot
using Distributions

function plot_histogram(distribution, n)
    epsilon_values = rand(distribution, n)  # n draws from distribution
    plt[:hist](epsilon_values)
end

lp = Laplace()
lg = Logistic()
cc = Cauchy()
prt = Pareto()
#plot_histogram(lp, 500)
#plot_histogram(lg, 500)
#plot_histogram(cc, 500)
plot_histogram(prt, 100)
legend()


text:
using Weave

# convert_doc("LegendGeSim_Tutorial.ipynb", "LegendGeSim_Tutorial.jmd")

convert_doc("LegendGeSim_Tutorial.jmd", "LegendGeSim_Tutorial.jl")
convert_doc("LegendGeSim_Tutorial.jmd", "LegendGeSim_Tutorial.ipynb")


text:
using Test
using QuantumLattices.Prerequisites

@testset "decimaltostr" begin
    @test decimaltostr(:a) == ":a"
    @test decimaltostr(1) == "1"
    @test decimaltostr(10^6) == "1000000"
    @test decimaltostr(1//7) == "1//7"
    @test decimaltostr(1.0) == "1.0"
    @test decimaltostr(10.0^6) == "1.0e+06"
    @test decimaltostr(10^-5) == "1.0e-05"
    @test decimaltostr(1/7) == "0.14286"
    @test decimaltostr(1/7, 8) == "0.14285714"
    @test decimaltostr(0im) == "0.0"
    @test decimaltostr(1//7+0im) == "1//7"
    @test decimaltostr(1.0im) == "1.0im"
    @test decimaltostr(-1.0im) == "-1.0im"
    @test decimaltostr(0.1+0.12im) == "0.1+0.12im"
    @test decimaltostr(0.1-0.12im) == "0.1-0.12im"
end

@testset "ordinal" begin
    @test ordinal(1) == "1st"
    @test ordinal(2) == "2nd"
    @test ordinal(3) == "3rd"
    @test ordinal(4) == "4th"
    @test ordinal(5) == "5th"
end

@testset "delta" begin
    @test delta(1, 2) == 0
    @test delta(1, 1) == 1
end

@testset "searchsortedfirst" begin
    idx = CartesianIndices((2, 2, 2))
    for (i, index) in enumerate(idx)
        @test searchsortedfirst(idx, index) == i
    end
    @test searchsortedfirst(idx, CartesianIndex(0, 0, 0)) == 1
    @test searchsortedfirst(idx, CartesianIndex(3, 3, 3)) == 9
end


text:
let

    vec_A = Array(Matrix{Float64}, 2)
    vec_b = Array(Vector{Float64}, 2)
    vec_A[1] = [0.5,]'
    vec_A[2] = [2.0,]'
    vec_b[1] = [2.0]
    vec_b[2] = [0.5]
    mix_obs = MixtureModel([MvNormal([0.0], [1.0]), MvNormal([1.0], [1.0])],[0.5,0.5])
    mix_AgO = MixtureModel([MvNormal([NaN], [1.0]), MvNormal([NaN], [3.0])])

    gmr = GMR(vec_A, vec_b, mix_obs, mix_AgO)

    show(IOBuffer(), gmr)

    pred = gmr([0.0])
    weights = probs(pred.prior)
    @test isapprox(weights[1], 0.6225, atol=1e-4)
    @test isapprox(weights[2], 0.3775, atol=1e-4)
    @test isapprox(pred.components[1].μ[1], 2.0, atol=1e-8)
    @test isapprox(pred.components[2].μ[1], 0.5, atol=1e-8)
    @test full(pred.components[1].Σ)[1] == 1.0
    @test full(pred.components[2].Σ)[1] == 9.0

    mix = MixtureModel([
            MvNormal([0.0, 0.0], [1.0, 1.0]),
            MvNormal([1.0, 1.0], [1.0, 1.0]),
        ], [0.25,0.75]
        )
    gmr = GMR(mix, 1)
    @test isapprox(gmr.vec_A[1][1], 0.0, atol=1e-8)
    @test isapprox(gmr.vec_A[2][1], 0.0, atol=1e-8)
    @test isapprox(gmr.vec_b[1][1], 0.0, atol=1e-8)
    @test isapprox(gmr.vec_b[2][1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Act_given_Obs.components[1].Σ)[1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Act_given_Obs.components[2].Σ)[1], 1.0, atol=1e-8)
    @test probs(gmr.mixture_Obs.prior) == [0.25, 0.75]
    @test isapprox(gmr.mixture_Obs.components[1].μ[1], 0.0, atol=1e-8)
    @test isapprox(gmr.mixture_Obs.components[2].μ[1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Obs.components[1].Σ)[1], 1.0, atol=1e-8)
    @test isapprox(full(gmr.mixture_Obs.components[2].Σ)[1], 1.0, atol=1e-8)
end