text:
#=
 Makie drivers etc. for Pseudospectra.jl

This file is part of Pseudospectra.jl.

Julia implementation
Copyright (c) 2020 Ralph A. Smith

Portions derived from EigTool:
 Copyright (c) 2002-2014, The Chancellor, Masters and Scholars
 of the University of Oxford, and the EigTool Developers. All rights reserved.
 EigTool is maintained on GitHub:  https://github.com/eigtool

SPDX-License-Identifier: BSD-3-Clause
License-Filename: LICENSES/BSD-3-Clause_Eigtool
=#
module PseudospectraMakie

using ..AbstractPlotting
using PlotUtils
using Pseudospectra, LinearAlgebra, Printf
using Colors: RGBA, RGB

export MakieGUIState, psa

# we implement specific methods for these here:
import Pseudospectra: redrawcontour, surfplot, arnoldiplotter!, ewsplotter
import Pseudospectra: plotmode, replzdlg, addmark, fillopts, isheadless
import Pseudospectra: mtxexpsplot, mtxpowersplot
import Pseudospectra: zoomin!, zoomout!, _portrait

# we use these internals here:
import Pseudospectra: vec2ax, expandlevels, oneeigcond, psmode_inv_lanczos
import Pseudospectra: dummyqdlg, replqdlg, transient_bestlb
import Pseudospectra: numrange!

# for convenience (so one could import just this module)
import Pseudospectra: new_matrix, driver!
export new_matrix, driver!

const default_opts = Dict{Symbol,Any}(
    :contourstyle => nothing, :markeig => true,
    :no_waitbar => false,
    # No satisfactory colormap choices are built into Plots.
    # So we bit the bullet and implemented the one from EigTool
    # (which is really good); its spec is added below.
    :contourkw => Dict{Symbol,Any}(:linewidth => 3,
                                   :fill => false)
)

mutable struct MakieGUIState <: GUIState
    mainph # opaque backend object for main plot
    mainfignum::Int
    drawcmd # function to display a plot object (pluggable for GUI use)
    ph2 # opaque backend object for secondary plot
    markerlist
    counter
    is_headless::Bool
    do_savefig::Bool
    figfile

    function MakieGUIState(ph=nothing,num=0,specialcmd=nothing;
                           headless=false, savefigs=true)
        if specialcmd == nothing
            dc = headless ? dcheadless : dcinteractive
        else
            dc = specialcmd
        end
        new(ph,num,dc,nothing,[],0,headless,savefigs,"")
    end
end

isheadless(gs::MakieGUIState) = gs.is_headless

function fillopts(gs::MakieGUIState,optsin::Dict{Symbol,Any}=Dict{Symbol,Any}())
    opts = merge(default_opts,optsin)
end

function drawp(gs::MakieGUIState, p, id)
    gs.drawcmd(gs,p,id)
end

isIJulia() = isdefined(Main, :IJulia) && Main.IJulia.inited

function dcinteractive(gs::MakieGUIState,p,n)
    # CHECKME: do we need special treatment for
    # if isIJulia()
    display(p)
    nothing
end

function dcheadless(gs::MakieGUIState,p,n)
    if gs.do_savefig
        if length(gs.figfile) > 0
            fn = gs.figfile
        else
            fn = tempname() * ".png"
        end
        AbstractPlotting.save(fn, p)
        @info("graph saved in $(fn)")
    end
    nothing
end

# Makie weirdness
# markersize is in data units (pace Makie docs) so we have to adjust it
const MARKER_RATIO = Ref(0.03)

################################################################
# The main plotting functions

