text:
# ---
# title: 655. Print Binary Tree
# id: problem655
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Tree
# link: <https://leetcode.com/problems/print-binary-tree/description/>
# hidden: true
# ---
#
# Print a binary tree in an m*n 2D string array following these rules:
#
#   1. The row number `m` should be equal to the height of the given binary tree.
#   2. The column number `n` should always be an odd number.
#   3. The root node's value (in string format) should be put in the exactly middle of the first row it can be put. The column and the row where the root node belongs will separate the rest space into two parts ( **left-bottom part and right-bottom part** ). You should print the left subtree in the left-bottom part and print the right subtree in the right-bottom part. The left-bottom part and the right-bottom part should have the same size. Even if one subtree is none while the other is not, you don't need to print anything for the none subtree but still need to leave the space as large as that for the other subtree. However, if two subtrees are none, then you don't need to leave space for both of them.
#   4. Each unused space should contain an empty string `""`.
#   5. Print the subtrees following the same rules.
#
# **Example 1:**
#
#
#
#     Input:
#          1
#         /
#        2
#     Output:
#     [["", "1", ""],
#      ["2", "", ""]]
#
#
# **Example 2:**
#
#
#
#     Input:
#          1
#         / \
#        2   3
#         \
#          4
#     Output:
#     [["", "", "", "1", "", "", ""],
#      ["", "2", "", "", "", "3", ""],
#      ["", "", "4", "", "", "", ""]]
#
#
# **Example 3:**
#
#
#
#     Input:
#           1
#          / \
#         2   5
#        /
#       3
#      /
#     4
#     Output:
#
#     [["",  "",  "", "",  "", "", "", "1", "",  "",  "",  "",  "", "", ""]
#      ["",  "",  "", "2", "", "", "", "",  "",  "",  "",  "5", "", "", ""]
#      ["",  "3", "", "",  "", "", "", "",  "",  "",  "",  "",  "", "", ""]
#      ["4", "",  "", "",  "", "", "", "",  "",  "",  "",  "",  "", "", ""]]
#
#
# **Note:** The height of binary tree is in the range of [1, 10].
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
# Autogenerated wrapper script for CSDP_jll for x86_64-unknown-freebsd
export libcsdp

using OpenBLAS32_jll
using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("CSDP")
JLLWrappers.@declare_library_product(libcsdp, "libcsdp.so")
function __init__()
    JLLWrappers.@generate_init_header(OpenBLAS32_jll, CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libcsdp,
        "lib/libcsdp.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()

ext: jl


text:
#=
Given the mapping a = 1, b = 2, ... z = 26, and an encoded message, count the number of ways it can be decoded.

For example, the message '111' would give 3, since it could be decoded as 'aaa', 'ka', and 'ak'.

You can assume that the messages are decodable. For example, '001' is not allowed.
=#
using Test
include("Solutions/problem7_decodeways.jl")

@test count_ways(111) == 3
@test count_ways(1111111111) == 89

max_stars_repo_name: DominiqueCaron/daily-coding-problem

max_issues_repo_name: DominiqueCaron/daily-coding-problem

max_forks_repo_name: DominiqueCaron/daily-coding-problem


text:
push!(LOAD_PATH, pwd()*"/../../src/")
include("utils-logit.jl")

using Distributions
using PyPlot
using RDatasets
using DataFrames
using BVB
using Optim: optimize, LBFGS

df = dataset("boot", "nodal")
y = convert(Vector{Float64}, df[:R])
y = 2 * y - 1.
X = convert(Matrix{Float64}, delete!(df, :R))
@show n, p = size(X)

@show λ = 1.
ll(β) = ll_logistic(β, y, X; λ=λ)
approx_d = MvNormal(zeros(p), 100.0 * eye(p))

@show max_iter = 30
α_vec = zeros(max_iter)

iter_vec = collect(1:max_iter)
eval_interval = 2
eval_on = map(i->(i%eval_interval==1), iter_vec)
eval_vec = Vector{Int64}()
elbo_vec = Vector{Real}()

for iter in 1:max_iter
    if iter%20==0
        println(@sprintf "iter %d" iter)
    end
    # get new component by matching to functional gradient
    h = BVB.get_new_component_laplacian(ll, approx_d, diagonal=true, verbose=false)
    println("\niter $iter")
    @show h.μ
    @show h.Σ.mat
    # get rid of bad h
    all(isfinite.(h.μ)) && all(isfinite.(h.Σ.mat)) || continue
    # α = BVB.search_α(approx_d, h, ll, n=10000, verbose=true)
    α = BVB.search_α_exclusive_sgd(approx_d, h, ll, n=100, min_iter=20, verbose=true)
    approx_d = BVB.add_to_mixture(approx_d, h, α)

    # tracing
    α_vec[iter] = α
    if eval_on[iter]
        push!(eval_vec, iter)
        push!(elbo_vec, BVB.ELBO(approx_d, ll; n=10000))
        println(@sprintf("ELBO (Monte Carlo) = %g", elbo_vec[end]))
    end
end

figure()
subplot(2,2,1)
plot(eval_vec, elbo_vec)
xlabel("iter"); ylabel("ELBO")

subplot(2,2,2)
xlabel("iter"); ylabel("alpha")
vlines(iter_vec, 0, α_vec)

subplot(2,2,3)
plot(1:length(probs(approx_d)), sort(probs(approx_d), rev=true), ".-")
xlabel("component"); ylabel("weight")

writetable("bayes-logit-nodal-elbo.csv",
           DataFrame(iter = eval_vec, elbo = elbo_vec))
writetable("bayes-logit-nodal-alpha.csv",
           DataFrame(iter = 1:max_iter, elbo = α_vec))
@show writedlm("./bayes-logit-nodal.out", rand(approx_d, 4000)')

show()
println("Done.")

max_stars_repo_name: richardkwo/boosting-VB

max_issues_repo_name: richardkwo/boosting-VB

max_forks_repo_name: richardkwo/boosting-VB


text:
## activate project environment
using Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

## load packages
using LinearAlgebra, Printf, Statistics
using BenchmarkTools
using Trixi

## set up benchmark code
function format_result(result, description)
  text = @sprintf("    & \\SI{%4.1f \\pm %.1f}{\\ns}", time(mean(result)), time(std(result)))
  if !isempty(description)
    text *= "  % " * description
  end
  return text
end

@inline function rotated_precomputed_flux(numerical_flux, u_ll, u_rr, norm_, normal_vector, equations)
  u_ll_rotated = Trixi.rotate_to_x(u_ll, normal_vector, equations)
  u_rr_rotated = Trixi.rotate_to_x(u_rr, normal_vector, equations)

  f = numerical_flux(u_ll_rotated, u_rr_rotated, 1, equations)

  return Trixi.rotate_from_x(f, normal_vector, equations) * norm_
end

@inline function rotated_precomputed_flux(numerical_flux, u_ll, u_rr, norm_, normal_vector, tangent1, tangent2, equations)
  u_ll_rotated = Trixi.rotate_to_x(u_ll, normal_vector, tangent1, tangent2, equations)
  u_rr_rotated = Trixi.rotate_to_x(u_rr, normal_vector, tangent1, tangent2, equations)

  f = numerical_flux(u_ll_rotated, u_rr_rotated, 1, equations)

  return Trixi.rotate_from_x(f, normal_vector, tangent1, tangent2, equations) * norm_
end

function run_benchmarks_2d(numerical_flux, write_or_append="w")
  println("Benchmarking ", numerical_flux, " in 2D")

  equations = CompressibleEulerEquations2D(1.4)

  u_ll = SVector(1.0, 0.5, -0.7, 1.0)
  u_rr = SVector(1.5, -0.2, 0.1, 5.0)

  normal_direction = SVector(1.0, 0.0)

  sleep(0.5)
  f = numerical_flux
  result = @benchmark $f($(Ref(u_ll))[], $(Ref(u_rr))[], $1, $equations)
  format_result(result, "Cartesian") |> println
  open(joinpath(@__DIR__, "Cartesian_2D.txt"), write_or_append) do io
    println(io, format_result(result, "Cartesian"))
  end

  sleep(0.5)
  f = numerical_flux
  result = @benchmark $f($(Ref(u_ll))[], $(Ref(u_rr))[], $(Ref(normal_direction))[], $equations)
  format_result(result, "Directional") |> println
  open(joinpath(@__DIR__, "Directional_2D.txt"), write_or_append) do io
    println(io, format_result(result, "Directional"))
  end

  sleep(0.5)
  f = FluxRotated(numerical_flux)
  result = @benchmark $f($(Ref(u_ll))[], $(Ref(u_rr))[], $(Ref(normal_direction))[], $equations)
  format_result(result, "Rotated") |> println
  open(joinpath(@__DIR__, "Rotated_2D.txt"), write_or_append) do io
    println(io, format_result(result, "Rotated"))
  end

  # uses code from `FluxRotated`
  sleep(0.5)
  norm_ = norm(normal_direction)
  normal_vector = normal_direction / norm_
  result = @benchmark rotated_precomputed_flux($numerical_flux, $(Ref(u_ll))[], $(Ref(u_rr))[],
    $(Ref(norm_))[], $(Ref(normal_vector))[], $equations)
  format_result(result, "Rotated (precomputed)") |> println
  open(joinpath(@__DIR__, "Rotated_precomputed_2D.txt"), "w") do io
    println(io, format_result(result, "Rotated (precomputed)"))
  end

  return nothing
end

function run_benchmarks_3d(numerical_flux, write_or_append="w")
  println("Benchmarking ", numerical_flux, " in 3D")

  equations = CompressibleEulerEquations3D(1.4)

  u_ll = SVector(1.0, 0.5, -0.7, 0.1, 1.0)
  u_rr = SVector(1.5, -0.2, 0.1, -0.3, 5.0)

  normal_direction = SVector(1.0, 0.0, 0.0)

  sleep(0.5)
  f = numerical_flux
  result = @benchmark $f($(Ref(u_ll))[], $(Ref(u_rr))[], $1, $equations)
  format_result(result, "Cartesian") |> println
  open(joinpath(@__DIR__, "Cartesian_3D.txt"), write_or_append) do io
    println(io, format_result(result, "Cartesian"))
  end

  sleep(0.5)
  f = numerical_flux
  result = @benchmark $f($(Ref(u_ll))[], $(Ref(u_rr))[], $(Ref(normal_direction))[], $equations)
  format_result(result, "Directional") |> println
  open(joinpath(@__DIR__, "Directional_3D.txt"), write_or_append) do io
    println(io, format_result(result, "Directional"))
  end

  sleep(0.5)
  f = FluxRotated(numerical_flux)
  result = @benchmark $f($(Ref(u_ll))[], $(Ref(u_rr))[], $(Ref(normal_direction))[], $equations)
  format_result(result, "Rotated") |> println
  open(joinpath(@__DIR__, "Rotated_3D.txt"), write_or_append) do io
    println(io, format_result(result, "Rotated"))
  end

  # uses code from `FluxRotated`
  sleep(0.5)
  norm_ = norm(normal_direction)
  normal_vector = normal_direction / norm_
  tangent1 = SVector(normal_direction[2], normal_direction[3], -normal_direction[1])
  tangent1 -= dot(normal_vector, tangent1) * normal_vector
  tangent1 = normalize(tangent1)
  tangent2 = normalize(cross(normal_direction, tangent1))
  result = @benchmark rotated_precomputed_flux($numerical_flux, $(Ref(u_ll))[], $(Ref(u_rr))[],
    $(Ref(norm_))[], $(Ref(normal_vector))[], $(Ref(tangent1))[], $(Ref(tangent2))[], $equations)
  format_result(result, "Rotated (precomputed)") |> println
  open(joinpath(@__DIR__, "Rotated_precomputed_3D.txt"), write_or_append) do io
    println(io, format_result(result, "Rotated (precomputed)"))
  end

  return nothing
end

## run benchmarks
run_benchmarks_2d(flux_shima_etal, "w")
run_benchmarks_2d(flux_ranocha, "a")
run_benchmarks_2d(flux_lax_friedrichs, "a")
run_benchmarks_2d(flux_hll, "a")

run_benchmarks_3d(flux_shima_etal, "w")
run_benchmarks_3d(flux_ranocha, "a")
run_benchmarks_3d(flux_lax_friedrichs, "a")
run_benchmarks_3d(flux_hll, "a")


text:
using RBNF
using MLStyle
using PrettyPrint

@as_record Token
@as_record Struct_mainprogram
@as_record Struct_ifstmt
@as_record Struct_gate
@as_record Struct_gatedecl
@as_record Struct_decl
@as_record Struct_barrier_ids
@as_record Struct_reset
@as_record Struct_measure
@as_record Struct_iduop
@as_record Struct_u
@as_record Struct_cx
@as_record Struct_h
@as_record Struct_x
@as_record Struct_y
@as_record Struct_z
@as_record Struct_s
@as_record Struct_sdg
@as_record Struct_t
@as_record Struct_tdg
@as_record Struct_rx
@as_record Struct_ry
@as_record Struct_rz
@as_record Struct_ch
@as_record Struct_crx
@as_record Struct_cry
@as_record Struct_crz
@as_record Struct_idlist
@as_record Struct_mixedlist
@as_record Struct_argument
@as_record Struct_explist
@as_record Struct_pi
@as_record Struct_fnexp
@as_record Struct_neg
@as_record Struct_bin

function rec(ctx_tokens)
    function app(op, args...)
        args = map(rec, args)
        op = Symbol(op)
        :($op($(args...)))
    end

    @match ctx_tokens begin
        Struct_pi(_) => Base.pi
        Token{:id}(str=str) => Symbol(str)
        Token{:nnreal}(str=str) => parse(Float64, str)
        Token{:nninteger}(str=str) => parse(Int64, str)
        Struct_neg(value=value) => :(-$(rec(value)))
        Struct_bin(l=l, op=Token(str=op), r=r) => app(op, l, r)
        Struct_idlist(hd=Token(str=hd), tl=nothing) => [Symbol(hd)]
        Struct_idlist(hd=Token(str=hd), tl=tl) => [Symbol(hd), rec(tl)...]

        Struct_explist(hd=hd, tl=nothing) => [rec(hd)]
        Struct_explist(hd=hd, tl=tl) => [rec(hd), rec(tl)...]

        Struct_mixedlist(hd=hd, tl=nothing) => [rec(hd)]
        Struct_mixedlist(hd=hd, tl=tl) => [rec(hd), rec(tl)...]

        Struct_argument(id=Token(str=id), arg=nothing) => Symbol(id)
        Struct_argument(id=Token(str=id), arg=Token(str=int)) =>
            let ind = parse(Int, int) + 1 # due to julia 1-based index
                :($(Symbol(id))[$ind])
            end
        Struct_fnexp(fn = Token(str=fn), arg=arg) =>
            let fn = @match fn begin
                        "sin" => sin
                        "cos" => cos
                        "tan" => tan
                        "exp" => exp
                        "ln"  => log
                        "sqrt"=> sqrt
                        _     => error("not impl yet")
                    end
                app(fn, arg)
            end
        _ => nothing
    end
end

function trans_reg(ctx_tokens, qregs)
    function app(op, args...)
        args = map(rec, args)
        op = Symbol(op)
        :($op($(args...)))
    end

    @match ctx_tokens begin
        Struct_decl(
            regtype = Token(str=regtype),
            id = Token(str=id),
            int = Token(str = n)
        ) =>
            let id = Symbol(id),
                n = parse(Int, n)

                return :($id = qreg($n); push!($qregs, $id))
            end

        Struct_mainprogram(
            prog = stmts
        ) =>
            let stmts = trans_reg.(stmts, qregs)
                stmts
            end
        _ => nothing
    end
end

function trans_gates(ctx_tokens, qasm_cgc,  N)
    function app(op, args...)
        args = map(rec, args)
        op = Symbol(op)
        :($op($(args...)))
    end
    @match ctx_tokens begin
        Struct_iduop(gate_name = Token(str=gate_name), args=nothing, outs=outs) =>
            let refs = rec(outs),
                gate_name = Symbol("custom_gate_"*gate_name)
                :($gate_name(nothing, $(refs...); qasm_cgc=$qasm_cgc, N=$N))
            end

        Struct_iduop(gate_name = Token(str=gate_name), args=exprlist, outs=outs) =>
            let refs = rec(outs),
                exprs = Expr(:tuple, rec(exprlist)...),
                gate_name = Symbol("custom_gate_"*gate_name)
                :($gate_name($exprs, $(refs...); qasm_cgc=$qasm_cgc, N=$N))
            end

        Struct_cx(out1=out1, out2=out2) =>
            let ref1 = rec(out1),
                ref2 = rec(out2)
                :(append!($qasm_cgc, [circuit_gate(($ref2), X, ($ref1))]))
            end

        Struct_ch(out1=out1, out2=out2) =>
            let ref1 = rec(out1),
                ref2 = rec(out2)
                :(append!($qasm_cgc, [circuit_gate(($ref2), HadamardGate(), ($ref1))]))
            end

        Struct_u(in1=in1, in2=in2, in3=in3, out=out) =>
            let (a, b, c) = map(rec, (in1, in2, in3)),
                ref = :($(rec(out))[1])
                :(append!($qasm_cgc, [circuit_gate(($ref), RzGate($a)),
                   circuit_gate(($ref), RyGate($b)),
                   circuit_gate(($ref), RzGate($c)),]))
            end

        Struct_x(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), X)]))
            end

        Struct_y(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), Y)]))
            end

        Struct_z(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), Z)]))
            end

        Struct_h(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), HadamardGate())]))
            end

        Struct_t(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), TGate())]))
            end

        Struct_tdg(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), TdagGate())]))
            end

        Struct_s(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), SGate())]))
            end

        Struct_sdg(out=out) =>
            let ref = :($(rec(out)))
                :(append!($qasm_cgc, [circuit_gate(($ref), SdagGate())]))
            end

        Struct_rx(in=in, out=out) =>
            let ref = :($(rec(out))),
                arg = :($(rec(in)))
                :(append!($qasm_cgc, [circuit_gate(($ref), RxGate($arg))]))
            end

        Struct_ry(in=in, out=out) =>
            let ref = :($(rec(out))),
                arg = :($(rec(in)))
                :(append!($qasm_cgc, [circuit_gate(($ref), RyGate($arg))]))
            end

        Struct_rz(in=in, out=out) =>
            let ref = :($(rec(out))),
                arg = :($(rec(in)))
                :(append!($qasm_cgc, [circuit_gate(($ref), RzGate($arg))]))
            end

        Struct_crx(in=in, out1=out1, out2=out2) =>
            let out = :($(rec(out2))),
                cntrl = :($(rec(out1))),
                arg = :($(rec(in)))

                :(append!($qasm_cgc, [circuit_gate(($out), RxGate($arg), ($cntrl))]))
            end

        Struct_cry(in=in, out1=out1, out2=out2) =>
            let out = :($(rec(out2))),
                cntrl = :($(rec(out1))),
                arg = :($(rec(in)))

                :(append!($qasm_cgc, [circuit_gate(($out), RyGate($arg), ($cntrl))]))
            end

        Struct_crz(in=in, out1=out1, out2=out2) =>
            let out = :($(rec(out2))),
                cntrl = :($(rec(out1))),
                arg = :($(rec(in)))

                :(append!($qasm_cgc, [circuit_gate(($out), RzGate($arg), ($cntrl))]))
            end

        Struct_gate(
            decl = Struct_gatedecl(
                id=Token(str=fid),
                args=args,
                outs=outs
            ),
            goplist=goplist
         ) =>
            let out_ids :: Vector{Symbol} = rec(outs),
                fid = Symbol("custom_gate_"*fid),
                args = rec(args),
                goplist = trans_gates.(goplist, Ref(qasm_cgc), Ref(N))

                if isnothing(goplist)
                    goplist=[]
                end
                if isnothing(args)
                    args=[]
                end
                quote
                    function $fid(($(args...), ), $(out_ids...); qasm_cgc=$qasm_cgc, N=$N)
                        $(goplist...)
                    end
                end
            end

        # Struct_ifstmt(l=Token(str=l), r=r, gate_name=Token(str=gate_name), args=nothing, outs=outs) =>
        #     let l = Symbol(l),
        #         r = rec(r),
        #         refs = rec(outs),
        #         gate_name = Symbol("custom_gate_"*gate_name)

        #         :($gate_name(nothing, reg_check($l,$r), $(refs...); qasm_cgc=$qasm_cgc, N=$N))
        #     end

        # Struct_ifstmt(l=Token(str=l), r=r, gate_name=Token(str=gate_name), args=exprlist, outs=outs) =>
        #     let l = Symbol(l),
        #         r = rec(r),
        #         refs = rec(outs),
        #         exprs = Expr(:tuple, rec(exprlist)...),
        #         gate_name = Symbol("custom_gate_"*gate_name)

        #         :($gate_name($exprs, reg_check($l,$r), $(refs...); qasm_cgc=$qasm_cgc, N=$N))
        #     end

        Struct_mainprogram(
            prog = stmts
        ) =>
            let stmts = trans_gates.(stmts, qasm_cgc, N)
                stmts
            end
        _ => nothing
    end
