text:
using SparseArrays, LinearAlgebra

##### Fermion creation & annihilation operators

"""
Syntax: cup,cdn = fermiops(ns). Form spin 1/2 fermi operators of ns sites.
Details:
Define d_j: locally anticommuting fermi operators which commute between
different sites.
To enforce anticommutation between different sites, attach a 'string' by
defining c_k=Pi_{j<k} (-1)^(n_j) d_k.
"""
function fermiops(ns::Int)
    cups = sparse([0 0 0 0; 1 0 0 0; 0 0 0 0; 0 0 1 0]')
    cdns = sparse([0 0 0 0; 0 0 0 0; 1 0 0 0; 0 -1 0 0]')
    N = sparsevec([0, 1, 1, 2])
    signN = sparse(Diagonal([1, -1, -1, 1]))

    cup = SparseMatrixCSC{Float64,Int64}[]
    cdn = SparseMatrixCSC{Float64,Int64}[]
    id = sparse(1I, 4, 4)

    for k in 1:ns
        matup = spzeros(1,1)
        matup .= 1.
        matdn = spzeros(1,1)
        matdn .= 1.
        for j in 1:ns
            if j == k
                matup = kron(matup, cups)
                matdn = kron(matdn, cdns)
            elseif j < k
                matup = kron(matup, signN)
                matdn = kron(matdn, signN)
            elseif j > k
                matup = kron(matup, id)
                matdn = kron(matdn, id)
            end
        end
        push!(cup, matup)
        push!(cdn, matdn)
    end
    return cup, cdn
end

"""
Syntax: c = fermiops(ns). Form fermi operators of ns sites.
Details:
Define d_j: locally anticommuting fermi operators which commute between
different sites.
To enforce anticommutation between different sites, attach a 'string' by
defining c_k=Pi_{j<k} (-1)^(n_j) d_k.
"""
function fermiops_spinless(ns::Int)
    cs = sparse([0 0; 1 0]')
    N = sparsevec([0, 1])
    signN = sparse(Diagonal([1, -1])) # even / uneven # of fermion operators

    c = SparseMatrixCSC{Int64,Int64}[]
    id = sparse(1I, 2, 2)

    for k in 1:ns
        mat = 1
        for j in 1:ns
            if j == k
                mat = kron(mat, cs)
            elseif j < k
                mat = kron(mat, signN)
            elseif j > k
                mat = kron(mat, id)
            end
        end
        push!(c, mat)
    end
    return c
end

# Note: the resulting cs are not really fermion operators
# i.e. they do not fullfill the usual anticommutation relations!
function fermiops_spinless_spbasis(ns::Int)
    n = ns + 1 # the one is the all zeros state
    c = Array{SparseMatrixCSC{Int64,Int64}, 1}(undef, ns)

    for k in 1:ns
        c[k] = spzeros(Int, n,n)
        c[k][1,k+1] = 1
    end

    return c
end

anticomm(a,b) = Matrix(a*b + b*a)
comm(a,b) = Matrix(a*b - b*a)

function check_fermiops(c::AbstractVector{T}) where T<:AbstractArray
    n = size(c[1], 1)
    id = sparse(1I, n, n)
    for (j, c1) in enumerate(c)
        for (k, c2) in enumerate(c)
             # cc anticommutator should vanish
            cc = c1 * c2 + c2 * c1
            @assert sum(abs.(cc)) == 0

            # c^dagger c  anticommutator
            # should be identity if j==k and zero if j!=k
            cdc = c1' * c2 + c2 * c1'
            if j == k
                @assert sum(abs.(cdc .- id)) == 0
            else
                @assert sum(abs.(cdc)) == 0
            end
        end
    end
    return true
end

check_fermiops(c::AbstractArray{T}) where T = check_fermiops([c])

#### Hamiltonians

function generate_H_tb_spinless(t, mu, ns, c)
    L = Int(sqrt(ns))
    sqlattice = reshape(collect(1:ns), (L,L))
    # up and right neighbors
    unn = circshift(sqlattice, (1,0))[:]
    rnn = circshift(sqlattice, (0,-1))[:]
    dim = Int(2^ns)

    # chemical potential

    # initialize N:
    N = spzeros(Float64, dim, dim)

    for i in 1:ns
        N = N + c[i]' * c[i]
    end

    H = - mu * N

    @assert ishermitian(H)

    # hoppings

    for i in 1:ns
        H += - t * (c[i]' * c[unn[i]]) - t * (c[unn[i]]' * c[i])
        H += - t * (c[i]' * c[rnn[i]]) - t * (c[rnn[i]]' * c[i])
    end

    @assert ishermitian(H)
    return H
end

function generate_H_Hubbard(t, U, mu, ns, cup, cdn)
    L = Int(sqrt(ns))
    sqlattice = reshape(collect(1:ns), (L,L))
    # up and right neighbors
    unn = circshift(sqlattice, (1,0))[:]
    rnn = circshift(sqlattice, (0,-1))[:]
    dim = Int(4^ns)

    # chemical potential
    N = spzeros(Float64, dim, dim)
    for i in 1:ns
        N = N + cup[i]' * cup[i]
        N = N + cdn[i]' * cdn[i]
    end

    H = - mu * N

    @assert ishermitian(H)

    # hoppings
    for i in 1:ns
        for c in (cup, cdn)
            H += - t * (c[i]' * c[unn[i]]) - t * (c[unn[i]]' * c[i])
            H += - t * (c[i]' * c[rnn[i]]) - t * (c[rnn[i]]' * c[i])
        end
    end

    # on-site interaction
    for i in 1:ns
        nup = cup[i]' * cup[i]
        ndn = cdn[i]' * cdn[i]
        H = H + U * (nup - 0.5I) * (ndn - 0.5I)
    end

    @assert ishermitian(H)
    return H
end

# normal = mu=0 doesn't imply half filling
function generate_H_Hubbard_normal(t, U, mu, ns, cup, cdn)
    L = Int(sqrt(ns))
    sqlattice = reshape(collect(1:ns), (L,L))
    # up and right neighbors
    unn = circshift(sqlattice, (1,0))[:]
    rnn = circshift(sqlattice, (0,-1))[:]
    dim = Int(4^ns)

    # chemical potential
    N = spzeros(Float64, dim, dim)
    for i in 1:ns
        N = N + cup[i]' * cup[i]
        N = N + cdn[i]' * cdn[i]
    end

    H = - mu * N

    @assert ishermitian(H)

    # hoppings
    for i in 1:ns
        for c in (cup, cdn)
            H += - t * (c[i]' * c[unn[i]]) - t * (c[unn[i]]' * c[i])
            H += - t * (c[i]' * c[rnn[i]]) - t * (c[rnn[i]]' * c[i])
        end
    end

    # on-site interaction
    for i in 1:ns
        nup = cup[i]' * cup[i]
        ndn = cdn[i]' * cdn[i]
        H = H + U * (nup * ndn)
    end

    @assert ishermitian(H)
    return H
end

function generate_H_tb_spinless_spbasis(t, mu, ns, c)
    L = Int(sqrt(ns))
    sqlattice = reshape(collect(1:ns), (L,L))
    # up and right neighbors
    unn = circshift(sqlattice, (1,0))[:]
    rnn = circshift(sqlattice, (0,-1))[:]
    dim = Int(ns+1)

    # chemical potential

    # initialize N:
    N = spzeros(Float64, dim, dim)

    for i in 1:ns
        N = N + c[i]' * c[i]
    end

    H = - mu * N

    @assert ishermitian(H)

    # hoppings

    for i in 1:ns
        H += - t * (c[i]' * c[unn[i]]) - t * (c[unn[i]]' * c[i])
        H += - t * (c[i]' * c[rnn[i]]) - t * (c[rnn[i]]' * c[i])
    end

    @assert ishermitian(H)
    return H
end

function peirls_phase(i,j,L,B)
    @assert i in 1:L^2
    @assert j in 1:L^2

    sqlattice = reshape(collect(1:L^2), (L,L))
    i2, i1 = ind2sub(sqlattice, i)
    j2, j1 = ind2sub(sqlattice, j)

    if (i1 in 1:L-1 && j1 == i1+1) || (i1 == L && j1 == 1)
        return - 2*pi * B * (i2 - 1)

    elseif (i1 in 2:L && j1 == i1-1) || (i1 == 1 && j1 == L)
        return 2*pi * B * (i2 - 1)

    elseif i2 == L && j2 == 1
        return 2*pi * B * L * (i1 - 1)

    elseif i2 == 1 && j2 == L
        return - 2*pi * B * L * (i1 - 1)

    else
        return 0.
    end
end

function generate_H_tb_spinless_bfield_spbasis(t, mu, B, ns, c)
    L = Int(sqrt(ns))
    sqlattice = reshape(collect(1:ns), (L,L))
    # up and right neighbors
    unn = circshift(sqlattice, (1,0))[:]
    rnn = circshift(sqlattice, (0,-1))[:]
    dim = Int(ns+1)

    # chemical potential

    # initialize N:
    N = spzeros(Float64, dim, dim)

    for i in 1:ns
        N = N + c[i]' * c[i]
    end

    H = - mu * N

    @assert ishermitian(H)

    # hoppings

    for i in 1:ns
        Au = exp(im * peirls_phase(i, unn[i], L, B))
        Aut = exp(im * peirls_phase(unn[i], i, L, B))
        Ar = exp(im * peirls_phase(i, rnn[i], L, B))
        Art = exp(im * peirls_phase(rnn[i], i, L, B))

        H += - t * Au * (c[i]' * c[unn[i]]) - t * Aut * (c[unn[i]]' * c[i])
        H += - t * Ar * (c[i]' * c[rnn[i]]) - t * Art * (c[rnn[i]]' * c[i])
    end

    @assert ishermitian(H)
    return H
end

function spinspin(u, j, k, cup, cdn)
    """
    Spin-spin and density-density interaction. Returns
    sum_i u(i) (cdag(j) sigma(i) c(k)+ H.C)^2.
    Sigma are Pauli matrices (NOT spin matrices, i.e no factor of 1/2).
    u is a list of coefficients for (0, x, y, z) components of spin, where 0 is
    just the identity. j, k are site indices.
    """
    v = (cup[j]' * cup[k] + cdn[j]' * cdn[k] +
         cup[k]' * cup[j] + cdn[k]' * cdn[j])
    V = u[1] * v * v

    v = (cup[j]' * cdn[k] + cdn[j]' * cup[k] +
         cup[k]' * cdn[j] + cdn[k]' * cup[j])
    V = V + u[2] * v * v

    v = (cup[j]' * cdn[k] - cdn[j]' * cup[k] +
         cup[k]' * cdn[j] - cdn[k]' * cup[j])
    V = V - u[3] * v * v

    v = (cup[j]' * cup[k] - cdn[j]' * cdn[k] +
         cup[k]' * cup[j] - cdn[k]' * cdn[j])
    V = V + u[4] * v * v

    return V
end

function thop(t, j, k, cup, cdn)
    T = t * (cup[j]' * cup[k]) + t * (cup[k]' * cup[j])
    T = T + t * (cdn[j]' * cdn[k]) + t * (cdn[k]' * cdn[j])
    return T
end

function generate_H_SDW(params, ns, cup, cdn)
    txh = params["txh"]
    txv = params["txv"]
    tyh = params["tyh"]
    tyv = params["tyv"]
    mu = params["mu"]
    r = params["r"]
    lambdax = params["lambdax"]
    lambday = params["lambday"]
    lambdaz = params["lambdaz"]
    lambda0 = params["lambda0"]

    lambdas = Float64[lambda0, lambdax, lambday, lambdaz]
    u = Float64[-lamb^2 / (2 * r) for lamb in lambdas]

    if ns != 8
        throw("Number of sites != 8")
    end

    # chemical potential

    # initialize N:
    dim = Int(4^ns)
    N1 = spzeros(Float64, dim, dim)
    N2 = spzeros(Float64, dim, dim)

    for i in 1:4
        N1 = N1 + cup[i]' * cup[i]
        N1 = N1 + cdn[i]' * cdn[i]
    end
    for i in 5:8
        N2 = N2 + cup[i]' * cup[i]
        N2 = N2 + cdn[i]' * cdn[i]
    end

    H = - mu * N1 - mu * N2

    # Hopping part:
    H = H - thop(2 * txv, 1, 2, cup, cdn)
    H = H - thop(2 * txv, 3, 4, cup, cdn)
    H = H - thop(2 * txh, 1, 3, cup, cdn)
    H = H - thop(2 * txh, 2, 4, cup, cdn)
    H = H - thop(2 * tyv, 5, 6, cup, cdn)
    H = H - thop(2 * tyv, 7, 8, cup, cdn)
    H = H - thop(2 * tyh, 5, 7, cup, cdn)
    H = H - thop(2 * tyh, 6, 8, cup, cdn)

    # Interaction
    for j in 1:4
        H = H + spinspin(u, j, j + 4, cup, cdn)
    end

    @assert ishermitian(H)
    return H
end

using BlockArrays
function generate_H_SDW_spbasis(params, ns, c)
    txh = params["txh"]
    txv = params["txv"]
    tyh = params["tyh"]
    tyv = params["tyv"]
    mu = params["mu"]

    p = copy(params)
    p["th"] = txh
    p["tv"] = txv
    p["mu"] = mu
    H1 = generate_H_SDW_spbasis_single_flavor(p, ns, c)

    p["th"] = tyh
    p["tv"] = tyv
    p["mu"] = mu
    H2 = generate_H_SDW_spbasis_single_flavor(p, ns, c)

    L = Int(sqrt(ns))
    N = L^2

    H = BlockArray(zeros(4*N, 4*N), [N,N,N,N], [N,N,N,N]) # xup, yup, xdown, ydown
    setblock!(H, H1, 1, 1)
    setblock!(H, H2, 2, 2)
    setblock!(H, H1, 3, 3)
    setblock!(H, H2, 4, 4)
    setblock!(H, zeros(N,N), 1, 2)
    setblock!(H, zeros(N,N), 2, 1)
    Matrix(H)
end

function generate_H_SDW_spbasis_single_flavor(params, ns, c)
    th = params["th"]
    tv = params["tv"]
    mu = params["mu"]

    L = Int(sqrt(ns))
    sql = reshape(1:ns, (L,L))
    # up and right neighbors
    unn = circshift(sql, (-1,0))[:]
    rnn = circshift(sql, (0,-1))[:]
    dim = Int(ns+1)

    # initialize N:
    N = spzeros(Float64, dim, dim)

    for i in 1:ns
        N = N + c[i]' * c[i]
    end

    H = - mu * N

    @assert ishermitian(H)

    # hoppings

    for i in 1:ns
        H += - tv * (c[i]' * c[unn[i]]) - tv * (c[unn[i]]' * c[i])
        H += - th * (c[i]' * c[rnn[i]]) - th * (c[rnn[i]]' * c[i])
    end

    @assert ishermitian(H)
    return Matrix(H[2:end, 2:end]) # cut off all zeros (no particles) state
end

function ϵ(params, ky, kx)
    th = params["th"]
    tv = params["tv"]
    mu = params["mu"]

    gamma = -2*th*cos(kx) - 2*tv*cos(ky)
    ϵ = gamma - mu
end

generate_H_SDW_k_space_single_flavor(params, ky, kx) = ϵ(params, ky, kx) # H is diagonal in k-space

# This probably doesn't work that way.
# We probably have to adjust the phases because e.g 1->2 once exists on the lattice and once via PBC.
function peirls_SDW(i::Int , j::Int, B::Float64, sql::Matrix{Int}, pbc::Bool)::ComplexF64
    # peirls_phase_factors e^{im*Aij}

    i2, i1 = ind2sub(sql, i)
    j2, j1 = ind2sub(sql, j)
    L = size(sql, 1)

    if pbc
        if (i1 == L && j1 == 1)
            return exp(im*(- 2*pi * B * (i2 - 1)))

        elseif (i1 == 1 && j1 == L)
            return exp(im*(2*pi * B * (i2 - 1)))

        elseif i2 == L && j2 == 1
            return exp(im*(2*pi * B * L * (i1 - 1)))

        elseif i2 == 1 && j2 == L
            return exp(im*(- 2*pi * B * L * (i1 - 1)))

        else
            return exp(im*0.)
        end
    else
        if (i1 in 1:L-1 && j1 == i1+1)
            return exp(im*(- 2*pi * B * (i2 - 1)))

        elseif (i1 in 2:L && j1 == i1-1)
            return exp(im*(2*pi * B * (i2 - 1)))

        else
            return exp(im*0.)
        end
    end
end

function generate_H_SDW_Bfield(params, ns, cup, cdn)
    B = zeros(2,2) # colidx = flavor, rowidx = spin up,down
    B[1,1] = B[2,2] = 1 / (ns/2) # ns == number of sites *2 because of 2 flavors
    B[1,2] = B[2,1] = - 1 / (ns/2)

    txh = params["txh"]
    txv = params["txv"]
    tyh = params["tyh"]
    tyv = params["tyv"]
    mu = params["mu"]
    r = params["r"]
    lambdax = params["lambdax"]
    lambday = params["lambday"]
    lambdaz = params["lambdaz"]
    lambda0 = params["lambda0"]

    lambdas = Float64[lambda0, lambdax, lambday, lambdaz]
    u = Float64[-lamb^2 / (2 * r) for lamb in lambdas]

    if ns != 8
        throw("Number of sites != 8")
    end

    # chemical potential

    # initialize N:
    dim = Int(4^ns)
    N1 = spzeros(Float64, dim, dim)
    N2 = spzeros(Float64, dim, dim)

    for i in 1:4
        N1 = N1 + cup[i]' * cup[i]
        N1 = N1 + cdn[i]' * cdn[i]
    end
    for i in 5:8
        N2 = N2 + cup[i]' * cup[i]
        N2 = N2 + cdn[i]' * cdn[i]
    end

    H = - mu * N1 - mu * N2

    # for linidx to cartesianidx
    sql = reshape(collect(1:Int(ns/2)), (2,2))

    cs = vcat(cup, cdn);
    t = [txh tyh; txv tyv]

    for f in 1:2 # flavor
        for s in 1:2 # spin
            c = cs[(s-1)*8+(f-1)*4+1:(s-1)*8+(f-1)*4+4]
            A =
            # hopping within the lattice
            H = H - peirls_SDW(1, 2, B[s,f], sql, false) * t[2,f] * (c[1]' * c[2]) - peirls_SDW(2, 1, B[s,f], sql, false) * t[2,f] * (c[2]' * c[1])
            H = H - peirls_SDW(3, 4, B[s,f], sql, false) * t[2,f] * (c[3]' * c[4]) - peirls_SDW(4, 3, B[s,f], sql, false) * t[2,f] * (c[4]' * c[3])
            H = H - peirls_SDW(1, 3, B[s,f], sql, false) * t[1,f] * (c[1]' * c[3]) - peirls_SDW(3, 1, B[s,f], sql, false) * t[1,f] * (c[3]' * c[1])
            H = H - peirls_SDW(2, 4, B[s,f], sql, false) * t[1,f] * (c[2]' * c[4]) - peirls_SDW(4, 2, B[s,f], sql, false) * t[1,f] * (c[4]' * c[2])

            # hopping using PBC
            H = H - peirls_SDW(1, 2, B[s,f], sql, true) * t[2,f] * (c[1]' * c[2]) - peirls_SDW(2, 1, B[s,f], sql, true) * t[2,f] * (c[2]' * c[1])
            H = H - peirls_SDW(3, 4, B[s,f], sql, true) * t[2,f] * (c[3]' * c[4]) - peirls_SDW(4, 3, B[s,f], sql, true) * t[2,f] * (c[4]' * c[3])
            H = H - peirls_SDW(1, 3, B[s,f], sql, true) * t[1,f] * (c[1]' * c[3]) - peirls_SDW(3, 1, B[s,f], sql, true) * t[1,f] * (c[3]' * c[1])
            H = H - peirls_SDW(2, 4, B[s,f], sql, true) * t[1,f] * (c[2]' * c[4]) - peirls_SDW(4, 2, B[s,f], sql, true) * t[1,f] * (c[4]' * c[2])
        end
    end

    # Interaction
    for j in 1:4
        H = H + spinspin(u, j, j + 4, cup, cdn)
    end

    @assert ishermitian(H)
    return H
end

#### Quantities

# Density of states
function DOS(evals; δ=0.01, ωmin=minimum(evals), ωmax=maximum(evals), N=1000)
    ω = LinRange(ωmin, ωmax, N)
    kernel = 1 ./ (ω .- evals' .+ im*δ);
    ν = -1 ./ pi * imag(vec(sum(kernel, 2)))
    return ν, ω
end

function EV(O::SparseMatrixCSC, evecs, evals, beta::Float64)
    V = sparse(evecs)
    O_trans = V' * (O * V)
    E = evals .- evals[1]
    Z = sum(exp.(- beta * E))
    return real(1.0 / Z * sum(exp.(- beta * E) .* diag(O_trans)))
end

function GF(cs, evecs, evals, beta)
    # single-particle Green's function
    g = zeros(ComplexF64, length(cs), length(cs))
    for (j1, c1) in enumerate(cs)
        for (j2, c2) in enumerate(cs)
            # println(j1, " ", j2)
            g[j1,j2] = EV(c1 * c2', evecs, evals, beta)
        end
    end
    return g
end

"""
Calculate GF directly from single-particle basis eigenvalue decomposition
via (stabilized) version of G_ij = [I / (I + e^(-βH)]_ij.

The result matches the output of function GF above when calculated from full basis.
"""
function GF_free_spbasis(evals, evecs, beta)
    U = evecs
    D = exp.(-beta .* evals)
    Vd = copy(U')

    inv_one_plus_udv_scalettar(U,D,Vd)
end

GF_element_k_space(ϵ, beta) = 1 / (1 + exp(-beta * ϵ))

function GF_k_space(params, L, beta)
    # momentum grid
    kxs = kys = LinRange(-pi, pi, L+1)[1:end-1]
    nk = length(kxs)

    gk = zeros(Float64, nk, nk)

    for j in 1:nk
        for i in 1:nk
            E = ϵ(params, kys[i], kxs[j])
            gk[i,j] = GF_element_k_space(E, beta)
        end
    end
    gk
end

TDGF_element_k_space(ϵ, beta, tau) = exp(-tau * ϵ) / (1 + exp(-beta * ϵ))
TDGF_element_k_space_G0t(ϵ, beta, tau) = -exp(+tau * ϵ) / (1 + exp(beta * ϵ))

function TDGF_k_space(params, L, beta, tau; G0t=false)
    kxs = kys = LinRange(-pi, pi, L+1)[1:end-1]
    nk = length(kxs)

    gk = zeros(Float64, nk, nk)

    for j in 1:nk
        for i in 1:nk
            E = ϵ(params, kys[i], kxs[j])
            if !G0t
                gk[i,j] = TDGF_element_k_space(E, beta, tau)
            else
                gk[i,j] = TDGF_element_k_space_G0t(E, beta, tau)
            end
        end
    end
    gk
end

"""
Calculates G(τ, 0)
"""
function TDGF(cs, tau, evecs, evals, beta)
    tdgf = zeros(length(cs), length(cs))
    for (j1, c1) in enumerate(cs)
        for (j2, c2) in enumerate(cs)
            tdgf[j1,j2] = corr(c1, c2', tau, evecs, evals, beta)
        end
    end
    return tdgf
end

"""
Calculates G(τ,0) as G_ij = [e^(-τH) * (I / (I + e^(-βH))]_ij (the multiplication with e^(-τH) is not stabilized)
"""
function TDGF_free_spbasis_naive(evals, evecs, beta, tau)
    @warn "This calculation isn't stabilized and might lead to wrong results! Use `TDGF_free_spbasis` instead."
    nom = evecs * (Diagonal(exp.(-tau .* evals)) * evecs')
    GF_free_spbasis(evals, evecs, beta) * nom
end

"""
Calculates G(τ,0) as G_ij = [e^(-τH) / (I + e^(-βH)]_ij = [I / {e^(τH) + e^((τ-β)H)}]_ij

Stabilized through inv_sum_udvs.
"""
function TDGF_free_spbasis(evals, evecs, beta, tau)
    Ua = evecs
    Da = exp.(tau .* evals)
    Vda = copy(Ua')

    Ub = evecs
    Db = exp.((tau-beta) .* evals)
    Vdb = copy(Ub')

    U, D, Vd = inv_sum_udvs(Ua,Da,Vda, Ub,Db,Vdb)
    U * Diagonal(D) * Vd
end

# using PyCall
# @pyimport numpy as np
"""
    compute correlation function of operators
    chi = < A(tau) B(0) > = 1/Z sum_n_m exp(-(beta -tau) En) Anm exp(-tau Em) Bmn
"""
function corr(A, B, tau, evecs, evals, beta)
    @assert 0. <= tau <= beta
    # tau <= 1e-5 && warn("for short times corr is way off")

    chi = zero(eltype(A))
    V = sparse(evecs)

    A_trans = V' * A * V
    B_trans = V' * B * V

    E = evals .- evals[1]

    Z = sum(exp.(- beta * E))

    u1 = exp.(-(beta - tau) * E)
    u2 = exp.(- tau * E)

    # chi = np.einsum("n,nm,m,mn", u1, A_trans, u2, B_trans)/Z
    chi = einsum_explicit(u1, A_trans, u2, B_trans)/Z
    return chi
end

# explicit and faster version of np.einsum("n,nm,m,mn", u1, A_trans, u2, B_trans)
function einsum_explicit(u1,A_trans,u2,B_trans)
    s = 0.
    @inbounds for n in 1:length(u1)
        for m in 1:length(u2)
            s += u1[n] * A_trans[n,m] * u2[m] * B_trans[m,n]
        end
    end
    s
end

"""
    compute correlation function of operators
    chi = < A(tau) B(tau) C(0) D(0) > = 1/Z sum_n_m exp(-beta En) exp(tau En) Anm exp(-tau Em) exp(tau Em) Bmn exp(tau En)
"""
function corr_fourpoint(A, B, C, D, tau, evecs, evals, beta)
    @assert 0. <= tau <= beta
    # tau <= 1e-5 && warn("for short times corr is way off")

    chi = zero(eltype(A))
    V = sparse(evecs)

    A_trans = V' * A * V
    B_trans = V' * B * V
    C_trans = V' * C * V
    D_trans = V' * D * V

    E = evals .- evals[1]

    Z = sum(exp.(- beta * E))

    u1 = exp.(-(beta - tau) * E)
    u2 = exp.(- tau * E)

    # chi = np.einsum("n,nm,mk,k,kp,pn", u1, A_trans, B_trans, u2, C_trans, D_trans)/Z
    chi = einsum_explicit(u1,A_trans,B_trans,u2,C_trans,D_trans)/Z
    return chi
end

# explicit and faster version of
# np.einsum("n,nm,mk,k,kp,pn", u1, A_trans, B_trans, u2, C_trans, D_trans)
function einsum_explicit(u1,A_trans,B_trans,u2,C_trans,D_trans)
    s = 0.
    @inbounds for n in 1:length(u1)
        for m in 1:size(A_trans, 2)
            for k in 1:length(u2)
                for p in 1:size(C_trans, 2)
                    s += u1[n] * A_trans[n,m] * B_trans[m,k] * u2[k] * C_trans[k,p] * D_trans[p,n]
                end
            end
        end
    end
    s
end

function QMC.siteidx(L,N,sql,x,y)
  xpbc = mod1(x, L)
  ypbc = mod1(y, L)
  sql[ypbc, xpbc] # to linear idx
end

(@isdefined UP) || (const UP, DOWN = 0,1)
(@isdefined X) || (const X, Y = 0,1)
getop(cs, i, band, spin) = cs[spin*8 + band*4 + i]

"""
Calculate ETPC from real-space ED -> etpc_minus, etpc_plus
"""
function etpc(cs, evecs, evals, beta)
    L = 2
    N = 4
    sql = reshape(1:N,L,L)

    pc_minus = zeros(L,L)
    pc_plus = zeros(L,L)

    for x in 0:(L-1), y in 0:(L-1) # r (displacement)
        for x0 in 1:L, y0 in 1:L # r0 (origin)
            ri = siteidx(L,N,sql, x0+x, y0+y)
            r0 = siteidx(L,N,sql, x0, y0)

            op1 = getop(cs, ri, X, DOWN) * getop(cs, ri, X, UP) * getop(cs, r0, X, UP)' * getop(cs, r0, X, DOWN)'
            op2 = getop(cs, ri, X, DOWN) * getop(cs, ri, X, UP) * getop(cs, r0, Y, UP)' * getop(cs, r0, Y, DOWN)'
            op3 = getop(cs, ri, Y, DOWN) * getop(cs, ri, Y, UP) * getop(cs, r0, X, UP)' * getop(cs, r0, X, DOWN)'
            op4 = getop(cs, ri, Y, DOWN) * getop(cs, ri, Y, UP) * getop(cs, r0, Y, UP)' * getop(cs, r0, Y, DOWN)'

            pc_minus[y+1, x+1] += EV(op1, evecs, evals, beta)
            pc_minus[y+1, x+1] -= EV(op2, evecs, evals, beta)
            pc_minus[y+1, x+1] -= EV(op3, evecs, evals, beta)
            pc_minus[y+1, x+1] += EV(op4, evecs, evals, beta)

            pc_plus[y+1, x+1] += EV(op1, evecs, evals, beta)
            pc_plus[y+1, x+1] += EV(op2, evecs, evals, beta)
            pc_plus[y+1, x+1] += EV(op3, evecs, evals, beta)
            pc_plus[y+1, x+1] += EV(op4, evecs, evals, beta)
        end
    end

    pc_minus ./=N
    pc_plus ./=N

    return pc_minus, pc_plus
end

"""
Calculate ETCDC from real-space ED -> etcdc_minus, etcdc_plus
"""
function etcdc(cs, evecs, evals, beta)
    L = 2
    N = 4
    sql = reshape(1:N,L,L)

    cdc_minus = zeros(L,L)
    cdc_plus = zeros(L,L)

    for x in 0:(L-1), y in 0:(L-1) # r (displacement)
        for s1 in (UP, DOWN), s2 in (UP, DOWN)
            for x0 in 1:L, y0 in 1:L # r0 (origin)
                ri = siteidx(L,N,sql, x0+x, y0+y)
                r0 = siteidx(L,N,sql, x0, y0)

                op1 = getop(cs, ri, X, s1)' * getop(cs, ri, X, s1) * getop(cs, r0, X, s2)' * getop(cs, r0, X, s2)
                op2 = getop(cs, ri, X, s1)' * getop(cs, ri, X, s1) * getop(cs, r0, Y, s2)' * getop(cs, r0, Y, s2)
                op3 = getop(cs, ri, Y, s1)' * getop(cs, ri, Y, s1) * getop(cs, r0, X, s2)' * getop(cs, r0, X, s2)
                op4 = getop(cs, ri, Y, s1)' * getop(cs, ri, Y, s1) * getop(cs, r0, Y, s2)' * getop(cs, r0, Y, s2)

                cdc_minus[y+1, x+1] += EV(op1, evecs, evals, beta)
                cdc_minus[y+1, x+1] -= EV(op2, evecs, evals, beta)
                cdc_minus[y+1, x+1] -= EV(op3, evecs, evals, beta)
                cdc_minus[y+1, x+1] += EV(op4, evecs, evals, beta)

                cdc_plus[y+1, x+1] += EV(op1, evecs, evals, beta)
                cdc_plus[y+1, x+1] += EV(op2, evecs, evals, beta)
                cdc_plus[y+1, x+1] += EV(op3, evecs, evals, beta)
                cdc_plus[y+1, x+1] += EV(op4, evecs, evals, beta)
            end
        end
    end

    cdc_minus ./=N
    cdc_plus ./=N

    return cdc_minus, cdc_plus
end

# ARCHIV

# Guess this is wrong because of o = c1_trans * c2dagger_trans
function TDGFelement(c1, c2, tau, evecs, evals, beta)
    E = evals - evals[1]
    V = sparse(evecs)
    c1_trans = V' * (c1 * V)
    c2dagger_trans = V' * (c2' * V)
    o = c1_trans * c2dagger_trans

    le = length(E)
    g = 0.0
    @inbounds for n in 1:le
        for m in 1:le
            g += exp(-beta*E[n] + tau*(E[n] - E[m])) * o[n,m]
        end
    end
    Z = sum(exp.(- beta * E))
    return g/Z
end
nothing

max_stars_repo_name: crstnbr/dqmc

max_issues_repo_name: crstnbr/dqmc

max_forks_repo_name: crstnbr/dqmc


text:
"""
```
measurement{T<:AbstractFloat}(m::RealBond{T}, TTT::Matrix{T}, RRR::Matrix{T},
                              CCC::Vector{T})
```

Assign measurement equation

```
y_t = ZZ*s_t + DD + u_t
```

where

```
Var(ϵ_t) = QQ
Var(u_t) = EE
Cov(ϵ_t, u_t) = 0
```
"""
function measurement(m::RealBond{T}, TTT::Matrix{T},
                     TTT_jump::Matrix{T},
                     RRR::Matrix{T}, CCC::Vector{T}) where T<:AbstractFloat
    endo      = m.endogenous_states_unnormalized
    exo       = m.exogenous_shocks
    obs       = m.observables

    _n_model_states = n_model_states(m)
    _n_states       = n_backward_looking_states(m)
    _n_jumps        = n_jumps(m)

    _n_observables = n_observables(m)
    _n_shocks_exogenous = n_shocks_exogenous(m)

    ZZ = zeros(_n_observables, _n_model_states)
    DD = zeros(_n_observables)
    EE = zeros(_n_observables, _n_observables)
    QQ = zeros(_n_shocks_exogenous, _n_shocks_exogenous)

    # Load in parameters, steady-state parameters, and grids
    γ::Float64     = m[:γ].value
    ν::Float64     = m[:ν].value
    R::Float64     = m[:R].value
    abar::Float64  = m[:abar].value
    aborrow        = abar/R

    ell::Vector{Float64}  = m[:lstar].value
    c::Vector{Float64}    = m[:cstar].value
    μ::Vector{Float64}    = m[:μstar].value
    η::Vector{Float64}    = m[:ηstar].value
    χss::Vector{Float64}  = m[:χstar].value

    xwts::Vector{Float64}  = m.grids[:xgrid].weights
    sgrid::Vector{Float64} = m.grids[:sgrid].points
    swts::Vector{Float64}  = m.grids[:sgrid].weights
    xgrid_total::Vector{Float64}  = m.grids[:xgrid_total]

    nx::Int = get_setting(m, :nx)
    ns::Int = get_setting(m, :ns)

    # This is just a snippet of code that constructs linearized GDP
    # It should be incorporated into a larger file that simulates, filters etc.
    # This is based on real_bond_hank.jl

    dGDP_dMU, GDP, dGDP_dELL, dGDP_dRR, dGDP_dWW, dGDP_dTT = construct_GDPfn_realbond(nx, ns, xgrid_total, sgrid, xwts, swts,
                                                                                γ, ν, abar, R, aborrow, μ,
                                                                                c, η, ell, χss)

    GDPfn = zeros(1, 2*nx*ns+7) # GDP as function of un-normalized MU Z MON ELL RR II WW PI TT
    # note: GDP is only a function of contemporaneous variables
    # we are using the indices (MUP,ZP, etc.) corresponding to date t+1 variables
    # simply because these indices happen to work here also
    # this does not mean that GDP is a function of date t+1 variables
    GDPfn[1, endo[:μ′_t]]  = dGDP_dMU
    GDPfn[1, endo[:z′_t]]  .= GDP
    GDPfn[1, endo[:l′_t]]  = dGDP_dELL
    GDPfn[1, endo[:R′_t]]  .= dGDP_dRR
    GDPfn[1, endo[:w′_t]]  .= dGDP_dWW
    GDPfn[1, endo[:t′_t]]  .= dGDP_dTT

    ########################################
    Qx, Qy, _, _ = compose_normalization_matrices(m)
    gx2  = Qy'*TTT_jump*Qx

    # now we need to create GDP as a function of the normalized states
    ZZ_states = (1/GDP)*GDPfn*[eye(nx*ns+2); gx2]*Qx' # this is for log GDP
                                                      # to use the level of gdp, remove (1/GDP)

    ZZ = Matrix{Float64}(undef, _n_observables, _n_model_states)
    ZZ[1:_n_states] = ZZ_states
    ZZ[_n_states+1:end] = zeros(_n_jumps)

    # Measurement error
    EE[obs[:obs_gdp], obs[:obs_gdp]] = m[:e_y]

    # Variance of innovations
    QQ[exo[:z_sh], exo[:z_sh]] = m[:σ_z]^2
    QQ[exo[:mon_sh], exo[:mon_sh]] = m[:σ_mon]^2

    # Adjustment to DD because measurement equation assumes CCC is the zero vector
    if any(CCC .!= 0)
        DD += ZZ*((UniformScaling(1) - TTT)\CCC)
    end

    return Measurement(ZZ, DD, QQ, EE)
end

function construct_GDPfn_realbond(nx::Int, ns::Int,
                                  xgrid_total::Vector{Float64}, sgrid::Vector{Float64},
                                  xwts::Vector{Float64}, swts::Vector{Float64},
                                  γ::Float64, ν::Float64, abar::Float64,
                                  R::Float64, aborrow::Float64,
                                  μ::Vector{Float64}, c::Vector{Float64}, η::Vector{Float64},
                                  ell::Vector{Float64}, χss::Vector{Float64})
    dGDP_dMU  = zeros(1,nx*ns)
    dGDP_dELL = zeros(1,nx*ns)
    dGDP_dWW  = 0.0
    dGDP_dRR  = 0.0
    dGDP_dTT  = 0.0
    GDP       = 0.0

    unc       = zeros(nx*ns)
    chipW, chipR, chipX = construct_chip_realbond(xgrid_total, γ, ν, aborrow,
                                                  abar, R, χss)
    for ix =1:nx
        for is=1:ns
            i = ix + nx*(is-1)
            unc[i] = ((ell[i]^(-1/γ))<=χss[i]) # =1 if unconstrained
            dGDP_dMU[i]  = xwts[ix]*swts[is]*sgrid[is]*η[i]
            dGDP_dELL[i] = xwts[ix]*swts[is]*μ[ix]*(sgrid[is]*γ*η[i]/(ν*c[i]))*(1/γ)*unc[i]*ell[i]^(-(1/γ)-1)
            dGDP_dWW    += swts[is]*xwts[ix]*μ[i]*((sgrid[is]*η[i])/ν - (sgrid[is]*γ*η[i]/(ν*c[i]))*(1-unc[i])*chipW[i])
            dGDP_dRR    += -swts[is]*xwts[ix]*μ[i]*(sgrid[is]*γ*η[i]/(ν*c[i]))*(1-unc[i])*chipR[i]
            dGDP_dTT    += -swts[is]*xwts[ix]*μ[i]*(sgrid[is]*γ*η[i]/(ν*c[i]))*(1-unc[i])*chipX[i]
            GDP         += swts[is]*xwts[ix]*sgrid[is]*η[i]*μ[i] # note, this is also dGDP_dZ. It is already created by real_bond_hank.jl
        end
    end
    return dGDP_dMU, GDP, dGDP_dELL, dGDP_dRR, dGDP_dWW, dGDP_dTT
end

ext: jl


text:
module TestSmoke

using Aqua
Aqua.test_all(
    Aqua;
    stale_deps = (; ignore = [:Compat]),  # conditionally loaded
)
Aqua.test_all(
    Aqua;
    ambiguities = false,
    unbound_args = false,
    undefined_exports = false,
    stale_deps = (; ignore = [:Compat]),
    deps_compat = true,
    project_toml_formatting = true,
)

end  # module


text:
# Helper functions

function savefig(mission::Mission, obs_row::DataFrames.DataFrameRow{DataFrames.DataFrame},
        e_range::Tuple{Float64,Float64}, kind::Symbol, plot_name::String;
        kwargs...
    )
    kwargs = Dict(kwargs)

    obs_log = _log_entry(; category=:images, e_range=e_range, kind=kind, file_name=string(plot_name, ".png"), kwargs...)

    plot_path = abspath(_obs_path_local(mission, obs_row; kind=:web), "JAXTAM", _log_query_path(; category=:images, kind=kind, e_range=e_range, file_name=string(plot_name, ".png"), kwargs...))

    mkpath(dirname(plot_path)); Plots.savefig(plot_path)
    @info "Saved $plot_path"

    _log_add(mission, obs_row, Dict("images"=>Dict(e_range=>obs_log)))

    return plot_path
end

function savefig(plot_data::JAXTAMData, obs_row::DataFrames.DataFrameRow{DataFrames.DataFrame},
        e_range::Tuple; plot_name::String="",
        kwargs...
    )
    mission = plot_data.mission

    types_unbinned = [ESpec]
    types_binned   = [BinnedData, PgramData]
    types_binsec   = [FFTData, FFTSpec, PulseC]

    plot_type_symbols = Dict(
        ESpec      => :espec,
        BinnedData => :lcurve,
        PgramData  => :pgram,
        FFTData    => :fspec,
        FFTSpec    => :sgram,
        PulseC     => :pulsec
    )

    plot_type = plot_type_symbols[typeof(plot_data)]

    if plot_name == ""
        plot_name = string(plot_type)
    end

    plot_path = if any([isa(plot_data, tu) for tu in types_unbinned])
        savefig(mission, obs_row, e_range, plot_type, plot_name; kwargs...)
    elseif any([isa(plot_data, tb) for tb in types_binned])
        savefig(mission, obs_row, e_range, plot_type, plot_name; bin_time=plot_data.bin_time, kwargs...)
    elseif any([isa(plot_data, tbs) for tbs in types_binsec])
        bss = plot_data.bin_time * plot_data.bin_size
        savefig(mission, obs_row, e_range, plot_type, plot_name; bin_time=plot_data.bin_time, bin_size_sec=bss, kwargs...)
    else
        @error "Data type not found in savefig configuration"
    end

    return plot_path
end

"""
    _recursive_first(instrument_data::Dict)

Recursively uses the `first()` function to pull out the first actual element
of a dictionary

Useful when dealing with nested `insturment_data` such as `Dict{Symbol,Dict{Int64,JAXTAM.FFTData}}`
as it will pull out a value of type `FFTData`, typically when getting other fields like `bin_time`
"""
function _recursive_first(instrument_data::Dict)
    f = first(instrument_data)[2]

    if typeof(f) <: Dict
        return _recursive_first(f)
    else
        return f
    end
end

function _plot_formatter!()
    # Plots.plot!(title_location=:left, titlefontsize=12, margin=2mm, xguidefontsize=10, yguidefontsize=10)
    return Plots.plot!(title_location=:left, margin=2mm, xguidefontsize=12, yguidefontsize=12)
end

function _log10_minor_ticks!(start=0.01, stop=4096)
    start < 1 ? start = 1/nextpow(10, 1/start) : start = prevpow(10, start)
    stop  < 1 ? stop  = 1/prevpow(10, 1/stop)  : stop  = nextpow(10, stop)

    segments = Int(log10(stop./start)+1)
    ranges   = repeat(1:9, inner=(1, segments)) # Skip powers of 10 (1, 10) to avoid redrawing

    segment_mults = 10 .^(0:segments-1) .* start
    segment_mults = repeat(segment_mults', outer=(9,1))

    log10_minors = ranges .* segment_mults

    vline!(log10_minors[2:end, :][:], color=:grey, alpha=0.2, lab="")

    return log10_minors[1, :]
end

function _save_plot_data_csv(obs_row, bin_time, subfolder, file_name; kwargs...)
    @warn "FIX THIS"
    file_dir  = joinpath(obs_row[1, :obs_path], "JAXTAM/lc/$bin_time/images/", subfolder)
    file_path = joinpath(file_dir, file_name)

    mkpath(file_dir)

    data = DataFrame(kwargs)
    data = hcat([data[i] for i in 1:size(data,2)]...)

    writedlm(file_path, data)

    println("CSV: $file_path")

    return data
end

# Lightcurve plotting functions

function plot!(data::BinnedData; lab="", size_in=(1140,400), title_append="", logy=false)
    bin_time_pow2 = Int(log2(data.bin_time))

    Plots.plot()

    plot_title = "Lightcurve - 2^$(bin_time_pow2) - $(data.bin_time) bt$title_append"

    Plots.plot!(data.times, data.counts,
        xlab="Time (s)",
        lab=lab, alpha=1, title=plot_title)

    # NOTE: Disabled GTI start/stop label
    Plots.vline!(data.gtis[2], lab="",  alpha=0.75)
    Plots.vline!(data.gtis[1], lab="", alpha=0.75)

    count_min = maximum([minimum(data.counts[data.counts .!= 0]), 0.1])
    count_max = maximum(data.counts)
    if logy
        log10_min = count_min > 1 ? prevpow(10, count_min) : 1/prevpow(10, 1/count_min)
        yticks = range(log10(log10_min), stop=log10(nextpow(10, count_max)), length=5)

        yticks = round.(exp10.(yticks), sigdigits=3)

        ylim = (log10_min, nextpow(10, count_max))
        yaxis!(yticks=yticks, ylims=ylim, yscale=:log10, ylab="Counts (log10)")
    else
        width = count_max - count_min
        yaxis!(ylims=(count_min-(width*0.1),count_max+(width*0.1)), ylab="Counts")
    end

    try
        yaxis!(yformatter=yi->round(yi, sigdigits=3))
    catch
        yaxis!(yformatter=yi->round(Int, yi))
    end

    Plots.plot!(size=size_in)

    _plot_formatter!()
    return Plots.plot!()
end

function plot(instrument_data::Dict{Symbol,JAXTAM.BinnedData}; size_in=(1140,400), save=false, title_append="")
    instruments = keys(instrument_data)

    example = _recursive_first(instrument_data)
    e_range = example.e_range
    title_append = string(" - $e_range keV", title_append)

    plt = Plots.plot()

    # NOTE: Disable instrument label
    for instrument in instruments
        plt = plot!(instrument_data[Symbol(instrument)]; lab="", title_append=title_append)
    end

    if save
        obs_row = master_query(example.mission, :obsid, example.obsid)
        savefig(example, obs_row, e_range)
    end

    return Plots.plot!(size=size_in)
end

function plot_groups(instrument_data::Dict{Symbol,JAXTAM.BinnedData}; size_in=(1140,400), save=false)
    instruments = keys(instrument_data)

    group_plots = Dict{Symbol,Dict{Int64,Plots.Plot}}()

    example = _recursive_first(instrument_data)
    e_range = example.e_range
    obs_row = master_query(example.mission, :obsid, example.obsid)

    for instrument in instruments
        instrument_group_data  = _group_return(instrument_data[instrument])
        instrument_group_plots = Dict{Int64,Plots.Plot}()

        availabel_groups = collect(keys(instrument_group_data))

        for group in availabel_groups
            group_data = instrument_group_data[group]

            title_append = " - $e_range keV - group $group/$(maximum(availabel_groups))"

            instrument_group_plots[group] = plot(Dict(instrument=>group_data); save=false,
                    title_append=title_append, size_in=size_in)

            if save
                savefig(group_data, obs_row, e_range; group=group)
            end
        end

        group_plots[instrument] = instrument_group_plots
    end

    return group_plots
end

# Power spectra plotting functions

function plot!(data::FFTData; title_append="", norm=:rms, rebin=(:log10, 0.01), freq_lims=(:start, :end),
        lab="", logx=true, logy=true, show_errors=true,
        size_in=(1140,900)
    )
    bin_time_pow2 = Int(log2(data.bin_time))

    avg_power = data.avg_power
    freq      = data.freq

    if freq_lims[1] == :start
        # Don't plot the 0Hz amplitude
        idx_min  = 2
        freq_min = freq[idx_min]
    else
        idx_min  = findfirst(freq .>= freq_lims[1])
        idx_min <= 0 ? idx_min=1 : idx_min=idx_min
        freq_min = freq[idx_min]
    end

    if freq_lims[2] == :end
        idx_max  = length(freq)
        freq_max = freq[end]
    else
        idx_max  = findfirst(freq .> freq_lims[2]) - 1 # Find first > max freq, take the freq before that
        freq_max = freq[idx_max]
    end

    avg_power = avg_power[idx_min:idx_max]
    freq   = freq[idx_min:idx_max]

    freq, avg_power, errors = _fspec_rebin(avg_power, freq, data.bin_count, data.bin_size, data.bin_time, rebin)
    ylab = ""

    if norm == :rms
        src_ctrate = mean(data.src_ctrate); bkg_ctrate = mean(data.bkg_ctrate)
        rms_factor = 1
        if bkg_ctrate == 0.0
            rms_factor = 1/src_ctrate
        else
            rms_factor = (src_ctrate + bkg_ctrate) ./ src_ctrate^2
        end

        avg_power = avg_power .- 2
        avg_power = (avg_power.*rms_factor).*freq
        errors = errors.*freq*rms_factor
        power_max = maximum(avg_power[2:end]); power_min = maximum([0.0001, minimum(abs.(avg_power[2:end]))])
        avg_power[avg_power .<=0] .= NaN
        ylab = "Amplitude (RMS*freq)"
    elseif norm == :leahym2
        errors = errors.*freq
        avg_power = (avg_power.-2).*freq
        power_max = maximum(avg_power[2:end]); power_min = minimum(abs.(avg_power[2:end]))
        avg_power[avg_power .<=0] .= NaN
        ylab = "Amplitude (Leahy - 2)*freq"
    elseif norm == :leahy
        power_max = maximum(avg_power[2:end]); power_min = minimum(avg_power[2:end])
        hline!([2], line=:dash, lab="")
        ylab = "Amplitude (Leahy)"
    else
        @error "Plot norm type '$norm' not found"
    end

    if show_errors
        Plots.plot!(freq, avg_power, color=:black,
            yerr=errors, lab=lab)
    else
        Plots.plot!(freq, avg_power, color=:black, lab=lab)
    end

    if logx
        xaxis!(xscale=:log10, xformatter=xi->xi, xlim=(freq_min, freq[end]), xlab="Freq (Hz) - log10")
        _log10_minor_ticks!(freq_min, freq[end])
    else
        xaxis!(xlab="Freq (Hz)", xlim=(freq_min,freq_max))
    end

    if logy
        # If power_min < 1, can't use prevpow10 for ylims, hacky little fix is 1/prevpow(10, 1/power_min)
        power_min > 1 ? ylim = (prevpow(10, power_min), nextpow(10, power_max)) : ylim = (1/prevpow(10, 1/power_min), nextpow(10, power_max))
        yaxis!(yscale=:log10, yformatter=yi->round(yi, sigdigits=3), ylims=ylim, ylab="$ylab - log10")
    else
        yaxis!(ylab=ylab)
    end

    Plots.plot!(alpha=1,
        title="FFT - $(data.obsid) - $(data.e_range) keV - 2^$(bin_time_pow2) bt - $(data.bin_size*data.bin_time) bs - $rebin rebin - $(data.bin_count) - sections averaged$title_append")

    _plot_formatter!()
    return Plots.plot!(size=size_in)
end

function plot(instrument_data::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}}; title_append="", freq_lims=(:start, :end),
        size_in=(1140,900), norm=:rms, rebin=(:log10, 0.01), logx=true, logy=true, show_errors=true, save=false, save_csv=false)
    instruments = keys(instrument_data)

    plt = Plots.plot()

    example = _recursive_first(instrument_data)
    e_range = example.e_range

    for instrument in instruments # NOTE: Disabled instrument label `String(instrument)`
        plt = JAXTAM.plot!(instrument_data[Symbol(instrument)][-1]; title_append=title_append, freq_lims=freq_lims,
            norm=norm, rebin=rebin, logx=logx, logy=logy, lab="", show_errors=show_errors)

        if save_csv
            obs_row = master_query(example.mission, :obsid, example.obsid)
            # _save_plot_data_csv(obs_row, example.bin_time, # TODO: RE-ENABLE THIS WHEN FIXED
            #     "fspec/$(example.bin_size*example.bin_time)", "fspec.csv";
            #     freq=plt.series_list[1].d[:x], power=plt.series_list[1].d[:y], power_error=plt.series_list[1].d[:yerror])
        end
    end

    if save
        obs_row = master_query(example.mission, :obsid, example.obsid)
        savefig(example, obs_row, e_range)
    end

    return Plots.plot!(size=size_in)
end

function plot_groups(instrument_data::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}};
    size_in=(1140,600), norm=:rms, rebin=(:log10, 0.01), freq_lims=(:start,:end), logx=true, logy=true, show_errors=true, save=false)

    instruments = keys(instrument_data)

    group_plots = Dict{Symbol,Dict{Int64,Plots.Plot}}()

    example = _recursive_first(instrument_data)
    obs_row = master_query(example.mission, :obsid, example.obsid)
    e_range = example.e_range

    for instrument in instruments
        availabel_groups = unique([gti.group for gti in values(instrument_data[instrument])])
        availabel_groups = availabel_groups[availabel_groups .>= 0] # Excluse -1, -2, etc... for scrunched/mean FFTData

        instrument_group_plots = Dict{Int64,Plots.Plot}()

        for group in availabel_groups
            instrument_data_group = Dict{Int64,JAXTAM.FFTData}()
            gtis_in_group = [gti.gti_index for gti in values(instrument_data[instrument]) if gti.group == group]

            for gti_no in gtis_in_group
                instrument_data_group[gti_no] = instrument_data[instrument][gti_no]
            end

            title_append = " - group $group/$(maximum(availabel_groups))"

            # Have to run _scrunch_sections on group data, since it lacks the -1 indexed average amplitudes
            instrument_group_plots[group] = plot(Dict(instrument=>_scrunch_sections(instrument_data_group));
                    size_in=size_in, norm=norm, rebin=rebin, freq_lims=freq_lims, logx=logx, logy=logy, save=false, show_errors=show_errors,
                    title_append=title_append)

            if save
                group_data = instrument_data_group[gtis_in_group[1]]
                savefig(group_data, obs_row, e_range; group=group)
            end

        end

        group_plots[instrument] = instrument_group_plots
    end

    return group_plots
end

# Energy Spectra plotting functions

struct ESpec end

function plot_energy!(data::JAXTAM.InstrumentData; e_lims=:auto, size_in=(1140,400), save=false, title_append="")
    if e_lims == :auto
        e_lims = _mission_good_e_range(data.mission)
    end

    exposure = sum(data.gtis[:STOP] .- data.gtis[:START])

    energies = data.events[:E]
    energies = energies[e_lims[1] .<= energies .<= e_lims[2]]

    energy_bins     = e_lims[1]:1e-2:e_lims[2]
    energy_hist_fit = OnlineStats.Hist(energy_bins)
    energy_hist     = fit!(energy_hist_fit, sort(energies))

    Plots.plot(value(energy_hist)[1], value(energy_hist)[2]./exposure;
        size=size_in, lab="",
        title="Energy Spectra - $(data.obsid) - $e_lims keV")

    Plots.xaxis!(xlab="Energy [keV]")
    Plots.yaxis!(yscale=:log10, yformatter=yi->yi, ylab="Normalised Counts - log10")

    _plot_formatter!()
    return Plots.plot!()
end

function plot_energy(instrument_data::Dict{Symbol,JAXTAM.InstrumentData}; size_in=(1140,400), save=false, title_append="")
    instruments = keys(instrument_data)

    example_lc = _recursive_first(instrument_data)

    if !haskey(example_lc.events, :E)
        @error "`instrument_data` does not contain the energy column. Are you passing data from read_cl or calibrate?"
    end

    for instrument in instruments
        plt = plot_energy!(instrument_data[instrument], save=false, title_append=title_append)
    end

    return Plots.plot!(size=size_in)
end

# Periodogram plotting functions

function plot!(data::PgramData; title_append="", rebin=(:linear, 1),
        lab="", logx=false, logy=true, size_in=(1140,600))
    bin_time_pow2 = Int(log2(data.bin_time))

    # Don't plot the 0Hz amplitude
    power = data.power
    freq  = data.freq
    power[1] = NaN
    freq[1]   = NaN

    if rebin == (:linear, 1)
        # Do... nothing
    else
        freq, power, errors = _fspec_rebin(power, freq, 1, 1, missing, rebin)
    end

    Plots.plot!(xlab="Freq (Hz)", alpha=1)

    if logy
        power_min = maximum([0.0001 minimum(power[2:end])])
        power_max = maximum(power[2:end])

        power = power[:]
        power[power .<= 0] .= NaN

        # If power_min < 1, can't use prevpow10 for ylims, hacky little fix is 1/prevpow(10, 1/power_min)
        # removed that anyway, set ylim to 1 if power_min < 1
        # TODO: Look at/fix the manual ylim settings, since it seems to... make things worse usually
        # power_min > 1 ? ylim = (prevpow(10, power_min), nextpow(10, power_max)) : ylim = (1, nextpow(10, power_max))
        yaxis!(yscale=:log10, yformatter=yi->round(yi, sigdigits=3), ylims=(power_min,power_max))
    else
        yaxis!(ylims=(maximum([0.0, minimum(power[2:end])]), maximum(power[2:end])*1.1))
    end

    xlim = (freq[2], freq[end])
    if logx
        xaxis!(xscale=:log10, xformatter=xi->xi, xlim=xlim, xlab="Freq (Hz) - log10")
        pow10_ticks = _log10_minor_ticks!(xlim[1], xlim[2])
        xticks!(pow10_ticks)
    else
        xaxis!(xlab="Freq (Hz)", xlim=xlim)
    end

    # NOTE: Disabled instrument and pgram type label `lab="$lab - $(data.pg_type)"`
    Plots.plot!(freq, power, color=:black, ylab="Amplitude", lab="",
        title="Periodogram - $(data.obsid) - $(data.e_range) keV - 2^$(bin_time_pow2) bt - $rebin rebin$title_append")

    _plot_formatter!()
    return Plots.plot!(size=size_in)
end

function plot(instrument_data::Dict{Symbol,JAXTAM.PgramData};
        rebin=(:linear, 1), logx=true, logy=false, save=false,
        size_in=(1140,600), title_append="")

    instruments = keys(instrument_data)

    plt = Plots.plot()

    for instrument in instruments # NOTE: Disabled instrument label `string(instrument)`
        plt = plot!(instrument_data[instrument], lab="",
            rebin=rebin, size_in=size_in, logx=logx, logy=logy, title_append=title_append)
    end

    if save
        example = _recursive_first(instrument_data)
        obs_row = master_query(example.mission, :obsid, example.obsid)
        e_range = example.e_range

        savefig(example, obs_row, e_range)
    end

    return Plots.plot!(size=size_in)
end

function plot_groups(instrument_data::Dict{Symbol,Dict{Int64,JAXTAM.PgramData}};
    rebin=(:linear, 1), logx=true, logy=false, save=false,
    size_in=(1140,600), title_append="")

    instruments = keys(instrument_data)

    example = _recursive_first(instrument_data)
    obs_row = master_query(example.mission, :obsid, example.obsid)
    e_range = example.e_range

    group_plots = Dict{Symbol,Dict{Int64,Plots.Plot}}()
    for instrument in instruments
        instrument_group_plots = Dict{Int64,Plots.Plot}()
        availabel_groups = unique([gti.group for gti in values(instrument_data[instrument])])
        availabel_groups = availabel_groups[availabel_groups .>= 0] # Excluse -1, -2, etc... for scrunched/mean FFTData

        for group in availabel_groups
            title_append = " - group $group/$(maximum(availabel_groups))"

            instrument_group_plots[group] = plot(Dict(instrument=>instrument_data[instrument][group]);
                title_append=title_append, size_in=size_in, save=false)

            if save
                savefig(example, obs_row, e_range; group=group)
            end
        end

        group_plots[instrument] = instrument_group_plots
    end

    return group_plots
end

# Spectrogram plotting functions

struct FFTSpec <: JAXTAMData
    mission    :: Mission
    instrument :: Symbol
    obsid      :: String
    e_range    :: Tuple{Float64,Float64}
    bin_time   :: Real
    bin_size   :: Int
end
FFTSpec(fs::JAXTAM.FFTData) = FFTSpec(fs.mission, fs.instrument, fs.obsid, fs.e_range, fs.bin_time, fs.bin_size)

function _plot_sgram(sgram_freq, sgram_power, sgram_bounds, sgram_groups,
        e_range, obsid, bin_time_pow2, bin_size, bin_time, rebin, size_in, disable_x=true)

    # if size(sgram_power) ==

    x_tick_approx = exp10.(-1:1:4)
    x_ticks = [findfirst(f->f>=x_lim, sgram_freq) for x_lim in x_tick_approx]
    x_ticks = unique(x_ticks)
    x_ticks = x_ticks[x_ticks.!=nothing]

    try
        heatmap(sgram_power, size=size_in, fill=true)
    catch err
        if typeof(err) != InterruptException
            throw(JAXTAMError("Error plotting spectrogram", :plot, err))
        else
            rethrow(err)
        end
    end
    xaxis!(xticks=x_ticks, xformatter=yi->sgram_freq[Int(yi)], xlab="Freq (Hz - log10)")

    title!("Spectrogram - $(obsid) - $e_range keV - 2^$(bin_time_pow2) bt - $(bin_size*bin_time) bs - $rebin rebin")

    hline!(sgram_bounds.+0.5, alpha=0.75, line=:dot, lab="")
    if length(sgram_bounds) < 25
        yaxis_bounds_to_group = Dict(diag([(bound,group) for bound in sgram_bounds, group in sgram_groups]))
        yaxis!(yticks=sgram_bounds, yformatter=yi->yaxis_bounds_to_group[Int(yi)], ylab="Group")
    else
        yaxis!(ylab="Groups")
    end

    _plot_formatter!()
end

function plot_sgram(fs::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}};
        rebin=(:log10, 0.01), size_in=(1140,600), save=false)
    instruments = keys(fs)

    sgram_instrument_plots = Dict{Symbol,Plots.Plot}()
    for instrument in instruments
        example = fs[instrument][-1]
        bin_time_pow2 = Int(log2(example.bin_time))

        sgram_freq, sgram_power, sgram_bounds, sgram_groups = fspec_rebin_sgram(fs[instrument], rebin=rebin)

        if size(sgram_power, 2) == 1
            @warn "sgram_power only has one dimension, cannot plot spectrogram"
            return sgram_instrument_plots
        end

        src_ctrate = mean(fs[instrument][-1].src_ctrate); bkg_ctrate = mean(fs[instrument][-1].bkg_ctrate)
        rms_factor = 1

        if bkg_ctrate == 0.0
            rms_factor = 1/src_ctrate
        else
            rms_factor = (src_ctrate + bkg_ctrate) ./ src_ctrate^2
        end

        sgram_power = sgram_power .- 2
        sgram_power = (sgram_power.*rms_factor).*sgram_freq
        sgram_power[sgram_power .<= 0] .= 0
        sgram_power = sgram_power'

        sgram_instrument_plots[instrument] = _plot_sgram(sgram_freq, sgram_power, sgram_bounds, sgram_groups,
            example.e_range, example.obsid, bin_time_pow2, example.bin_size, example.bin_time, rebin,
            size_in)

        if save
            obs_row = master_query(example.mission, :obsid, example.obsid)
            savefig(FFTSpec(example), obs_row, example.e_range)
        end
    end

    return sgram_instrument_plots
end

# Pulsation Check Candle Plotting

struct PulseC  <: JAXTAMData
    mission    :: Mission
    instrument :: Symbol
    obsid      :: String
    e_range    :: Tuple{Float64,Float64}
    bin_time   :: Real
    bin_size   :: Int
end
PulseC(fs::JAXTAMData) = PulseC(fs.mission, fs.instrument, fs.obsid, fs.e_range, fs.bin_time, fs.bin_size)

function _plot_pulses_candle(power, freq, power_limit, f_min, f_max)
    base_freq = freq[:, 1]

    if f_min == :start
        f_min = base_freq[2]
    end

    if f_max == :end
        f_max = base_freq[end]
    end

    # Mask for power limits, and fo excluding zero-freq
    mask = (power .>= power_limit) .* (freq .!= 0)

    power = power[mask]
    freq  = freq[mask]

    Plots.plot()
    _log10_minor_ticks!(0.01, 4096)
    Plots.plot!(freq, power, line=:sticks, lab="")

    return Plots.scatter!(freq, power, lab="", alpha=0.5)
end

function plot_pulses_candle(fs::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}};
    power_limit=30, size_in=(1140,600), save=false)

    plots_candle = Dict{Symbol,Plots.Plot}()

    example = _recursive_first(fs)

    e_range = example.e_range
    bin_time_pow2 = Int(log2(example.bin_time))
    bin_time = example.bin_time
    bin_size = example.bin_size
    obsid = example.obsid

    instruments = keys(fs)
    for instrument in instruments
        power = hcat([f[2].power for f in fs[instrument] if f[1] > 0]...)
        freq  = repeat(fs[instrument][-1].freq, inner=(1, size(power, 2)))

        _plot_pulses_candle(power, freq, power_limit, 0.01, :end)

        Plots.title!("Pulsations - $obsid - $e_range keV - 2^$(bin_time_pow2) bt - $(bin_size*bin_time) bs")

        xaxis!(xscale=:log10, xformatter=xi->xi, xlim=(0.01, freq[end, 1]), xlab="Freq [Hz] - log10")
        yaxis!(ylab=("Amplitude (Leahy) >= $power_limit"), ylims=(power_limit, maximum(power[2:end, :])*1.1))

        _plot_formatter!()

        plots_candle[instrument] = Plots.plot!(size=size_in)

        if save
            obs_row = master_query(example.mission, :obsid, example.obsid)
            savefig(PulseC(example), obs_row, e_range)
        end
    end

    return plots_candle
end

function plot_pulses_candle_groups(fs::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}};
    power_limit=30, size_in=(1140,600), save=false, f_lims=(0.01, :end))

    instruments = keys(fs)

    example = _recursive_first(fs)
    obs_row = master_query(example.mission, :obsid, example.obsid)
    e_range = example.e_range
    bin_time_pow2 = Int(log2(example.bin_time))
    bin_time = example.bin_time
    bin_size = example.bin_size
    obsid = example.obsid

    group_plots = Dict{Symbol,Dict{Int64,Plots.Plot}}()
    for instrument in instruments
        instrument_group_plots = Dict{Int64,Plots.Plot}()
        availabel_groups = unique([f[2].group for f in fs[instrument] if f[1] > 0])

        for group in availabel_groups
            title_append = " - group $group/$(maximum(availabel_groups))"

            power = hcat([f[2].power for f in fs[instrument] if f[2].group == group]...)
            freq  = repeat(fs[instrument][-1].freq, inner=(1, size(power, 2)))

            f_lims[1] == :start ? f_lims = (freq[1, 1], f_lims[2])    : ""
            f_lims[2] == :end   ? f_lims = (f_lims[1] , freq[end, 1]) : ""

            _plot_pulses_candle(power, freq, power_limit, f_lims[1], f_lims[2])

            Plots.title!("Pulsations - $obsid - $e_range keV - 2^$(bin_time_pow2) bt - $(bin_size*bin_time) bs$title_append")

            xaxis!(xscale=:log10, xformatter=xi->xi, xlim=f_lims, xlab="Freq [Hz] - log10")
            yaxis!(ylab=("Amplitude (Leahy) >= $power_limit"), ylims=(power_limit, maximum(power[2:end, :])*1.1))

            _plot_formatter!()

            instrument_group_plots[group] = Plots.plot!(size=size_in)

            if save
                savefig(PulseC(example), obs_row, e_range; group=group)
            end
        end

        group_plots[instrument] = instrument_group_plots
    end

    return group_plots