function redrawcontour(gs::MakieGUIState, ps_data::PSAStruct, opts)
    zoom = ps_data.zoom_list[ps_data.zoom_pos]
    ps_dict = ps_data.ps_dict
    levels = expandlevels(zoom.levels)
    Z,x,y = zoom.Z, zoom.x, zoom.y
    if !zoom.computed || isempty(zoom.Z)
        error("should not get here w/ uncomputed zoom")
    end
    eigA = ps_dict[:proj_ews]
    fig = Figure(resolution = (500, 500))
    ax1 = fig[1,1] = Axis(fig)
    ctx = ax1
    gs.mainph = fig

    # Before layout was working, we did the equivalent of this:
    # gs.mainph = Scene()
    # ctx = gs.mainph
        if isempty(levels)
            contour!(ctx,x,y,log10.(Z');opts[:contourkw]...)
        else
            kwargs = merge(Dict(:levels => levels),opts[:contourkw])
            contour!(ctx,x,y,log10.(Z'); kwargs...)
            # gs.mainph = quantour(x,y,log10.(Z), levels; opts[:contourkw]...)
        end
        if !isempty(eigA)
            scatter!(ctx,real(eigA),imag(eigA),color=:black)
        end
        if get(opts,:showimagax,false)
            plot!(ctx,[0,0],zoom.ax[3:4],color=:black)
        end
        if get(opts,:showunitcircle,false)
            plot!(ctx,cos.((0:0.01:2)*π),sin.((0:0.01:2)*π),
                  color=:black)
        end
        if get(opts,:showfov,false)
            if isempty(get(ps_dict,:schur_mtx,[]))
                @warn("showfov set in opts but unavailable")
            else
                if isempty(get(ps_dict,:fov,[]))
                    numrange!(ps_data,get(opts,:fov_npts,20))
                end
                plot!(ctx,real.(ps_dict[:fov]),imag.(ps_dict[:fov]),
                      color=:black)
            end
        end
        setxylims!(ctx,zoom.ax)
        drawp(gs,gs.mainph,1)
    nothing
end

"""
make a surface plot of the current spectral portrait
"""
function surfplot(gs::MakieGUIState, ps_data::PSAStruct, opts)
    zoom = ps_data.zoom_list[ps_data.zoom_pos]
    zoom.computed || return
    nx,ny = size(zoom.Z)
    # TODO: allow for user options
    # TODO: control plotting of mesh, e.g.
    # line_freq = floor(Int,min(nx,ny)/10)
    # TODO: add cylinders for eigenvalues, like upstream
    ph = Scene()
    surface!(ph,zoom.x,zoom.y,-log10.(zoom.Z'))
    drawp(gs,ph,2)
end

struct ArnoldiPlotterState{T}
    axis::Scene
    ewsnode::Node{T}
    shiftsnode::Node{T}
end

function arnoldiplotter!(gs::MakieGUIState,old_ax,opts,dispvec,infostr,
                         ews,shifts,state)
    update_ax = false

    if !get(opts,:ARPACK_auto_ax,true)
        @warn("non-auto axes for ARPACK not implemented")
        opts[:ARPACK_auto_ax] = true
    end
    if true # get(opts,:ARPACK_auto_ax,true)
        ax = vec2ax(dispvec)
        if  !(state === nothing)
            if isempty(old_ax)
                append!(old_ax,getxylims(state.axis))
                update_ax = true
            else
#                copy!(old_ax,getxylims(state.axis))
                ax[1] = min(ax[1],old_ax[1])
                ax[3] = min(ax[3],old_ax[3])
                ax[2] = max(ax[2],old_ax[2])
                ax[4] = max(ax[4],old_ax[4])
                update_ax = !all(ax .== old_ax)
            end
        else
            update_ax = true
        end
    end # ax setting

    if state === nothing
        ewsnode = Node(Point2f0.(real.(ews),imag.(ews)))
        shiftsnode = Node(Point2f0.(real.(shifts),imag.(shifts)))

        ax1 = Scene()
        ms = MARKER_RATIO[] * (ax[2] - ax[1])
        s1 = scatter!(ax1,ewsnode,color=:black,markersize=ms)
        title(ax1, infostr)

        s2 = scatter!(ax1,shiftsnode,
                      color=:red,marker=:+,markersize=ms,strokewidth=0.1*ms)
        setxylims!(ax1,ax)
        leg = legend([s1,s2],["ews","shifts"])
        scene = Scene()
        vbox(ax1,leg; parent=scene)

        state = ArnoldiPlotterState(ax1, ewsnode, shiftsnode)
        gs.mainph = scene

        # drawcmd() should force a redraw
        drawp(gs,gs.mainph,1)
    else
        state.ewsnode[] = Point2f0.(real.(ews),imag.(ews))
        state.shiftsnode[] = Point2f0.(real.(shifts),imag.(shifts))
        if update_ax
            ms = MARKER_RATIO[] * (ax[2] - ax[1])
            for j in 1:length(state.axis)
                if state.axis[j] isa Scatter
                    state.axis[j].attributes[:markersize] = ms
                end
            end
            setxylims!(state.axis, ax)
            copy!(old_ax, ax)
        end
    end
    sleep(0.05)
    return state
end

function ewsplotter(gs::MakieGUIState, ews::Vector, zoom)
    if isIJulia()
        # unless we can figure out how to overplot
        return nothing
    end

    fig = Figure()
    gs.mainph = fig
    ax1 = fig[1,1] = Axis(fig)
    scatter!(ax1, real(ews),imag(ews),color=:black)
    setxylims!(ax1,zoom.ax)
    drawp(gs,gs.mainph,1)
end

function plotmode(gs::MakieGUIState,z,A,U,pseudo::Bool,approx::Bool,verbosity)
    m,n = size(A)
    if pseudo
        modestr = "Pseudo"
        niter = 20
        q = normalize!(randn(n) + randn(n)*im)
        Ss,q = psmode_inv_lanczos(A,q,z,1e-15,niter)
        the_col = :magenta
    else
        modestr = "Eigen"
        Ss,q = oneeigcond(A,z,verbosity,
                          dlg=ifelse(isheadless(gs),dummyqdlg,replqdlg))
        the_cond = Ss
        the_col = :cyan
    end
    # FIXME: bail out on error

    # transform to account for Schur decomposition
    q = U * q

    # upstream tries to fetch psmode_x_points_ from base workspace
    x = collect(1:length(q))
    showlog = sum(abs.(diff(q))) * (1/length(q)) >= 10*eps()
    if pseudo
        the_str = @sprintf("ϵ = %15.2e",Ss[end])
    else
        the_rel = approx ? "≈" : "="
        the_str = "κ(λ) " * the_rel *
            @sprintf("%15.2e",the_cond)
    end
    λ_str = @sprintf("%12g%+12gi",real(z),imag(z))
    ax1_title = "$(modestr)mode: " * the_str * "\nλ=" * λ_str
    fig = Figure(resolution = (500, 500))
    ax1 = fig[1,1] = Axis(fig, title = ax1_title)
    l1=lines!(ax1,x,real.(q),color=the_col)
    l2=lines!(ax1,x,abs.(q),color=:black)
    l3=lines!(ax1,x,-abs.(q),color=:black)
    xlims!(ax1,(x[1],x[end]))
    # this should not be needed
    ylims!(ax1,(-1.0,1.0))
    leg = fig[1,2] = Legend(fig, [l1,l2],["realpt", "abs"])

    if showlog
        # TODO: log ticks
        ax2, l2b = lines(fig[2,1],x,log10.(abs.(q)),color=:black)
        xlims!(ax2,x[1],x[end])
        # this should not be needed
        ylims!(ax2,(-10.0,0.0))
        ax2.ylabel = "log₁₀(|xⱼ|)"
        ax2.xlabel = "j"
    else
        ax1.xlabel = "j"
    end
    drawp(gs, fig, 2)
    nothing
end

function mtxpowersplot(gs::MakieGUIState,ps_data::PSAStruct,nmax=50;
                       gradual=false, lbmethod=:none, lbdk=0.25)

    stop_trans = false

    sp_step = 1 # basic stepsize

    A = ps_data.input_matrix^sp_step

    trans = zeros(nmax+1)
    the_pwr = similar(trans)

    trans[1] = 1.0
    the_pwr[1] = 0.0
    ax = [0,20,0,1.5]
    ax2 = copy(ax)
    ax2[3] = 1e-2
    pos = 2
    mtx = Matrix(1.0I,size(A)...)
    max_tp = -Inf
    min_tp = Inf
    ax_factor = 3

    # spectral radius
    alp = maximum(abs.(ps_data.ps_dict[:ews]))
    newfig = true
    pts,bnds = zeros(0),zeros(0)

    function doplot()
        fig = Figure(resolution = (500, 500))
        p1 = fig[1,1] = Axis(fig)
        l1a = lines!(p1,the_pwr[1:pos],trans[1:pos],color=:blue) # dot-dash
        l2a = lines!(p1,the_pwr[1:pos],alp.^the_pwr[1:pos],color=:black, linestyle=:dash)
        xlims!(p1,(ax[1],ax[2]))
        ylims!(p1,(ax[3],ax[4]))
        p1.ylabel = "∥Aᵏ∥"
        p1.xlabel = "k"
        p2, l1b = lines(fig[2,1], the_pwr[1:pos], log10.(trans[1:pos]), color=:blue)
        xlims!(p2,(ax2[1],ax2[2]))
        ylims!(p2,(log10(ax2[3]),log10(ax2[4])))
        p2.ylabel = "log₁₀"
        # TODO: log ticks

        # spectral growth
        l2b = lines!(p2,the_pwr[1:pos],log10.(alp.^the_pwr[1:pos]), color=:black, linestyle=:dash)

        if isempty(bnds)
            leg = fig[3,1] = Legend(fig, [l1b,l2b],["norms", "spectral"])
        else
            l3b = lines!(p2,pts,bnds, linestyle=:dashdot, color=:green)
            leg = fig[3,1] = Legend(fig, [l1b,l2b,l3b],["norms", "spectral", "bound"])
        end
        leg.tellheight = true
        leg.orientation = :horizontal
        # yspan = log10(ax2[4]) - log10(ax2[3])
        # step = max(1,floor(yspan/3))
        #    plot!(yticks = log10(ax2[3]):step:floor(log10(ax2[4])))
        drawp(gs,fig,2)
    end
    if lbmethod != :none
        pts,bnds,sel_pt = transient_bestlb(ps_data,:powers,0:nmax,
                                           method=lbmethod, dk=lbdk)
    end
    while !stop_trans
        mtx = A * mtx
        trans[pos] = norm(mtx)
        the_pwr[pos] = sp_step*(pos-1)
        max_tp = max(max_tp,trans[pos])
        min_tp = min(min_tp,trans[pos])

        if the_pwr[pos] > ax[2]*0.9
            ax[2] *= ax_factor
            ax2[2] *= ax_factor
        end
        if max_tp > ax[4] * 0.9
            ax[4] = max(max_tp * 1.5,ax_factor * ax[4])
        end
        if max_tp > ax2[4] * 0.9
            ax2[4] = ax2[4]^ax_factor
        end
        if (log10(min_tp)-log10(ax2[3])) / (log10(ax2[4]) - log10(ax2[3])) < 0.1
            ax2[3] = ax2[3]^ax_factor
        end

        if gradual
            doplot()
            sleep(0.02)
        end
        if (max_tp > 1e130) || ((min_tp < 1e-130) && (min_tp != 0))
            @warn("stopping: numbers going out of range")
            stop_trans = true
        end
        if min_tp == 0
            @warn("stopping: exactly zero matrix")
            stop_trans = true
        end
        (pos > nmax) && (stop_trans = true)
        pos += 1
    end
    if !gradual
        pos -= 1
        doplot()
    end
end

function mtxexpsplot(gs::MakieGUIState, ps_data::PSAStruct,dt=0.1,nmax=50;
                     gradual=false, lbmethod=:none)

    stop_trans = false
    A = ps_data.input_matrix
    eAdt = exp(dt*A)

    trans = zeros(nmax+1)
    the_time = similar(trans)

    trans[1] = 1.0
    the_time[1] = 0.0
    ax = [0,20*dt,0,2]
    ax2 = copy(ax)
    ax2[3] = 1e-2
    pos = 2
    eAt = Matrix(1.0I,size(A)...)
    max_tp = -Inf
    min_tp = Inf
    ax_factor = 3

    # spectral abscissa
    alp = maximum(real.(ps_data.ps_dict[:ews]))
    newfig = true
    if lbmethod != :none
        pts,bnds,sel_pt = transient_bestlb(ps_data,:exp,(0:nmax)*dt,
                                           method=lbmethod)
    else
        pts,bnds = zeros(0),zeros(0)
    end

    function doplot()
        fig = Figure(resolution = (500, 500))
        p1 = fig[1,1] = Axis(fig)
        l1a = lines!(p1,the_time[1:pos], trans[1:pos],color=:blue)
        l2a = lines!(p1,the_time[1:pos], exp.(alp*the_time[1:pos]), color=:black,
                     linestyle=:dash)
        xlims!(p1,(ax[1],ax[2]))
        ylims!(p1,(ax[3],ax[4]))
        p1.ylabel = "∥exp(At)∥"
        p1.xlabel = "t"

        p2, l1b = lines(fig[2,1], the_time[1:pos], log10.(trans[1:pos]), color=:blue)
        xlims!(p2,(ax2[1],ax2[2]))
        ylims!(p2,(log10(ax2[3]),log10(ax2[4])))
        p2.ylabel = "log₁₀"

        # TODO: logarithmic ticks for y-axis
        l2b = lines!(p2,the_time[1:pos],(1/log(10.0)) * (alp*the_time[1:pos]), color=:black,
                     linestyle=:dash)
        if isempty(bnds)
            leg = fig[3,1] = Legend(fig, [l1b,l2b],["norms", "spectral"])
        else
            l3b = lines!(p2,pts[1:pos],log10.(bnds[1:pos]), color=:green, linestyle=:dashdot)
            leg = fig[3,1] = Legend(fig, [l1b,l2b,l3b],["norms", "spectral", "bound"])
        end
        leg.tellheight = true
        leg.orientation = :horizontal

        # yspan = log10(ax2[4]) - log10(ax2[3])
        #    step = max(1,floor(yspan/3))
        #    plot!(yticks = log10(ax2[3]):step:floor(log10(ax2[4])))
        drawp(gs,fig,2)
    end

    while !stop_trans
        eAt = eAt * eAdt
        trans[pos] = norm(eAt)
        the_time[pos] = (pos-1)*dt
        max_tp = max(max_tp,trans[pos])
        min_tp = min(min_tp,trans[pos])

        if the_time[pos] > ax[2]*0.9
            ax[2] *= ax_factor
            ax2[2] *= ax_factor
        end
        if max_tp > ax[4] * 0.9
            ax[4] = max(max_tp * 1.5,ax_factor * ax[4])
        end
        if max_tp > ax2[4] * 0.9
            ax2[4] = ax2[4]^ax_factor
        end
        if (log10(min_tp)-log10(ax2[3])) / (log10(ax2[4]) - log10(ax2[3])) < 0.1
            ax2[3] = ax2[3]^ax_factor
        end
        if gradual
            doplot()
            sleep(0.02)
        end

        if (max_tp > 1e130) || ((min_tp < 1e-130) && (min_tp != 0))
            @warn("stopping: numbers going out of range")
            stop_trans = true
        end
        (pos > nmax) && (stop_trans = true)
        pos += 1
    end
    if !gradual
        pos -= 1
        doplot()
    end
end

function zoomin!(gs::MakieGUIState, ps_data::PSAStruct,
                 optsin=Dict{Symbol,Any}(); zkw=zeros(0))
    opts = fillopts(gs,optsin)
    if isempty(zkw)
        # User might get z from some other figure, but so what?
        z = replzdlg(gs)
    else
        z = zkw
    end
    ax = getxylims(gs.mainph)
    res = zoomin!(ps_data,z,ax)
    if res < 0
        @warn("unable to zoom based on requested point")
    elseif res == 0
        # redraw existing portrait
        redrawcontour(gs, ps_data, opts)
    else
        # compute and draw new portrait
        driver!(ps_data, opts, gs)
    end
    res
end

function zoomout!(gs::MakieGUIState, ps_data::PSAStruct,
                 optsin=Dict{Symbol,Any}(); include_fov=false, zkw=zeros(0))
    opts = fillopts(gs,optsin)
    if isempty(zkw)
        z = replzdlg(gs)
    else
        z = zkw
    end
    ax = getxylims(gs.mainph); # current plot axes
    res = zoomout!(ps_data,z,ax,include_fov=include_fov)
    if res < 0
        @warn("unable to zoom based on requested point")
    elseif res == 0
        # redraw existing portrait
        redrawcontour(gs, ps_data, opts)
    else
        # compute and draw new portrait
        driver!(ps_data, opts, gs)
    end
    res
end

################################################################
# Wrappers

function _portrait(xs,ys,Z,eigA)
    fig = Figure()
    ax1 = fig[1,1] = Axis(fig)
    c = contour!(ax1,xs,ys,log10.(Z'))
    ax = getxylims(fig)
    ms = MARKER_RATIO[] * (ax[2] - ax[1])
    scatter!(ax1, real(eigA), imag(eigA), color=:black, # label="eigvals",
             markersize=ms)
    fig
end

"""
    psa(A [,opts]) -> ps_data, graphics_state

Compute and plot pseudospectra of a matrix.

This is a rudimentary command-line driver for the Pseudospectra package.
"""
function psa(A::AbstractMatrix, opts::Dict{Symbol,Any}=Dict{Symbol,Any}())
    allopts = merge(default_opts,optsin)
    ps_data = new_matrix(A,allopts)
        gs = MakieGUIState()

    driver!(ps_data,allopts,gs)
    ps_data, gs
end
#############################################################
# Utilities

function setxylims!(ph,ax)
    xlims!(ph,(ax[1],ax[2]))
    ylims!(ph,(ax[3],ax[4]))
end

function getxylims(ph)
    if ph isa Scene
        sl = scene_limits(ph)
    else
        ctx = content(ph[1,1])
        sl = ctx.finallimits[]
    end
    if sl === nothing
        return fill(0.0,4)
    end
    v0 = sl.origin
    v1 = v0 + sl.widths
    return [v0[1],v1[1],v0[2],v1[2]]
end

# this is here because other backends allow picking with cursor
function replzdlg(gs::MakieGUIState; what="a z-value")
    x,y = NaN,NaN

    println("Specify " * what *
            "; enter real and imaginary components separated by whitespace")
    print("z = ")
    l = readline()
    strs = split(l)
    try
        x = parse(Float64,strs[1])
        y = parse(Float64,strs[2])
    catch
        println("unable to parse input")
        x = NaN
    end
    z = x + y*im
    return z
end

function addmark(gs::MakieGUIState,z,mykey)
    x,y = real(z),imag(z)
    if mykey == :pseudo
        mkey = :circle
        ckey = :magenta
    elseif mykey == :eigen
        mkey = :circle
        ckey = :cyan
    else
        mkey = :x
        ckey = :black
    end
    # FIXME: add marker to list (replacing if any)
    # then trigger a redraw that handles everything
    ctx = gs.mainph[1,1]
    scatter!(ctx,[x],[y],color=ckey,marker=mkey,markersize=7)
    drawp(gs,gs.mainph,1)
    nothing
end

"""
    The colormap from EigTool. This provides clearly
    distinguishable colors for a line-contour plot.
"""
function etcgrad()
    cm = [
        0   1.00000000000000   0.10000000000000
        0   0.92500000000000   0.17500000000000
        0   0.85000000000000   0.25000000000000
        0   0.77500000000000   0.32500000000000
        0   0.70000000000000   0.40000000000000
        0   0.62500000000000   0.47500000000000
        0   0.55000000000000   0.55000000000000
        0   0.47500000000000   0.62500000000000
        0   0.40000000000000   0.70000000000000
        0   0.32500000000000   0.77500000000000
        0   0.25000000000000   0.85000000000000
        0   0.17500000000000   0.92500000000000
        0   0.10000000000000   1.00000000000000
        0.06153846153846   0.09230769230769   1.00000000000000
        0.12307692307692   0.08461538461538   1.00000000000000
        0.18461538461538   0.07692307692308   1.00000000000000
        0.24615384615385   0.06923076923077   1.00000000000000
        0.30769230769231   0.06153846153846   1.00000000000000
        0.36923076923077   0.05384615384615   1.00000000000000
        0.43076923076923   0.04615384615385   1.00000000000000
        0.49230769230769   0.03846153846154   1.00000000000000
        0.55384615384615   0.03076923076923   1.00000000000000
        0.61538461538462   0.02307692307692   1.00000000000000
        0.67692307692308   0.01538461538462   1.00000000000000
        0.73846153846154   0.00769230769231   1.00000000000000
        0.80000000000000                  0   1.00000000000000
        0.80769230769231                  0   0.92307692307692
        0.81538461538462                  0   0.84615384615385
        0.82307692307692                  0   0.76923076923077
        0.83076923076923                  0   0.69230769230769
        0.83846153846154                  0   0.61538461538462
        0.84615384615385                  0   0.53846153846154
        0.85384615384615                  0   0.46153846153846
        0.86153846153846                  0   0.38461538461538
        0.86923076923077                  0   0.30769230769231
        0.87692307692308                  0   0.23076923076923
        0.88461538461538                  0   0.15384615384615
        0.89230769230769                  0   0.07692307692308
        0.90000000000000                  0                  0
        0.86923076923077                  0                  0
        0.83846153846154                  0                  0
        0.80769230769231                  0                  0
        0.77692307692308                  0                  0
        0.74615384615385                  0                  0
        0.71538461538462                  0                  0
        0.68461538461538                  0                  0
        0.65384615384615                  0                  0
        0.62307692307692                  0                  0
        0.59230769230769                  0                  0
        0.56153846153846                  0                  0
        0.53076923076923                  0                  0
        0.50000000000000                  0                  0
        0.54166666666667   0.05000000000000                  0
        0.58333333333333   0.10000000000000                  0
        0.62500000000000   0.15000000000000                  0
        0.66666666666667   0.20000000000000                  0
        0.70833333333333   0.25000000000000                  0
        0.75000000000000   0.30000000000000                  0
        0.79166666666667   0.35000000000000                  0
        0.83333333333333   0.40000000000000                  0
        0.87500000000000   0.45000000000000                  0
        0.91666666666667   0.50000000000000                  0
        0.95833333333333   0.55000000000000                  0
        1.00000000000000   0.60000000000000                  0
    ]
    nc = size(cm,1)
    cc = Vector{RGBA{Float64}}()
    vv = zeros(nc)
    for i in 1:nc
        col = RGBA(RGB(cm[i,:]...))
        push!(cc,col)
        vv[i]=(i-1)/(nc-1)
    end
    if :ContinuousColorGradient in names(PlotUtils, all=true)
        return PlotUtils.ContinuousColorGradient(cc,vv)
    else
        return PlotUtils.ColorGradient(cc,vv)
    end
end

const eigtool_cgrad = etcgrad()

default_opts[:contourkw][:linecolor] = eigtool_cgrad;

end # module


text:
############################
# Convenience Constructors #
############################
    statevec(i::Int, fb::FiniteBasis) = QuArray(single_coeff(i, length(fb)), fb)
    statevec(i::Int, lens::Int...=i) = statevec(i, FiniteBasis(lens))

    ###########################
    # Array-like Constructors #
    ###########################
    function Base.zeros(qa::AbstractQuArray)
        fc = zeros(coeffs(qa))
        QAT = similar_type(typeof(qa))
        return QAT(fc, bases(qa))
    end
    function Base.eye(qa::AbstractQuArray)
        fc = eye(coeffs(qa))
        QAT = similar_type(typeof(qa))
        return QAT(fc, bases(qa))
    end

    ####################
    # Helper Functions #
    ####################
    one_at_ind!(arr, i) = setindex!(arr, one(eltype(arr)), i)
    single_coeff(i, lens...) = one_at_ind!(zeros(lens), i)

# Reference :
# Section - 1.1, http://cds.cern.ch/record/331607/files/9708012.pdf
function coherentstatevec_inf(n::Int,alpha::Number)
    s = zeros(typeof(float(alpha)),n)
    s[1] = one(alpha)
    for i in 2:n
        s[i] = alpha/sqrt(i-1)*s[i-1]
    end
    z = QuArray(s)
    return scale!(exp(-abs2(alpha)/2),z)
end

# Reference :
# 2nd Defintion : http://en.wikipedia.org/wiki/Coherent_states#Mathematical_features_of_the_canonical_coherent_states
coherentstatevec(n::Int, alpha::Number) = displaceop(n,alpha)*statevec(1,FiniteBasis(n))

export statevec,
    coherentstatevec

ext: jl


text:
# Train a truncated SVD model of rank max_rank on the ratings in rating_set.
# min_epochs and max_epochs control the range of possible number of epochs used
# in the gradient descent. Since we train features sequentially, overfitting
# can happen on early features. If min_epochs < max_epochs, a heuristic for
# early exit is used that tends to avoid overfitting, tracking the rate of
# change in the error between the model and the known ratings. This rate of
# change may decrease sharply initially, level off and increase, then finally
# decrease again. We try to catch it on it's second (final) decrease and
# terminate training the current feature when we do. learning_rate and
# regularizer are the two standard gradient descent parameters that control
# how fast the descent happens and prevent the descent from over-aggressive
# training in any direction.

function train(rating_set::RatingSet,
               max_rank;
               min_epochs=0,
               max_epochs=100,
               learning_rate=0.001,
               regularizer=0.02)
    user_features = [0.1 for r=1:length(rating_set.user_to_index), c=1:max_rank]
    item_features = [0.1 for r=1:length(rating_set.item_to_index), c=1:max_rank]
    residuals = [Residual(rating.value, 0.0, 0.0) for rating in rating_set.training_set]
    num_ratings = length(rating_set.training_set)
    p = Progress(max_rank, 1, "Computing truncated rank $(max_rank) SVD ")
    for rank=1:max_rank
        errors = [0.0, Inf, Inf]
        for i=1:max_epochs
            for j=1:num_ratings
                rating, residual = rating_set.training_set[j], residuals[j]
                item_feature = item_features[rating.item, rank]
                user_feature = user_features[rating.user, rank]
                residual.curr_error = residual.value - user_feature * item_feature
                error_diff = residual.prev_error - residual.curr_error
                errors[1] += error_diff * error_diff
                residual.prev_error = residual.curr_error
                item_features[rating.item, rank] += learning_rate * (residual.curr_error * user_feature - regularizer * item_feature)
                user_features[rating.user, rank] += learning_rate * (residual.curr_error * item_feature - regularizer * user_feature)
            end
            # distance decreases, then increases, then decreases. we want to catch it on second decrease
            if i > min_epochs && errors[1] < errors[2] && errors[2] > errors[3]
                break
            end
            errors[1], errors[2], errors[3] = 0.0, errors[1], errors[2]
        end
        for residual in residuals
            residual.value = residual.curr_error
            residual.prev_error = 0.0
        end
        next!(p)
    end

    # We end up with just the U and V matrices of the singular value
    # decomposition, but S can be extracted by normalizing both U and V.
    singular_values = [norm(user_features[:,rank]) * norm(item_features[:,rank]) for rank=1:max_rank]
    for rank=1:max_rank
        user_features[:,rank] /= norm(user_features[:,rank])
        item_features[:,rank] /= norm(item_features[:,rank])
    end
    return RatingsModel(rating_set.user_to_index,
                        rating_set.item_to_index,
                        user_features,
                        singular_values,
                        item_features)
end


text:
isharshad(x)   = x % sum(digits(x)) == 0
nextharshad(x) = begin while !isharshad(x+1) x += 1 end; return x + 1 end

function harshads(n::Integer)
	h = Vector{typeof(n)}(n)
	h[1] = 1
	for j in 2:n
		h[j] = nextharshad(h[j-1])
	end
	return h
end

println("First 20 harshad numbers: ", join(harshads(20), ", "))
println("First harshad number after 1001: ", nextharshad(1000))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
import Base.append!

struct ElasticCovStrat <: CovarianceStrategy
    capacity::Int
    stepsize::Int
end
function init_precompute(covstrat::ElasticCovStrat, X, y, k)
    nobs = size(X, 2)
    FullCovariancePrecompute(nobs)
end

append!(gp, x::AbstractVector, y::Float64) = append!(gp, reshape(x, :, 1), [y])
function append!(gp::GPE{X,Y,M,K,CS,D,P}, x::AbstractMatrix, y::AbstractVector) where {X,Y,M,K,CS,D,P <: ElasticPDMat}
    size(x, 2) == length(y) || error("$(size(x, 2)) observations, but $(length(y)) targets.")
    newcov = [cov(gp.kernel, gp.x, x); cov(gp.kernel, x, x) + (noise_variance(gp) + eps())*I]
    append!(gp.data, gp.kernel, gp.x, x)
    append!(gp.x, x)
    append!(gp.cK, newcov)
    gp.nobs += length(y)
    append!(gp.y, y)
    update_target!(gp, kern = false, noise = false)
end

LinearAlgebra.ldiv!(cK::ElasticPDMat, x) = ldiv!(cK.chol, x)
wrap_cK(cK::ElasticPDMat, Σbuffer, chol::Cholesky) = cK
mat(cK::ElasticPDMat) = view(cK.mat)
cholfactors(cK::ElasticPDMat) = view(cK.chol).factors

function ElasticGPE(dim; mean::Mean = MeanZero(), kernel = SE(0.0, 0.0),
                    logNoise = -2.0, kwargs...)
    x = ElasticArray(Array{Float64}(undef, dim, 0))
    y = ElasticArray(Array{Float64}(undef, 0))
    ElasticGPE(x, y, mean, kernel, logNoise; kwargs...)
end
function alloc_cK(covstrat::ElasticCovStrat, nobs)
    # create placeholder PDMat
    m    = Matrix{Float64}(undef, nobs, nobs)
    chol = Matrix{Float64}(undef, nobs, nobs)
    cK = ElasticPDMat(m, Cholesky(chol, :U, 0), capacity=covstrat.capacity, stepsize=covstrat.stepsize)
    return cK
end
function ElasticGPE(x::AbstractMatrix, y::AbstractVector, mean::Mean, kernel::Kernel,
                    logNoise = -2.0;
                    capacity = 10^3, stepsize = 10^3)
    data = ElasticKernelData(kernel, x, x, capacity=capacity, stepsize=stepsize)
    nobs = length(y)
    # create placeholder PDMat
    covstrat = ElasticCovStrat(capacity, stepsize)
    cK = alloc_cK(covstrat, nobs)
    gp = GPE(ElasticArray(x), ElasticArray(y), mean, kernel, logNoise, covstrat, data, cK)
    initialise_target!(gp)
end

function prepareappend!(kd, Xnew)
    dim, nobs_new = size(Xnew)
    nobs = kd.dims[1]
    if nobs + nobs_new > kd.capacity[1]
        kd.capacity = kd.capacity .+ kd.stepsize
        ElasticPDMats.resize!(kd)
    end
    kd, dim, nobs, nobs_new
end

#===========================
    Elastic Kernel Data
===========================#

function ElasticKernelData(k::Isotropic, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    kerneldata = IsotropicData(AllElasticArray(2; capacity = (capacity, capacity), stepsize = (stepsize, stepsize)))
    nobs = size(X, 2)
    distance!(view(kerneldata.R, 1:nobs, 1:nobs), k, X)
    setdimension!(kerneldata.R, nobs, 1:2)
    kerneldata
end
function append!(kerneldata::IsotropicData{<:AllElasticArray}, k::Isotropic, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.R, Xnew)
    distance!(view(kd, 1:nobs, nobs + 1:nobs + nobs_new), k, X, Xnew)
    copyto!(view(kd, nobs + 1:nobs + nobs_new, 1:nobs), transpose(view(kd, 1:nobs, nobs + 1:nobs + nobs_new)))
    distance!(view(kd, nobs + 1:nobs + nobs_new, nobs + 1:nobs + nobs_new), k, Xnew)
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end
append!(kerneldata, k, X, Xnew::AbstractVector) = append!(kerneldata, k, X, reshape(Xnew, :, 1))

function ElasticKernelData(k::StationaryARD, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    dim, nobs = size(X)
    dist_stack = AllElasticArray(3; capacity = (capacity, capacity, size(X, 1)),
                                    stepsize = (stepsize, stepsize, 0))
    for d in 1:dim
        grad_ls = view(dist_stack, 1:nobs, 1:nobs, d)
        distance!(grad_ls, SqEuclidean(), view(X, d:d, :))
    end
    setdimension!(dist_stack, nobs, 1:2)
    setdimension!(dist_stack, dim, 3)
    StationaryARDData(dist_stack)
end
function append!(kerneldata::StationaryARDData, kernel::StationaryARD, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.dist_stack, Xnew)
    for d in 1:dim
        grad_ls = view(kd, 1:nobs, nobs + 1:nobs + nobs_new, d)
        distance!(grad_ls, SqEuclidean(), view(X, d:d, :), view(Xnew, d:d, :))
        copyto!(view(kd, nobs + 1:nobs + nobs_new, 1:nobs, d),
                transpose(grad_ls))
        distance!(view(kd, nobs + 1:nobs + nobs_new,
                       nobs + 1:nobs + nobs_new, d), SqEuclidean(), view(Xnew, d:d, :))
    end
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end

function ElasticKernelData(k::LinArd, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    dim, nobs = size(X)
    XtX_d = AllElasticArray(3; capacity = (capacity, capacity, size(X, 1)),
                                    stepsize = (stepsize, stepsize, 0))
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:i
                XtX_d[i, j, d] = XtX_d[j, i, d] = X[d, i] * X[d, j]
            end
        end
    end
    setdimension!(XtX_d, nobs, 1:2)
    setdimension!(XtX_d, dim, 3)
    LinArdData(XtX_d)
end
function append!(kerneldata::LinArdData, kernel::LinArd, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.XtX_d, Xnew)
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:nobs_new
                kd[i, nobs + j, d] = kd[nobs + j, i, d] = X[d, i] * Xnew[d, j]
            end
        end
        for i in 1:nobs_new
            for j in 1:i
                kd[nobs + i, nobs + j, d] = kd[nobs + j, nobs + i, d] = Xnew[d, i] * Xnew[d, j]
            end
        end
    end
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end

function ElasticKernelData(k::LinIso, X1::AbstractMatrix, X2::AbstractMatrix; capacity = 10^3, stepsize = 10^3)
    @assert X1==X2 # TODO: extend to X1 != X2
    X = X1
    dim, nobs = size(X)
    XtX = AllElasticArray(2, capacity = (capacity, capacity), stepsize = (stepsize, stepsize))
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:i
                XtX[i, j] = XtX[j, i] += X[d, i] * X[d, j]
            end
        end
    end
    setdimension!(XtX, nobs, 1:2)
    LinIsoData(XtX)
end
function append!(kerneldata::LinIsoData, kernel::LinIso, X::AbstractMatrix, Xnew::AbstractMatrix)
    kd, dim, nobs, nobs_new = prepareappend!(kerneldata.XtX, Xnew)
    @inbounds @simd for d in 1:dim
        for i in 1:nobs
            for j in 1:nobs_new
                kd[i, j + nobs] = kd[j + nobs, i] += X[d, i] * Xnew[d, j]
            end
        end
        for i in 1:nobs_new
            for j in 1:i
                kd[i + nobs, j + nobs] = kd[j + nobs, i + nobs] += Xnew[d, i] * Xnew[d, j]
            end
        end
    end
    setdimension!(kd, nobs + nobs_new, 1:2)
    kerneldata
end


text:
# Canvas: a thing upon which other things are placed.

export Canvas, Units, Rotation, canvas, deferredcanvas,
       draw, drawpart, set_unit_box, set_box

# A box giving the coordinate system used by a canvas.
type Units
    x0::Float64
    y0::Float64
    width::Float64
    height::Float64

    function Units()
        new(0.0, 0.0, 1.0, 1.0)
    end

    function Units(width::Number,
                   height::Number)
        new(0.0, 0.0, width, height)
    end

    function Units(x0::Number,
                   y0::Number,
                   width::Number,
                   height::Number)
        new(x0, y0, width, height)
    end
end

# Note: Units has almost the same representation as BoundingBox, except no unit
# is associated with the numbers.
function convert(::Type{BoundingBox}, u::Units)
    BoundingBox(u.x0, u.y0, u.width, u.height)
end

function convert(::Type{Units}, u::BoundingBox)
    Units(u.x0.value, u.y0.value, u.width.value, u.height.value)
end

# Canvases are containers for forms and other canvases with an associated
# coordinate transform.
abstract Canvas

# An identity element for the canvas monoid.
type EmptyCanvas <: Canvas end
const empty_canvas = EmptyCanvas()

# A deferred canvas is a function that produces a canvas and is evaluated when
# the graphic is drawn. This is useful when the layout of the contents of a
# canvas depends on its size in absolute coordinates, which is only known, in
# general at draw time.
#
# The function `f` must have the signature:
#   f(box::BoundingBox, unit_box::BoundingBox)
#
# Where box is the bounding box in absolute coordinates of the parent canvas,
# and unit_box is its unit box.
#
# Another way to think of this: a canvas in to an expression what a
# deferredcanvas is to a macro.
#
type DeferredCanvas <: Canvas
    f::Function
    order::Integer
end

deferredcanvas(f::Function) = DeferredCanvas(f, 0)
deferredcanvas(f::Function, order::Integer) = DeferredCanvas(f, order)

# non-empty tree of canvases.
type CanvasTree <: Canvas
    box::BoundingBox
    form::Form
    property::Property
    unit_box::Units
    rot::Rotation
    children::List{Canvas}

    # Z-order of this canvas relative to its siblings.
    order::Integer

    # True if this canvas should use its parent's coordinate system.
    units_inherited::Bool

    # True if children of the canvas should be clipped by its bounding box.
    clip::Bool

    function CanvasTree(box::BoundingBox,
                        form::Form,
                        property::Property,
                        unit_box::Units,
                        rot::Rotation,
                        children::List{Canvas},
                        order::Integer,
                        units_inherited::Bool,
                        clip::Bool)

        new(box, form, property, unit_box, rot,
            children, order, units_inherited, clip)
    end

    function CanvasTree(x0::MeasureOrNumber=0.0w,
                        y0::MeasureOrNumber=0.0h,
                        width::MeasureOrNumber=1.0w,
                        height::MeasureOrNumber=1.0h;
                        unit_box=Units(),
                        rotation=Rotation(),
                        units_inherited=false,
                        order=0,
                        clip=false)
        new(BoundingBox(x0, y0, width, height),
                empty_form,
                empty_property,
                unit_box,
                rotation,
                ListNil{Canvas}(),
                order,
                units_inherited,
                clip)
    end

    # shallow copy constructor
    function CanvasTree(c::CanvasTree)
        new(c.box, c.form, c.property, c.unit_box,
            c.rot, c.children, c.order, c.units_inherited, c.clip)
    end
end

# Alias constructor functions
const canvas = CanvasTree

copy(c::CanvasTree) = CanvasTree(c)

# Return a new Canvas with unit_box substituted.
function set_unit_box(c::CanvasTree, unit_box::Units)
    c = copy(c)
    c.unit_box = unit_box
    c
end

# Return a new Canvas with the box substituted
function set_box(c::CanvasTree, box::BoundingBox)
    c = copy(c)
    c.box = box
    c
end

function show(io::IO, c::CanvasTree)
    print(io, "Canvas(")
    show(io, c.children)
    print(io, ")")
end

# Similar to dump(), but briefer
function contents(io, c::CanvasTree, n::Int, indent)
    lc = length(c.children)
    lf = length(c.form)
    println(indent, "Canvas with ", lc, lc == 1 ? " child" : " children",
            " and ", lf, lf == 1 ? " form" : " forms")
    if n > 0
        # Children
        i = 1
        for cc in c.children
            contents(io, cc, n - 1, string(indent, "  "))
            if i > 10
                println(io, indent, "  ...")
                break
            end
            i += 1
        end
        # Form
        contents(io, c.form, n, indent)
    end
end
contents(io, c::CanvasTree, n::Int) = contents(io, c, n, "")
contents(io, c::CanvasTree) = contents(io, c, 10, "")
contents(c::CanvasTree, args...) = contents(OUTPUT_STREAM, c, args...)

compose(canvas::Canvas) = canvas

# Compositions of canvases is tree joining by making b a subtree of a.
#     a      b             a___
#    / \    / \    --->   / \  \
#   X   Y  U   V         X   Y  b
#                              / \
#                             U   V
function compose(canvases::Canvas...)
    root = empty_canvas
    for canvas in canvases
        if canvas === empty_canvas
            continue
        elseif root === empty_canvas
            root = copy(canvas)
        else
            root.children = cons(canvas, root.children)
        end
    end
    root
end

# Composition of forms into canvases
function compose(a::CanvasTree, b::Form)
    CanvasTree(a.box, combine(a.form, b), a.property,
               a.unit_box, a.rot, a.children, a.order, a.units_inherited, a.clip)
end

function compose(a::EmptyCanvas, b::Form)
    a
end

function compose(a::CanvasTree, b::Property)
    CanvasTree(a.box, a.form, combine(a.property, b),
               a.unit_box, a.rot, a.children, a.order, a.units_inherited, a.clip)
end

function compose(a::EmptyCanvas, b::Property)
    a
end

function compose(a::DeferredCanvas, b::Union(Canvas, Form, Property))
    error("Error: Deferred canvases cannot be composed with.")
end

# Draw a realization of the graphic on a given backend.
#
# Args:
#   backend: Backend on which to draw the graphic.
#   root_canvas: Root node of the graphic.
#
# Returns:
#   nothing
#
# Effects:
#   The root_canvas will be finished after the call to draw. If you wish to draw
#   multiple canvases, use `drawpart`.
#
function draw(backend::Backend, root_canvas::Canvas)
    drawpart(backend, root_canvas)
    finish(backend)
end

# Draw without finishing the backend.
function drawpart(backend::Backend, root_canvas::Canvas)
    S = {(root_canvas, NativeTransform(), BoundingBox(), root_box(backend))}
    while !isempty(S)
        s = pop!(S)

        if s == :POP_PROPERTY
            pop_property(backend)
            continue
        end

        (canvas, parent_t, unit_box, parent_box) = s

        if canvas === empty_canvas
            continue
        end

        if typeof(canvas) == DeferredCanvas
            abs_parent_box =
                BoundingBox(convert(SimpleMeasure{MillimeterUnit}, parent_box.x0),
                            convert(SimpleMeasure{MillimeterUnit}, parent_box.y0),
                            convert(SimpleMeasure{MillimeterUnit}, parent_box.width),
                            convert(SimpleMeasure{MillimeterUnit}, parent_box.height))

            canvas = canvas.f(abs_parent_box, unit_box)
            if !(typeof(canvas) <: Canvas)
                error("Error: A deferred canvas function did not evaluate to a canvas.")
            end
        end

        box = native_measure(canvas.box, parent_t, unit_box, parent_box, backend)
        rot = native_measure(canvas.rot, parent_t, unit_box, box, backend)
        t = combine(rot, parent_t)

        if !canvas.units_inherited
            unit_box = convert(BoundingBox, canvas.unit_box)
        end

        property = canvas.property
        if canvas.clip
            property |= clip(Point(box.x0, box.y0),
                             Point(box.x0 + box.width, box.y0),
                             Point(box.x0 + box.width, box.y0 + box.height),
                             Point(box.x0, box.y0 + box.height))
        end

        if !is(property, empty_property)
            push!(S, :POP_PROPERTY)
            push_property(backend, native_measure(property, t, unit_box,
                                                  box, backend))
        end

        draw(backend, t, unit_box, box, canvas.form)

        ordered_children = Array((Integer, Integer, Canvas), length(canvas.children))
        for (i, child) in enumerate(canvas.children)
            # The tuple here includes i so canvases are never compared to break
            # ties.
            ordered_children[i] = (child.order, i, child)
        end
        sort!(ordered_children)
        reverse!(ordered_children)

        for (_, _, child) in ordered_children
            push!(S, (child, t, unit_box, box))
        end
    end
end


text:
using SpecialFunctions: erf
using SpecialFunctions: gamma
using Polynomials

struct PspHgh{T} <: NormConservingPsp
    Zion::Int             # Ionic charge (Z - valence electrons)
    rloc::T               # Range of local Gaussian charge distribution
    cloc::SVector{4,T}    # Coefficients for the local part
    lmax::Int             # Maximal angular momentum in the non-local part
    rp::Vector{T}         # Projector radius for each angular momentum
    h::Vector{Matrix{T}}  # Projector coupling coefficients per AM channel: h[l][i1,i2]
    identifier::String    # String identifying the PSP
    description::String   # Descriptive string
end
charge_ionic(psp::PspHgh) = psp.Zion

"""
    PspHgh(Zion::Number, rloc::Number, cloc::Vector, rp::Vector, h::Vector;
           identifier="", description="")

Construct a Hartwigsen, Goedecker, Teter, Hutter separable dual-space
Gaussian pseudopotential (1998). The required parameters are the ionic
charge `Zion` (total charge - valence electrons), the range for the local
Gaussian charge distribution `rloc`, the coefficients for the local part
`cloc`, the projector radius `rp` (one per AM channel) and the non-local
coupling coefficients between the projectors `h` (one matrix per AM channel).
"""
function PspHgh(Zion, rloc::T, cloc::AbstractVector, rp, h; identifier="", description="") where {T}
    length(rp) == length(h) || error("Length of rp and h do not agree.")
    length(cloc) <= 4 || error("length(cloc) > 4 not supported.")
    if length(cloc) < 4
        n_extra = 4 - length(cloc)
        cloc = [cloc; zeros(T, n_extra)]
    end

    lmax = length(h) - 1
    PspHgh{T}(Zion, rloc, cloc, lmax, rp, h, identifier, description)
end

function parse_hgh_file(path; identifier=path)
    lines = readlines(path)
    description = lines[1]

    # lines[2] contains the number of electrons (and the AM channel in which they sit)
    m = match(r"^ *(([0-9]+ *)+)", lines[2])
    n_elec = [parse(Int, part) for part in split(m[1])]
    Zion = sum(n_elec)

    # lines[3] contains rloc nloc and coefficients for it
    m = match(r"^ *([-.0-9]+) +([0-9]+)( +([-.0-9]+ *)+)? *", lines[3])
    rloc = parse(Float64, m[1])
    nloc = parse(Int, m[2])

    cloc = []
    m[3] !== nothing && (cloc = [parse(Float64, part) for part in split(m[3])])
    @assert length(cloc) == nloc

    # lines[4] contains the maximal AM channel
    m = match(r"^ *([0-9]+)", lines[4])
    lmax = parse(Int, m[1]) - 1

    rp = Vector{Float64}(undef, lmax + 1)
    h = Vector{Matrix{Float64}}(undef, lmax + 1)
    cur = 5  # Current line to parse

    for l in 0:lmax
        # loop over all AM channels and extract projectors,
        # these are given in blocks like
        #
        #    0.42273813    2     5.90692831    -1.26189397
        #                                       3.25819622
        #    0.48427842    1     2.72701346
        #
        # In each such blocks, the first number is the rp, the next is the number
        # of projectors and the following numbers (including the indented continuation
        # in the next lines) is the upper triangle of the coupling matrix h for this AM.
        # Then the next block commences unindented.

        # Match the first line of a block:
        m = match(r"^ *([-.0-9]+) +([0-9]+)( +([-.0-9]+ *)+)? *", lines[cur])
        rp[l + 1] = parse(Float64, m[1])
        nproj = parse(Int, m[2])
        h[l + 1] = Matrix{Float64}(undef, nproj, nproj)

        # If there are no projectors for this AM channel nproj is zero
        # and we can just increase cur and move on to the next block.
        if nproj == 0
            cur += 1
            continue
        end

        # Else we have to parse the extra parts of the hcoeff matrix.
        # This is done here.
        hcoeff = [parse(Float64, part) for part in split(m[3])]
        for i in 1:nproj
            for j in i:nproj
                h[l + 1][j, i] = h[l + 1][i, j] = hcoeff[j - i + 1]
            end

            # Parse the next (indented) line
            cur += 1
            cur > length(lines) && break
            m = match(r"^ *(([-.0-9]+ *)+)", lines[cur])
            hcoeff = [parse(Float64, part) for part in split(m[1])]
        end
    end

    PspHgh(Zion, rloc, cloc, rp, h; identifier, description)
end

@doc raw"""
The local potential of a HGH pseudopotentials in reciprocal space
can be brought to the form ``Q(t) / (t^2 exp(t^2 / 2))``
where ``t = r_\text{loc} q`` and `Q`
is a polynomial of at most degree 8. This function returns `Q`.
"""
@inline function psp_local_polynomial(T, psp::PspHgh, t=Polynomial(T[0, 1]))
    rloc::T = psp.rloc
    Zion::T = psp.Zion

    # The polynomial prefactor P(t) (as used inside the { ... } brackets of equation
    # (5) of the HGH98 paper)
    P = (  psp.cloc[1]
         + psp.cloc[2] * (  3 -    t^2              )
         + psp.cloc[3] * ( 15 -  10t^2 +   t^4      )
         + psp.cloc[4] * (105 - 105t^2 + 21t^4 - t^6))

    4T(π) * rloc^2 * (-Zion + sqrt(T(π) / 2) * rloc * t^2 * P)
end

# [GTH98] (6) except they do it with plane waves normalized by 1/sqrt(Ω).
function eval_psp_local_fourier(psp::PspHgh, q::T) where {T <: Real}
    t::T = q * psp.rloc
    psp_local_polynomial(T, psp, t) * exp(-t^2 / 2) / t^2
end

@doc raw"""
Estimate an upper bound for the argument `q` after which
`abs(eval_psp_local_fourier(psp, q))` is a strictly decreasing function.
"""
function qcut_psp_local(psp::PspHgh{T}) where {T}
    Q = DFTK.psp_local_polynomial(T, psp)  # polynomial in t = q * rloc

    # Find the roots of the derivative polynomial:
    res = roots(Polynomial([0, 1]) * derivative(Q) - Polynomial([2, 0, 1]) * Q)
    res = T[r for r in res if abs(imag(r)) < 100eps(T)]
    maximum(res, init=zero(T)) / psp.rloc
end

# [GTH98] (1)
function eval_psp_local_real(psp::PspHgh, r::T) where {T <: Real}
    r == 0 && return eval_psp_local_real(psp, eps(T)) # quick hack for the division by zero below
    cloc = psp.cloc
    rr = r / psp.rloc
    convert(T,
        - psp.Zion / r * erf(rr / sqrt(T(2)))
        + exp(-rr^2 / 2) * (cloc[1] + cloc[2] * rr^2 + cloc[3] * rr^4 + cloc[4] * rr^6)
    )
end

@doc raw"""
The nonlocal projectors of a HGH pseudopotentials in reciprocal space
can be brought to the form ``Q(t) exp(-t^2 / 2)`` where ``t = r_l q``
and `Q` is a polynomial. This function returns `Q`.
"""
@inline function psp_projector_polynomial(T, psp::PspHgh, i, l, t=Polynomial(T[0, 1]))
    @assert 0 <= l <= length(psp.rp) - 1
    @assert i > 0
    rp::T = psp.rp[l + 1]
    common::T = 4T(π)^(5 / T(4)) * sqrt(T(2^(l + 1)) * rp^3)

    # Note: In the (l == 0 && i == 2) case the HGH paper has an error.
    #       The first 8 in equation (8) should not be under the sqrt-sign
    #       This is the right version (as shown in the GTH paper)
    (l == 0 && i == 1) && return convert(typeof(t), common)
    (l == 0 && i == 2) && return common * 2 /  sqrt(T(  15))       * ( 3 -   t^2      )
    (l == 0 && i == 3) && return common * 4 / 3sqrt(T( 105))       * (15 - 10t^2 + t^4)
    #
    (l == 1 && i == 1) && return common * 1 /  sqrt(T(   3)) * t
    (l == 1 && i == 2) && return common * 2 /  sqrt(T( 105)) * t   * ( 5 -   t^2)
    (l == 1 && i == 3) && return common * 4 / 3sqrt(T(1155)) * t   * (35 - 14t^2 + t^4)
    #
    (l == 2 && i == 1) && return common * 1 /  sqrt(T(  15)) * t^2
    (l == 2 && i == 2) && return common * 2 / 3sqrt(T( 105)) * t^2 * ( 7 -   t^2)
    #
    (l == 3 && i == 1) && return common * 1 /  sqrt(T( 105)) * t^3

    error("Not implemented for l=$l and i=$i")
end

@doc raw"""
Estimate an upper bound for the argument `q` after which
`eval_psp_projector_fourier(psp, q)` is a strictly decreasing function.
"""
function qcut_psp_projector(psp::PspHgh{T}, i, l) where {T}
    Q = DFTK.psp_projector_polynomial(T, psp, i, l)  # polynomial in q * rp[l + 1]

    # Find the roots of the derivative polynomial:
    res = roots(derivative(Q) - Polynomial([0, 1]) * Q)
    res = T[r for r in res if abs(imag(r)) < 100eps(T)]
    maximum(res, init=zero(T)) / psp.rp[l + 1]
end

# [HGH98] (7-15) except they do it with plane waves normalized by 1/sqrt(Ω).
function eval_psp_projector_fourier(psp::PspHgh, i, l, q::T) where {T <: Real}
    t::T = q * psp.rp[l + 1]
    psp_projector_polynomial(T, psp, i, l, t) * exp(-t^2 / 2)
end

# [HGH98] (3)
function eval_psp_projector_real(psp::PspHgh, i, l, r::T) where {T <: Real}
    rp = T(psp.rp[l + 1])
    ired = (4i - 1) / T(2)
    sqrt(T(2)) * r^(l + 2(i - 1)) * exp(-r^2 / 2rp^2) / rp^(l + ired) / sqrt(gamma(l + ired))
end

function eval_psp_energy_correction(T, psp::PspHgh, n_electrons)
    # By construction we need to compute the DC component of the difference
    # of the Coulomb potential (-Z/G^2 in Fourier space) and the pseudopotential
    # i.e. -Z/(ΔG)^2 -  eval_psp_local_fourier(psp, ΔG) for ΔG → 0. This is:
    cloc_coeffs = T[1, 3, 15, 105]
    difference_DC = (T(psp.Zion) * T(psp.rloc)^2 / 2
                     + sqrt(T(π)/2) * T(psp.rloc)^3 * T(sum(cloc_coeffs .* psp.cloc)))

    # Multiply by number of electrons and 4π (spherical Hankel prefactor)
    # to get energy per unit cell
    4T(π) * n_electrons * difference_DC
end


text:
## extensions for Winston

## Pkg.installed() is _slow_
if "Winston" in readdir(Pkg.dir())
    using Winston

    ## add display method with graphics device as first argument
    ## not in multimedia spec, but is in winston
    function Base.display(cg::CairoGraphics, pc::Winston.PlotContainer)
        c = cg[:widget]
        c.draw = function(_)
            ctx = Base.Graphics.getgc(c)
            Base.Graphics.set_source_rgb(ctx, 1, 1, 1)
            Cairo.paint(ctx)
            Winston.page_compose(pc, Gtk.cairo_surface(c))
        end
            Gtk.draw(c)
        end

        function DisplayPlot(self::ManipulateObject, x::Winston.FramedPlot; kwargs...)
            if isa(x, Nothing) return end
            oa = self.output_area

            if length(children(oa)) > 0 && isa(oa.children[1], CairoGraphics)
                cnv = oa.children[1]
                display(cnv, x)
#                Winston.display(cnv.o, x)
                return
            elseif length(children(oa)) > 0
                pop!(oa)
            end
            ## add one
            cnv = cairographic(oa, width=480, height=480)
            cnv[:sizepolicy] = (:expand, :expand)
            push!(oa, cnv)
            display(cnv, x)
#            Winston.display(cnv.o, x)
        end

        ## union has some ambiguity warnings
        Display(::MIME"application/x-tcltk", self::ManipulateObject, x::Winston.FramedPlot; kwargs...) = DisplayPlot(self, x; kwargs...)
        Display(::MIME"application/x-gtk", self::ManipulateObject, x::Winston.FramedPlot; kwargs...) = DisplayPlot(self, x; kwargs...)

end


text:
"""
    readSOPFile(filePath::String)

Read a properly formatted (SOP according to TSPLIB) .sop file and return a SequencingModel

# Arguments
- `filePath::String`: The location of the sop file
"""
function readSOPFile(filePath::String)
    model = nothing
    open(filePath) do openedFile
        model = parseSOPInput!(read(openedFile, String))
    end
    return model
end

"""
    parseSOPInput!(input_data::String)

Convert the contents of a file, and return a SequencingModel

# Arguments
- `input_data::String`: The contents of file
"""
function parseSOPInput!(input_data::String)
    #split the file into lines and store each line as a node object
    lines = split(input_data, '\n')

    #values
    dimension = 0
    objectiveValues = Dict{Tuple{Int, Int}, Int}()
    precedenceP2F = Dict{Int, Vector{Int}}()#map of priors to followers
    precedenceF2P = Dict{Int, Vector{Int}}()#map of followers to priors

    #bools for reaing file
    matrixSectionHeading = false
    matrixSection = false

    #counter for reading file
    n = 1

    #main loop
    for i in 1:length(lines)
        splitLine = split(lines[i], (' ','\t','\r'))
        filter!(x -> x!="", splitLine)
        if length(splitLine)>=1
            if occursin("EDGE_WEIGHT_SECTION",splitLine[1])
                matrixSectionHeading = true
            elseif matrixSectionHeading
                matrixSectionHeading = false
                matrixSection = true
                dimension = parse(Int,splitLine[1])
            elseif matrixSection && n <= dimension
                for j in 1:length(splitLine)
                    value  = splitLine[j]
                    if length(value)>=1
                        val = parse(Int,value)
                        #if precedence constraint
                        if val == -1 && n!=1 && j!=1
                            if haskey(precedenceP2F,j)
                                push!(precedenceP2F[j],n)
                            else
                                precedenceP2F[j] = [n]
                            end
                            if haskey(precedenceF2P,n)
                                push!(precedenceF2P[n],j)
                            else
                                precedenceF2P[n] = [j]
                            end
                        elseif n != j && val!= -1 #else if it does not point to itself
                            objectiveValues[(n, j)] = val
                        end
                    end
                end
                n = n+1
            end
        end
    end
    delete!(objectiveValues, (1, dimension))
    precedenceNumbers = getPrecedenceNumbers(precedenceF2P)
    for i in 2:dimension
        if !haskey(precedenceNumbers,i)
            precedenceNumbers[i] = 1
        else
            precedenceNumbers[i] += 1
        end
    end
    return SequencingModel(
        ExtensionalArcObjective(ExtensionalArcFunction(objectiveValues),minimization),
        dimension,
        1,
        dimension,
        true,
        precedenceP2F,
        precedenceF2P,
        precedenceNumbers
    )
end

"""
    printSOPSolution(bestKnownSolution::Vector{Int}, bestKnownValue::T) where T<:Real

Print a SOP solution to the terminal

# Arguments
- `bestKnownSolution::Vector{Int}`: The solution as an ordered list of nodes
- `bestKnownValue::T`: The cost of the solution
"""
function printSOPSolution(bestKnownSolution::Vector{Int}, bestKnownValue::T) where T<:Real
    println()
    print("Sequence: ",)
    for i in 1:length(bestKnownSolution)
        print(bestKnownSolution[i], " ")
    end
    if isinteger(bestKnownValue)
        println("\nCost: ", Int(bestKnownValue))
    else
        println("\nCost: ", bestKnownValue)
    end
end

"""
    writeSOPSolution(bestKnownSolution::Vector{Int}, bestKnownValue::T, fileName::String; time::Union{U,Nothing}=nothing, timeUnit::Union{String,Nothing}=nothing) where{T<:Real, U<:Real}

Print a SOP solution to the terminal

# Arguments
- `bestKnownSolution::Vector{Int}`: The solution as an ordered list of nodes
- `bestKnownValue::T`: The cost of the solution
- `bestKnownValue::T`: The best known relaxed bound
- `fileName::String`: The file to write to
- `time::U`: The time elapsed
- `timeUnit`: The unit of time being used
"""
function writeSOPSolution(bestKnownSolution::Vector{Int}, bestKnownValue::T, bestKnownBound::T, fileName::String; time::Union{U,Nothing}=nothing, timeUnit::Union{String,Nothing}=nothing) where{T<:Real, U<:Real}
    f = open(fileName,"a")
        write(f, "Seqeunce: ")
        for i in 1:length(bestKnownSolution)
            write(f,string(bestKnownSolution[i], " "))
        end
        if isinteger(bestKnownValue)
            write(f,"\nCost: ", string(Int(bestKnownValue)))
        else
            write(f,"\nCost: ", string(bestKnownValue))
        end
        if isinteger(bestKnownBound)
            write(f,"\nRelaxed Bound: ", string(Int(bestKnownBound)))
        else
            write(f,"\nRelaxed Bound: ", string(bestKnownBound))
        end
        write(f,"\n")
        if !isnothing(time)
            write(f, "Time Elapsed: ", string(time), " $timeUnit\n")
        end
    close(f)
end


text:
using DiscreteMarkovChains
using Test
using LinearAlgebra

@testset "Converting To Discrete" begin
    # Empty test
    T = [][:,:]
    X = ContinuousMarkovChain(T)
    @test DiscreteMarkovChain(X) == DiscreteMarkovChain(T)

    state_space = [1, 2, 3]
    T = [
        0 1 0;
        1 0 0;
        0 0 1;
    ] - I
    X = ContinuousMarkovChain(state_space, T)

    T1 = transition_matrix(DiscreteMarkovChain(X))
    T2 = transition_matrix(DiscreteMarkovChain(state_space, exp(T)))
    @test T1 == T2
end

@testset "Converting To Continuous" begin
    # Empty test
    T = [][:,:]
    X = DiscreteMarkovChain(T)
    @test ContinuousMarkovChain(X) == ContinuousMarkovChain(T)

    state_space = [1, 2, 3]
    T = [
        0.2 0.8 0;
        0 0.2 0.8;
        0 0 1;
    ]
    X = DiscreteMarkovChain(state_space, T)

    Q1 = transition_matrix(ContinuousMarkovChain(X))
    Q2 = transition_matrix(ContinuousMarkovChain(state_space, log(T)))
    @test Q1 == Q2
end

@testset "Embedded Chains" begin
    # Empty test
    T = [][:,:]
    X = ContinuousMarkovChain(T)
    @test embedded(X) == DiscreteMarkovChain(T)

    X = DiscreteMarkovChain(T)
    @test embedded(X) == X

    T = [
        -2 1 1;
        1 -1 0;
        2 3 -5;
    ]
    X = ContinuousMarkovChain(T)
    desired_new_matrix = [
        0 0.5 0.5;
        1 0 0;
        0.4 0.6 0;
    ]
    @test transition_matrix(embedded(X)) ≈ desired_new_matrix
end

@testset "Communication Classes" begin
    # Empty test
    T = [][:,:]
    X = ContinuousMarkovChain(T)
    @test communication_classes(X) == ([], [])

    state_space = [1, 2, 3]
    T = [
        0 1 0;
        1 0 0;
        0 0 1;
    ] - I
    X = ContinuousMarkovChain(state_space, T)
    @test communication_classes(X) == ([[1, 2], [3]], [true, true])

    T = [
        0 5 5 0 0;
        0 0 0 10 0;
        5 0 5 0 0;
        0 10 0 0 0;
        0 3 0 3 4;
    ]/10 - I
    X = ContinuousMarkovChain(0:4, T)
    classes, recurrence = communication_classes(X)
    @test classes == [[1, 3], [0, 2], [4]]
    @test recurrence == [true, false, false]

    T = [
        0 0 0 10 0 0;
        5 0 5 0 0 0;
        0 4 0 0 0 6;
        10 0 0 0 0 0;
        0 10 0 0 0 0;
        0 0 0 5 5 0;
    ]/10 - I
    X = ContinuousMarkovChain(0:5, T)
    classes, recurrence = communication_classes(X)
    @test classes == [[0, 3], [1, 2, 5, 4]]
    @test recurrence == [true, false]

    T = [
        2 0 0 3 0 0 3 2 0 0;
        0 10 0 0 0 0 0 0 0 0;
        0 2 2 0 0 0 0 0 3 3;
        0 0 0 3 0 0 6 1 0 0;
        0 0 0 0 5 5 0 0 0 0;
        0 0 0 0 0 10 0 0 0 0;
        4 0 0 5 0 0 1 0 0 0;
        2 0 0 4 0 0 2 2 0 0;
        3 0 1 0 0 0 0 0 4 2;
        0 0 4 0 0 0 0 0 3 3;
    ]/10 - I
    X = ContinuousMarkovChain(0:9, T)
    classes, recurrence = communication_classes(X)
    @test classes == [[0, 3, 6, 7], [1], [2, 8, 9], [5], [4]]
    @test recurrence == [true, true, false, true, false]
end

@testset "Decompose and Canonical Form" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test decompose(X) == ([], T, T, T)

    T = [
        1 0 0;
        1//3 1//3 1//3;
        0 1//4 3//4;
    ] - I
    X = ContinuousMarkovChain(["Sunny", "Cloudy", "Rainy"], T)
    @test decompose(X) == (
        ["Sunny", "Cloudy", "Rainy"],
        T[1:1, 1:1], T[2:3, 1:1], T[2:3, 2:3],
    )
    @test canonical_form(X) == (["Sunny", "Cloudy", "Rainy"], T)

    T = [
        1//4 3//4 0;
        1//3 1//3 1//3;
        0 1//4 3//4;
    ] - I
    X = ContinuousMarkovChain([1, 2, 3], T)
    @test decompose(X) == ([1, 2, 3], T, Array{Any}(undef, 0, 3), Array{Any}(undef, 0, 0))
    @test canonical_form(X) == ([1, 2, 3], T)

    T = [
        1 0 0 0 0;
        1//2 0 1//2 0 0;
        0 1//2 0 1//2 0;
        0 0 1//2 0 1//2;
        0 0 0 0 1;
    ] - I
    X = ContinuousMarkovChain(0:4, T)
    states, A, B, C = decompose(X)
    @test states == [0, 4, 1, 2, 3]
    @test A == [1 0; 0 1] - I
    @test B == [1//2 0; 0 0; 0 1//2]
    @test C == [0 1//2 0; 1//2 0 1//2; 0 1//2 0] - I
    states, new_matrix = canonical_form(X)
    @test states == [0, 4, 1, 2, 3]
    desired_new_matrix = [
        1 0 0 0 0;
        0 1 0 0 0;
        1//2 0 0 1//2 0;
        0 0 1//2 0 1//2;
        0 1//2 0 1//2 0;
    ] - I
    @test new_matrix == desired_new_matrix
end

@testset "Regular, Ergodic, Absorbing" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test !is_ergodic(X)
    @test !is_absorbing(X)

    T = [
        0 4 0 0 0;
        1 0 3 0 0;
        0 2 0 2 0;
        0 0 3 0 1;
        0 0 0 4 0;
    ]/4.0 - I
    X = ContinuousMarkovChain(1:5, T)
    @test is_ergodic(X)
    @test !is_absorbing(X)

    T = [
        0 1;
        1 0;
    ] - I
    X = ContinuousMarkovChain(1:2, T)
    @test is_ergodic(X)
    @test !is_absorbing(X)

    T = [
        2 1 1;
        2 0 2;
        1 1 2;
    ]/4 - I
    X = ContinuousMarkovChain(1:3, T)
    @test is_ergodic(X)
    @test !is_absorbing(X)

    T = [
        1 1;
        1 1;
    ]/2 - I
    X = ContinuousMarkovChain(1:2, T)
    @test is_ergodic(X)
    @test !is_absorbing(X)
end

@testset "Reversible" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = DiscreteMarkovChain(T)
    @test is_reversible(X)
    X = ContinuousMarkovChain(T)
    @test is_reversible(X)

    T = [
        .425 .000 .075 .500;
        .000 .550 .250 .200;
        .300 .250 .450 .000;
        .500 .050 .000 .450;
    ]
    X = DiscreteMarkovChain(T)
    @test is_reversible(X)
    X = ContinuousMarkovChain(T-I)
    @test is_reversible(X)

    T = [
        0.5 0.5 0.0;
        0.0 0.2 0.8;
        0.3 0.3 0.4;
    ]
    X = DiscreteMarkovChain(T)
    @test !is_reversible(X)
    X = ContinuousMarkovChain(T-I)
    @test !is_reversible(X)
end

@testset "Stationary distribution" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test stationary_distribution(X) == Any[]

    T = [-2 2 0; 0 -4 4; 0 0 0]
    X = ContinuousMarkovChain(T)
    @test stationary_distribution(X) == [0, 0, 1]

    T = [
        -25 20 5;
        300 -500 200;
        20 400 -420;
    ]//1000
    X = ContinuousMarkovChain(T)
    @test stationary_distribution(X) == [100//113, 8//113, 5//113]
end

@testset "Fundamental Matrix" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test fundamental_matrix(X) == Array{Any}(undef, 0, 0)

    # http://www.mat.ufrgs.br/~giacomo/Softwares/R/Dobrow%20-%20Stochastic%20with%20R/ch7.pdf
    # Example 7.16
    T = [
        0 0 0;
        1 -1 0;
        1 1 -2;
    ]
    X = ContinuousMarkovChain(T)
    desired_new_matrix = [
        1 0;
        1/2 1/2;
    ]
    @test fundamental_matrix(X) == desired_new_matrix
end

@testset "Expected Time To Absorption" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test mean_time_to_absorption(X) == Any[]

    T = [
        0 0 0;
        4 -4 0;
        0 2 -2;
    ]
    X = ContinuousMarkovChain(T)
    @test mean_time_to_absorption(X) == [0.25, 0.75]
end

@testset "Exit Probability" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test exit_probabilities(X) == Array{Any}(undef, 0, 0)

    T = [
        0 0 0;
        0 -1 1;
        2 2 -4;
    ]
    X = ContinuousMarkovChain(T)
    @test exit_probabilities(X) == [1; 1][:, :]

    T = [
        0 0 0;
        0 0 0;
        8 8 -16;
    ]
    X = ContinuousMarkovChain(T)
    @test exit_probabilities(X) == [0.5 0.5]
end

@testset "Mean Recurrence Time" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test mean_recurrence_time(X) == Any[]

    T = [
        -1 1 0;
        0 -1 1;
        1 0 -1;
    ]
    X = ContinuousMarkovChain(T)
    @test mean_recurrence_time(X) == [3, 3, 3]

    T = [
        -1 1 0;
        0 -1 1;
        2 2 -4;
    ]
    X = ContinuousMarkovChain(T)
    @test mean_recurrence_time(X) == [3.5, 1.75, 1.75]
end

@testset "Mean First Passage Time" begin
    # Empty test
    T = Array{Any}(undef, 0, 0)
    X = ContinuousMarkovChain([], T)
    @test mean_first_passage_time(X) == Array{Any}(undef, 0, 0)

    T = [
        -1 1 0;
        0 -1 1;
        1 0 -1;
    ]
    X = ContinuousMarkovChain(T)
    @test mean_first_passage_time(X) == [
        0.0 1.0 2.0;
        2.0 0.0 1.0;
        1.0 2.0 0.0;
    ]

    T = [
        -1 1 0;
        0 -1 1;
        2 2 -4;
    ]
    X = ContinuousMarkovChain(T)
    @test mean_first_passage_time(X) == [
        0.0 1.0 2.0;
        2.5 0.0 1.0;
        1.5 0.75 0.0;
    ]
end


text:
#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct Bus <: Topology
        number::Int64
        name::String
        bustype::Union{Nothing, BusTypes.BusType}
        angle::Union{Nothing, Float64}
        voltage::Union{Nothing, Float64}
        voltagelimits::Union{Nothing, Min_Max}
        basevoltage::Union{Nothing, Float64}
        area::Union{Nothing, Area}
        load_zone::Union{Nothing, LoadZone}
        ext::Dict{String, Any}
        internal::InfrastructureSystemsInternal
    end

A power-system bus.

# Arguments
- `number::Int64`: number associated with the bus
- `name::String`: the name of the bus
- `bustype::Union{Nothing, BusTypes.BusType}`: bus type
- `angle::Union{Nothing, Float64}`: angle of the bus in radians, validation range: (-1.571, 1.571), action if invalid: error
- `voltage::Union{Nothing, Float64}`: voltage as a multiple of basevoltage, validation range: voltagelimits, action if invalid: warn
- `voltagelimits::Union{Nothing, Min_Max}`: limits on the voltage variation as multiples of basevoltage
- `basevoltage::Union{Nothing, Float64}`: the base voltage in kV, validation range: (0, nothing), action if invalid: error
- `area::Union{Nothing, Area}`: the area containing the bus
- `load_zone::Union{Nothing, LoadZone}`: the load zone containing the bus
- `ext::Dict{String, Any}`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct Bus <: Topology
    "number associated with the bus"
    number::Int64
    "the name of the bus"
    name::String
    "bus type"
    bustype::Union{Nothing, BusTypes.BusType}
    "angle of the bus in radians"
    angle::Union{Nothing, Float64}
    "voltage as a multiple of basevoltage"
    voltage::Union{Nothing, Float64}
    "limits on the voltage variation as multiples of basevoltage"
    voltagelimits::Union{Nothing, Min_Max}
    "the base voltage in kV"
    basevoltage::Union{Nothing, Float64}
    "the area containing the bus"
    area::Union{Nothing, Area}
    "the load zone containing the bus"
    load_zone::Union{Nothing, LoadZone}
    ext::Dict{String, Any}
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal

    function Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, internal, )
        (number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, internal, ) = CheckBusParams(
            number,
            name,
            bustype,
            angle,
            voltage,
            voltagelimits,
            basevoltage,
            area,
            load_zone,
            ext,
            internal,
        )
        new(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, internal, )
    end
end

function Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area=nothing, load_zone=nothing, ext=Dict{String, Any}(), )
    Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, InfrastructureSystemsInternal(), )
end

function Bus(; number, name, bustype, angle, voltage, voltagelimits, basevoltage, area=nothing, load_zone=nothing, ext=Dict{String, Any}(), )
    Bus(number, name, bustype, angle, voltage, voltagelimits, basevoltage, area, load_zone, ext, )
end

# Constructor for demo purposes; non-functional.
function Bus(::Nothing)
    Bus(;
        number=0,
        name="init",
        bustype=nothing,
        angle=0.0,
        voltage=0.0,
        voltagelimits=(min=0.0, max=0.0),
        basevoltage=nothing,
        area=nothing,
        load_zone=nothing,
        ext=Dict{String, Any}(),
    )
end

"""Get Bus number."""
get_number(value::Bus) = value.number

InfrastructureSystems.get_name(value::Bus) = value.name
"""Get Bus bustype."""
get_bustype(value::Bus) = value.bustype
"""Get Bus angle."""
get_angle(value::Bus) = value.angle
"""Get Bus voltage."""
get_voltage(value::Bus) = value.voltage
"""Get Bus voltagelimits."""
get_voltagelimits(value::Bus) = value.voltagelimits
"""Get Bus basevoltage."""
get_basevoltage(value::Bus) = value.basevoltage
"""Get Bus area."""
get_area(value::Bus) = value.area
"""Get Bus load_zone."""
get_load_zone(value::Bus) = value.load_zone
"""Get Bus ext."""
get_ext(value::Bus) = value.ext
"""Get Bus internal."""
get_internal(value::Bus) = value.internal

"""Set Bus number."""
set_number!(value::Bus, val::Int64) = value.number = val

InfrastructureSystems.set_name!(value::Bus, val::String) = value.name = val
"""Set Bus bustype."""
set_bustype!(value::Bus, val::Union{Nothing, BusTypes.BusType}) = value.bustype = val
"""Set Bus angle."""
set_angle!(value::Bus, val::Union{Nothing, Float64}) = value.angle = val
"""Set Bus voltage."""
set_voltage!(value::Bus, val::Union{Nothing, Float64}) = value.voltage = val
"""Set Bus voltagelimits."""
set_voltagelimits!(value::Bus, val::Union{Nothing, Min_Max}) = value.voltagelimits = val
"""Set Bus basevoltage."""
set_basevoltage!(value::Bus, val::Union{Nothing, Float64}) = value.basevoltage = val
"""Set Bus area."""
set_area!(value::Bus, val::Union{Nothing, Area}) = value.area = val
"""Set Bus load_zone."""
set_load_zone!(value::Bus, val::Union{Nothing, LoadZone}) = value.load_zone = val
"""Set Bus ext."""
set_ext!(value::Bus, val::Dict{String, Any}) = value.ext = val
"""Set Bus internal."""
set_internal!(value::Bus, val::InfrastructureSystemsInternal) = value.internal = val


text:
matchvertex2edges(x, edges) = x[:, [e.src for e in edges]]
function ∇matchvertex2edges(Δ, x, edges)
    o = similar(x) .= 0
    for (i,e) in enumerate(edges)
        o[:, e.src] .+= Δ[:,i]
    end
    o
end
matchvertex2edges(x::Flux.Tracker.TrackedMatrix, edges) = Flux.Tracker.track(matchvertex2edges, x, edges)
Flux.Tracker.@grad function matchvertex2edges(x, edges)
    return(matchvertex2edges(Flux.data(x), edges), Δ -> (∇matchvertex2edges(Δ, x, edges), nothing))
end


text:
# syntax: proto2
# COV_EXCL_START
using ProtoBuf
import ProtoBuf.meta
import ..hadoop

# service methods for ApplicationHistoryProtocolService
const _ApplicationHistoryProtocolService_methods = MethodDescriptor[
        MethodDescriptor("getApplicationReport", 1, hadoop.yarn.GetApplicationReportRequestProto, hadoop.yarn.GetApplicationReportResponseProto),
        MethodDescriptor("getApplications", 2, hadoop.yarn.GetApplicationsRequestProto, hadoop.yarn.GetApplicationsResponseProto),
        MethodDescriptor("getApplicationAttemptReport", 3, hadoop.yarn.GetApplicationAttemptReportRequestProto, hadoop.yarn.GetApplicationAttemptReportResponseProto),
        MethodDescriptor("getApplicationAttempts", 4, hadoop.yarn.GetApplicationAttemptsRequestProto, hadoop.yarn.GetApplicationAttemptsResponseProto),
        MethodDescriptor("getContainerReport", 5, hadoop.yarn.GetContainerReportRequestProto, hadoop.yarn.GetContainerReportResponseProto),
        MethodDescriptor("getContainers", 6, hadoop.yarn.GetContainersRequestProto, hadoop.yarn.GetContainersResponseProto),
        MethodDescriptor("getDelegationToken", 7, hadoop.common.GetDelegationTokenRequestProto, hadoop.common.GetDelegationTokenResponseProto),
        MethodDescriptor("renewDelegationToken", 8, hadoop.common.RenewDelegationTokenRequestProto, hadoop.common.RenewDelegationTokenResponseProto),
        MethodDescriptor("cancelDelegationToken", 9, hadoop.common.CancelDelegationTokenRequestProto, hadoop.common.CancelDelegationTokenResponseProto)
    ] # const _ApplicationHistoryProtocolService_methods
const _ApplicationHistoryProtocolService_desc = ServiceDescriptor("hadoop.yarn.ApplicationHistoryProtocolService", 1, _ApplicationHistoryProtocolService_methods)

ApplicationHistoryProtocolService(impl::Module) = ProtoService(_ApplicationHistoryProtocolService_desc, impl)

mutable struct ApplicationHistoryProtocolServiceStub <: AbstractProtoServiceStub{false}
    impl::ProtoServiceStub
    ApplicationHistoryProtocolServiceStub(channel::ProtoRpcChannel) = new(ProtoServiceStub(_ApplicationHistoryProtocolService_desc, channel))
end # mutable struct ApplicationHistoryProtocolServiceStub

mutable struct ApplicationHistoryProtocolServiceBlockingStub <: AbstractProtoServiceStub{true}
    impl::ProtoServiceBlockingStub
    ApplicationHistoryProtocolServiceBlockingStub(channel::ProtoRpcChannel) = new(ProtoServiceBlockingStub(_ApplicationHistoryProtocolService_desc, channel))
end # mutable struct ApplicationHistoryProtocolServiceBlockingStub

getApplicationReport(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationReportRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[1], controller, inp, done)
getApplicationReport(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationReportRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[1], controller, inp)

getApplications(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationsRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[2], controller, inp, done)
getApplications(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationsRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[2], controller, inp)

getApplicationAttemptReport(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationAttemptReportRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[3], controller, inp, done)
getApplicationAttemptReport(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationAttemptReportRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[3], controller, inp)

getApplicationAttempts(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationAttemptsRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[4], controller, inp, done)
getApplicationAttempts(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.yarn.GetApplicationAttemptsRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[4], controller, inp)

getContainerReport(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.yarn.GetContainerReportRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[5], controller, inp, done)
getContainerReport(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.yarn.GetContainerReportRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[5], controller, inp)

getContainers(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.yarn.GetContainersRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[6], controller, inp, done)
getContainers(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.yarn.GetContainersRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[6], controller, inp)

getDelegationToken(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.common.GetDelegationTokenRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[7], controller, inp, done)
getDelegationToken(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.common.GetDelegationTokenRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[7], controller, inp)

renewDelegationToken(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.common.RenewDelegationTokenRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[8], controller, inp, done)
renewDelegationToken(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.common.RenewDelegationTokenRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[8], controller, inp)

cancelDelegationToken(stub::ApplicationHistoryProtocolServiceStub, controller::ProtoRpcController, inp::hadoop.common.CancelDelegationTokenRequestProto, done::Function) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[9], controller, inp, done)
cancelDelegationToken(stub::ApplicationHistoryProtocolServiceBlockingStub, controller::ProtoRpcController, inp::hadoop.common.CancelDelegationTokenRequestProto) = call_method(stub.impl, _ApplicationHistoryProtocolService_methods[9], controller, inp)

export ApplicationHistoryProtocolService, ApplicationHistoryProtocolServiceStub, ApplicationHistoryProtocolServiceBlockingStub, getApplicationReport, getApplications, getApplicationAttemptReport, getApplicationAttempts, getContainerReport, getContainers, getDelegationToken, renewDelegationToken, cancelDelegationToken
# COV_EXCL_STOP


text:
function removesubtype(typ)
    if typ isa Expr && typ.head == :(<:)
        typ.args[1]
    else
        typ
    end
end

function adjtype(Atyp)
    if Atyp isa Expr && Atyp.args[1] == :QuasiAdjoint
        removesubtype(Atyp.args[3])
    else
        :(QuasiAdjoint{<:Any,<:$Atyp})
    end
end

macro simplify(qt)
    qt.head == :function || qt.head == :(=) || error("Must start with a function")
    @assert qt.args[1].head == :call
    if qt.args[1].args[1] == :(*)
        mat = qt.args[2]
        @assert qt.args[1].args[2].head == :(::)
        Aname,Atyp = qt.args[1].args[2].args
        Aadj = adjtype(Atyp)
        Bname,Btyp = qt.args[1].args[3].args
        Badj = adjtype(Btyp)
        if length(qt.args[1].args) == 3
            ret = quote
                ContinuumArrays.simplifiable(::typeof(*), A::$Atyp, B::$Btyp) = Val(true)
                function ContinuumArrays.mul($Aname::$Atyp, $Bname::$Btyp)
                    $mat
                end
            end
            if Aadj ≠ Btyp
                @assert Atyp ≠ Badj
                ret = quote
                    $ret
                    ContinuumArrays.simplifiable(::typeof(*), ::$Badj, A::$Aadj) = Val(true)
                    ContinuumArrays.mul(Bc::$Badj, Ac::$Aadj) = ContinuumArrays.mul(Ac', Bc')'
                end
            end

            esc(ret)
        else
            @assert length(qt.args[1].args) == 4
            Cname,Ctyp = qt.args[1].args[4].args
            esc(quote
                ContinuumArrays.simplifiable(::typeof(*), ::$Atyp, ::$Btyp, ::$Ctyp) = Val(true)
                function ContinuumArrays._simplify(::typeof(*), $Aname::$Atyp, $Bname::$Btyp, $Cname::$Ctyp)
                    $mat
                end
                Base.copy(M::ContinuumArrays.QMul3{<:$Atyp,<:$Btyp,<:$Ctyp}) = ContinuumArrays.simplify(M)
            end)
        end
    end
end

struct DiracDelta{T,A} <: LazyQuasiVector{T}
    x::T
    axis::A
    DiracDelta{T,A}(x, axis) where {T,A} = new{T,A}(x,axis)
end

DiracDelta{T}(x, axis::A) where {T,A<:AbstractQuasiVector} = DiracDelta{T,A}(x, axis)
DiracDelta{T}(x, domain) where T = DiracDelta{T}(x, Inclusion(domain))
DiracDelta(x, domain) = DiracDelta{eltype(x)}(x, Inclusion(domain))
DiracDelta(axis::AbstractQuasiVector) = DiracDelta(zero(float(eltype(axis))), axis)
DiracDelta(domain) = DiracDelta(Inclusion(domain))

axes(δ::DiracDelta) = (δ.axis,)
IndexStyle(::Type{<:DiracDelta}) = IndexLinear()

==(a::DiracDelta, b::DiracDelta) = a.axis == b.axis && a.x == b.x

function getindex(δ::DiracDelta{T}, x::Number) where T
    x ∈ δ.axis || throw(BoundsError())
    x == δ.x ? inv(zero(T)) : zero(T)
end

@simplify *(A::QuasiAdjoint{<:Any,<:DiracDelta}, B::AbstractQuasiVector) = B[parent(A).x]
@simplify *(A::QuasiAdjoint{<:Any,<:DiracDelta}, B::AbstractQuasiMatrix) = B[parent(A).x,:]

show(io::IO, δ::DiracDelta) = print(io, "δ at $(δ.x) over $(axes(δ,1))")
show(io::IO, ::MIME"text/plain", δ::DiracDelta) = show(io, δ)

#########
# Derivative
#########

struct Derivative{T,D} <: LazyQuasiMatrix{T}
    axis::Inclusion{T,D}
end

Derivative{T}(axis::A) where {T,A<:Inclusion} = Derivative{T,A}(axis)
Derivative{T}(domain) where T = Derivative{T}(Inclusion(domain))
Derivative(axis) = Derivative{Float64}(axis)

axes(D::Derivative) = (D.axis, D.axis)
==(a::Derivative, b::Derivative) = a.axis == b.axis
copy(D::Derivative) = Derivative(copy(D.axis))

function diff(d::AbstractQuasiVector)
    x = axes(d,1)
    Derivative(x)*d
end

^(D::Derivative, k::Integer) = ApplyQuasiArray(^, D, k)

# struct Multiplication{T,F,A} <: AbstractQuasiMatrix{T}
#     f::F
#     axis::A
# end

const Identity{T,D} = QuasiDiagonal{T,Inclusion{T,D}}

Identity(d::Inclusion) = QuasiDiagonal(d)


text:
@testset "ARN" begin
    @test arn(aws,"s3","foo/bar") == "arn:aws:s3:::foo/bar"
    @test arn(aws,"s3","foo")     == "arn:aws:s3:::foo"
    @test arn(aws,"sqs", "au-test-queue", "ap-southeast-2", "1234") ==
          "arn:aws:sqs:ap-southeast-2:1234:au-test-queue"
    @test arn(aws,"sns","*","*",1234) == "arn:aws:sns:*:1234:*"
    @test arn(aws,"iam","role/foo-role", "", 1234) == "arn:aws:iam::1234:role/foo-role"
end


text:
using QuantumCumulants
using Test
using OrdinaryDiffEq
using ModelingToolkit

@testset "scaling" begin

order = 2
N_c = 1 #number of clusters
@cnumbers Δc κ Γ2 Γ3 Γ23 η ν3 ν2
Δ2 = [Parameter(Symbol(:Δ2_, i)) for i=1:N_c]
Δ3 = [Parameter(Symbol(:Δ3_, i)) for i=1:N_c]
Ω3 = [Parameter(Symbol(:Ω3_, i)) for i=1:N_c]
g = [Parameter(Symbol(:g_, i)) for i=1:N_c]
N = [Parameter(Symbol(:N_, i)) for i=1:N_c]

# Define hilbert space
hf = FockSpace(:cavity)
ha = [NLevelSpace(Symbol(:atoms, j),3) for j=1:N_c]
ha_c = [ClusterSpace(ha[j],N[j],order) for j=1:N_c]
h = ⊗(hf, ha_c...)
# Define the fundamental operators
@qnumbers a::Destroy(h,1)
S(i,j,c) = Transition(h,Symbol(:σ, c),i, j, 1+c) #c=cluster

# Hamiltonian
H = Δc*a'*a + sum(Δ2[c]*sum(S(2,2,c)) for c=1:N_c) + sum(Δ3[c]*sum(S(3,3,c)) for c=1:N_c) +
    sum(Ω3[c]*(sum(S(3,1,c)) + sum(S(1,3,c))) for c=1:N_c) + sum(g[c]*(a'*sum(S(1,2,c)) + a*sum(S(2,1,c))) for c=1:N_c)
# Collapse operators
J = [a,[S(1,2,c) for c=1:N_c]..., [S(1,3,c) for c=1:N_c]..., [S(2,3,c) for c=1:N_c]..., [S(3,3,c) for c=1:N_c]..., [S(2,2,c) for c=1:N_c]..., a'a]
rates = [κ,[Γ2 for i=1:N_c]...,[Γ3 for i=1:N_c]...,[Γ23 for i=1:N_c]...,[ν3 for i=1:N_c]...,[ν2 for i=1:N_c]...,η]

# Derive equation for average photon number
ops = [a'a, S(2,2,1)[1], a'*S(1,2,1)[1]]
he = meanfield(ops,H,J;rates=rates,order=2)

# Custom filter function -- include only phase-invariant terms
ϕ(x) = 0
ϕ(x::Destroy) = -1
ϕ(x::Create) = 1
function ϕ(t::Transition)
    if (t.i==1 && t.j==2) || (t.i==3 && t.j==2) || (t.i==:g && t.j==:e)
        -1
    elseif (t.i==2 && t.j==1) || (t.i==2 && t.j==3) || (t.i==:e && t.j==:g)
        1
    else
        0
    end
end
ϕ(avg::Average) = ϕ(avg.arguments[1])
function ϕ(t::QuantumCumulants.QMul)
    p = 0
    for arg in t.args_nc
        p += ϕ(arg)
    end
    return p
end
phase_invariant(x) = iszero(ϕ(x))

he_avg = cumulant_expansion(he,2)
he_scale = complete(he_avg;filter_func=phase_invariant, order=order, multithread=false)
@test length(he_scale) == 9

ps = [Δc; κ; Γ2; Γ3; Γ23; η; ν3; ν2; Δ2; Δ3; Ω3; g; N]
@named sys = ODESystem(he_scale)
u0 = zeros(ComplexF64, length(he_scale))

N0 = 1000/N_c
N_ = N0*[1.0 for c=1:N_c]

p0 = ps .=> [ones(length(ps)-1); N_]
prob1 = ODEProblem(sys,u0,(0.0,1.0),p0)
sol1 = solve(prob1, Tsit5(), reltol=1e-12, abstol=1e-12)

@test sol1.u[end][1] ≈ 0.0758608728203

avg = average(a*S(2,1,1)[2])
# TODO: @test sol1[avg], sol1, he_scale) == map(conj, get_solution(average(a'*S(1,2,1)[1]), sol1, he_scale)) == conj.(getindex.(sol1.u, 3))

## Two-level laser
M = 2
@cnumbers N
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom, (:g,:e))
hc = ClusterSpace(ha, N, M)
h = tensor(hf, hc)

@qnumbers a::Destroy(h)
σ(i,j) = Transition(h,:σ,i,j)

@cnumbers Δ g κ γ ν

H = Δ*a'*a + g*sum(a'*σ(:g,:e)[i] + a*σ(:e,:g)[i] for i=1:M)
J = [a;[σ(:g,:e)[i] for i=1:M];[σ(:e,:g)[i] for i=1:M]]
rates = [κ; [γ for i=1:M]; [ν for i=1:M]]

he = meanfield(a'*a, H, J; rates=rates, order=2)

# Complete
he_scaled = complete(he;filter_func=phase_invariant)

names = he_scaled.names
avg = average(σ(:e,:g)[1]*σ(:e,:e)[2])
@test isequal(average(σ(:e,:e)[1]*σ(:e,:g)[2]), QuantumCumulants.substitute_redundants(avg,[QuantumCumulants.ClusterAon(2,1),QuantumCumulants.ClusterAon(2,2)],names))

@test QuantumCumulants.lt_reference_order(σ(:e,:g)[1],σ(:g,:e)[2])
@test !QuantumCumulants.lt_reference_order(σ(:g,:e)[1],σ(:e,:g)[2])

he_avg = cumulant_expansion(he_scaled,2)
@test isempty(find_missing(he_avg))

ps = (Δ, g, γ, κ, ν, N)
@named sys = ODESystem(he_avg)
p0 = ps .=> (0, 1.5, 0.25, 1, 4, 7)
u0 = zeros(ComplexF64, length(he_scaled))
prob = ODEProblem(sys, u0, (0.0, 50.0), p0)
sol = solve(prob, RK4(), abstol=1e-10, reltol=1e-10)

@test sol.u[end][1] ≈ 12.601868534

# Spectrum
corr = CorrelationFunction(a',a,he_avg;steady_state=true,filter_func=phase_invariant)
Spec = Spectrum(corr,ps)
s = Spec(range(-π, π, length=301), sol.u[end], getindex.(p0,2))
@test all(s .>= 0.0)

## Some abstract tests
M = 4
@cnumbers N
hc = FockSpace(:cavity)
hvib = FockSpace(:mode)
hcluster = ClusterSpace(hvib,N,M)
h = ⊗(hc, hcluster)
a = Destroy(h,:a,1)
b = Destroy(h,:b,2)

names = [:a,[Symbol(:b_,i) for i=1:M]]
scale_aons = [QuantumCumulants.ClusterAon(2,i) for i=1:M]

avg = average(a*b[1]*b[2]*b[2])
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(average(a*b[1]*b[1]*b[2]), avg_sub)

avg = average(b[1]'*b[2]'*b[2])
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(avg_sub, average(b[1]'*b[1]*b[2]'))

avg = average(b[1]*b[1]*b[2]'*b[2])
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(avg_sub, average(b[1]'*b[1]*b[2]*b[2]))

avg = average(b[1]*b[2]'*b[3]*b[4]')
avg_sub = QuantumCumulants.substitute_redundants(avg, scale_aons, names)
@test isequal(avg_sub, average(b[1]'*b[2]'*b[3]*b[4]))

# Test Holstein
M = 2
hc = FockSpace(:cavity)
hvib = FockSpace(:mode)
hcluster = ClusterSpace(hvib,N,M)
h = ⊗(hc, hcluster)
@cnumbers G Δ κ γ Ω
a = Destroy(h,:a,1)
b = Destroy(h,:b,2)

H = Δ*a'*a + G*sum(b[i] + b[i]' for i=1:M)*a'*a + Ω*(a+a')
J = [a,b]
rates = [κ,γ]
ops = [a,a'*a,a*a,b[1],a*b[1],a'*b[1],b[1]'*b[1],b[1]*b[1],b[1]'*b[2],b[1]*b[2]]
he = meanfield(ops,H,J;rates=rates)

he_avg = cumulant_expansion(he,2)
@test isempty(find_missing(he_avg))

ps = (G,Δ,κ,γ,Ω,N)
@named sys = ODESystem(he_avg)

u0 = zeros(ComplexF64, length(he_avg))
p0 = ps .=> ones(length(ps))
prob = ODEProblem(sys, u0, (0.0,1.0), p0)
sol = solve(prob, Tsit5())

# Test molecule
M = 2 # Order
# Prameters
@cnumbers λ ν Γ η Δ γ N
# Hilbert space
h_in = NLevelSpace(:internal, 2)
hv = FockSpace(:vib)
hc = ClusterSpace(hv, N, M)
h = ⊗(h_in, hc)
# Operators
σ(i,j) = Transition(h, :σ, i, j)
b = Destroy(h, :b)
# Hamiltonian
H0 = Δ*σ(2,2) + ν*sum(b_'*b_ for b_ in b)
H_holstein = -1*λ*sum((b_' + b_ for b_ in b))*σ(2,2)
Hl = η*(σ(1,2) + σ(2,1))
H = H0 + H_holstein + Hl
# Jumps
J = [σ(1,2), b]
rates = [γ,Γ]
# Equations
ops = [σ(2,2),σ(1,2),b[1],σ(1,2)*b[1],σ(2,1)*b[1],σ(2,2)*b[1],b[1]'*b[1],b[1]*b[1],b[1]'*b[2],b[1]*b[2]]
he = meanfield(ops, H, J; rates=rates)
he_avg = cumulant_expansion(he,2)
@test isempty(find_missing(he_avg))
ps = (Δ,η,γ,λ,ν,Γ,N)
# Generate function
@named sys = ODESystem(he_avg)
p0 = ps .=> [ones(length(ps)-1); 4]
u0 = zeros(ComplexF64,length(he_avg))
prob1 = ODEProblem(sys,u0,(0.0,1.0),p0)
sol1 = solve(prob1,Tsit5(),abstol=1e-12,reltol=1e-12)
bdb1 = sol1[b[1]'*b[1]][end]
σ22_1 = sol1[σ(2,2)][end]
σ12_1 = sol1[σ(1,2)][end]

## Two clusters
N_c = 2
N = cnumbers([Symbol(:N_, i) for i=1:N_c]...)
M = 2
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom, (:g,:e))
hc = [ClusterSpace(ha, N[i], M) for i=1:N_c]
h = tensor(hf, hc...)

@qnumbers a::Destroy(h)
σ(i,j,c) = Transition(h,Symbol(:σ_, c),i,j,c+1)

@cnumbers κ
ν = cnumbers([Symbol(:ν_, c) for c=1:N_c]...)
γ = cnumbers([Symbol(:γ_, c) for c=1:N_c]...)
Δ = cnumbers([Symbol(:Δ_, c) for c=1:N_c]...)
g = cnumbers([Symbol(:g_, c) for c=1:N_c]...)

H = sum(Δ[c]*σ(:e,:e,c)[k] for c=1:N_c, k=1:M) + sum(g[c]*(a'*σ(:g,:e,c)[i] + a*σ(:e,:g,c)[i]) for i=1:M, c=1:N_c)
J = [a;[σ(:g,:e,c) for c=1:N_c];[σ(:e,:g,c) for c=1:N_c]]
rates = [κ,γ...,ν...]

ops = [a'*a]
he = meanfield(ops, H, J; rates=rates, order=2)

# Scale
he_scaled = complete(he; filter_func=phase_invariant)

@test isempty(find_missing(he_scaled))

ps = (κ, Δ..., g..., γ..., ν..., N...)
@named sys = ODESystem(he_scaled)
if N_c==2
    p0 = ps .=> (1, [0 for i=1:N_c]..., [1.5 for i=1:N_c]..., [0.25 for i=1:N_c]..., [4 for i=1:N_c]..., 4, 3)
elseif N_c==3
    p0 = ps .=> (1, [0 for i=1:N_c]..., [1.5 for i=1:N_c]..., [0.25 for i=1:N_c]..., [4 for i=1:N_c]..., 2, 3, 2)
end
u0 = zeros(ComplexF64, length(he_scaled))
prob = ODEProblem(sys, u0, (0.0, 50.0), p0)
sol = solve(prob, RK4(), abstol=1e-10, reltol=1e-10)

@test sol.u[end][1] ≈ 12.601868534

### 6-level laser

order = 2
# Define parameters
@cnumbers Δc κ g Γ12 Γ13 Γ24 Γ34 Γ54 Γ16 Γ64 Δ3 Δ4 Δ5 Δ6 Ω13 Ω34 Ω54 Ω64 η ν12 ν13 ν34 ν54 ν64 N

# Define hilbert space
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom,6)
ha_c = ClusterSpace(ha,N,order)
h = ⊗(hf, ha_c)
# Define the fundamental operators
a = Destroy(h,:a,1)
σ(i,j) = Transition(h,:σ,i,j)
@test length(σ(6,6)) == order

# Hamiltonian
H = Δc*a'a + Δ3*sum(σ(3,3)) + Δ4*sum(σ(4,4)) + Δ5*sum(σ(5,5)) + Δ6*sum(σ(6,6)) +
    Ω13*(sum(σ(3,1)) + sum(σ(1,3))) + Ω34*(sum(σ(3,4)) + sum(σ(4,3))) + Ω54*(sum(σ(5,4)) + sum(σ(4,5))) + Ω64*(sum(σ(6,4)) + sum(σ(4,6))) +
    g*(a'*sum(σ(1,2)) + a*sum(σ(2,1)))
# Collapse operators
J = [a, σ(1,2), σ(1,3), σ(2,4), σ(3,4), σ(5,4), σ(6,4), σ(1,6), σ(2,2), σ(3,3)+σ(4,4)+σ(5,5)+σ(6,6), σ(4,4)+σ(5,5)+σ(6,6), σ(5,5), σ(6,6), a'a]
rates = [κ, Γ12, Γ13, Γ24, Γ34, Γ54, Γ64, Γ16, ν12, ν13, ν34, ν54, ν64, η]

# Derive equation for average photon number
ops = [a'a, σ(2,2)[1], σ(3,3)[1], σ(4,4)[1], σ(5,5)[1], σ(6,6)[1]]
eqs_ops = meanfield(ops,H,J;rates=rates, order=order, multithread=true)

@test length(eqs_ops) == length(ops)

### 4th order 2-level laser
order = 4
# Define parameters
@cnumbers δA ΩA wA νA γ κ δc NA
# Define hilbert space
hf = FockSpace(:cavity)
haA = NLevelSpace(:atomA,2)
heA = ClusterSpace(haA, NA, order) #atom ensemble 1
h = ⊗(hf, heA)
a = Destroy(h,:a,1)
σA(i,j) = Transition(h,:σA,i,j,2)
# Hamiltonian
H = δc*a'a + δA*sum(σA(2,2)) + ΩA/2*(a'sum(σA(1,2)) + a*sum(σA(2,1)))
# dissipative processes
J = [a, σA(1,2), σA(2,1), σA(2,2)]
rates = [κ, γ, wA, γ, νA]
# Derive equation for average photon number
ops = [a'a, σA(2,2)[1]]
he_ops = meanfield(ops,H,J;rates=rates, multithread=true, order=order)

he_scale = complete(he_ops; filter_func=phase_invariant, order=order, multithread=true)
@test length(he_scale) == 18
@test isempty(find_missing(he_scale))

### 4th order synchronization

order = 4
# Define parameters
@cnumbers δA δB ΩA ΩB wA wB νA νB γ κ δc NA NB
# Define hilbert space
hf = FockSpace(:cavity)
haA = NLevelSpace(:atomA,2)
haB = NLevelSpace(:atomB,2)
heA = ClusterSpace(haA, NA, order) #atom ensemble 1
heB = ClusterSpace(haB, NB, order) #atom ensemble 2
h = ⊗(hf, heA, heB)
# Define the fundamental operators
a = Destroy(h,:a,1)
σA(i,j) = Transition(h,:σA,i,j,2)
σB(i,j) = Transition(h,:σB,i,j,3)
# Hamiltonian
H = δc*a'a + δA*sum(σA(2,2)) + δB*sum(σB(2,2)) +
    ΩA/2*(a'sum(σA(1,2)) + a*sum(σA(2,1))) + ΩB/2*(a'sum(σB(1,2)) + a*sum(σB(2,1)))
# dissipative processes
J = [a, σA(1,2), σA(2,1), σB(1,2), σB(2,1), σA(2,2), σB(2,2)]
rates = [κ, γ, wA, γ, wB, νA, νB]
# Derive equation for average photon number
ops = [a'a, σA(2,2)[1], σB(2,2)[1]]
he_ops = meanfield(ops,H,J;rates=rates, multithread=true, order=order)

he_scale = complete(he_ops; filter_func=phase_invariant, order=order, multithread=true)
@test length(he_scale) == 66
@test isempty(find_missing(he_scale))

@named sys = ODESystem(he_scale)
u0 = zeros(ComplexF64, length(he_scale))
ps = [δA, δB, ΩA, ΩB, wA, wB, νA, νB, γ, κ, δc, NA, NB]
p0 = ps.=>[1.0 + i/20 for i=1:length(ps)]
prob = ODEProblem(sys,u0,(0.0,1.0),p0)
sol = solve(prob,RK4())

uend = copy(sol.u[end])
@test length(unique(uend)) == 66
uend_f = filter(x->imag(x) != 0, uend)
@test 2*length(uend_f) == length(unique([uend_f; adjoint.(uend_f)]))

end # testset