end

function transform_qasm(ctx_tokens)

    qregs = Qaintessent.QRegister[]
    reg_declr = trans_reg(ctx_tokens, Ref(qregs))
    eval.(reg_declr)

    qasm_cgc = Circuit(qregs...)
    N = num_wires(qasm_cgc)

    gates_declr = trans_gates(ctx_tokens, Ref(qasm_cgc), Ref(N))

    eval.(gates_declr)

    qasm_cgc
end

"""
    qasm2cgc(txt::String)

converts OpenQASM 2.0 text to Circuit{N} object
"""
function qasm2cgc(txt::String)
    qasmlex = lex(txt)
    qasmparse = parse_qasm(qasmlex)
    qasm_cgc = transform_qasm(qasmparse)
end

function Base.:(==)(dcl1::Struct_gate, dcl2::Struct_gate)
    dcl1.decl == dcl2.decl && dcl1.goplist == dcl2.goplist
end


text:
using Compat

import Base:
    linearindexing, push!, size, sort, sort!, permute!, issorted, sortperm,
    summary, resize!, vcat, serialize, deserialize, append!, copy!, view

export Columns, colnames, ncols, ColDict, insertafter!, insertbefore!, @cols, setcol, pushcol, popcol, insertcol, insertcolafter, insertcolbefore, renamecol
export map_rows
export All, Not, Between, Keys

"""
A type that stores an array of tuples as a tuple of arrays.

# Fields:

- `columns`: a tuple or named tuples of arrays. Also `columns(x)`
"""
struct Columns{D<:Union{Tup, Pair}, C<:Union{Tup, Pair}} <: AbstractVector{D}
    columns::C

    function Columns{D,C}(c) where {D<:Tup,C<:Tup}
        length(c) > 0 || error("must have at least one column")
        n = length(c[1])
        for i = 2:length(c)
            length(c[i]) == n || error("all columns must have same length")
        end
        new{D,C}(c)
    end

    function Columns{D,C}(c::Pair) where {D<:Pair,C<:Pair{<:AbstractVector, <:AbstractVector}}
        length(c.first) == length(c.second) || error("all columns must have same length")
        new{D,C}(c)
    end
end

function Columns(cols::AbstractVector...; names::Union{Vector,Tuple{Vararg{Any}},Void}=nothing)
    if isa(names, Void) || any(x->!(x isa Symbol), names)
        Columns{eltypes(typeof(cols)),typeof(cols)}(cols)
    else
        dt = eval(:(@NT($(names...)))){map(eltype, cols)...}
        ct = eval(:(@NT($(names...)))){map(typeof, cols)...}
        Columns{dt,ct}(ct(cols...))
    end
end

Columns(; pairs...) = Columns(map(x->x[2],pairs)..., names=Symbol[x[1] for x in pairs])

Columns(c::Union{Tup, Pair}) = Columns{eltypes(typeof(c)),typeof(c)}(c)

# IndexedTable-like API

"""
    colnames(itr)

Returns the names of the "columns" in `itr`.

# Examples:

```jldoctest
julia> colnames([1,2,3])
1-element Array{Int64,1}:
 1

julia> colnames(Columns([1,2,3], [3,4,5]))
2-element Array{Int64,1}:
 1
 2

julia> colnames(table([1,2,3], [3,4,5]))
2-element Array{Int64,1}:
 1
 2

julia> colnames(Columns(x=[1,2,3], y=[3,4,5]))
2-element Array{Symbol,1}:
 :x
 :y

julia> colnames(table([1,2,3], [3,4,5], names=[:x,:y]))
2-element Array{Symbol,1}:
 :x
 :y

julia> colnames(ndsparse(Columns(x=[1,2,3]), Columns(y=[3,4,5])))
2-element Array{Symbol,1}:
 :x
 :y

julia> colnames(ndsparse(Columns(x=[1,2,3]), [3,4,5]))
2-element Array{Any,1}:
 :x
 1
 2

julia> colnames(ndsparse(Columns(x=[1,2,3]), [3,4,5]))
2-element Array{Any,1}:
 :x
 2

julia> colnames(ndsparse(Columns([1,2,3], [4,5,6]), Columns(x=[6,7,8])))
3-element Array{Any,1}:
 1
 2
 :x

julia> colnames(ndsparse(Columns(x=[1,2,3]), Columns([3,4,5],[6,7,8])))
3-element Array{Any,1}:
 :x
 2
 3
```
"""
function colnames end

Base.@pure colnames(t::AbstractVector) = [1]
columns(v::AbstractVector) = v

Base.@pure colnames(t::Columns) = fieldnames(eltype(t))
Base.@pure colnames(t::Columns{<:Pair, <:Pair}) = colnames(t.columns.first) => colnames(t.columns.second)

"""
`columns(itr[, select::Selection])`

Select one or more columns from an iterable of rows as a tuple of vectors.

`select` specifies which columns to select. See [`Selection convention`](@ref select) for possible values. If unspecified, returns all columns.

`itr` can be `NDSparse`, `Columns` and `AbstractVector`, and their distributed counterparts.

# Examples

```jldoctest
julia> t = table([1,2],[3,4], names=[:x,:y])
Table with 2 rows, 2 columns:
x  y
────
1  3
2  4

julia> columns(t)
(x = [1, 2], y = [3, 4])

julia> columns(t, :x)
2-element Array{Int64,1}:
 1
 2

julia> columns(t, (:x,))
(x = [1, 2])

julia> columns(t, (:y,:x=>-))
(y = [3, 4], x = [-1, -2])
```
"""
function columns end

columns(c) = error("no columns defined for $(typeof(c))")
columns(c::Columns) = c.columns

# Array-like API

eltype{D,C}(::Type{Columns{D,C}}) = D
length(c::Columns) = length(c.columns[1])
length(c::Columns{<:Pair, <:Pair}) = length(c.columns.first)
ndims(c::Columns) = 1

"""
`ncols(itr)`

Returns the number of columns in `itr`.

```jldoctest
julia> ncols([1,2,3])
1

julia> d = ncols(rows(([1,2,3],[4,5,6])))
2

julia> ncols(table(([1,2,3],[4,5,6])))
2

julia> ncols(table(@NT(x=[1,2,3],y=[4,5,6])))
2

julia> ncols(ndsparse(d, [7,8,9]))
3
```
"""
function ncols end
ncols(c::Columns) = nfields(typeof(c.columns))
ncols(c::Columns{<:Pair, <:Pair}) = ncols(c.columns.first) => ncols(c.columns.second)
ncols(c::AbstractArray) = 1