end

# Pulsation Check Spectrogram Plotting

function plot_pulses(fs::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}};
        freq_bin=10, power_limits=[10, 25, 50], size_in=(1140,600), save=false)

    pulsation_instrument_plots = Dict{Symbol,Plots.Plot}()
    instruments = keys(fs)
    for instrument in instruments
        example_data = fs[instrument][-1]
        bin_time_pow2 = Int(log2(example_data.bin_time))
        e_range = (config(example_data.mission).good_energy_min, config(example_data.mission).good_energy_max)

        pulsation_freq, pulsation_power, pulsation_bounds, pulsation_groups = 0, 0, 0, 0
        for p in power_limits
            rebin = (:freq_binary, freq_bin, p)

            pulsation_freq, pulsation_power_new, pulsation_bounds, pulsation_groups = fspec_rebin_sgram(fs[instrument]; rebin=rebin)

            if pulsation_power == 0
                pulsation_power = pulsation_power_new
            else
                pulsation_power[pulsation_power_new .!= 0] .= p
            end
        end

        rebin = (:freq_binary, freq_bin, power_limits)

        pulsation_instrument_plots[instrument] = _plot_sgram(pulsation_freq, pulsation_power', pulsation_bounds, pulsation_groups,
            e_min, e_max, example_data.obsid, bin_time_pow2, example_data.bin_size, example_data.bin_time,
            rebin, size_in, false)

        if save
            obs_row = master_query(example_data.mission, :obsid, example_data.obsid)
            _savefig_obsdir(example_data.mission, obs_row, example_data.bin_time, "pulse/$(example_data.bin_size*example_data.bin_time)", "pulse.png")
        end
    end

    return pulsation_instrument_plots