size(c::Columns) = (length(c),)
Base.IndexStyle(::Type{<:Columns}) = IndexLinear()
summary(c::Columns{D}) where {D<:Tuple} = "$(length(c))-element Columns{$D}"

empty!(c::Columns) = (foreach(empty!, c.columns); c)
empty!(c::Columns{<:Pair, <:Pair}) = (foreach(empty!, c.columns.first.columns); foreach(empty!, c.columns.second.columns); c)

function similar(c::Columns{D,C}) where {D,C}
    cols = _map(similar, c.columns)
    Columns{D,typeof(cols)}(cols)
end

function similar(c::Columns{D,C}, n::Integer) where {D,C}
    cols = _map(a->similar(a,n), c.columns)
    Columns{D,typeof(cols)}(cols)
end

function Base.similar{T<:Columns}(::Type{T}, n::Int)::T
    T_cols = T.parameters[2]
    f = T_cols <: Tuple ? tuple : T_cols
    T(f(map(t->similar(t, n), T.parameters[2].parameters)...))
end

function convert{N}(::Type{Columns}, x::AbstractArray{<:NTuple{N,Any}})
    eltypes = (eltype(x).parameters...)
    copy!(Columns(map(t->Vector{t}(length(x)), eltypes)), x)
end

function convert(::Type{Columns}, x::AbstractArray{<:NamedTuple})
    eltypes = (eltype(x).parameters...)
    copy!(Columns(map(t->Vector{t}(length(x)), eltypes)..., names=fieldnames(eltype(x))), x)
end

getindex(c::Columns{D}, i::Integer) where {D<:Tuple} = ith_all(i, c.columns)
getindex(c::Columns{D}, i::Integer) where {D<:NamedTuple} = D(ith_all(i, c.columns)...)
getindex(c::Columns{D}, i::Integer) where {D<:Pair} = getindex(c.columns.first, i) => getindex(c.columns.second, i)

getindex(c::Columns, p::AbstractVector) = Columns(_map(c->c[p], c.columns))

view(c::Columns, I) = Columns(_map(a->view(a, I), c.columns))

@inline setindex!(I::Columns, r::Union{Tup, Pair}, i::Integer) = (foreach((c,v)->(c[i]=v), I.columns, r); I)

@inline push!(I::Columns, r::Union{Tup, Pair}) = (foreach(push!, I.columns, r); I)

append!(I::Columns, J::Columns) = (foreach(append!, I.columns, J.columns); I)

copy!(I::Columns, J::Columns) = (foreach(copy!, I.columns, J.columns); I)

resize!(I::Columns, n::Int) = (foreach(c->resize!(c,n), I.columns); I)

_sizehint!(c::Columns, n::Integer) = (foreach(c->_sizehint!(c,n), c.columns); c)

function ==(x::Columns, y::Columns)
    nc = length(x.columns)
    length(y.columns) == nc || return false
    fieldnames(eltype(x)) == fieldnames(eltype(y)) || return false
    n = length(x)
    length(y) == n || return false
    for i in 1:nc
        x.columns[i] == y.columns[i] || return false
    end
    return true
end

==(x::Columns{<:Pair}, y::Columns) = false
==(x::Columns, y::Columns{<:Pair}) = false
==(x::Columns{<:Pair}, y::Columns{<:Pair}) = (x.columns.first == y.columns.first) && (x.columns.second == y.columns.second)

function _strip_pair(c::Columns{<:Pair})
    f, s = map(columns, c.columns)
    (f isa AbstractVector) && (f = (f,))
    (s isa AbstractVector) && (s = (s,))
    Columns(f..., s...)
end

function sortperm(c::Columns)
    cols = c.columns
    x = cols[1]
    if (eltype(x) <: AbstractString && !(x isa PooledArray)) || length(cols) > 1
        pa = PooledArray(x)
        p = sortperm_fast(pa)
    else
        p = sortperm_fast(x)
    end
    if length(cols) > 1
        y = cols[2]
        refine_perm!(p, cols, 1, x, y, 1, length(x))
    end
    return p
end

sortperm(c::Columns{<:Pair}) = sortperm(_strip_pair(c))

issorted(c::Columns) = issorted(1:length(c), lt=(x,y)->rowless(c, x, y))
issorted(c::Columns{<:Pair}) = issorted(_strip_pair(c))

# assuming x[p] is sorted, sort by remaining columns where x[p] is constant
function refine_perm!(p, cols, c, x, y, lo, hi)
    temp = similar(p, 0)
    order = Base.Order.By(j->(@inbounds k=y[j]; k))
    nc = length(cols)
    i = lo
    while i < hi
        i1 = i+1
        @inbounds while i1 <= hi && roweq(x, p[i1], p[i])
            i1 += 1
        end
        i1 -= 1
        if i1 > i
            sort_sub_by!(p, i, i1, y, order, temp)
            if c < nc-1
                z = cols[c+2]
                refine_perm!(p, cols, c+1, y, z, i, i1)
            end
        end
        i = i1+1
    end
end

function permute!(c::Columns, p::AbstractVector)
    for v in c.columns
        copy!(v, v[p])
    end
    return c
end
permute!(c::Columns{<:Pair}, p::AbstractVector) = (permute!(c.columns.first, p); permute!(c.columns.second, p); c)
sort!(c::Columns) = permute!(c, sortperm(c))
sort(c::Columns) = c[sortperm(c)]

function Base.vcat(c::Columns, cs::Columns...)
    fns = map(fieldnames, (map(x->x.columns, (c, cs...))))
    f1 = fns[1]
    for f2 in fns[2:end]
        if f1 != f2
            errfields = join(map(string, fns), ", ", " and ")
            throw(ArgumentError("Cannot concatenate columns with fields $errfields"))
        end
    end
    Columns(map(vcat, map(x->x.columns, (c,cs...))...))
end

function Base.vcat(c::Columns{<:Pair}, cs::Columns{<:Pair}...)
    Columns(vcat(c.columns.first, (x.columns.first for x in cs)...) =>
            vcat(c.columns.second, (x.columns.second for x in cs)...))
end

abstract type SerializedColumns end

function serialize(s::AbstractSerializer, c::Columns)
    Base.Serializer.serialize_type(s, SerializedColumns)
    serialize(s, eltype(c) <: NamedTuple)
    serialize(s, isa(c.columns, NamedTuple))
    serialize(s, fieldnames(c.columns))
    for col in c.columns
        serialize(s, col)
    end
end

function deserialize(s::AbstractSerializer, ::Type{SerializedColumns})
    Dnamed = deserialize(s)
    Cnamed = deserialize(s)
    fn = deserialize(s)
    cols = Any[ deserialize(s) for i = 1:length(fn) ]
    if Cnamed
        c = Columns(cols..., names = fn)
        if !Dnamed
            dt = eltypes(typeof((c.columns...,)))
            c = Columns{dt,typeof(c.columns)}(c.columns)
        end
    else
        c = Columns(cols...)
    end
    return c
end

# fused indexing operations
# these can be implemented for custom vector types like PooledVector where
# you can get big speedups by doing indexing and an operation in one step.

@inline cmpelts(a, i, j) = (@inbounds x=cmp(a[i], a[j]); x)
@inline copyelt!(a, i, j) = (@inbounds a[i] = a[j])
@inline copyelt!(a, i, b, j) = (@inbounds a[i] = b[j])

@inline cmpelts(a::PooledArray, i, j) = (x=cmp(a.refs[i],a.refs[j]); x)
@inline copyelt!(a::PooledArray, i, j) = (a.refs[i] = a.refs[j])

# row operations

copyrow!(I::Columns, i, src) = foreach(c->copyelt!(c, i, src), I.columns)
copyrow!(I::Columns, i, src::Columns, j) = foreach((c1,c2)->copyelt!(c1, i, c2, j), I.columns, src.columns)
copyrow!(I::AbstractArray, i, src::AbstractArray, j) = (@inbounds I[i] = src[j])
pushrow!(to::Columns, from::Columns, i) = foreach((a,b)->push!(a, b[i]), to.columns, from.columns)
pushrow!(to::AbstractArray, from::AbstractArray, i) = push!(to, from[i])

@generated function rowless{D,C}(c::Columns{D,C}, i, j)
    N = length(C.parameters)
    ex = :(cmpelts(getfield(c.columns,$N), i, j) < 0)
    for n in N-1:-1:1
        ex = quote
            let d = cmpelts(getfield(c.columns,$n), i, j)
                (d == 0) ? ($ex) : (d < 0)
            end
        end
    end
    ex
end

@generated function roweq{D,C}(c::Columns{D,C}, i, j)
    N = length(C.parameters)
    ex = :(cmpelts(getfield(c.columns,1), i, j) == 0)
    for n in 2:N
        ex = :(($ex) && (cmpelts(getfield(c.columns,$n), i, j)==0))
    end
    ex
end

@inline roweq(x::AbstractVector, i, j) = x[i] == x[j]

# uses number of columns from `d`, assuming `c` has more or equal
# dimensions, for broadcast joins.
@generated function rowcmp{D}(c::Columns, i, d::Columns{D}, j)
    N = length(D.parameters)
    ex = :(cmp(getfield(c.columns,$N)[i], getfield(d.columns,$N)[j]))
    for n in N-1:-1:1
        ex = quote
            let k = cmp(getfield(c.columns,$n)[i], getfield(d.columns,$n)[j])
                (k == 0) ? ($ex) : k
            end
        end
    end
    ex
end

@inline function rowcmp(c::AbstractVector, i, d::AbstractVector, j)
    cmp(c[i], d[j])
end

# test that the row on the right is "as of" the row on the left, i.e.
# all columns are equal except left >= right in last column.
# Could be generalized to some number of trailing columns, but I don't
# know whether that has applications.
@generated function row_asof{D,C}(c::Columns{D,C}, i, d::Columns{D,C}, j)
    N = length(C.parameters)
    if N == 1
        ex = :(!isless(getfield(c.columns,1)[i], getfield(d.columns,1)[j]))
    else
        ex = :(isequal(getfield(c.columns,1)[i], getfield(d.columns,1)[j]))
    end
    for n in 2:N
        if N == n
            ex = :(($ex) && !isless(getfield(c.columns,$n)[i], getfield(d.columns,$n)[j]))
        else
            ex = :(($ex) && isequal(getfield(c.columns,$n)[i], getfield(d.columns,$n)[j]))
        end
    end
    ex
end

# map

"""
`map_rows(f, c...)`

Transform collection `c` by applying `f` to each element. For multiple collection arguments, apply `f`
elementwise. Collect output as `Columns` if `f` returns
`Tuples` or `NamedTuples` with constant fields, as `Array` otherwise.

# Examples

```jldoctest map_rows
julia> map_rows(i -> @NT(exp = exp(i), log = log(i)), 1:5)
5-element IndexedTables.Columns{NamedTuples._NT_exp_log{Float64,Float64},NamedTuples._NT_exp_log{Array{Float64,1},Array{Float64,1}}}:
 (exp = 2.71828, log = 0.0)
 (exp = 7.38906, log = 0.693147)
 (exp = 20.0855, log = 1.09861)
 (exp = 54.5982, log = 1.38629)
 (exp = 148.413, log = 1.60944)
```
"""
function map_rows(f, iters...)
    collect_columns(f(i...) for i in zip(iters...))
end

# 1-arg case
map_rows(f, iter) = collect_columns(f(i) for i in iter)

## Special selectors to simplify column selector

"""
`All(cols)`

Select the union of the selections in `cols`. If `cols == ()`, select all columns.

# Examples

```jldoctest All
julia> t = table([1,1,2,2], [1,2,1,2], [1,2,3,4], [0, 0, 0, 0],
                        names=[:a,:b,:c,:d])
Table with 4 rows, 4 columns:
a  b  c  d
──────────
1  1  1  0
1  2  2  0
2  1  3  0
2  2  4  0

julia> select(t, All(:a, (:b, :c)))
Table with 4 rows, 3 columns:
a  b  c
───────
1  1  1
1  2  2
2  1  3
2  2  4

julia> select(t, All())
Table with 4 rows, 4 columns:
a  b  c  d
──────────
1  1  1  0
1  2  2  0
2  1  3  0
2  2  4  0
```
"""
struct All{T}
    cols::T
end

All(args...) = All(args)

"""
`Not(cols)`

Select the complementary of the selection in `cols`. `Not` can accept several arguments,
in which case it returns the complementary of the union of the selections.

# Examples

```jldoctest Not
julia> t = table([1,1,2,2], [1,2,1,2], [1,2,3,4],
                        names=[:a,:b,:c], pkey = (:a, :b))
Table with 4 rows, 3 columns:
a  b  c
───────
1  1  1
1  2  2
2  1  3
2  2  4

julia> select(t, Keys())
Table with 4 rows, 2 columns:
b  c
────
1  1
2  2
1  3
2  4

julia> select(t, Not(:a, (:a, :b)))
Table with 4 rows, 1 columns:
c
─
1
2
3
4
```
"""
struct Not{T}
    cols::T
end

Not(args...) = Not(All(args))

"""
`Keys()`

Select the primary keys.

# Examples

```jldoctest Keys
julia> t = table([1,1,2,2], [1,2,1,2], [1,2,3,4],
                               names=[:a,:b,:c], pkey = (:a, :b))
Table with 4 rows, 3 columns:
a  b  c
───────
1  1  1
1  2  2
2  1  3
2  2  4

julia> select(t, Keys())
Table with 4 rows, 2 columns:
a  b
────
1  1
1  2
2  1
2  2
```
"""
struct Keys; end

"""
`Between(first, last)`

Select the columns between `first` and `last`.

# Examples

```jldoctest Between
julia> t = table([1,1,2,2], [1,2,1,2], [1,2,3,4], ["a", "b", "c", "d"],
                                      names=[:a,:b,:c, :d])
Table with 4 rows, 4 columns:
a  b  c  d
────────────
1  1  1  "a"
1  2  2  "b"
2  1  3  "c"
2  2  4  "d"

julia> select(t, Between(:b, :d))
Table with 4 rows, 3 columns:
b  c  d
─────────
1  1  "a"
2  2  "b"
1  3  "c"
2  4  "d"
```
"""
struct Between{T1 <: Union{Int, Symbol}, T2 <: Union{Int, Symbol}}
    first::T1
    last::T2
end

const SpecialSelector = Union{Not, All, Keys, Between, Function, Regex}

hascolumns(t, s) = true
hascolumns(t, s::Symbol) = s in colnames(t)
hascolumns(t, s::Int) = s in 1:length(columns(t))
hascolumns(t, s::Tuple) = all(hascolumns(t, x) for x in s)
hascolumns(t, s::Not) = hascolumns(t, s.cols)
hascolumns(t, s::Between) = hascolumns(t, s.first) && hascolumns(t, s.last)
hascolumns(t, s::All) = all(hascolumns(t, x) for x in s.cols)

lowerselection(t, s)                     = s
lowerselection(t, s::Union{Int, Symbol}) = colindex(t, s)
lowerselection(t, s::Tuple)              = map(x -> lowerselection(t, x), s)
lowerselection(t, s::Not)                = excludecols(t, lowerselection(t, s.cols))
lowerselection(t, s::Keys)               = lowerselection(t, IndexedTables.pkeynames(t))
lowerselection(t, s::Between)            = Tuple(colindex(t, s.first):colindex(t, s.last))
lowerselection(t, s::Function)           = colindex(t, Tuple(filter(s, colnames(t))))
lowerselection(t, s::Regex)              = lowerselection(t, x -> ismatch(s, string(x)))

function lowerselection(t, s::All)
    s.cols == () && return lowerselection(t, valuenames(t))
    ls = (isa(i, Tuple) ? i : (i,) for i in lowerselection(t, s.cols))
    ls |> Iterators.flatten |> union |> Tuple
end

### Iteration API

# For `columns(t, names)` and `rows(t, ...)` to work, `t`
# needs to support `colnames` and `columns(t)`

Base.@pure function colindex(t, col::Tuple)
    fns = colnames(t)
    map(x -> _colindex(fns, x), col)
end

Base.@pure function colindex(t, col)
    _colindex(colnames(t), col)
end

function colindex(t, col::SpecialSelector)
    colindex(t, lowerselection(t, col))
end

function _colindex(fnames::AbstractArray, col, default=nothing)
    if isa(col, Int) && 1 <= col <= length(fnames)
        return col
    elseif isa(col, Symbol)
        idx = findfirst(fnames, col)
        idx > 0 && return idx
    elseif isa(col, Pair{<:Any, <:AbstractArray})
        return 0
    elseif isa(col, Tuple)
        return 0
    elseif isa(col, Pair{Symbol, <:Pair}) # recursive pairs
        return _colindex(fnames, col[2])
    elseif isa(col, Pair{<:Any, <:Any})
        return _colindex(fnames, col[1])
    elseif isa(col, AbstractArray)
        return 0
    end
    default !== nothing ? default : error("column $col not found.")
end

# const ColPicker = Union{Int, Symbol, Pair{Symbol=>Function}, Pair{Symbol=>AbstractVector}, AbstractVector}
column(c, x) = columns(c)[colindex(c, x)]

# optimized method
@inline function column(c::Columns, x::Union{Int, Symbol})
    getfield(c.columns, x)
end

column(t, a::AbstractArray) = a
column(t, a::Pair{Symbol, <:AbstractArray}) = column(t, a[2])
column(t, a::Pair{Symbol, <:Pair}) = rows(t, a[2]) # renaming a selection
column(t, a::Pair{<:Any, <:Any}) = map(a[2], rows(t, a[1]))
column(t, s::SpecialSelector) = rows(t, lowerselection(t, s))

function columns(c, sel::Union{Tuple, SpecialSelector})
    which = lowerselection(c, sel)
    cnames = colnames(c, which)
    if all(x->isa(x, Symbol), cnames)
        tuplewrap = namedtuple(cnames...)
    else
        tuplewrap = tuple
    end
    tuplewrap((rows(c, w) for w in which)...)
end

"""
`columns(itr, which)`

Returns a vector or a tuple of vectors from the iterator.

"""
columns(t, which) = column(t, which)

function colnames(c, cols::Union{Tuple, AbstractArray})
    map(x->colname(c, x), cols)
end

colnames(c, cols::SpecialSelector) = colnames(c, lowerselection(c, cols))

function colname(c, col)
    if isa(col, Union{Int, Symbol})
        col == 0 && return 0
        i = colindex(c, col)
        return colnames(c)[i]
    elseif isa(col, Pair{<:Any, <:Any})
        return col[1]
    elseif isa(col, Tuple)
        #ns = map(x->colname(c, x), col)
        return 0
    elseif isa(col, SpecialSelector)
        return 0
    elseif isa(col, AbstractVector)
        return 0
    end
    error("column named $col not found")
end

"""
`rows(itr[, select::Selection])`

Select one or more fields from an iterable of rows as a vector of their values.

`select` specifies which fields to select. See [`Selection convention`](@ref select) for possible values. If unspecified, returns all columns.

`itr` can be `NDSparse`, `Columns` and `AbstractVector`, and their distributed counterparts.

# Examples

```jldoctest
julia> t = table([1,2],[3,4], names=[:x,:y])
Table with 2 rows, 2 columns:
x  y
────
1  3
2  4

julia> rows(t)
2-element IndexedTables.Columns{NamedTuples._NT_x_y{Int64,Int64},NamedTuples._NT_x_y{Array{Int64,1},Array{Int64,1}}}:
 (x = 1, y = 3)
 (x = 2, y = 4)

julia> rows(t, :x)
2-element Array{Int64,1}:
 1
 2

julia> rows(t, (:x,))
2-element IndexedTables.Columns{NamedTuples._NT_x{Int64},NamedTuples._NT_x{Array{Int64,1}}}:
 (x = 1)
 (x = 2)

julia> rows(t, (:y,:x=>-))
2-element IndexedTables.Columns{NamedTuples._NT_y_x{Int64,Int64},NamedTuples._NT_y_x{Array{Int64,1},Array{Int64,1}}}:
 (y = 3, x = -1)
 (y = 4, x = -2)
```
Note that vectors of tuples returned are `Columns` object and have columnar internal storage.
"""
function rows end

rows(x::AbstractVector) = x
function rows(cols::Tup)
    if nfields(cols) === 0
        error("Cannot construct rows with 0 columns")
    else
        Columns(cols)
    end
end

rows(t, which...) = rows(columns(t, which...))

_cols_tuple(xs::Columns) = columns(xs)
_cols_tuple(xs::AbstractArray) = (xs,)
concat_cols(xs, ys) = rows(concat_tup(_cols_tuple(xs), _cols_tuple(ys)))

## Mutable Columns Dictionary

mutable struct ColDict{T}
    pkey::Vector{Int}
    src::T
    names::Vector
    columns::Vector
    copy::Union{Void, Bool}
end

"""
    d = ColDict(t)

Create a mutable dictionary of columns in `t`.

To get the immutable iterator of the same type as `t`
call `d[]`
"""
function ColDict(t; copy=nothing)
    ColDict(Int[], t, convert(Array{Any}, Base.copy(colnames(t))), Any[columns(t)...], copy)
end

function Base.getindex(d::ColDict{<:Columns})
    Columns(d.columns...; names=d.names)
end

Base.getindex(d::ColDict, key) = rows(d[], key)
Base.getindex(d::ColDict, key::AbstractArray) = key

function Base.setindex!(d::ColDict, x, key::Union{Symbol, Int})
    k = _colindex(d.names, key, 0)
    col = d[x]
    if k == 0
        push!(d.names, key)
        push!(d.columns, col)
    elseif k in d.pkey
        # primary key column has been modified.
        # copy the table as this results in a shuffle
        if d.copy === nothing
            d.copy = true
        end
        d.columns[k] = col
    else
        d.columns[k] = col
    end
end

set!(d::ColDict, key::Union{Symbol, Int}, x) = setindex!(d, x, key)

function Base.haskey(d::ColDict, key)
    _colindex(d.names, key, 0) != 0
end

function Base.insert!(d::ColDict, index, key, col)
    if haskey(d, key)
        error("Key $key already exists. Use dict[key] = col instead of inserting.")
    else
        insert!(d.names, index, key)
        insert!(d.columns, index, rows(d.src, col))
        for (i, pk) in enumerate(d.pkey)
            if pk >= index
                d.pkey[i] += 1 # moved right
            end
        end
    end
end

function insertafter!(d::ColDict, i, key, col)
    k = _colindex(d.names, i, 0)
    if k == 0
        error("$i not found. Cannot insert column after $i")
    end
    insert!(d, k+1, key, col)
end

function insertbefore!(d::ColDict, i, key, col)
    k = _colindex(d.names, i, 0)
    if k == 0
        error("$i not found. Cannot insert column after $i")
    end
    insert!(d, k, key, col)
end

function Base.pop!(d::ColDict, key::Union{Symbol, Int}=length(d.names))
    k = _colindex(d.names, key, 0)
    local col
    if k == 0
        error("Column $key not found")
    else
        col = d.columns[k]
        deleteat!(d.names, k)
        deleteat!(d.columns, k)
        idx = [pk[1] for pk in enumerate(d.pkey) if pk[2] == k]
        deleteat!(d.pkey, idx)
        for i in 1:length(d.pkey)
            if d.pkey[i] > k
                d.pkey[i] -= 1
            end
        end
        if !isempty(idx) && d.copy === nothing
            # set copy to true
            d.copy = true
        end
    end
    col
end

function rename!(d::ColDict, col::Union{Symbol, Int}, newname)
    k = _colindex(d.names, col, 0)
    if k == 0
        error("$col not found. Cannot rename it.")
    end
    d.names[k] = newname
end

function Base.push!(d::ColDict, key::Union{Symbol, Int}, x)
    push!(d.names, key)
    push!(d.columns, rows(d.src, x))
end

for s in [:(Base.pop!), :(Base.push!), :(rename!), :(set!)]
    if s == :(Base.pop!)
        typ = :(Union{Symbol, Int})
    else
        typ = :Pair
        @eval $s(t::ColDict, x::Pair) = $s(t, x.first, x.second)
    end
    @eval begin
        function $s(t::ColDict, args)
            for i in args
                $s(t, i)
            end
        end
        $s(t::ColDict, args::Vararg{$typ}) = $s(t, args)
    end
end

function _cols(expr)
    if expr.head == :call
        dict = :(dict = ColDict($(expr.args[2])))
        expr.args[2] = :dict
        quote
            let $dict
                $expr
                dict[]
            end
        end |> esc
    else
        error("This form of @cols is not implemented. Use `@cols f(t,args...)` where `t` is the collection.")
    end
end

macro cols(expr)
    _cols(expr)
end

# Modifying a columns

"""
`setcol(t::Table, col::Union{Symbol, Int}, x::Selection)`

Sets a `x` as the column identified by `col`. Returns a new table.

`setcol(t::Table, map::Pair...)`

Set many columns at a time.

# Examples:

```jldoctest setcol
julia> t = table([1,2], [3,4], names=[:x, :y])
Table with 2 rows, 2 columns:
x  y
────
1  3
2  4

julia> setcol(t, 2, [5,6])
Table with 2 rows, 2 columns:
x  y
────
1  5
2  6

```

`x` can be any selection that transforms existing columns.

```jldoctest setcol
julia> setcol(t, :x, :x => x->1/x)
Table with 2 rows, 2 columns:
x    y
──────
1.0  5
0.5  6

```

`setcol` will result in a re-sorted copy if a primary key column is replaced.

```jldoctest setcol
julia> t = table([0.01, 0.05], [1,2], [3,4], names=[:t, :x, :y], pkey=:t)
Table with 2 rows, 3 columns:
t     x  y
──────────
0.01  1  3
0.05  2  4

julia> t2 = setcol(t, :t, [0.1,0.05])
Table with 2 rows, 3 columns:
t     x  y
──────────
0.05  2  4
0.1   1  3

julia> t == t2
false

```

If `col` is not an existing column, `setcol` will add it:

```jldoctest setcol
julia> t = table([1,2], [2,3], names = [:a,:b])
Table with 2 rows, 2 columns:
a  b
────
1  2
2  3

julia> setcol(t, :c, [1,2])
Table with 2 rows, 3 columns:
a  b  c
───────
1  2  1
2  3  2
```
"""
setcol(t, args...) = @cols set!(t, args...)

"""
`pushcol(t, name, x)`

Push a column `x` to the end of the table. `name` is the name for the new column. Returns a new table.

`pushcol(t, map::Pair...)`

Push many columns at a time.

# Example:

```jldoctest
julia> t = table([0.01, 0.05], [2,1], [3,4], names=[:t, :x, :y], pkey=:t)
Table with 2 rows, 3 columns:
t     x  y
──────────
0.01  2  3
0.05  1  4

julia> pushcol(t, :z, [1//2, 3//4])
Table with 2 rows, 4 columns:
t     x  y  z
────────────────
0.01  2  3  1//2
0.05  1  4  3//4

```
"""
pushcol(t, args...) = @cols push!(t, args...)