end

# Covariance plotting

function plot_fspec_cov1d(fs::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}}; size_in=(1140,600), rebin=(:log10, 0.01))
    instruments = keys(fs)

    example_data = _recursive_first(fs)
    obsid = example_data.obsid
    bin_time = example_data.bin_time
    bin_time_pow2 = Int(log2(example_data.bin_time))
    bin_size = example_data.bin_size

    cov1d_plots = Dict{Symbol,Plots.Plot}()
    for instrument in instruments
        fspec_freq, fspec_power = JAXTAM.fspec_rebin_sgram(fs[instrument]; rebin=rebin)

        fspec_power[1, :] .= 0
        src_ctrate = mean(example_data.src_ctrate); bkg_ctrate = mean(example_data.bkg_ctrate)
        rms_factor = 1
        if bkg_ctrate == 0.0
            rms_factor = 1/src_ctrate
        else
            rms_factor = (src_ctrate + bkg_ctrate) ./ src_ctrate.^2
        end

        fspec_power = fspec_power .- 2
        fspec_power = (fspec_power.*rms_factor).*fspec_freq

        fspec_diag = diag(cov(fspec_power, dims=2))
        yaxis_max  = nextpow(10, maximum(fspec_diag))

        bin_count = size(fspec_power, 2)

        fspec_diag[fspec_diag .<= 1e-8] .= NaN

        Plots.plot(fspec_freq, fspec_diag, lab="", #NOTE: Disabled instrument label `lab=instrument`
            color=:black, size=size_in,
            title="FFT 1D Covariance - $(obsid) - 2^$(bin_time_pow2) bt - $(bin_size*bin_time) bs - $rebin rebin - $(bin_count) sections averaged")

        xaxis!(xscale=:log10, xformatter=xi->xi, xlab="Freq (Hz - log10)", xlims=(0.01, nextpow(2, maximum(fspec_freq))))
        yaxis!(yscale=:log10, yformatter=xi->xi, ylab="Cov (diag - log10)") # , ylims=(10, yaxis_max)
        # hline!([4000], lab="4000 - Threshold")

        cov1d_plots[instrument] = _plot_formatter!()
    end

    return cov1d_plots