"""
`popcol(t, col)`

Remove the column `col` from the table. Returns a new table.

`popcol(t, cols...)`

Remove many columns at a time.

```jldoctest
julia> t = table([0.01, 0.05], [2,1], [3,4], names=[:t, :x, :y], pkey=:t)
Table with 2 rows, 3 columns:
t     x  y
──────────
0.01  2  3
0.05  1  4

julia> popcol(t, :x)
Table with 2 rows, 2 columns:
t     y
───────
0.01  3
0.05  4
```
"""
popcol(t, args...) = @cols pop!(t, args...)

"""
`insertcol(t, position::Integer, name, x)`

Insert a column `x` named `name` at `position`. Returns a new table.

```jldoctest
julia> t = table([0.01, 0.05], [2,1], [3,4], names=[:t, :x, :y], pkey=:t)
Table with 2 rows, 3 columns:
t     x  y
──────────
0.01  2  3
0.05  1  4

julia> insertcol(t, 2, :w, [0,1])
Table with 2 rows, 4 columns:
t     w  x  y
─────────────
0.01  0  2  3
0.05  1  1  4

```
"""
insertcol(t, i::Integer, name, x) = @cols insert!(t, i, name, x)

"""
`insertcolafter(t, after, name, col)`

Insert a column `col` named `name` after `after`. Returns a new table.

```jldoctest
julia> t = table([0.01, 0.05], [2,1], [3,4], names=[:t, :x, :y], pkey=:t)
Table with 2 rows, 3 columns:
t     x  y
──────────
0.01  2  3
0.05  1  4

julia> insertcolafter(t, :t, :w, [0,1])
Table with 2 rows, 4 columns:
t     w  x  y
─────────────
0.01  0  2  3
0.05  1  1  4
```
"""
insertcolafter(t, after, name, x) = @cols insertafter!(t, after, name, x)

"""
`insertcolbefore(t, before, name, col)`

Insert a column `col` named `name` before `before`. Returns a new table.

```jldoctest
julia> t = table([0.01, 0.05], [2,1], [3,4], names=[:t, :x, :y], pkey=:t)
Table with 2 rows, 3 columns:
t     x  y
──────────
0.01  2  3
0.05  1  4

julia> insertcolbefore(t, :x, :w, [0,1])
Table with 2 rows, 4 columns:
t     w  x  y
─────────────
0.01  0  2  3
0.05  1  1  4
```
"""
insertcolbefore(t, before, name, x) = @cols insertbefore!(t, before, name, x)

"""
`renamecol(t, col, newname)`

Set `newname` as the new name for column `col` in `t`. Returns a new table.

`renamecol(t, map::Pair...)`

Rename many columns at a time.

```jldoctest
julia> t = table([0.01, 0.05], [2,1], names=[:t, :x])
Table with 2 rows, 2 columns:
t     x
───────
0.01  2
0.05  1

julia> renamecol(t, :t, :time)
Table with 2 rows, 2 columns:
time  x
───────
0.01  2
0.05  1
```
"""
renamecol(t, args...) = @cols rename!(t, args...)

## Utilities for mapping and reduction with many functions / OnlineStats

using OnlineStats

@inline _apply(f::OnlineStat, g, x) = (fit!(g, x); g)
@inline _apply(f::Tup, y::Tup, x::Tup) = map(_apply, f, y, x)
@inline _apply(f, y, x) = f(y, x)
@inline _apply(f::Tup, x::Tup) = map(_apply, f, x)
@inline _apply(f, x) = f(x)

@inline init_first(f, x) = x
@inline init_first(f::OnlineStat, x) = (g=copy(f); fit!(g, x); g)
@inline init_first(f::Tup, x::Tup) = map(init_first, f, x)

# Initialize type of output, functions to apply, input and output vectors

function reduced_type(f, x, isvec, key = nothing)
    if key !== nothing
        _promote_op(f, eltype(key), typeof(x))
    elseif isvec
        _promote_op(f, typeof(x))
    else
        _promote_op((a,b)->_apply(f, init_first(f, a), b),
                    eltype(x), eltype(x))
    end
end

function init_inputs(f, x, gettype, isvec) # normal functions
    f, x, gettype(f, x, isvec)
end

nicename(f) = Symbol(f)
nicename(o::OnlineStat) = Symbol(typeof(o).name.name)

function mapped_type(f, x, isvec)
    _promote_op(f, eltype(x))
end

init_funcs(f, isvec) = init_funcs((f,), isvec)

function init_funcs(f::Tup, isvec)
    if isa(f, NamedTuple)
        return init_funcs((map(Pair, fieldnames(f), f)...), isvec)
    end

    funcmap = map(f) do g
        if isa(g, Pair)
            name = g[1]
            if isa(g[2], Pair)
                sel, fn = g[2]
            else
                sel = nothing
                fn = g[2]
            end
            (name, sel, fn)
        else
            (nicename(g), nothing, g)
        end
    end

    ns = map(x->x[1], funcmap)
    ss = map(x->x[2], funcmap)
    fs = map(map(x->x[3], funcmap)) do f
        f
    end

    namedtuple(ns...)(fs...), ss
end

function init_inputs(f::Tup, input, gettype, isvec)
    if isa(f, NamedTuple)
        return init_inputs((map(Pair, fieldnames(f), f)...), input, gettype, isvec)
    end
    fs, selectors = init_funcs(f, isvec)

    xs = map(s->s === nothing ? input : rows(input, s), selectors)

    output_eltypes = map((f,x) -> gettype(f, x, isvec), fs, xs)

    ns = fieldnames(fs)
    NT = namedtuple(ns...)

    # functions, input, output_eltype
    NT(fs...), rows(NT(xs...)), NT{output_eltypes...}
end


text:
using ZChop

"""
    QFI(ρ, dρ [, abstol])

Numerically evaluate the quantum Fisher information for the matrix ρ given its derivative dρ wrt the parameter

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `ρ`:  Density matrix
    * `dhro`: Derivative wrt the parameter to be estimated
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI(ρ, dρ; abstol = 1e-5)

    # Get the eigenvalues and eigenvectors of the density matrix
    # We enforce its Hermiticity so that the algorithm is more efficient and returns real values

    eigval, eigvec = eigen(Hermitian(Matrix(zchop(ρ, 1e-10))))

    dim = length(eigval)
    tmpvec = Array{eltype(ρ)}(undef, size(ρ, 1))
    res = 0.
    tmp = 0.
    for m = 1:dim
        mul!(tmpvec, dρ, view(eigvec, :, m))
        for n = 1:dim
            tmp = eigval[n] + eigval[m]
            if tmp > abstol
                res += 2 * (1. / tmp) * abs2(dot(view(eigvec, :, n), tmpvec))
            end
        end
    end
    return res
end

"""
    QFI!(ρ, dρ [, abstol])

Numerically evaluate the quantum Fisher information for the matrix ρ given its derivative dρ wrt the parameter.
Can overwrite ρ.

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `ρ`:  Density matrix (gets overwritten)
    * `dhro`: Derivative wrt the parameter to be estimated
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI!(ρ, dρ; abstol = 1e-5)
    # Get the eigenvalues and eigenvectors of the density matrix
    # We enforce its Hermiticity so that the algorithm is more efficient and returns real values
    eigval, eigvec = eigen!(Hermitian(ρ))
    dim = length(eigval)
    tmpvec = Array{eltype(ρ)}(undef, size(ρ, 1))
    res = 0.
    tmp = 0.
    for m = 1:dim
        mul!(tmpvec, dρ, view(eigvec, :, m))
        for n = 1:dim
            tmp = eigval[n] + eigval[m]
            if tmp > abstol
                res += 2 * (1. / tmp) * abs2(dot(view(eigvec, :, n), tmpvec))
            end
        end
    end
    return res
end

function QFI(ρ::BlockDiagonal, dρ::BlockDiagonal; abstol = 1e-5)
    qfi = 0.
    for i in 1:nblocks(ρ)
        qfi += QFI(ρ.blocks[i], dρ.blocks[i])
    end
    return qfi
end

function QFI!(ρ::BlockDiagonal, dρ::BlockDiagonal; abstol = 1e-5)
    qfi = 0.
    for i in 1:nblocks(ρ)
        qfi += QFI!(ρ.blocks[i], dρ.blocks[i])
    end
    return qfi
end

"""
    QFI(state::State, [, abstol])

Numerically evaluate the quantum Fisher information for the State state.

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `state`: A State object
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI(state::State, abstol = 1e-5)
    QFI(state.ρ, state.dρ)
end

"""
    QFI!(state::State, [, abstol])

Numerically evaluate the quantum Fisher information for the State state,
overwriting one of its temporary variables.

This function is the implementation of Eq. (13) in Paris, Int. J. Quantum Inform. 7, 125 (2009).

# Arguments
    * `state`: A State object
    * `abstol = 1e-5`: tolerance in the denominator of the formula
"""
function QFI!(state::State, abstol = 1e-5)
    copy!(state._tmp1, state.ρ)
    QFI!(state._tmp1, state.dρ)
end

"""
    FI(ρ, dρ, op, [, tol])

Evaluates the Fisher information for the measurement operator op, for density matrix ρ, with
derivative dρ with respect to the parameter.