end

function _plot_cov2d(fs::Dict{Int64,JAXTAM.FFTData}, rebin::Tuple, zoom_log10=false)
    fspec_freq, fspec_power = JAXTAM.fspec_rebin_sgram(fs; rebin=rebin)

    fspec_power[1, :] .= 0
    src_ctrate = mean(fs[-1].src_ctrate); bkg_ctrate = mean(fs[-1].bkg_ctrate)
    rms_factor = 1
    if bkg_ctrate == 0.0
        rms_factor = 1/src_ctrate
    else
        rms_factor = (src_ctrate + bkg_ctrate) ./ src_ctrate.^2
    end

    fspec_power = fspec_power .- 2
    fspec_power = (fspec_power.*rms_factor).*fspec_freq

    if rebin[1] == :linear
        fspec_cov_2d = cov(reverse(reverse(fspec_power, dims=2), dims=1), dims=2)
    elseif rebin[1] == :log10
        fspec_cov_2d = cov(fspec_power, dims=2)
    else
        error("Invalid rebin type: $(rebin[1])")
    end

    fspec_cov_2d[isnan.(fspec_cov_2d)] .= 0

    max_ind_2d = findmax(fspec_cov_2d)[2]
    max_freq_x = fspec_freq[max_ind_2d[1]]
    max_freq_y = fspec_freq[max_ind_2d[2]]

    if zoom_log10!=false && rebin[1]==:log10
        zoom_ind_start = findfirst(fspec_freq .>= max_freq_x - max_freq_x/2)
        zoom_ind_stop  = findfirst(fspec_freq .>= max_freq_x + max_freq_x/2)
        println(zoom_log10)
        println(zoom_ind_start)
        println(zoom_ind_stop)
        fspec_cov_2d = fspec_cov_2d[zoom_ind_start:zoom_ind_stop, zoom_ind_start:zoom_ind_stop]

        max_ind_2d = findmax(fspec_cov_2d)[2]
        max_freq_x = fspec_freq[max_ind_2d[1]]
        max_freq_y = fspec_freq[max_ind_2d[2]]

        rebin = (rebin[1], rebin[2], "$(zoom_log10) zoom")
    end

    heatmap(fspec_cov_2d, legend=false, aspect_ratio=:equal, xlab="$rebin cov")

    vline!([max_ind_2d[1]], color=:cyan, alpha=0.25, line=:dot)
    hline!([max_ind_2d[2]], color=:cyan, alpha=0.25, line=:dot)
    xticks!([max_ind_2d[1]]); yticks!([0])
    xaxis!(xformatter=xi->"$(round(max_freq_x, sigdigits=3)) Hz")
    yaxis!(yformatter=xi->"$(round(max_freq_y, sigdigits=3)) Hz")
end

function plot_fspec_cov2d(fs::Dict{Symbol,Dict{Int64,JAXTAM.FFTData}}; size_in=(1140,600*2))
    instruments = keys(fs)

    example_data = _recursive_first(fs)
    obsid = example_data.obsid
    bin_time = example_data.bin_time
    bin_time_pow2 = Int(log2(example_data.bin_time))
    bin_size = example_data.bin_size
    fs_length = length(example_data.avg_power)

    rebin_lin = (:linear, maximum([floor(Int, fs_length/1024), 1]))
    rebin_log = (:log10, 0.01)

    for instrument in instruments
        fspec_freq, fspec_power = JAXTAM.fspec_rebin_sgram(fs[instrument]; rebin=rebin_lin)

        cov2d_linear_plot = _plot_cov2d(fs[instrument], rebin_lin)

        cov2d_log_plot = _plot_cov2d(fs[instrument], rebin_log, false)

        cov2d_log_plot_x1 = _plot_cov2d(fs[instrument], rebin_log, 1)
        cov2d_log_plot_x2 = _plot_cov2d(fs[instrument], rebin_log, 0.5)

        bin_count = size(fspec_power, 2)

        l = @layout [b c; d e] # [a{.001h}; [b c; d e]]
        dual_cov_plot = Plots.plot(
            # Plots.plot(title="FFT 2D Covariances - $(obsid) - 2^$(bin_time_pow2) bt - $(bin_size*bin_time) bs - $(bin_count) sections averaged",
            #     #annotation=(0.25, 0.5,
            #     #"FFT 2D Covariances - $(obsid) - 2^$(bin_time_pow2) bt - $(bin_size*bin_time) bs - $(bin_count) sections averaged", 12),
            #     framestyle = :none),
            cov2d_linear_plot,
            cov2d_log_plot,
            cov2d_log_plot_x1,
            cov2d_log_plot_x2,
            layout=l,
            size=size_in)

        _plot_formatter!()

        return dual_cov_plot
    end