tol gives the tolerance to treat the denominator as 0 (and avoid NaNs). Default is `tol=1e-10`.
"""
function FI(ρ::BlockDiagonal, dρ::BlockDiagonal, op::BlockDiagonal)
    ed = eigen(op)
    return FI(ρ, dρ, ed)
end

function FI(ρ::Matrix, dρ::Matrix, measurement::Eigen; tol=1e-10)
    # We use Eq. (4) of Paris, Int. J. Quantum Inf. 7, 125 (2009)
    # We need the projectors onto the subspaces of the operator op
    fi = 0.0
    for i in 1:size(measurement.vectors, 1)
        evec = measurement.vectors[:, i]
        denominator = real(evec' * ρ * evec)

        # Empty blocks would cause a division by 0
        # We use a tolerance
        if denominator > tol
            fi += real(evec' * dρ * evec)^2 / denominator
        end
    end
    return fi
end

function FI(ρ::BlockDiagonal, dρ::BlockDiagonal, measurement::Array{Eigen}; tol=1e-10)
    # We use Eq. (4) of Paris, Int. J. Quantum Inf. 7, 125 (2009)
    # We need the projectors onto the subspaces of the operator op
    fi = 0.0
    for i in 1:nblocks(ρ)
        fi += FI(ρ.blocks[i], dρ.blocks[i], measurement[i], tol=tol)
    end
    return fi
end

FI(state::State, measurement::Eigen) = FI(state.ρ, state.dρ, measurement)
FI(state::State, measurement::Array{Eigen}) = FI(state.ρ, state.dρ, measurement)


text:
@testset "Scheduler Checkpointing" begin
    ctx = Context([1,workers()...])
    d = Ref(false)
    a = delayed(()->begin
        d[] = true
        1
    end)()
    f = Ref(0)
    opts = Dagger.Sch.SchedulerOptions(;
    single=1,
    checkpoint=(result)->begin
        f[] = collect(result)
    end)
    @test collect(ctx, a; options=opts) == 1
    @test d[]
    @test f[] == 1

    opts = Dagger.Sch.SchedulerOptions(;
    single=1,
    restore=()->begin
        Dagger.tochunk(f[])
    end)
    f[] = 2
    d[] = false
    @test collect(ctx, a; options=opts) == 2
    @test !(d[]) # thunks aren't run again
    @test f[] == 2 # checkpoint isn't regenerated

    @testset "Checkpoint Failure" begin
        d = Ref(false)
        a = delayed(()->begin
            d[] = true
            1
        end)()
        e = Ref(false)
        opts = Dagger.Sch.SchedulerOptions(;
        single=1,
        checkpoint=(result)->begin
            e[] = true
            error("Test")
        end)
        _, err = @grab_output begin
            @test collect(ctx, a; options=opts) == 1
        end
        @test occursin("Scheduler checkpoint failed", err)
        @test d[] # thunks executed
        @test e[] # checkpoint executed
    end
    @testset "Restore Failure" begin
        d = Ref(false)
        a = delayed(()->begin
            d[] = true
            1
        end)()
        e = Ref(false)
        opts = Dagger.Sch.SchedulerOptions(;
        single=1,
        restore=()->begin
            e[] = true
            error("Test")
        end)
        _, err = @grab_output begin
            @test collect(ctx, a; options=opts) == 1
        end
        @test occursin("Scheduler restore failed", err)
        @test occursin("Test", err)
        @test d[] # thunks executed
        @test e[] # restore executed
    end
    @testset "Restore Failure (quiet)" begin
        a = delayed(()->begin
            1
        end)()
        opts = Dagger.Sch.SchedulerOptions(;
        single=1,
        restore=()->begin
            nothing
        end)
        _, err = @grab_output begin
            @test collect(ctx, a; options=opts) == 1
        end
        @test !occursin("Scheduler restore failed", err)
    end
end

@testset "Thunk Checkpointing" begin
    ctx = Context([1,workers()...])
    d = Ref(false)
    opts = Dagger.Sch.ThunkOptions(;
    single=1,
    checkpoint=(thunk, result)->begin
        @assert thunk.f != Base.:*
        f[] = collect(result)
    end)
    a = delayed(()->begin
        d[] = true
        1
    end; options=opts)()
    b = delayed(*)(a,2)
    f = Ref(0)
    @test collect(ctx, b) == 2
    @test d[]
    @test f[] == 1

    opts = Dagger.Sch.ThunkOptions(;
    single=1,
    restore=(thunk)->begin
        @assert thunk.f != Base.:*
        Dagger.tochunk(f[])
    end)
    a = delayed(()->begin
        d[] = true
        1
    end; options=opts)()
    b = delayed(*)(a,2)
    f[] = 2
    d[] = false
    @test collect(ctx, b) == 4 # thunk b is run
    @test !(d[]) # thunk a isn't run again
    @test f[] == 2 # checkpoint isn't regenerated

    @testset "Checkpoint Failure" begin
        d = Ref(false)
        e = Ref(false)
        opts = Dagger.Sch.ThunkOptions(;
        single=1,
        checkpoint=(thunk, result)->begin
            e[] = true
            error("Test")
        end)
        a = delayed(()->begin
            d[] = true
            1
        end; options=opts)()
        _, err = @grab_output begin
            @test collect(ctx, a) == 1
        end
        @test occursin("Thunk checkpoint failed", err)
        @test d[] # thunk executed
        @test e[] # checkpoint executed
    end
    @testset "Restore Failure" begin
        d = Ref(false)
        e = Ref(false)
        opts = Dagger.Sch.ThunkOptions(;
        single=1,
        restore=(thunk)->begin
            e[] = true
            error("Test")
        end)
        a = delayed(()->begin
            d[] = true
            1
        end; options=opts)()
        _, err = @grab_output begin
            @test collect(ctx, a) == 1
        end
        @test occursin("Thunk restore failed", err)
        @test occursin("Test", err)
        @test d[] # thunk executed
        @test e[] # restore executed
    end
    @testset "Restore Failure (quiet)" begin
        opts = Dagger.Sch.ThunkOptions(;
        single=1,
        restore=(thunk)->begin
            nothing
        end)
        a = delayed(()->begin
            1
        end; options=opts)()
        _, err = @grab_output begin
            @test collect(ctx, a) == 1
        end
        @test !occursin("Thunk restore failed", err)
    end
end


text:
@testset "Manopt Cyclic Proximal Point" begin
    using Dates
    n = 100
    N = PowerManifold(Circle(), n)
    f = artificial_S1_signal(n)
    F = x -> costL2TV(N, f, 0.5, x)
    proxes = ((λ, x) -> prox_distance(N, λ, f, x), (λ, x) -> prox_TV(N, 0.5 * λ, x))
    o = cyclic_proximal_point(
        N,
        F,
        proxes,
        f;
        λ=i -> π / (2 * i),
        stopping_criterion=StopAfterIteration(100),
        debug=[
            DebugIterate(), " ", DebugCost(), " ", DebugProximalParameter(), "\n", 10000
        ],
        record=[RecordProximalParameter(), RecordIterate(f), RecordCost()],
        return_options=true,
    )
    fR = get_solver_result(o)
    fR2 = cyclic_proximal_point(
        N, F, proxes, f; λ=i -> π / (2 * i), stopping_criterion=StopAfterIteration(100)
    )
    @test fR == fR2
    rec = get_record(o)
    @test F(f) > F(fR)
    #
    o = CyclicProximalPointOptions(f, StopAfterIteration(1), i -> π / (2 * i))
    p = ProximalProblem(N, F, proxes, [1, 2])
    @test_throws ErrorException get_proximal_map(p, 1.0, f, 3)
    @test_throws ErrorException ProximalProblem(N, F, proxes, [1, 2, 2])
    order1 = collect(1:3)
    Manopt.update_cpp_order!(order1, 3, 0, RandomEvalOrder())
    @test all(isinteger.(order1))
    @test minimum(order1) == 1
    @test maximum(order1) == 3

    # i=0 introduce new fixed ranom order
    order2 = collect(1:3)
    Manopt.update_cpp_order!(order2, 3, 0, FixedRandomEvalOrder())
    @test all(isinteger.(order2))
    @test minimum(order2) == 1
    @test maximum(order2) == 3
    order3 = copy(order2)
    Manopt.update_cpp_order!(order3, 3, 1, FixedRandomEvalOrder())
    @test order3 == order2 # only update on i=0

    @test Manopt.update_cpp_order!(collect(1:3), 3, 0, LinearEvalOrder()) == 1:3
    @test Manopt.update_cpp_order!(collect(1:3), 3, 1, LinearEvalOrder()) == 1:3
end


text:
module IterativeSolvers

using LinearAlgebra
using Parameters
using Printf

import LinearAlgebra: ldiv!

export Continuous, IterativeSolver, ConjugateGradient, BiCGStab, GMRES, solve!

"""
Allows the identity operator `I` to be used as a preconditioner.
"""
function ldiv!(vout::AbstractVector,I::UniformScaling,vin::AbstractVector)
    copyto!(vout,vin)
    return nothing
end

"""
Abstract type to represent continuous real or complex numbers.
"""
Continuous = Union{AbstractFloat,Complex{<:AbstractFloat}}

"""
Abstract type to represent Iterative Solver methods.
"""
abstract type IterativeSolver{Ttol<:AbstractFloat,Tdata<:Continuous} end

#######################################
## Conjugate Gradient Implementation ##
#######################################

"""
To represent Conjugate Gradient algorithm.
"""
mutable struct ConjugateGradient{Ttol,Tdata} <: IterativeSolver{Ttol,Tdata}

    tol::Ttol
    maxiter::Int
    κmax::Ttol
    N::Int
    r::Vector{Tdata}
    p::Vector{Tdata}
    z::Vector{Tdata}

    function ConjugateGradient(z::AbstractVector{Tdata};tol::Ttol=1e-4,maxiter::Int=0,κmax::Ttol=1e12) where {Ttol<:AbstractFloat,Tdata<:Continuous}

        N = length(z)
        if maxiter<1
            maxiter = N
        end
        r = zeros(Tdata,N)
        p = zeros(Tdata,N)
        z = zeros(Tdata,N)
        return new{Ttol,Tdata}(tol,maxiter,κmax,N,r,p,z)
    end
end

"""
Solve `A⋅x=b` using the Conjugate Gradient method with a split preconditioner
such that `[L⁻¹⋅A⋅L⁻ᵀ]⋅u=L⁻¹⋅b` where `u=L⁻ᵀ⋅x`.
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},cg::ConjugateGradient{Ttol,Tdata},L,Lt;
                maxiter::Int=0,tol::Ttol=0.0,κmax::Ttol=0.0)::Int where {Ttol,Tdata}

    r = cg.r
    p = cg.p
    z = cg.z

    # set to default maxiter
    if iszero(maxiter)
        maxiter = cg.maxiter
    end

    # set to default tolerance
    if iszero(tol)
        tol = cg.tol
    end

    # set default max condition number
    if iszero(κmax)
        κmax = cg.κmax
    end

    # r₀ = b - A⋅x₀
    mul!(r,A,x)
    axpby!(1.0,b,-1.0,r)

    # r₀ = L⁻¹⋅r₀ = L⁻¹⋅(b - A⋅x₀)
    ldiv!(L,r)

    # p₀ = L⁻ᵀ⋅r₀ = L⁻ᵀ⋅L⁻¹⋅(b - A⋅x₀)
    ldiv!(p,Lt,r)

    # |L⁻ᵀ⋅L⁻¹⋅b|
    ldiv!(z,L,b)
    ldiv!(Lt,z)
    normL⁻ᵀL⁻¹b = norm(z)

    # calcualte initial tolerance
    ϵ₀ = norm(p)/normL⁻ᵀL⁻¹b
    ϵ  = ϵ₀

    # initial lower bound for condition number
    κmin  = 0.0

    # r₀⋅r₀
    rdotr = dot(r,r)

    @fastmath @inbounds for j in 1:maxiter

        # αⱼ = (rⱼ⋅rⱼ)/(pⱼ⋅A⋅pⱼ)
        mul!(z,A,p)
        α = rdotr/dot(p,z)

        # xⱼ₊₁ = xⱼ + αⱼ⋅pⱼ
        axpy!(α,p,x)

        # rⱼ₊₁ = rⱼ - αⱼ⋅L⁻¹⋅A⋅pⱼ
        ldiv!(L,z)
        axpy!(-α,z,r)

        # βⱼ = (rⱼ₊₁⋅rⱼ₊₁)/(rⱼ⋅rⱼ)
        new_rdotr = dot(r,r)
        β     = new_rdotr/rdotr
        rdotr = new_rdotr

        # pⱼ₊₁ = L⁻ᵀ⋅rⱼ₊₁ + βⱼ⋅pⱼ
        ldiv!(z,Lt,r)
        axpby!(1.0,z,β,p)

        # ϵ = |pⱼ₊₁|/|L⁻ᵀ⋅L⋅b| = |L⁻ᵀ⋅L⋅(A⋅xⱼ₊₁-b)|/|L⁻ᵀ⋅L⋅b|
        ϵ = norm(p)/normL⁻ᵀL⁻¹b

        # approximate lower bound for condition numbers
        κmin = max( κmin , (2*j/log(2*ϵ₀/ϵ))^2 )

        # check stop criteria
        if ϵ < tol  || κmin > κmax
            return j
        end
    end

    return maxiter

end

"""
Solve `A⋅x=b` using the Conjugate Gradient method with a preconditioner `P`.
Based on pseudocode from: https://www-users.cs.umn.edu/~saad/Calais/PREC.pdf
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},cg::ConjugateGradient{Ttol,Tdata},P;
                maxiter::Int=0,tol::Ttol=0.0,κmax::Ttol=0.0)::Int where {Ttol,Tdata}

    r = cg.r
    p = cg.p
    z = cg.z

    # set to default maxiter
    if iszero(maxiter)
        maxiter = cg.maxiter
    end

    # set to default tolerance
    if iszero(tol)
        tol = cg.tol
    end

    # set default max condition number
    if iszero(κmax)
        κmax = cg.κmax
    end

    # |b|
    normb = norm(b)

    # r₀ = b - A⋅x₀
    mul!(r,A,x)
    axpby!(1.0,b,-1.0,r)

    # z₀ = P⁻¹⋅r₀ = P⁻¹⋅(b - A⋅x₀)
    ldiv!(z,P,r)

    # p₀ = z₀
    copyto!(p,z)

    # r₀⋅z₀
    rdotz = dot(r,z)

    # calcualte initial tolerance
    ϵ₀ = norm(r)/normb
    ϵ  = ϵ₀

    # initial lower bound for condition number
    κmin  = 0.0

    @fastmath @inbounds for j in 1:maxiter

        # αⱼ = (rⱼ⋅zⱼ)/(pⱼ⋅A⋅pⱼ)
        mul!(z,A,p)
        α = rdotz/dot(p,z)

        # xⱼ₊₁ = xⱼ + αⱼ⋅pⱼ
        axpy!(α,p,x)

        # rⱼ₊₁ = rⱼ - αⱼ⋅A⋅pⱼ
        axpy!(-α,z,r)

        # ϵ = |rⱼ₊₁|/|b|
        ϵ = norm(r)/normb

        # approximate lower bound for condition numbers
        κmin = max( κmin , (2*j/log(2*ϵ₀/ϵ))^2 )

        # check stop criteria
        if ϵ < tol  || κmin > κmax
            return j
        end

        # zⱼ₊₁ = P⁻¹⋅rⱼ₊₁
        ldiv!(z,P,r)

        # βⱼ = (rⱼ₊₁⋅zⱼ₊₁)/(rⱼ⋅zⱼ)
        new_rdotz = dot(r,z)
        β     = new_rdotz/rdotz
        rdotz = new_rdotz

        # pⱼ₊₁ = zⱼ₊₁ + βⱼ⋅pⱼ
        axpby!(1.0,z,β,p)
    end

    return maxiter
end

"""
Solve `A⋅x=b` using the Conjugate Gradient method with no preconditioning.
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},cg::ConjugateGradient{Ttol,Tdata};
                maxiter::Int=0,tol::Ttol=0.0,κmax::Ttol=0.0)::Int where {Ttol,Tdata}

    r = cg.r
    p = cg.p
    z = cg.z

    # set to default maxiter
    if iszero(maxiter)
        maxiter = cg.maxiter
    end

    # set to default tol
    if iszero(tol) || κmin > κmax
        tol = cg.tol
    end

    # set default max condition number
    if iszero(κmax)
        κmax = cg.κmax
    end

    # |b|
    normb = norm(b)

    # r₀ = b - A⋅x₀
    mul!(r,A,x)
    axpby!(1.0,b,-1.0,r)

    # p₀ = r₀
    copyto!(p,r)

    # r₀⋅r₀
    rdotr = dot(r,r)

    # calcualte initial tolerance
    ϵ₀ = norm(r)/normb
    ϵ  = ϵ₀

    # initial lower bound for condition number
    κmin = 0.0

    @fastmath @inbounds for j in 1:maxiter

        # αⱼ = (rⱼ⋅rⱼ)/(pⱼ⋅A⋅pⱼ)
        mul!(z,A,p)
        α = rdotr/dot(p,z)

        # xⱼ₊₁ = xⱼ + αⱼ⋅pⱼ
        axpy!(α,p,x)

        # rⱼ₊₁ = rⱼ - αⱼ⋅A⋅pⱼ
        axpy!(-α,z,r)

        # ϵ = |rⱼ₊₁|/|b| = |b-A⋅xⱼ₊₁|/|b|
        ϵ = norm(r)/normb

        # approximate lower bound for condition numbers
        κmin = max( κmin , (2*j/log(2*ϵ₀/ϵ))^2 )

        # check stop criteria
        if ϵ < tol || κmin > κmax
            return j
        end

        # βⱼ = (rⱼ₊₁⋅rⱼ₊₁)/(rⱼ⋅rⱼ)
        new_rdotr = dot(r,r)
        β     = new_rdotr/rdotr
        rdotr = new_rdotr

        # pⱼ₊₁ = rⱼ₊₁ + βⱼ⋅pⱼ
        axpby!(1.0,r,β,p)
    end

    return maxiter
end

#############################
## BiCGStab Implementation ##
#############################