end


text:
using Plots
using OceanParameterizations

files =  ["free_convection", "strong_wind", "strong_wind_no_coriolis", "weak_wind_strong_cooling",
          "strong_wind_weak_cooling", "strong_wind_weak_heating"]

## T profiles

Ts = Dict() # maps file name to T array
Ts2 = Dict() # maps file name to T array
for file in files
    les = OceanParameterizations.DataWrangling.read_les_output(file) # <: OceananigansData
    Ts[file] = ZeroMeanUnitVarianceScaling(les.T).(les.T)
    Ts2[file] = MinMaxScaling(les.T).(les.T)
end

Ts
Ts2

anim = @animate for i in 1:289
    plot(legend=false, xlabel="T", ylabel="z", xlims=(-1.7,2.3))
    for (file, T) in Ts
        plot!(T[:,i], collect(-255:2:0), title = "ZeroMeanUnitVarianceScaling", label="$(file)")
    end
end
gif(anim, pwd()*"/ZeroMeanUnitVarianceScaling.gif", fps=20)

anim = @animate for i in 1:289
    plot(legend=false, xlabel="T", ylabel="z", xlims=(0.0,1.0))
    for (file, T) in Ts
        plot!(T[:,i], collect(-255:2:0), title = "MinMaxScaling", label="$(file)")
    end
end
gif(anim, pwd()*"/MinMaxScaling.gif", fps=20)

## uw profiles unscaled

Ts3 = Dict() # maps file name to T array
for file in files
    les = OceanParameterizations.DataWrangling.read_les_output(file) # <: OceananigansData
    Ts3[file] = les.wu
end

Ts3

anim = @animate for i in 1:289
    plot(legend=false, xlabel="uw", ylabel="z", xlims=(-0.0007,0.0001))
    for (file, T) in Ts3
        plot!(T[:,i], collect(-255:2:1), title = "Unscaled", label="$(file)")
    end
end
gif(anim, pwd()*"/uw_unscaled.gif", fps=20)

## uw profiles

Ts = Dict() # maps file name to T array
Ts2 = Dict() # maps file name to T array
for file in files
    les = OceanParameterizations.DataWrangling.read_les_output(file) # <: OceananigansData
    Ts[file] = ZeroMeanUnitVarianceScaling(les.wu).(les.wu)
    Ts2[file] = MinMaxScaling(les.wu).(les.wu)
end

Ts
Ts2

anim = @animate for i in 1:289
    plot(legend=false, xlabel="uw", ylabel="z", xlims=(-3,3))
    for (file, T) in Ts
        plot!(T[:,i], collect(-255:2:1), title = "ZeroMeanUnitVarianceScaling", label="$(file)")
    end