"""
To represent BiCGStab algorithm.
"""
mutable struct BiCGStab{Ttol,Tdata} <: IterativeSolver{Ttol,Tdata}

    tol::Ttol
    maxiter::Int
    r::Vector{Tdata}
    r̃::Vector{Tdata}
    p::Vector{Tdata}
    p̂::Vector{Tdata}
    s::Vector{Tdata}
    ŝ::Vector{Tdata}
    v::Vector{Tdata}
    t::Vector{Tdata}

    function BiCGStab(x::AbstractVector{Tdata};tol::Ttol=1e-4,maxiter::Int=0) where {Ttol<:AbstractFloat,Tdata<:Continuous}

        N = length(x)
        r = zeros(Tdata,N)
        r̃ = zeros(Tdata,N)
        p = zeros(Tdata,N)
        p̂ = zeros(Tdata,N)
        s = zeros(Tdata,N)
        ŝ = zeros(Tdata,N)
        v = zeros(Tdata,N)
        t = zeros(Tdata,N)
        return new{Ttol,Tdata}(tol,maxiter,r,r̃,p,p̂,s,ŝ,v,t)
    end
end

"""
Solve linear system using preconditioned BiCGStab.
"""
function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},bicgstab::BiCGStab{Ttol,Tdata},P=I;
                maxiter::Int=0,tol::Ttol=0.0)::Int where {Ttol,Tdata}

    @unpack r, r̃, p, p̂, s, ŝ, v, t = bicgstab

    # r = b - A⋅x
    mul!(r,A,x)
    @. r = b - r

    # r̃ = r
    copyto!(r̃,r)

    # |b|
    b̄ = norm(b)

    # intializing values
    β    = 0.0
    α    = 0.0
    ρᵢ₋₁ = 1.0
    ρᵢ₋₂ = 1.0
    ω    = 1.0

    if iszero(maxiter)
        maxiter = bicgstab.maxiter
    end

    if iszero(tol)
        tol = bicgstab.tol
    end

    # iterate to convergence
    @fastmath @inbounds for i in 1:maxiter
        ρᵢ₋₂ = ρᵢ₋₁
        ρᵢ₋₁ = dot(r̃,r)
        if ρᵢ₋₁==0.0
            break
        end
        β    = (ρᵢ₋₁/ρᵢ₋₂)*(α/ω)
        @. p = r + β*(p-ω*v)
        ldiv!(p̂,P,p)
        mul!(v,A,p̂)
        α    = ρᵢ₋₁/dot(r̃,v)
        @. s = r - α*v
        ϵ    = norm(s)/b̄
        if ϵ<tol
            @. x += α*p̂
            return i
        end
        ldiv!(ŝ,P,s)
        mul!(t,A,ŝ)
        ω     = dot(t,s)/dot(t,t)
        @. x += α*p̂ + ω*ŝ
        @. r  = s - ω*t
        ϵ     = norm(r)/b̄
        if ϵ<tol
            return i
        end
        if ω==0.0
            break
        end
    end

    return maxiter
end

##########################
## GMRES Implementation ##
##########################

"""
To represent GMRES algorithm.
"""

mutable struct GMRES{Ttol,Tdata} <: IterativeSolver{Ttol,Tdata}

    ndim::Int
    maxiter::Int
    restart::Int
    tol::Ttol
    H::Matrix{Tdata}
    V::Matrix{Tdata}
    r::Vector{Tdata}
    w::Vector{Tdata}
    s::Vector{Tdata}
    y::Vector{Tdata}
    sn::Vector{Tdata}
    cs::Vector{Tdata}

    function GMRES(x::AbstractVector{Tdata}; maxiter::Int=-1,restart::Int=-1,tol::Ttol=1e-4) where {Ttol<:AbstractFloat,Tdata<:Continuous}

        ndim = length(x)
        if maxiter<0
            maxiter = ndim
        end
        if restart<0
            restart = min(20,ndim)
        end
        H  = zeros(Tdata,restart+1,restart)
        V  = zeros(Tdata,ndim,restart+1)
        r  = zeros(Tdata,ndim)
        w  = zeros(Tdata,ndim)
        s  = zeros(Tdata,restart+1)
        y  = zeros(Tdata,restart+1)
        sn = zeros(Tdata,restart+1)
        cs = zeros(Tdata,restart+1)
        return new{Ttol,Tdata}(ndim,maxiter,restart,tol,H,V,r,w,s,y,sn,cs)
    end
end

function solve!(x::AbstractVector{Tdata},A,b::AbstractVector{Tdata},gmres::GMRES{Ttol,Tdata},M=I;
                maxiter::Int=0,tol::Ttol=0.0)::Int where {Ttol,Tdata}

    H  = gmres.H
    V  = gmres.V
    r  = gmres.r
    w  = gmres.w
    s  = gmres.s
    y  = gmres.y
    sn = gmres.sn
    cs = gmres.cs

    if iszero(maxiter)
        maxiter = gmres.maxiter
    end

    if iszero(tol)
        tol = gmres.tol
    end

    # calculating norm of b
    copyto!(r,b)
    ldiv!(M,r)
    normb = norm(r)
    if normb==0.0
        normb = 1.0
    end

    # calculating residual vector r and getting its norm
    mul!(r,A,x)  # r = A⋅x
    @. r = b - r # r = b - A⋅x
    ldiv!(M,r)   # r = M \ (b - A⋅x)
    β = norm(r)

    # iteration counter
    iter = 0

    # initialize error
    ϵ = β/normb
    if ϵ < tol
        return iter
    end

    @fastmath @inbounds while iter < maxiter
        @. V[:,1] = r/β
        fill!(s,0.0)
        s[1]  = β
        for i in 1:gmres.restart
            iter += 1
            vi    = @view V[:,i]
            mul!(w,A,vi)
            ldiv!(M,w)
            for k in 1:i
                vk     = @view V[:,k]
                H[k,i] = dot(vk,w)
                @. w  -= H[k,i] * vk
            end
            H[i+1,i]    = norm(w)
            @. V[:,i+1] = w / H[i+1,i]
            for k in 1:i-1
                H[k,i], H[k+1,i] = apply_plane_rotation(H[k,i], H[k+1,i], cs[k], sn[k])
            end
            cs[i], sn[i]     = generate_plane_rotation(H[i,i], H[i+1,i])
            H[i,i], H[i+1,i] = apply_plane_rotation(H[i,i], H[i+1,i], cs[i], sn[i])
            s[i], s[i+1]     = apply_plane_rotation(s[i], s[i+1], cs[i], sn[i])
            ϵ                = abs(s[i+1])/normb
            if ϵ < tol
                update!(x,i,H,s,y,V)
                return iter
            end
            if iter==gmres.maxiter
                break
            end
        end
        update!(x,gmres.restart,H,s,y,V)
        mul!(r,A,x)  # r = A⋅x
        @. r = b - r # r = b - A⋅x
        ldiv!(M,r)   # r = M \ (b - A⋅x) = M⁻¹⋅(b - A⋅x)
        β    = norm(r)
        ϵ    = β/normb
        if ϵ < tol
            return iter
        end
    end

    return iter
end

function update!(x::AbstractVector,k::Int,H::Matrix,s::AbstractVector,y::AbstractVector,V::AbstractMatrix)

    copyto!(y,s)

    @fastmath @inbounds for i in k:-1:1
        y[i] /= H[i,i]
        for j in i-1:-1:1
            y[j] -= H[j,i] * y[i]
        end
    end

    @fastmath @inbounds for j in 1:k
        for i in 1:length(x)
            x[i] += y[j] * V[i,j]
        end
    end

    return nothing
end

function apply_plane_rotation(dx::T,dy::T,cs::T,sn::T)::Tuple{T,T} where {T<:Continuous}

    dxp =  cs * dx + conj(sn) * dy
    dyp = -sn * dx +      cs  * dy
    return dxp, dyp
end

function generate_plane_rotation(dx::T,dy::T)::Tuple{T,T} where {T<:Real}

    dz = dx + im*dy
    θ  = angle(dz)
    return cos(θ), sin(θ)
end

function generate_plane_rotation(dx::T,dy::T)::Tuple{T,T} where {T<:Complex}

    c = 0.0
    s = 1.0
    if dx != 0.0
        c = abs(dx)/sqrt(abs2(dx)+abs2(dy))
        s = (dy/dx)*c
    end
    return c, s
end

end

max_stars_repo_name: cohensbw/Langevin

max_issues_repo_name: cohensbw/ElPhDynamics

max_forks_repo_name: cohensbw/ElPhDynamics


text:
using Test, YaoBlocks, YaoArrayRegister

@testset "test constructor" for T in [Float16, Float32, Float64]
    # NOTE: type should follow the axis
    @test RotationGate(X, 0.1) isa PrimitiveBlock{1}

    @test Rx(T(0.1)) isa RotationGate{1,T,XGate}
    @test Ry(T(0.1)) isa RotationGate{1,T,YGate}
    @test Rz(T(0.1)) isa RotationGate{1,T,ZGate}
end

@testset "test matrix" begin
    theta = 2.0
    for (DIRECTION, MAT) in [
        (X, [cos(theta / 2) -im * sin(theta / 2); -im * sin(theta / 2) cos(theta / 2)]),
        (Y, [cos(theta / 2) -sin(theta / 2); sin(theta / 2) cos(theta / 2)]),
        (Z, [exp(-im * theta / 2) 0; 0 exp(im * theta / 2)]),
        (CNOT, exp(-mat(CNOT) / 2 * theta * im |> Matrix)),
        (control(2, (1,), 2 => X), exp(-mat(CNOT) / 2 * theta * im |> Matrix)),
    ]
        @test mat(RotationGate(DIRECTION, theta)) ≈ MAT
    end
end

@testset "test apply" begin
    r = rand_state(1)
    @test state(apply!(copy(r), Rx(0.1))) ≈ mat(Rx(0.1)) * state(r)
end

@testset "test dispatch" begin
    @test dispatch!(Rx(0.1), 0.3) == Rx(0.3)
    @test nparameters(Rx(0.1)) == 1

    @testset "test $op" for op in [+, -, *, /]
        @test dispatch!(op, Rx(0.1), π) == Rx(op(0.1, π))
    end

    @test_throws AssertionError dispatch!(Rx(0.1), (0.2, 0.3))
end

@testset "adjoints" begin
    @test Rx(0.1)' == Rx(-0.1)
    @test Rx(0.2)' == Rx(-0.2)
    @test copy(Rx(0.1)) == Rx(0.1)

    g = Rx(0.1) # creates a new one
    @test copy(g) !== g
end

@testset "isunitary" begin
    g = Rx(0.1 + 0im)
    @test @test_nowarn isunitary(g) == true

    g = Rx(0.1 + 1im)
    @test @test_logs (
        :warn,
        "θ in RotationGate is not real, got θ=$(g.theta), fallback to matrix-based method",
    ) isunitary(g) == false
end


text:
#=
nothing fancy - only used internally (and therefore cryptic) - stores some partially contracted things
seperates out this bit of logic from effective_excitation_hamiltonian (now more readable)
can also - potentially - partially reuse this in other algorithms
=#
struct QPEnv{A,B} <: Cache
    lBs::PeriodicArray{A,1}
    rBs::PeriodicArray{A,1}

    lenvs::B
    renvs::B
end

function environments(exci::Union{InfiniteQP,Multiline{<:InfiniteQP}}, H; solver=Defaults.solver)
    # Explicitly define optional arguments as these depend on solver,
    # which needs to come after these arguments.
    lenvs = environments(exci.left_gs, H; solver=solver)

    return environments(exci, H, lenvs; solver=solver)
end

function environments(exci::Union{InfiniteQP,Multiline{<:InfiniteQP}}, H, lenvs; solver=Defaults.solver)
    # Explicitly define optional arguments as these depend on solver,
    # which needs to come after these arguments.
    renvs = exci.trivial ? lenvs : environments(exci.right_gs, H; solver=solver)

    return environments(exci, H, lenvs, renvs; solver=solver)
end

function environments(exci::InfiniteQP, ham::MPOHamiltonian, lenvs, renvs;solver=Defaults.solver)
    ids = collect(Iterators.filter(x->isid(ham,x),2:ham.odim-1));

    #build lBs(c)
    lB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.left_gs,0)*ham.domspaces[1,k]',
                    space(exci[1],3)'*virtualspace(exci.right_gs,0)) for k in 1:ham.odim]
    lBs = typeof(lB_cur)[]

    for pos = 1:length(exci)
        lB_cur = exci_transfer_left(lB_cur,ham[pos],exci.right_gs.AR[pos],exci.left_gs.AL[pos])*exp(conj(1im*exci.momentum))
        lB_cur += exci_transfer_left(leftenv(lenvs,pos,exci.left_gs),ham[pos],exci[pos],exci.left_gs.AL[pos])*exp(conj(1im*exci.momentum))

        exci.trivial && for i in ids
            @plansor lB_cur[i][-1 -2;-3 -4] -= lB_cur[i][1 4;-3 2]*r_RL(exci.left_gs,pos)[2;3]*τ[3 4;5 1]*l_RL(exci.left_gs,pos+1)[-1;6]*τ[5 6;-4 -2]
        end

        push!(lBs,lB_cur)
    end

    #build rBs(c)
    rB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.left_gs,length(exci))*ham.imspaces[length(exci),k]',
                    space(exci[1],3)'*virtualspace(exci.right_gs,length(exci))) for k in 1:ham.odim]
    rBs = typeof(rB_cur)[]

    for pos=length(exci):-1:1
        rB_cur = exci_transfer_right(rB_cur,ham[pos],exci.left_gs.AL[pos],exci.right_gs.AR[pos])*exp(1im*exci.momentum)
        rB_cur += exci_transfer_right(rightenv(renvs,pos,exci.right_gs),ham[pos],exci[pos],exci.right_gs.AR[pos])*exp(1im*exci.momentum)

        exci.trivial && for i in ids
            @plansor rB_cur[i][-1 -2;-3 -4] -= τ[6 4;1 3]*rB_cur[i][1 3;-3 2]*l_LR(exci.left_gs,pos)[2;4]*r_LR(exci.left_gs,pos-1)[-1;5]*τ[-2 -4;5 6]
        end

        push!(rBs,rB_cur)
    end
    rBs = reverse(rBs)

    local lBE::typeof(rB_cur);
    local rBE::typeof(rB_cur);

    @sync begin
        @Threads.spawn lBE = left_excitation_transfer_system(lB_cur,ham,exci,solver=solver)
        @Threads.spawn rBE = right_excitation_transfer_system(rB_cur,ham,exci, solver=solver)
    end

    lBs[end] = lBE;

    for i=1:length(exci)-1
        lBE = exci_transfer_left(lBE,ham[i],exci.right_gs.AR[i],exci.left_gs.AL[i])*exp(conj(1im*exci.momentum))

        exci.trivial && for k in ids
            @plansor lBE[k][-1 -2;-3 -4] -= lBE[k][1 4;-3 2]*r_RL(exci.left_gs,i)[2;3]*τ[3 4;5 1]*l_RL(exci.left_gs,i+1)[-1;6]*τ[5 6;-4 -2]
        end

        lBs[i] += lBE;
    end

    rBs[1] = rBE;

    for i=length(exci):-1:2
        rBE = exci_transfer_right(rBE,ham[i],exci.left_gs.AL[i],exci.right_gs.AR[i])*exp(1im*exci.momentum)

        exci.trivial && for k in ids
            @plansor rBE[k][-1 -2;-3 -4] -= τ[6 4;1 3]*rBE[k][1 3;-3 2]*l_LR(exci.left_gs,i)[2;4]*r_LR(exci.left_gs,i-1)[-1;5]*τ[-2 -4;5 6]
        end

        rBs[i] += rBE
    end

    return QPEnv(PeriodicArray(lBs),PeriodicArray(rBs),lenvs,renvs)