end
gif(anim, pwd()*"/ZeroMeanUnitVarianceScaling_uw.gif", fps=20)

anim = @animate for i in 1:289
    plot(legend=false, xlabel="uw", ylabel="z", xlims=(0,1))
    for (file, T) in Ts2
        plot!(T[:,i], collect(-255:2:1), title = "MinMaxScaling", label="$(file)")
    end
end
gif(anim, pwd()*"/MinMaxScaling_uw.gif", fps=20)

## vw profiles unscaled

Ts3 = Dict() # maps file name to T array
for file in files
    les = OceanParameterizations.DataWrangling.read_les_output(file) # <: OceananigansData
    Ts3[file] = les.wv
end

Ts3

anim = @animate for i in 1:289
    plot(legend=false, xlabel="vw", ylabel="z", xlims=(-0.0001,0.0005))
    for (file, T) in Ts3
        plot!(T[:,i], collect(-255:2:1), title = "Unscaled", label="$(file)")
    end
end
gif(anim, pwd()*"/vw_unscaled.gif", fps=20)

## wT profiles unscaled

Ts3 = Dict() # maps file name to T array
for file in files
    les = OceanParameterizations.DataWrangling.read_les_output(file) # <: OceananigansData
    Ts3[file] = les.wT
end

Ts3

anim = @animate for i in 1:289
    plot(legend=false, xlabel="wT", ylabel="z", xlims=(-0.00009,0.00006))
    for (file, T) in Ts3
        plot!(T[:,i], collect(-255:2:1), title = "Unscaled", label="$(file)")
    end
end
gif(anim, pwd()*"/wT_unscaled.gif", fps=20)


text:
# this is a little Monte Carlo exercise that illustrates that
# the OLS estimator is biased when we have an autoregressive model
# with weak exogeneity
using Plots
reps = 1000 # number of Monte Carlo reps.
n = 20 # sample size
betas = zeros(reps)
truebetas = [0, 0.9]
for i = 1:reps
	x = zeros(n+1)
	x[1] = 0.0
	# generate AR(1) data
	for t = 2:n+1
		x[t] = truebetas[1] + truebetas[2]*x[t-1] + randn()
	end
	y = x[2:n+1]    # dependent variable
	x = x[1:n]      # explanatory variable is the lagged dep var.
	x = [ones(n,1) x]
	beta = y\x
	betas[i] = beta[2]
end

betas = betas .- truebetas[2]
histogram(betas, label="")
gui()
#savefig("Biased.svg")


text:
using Test
using testrepo

@test testrepo.greet() == "Hello"

max_stars_repo_name: mfherbst/testrepo

max_issues_repo_name: mfherbst/testrepo

max_forks_repo_name: mfherbst/testrepo


text:
using Base.Test
using BackgroundSegmenter

tic()
@time @testset "Test queue" begin include("test_queue.jl") end
@time @testset "Test mixture model" begin include("test_mixture.jl") end
@time @testset "Test component finder" begin include("test_components.jl") end
@time @testset "Test MinCut algorithm" begin include("test_mincut.jl") end
@time @testset "Test Path data structure" begin include("test_path.jl") end
toc()