end

function environments(exci::FiniteQP,ham::MPOHamiltonian,lenvs=environments(exci.left_gs,ham),renvs=exci.trivial ? lenvs : environments(exci.right_gs,ham))
    #construct lBE
    lB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.left_gs,0)*ham.domspaces[1,k]',
                    space(exci[1],3)'*virtualspace(exci.left_gs,0)) for k in 1:ham.odim]
    lBs = typeof(lB_cur)[]
    for pos = 1:length(exci)
        lB_cur = exci_transfer_left(lB_cur,ham[pos],exci.right_gs.AR[pos],exci.left_gs.AL[pos])
        lB_cur += exci_transfer_left(leftenv(lenvs,pos,exci.left_gs),ham[pos],exci[pos],exci.left_gs.AL[pos])
        push!(lBs,lB_cur)
    end

    #build rBs(c)
    rB_cur = [ TensorMap(zeros,eltype(exci),
                    virtualspace(exci.right_gs,length(exci))*ham.imspaces[length(exci),k]',
                    space(exci[1],3)'*virtualspace(exci.right_gs,length(exci))) for k in 1:ham.odim]
    rBs = typeof(rB_cur)[]
    for pos=length(exci):-1:1
        rB_cur = exci_transfer_right(rB_cur,ham[pos],exci.left_gs.AL[pos],exci.right_gs.AR[pos])
        rB_cur += exci_transfer_right(rightenv(renvs,pos,exci.right_gs),ham[pos],exci[pos],exci.right_gs.AR[pos])
        push!(rBs,rB_cur)
    end
    rBs=reverse(rBs)

    return QPEnv(PeriodicArray(lBs),PeriodicArray(rBs),lenvs,renvs)
end

function environments(exci::Multiline{<:InfiniteQP}, ham::MPOMultiline, lenvs, renvs;solver=Defaults.solver)
    exci.trivial || @warn "there is a phase ambiguity in topologically nontrivial statmech excitations"

    left_gs = exci.left_gs;
    right_gs = exci.right_gs;
    exci_space = space(exci[1][1],3);
    lBs = map(1:size(left_gs,1)) do row
        renorms = (map(1:size(left_gs,2)) do col
            v = transfer_left(leftenv(lenvs,row,col,left_gs),ham[row,col],left_gs.AC[row,col],left_gs.AC[row+1,col]);
            @plansor v[1 2;3]*rightenv(lenvs,row,col,left_gs)[3 2;1]
        end).^-1;

        lB_cur = TensorMap(zeros,eltype(left_gs.AL[1,1]),virtualspace(left_gs,row+1,0)*_firstspace(ham[row,1])',exci_space'*virtualspace(left_gs,row,0));

        c_lBs = typeof(lB_cur)[];
        for col in 1:size(left_gs,2)
            lB_cur = renorms[col]*exci_transfer_left(lB_cur,ham[row,col],right_gs.AR[row,col],left_gs.AL[row+1,col])*exp(-1im*exci.momentum)
            lB_cur += renorms[col]*exci_transfer_left(leftenv(lenvs,row,col,left_gs),ham[row,col],exci[row][col],left_gs.AL[row+1,col])*exp(-1im*exci.momentum)
            push!(c_lBs,lB_cur);
        end

        (c_lBs[end],convhist) = linsolve(lB_cur,lB_cur,solver) do v
            y = exp(-1im*size(left_gs,2)*exci.momentum)*prod(renorms)*exci_transfer_left(v,ham[row,:],right_gs.AR[row,:],left_gs.AL[row+1,:])
            if exci.trivial
                @plansor y[-1 -2;-3 -4] -= y[4 2;-3 3]*conj(left_gs.CR[row+1,0][4;1])*rightenv(lenvs,row,0,left_gs)[3 2;1]*leftenv(lenvs,row,1,left_gs)[-1 -2;5]*left_gs.CR[row,0][5;-4]
            end
            v - y
        end
        convhist.converged == 0 && @warn "lbe failed to converge $(convhist.normres)"

        cur = c_lBs[end];
        for col in 1:size(left_gs,2)-1
            cur = renorms[col]*exci_transfer_left(cur,ham[row,col],right_gs.AR[row,col],left_gs.AL[row+1,col])*exp(conj(1im*exci.momentum))
            c_lBs[col] += cur
        end

        PeriodicArray(c_lBs)
    end

    rBs = map(1:size(left_gs,1)) do row
        renorms = (map(1:size(right_gs,2)) do col
            v = transfer_left(leftenv(renvs,row,col,right_gs),ham[row,col],right_gs.AC[row,col],right_gs.AC[row+1,col]);
            @plansor v[1 2;3]*rightenv(renvs,row,col,right_gs)[3 2;1]
        end).^-1;

        rB_cur = TensorMap(zeros,eltype(left_gs.AL[1,1]),virtualspace(left_gs,row,0)*_firstspace(ham[row,1]),exci_space'*virtualspace(left_gs,row+1,0));

        c_rBs = typeof(rB_cur)[];
        for col in size(left_gs,2):-1:1
            rB_cur = renorms[col]*exci_transfer_right(rB_cur,ham[row,col],left_gs.AL[row,col],right_gs.AR[row+1,col])*exp(1im*exci.momentum)
            rB_cur += renorms[col]*exci_transfer_right(rightenv(renvs,row,col,right_gs),ham[row,col],exci[row][col],right_gs.AR[row+1,col])*exp(1im*exci.momentum)
            push!(c_rBs,rB_cur);
        end
        c_rBs = reverse(c_rBs);

        (c_rBs[1],convhist) = linsolve(rB_cur,rB_cur,GMRES()) do v
            y = exp(1im*size(left_gs,2)*exci.momentum)*prod(renorms)*exci_transfer_right(v,ham[row,:],left_gs.AL[row,:],right_gs.AR[row+1,:])
            if exci.trivial
                @plansor y[-1 -2;-3 -4] -= y[1 2;-3 3]*conj(left_gs.CR[row+1,0][1;4])*leftenv(lenvs,row,1,left_gs)[3 2;4]*rightenv(lenvs,row,0,left_gs)[5 -2;-4]*left_gs.CR[row,0][-1;5]
            end
            v - y
        end
        convhist.converged == 0 && @warn "rbe failed to converge $(convhist.normres)"

        cur = c_rBs[1];
        for col in size(left_gs,2):-1:2
            cur = renorms[col]*exci_transfer_right(cur,ham[row,col],left_gs.AL[row,col],right_gs.AR[row+1,col])*exp(1im*exci.momentum)
            c_rBs[col] += cur
        end

        PeriodicArray(c_rBs)
    end

    QPEnv(PeriodicArray(lBs),PeriodicArray(rBs),lenvs,renvs)
end


text:
"""
A modular package that allows users to define a hierachical model of epidemic
spreading from multiple building blocks:

1. EpiRt

2. EpiInf

3. EpiObs

"""
module Epidemia
__precompile__(false)

# import Random
# using Distributions
# using Dates: DateTime, Day
# using UnPack: @unpack
# import Bijectors

include("epirt.jl")

export
  RandomWalk,
  RandomWalkModel,
  ConstantPhaseRandomWalkModel,
  logpdf,
  rand

end # module


text:
module IBMQJulia

using Yao, YaoBlocks

include("api.jl")
include("qobj.jl")
include("extensions.jl")

export authenticate, createreg, status, getresult
end # module


text:
using GasChemMTK
using Test

@testset "GasChemMTK.jl" begin
    # Write your tests here.
end


text:
"""
    val_interp = gaugeinterp(gauge::VisClaw.Gauge, time_interp; varname=:eta::Symbol)

interp gauge values on a non-uniform time
"""
function gaugeinterp(gauge::VisClaw.Gauge, time_interp; varname=:eta::Symbol)
    val = getfield(gauge,varname)
    itp = Interpolations.LinearInterpolation(gauge.time, val, extrapolation_bc=Interpolations.Flat())
    val_interp = itp(time_interp)
    return  val_interp
end


text:
|(x::String, y::String)="$x|$y"

GET="GET"
POST="POST"
PUT="PUT"
DELETE="DELETE"
HEAD = "HEAD"
OPTIONS = "OPTIONS"
PATCH = "PATCH"

routes=Dict()
routes_array=[]
routes_patterns=Dict()
routes_patterns_array=[]

function addnotfound(message::String)
  routes["notfound"] = (req,res) -> begin
    res.status = 404
    return message
  end
end

addnotfound("NotFound")

function createurl(url::String,funtion::Function)
  if occursin(":",url)||occursin("(",url)
      url_ = "^"*url*"\$"
      url_ = replace(url_,":" => "(?<")
      url_ = replace(url_,">" => ">\\w+)")
      routes_patterns[Regex(url_)] = funtion
      push!(routes_patterns_array,Regex(url_))
      @info("Url added",Regex(url_))
  else
    routes[url] = funtion
    push!(routes_array,url)
    @info("Url added",url)
  end
end

macro page(exp1,exp2)
  quote
    createurl("GET"*$exp1,(req,res)->$exp2)
  end
end

macro route(exp1,exp2,exp3)
  quote
    verbs= split($exp1,"|")
    for i=verbs
      createurl(i*$exp2,(req,res)->$exp3)
    end
  end
end

function Get(URL::String, fun::Function)
  createurl("GET"*URL,fun)
end

function Post(URL::String, fun::Function)
  createurl("POST"*URL,fun)
end

function Put(URL::String, fun::Function)
  createurl("PUT"*URL,fun)
end

function Delete(URL::String, fun::Function)
  createurl("DELETE"*URL,fun)
end


text:
### Read the 3D Compass slice from SLIM website and take out the 2D slice
## Please refer to the details in the MIT license of this repository and in the license of the Compass model
## author: Ziyi Yin (ziyi.yin@gatech.edu)

using DrWatson
@quickactivate "CompassTimeLapseCCS"

using PyPlot
using SegyIO
run(`wget -r ftp://slim.gatech.edu/data//synth/Compass`) # this might take a while

# get velocity
block = segy_read("slim.gatech.edu/data/synth/Compass/final_velocity_model_ieee_6m.sgy")

# original compass model is in 25m*25m*6m
n1 = (1911,2730,341)
d = (25f0,25f0,6f0)

n = (637,910,341)

sx = get_header(block, "SourceX")
sy = get_header(block, "SourceY")

v_nogas = zeros(Float32,n)
v_gas = zeros(Float32,n)

for i = 1:n[1]
    x = d[1].*(i-1)
    inds = findall(sx.==x)
    slice = block.data[:,inds[sortperm(sy[inds])]]

    v_nogas[i,:,:] = transpose(slice[:,1:Int(end/3)])
end

for i = 1:n[1]
    x = d[1].*(i-1+n[1])
    inds = findall(sx.==x)
    slice = block.data[:,inds[sortperm(sy[inds])]]

    v_gas[n[1]-i+1,:,:] = transpose(slice[:,Int(end/3)+1:2*Int(end/3)][:,end:-1:1])
end

v = v_nogas[:,300,:]

# get density
block1 = segy_read("slim.gatech.edu/data/synth/Compass/final_density_model_ieee_6m.sgy")

sx = get_header(block1, "SourceX")
sy = get_header(block1, "SourceY")

rho_nogas = zeros(Float32,n)
rho_gas = zeros(Float32,n)

for i = 1:n[1]
    x = d[1].*(i-1)
    inds = findall(sx.==x)
    slice = block1.data[:,inds[sortperm(sy[inds])]]

    rho_nogas[i,:,:] = transpose(slice[:,1:Int(end/3)])
end

for i = 1:n[1]
    x = d[1].*(i-1+n[1])
    inds = findall(sx.==x)
    slice = block1.data[:,inds[sortperm(sy[inds])]]

    rho_gas[n[1]-i+1,:,:] = transpose(slice[:,Int(end/3)+1:2*Int(end/3)][:,end:-1:1])
end

rho = rho_nogas[:,300,:]

d = (25f0,6f0)

JLD2.@save "../model/CompassSlice.jld2" v rho d