text:
#=
devel_profile:
- Julia version: 1.1.0
- Author: Jirik
- Date: 2019-04-26
=#

println("start")
function profile_test(n)
    for i = 1:n
        A = randn(100,100,20)
        m = maximum(A)
        Am = mapslices(sum, A; dims=2)
        B = A[:,:,5]
        Bsort = mapslices(sort, B; dims=1)
        b = rand(100)
        C = B.*b
    end
end

profile_test(1)  # run once to trigger compilation
using Profile
Profile.clear()  # in case we have any previous profiling data
@profile profile_test(10)

println("after profile")
using ProfileView
ProfileView.view()

ProfileView.svgwrite("profile_results.svg")


text:
module myJuliaUtils

using Random
using LinearAlgebra
using Statistics
using Distributions
using StatsBase: counts, Histogram, fit

export

# extfunc.jl
lag0, eye, vech, vechRshp, vec2sym, vec2ltri, vec2ltriW1, stabcheck, stabcheckC, preparexy, sumsqr, nanmean,
wish, iwish, cholPSD, regMat2PD!, randnPSD, genPSDmatStrict, genPDmat, genPSDmat, quantileArr, companionf,
ismyapprox, normpdf, acf, decVCV, ols1, VARols1, inbetween, getmultdiag!,transf1To,transf1Back,split_data,standardize,unstandardize, showprog,
# evalF.jl
useN2fit, useMvN2fit, pspit, pscrps,
# data_manipulation.jl
dataTransf, detrend,
# generate.jl
genLPM, genEndoMC, simMC, genMC

include("extfunc.jl");
include("evalF.jl");
include("data_manipulation.jl");
include("generate.jl");

end # module

ext: jl

max_stars_repo_name: zymbuzz/myJuliaUtils

max_issues_repo_name: zymbuzz/myJuliaUtils

max_forks_repo_name: zymbuzz/myJuliaUtils


text:
# Using ffmpeg, create a movie

# following https://trac.ffmpeg.org/wiki/Create%20a%20video%20slideshow%20from%20images
using consonance

ins = "scratch/vis%04d.png"
outs = "vis_out.mp4"
run(`ffmpeg -framerate 20.0 -i $ins -c:v libx264 -r 30 -pix_fmt yuv420p $outs`)
# run(`ffmpeg -framerate 3.0 -i %04d.png -c:v libx264 -r 30 -pix_fmt yuv420p out.mp4`)

max_stars_repo_name: elnjensen/DiskJockey

max_issues_repo_name: elnjensen/DiskJockey

max_forks_repo_name: elnjensen/DiskJockey


text:
using UnitfulAstro
using Test

const u = UnitfulAstro

@testset "UnitfulAstro.jl" begin
    @testset "unit tests" begin
        @test 10^7*u.erg == 1*u.J
        @test 10^5*u.dyn == 1*u.N
        @test 60*u.arcminute == 1*u.°
        @test 60*u.arcsecond == 1*u.arcminute
        @test 3600*u.arcsecond == 1*u.°
        @test 365.25*u.d == 1*u.yr
        @test 1*u.ly == 1*u.c*u.yr
        @test 1*u.pc == 1*u.AU/u.arcsecond
        @test 1*u.angstrom == 1e-10*u.m
        @test 1*u.SFU ≈ 10000*u.Jy
        @test 1*u.TECU == 1e12*u.cm^-2
    end

    @testset "simple calculations" begin
        # gravitational acceleration at Earth's surface
        @test isapprox(1*u.GMearth/u.Rearth^2, 1*u.gn, atol=0.001*u.gn)
        # Solar flux at Earth's surface
        @test isapprox(1*u.Lsun/(4π*u.AU^2), 1*u.Ssun, atol=0.001*u.Ssun)
        # rest mass energy of the sun
        @test isapprox(1*u.Msun*u.c^2, 1787*u.foe, atol=0.5*u.foe)
    end

    @testset "magnitudes" begin
        @test isapprox(3631*u.Jy, 0*u.AB_mag, atol=0.001*u.AB_mag)
        @test isapprox(36.31*u.Jy, 5*u.AB_mag, atol=0.001*u.AB_mag)
        @test isapprox(363.1*u.mJy, 10*u.AB_mag, atol=0.001*u.AB_mag)
        @test isapprox(3.631*u.mJy, 15*u.AB_mag, atol=0.001*u.AB_mag)
        @test isapprox(5*u.AB_mag + 5*u.AB_mag,  4.247425010840047*u.AB_mag, atol=0.001*u.AB_mag)
        @test isapprox(5*u.AB_mag / 100, 10*u.AB_mag, atol=0.001*u.AB_mag)
        @test isapprox(5*u.AB_mag + 10*u.Jy, 46.31*u.Jy, atol=0.001*u.AB_mag)

        @test isapprox(3.0128e28*u.W, 1u.bol_Mag, atol=0.001*u.bol_Mag)
        @test isapprox(2.518_021_002e-8*u.W*u.m^-2, 1u.bol_mag, atol=0.001*u.bol_mag)

        @test isapprox(1u.U_mag, 1810*u.Jy, atol=0.001*u.U_mag)
        @test isapprox(1u.B_mag, 4260*u.Jy, atol=0.001*u.B_mag)
        @test isapprox(1u.V_mag, 3640*u.Jy, atol=0.001*u.V_mag)
        @test isapprox(1u.R_mag, 3080*u.Jy, atol=0.001*u.R_mag)
        @test isapprox(1u.I_mag, 2550*u.Jy, atol=0.001*u.I_mag)
        @test isapprox(1u.J_mag, 1600*u.Jy, atol=0.001*u.J_mag)
        @test isapprox(1u.H_mag, 1080*u.Jy, atol=0.001*u.H_mag)
        @test isapprox(1u.K_mag,  670*u.Jy, atol=0.001*u.K_mag)
        @test isapprox(1u.g_mag, 3730*u.Jy, atol=0.001*u.g_mag)
        @test isapprox(1u.r_mag, 4490*u.Jy, atol=0.001*u.r_mag)
        @test isapprox(1u.i_mag, 4760*u.Jy, atol=0.001*u.i_mag)
        @test isapprox(1u.z_mag, 4810*u.Jy, atol=0.001*u.z_mag)

        @test iszero(1u.B_mag - 1u.V_mag)
        @test isapprox(5*u.B_mag - 1u.V_mag, 4)
        @test isapprox(1u.B_mag - 5*u.V_mag, -4)
        @test isapprox(1u.B_mag - 0.5u.V_mag, 0.5)
        @test_throws u.MagnitudeError 1u.B_mag + 2u.V_mag
        @test_throws u.MagnitudeError 1u.B_mag * 2u.V_mag
        @test_throws u.MagnitudeError 1u.B_mag / 2u.V_mag
        @test_throws u.MagnitudeError 1u.B_mag // 2u.V_mag

        # Extra tests for rootpower calculations
        flux = 10u.erg
        mag = 1u.mag
        @test isapprox(flux * mag, 3.981071705534972u.erg)
        flux = 10u.erg/u.s
        @test isapprox(flux * mag, 3.981071705534972u.erg/u.s)
        flux = 10u.erg/u.s/u.cm^2
        @test isapprox(flux * mag, 3.981071705534972u.erg/u.s/u.cm^2)
        flux = 10u.erg/u.s/u.cm^2/u.angstrom
        @test isapprox(flux * mag, 3.981071705534972u.erg/u.s/u.cm^2/u.angstrom)
    end
end


text:
using LoopVectorization
using BenchmarkTools

function boundary_cyclic_kernel!(
    a, is, ie, js, je, igc, jgc)
    # East-west BCs
    @inbounds @. a[1:igc, :, :] = a[ie-igc+1:ie, :, :]
    @inbounds @. a[ie+1:end, :, :] = a[is:is+igc-1, :, :]

    # North-south BCs
    @inbounds @. a[:, 1:jgc, :] = a[:, je-jgc+1:je, :]
    @inbounds @. a[:, je+1:end, :] = a[:, js:js+igc-1, :]
end

function boundary_cyclic_kernel_loop!(
    a, is, ie, js, je, igc, jgc)
    # East-west BCs
    @tturbo for k in 1:size(a, 3)
        for j in 1:size(a, 2)
            for i in 1:igc
                a[i, j, k] = a[ie-i+1, j, k]
                a[ie+i, j, k] = a[is+i-1, j, k]
            end
        end
    end

    # North-south BCs
    @tturbo for k in 1:size(a, 3)
        for j in 1:jgc
            for i in 1:size(a, 1)
                a[i, j, k] = a[i, je-j+1, k]
                a[i, je+j, k] = a[i, js+j-1, k]
            end
        end
    end
end

function boundary_cyclic_kernel_turbo!(
    a, is, ie, js, je, igc, jgc)
    # East-west BCs
    @turbo @. a[1:igc, :, :] = a[ie-igc+1:ie, :, :]
    @turbo @. a[ie+1:end, :, :] = a[is:is+igc-1, :, :]

    # North-south BCs
    @turbo @. a[:, 1:jgc, :] = a[:, je-jgc+1:je, :]
    @turbo @. a[:, je+1:end, :] = a[:, js:js+igc-1, :]
end

a = rand(258, 3, 258)
is = 2; ie = 257; js = 2; je = 2; igc = 1; jgc = 1;

@btime boundary_cyclic_kernel!($a, $is, $ie, $js, $je, $igc, $jgc)
@btime boundary_cyclic_kernel_loop!($a, $is, $ie, $js, $je, $igc, $jgc)
@btime boundary_cyclic_kernel_turbo!($a, $is, $ie, $js, $je, $igc, $jgc)


text:
module TestFairTensor

using Test

include(joinpath("data", "data.jl"))

@testset "fair_tensor" begin
    ft = job_fairtensor()
    @test ft.mat == cat([2 2; 0 0; 0 2], [0 0; 2 1; 1 0], dims=3)
    @test Set(ft.labels) == Set(["Board", "Education", "Healthcare"])

    # Fairness Tensor Addition
    ft_1 = ft+ft
    @test ft_1.mat == ft.mat+ft.mat
end

end  # module


text:
@recipe function f(Vol::Tube{T}, coloring = missing) where T <: SSDFloat
    rStart = Vol.r_interval.left
    rStop = Vol.r_interval.right
    φStart = Vol.φ_interval.left
    φStop = Vol.φ_interval.right
    zStart = Vol.z_interval.left
    zStop = Vol.z_interval.right

    ismissing(Vol.translate) ? translate = [0.0,0.0,0.0] : translate = Vol.translate
    ismissing(coloring) ? c --> :green : c --> coloring
    @series begin
        label --> "Tube"
        partialcircle_3d(rStop,φStart,φStop,[0,0,zStart] .+ translate)
    end
    @series begin
        label:= ""
        partialcircle_3d(rStop,φStart,φStop,[0,0,zStop] .+ translate)
    end
    @series begin
        label:= ""
        partialcircle_3d(rStart,φStart,φStop,[0,0,zStart] .+ translate)
    end
    @series begin
        label:= ""
        partialcircle_3d(rStart,φStart,φStop,[0,0,zStop] .+ translate)
    end
    ## Vertical Lines

    if !iszero(rStart)
        @series begin
            label:= ""
            line_3d(rStart, rStart, φStart, φStart, zStart, zStop, translate = translate)
        end
    end
    if !iszero(rStart)
        @series begin
            label:= ""
            line_3d(rStart, rStart, φStop, φStop, zStart, zStop, translate = translate)
        end
    end
    @series begin
        label:= ""
        line_3d(rStop, rStop, φStart, φStart, zStart, zStop, translate = translate)
    end
    @series begin
        label:= ""
        line_3d(rStop, rStop, φStop, φStop, zStart, zStop, translate = translate)
    end

    ##Horizontal Lines

    if !isapprox((φStop - φStart)%2π , 0.0,atol=0.00001)
        @series begin
            label:= ""
            line_3d(rStart, rStop, φStart, φStart, zStart, zStart, translate = translate)
        end
        @series begin
            label:= ""
            line_3d(rStart, rStop, φStop, φStop, zStart, zStart, translate = translate)
        end
        @series begin
            label:= ""
            line_3d(rStart, rStop, φStart, φStart, zStop, zStop, translate = translate)
        end
        @series begin
            label:= ""
            line_3d(rStart, rStop, φStop, φStop, zStop, zStop, translate = translate)
        end
    end
end

@recipe function f(vol::Cone{T},n_aux_lines = 0, coloring = missing) where T
    ismissing(vol.translate) ? translate = [0.0,0.0,0.0] : translate = vol.translate
    ismissing(coloring) ? c --> :orange : c --> coloring
    @series begin
        label --> "Cone"
        partialcircle_3d(vol.rStop1,vol.φStart,vol.φStop,[0,0,vol.zStart]+translate)
    end
    @series begin
        label:= ""
        partialcircle_3d(vol.rStop2,vol.φStart,vol.φStop,[0,0,vol.zStop]+translate)
    end
    @series begin
        label:= ""
        partialcircle_3d(vol.rStart1,vol.φStart,vol.φStop,[0,0,vol.zStart]+translate)
    end
    @series begin
        label:= ""
        c-->coloring
        partialcircle_3d(vol.rStart2,vol.φStart,vol.φStop,[0,0,vol.zStop]+translate)
    end

    ## Vertical Lines

    @series begin
        label:= ""
        line_3d(vol.rStart1,vol.rStart2,vol.φStart,vol.φStart,vol.zStart,vol.zStop, translate = translate)
    end

    @series begin
        label:= ""
        line_3d(vol.rStart1,vol.rStart2,vol.φStop,vol.φStop,vol.zStart,vol.zStop, translate = translate)
    end

    @series begin
        label:= ""
        line_3d(vol.rStop1,vol.rStop2,vol.φStart,vol.φStart,vol.zStart,vol.zStop, translate = translate)
    end
    @series begin
        label:= ""
        line_3d(vol.rStop1,vol.rStop2,vol.φStop,vol.φStop,vol.zStart,vol.zStop, translate = translate)
    end

    ##Horizontal Lines

    if !isapprox((vol.φStop - vol.φStart)%2π , 0.0,atol=0.00001)
        @series begin
            label:= ""
            line_3d(vol.rStart1,vol.rStop1,vol.φStart,vol.φStart,vol.zStart,vol.zStart, translate = translate)
        end
        @series begin
            label:= ""
            line_3d(vol.rStart1,vol.rStop1,vol.φStop,vol.φStop,vol.zStart,vol.zStart, translate = translate)
        end
        @series begin
            label:= ""
            line_3d(vol.rStart2,vol.rStop2,vol.φStart,vol.φStart,vol.zStop,vol.zStop, translate = translate)
        end
        @series begin
            label:= ""
            line_3d(vol.rStart2,vol.rStop2,vol.φStop,vol.φStop,vol.zStop,vol.zStop, translate = translate)
        end
    end

end
function partialcircle_3d(radius,phiStart,phiStop,Translate::AbstractVector;nSteps=400)
    phirange = mylinspace(phiStart,phiStop,nSteps)

    x::Vector{AbstractFloat}=map(x->radius*cos.(x) , phirange)
    y::Vector{AbstractFloat}=map(x->radius*sin.(x) , phirange)
    z::Vector{AbstractFloat}=map(x->Translate[3], phirange)
    x.+=Translate[1]
    y.+=Translate[2]

    return x,y,z
end


text:
module SpeechAugment

using WAV: wavread
using FFTW: fft, ifft

function veclen(a::Array{T}) where T
    TYPE = typeof(a)
    if TYPE <: Array{T,1}
        return length(a)
    end
    if TYPE <: Array{T,2}
        @assert 1 in size(a)
        return length(a)
    end
end

include("./addecho.jl")
include("./addnoise.jl")
include("./distort.jl")
include("./drop.jl")
include("./farfield.jl")
include("./speed.jl")
include("./augment.jl")

end

max_stars_repo_name: sonosole/SpeechAugment

max_issues_repo_name: sonosole/SpeechAugment

max_forks_repo_name: sonosole/SpeechAugment


text:
@testset "Download Wordles" begin
    digest_wordle_number = data["wordle"]
    wordles, valid_words = Wordle.download_word_lists()

    @test begin
        wordle_digest = join(wordles[1:digest_wordle_number], ',') |> sha256 |> bytes2hex
        wordle_digest == data["data"]["known-wordles"]
    end

    @test begin
        wordle_digest = join(valid_words, ',') |> sha256 |> bytes2hex
        wordle_digest == data["data"]["valid-words"]
    end
end


text:
## DotProduct is not following the PreMetric rules since d(x, x) != 0 and d(x, y) >= 0 for all x, y
struct DotProduct <: Distances.UnionPreMetric end

@inline function Distances._evaluate(::DotProduct, a::AbstractVector, b::AbstractVector)
    @boundscheck if length(a) != length(b)
        throw(
            DimensionMismatch(
                "first array has length $(length(a)) which does not match the length of the second, $(length(b)).",
            ),
        )
    end
    return dot(a, b)
end

Distances.result_type(::DotProduct, Ta::Type, Tb::Type) = promote_type(Ta, Tb)

@inline Distances.eval_op(::DotProduct, a::Real, b::Real) = a * b
@inline function (dist::DotProduct)(a::AbstractArray, b::AbstractArray)
    return Distances._evaluate(dist, a, b)
end
@inline (dist::DotProduct)(a::Number, b::Number) = a * b


text:
module HMC

using LinearAlgebra: fill!
using Random
using LinearAlgebra
using Printf
using Parameters
using Printf
using Statistics
using Logging

using ..Utilities: get_index, reshaped, δ
using ..Models: AbstractModel, HolsteinModel, SSHModel, update_model!, mulM!, muldMdx!, mulMᵀ!
using ..PhononAction: calc_dSbdx!, calc_Sb
using ..FourierAcceleration: FourierAccelerator, fourier_accelerate!
import ..KPMPreconditioners

export HybridMonteCarlo, update!, refresh_ϕ!, calc_S

mutable struct HybridMonteCarlo{T<:AbstractFloat}

    """
    Number of degrees of freedom to update.
    """
    Ndof::Int

    """
    Dimension of M matrix.
    """
    Ndim::Int

    """
    Store initial phonon fields.
    """
    x0::Vector{T}

    """
    Time between refreshing the momentum p and auxialiary fields ϕ.
    """
    tr::T

    """
    Timestep.
    """
    Δt::T

    """
    Number of time steps.
    """
    Nt::Int

    """
    Smaller timestep used to evolve Sb for multi-timestep algorithm.
    """
    Δt′::T

    """
    Number of steps for Sb per one timestep of Sf for multi-timestep algorithm.
    """
    Nb::Int

    """
    Partial momentum refresh parameter.
    """
    α::T

    """
    Derivative of the action.
    """
    dSdx::Vector{T}

    """
    Velocity.
    """
    v::Vector{T}

    """
    Stores initial velocity vector.
    """
    v0::Vector{T}

    """
    Random gaussian vector.
    """
    R₊::Vector{T}

    """
    Random gaussian vector.
    """
    R₋::Vector{T}

    """
    Auxiliary fields for spin up electrons.
    """
    ϕ₊::Vector{T}

    """
    Auxiliary fields for spin down electrons.
    """
    ϕ₋::Vector{T}

    """
    The vector Λ⋅ϕ₊
    """
    Λϕ₊::Vector{T}

    """
    The vector Λ⋅ϕ₋
    """
    Λϕ₋::Vector{T}

    """
    O⁻¹⋅ϕ₊ where O=MᵀM
    """
    O⁻¹Λϕ₊::Vector{T}

    """
    O⁻¹⋅ϕ₋ where O=MᵀM
    """
    O⁻¹Λϕ₋::Vector{T}

    """
    Diagonal Λ matrix. Definition varies for SSH and Holstein models.
    """
    Λ::Vector{T}

    """
    Temporary storage vector of length Ndim.
    """
    u::Vector{T}

    """
    Temporary storage vector of length Ndof.
    """
    y::Vector{T}

    ######################
    ## Status Variables ##
    ######################

    """
    Whether to log HMC data to file.
    """
    log::Bool

    """
    Whether to log HMC data at every timestep.
    """
    verbose::Bool

    """
    log filename
    """
    logfile::IOStream

    """
    How many HMC updates have been performed.
    """
    updates::Int

    """
    Timestep of current HMC update.
    """
    t::Int

    """
    Whether most recent HMC update was accepted or rejected.
    """
    accepted::Bool

    """
    Total energy.
    """
    H::T

    """
    Action.
    """
    S::T

    """
    Kintetic energy
    """
    K::T

    """
    Iteration Count
    """
    iters::Int

    function HybridMonteCarlo(model::AbstractModel,Δt::T,tr::T,α::T,Nb::Int;
                              log::Bool=false, verbose::Bool=false, logfilename::String="",updates::Int=1) where {T<:AbstractFloat}

        # partial momentum refresh parameter
        @assert 0.0 <= α < 1.0

        Ndof   = model.Ndof
        Ndim   = model.Ndim

        x0     = zeros(T,Ndof)
        dSdx   = zeros(T,Ndof)
        v      = zeros(T,Ndof)
        v0     = zeros(T,Ndof)
        Λ      = ones(T,Ndim)
        R₊     = zeros(T,Ndim)
        R₋     = zeros(T,Ndim)
        ϕ₊     = zeros(T,Ndim)
        Λϕ₊    = zeros(T,Ndim)
        O⁻¹Λϕ₊ = zeros(T,Ndim)
        ϕ₋     = zeros(T,Ndim)
        Λϕ₋    = zeros(T,Ndim)
        O⁻¹Λϕ₋ = zeros(T,Ndim)

        u      = zeros(T,Ndim)
        y      = zeros(T,Ndof)

        # the action
        H = 0.0::T

        # number of timesteps
        Nt = round(Int,tr/Δt)

        # size of smaller timestep for Sb
        Δt′ = Δt/Nb

        # checking conditions on parameters
        @assert 0.0 <= α < 1.0

        t        = 0
        accepted = false
        H        = 0.0
        S        = 0.0
        K        = 0.0
        iters    = 0

        if isfile(logfilename)
            logfile = open(logfilename,"a")
        else
            logfile = open(logfilename,"w")
            write(logfile,"updates accepted timestep tot_energy action kin_energy iters\n")
        end
        if !log
            close(logfile)
        end

        return new{T}(Ndof, Ndim, x0, tr, Δt, Nt, Δt′, Nb, α, dSdx, v, v0, R₊, R₋,
                      ϕ₊, ϕ₋, Λϕ₊, Λϕ₋, O⁻¹Λϕ₊, O⁻¹Λϕ₋, Λ,
                      u, y, log, verbose, logfile, updates, t, accepted, H, S, K, iters)
    end

    function HybridMonteCarlo(hmc::HybridMonteCarlo{T},Δt::T,tr::T,α::T,Nb::Int;
                              log::Bool=false, verbose::Bool=false, logfilename::String="",updates::Int=1) where {T<:AbstractFloat}

        @unpack Ndof, Ndim, x0, H, dSdx, v, v0, R₊, R₋, Λ, ϕ₊, ϕ₋, Λϕ₊, Λϕ₋, O⁻¹Λϕ₊, O⁻¹Λϕ₋, u, y = hmc
        Nt  = round(Int,tr/Δt)
        Δt′ = Δt/Nb

        t        = 0
        accepted = false
        H        = 0.0
        S        = 0.0
        K        = 0.0
        iters    = 0

        if isfile(logfilename)
            logfile = open(logfilename,"a")
        else
            logfile = open(logfilename,"w")
            write(logfile,"updates accepted timestep tot_energy action kin_energy iters\n")
        end
        if !log
            close(logfile)
        end

        return new{T}(Ndof, Ndim, x0, tr, Δt, Nt, Δt′, Nb, α, dSdx, v, v0, R₊, R₋,
                      ϕ₊, ϕ₋, Λϕ₊, Λϕ₋, O⁻¹Λϕ₊, O⁻¹Λϕ₋, Λ,
                      u, y, log, verbose, logfile, updates, t, accepted, H, S, K, iters)
    end
end

"""
Write status of HMC to log file.
"""
function update_log(hmc::HybridMonteCarlo{T},model::AbstractModel{T},fa::FourierAccelerator{T}) where {T<:AbstractFloat}

    @unpack updates, t, accepted, iters, logfile = hmc

    # get energies
    H, S, K = calc_H(hmc, model, fa)

    if t==-1
        # outcome of HMC update accept/reject decision
        outcome = Int(accepted)
    else
        # trajectory ongoing, accept/reject decision comes after HMC trajectory
        outcome = -1
    end

    write(logfile, @sprintf("%d %d %d %.8f %.8f %.8f %d\n", updates, outcome, t, H, S, K, iters))
    flush(logfile)

    return nothing
end

"""
Do a Hybrid Monte Carlo update to the phonon fields.
"""
function update!(model::AbstractModel{T1,T2}, hmc::HybridMonteCarlo{T1}, fa::FourierAccelerator{T1}, preconditioner=I)::Tuple{Bool,T1}  where {T1,T2}

    # only do an update if the model has a non-zero number of degrees of freedom.
    if hmc.Ndof > 0

        # set HMC timestep to zero
        hmc.t = 0

        if hmc.Nb==1
            accepted, iters = standard_update!(model,hmc,fa,preconditioner)
        else
            accepted, iters = multitimestep_update!(model,hmc,fa,preconditioner)
        end

        # write output of logfile
        if hmc.log
            hmc.t = -1
            update_log(hmc,model,fa)
        end

        # increment HMC update counter
        hmc.updates += 1

        return accepted, iters
    else
        return true, 0
    end
end

"""
Standard HMC update.
"""
function standard_update!(model::AbstractModel{T1,T2}, hmc::HybridMonteCarlo{T1}, fa::FourierAccelerator{T1}, preconditioner=I)::Tuple{Bool,T1}  where {T1,T2}

    x     = model.x
    x0    = hmc.x0
    v0    = hmc.v0
    dSdx  = hmc.dSdx
    QdSdx = hmc.dSdx
    v     = hmc.v
    Nt    = hmc.Nt
    Δt    = hmc.Δt

    # keep track for linear solve flag
    flag = 0

    # keep track of iterations
    iters = 0

    # update exp{-Δτ⋅V[x]}
    update_model!(model)

    # refresh the velocity v
    refresh_v!(hmc,model,fa)

    # record intial state
    copyto!(x0,x)
    copyto!(v0,v)

    # refresh ϕ
    refresh_ϕ!(hmc,model)

    # calculate O⁻¹⋅ϕ₊ and O⁻¹⋅ϕ₋
    iters, flag = calc_O⁻¹Λϕ!(hmc,model,preconditioner,2.0)

    if iszero(flag)

        # calculate initial energy
        H₀, S, K = calc_H(hmc, model, fa)

        # calculate the initial dS/dx value
        fill!(dSdx,0.0)
        calc_dSdx!(hmc, model)

        # dS/dx(t+Δt) ==> Q⋅dS/dx(t+Δt)
        fourier_accelerate!(QdSdx,fa,dSdx,-1.0,use_mass=true)

        # log HMC state
        if hmc.log && hmc.verbose
            update_log(hmc,model,fa)
        end

        # iterate over time steps, doing leapfrog updates to the phonon fields
        for hmc.t in 1:Nt

            # v(t+Δt/2) = v(t) - Δt/2⋅Q⋅dS/dx(t)
            @. v = v - Δt/2*QdSdx

            # x(t+Δt) = x(t) + Δt⋅v(t+Δt/2)
            @. x = x + Δt*v

            # update exp{-Δτ⋅V[x]}
            update_model!(model)

            # calculate O⁻¹⋅ϕ₊ and O⁻¹⋅ϕ₋
            itrs, flag = calc_O⁻¹Λϕ!(hmc,model,preconditioner,1.0)
            iters += itrs

            # kill trajectory if error flag
            if flag > 0
                break
            end

            # calculate dS/dx(t+Δt) value
            fill!(dSdx,0.0)
            calc_dSdx!(hmc, model)

            # dS/dx(t+Δt) ==> Q⋅dS/dx(t+Δt)
            fourier_accelerate!(QdSdx,fa,dSdx,-1.0,use_mass=true)

            # v(t+Δt) = v(t+Δt/2) - Δt/2⋅Q⋅dS/dx(t+Δt)
            @. v = v - Δt/2*QdSdx

            # log HMC state
            if hmc.log && hmc.verbose
                update_log(hmc,model,fa)
            end
        end
    end

    # calcualte acceptance probability
    P = 0.0
    if iszero(flag)

        # calculate O⁻¹⋅ϕ₊ and O⁻¹⋅ϕ₋
        itrs, flag = calc_O⁻¹Λϕ!(hmc,model,preconditioner,2.0)
        iters += itrs

        if iszero(flag)

            # calculate final energy
            H₁, S, K = calc_H(hmc, model, fa)

            # calculate change in energy
            ΔH = H₁ - H₀

            # calculate probability of acceptance
            P = min(1.0, exp(-ΔH))
        end
    end

    # Metropolis-Hasting Accept/Reject Step
    if rand(model.rng) < P && iszero(flag) # if accepted

        hmc.accepted = true
        return hmc.accepted, T1(cld(iters,Nt+2))

    else # if rejected

        # reset to original phonon field
        copyto!(x,x0)

        # reset to reflected original velocities.
        # note: this does not do anything unless doing partial momentum refreshes.
        @. v = -v0

        # update exp{-Δτ⋅V[x]}
        update_model!(model)

        hmc.accepted = false
        return hmc.accepted, T1(cld(iters,Nt+2))
    end
end

"""
Multi-timestepping HMC update.
"""
function multitimestep_update!(model::AbstractModel{T1,T2}, hmc::HybridMonteCarlo{T1}, fa::FourierAccelerator{T1}, preconditioner=I)::Tuple{Bool,T1}  where {T1,T2}

    x      = model.x
    x0     = hmc.x0
    v0     = hmc.v0
    dSbdx  = hmc.dSdx
    QdSbdx = hmc.dSdx
    dSfdx  = hmc.dSdx
    QdSfdx = hmc.dSdx
    v      = hmc.v
    Nt     = hmc.Nt
    Δt     = hmc.Δt
    Nb     = hmc.Nb
    Δt′    = hmc.Δt′

    # keep track for linear solve flag
    flag = 0

    # keep track of iterations
    iters = 0

    # update exp{-Δτ⋅V[x]}
    update_model!(model)

    # refresh the velocity v
    refresh_v!(hmc,model,fa)

    # record intial phonon configuration
    copyto!(x0,x)
    copyto!(v0,v)

    # refresh ϕ
    refresh_ϕ!(hmc,model)

    # calculate O⁻¹⋅ϕ₊ and O⁻¹⋅ϕ₋
    itrs, flag = calc_O⁻¹Λϕ!(hmc,model,preconditioner,2.0)
    iters     += iters

    if iszero(flag)

        # calculate energy
        H₀, S, K = calc_H(hmc, model, fa)

        # calculate the initial dSf/dx value
        fill!(dSfdx,0.0)
        calc_dSfdx!(hmc, model)

        # dSf/dx(t+Δt) ==> Q⋅dSf/dx(t+Δt)
        fourier_accelerate!(QdSfdx,fa,dSfdx,-1.0,use_mass=true)

        # log HMC state
        if hmc.log && hmc.verbose
            update_log(hmc,model,fa)
        end

        # iterate over timesteps
        for hmc.t in 1:Nt

            # v(t+Δt/2) = v(t) - Δt/2⋅Q⋅dSf/dx(t)
            @. v = v - Δt/2*QdSfdx

            # calculate the initial dSb/dx value
            fill!(dSbdx,0.0)
            calc_dSbdx!(dSbdx,model)

            # dSb/dx(t+Δt) ==> Q⋅dSb/dx(t+Δt)
            fourier_accelerate!(QdSbdx,fa,dSbdx,-1.0,use_mass=true)

            # evolve Sb using Nb smaller timesteps of size Δt′=Δt/Nb
            for t′ in 1:Nb

                # v(t+Δt/2) = v(t) - Δt′/2⋅Q⋅dSb/dx(t)
                @. v = v - Δt′/2*QdSbdx

                # x(t+Δt) = x(t) + Δt⋅v(t+Δt/2)
                @. x = x + Δt′*v

                # calculate dSb/dx(t+Δt) value
                fill!(dSbdx,0.0)
                calc_dSbdx!(dSbdx,model)

                # dS/dx(t+Δt) ==> Q⋅dS/dx(t+Δt)
                fourier_accelerate!(QdSbdx,fa,dSbdx,-1.0,use_mass=true)

                # v(t+Δt) = v(t+Δt/2) - Δt/2⋅Q⋅dSb/dx(t+Δt)
                @. v = v - Δt′/2*QdSbdx
            end

            # update exp{-Δτ⋅V[x]}
            update_model!(model)

            # calculate O⁻¹⋅ϕ₊ and O⁻¹⋅ϕ₋
            itrs, flag = calc_O⁻¹Λϕ!(hmc,model,preconditioner,1.0)
            iters += itrs

            # kill trajectory if error occurred
            if flag > 0
                break
            end

            # calculate dSf/dx(t+Δt) value
            fill!(dSfdx,0.0)
            calc_dSfdx!(hmc, model)

            # dSf/dx(t+Δt) ==> Q⋅dSf/dx(t+Δt)
            fourier_accelerate!(QdSfdx,fa,dSfdx,-1.0,use_mass=true)

            # v(t+Δt) = v(t) - Δt/2⋅Q⋅dSf/dx(t)
            @. v = v - Δt/2*QdSfdx

            # log HMC state
            if hmc.log && hmc.verbose
                update_log(hmc,model,fa)
            end
        end
    end

    # calcualte acceptance probability
    P = 0.0
    if iszero(flag)

        # calculate O⁻¹⋅ϕ₊ and O⁻¹⋅ϕ₋
        itrs, flag = calc_O⁻¹Λϕ!(hmc,model,preconditioner,2.0)
        iters += itrs

        if iszero(flag)

            # calculate final energy
            H₁, S, K = calc_H(hmc,model,fa)

            # calculate change in energy
            ΔH = H₁ - H₀

            # calculate probability of acceptance
            P = min(1.0, exp(-ΔH))
        end
    end

    # Metropolis-Hasting Accept/Reject Step
    if rand(model.rng) < P && iszero(flag) # if accepted

        hmc.accepted = true
        return hmc.accepted, T1(cld(iters,Nt+2))

    else # if rejected

        # reset to original phonon field
        copyto!(x,x0)

        # reset to reflected original velocities.
        # note: this does not do anything unless doing partial momentum refreshes.
        @. v = -v0

        # update exp{-Δτ⋅V[x]}
        update_model!(model)

        hmc.accepted = false
        return hmc.accepted, T1(cld(iters,Nt+2))
    end
end

"""
The kinetics energy in our dyanics is given by `K=v⋅Q⁻¹⋅v/2=v⋅M⋅v/2`, so therefore it is possible to
refresh the velocity `v` according `v=√(Q)⋅R` where `R` is a vector of normal random numbers
and `Q` is the fourier acceleration matrix and `M` is the dynamical mass matrix.
More specifically, this function supports partial momentum refreshes of the
form `v = α⋅v + √(1-α²)⋅v′` where `v′=√(Q)⋅R=√(M⁻¹)⋅R`⋅
"""
function refresh_v!(hmc::HybridMonteCarlo{T},model::AbstractModel{T},fa::FourierAccelerator{T}) where {T<:AbstractFloat}

    R       = hmc.y
    sqrtQR  = hmc.y
    v       = hmc.v
    α       = hmc.α

    randn!(R,model)
    fourier_accelerate!(sqrtQR,fa,R,-0.5,use_mass=true)
    @. v = α*v + sqrt(1.0-α^2)*sqrtQR

    return nothing
end

"""
Refresh `ϕ` according to the relationship `ϕ ~ Λ⁻¹⋅Mᵀ⋅R` where `R` is a vector of normal random numbers.
"""
function refresh_ϕ!(hmc::HybridMonteCarlo{T},model::AbstractModel{T};sample_R::Bool=true)::T where {T}

    @unpack R₊, R₋, ϕ₊, ϕ₋, Λϕ₊, Λϕ₋, Λ = hmc

    # update Λ
    update_Λ!(hmc,model)

    # sample new random vectors
    if sample_R
        randn!(model.rng,R₊)
        randn!(model.rng,R₋)
    end

    # ϕ₊ = Λ⁻¹⋅Mᵀ⋅R₊
    mulMᵀ!(Λϕ₊,model,R₊)
    mulΛ⁻¹!(ϕ₊,Λϕ₊,hmc,model)

    # ϕ₋ = Λ⁻¹⋅Mᵀ⋅R₋
    mulMᵀ!(Λϕ₋,model,R₋)
    mulΛ⁻¹!(ϕ₋,Λϕ₋,hmc,model)

    # calculate total action
    hmc.S  = dot(R₊,R₊)/2 + dot(R₋,R₋)/2
    hmc.S += calc_Sb(model)

    return hmc.S
end

"""
Calculate the total energy `H = K + S`.
"""
function calc_H(hmc::HybridMonteCarlo{T1}, model::AbstractModel{T1,T2}, fa::FourierAccelerator{T1}) where {T1,T2}

    S = calc_S(hmc, model)
    K = calc_K(hmc, model, fa)
    hmc.H = S + K

    return hmc.H, S, K
end

"""
Calculate the kintetic energy `K = v⋅Q⁻¹⋅v/2` where `Q` is the acceleration matrix.
"""
function calc_K(hmc::HybridMonteCarlo{T1}, model::HolsteinModel{T1,T2,T3}, fa::FourierAccelerator{T1})::T1 where {T1,T2,T3}

    v    = hmc.v
    Q⁻¹v = hmc.y
    fourier_accelerate!(Q⁻¹v,fa,v,1.0,use_mass=true)
    hmc.K = dot(v,Q⁻¹v)/2

    return hmc.K
end

function calc_K(hmc::HybridMonteCarlo{T1}, model::SSHModel{T1,T2,T3}, fa::FourierAccelerator{T1})::T1 where {T1,T2,T3}

    v  = hmc.v
    mv = hmc.y
    fourier_accelerate!(mv,fa,v,1.0,use_mass=true)

    hmc.K = 0.0

    # iterate over fields
    for field in 1:hmc.Ndof
        # if a primary field
        if model.primary_field[field]==field
            # increment total kinetic energy accordingly
            hmc.K += v[field]*mv[field]/2
        end
    end

    return hmc.K
end

"""
Calcualte the action S = Sb + ϕ₊ᵀ⋅O⁻¹⋅ϕ₊/2 + ϕ₋ᵀ⋅O⁻¹⋅ϕ/2
"""
function calc_S(hmc::HybridMonteCarlo{T1}, model::AbstractModel{T1,T2})::T1 where {T1,T2}

    # S = ϕ₊ᵀ⋅O⁻¹⋅ϕ₊/2 + ϕ₋ᵀ⋅O⁻¹⋅ϕ₋/2
    hmc.S  = calc_Sf(hmc,model)

    # S = Sb + ϕ₊ᵀ⋅O⁻¹⋅ϕ₊/2 + ϕ₋ᵀ⋅O⁻¹⋅ϕ₋/2
    hmc.S += calc_Sb(model)

    return hmc.S
end

"""
Calculate the derivative of the action dS/dx = dSb/dx - ϕ₊ᵀ⋅O⁻ᵀ⋅[Mᵀ⋅dM/dx]⋅O⁻¹⋅ϕ₊ - ϕ₋ᵀ⋅O⁻ᵀ⋅[Mᵀ⋅dM/dx]⋅O⁻¹⋅ϕ₋
"""
function calc_dSdx!(hmc::HybridMonteCarlo{T1}, model::AbstractModel{T1,T2}) where {T1,T2}

    # dS/dx = -ϕ₊ᵀ⋅O⁻ᵀ⋅[dMᵀ/dx⋅M]⋅O⁻¹⋅ϕ₊ - ϕ₋ᵀ⋅O⁻ᵀ⋅[dMᵀ/dx⋅M]⋅O⁻¹⋅ϕ₋
    calc_dSfdx!(hmc, model)

    # dS/dx = dSb/dx - ϕ₊ᵀ⋅O⁻ᵀ⋅[dMᵀ/dx⋅M]⋅O⁻¹⋅ϕ₊ - ϕ₋ᵀ⋅O⁻ᵀ⋅[dMᵀ/dx⋅M]⋅O⁻¹⋅ϕ₋
    calc_dSbdx!(hmc.dSdx, model)

    return nothing
end

"""
Calculate the fermionic action S = [ϕ₊ᵀ⋅Λ⋅O⁻¹⋅Λ⋅ϕ₊ + ϕ₋ᵀ⋅Λ⋅O⁻¹⋅Λ⋅ϕ₋]/2 = [ϕ₊ᵀ⋅Λ⋅[Mᵀ⋅M]⁻¹⋅Λ⋅ϕ₊ + ϕ₋ᵀ⋅Λ⋅[Mᵀ⋅M]⁻¹⋅Λ⋅ϕ₋]/2
"""
function calc_Sf(hmc::HybridMonteCarlo{T},model::AbstractModel{T})::T where {T<:AbstractFloat}

    @unpack Λϕ₊, Λϕ₋, O⁻¹Λϕ₊, O⁻¹Λϕ₋ = hmc

    Sf  = dot(Λϕ₊,O⁻¹Λϕ₊)/2
    Sf += dot(Λϕ₋,O⁻¹Λϕ₋)/2

    return Sf
end

"""
Calculate the derivative of the fermionic action.
Each partial derivative `∂S/∂xᵢ(τ)` will be stored to the corresponding element in the array dSdx.
"""
function calc_dSfdx!(hmc::HybridMonteCarlo{T1}, model::AbstractModel{T1,T2}) where {T1,T2}

    @unpack O⁻¹Λϕ₊, O⁻¹Λϕ₋, Λϕ₊, Λϕ₋, ϕ₊, ϕ₋ = hmc
    dSfdx   = hmc.dSdx
    dMdx    = hmc.y
    MO⁻¹Λϕ  = hmc.u

    # dSf/dx += -[M⋅O⁻¹⋅Λ⋅ϕ₊]ᵀ⋅dM/dx⋅[O⁻¹⋅Λ⋅ϕ₊]
    mulM!(MO⁻¹Λϕ,model,O⁻¹Λϕ₊)
    muldMdx!(dMdx,MO⁻¹Λϕ,model,O⁻¹Λϕ₊)
    @. dSfdx += -dMdx

    # dSf/dx += -[M⋅O⁻¹⋅Λ⋅ϕ₋]ᵀ⋅dM/dx⋅[O⁻¹⋅Λ⋅ϕ₋]
    mulM!(MO⁻¹Λϕ,model,O⁻¹Λϕ₋)
    muldMdx!(dMdx,MO⁻¹Λϕ,model,O⁻¹Λϕ₋)
    @. dSfdx += -dMdx

    # dSf/dx += [ϕ₊]ᵀ⋅dΛ/dx⋅[O⁻¹⋅Λ⋅ϕ₊]
    muldΛdx!(dSfdx,ϕ₊,O⁻¹Λϕ₊,hmc,model)

    # dSf/dx += [ϕ₋]ᵀ⋅dΛ/dx⋅[O⁻¹⋅Λ⋅ϕ₋]
    muldΛdx!(dSfdx,ϕ₋,O⁻¹Λϕ₋,hmc,model)

    return nothing
end

"""
Solve `O⋅x=Λ⋅ϕ₊ ==> x=O⁻¹⋅Λ⋅ϕ₊` and `O⋅x=Λ⋅ϕ₋ ==> x=O⁻¹⋅Λ⋅ϕ₋` where `O = Mᵀ⋅M`.
"""
function calc_O⁻¹Λϕ!(hmc::HybridMonteCarlo{T1}, model::AbstractModel{T1,T2}, preconditioner=I, power::T1=1.0)::Tuple{Int,Int} where {T1,T2}

    @unpack Λ, ϕ₊, ϕ₋, Λϕ₊, Λϕ₋, O⁻¹Λϕ₊, O⁻¹Λϕ₋ = hmc
    M⁻ᵀΛϕ₊ = hmc.u
    M⁻ᵀΛϕ₋ = hmc.u

    # udpate tolerance used by solver
    tol = model.solver.tol::T1
    model.solver.tol = tol^power

    # count iterative solver iterations
    hmc.iters = 0

    # setup the precontioer
    KPMPreconditioners.setup!(preconditioner)

    #############################
    ## CALCULATE Λ⋅ϕ₊ AND Λ⋅ϕ₋ ##
    #############################

    update_Λ!(hmc,model)
    mulΛ!(Λϕ₊,ϕ₊,hmc,model)
    mulΛ!(Λϕ₋,ϕ₋,hmc,model)

    #########################
    ## CALCULATE  O⁻¹⋅Λ⋅ϕ₊ ##
    #########################

    # linear solve status flag
    flag = 0

    if !model.mul_by_M # if using Conjugate Gradient

        # solve linear system
        fill!(O⁻¹Λϕ₊,0.0)
        model.transposed = false
        iters, err, flag = ldiv!(O⁻¹Λϕ₊,model,Λϕ₊,preconditioner)
        hmc.iters += iters

    else # if using GMRES

        # solve linear system
        fill!(M⁻ᵀΛϕ₊,0.0)
        model.transposed  = true
        iters, err, flag1 = ldiv!(M⁻ᵀΛϕ₊,model,Λϕ₊,preconditioner)
        hmc.iters        += iters
        flag              = max(flag,flag1)

        # solve linear system
        fill!(O⁻¹Λϕ₊,0.0)
        model.transposed  = false
        iters, err, flag2 = ldiv!(O⁻¹Λϕ₊,model,M⁻ᵀΛϕ₊,preconditioner)
        hmc.iters        += iters
        flag              = max(flag,flag2)
    end

    #########################
    ## CALCULATE  O⁻¹⋅Λ⋅ϕ₋ ##
    #########################

    if !model.mul_by_M && iszero(flag)

        # solve linear system
        fill!(O⁻¹Λϕ₋,0.0)
        model.transposed =false
        iters, err, flag = ldiv!(O⁻¹Λϕ₋,model,Λϕ₋,preconditioner)
        hmc.iters += iters

    elseif iszero(flag)

        # solve linear system
        fill!(M⁻ᵀΛϕ₋,0.0)
        model.transposed  = true
        iters, err, flag1 = ldiv!(M⁻ᵀΛϕ₋,model,Λϕ₋,preconditioner)
        hmc.iters        += iters
        flag              = max(flag,flag1)

        # solve linear system
        fill!(O⁻¹Λϕ₋,0.0)
        model.transposed  = false
        iters, err, flag2 = ldiv!(O⁻¹Λϕ₋,model,M⁻ᵀΛϕ₋,preconditioner)
        hmc.iters        += iters
        flag              = max(flag,flag2)
    end

    # accounting for the fact that there is a spin up and spin down
    # linear system that needs to be solved
    if iszero(flag)
        hmc.iters = cld(hmc.iters,2)
    end

    # revert solvers tolerance to original value
    model.solver.tol = tol

    return hmc.iters, flag
end

"""
Construct the Λ matrix.
"""
function update_Λ!(hmc::HybridMonteCarlo{T1}, model::HolsteinModel{T1,T2}) where {T1,T2}

    @unpack λ, λ₂, Δτ, Lτ, Nph = model
    x = reshaped(model.x, Lτ, Nph)
    Λ = reshaped(hmc.Λ,   Lτ, Nph)

    @fastmath @inbounds for i in 1:Nph

        xᵢ  = @view x[:,i]
        Λᵢ  = @view Λ[:,i]
        λᵢ  = λ[i]
        λ₂ᵢ = λ₂[i]

        for τ in 1:Lτ
            xᵢτ   = xᵢ[τ]
            Λᵢ[τ] = exp(-Δτ*(λᵢ*xᵢτ + λ₂ᵢ*xᵢτ^2)/2)
        end
    end

    return nothing
end

function update_Λ!(hmc::HybridMonteCarlo{T1}, model::SSHModel{T1,T2}) where {T1,T2}

    return nothing
end

"""
Calculate v′=Λ⋅v.
"""
function mulΛ!(v′::AbstractVector{T}, v::AbstractVector{T}, hmc::HybridMonteCarlo{T}, model::HolsteinModel{T}) where {T}

    @unpack x, λ, Δτ, Lτ, Nph, Nsites = model

    u  = reshaped(v ,Lτ,Nsites)
    u′ = reshaped(v′,Lτ,Nsites)
    Λ  = reshaped(hmc.Λ,Lτ,Nsites)

    @fastmath @inbounds for i in 1:Nsites
        u1i = u[1,i]
        for τ in 1:Lτ-1
            u′[τ,i] = -Λ[τ+1,i] * u[τ+1,i]
        end
        u′[Lτ,i] = Λ[1,i] * u1i
    end

    return nothing
end

function mulΛ!(v′::AbstractVector{T}, v::AbstractVector{T}, hmc::HybridMonteCarlo{T}, model::AbstractModel{T}) where {T}

    return nothing
end

"""
Calculate v′=Λ⋅v.
"""
function mulΛ⁻¹!(v′::AbstractVector{T}, v::AbstractVector{T}, hmc::HybridMonteCarlo{T}, model::HolsteinModel{T}) where {T}

    @unpack x, λ, Δτ, Lτ, Nph, Nsites = model

    u  = reshaped(v ,Lτ,Nsites)
    u′ = reshaped(v′,Lτ,Nsites)
    Λ  = reshaped(hmc.Λ,Lτ,Nsites)

    @fastmath @inbounds for i in 1:Nsites
        uLi = u[Lτ,i]
        for τ in 2:Lτ
            u′[τ,i] = -inv(Λ[τ,i]) * u[τ-1,i]
        end
        u′[1,i] = inv(Λ[1,i]) * uLi
    end

    return nothing
end

function mulΛ⁻¹!(v′::AbstractVector{T}, v::AbstractVector{T}, hmc::HybridMonteCarlo{T}, model::AbstractModel{T}) where {T}

    return nothing
end

"""
Calculate ⟨vₗ|∂Λ/∂x(τ)|vᵣ⟩ for all τ, adding each result to the corresponding element in the vector dΛdx.
"""
function muldΛdx!(dΛdx::Vector{T1},vₗ::Vector{T1},vᵣ::Vector{T1},hmc::HybridMonteCarlo{T1},model::HolsteinModel{T1,T2}) where {T1,T2}

    @unpack x, λ, λ₂, Δτ, Lτ, Nph = model
    @unpack Λ = hmc

    @fastmath @inbounds for i in 1:Nph
        λᵢ       = λ[i]
        λ₂ᵢ      = λ₂[i]
        n′       = get_index(Lτ,i,Lτ)
        n        = get_index(1,i,Lτ)
        dΛdx[n] += vₗ[n] * (-Δτ*(λᵢ/2+λ₂ᵢ*x[n]))*Λ[n] * vᵣ[n′]
        for τ in 2:Lτ
            τm1      = mod1(τ-1,Lτ)
            n′       = get_index(τm1,i,Lτ)
            n        = get_index(τ,i,Lτ)
            dΛdx[n] += vₗ[n] * (Δτ*(λᵢ/2+λ₂ᵢ*x[n]))*Λ[n] * vᵣ[n′]
        end
    end

    return nothing
end

function muldΛdx!(dΛdx::Vector{T1},vₗ::Vector{T1},vᵣ::Vector{T1},hmc::HybridMonteCarlo{T1}, model::AbstractModel{T1,T2}) where {T1,T2}

    return nothing
end

end

max_stars_repo_name: cohensbw/Langevin

max_issues_repo_name: cohensbw/ElPhDynamics

max_forks_repo_name: cohensbw/ElPhDynamics


text:
master_formulation(rd::RoutingData, ::Any) =
    error("Not implemented: $(rd.model_type)")

# TODO: like _iter, add solve_master_problem that in turn calls this.

function get_traffic_matrices(rm::RoutingModel)
    rd = rm.data
    dual_constraints = rm.constraints_uncertainty_set

    l_tm = Set{Dict{Edge{Int}, Float64}}()
    for e in edges(rd)
        push!(l_tm, Dict{Edge{Int}, Float64}(k => dual(dual_constraints[e][k]) for k in keys(dual_constraints[e])))
    end
    return collect(l_tm)
end

function compute_routing(rd::RoutingData, ::Load, ::MinimumMaximum, formulation::FormulationType, ::Val{false}, ::DualReformulation, ::ObliviousUncertainty, ::UncertainDemand)
    start = time_ns()
    rm = master_formulation(rd, formulation)
    time_create_master_model_ms = (time_ns() - start) / 1_000_000.

    # Solve the problem and generate the output data structure.
    start = time_ns()
    optimize!(rm.model)
    rd.logmessage(objective_value(rm.model))
    status = termination_status(rm.model)

    # Export if needed.
    _export_lp_if_allowed(rd, rm.model, "lp_master")
    _export_lp_if_failed(rd, status, rm.model, "error_master", "Subproblem could not be solved!")

    # Retrieve the traffic matrices, if asked.
    matrices = if rd.model_robust_reformulation_traffic_matrices
        get_traffic_matrices(rm)
    else
        Dict{Edge{Int}, Float64}[]
    end

    stop = time_ns()

    # TODO: Export things like the normal case.

    return RoutingSolution(rd,
                           time_precompute_ms=rd.time_precompute_ms,
                           time_create_master_model_ms=time_create_master_model_ms,
                           time_solve_ms=(stop - start) / 1_000_000.,
                           objectives=[objective_value(rm.model)],
                           matrices=matrices,
                           routings=Routing(rd, value.(rm.routing)),
                           master_model=rm)
end


text:
# Autogenerated wrapper script for psurface_jll for armv7l-linux-gnueabihf-cxx03
export libpsurface, psurface_convert, psurface_simplify, psurface_smooth

JLLWrappers.@generate_wrapper_header("psurface")
JLLWrappers.@declare_library_product(libpsurface, "libpsurface.so.0")
JLLWrappers.@declare_executable_product(psurface_convert)
JLLWrappers.@declare_executable_product(psurface_simplify)
JLLWrappers.@declare_executable_product(psurface_smooth)
function __init__()
    JLLWrappers.@generate_init_header()
    JLLWrappers.@init_library_product(
        libpsurface,
        "lib/libpsurface.so",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_executable_product(
        psurface_convert,
        "bin/psurface-convert",
    )

    JLLWrappers.@init_executable_product(
        psurface_simplify,
        "bin/psurface-simplify",
    )

    JLLWrappers.@init_executable_product(
        psurface_smooth,
        "bin/psurface-smooth",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
# PREAMBLE

# PKG_SETUP

# ## Setup

using LinearAlgebra
using Plots
using DirectTrajectoryOptimization

# ## horizon
T = 101

# ## acrobot
num_state = 4
num_action = 1
num_parameter = 0

function acrobot(x, u, w)
    mass1 = 1.0
    inertia1 = 0.33
    length1 = 1.0
    lengthcom1 = 0.5

    mass2 = 1.0
    inertia2 = 0.33
    length2 = 1.0
    lengthcom2 = 0.5

    gravity = 9.81
    friction1 = 0.1
    friction2 = 0.1

    function M(x, w)
        a = (inertia1 + inertia2 + mass2 * length1 * length1
            + 2.0 * mass2 * length1 * lengthcom2 * cos(x[2]))

        b = inertia2 + mass2 * length1 * lengthcom2 * cos(x[2])

        c = inertia2

       return [a b; b c]
    end

    function Minv(x, w)
        a = (inertia1 + inertia2 + mass2 * length1 * length1
            + 2.0 * mass2 * length1 * lengthcom2 * cos(x[2]))

        b = inertia2 + mass2 * length1 * lengthcom2 * cos(x[2])

        c = inertia2

        return 1.0 / (a * c - b * b) * [c -b; -b a]
    end

    function τ(x, w)
        a = (-1.0 * mass1 * gravity * lengthcom1 * sin(x[1])
            - mass2 * gravity * (length1 * sin(x[1])
            + lengthcom2 * sin(x[1] + x[2])))

        b = -1.0 * mass2 * gravity * lengthcom2 * sin(x[1] + x[2])

        return [a; b]
    end

    function C(x, w)
        a = -2.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
        b = -1.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
        c = mass2 * length1 * lengthcom2 * sin(x[2]) * x[3]
        d = 0.0

        return [a b; c d]
    end

    function B(x, w)
        [0.0; 1.0]
    end

    q = view(x, 1:2)
    v = view(x, 3:4)

    qdd = Minv(q, w) * (-1.0 * C(x, w) * v
            + τ(q, w) + B(q, w) * u[1] - [friction1; friction2] .* v)

    return [x[3]; x[4]; qdd[1]; qdd[2]]
end

function midpoint_implicit(y, x, u, w)
    h = 0.05 # timestep
    y - (x + h * acrobot(0.5 * (x + y), u, w))
end

# ## model
dt = Dynamics(midpoint_implicit, num_state, num_state, num_action, num_parameter=num_parameter, evaluate_hessian=true)
dynamics = [dt for t = 1:T-1]

# ## initialization
x1 = [0.0; 0.0; 0.0; 0.0]
xT = [π; 0.0; 0.0; 0.0]

# ## objective
ot = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4]) + 0.1 * dot(u, u)
oT = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4])
ct = Cost(ot, num_state, num_action, num_parameter=num_parameter)
cT = Cost(oT, num_state, 0, num_parameter=num_parameter)
objective = [[ct for t = 1:T-1]..., cT]

# ## constraints
bnd1 = Bound(num_state, num_action)
bndt = Bound(num_state, num_action)
bndT = Bound(num_state, 0)
bounds = [bnd1, [bndt for t = 2:T-1]..., bndT]

constraints = [
        Constraint((x, u, w) -> x - x1, num_state, num_action),
        [Constraint() for t = 2:T-1]...,
        Constraint((x, u, w) -> x - xT, num_state, 0)
       ]

# ## problem
p = Solver(dynamics, objective, constraints, bounds,
    options=Options{Float64}())

# ## initialize
x_interpolation = linear_interpolation(x1, xT, T)
u_guess = [1.0 * randn(num_action) for t = 1:T-1]

initialize_states!(p, x_interpolation)
initialize_controls!(p, u_guess)

# ## solve
@time solve!(p)

# ## solution
x_sol, u_sol = get_trajectory(p)

@show x_sol[1]
@show x_sol[T]

# ## state
plot(hcat(x_sol...)')

# ## control
plot(hcat(u_sol[1:end-1]..., u_sol[end-1])', linetype = :steppost)


text:
# Experiment 4.1 from:
# C.-H. Chen, L.H. Lee: Stochastic Simulation Optimization, World Scientific, 2011

using Random
using EventSimulation
using Distributions

mutable struct SeqQueue <: AbstractState
    q1::SimResource{Int}
    q2::SimResource{Int}
    arrive::Vector{Float64}
    depart::Vector{Float64}
end

arrival(s) = (provide!(s, s.state.q1, 1); push!(s.state.arrive, s.now))
stage1(s) = register!(s, leave1, rand(Uniform(1, 39)))
leave1(s) = (provide!(s, s.state.q2, 1); request!(s, s.state.q1, 1, stage1))
stage2(s) = register!(s, leave2, rand(Uniform(5, 45)))
leave2(s) = (push!(s.state.depart, s.now); request!(s, s.state.q2, 1, stage2))

function exec(n1, n2)
    state = SeqQueue(SimResource{Int}(), SimResource{Int}(), Float64[], Float64[])
    s = Scheduler(state)
    repeat_register!(s, arrival, (x -> length(s.state.arrive) < 100 ? randexp() : nothing))
    foreach(i -> request!(s, state.q1, 1, stage1), 1:n1)
    foreach(i -> request!(s, state.q2, 1, stage2), 1:n2)
    go!(s)
    @assert length(state.depart) == length(state.arrive)
    return mean(state.depart) - mean(state.arrive)
end

for i in 1:10
    v = mean(exec(10+i, 21-i) for j in 1:2048)
    println((10+i, 21-i),":\t", v)
end


text:
export ConstantHomotopy

"""
    ConstantHomotopy(F::AbstractSystem)

Intepret the system `F` as the (constant) homotopy ``H(x,t)=F(x)``.
"""
struct ConstantHomotopy{S<:AbstractSystem} <: AbstractHomotopy
    system::S
end

Base.size(H::ConstantHomotopy) = size(H.system)

struct ConstantHomotopyCache{HC<:AbstractSystemCache} <: AbstractHomotopyCache
    cache::HC
end

cache(H::ConstantHomotopy, x, t) = ConstantHomotopyCache(cache(H.system, x))

const CH = ConstantHomotopy
const CHC = ConstantHomotopyCache

@propagate_inbounds evaluate!(u, H::CH, x, t, c::CHC) = evaluate!(u, H.system, x, c.cache)
@propagate_inbounds jacobian!(U, H::CH, x, t, c::CHC) = jacobian!(U, H.system, x, c.cache)
@propagate_inbounds function dt!(u, H::CH, x, t, c::CHC)
    u .= zero(eltype(u))
    u
end
@propagate_inbounds function evaluate_and_jacobian!(u, U, H::CH, x, t, c::CHC)
    evaluate_and_jacobian!(u, U, H.system, x, c.cache)
end
@propagate_inbounds function jacobian_and_dt!(U, u, H::CH, x, t, c::CHC)
    jacobian!(U, H.system, x, c.cache)
    u .= zero(eltype(u))
    nothing
end
evaluate(H::CH, x, t, c::CHC) = evaluate(H.system, x, c.cache)
jacobian(H::CH, x, t, c::CHC) = jacobian(H.system, x, c.cache)
function dt(H::CH, x, t, c::CHC)
    u = evaluate(H.system, x, c.cache)
    u .= zero(eltype(u))
    u
end


text:
#-------------------------------------------------------------------
#* EMSO Model Library (EML) Copyright (C) 2004 - 2007 ALSOC.
#*
#* This LIBRARY is free software; you can distribute it and/or modify
#* it under the therms of the ALSOC FREE LICENSE as available at
#* http://www.enq.ufrgs.br/alsoc.
#*
#* EMSO Copyright (C) 2004 - 2007 ALSOC, original code
#* from http://www.rps.eng.br Copyright (C) 2002-2004.
#* All rights reserved.
#*
#* EMSO is distributed under the therms of the ALSOC LICENSE as
#* available at http://www.enq.ufrgs.br/alsoc.
#*--------------------------------------------------------------------
#* Author: Gerson Balbueno Bicca
#* $Id: PHE.mso 250 2007-04-27 16:32:02Z bicca $
#*------------------------------------------------------------------
type Thermal_PHE
	Thermal_PHE()=begin
		new(
			positive(Dict{Symbol,Any}(
				:Brief=>"Heat Capacity Ratio",
				:Default=>0.5,
				:Lower=>1E-6
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Minimum Heat Capacity",
				:Lower=>1E-10,
				:Default=>1E3,
				:Unit=>"W/K"
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Maximum Heat Capacity",
				:Lower=>1E-10,
				:Default=>1E3,
				:Unit=>"W/K"
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Number of Units Transference",
				:Default=>0.05,
				:Lower=>1E-10
			)),
			positive(Dict{Symbol,Any}(
				:Brief=>"Effectiveness",
				:Default=>0.5,
				:Lower=>0.1,
				:Upper=>1.1,
				:Symbol=>"\\varepsilon"
			)),
			power(Dict{Symbol,Any}(
				:Brief=>"Heat Transfer",
				:Default=>7000,
				:Lower=>1E-6,
				:Upper=>1E10
			)),
			heat_trans_coeff(Dict{Symbol,Any}(
				:Brief=>"Overall Heat Transfer Coefficient Clean",
				:Default=>1,
				:Lower=>1E-6,
				:Upper=>1E10
			)),
			heat_trans_coeff(Dict{Symbol,Any}(
				:Brief=>"Overall Heat Transfer Coefficient Dirty",
				:Default=>1,
				:Lower=>1E-6,
				:Upper=>1E10
			)),
			[:Cr,:Cmin,:Cmax,:NTU,:Eft,:Q,:Uc,:Ud,]
		)
	end
	Cr::positive
	Cmin::positive
	Cmax::positive
	NTU::positive
	Eft::positive
	Q::power
	Uc::heat_trans_coeff
	Ud::heat_trans_coeff
	variables::Array{Symbol,1}
	attributes::Dict{Symbol,Any}
end
export Thermal_PHE
function atributes(in::Thermal_PHE,_::Dict{Symbol,Any})
	fields::Dict{Symbol,Any}=Dict{Symbol,Any}()
	fields[:Pallete]=false
	fields[:Brief]="to be documented"
	fields[:Info]="to be documented"
	drive!(fields,_)
	return fields
end
Thermal_PHE(_::Dict{Symbol,Any})=begin
	newModel=Thermal_PHE()
	newModel.attributes=atributes(newModel,_)
	newModel
end


text:
"""
A StateSpace model with a partioning imposed according to

A  | B1  B2
——— ————————
C1 | D11 D12
C2 | D21 D22

It corresponds to partioned input and output signals
u = [u1 u2]^T
y = [y1 y2]^T

"""
struct PartionedStateSpace{S<:AbstractStateSpace} <: LTISystem
    P::S
    nu1::Int
    ny1::Int
end
# For converting between different S
PartionedStateSpace{S}(partsys::PartionedStateSpace) where {S<:StateSpace} =
    PartionedStateSpace{S}(S(partsys.P), partsys.nu1, partsys.ny1)

function getproperty(sys::PartionedStateSpace, d::Symbol)
    P = getfield(sys, :P)
    nu1 = getfield(sys, :nu1)
    ny1 = getfield(sys, :ny1)

    if d === :Ts
        return P.Ts # Will throw deprecation until removed # DEPRECATED
    elseif d === :P
        return P
    elseif d === :nu1
        return nu1
    elseif d === :ny1
        return ny1
    elseif d === :A
        return P.A
    elseif d === :B1
        return P.B[:, 1:nu1]
    elseif d === :B2
        return P.B[:, nu1+1:end]
    elseif d === :C1
        return P.C[1:ny1, :]
    elseif d === :C2
        return P.C[ny1+1:end, :]
    elseif d === :D11
        return P.D[1:ny1, 1:nu1]
    elseif d === :D12
        return P.D[1:ny1, nu1+1:end]
    elseif d === :D21
        return P.D[ny1+1:end, 1:nu1]
    elseif d === :D22
        return P.D[ny1+1:end, nu1+1:end]
    else
        return getfield(P, d)
    end
end

timeevol(sys::PartionedStateSpace) = timeevol(sys.P)

function +(s1::PartionedStateSpace, s2::PartionedStateSpace)
    timeevol = common_timeevol(s1,s2)

    A = blockdiag(s1.A, s2.A)

    B = [[s1.B1; s2.B1] blockdiag(s1.B2, s2.B2)]

    C = [[s1.C1 s2.C1];
    blockdiag(s1.C2, s2.C2)]

    D = [(s1.D11 + s2.D11) s1.D12 s2.D12;
    [s1.D21; s2.D21] blockdiag(s1.D22, s2.D22)]

    P = StateSpace(A, B, C, D, timeevol) # How to handle discrete?
    PartionedStateSpace(P, s1.nu1 + s2.nu1, s1.ny1 + s2.ny1)
end

"""
    Series connection of partioned StateSpace systems.
"""
function *(s1::PartionedStateSpace, s2::PartionedStateSpace)
    timeevol = common_timeevol(s1,s2)

    A = [s1.A                           s1.B1*s2.C1;
    zeros(size(s2.A,1),size(s1.A,2))      s2.A]

    B = [s1.B1*s2.D11                         s1.B2           s1.B1*s2.D12;
    s2.B1              zeros(size(s2.B2,1),size(s1.B2,2))          s2.B2]

    C = [s1.C1                       s1.D11*s2.C1;
    s1.C2                        s1.D21*s2.C1;
    zeros(size(s2.C2,1),size(s1.C2,2))         s2.C2]

    D = [s1.D11*s2.D11           s1.D12        s1.D11*s2.D12;
    s1.D21*s2.D11           s1.D22        s1.D21*s2.D12;
    s2.D21          zeros(size(s2.D22,1),size(s1.D22,2))          s2.D22        ]

    P = StateSpace(A, B, C, D, timeevol)
    PartionedStateSpace(P, s2.nu1, s1.ny1)
end

# QUESTION: What about algebraic loops and well-posedness?! Perhaps issue warning if P1(∞)*P2(∞) > 1
function feedback(s1::PartionedStateSpace, s2::PartionedStateSpace)
    timeevol = common_timeevol(s1,s2)
    X_11 = (I + s2.D11*s1.D11)\[-s2.D11*s1.C1  -s2.C1]
    X_21 = (I + s1.D11*s2.D11)\[s1.C1  -s1.D11*s2.C1]

    # For the case of two outputs
    #    X_12 = [I   -s2.D11   -s2.D11*s1.D12   -s2.D12]
    #    X_22 = [s1.D11  I     s1.D12          -s1.D11*s2.D12]
    X_12 = (I + s2.D11*s1.D11)\[I      -s2.D11*s1.D12   -s2.D12]
    X_22 = (I + s1.D11*s2.D11)\[s1.D11   s1.D12          -s1.D11*s2.D12]

    A = [s1.B1 * X_11 ; s2.B1 * X_21] + blockdiag(s1.A, s2.A)

    B = [s1.B1 * X_12 ; s2.B1 * X_22]
    tmp = blockdiag(s1.B2, s2.B2)
    B[:, end-size(tmp,2)+1:end] .+= tmp

    C = [s1.D11 * X_11 ;
         s1.D21 * X_11 ;
         s2.D21 * X_21 ] + [s1.C1 zeros(size(s1.C1,1),size(s2.C1,2)); blockdiag(s1.C2, s2.C2)]

    D = [s1.D11 * X_12 ;
        s1.D21 * X_12 ;
        s2.D21 * X_22 ]
    tmp = [s1.D12 zeros(size(s1.D12,1),size(s2.D12,2)); blockdiag(s1.D22, s2.D22)]
    D[:, end-size(tmp,2)+1:end] .+= tmp

    # in case it is desired to consider both outputs
    # C = [s1.D11 * X_11 ;
    #      s2.D11 * X_21 ;
    #      s1.D21 * X_11 ;
    #      s2.D21 * X_21 ] + [blockdiag(s1.C1, s2.C1); blockdiag(s1.C2, s2.C2)]
    #
    # D = [s1.D11 * X_12 ;
    #     s2.D11 * X_22 ;
    #     s1.D21 * X_12 ;
    #     s2.D21 * X_22 ]
    #tmp = [blockdiag(s1.D12, s2.D12); blockdiag(s1.D22, s2.D22)]
    #D[:, end-size(tmp,2)+1:end] .+= tmp

    P = StateSpace(A, B, C, D, timeevol)
    PartionedStateSpace(P, s2.nu1, s1.ny1)
end

""" Concatenate systems vertically with
    same first input u1
    second input [u2_1; u2_2 ...]
    and output  y1 = [y1_1; y1_2, ...]
                y2 = [y2_1; y2_2, ...]
    for u1_i, u2_i, y1_i, y2_i, where i denotes system i
"""
function vcat_1(systems::PartionedStateSpace...)
    # Perform checks
    timeevol = common_timeevol(systems...)

    nu1 = systems[1].nu1
    if !all(s.nu1 == nu1 for s in systems)
        error("All systems must have same first input dimension")
    end

    A = blockdiag([s.A for s in systems]...)

    B1 = reduce(vcat, [s.B1 for s in systems])
    B2 = blockdiag([s.B2 for s in systems]...)

    C1 = blockdiag([s.C1 for s in systems]...)
    C2 = blockdiag([s.C2 for s in systems]...)

    D11 = reduce(vcat, [s.D11 for s in systems])
    D12 = blockdiag([s.D12 for s in systems]...)
    D21 = reduce(vcat, [s.D21 for s in systems])
    D22 = blockdiag([s.D22 for s in systems]...)

    sysnew = StateSpace(A, [B1 B2], [C1; C2], [D11 D12; D21 D22], timeevol)
    return PartionedStateSpace(sysnew, nu1, sum(s -> s.ny1, systems))
end

""" Concatenate systems horizontally with
    same first ouput y1 being the sum
    second output y2 = [y2_1; y2_2, ...]
    and inputs  u1 = [u1_1; u1_2, ...]
    u2 = [u2_1; u2_2, ...]
    for u1_i, u2_i, y1_i, y2_i, where i denotes system i
"""
function hcat_1(systems::PartionedStateSpace...)
    # Perform checks
    timeevol = common_timeevol(systems...)

    ny1 = systems[1].ny1
    if !all(s.ny1 == ny1 for s in systems)
        error("All systems must have same first ouput dimension")
    end

    A = blockdiag([s.A for s in systems]...)

    B1 = blockdiag([s.B1 for s in systems]...)
    B2 = blockdiag([s.B2 for s in systems]...)

    C1 = reduce(hcat, [s.C1 for s in systems])
    C2 = blockdiag([s.C2 for s in systems]...)

    D11 = reduce(hcat, [s.D11 for s in systems])
    D12 = reduce(hcat, [s.D12 for s in systems])
    D21 = blockdiag([s.D21 for s in systems]...)
    D22 = blockdiag([s.D22 for s in systems]...)

    sysnew = StateSpace(A, [B1 B2], [C1; C2], [D11 D12; D21 D22], timeevol)
    return PartionedStateSpace(sysnew, sum(s -> s.nu1, systems), ny1)
end

function Base.convert(::Type{<:PartionedStateSpace}, sys::T) where T<: StateSpace
    PartionedStateSpace(sys,sys.nu,sys.ny)
end


text:
pushfirst!(LOAD_PATH, joinpath(@__DIR__, "..", ".."))

ENV["GKSwstype"] = "nul"

using Plots
using ForwardDiff
using Measures
using LinearAlgebra

using Statistics
using Printf

using Oceananigans
using Oceananigans.Fields: interpolate
using Oceananigans.BoundaryConditions
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBoundary

#####
##### The Steady State Cylinder
#####

const R = 1.      # radius
const Re = 40.    # Reynolds number, Re < 40 --> Steady State
const nu = 2*R/Re  # viscosity

"""
    run_cylinder_steadyflow(output_time_interval = 1, stop_time = 100, arch = CPU(), Nh = 250, ν = 1/20,
                    momentum_advection = UpwindBiasedFifthOrder(), radius = R)

Run the steady state cylinder validation experiment until `stop_time` using `momentum_advection`
scheme or formulation, with horizontal resolution `Nh`, viscosity `ν`, and cylinder radius `R`, on `arch`itecture.
"""

function run_cylinder_steadystate(; output_time_interval = 1, stop_time = 100, arch = CPU(), Nh = 250, ν = nu,
                                    advection = UpwindBiasedFifthOrder(), radius = R)

    inside_cylinder(x, y, z) = (x^2 + y^2) <= radius # immersed solid

    underlying_grid = RectilinearGrid(arch, size=(Nh, Int(3*Nh/2),1), halo=(3, 3, 3),
                                           x = (-10, 10), y=(-10, 20), z = (0,1),
                                           topology = (Periodic, Bounded, Bounded))

    immersed_grid = ImmersedBoundaryGrid(arch, underlying_grid, GridFittedBoundary(inside_cylinder))

    # boundary conditions: inflow and outflow in y
    v_bc = OpenBoundaryCondition(1.0)
    v_bcs = FieldBoundaryConditions(north = v_bc, south = v_bc)

    immersed_model = NonhydrostaticModel(advection = advection,
                                         timestepper = :RungeKutta3,
                                         grid = immersed_grid,
                                         tracers = :mass,
                                         closure = IsotropicDiffusivity(ν=ν, κ=ν),
                                         boundary_conditions = (v=v_bcs,),
                                         coriolis = nothing,
                                         buoyancy = nothing)

    # ** Initial conditions **
    #
    # u: no flow
    # v: uniform flow of U∞ = 1, except inside cylinder (no flow)

    # Parameters
    U∞ = 1.0 # uniform flow

    # Total initial conditions
    uᵢ(x, y, z) = 0
    vᵢ(x, y, z) = U∞
    mᵢ(x, y, z) = 1.0
    #mᵢ(x, y, z) = ifelse(inside_cylinder(x,y,z), 0., 1.0)
    set!(immersed_model, u=uᵢ, v=vᵢ, mass=mᵢ)

    wall_clock = [time_ns()]

    function progress(sim)
        @info(@sprintf("Iter: %d, time: %.1f, Δt: %.3f, wall time: %s, max|v|: %.2f",
                       sim.model.clock.iteration,
                       sim.model.clock.time,
                       sim.Δt.Δt,
                       prettytime(1e-9 * (time_ns() - wall_clock[1])),
                       maximum(abs, sim.model.velocities.v.data.parent)))

        wall_clock[1] = time_ns()

        return nothing
    end

    @show experiment_name = "cylinder_tracer_Nh_$(Nh)_$(typeof(immersed_model.advection).name.wrapper)"

    wizard = TimeStepWizard(cfl=0.07, Δt=0.07 * underlying_grid.Δxᶜᵃᵃ, max_change=1.1, max_Δt=10.0, min_Δt=0.0001)

    simulation = Simulation(immersed_model, Δt=wizard, stop_time=stop_time, iteration_interval=100, progress=progress)

    # Output: primitive fields + computations
    u, v, w, pHY′, pNHS, mass  = merge(immersed_model.velocities, immersed_model.pressures, immersed_model.tracers)
    outputs = merge(immersed_model.velocities, immersed_model.pressures, immersed_model.tracers)

    data_path = experiment_name

    simulation.output_writers[:fields] =
            JLD2OutputWriter(immersed_model, outputs,
                             schedule = TimeInterval(output_time_interval),
                             prefix = data_path,
                             field_slicer = nothing,
                             force = true)

    @info "Running a simulation of an steady state cylinder..."

    start_time = time_ns()

    run!(simulation)
    return experiment_name
end

"""
    visualize_cylinder_steadystate(experiment_name)

Visualize the steady state cylinder data associated with `experiment_name`.
"""

# defining a function to mark boundary in plots
function circle_shape(h, k, r)
            θ = LinRange(0,2*π,500)
            h.+r*sin.(θ),k.+r*cos.(θ)
end

function visualize_cylinder_steadystate(experiment_name)

    @info "Making a fun movie about steady state cylinder flow..."

    filepath = experiment_name * ".jld2"

    u_timeseries = FieldTimeSeries(filepath,  "u")

    v_timeseries = FieldTimeSeries(filepath,  "v")

    p_timeseries = FieldTimeSeries(filepath,  "pNHS")

    m_timeseries = FieldTimeSeries(filepath, "mass")

    regular_grid = p_timeseries.grid

    xu, yu, zu = nodes(u_timeseries)
    xv, yv, zv = nodes(v_timeseries)
    xp, yp, zp = nodes(p_timeseries)
    xm, ym, zm = nodes(m_timeseries)

    anim = @animate for (i, t) in enumerate(p_timeseries.times)

        @info "Plotting frame $i of $(length(p_timeseries.times))..."

        Nx, Ny, Nz = size(regular_grid)

        u = u_timeseries[i]
        v = v_timeseries[i]
        p = p_timeseries[i]
        m = m_timeseries[i]

        ui = interior(u)[:, :, 1]
        vi = interior(v)[:, :, 1]
        pin = interior(p)[:, :, 1] # don't want to confuse with pi
        mi = interior(m)[:, :, 1]
        Nx, Ny = size(ui)

        kwargs = Dict(:aspectratio => 1,
                      :linewidth => 0,
                      :colorbar => :true,
                      :ticks => nothing,
                      :clims => (-1, 1),
                      :xlims => (-regular_grid.Lx/2, regular_grid.Lx/2),
                      :ylims => (-regular_grid.Ly/3, 2*regular_grid.Ly/3))

        u_plot = heatmap(xu, yu, clamp.(ui, -1, 1)'; color = :balance, size = (500,700),
                            guidefontsize = 14, titlefont=14, top_margin = -10.0mm, bottom_margin = -10.0mm,
                            legendfont = 10, tickfont = 8, kwargs...)
        plot!(circle_shape(0,0,1),seriestype=[:shape,],linecolor=:black,
                            legend=false, fillalpha=0)

        v_plot = heatmap(xv, yv, clamp.(vi, -1, 1)'; color = :balance, size = (500,700),
                            guidefontsize = 14, titlefont=14, top_margin = -10.0mm, bottom_margin = -10.0mm,
                            legendfont = 10, tickfont = 8, kwargs...)
        plot!(circle_shape(0,0,1),seriestype=[:shape,],linecolor=:black,
                            legend=false, fillalpha=0)

        p_plot = heatmap(xp, yp, clamp.(pin, -1, 1)'; color = :balance, size = (500,700),
                            guidefontsize = 14, titlefont=14, top_margin = -10.0mm, bottom_margin = -10.0mm,
                            legendfont = 10, tickfont = 8, kwargs...)
        plot!(circle_shape(0,0,1),seriestype=[:shape,],linecolor=:black,
                            legend=false, fillalpha=0)

        m_plot = heatmap(xm, ym, clamp.(mi, -1, 1)'; color = :thermal, size = (500,700),
                            guidefontsize = 14, titlefont=14, top_margin = -10.0mm, bottom_margin = -10.0mm,
                            legendfont = 10, tickfont = 8, kwargs...)
        plot!(circle_shape(0,0,1),seriestype=[:shape,],linecolor=:black,
                            legend=false, fillalpha=0)

        δlim = 1.0
        kwargs[:clims] = (-δlim, δlim)

        u_title = @sprintf("u at t = %.1f", t)
        v_title = @sprintf("v at t = %.1f", t)
        p_title = @sprintf("p at t = %.1f", t)
        m_title = @sprintf("m at t = %.1f", t)

        plot(u_plot, v_plot, p_plot, m_plot,
             title = [u_title v_title p_title m_title],
             layout = (2, 2))
    end
    mp4(anim, "contours_" * experiment_name * ".mp4", fps = 8)
end

function analyze_cylinder_steadystate(experiment_name)

    @info "Analyzing data from steady state cylinder flow..."
    # Analysis of coefficients is for no slip cylinder BC

    # distance function for solid cylinder
    @inline dist_cylinder(v) = sqrt((v[1])^2+(v[2])^2)-R
    # finding the sfc normal vector
    normalDist = x-> ForwardDiff.gradient(dist_cylinder,x)

    # function creating the projection matrix into surface normal and tangential
    function projection_matrix(N)
        if abs(N[1]) == 1 # if normal vector is entirely in the x direction
            v1 = [0; -sign(N[1]); 0]  # we want v1 = [0,0,1]
            v2 = [0; 0; -1]
        elseif abs(N[2]) == 1 || N[3]==0 # if normal vector is entirely in the y direction or 0 z comp
            v1 = [sign(N[2]); -N[1]*sign(N[2])/N[2]; 0] # we want v1 = [1,0,0]
            v1 = v1./norm(v1)
            v2 = cross(N,v1)
        else
            v1 = [sign(N[3]); 0; -N[1]*sign(N[3])/N[3]] # else we want v1 = [1, 0 , ?]
            v1 = v1./norm(v1)
            v2 = cross(N,v1)
        end
        transpose(hcat(v1,v2,N)) # creating a matrix of all 3 vectors
    end

    # function interpolating velocity vector at vector location xvec
    function interp_values(xvec, Us, Vs, Ws)
        xI = xvec; # forced node is the x argument
        uI = interpolate(Us, xI[1], xI[2], xI[3])
        vI = interpolate(Vs, xI[1], xI[2], xI[3])
        wI = interpolate(Ws, xI[1], xI[2], xI[3])
        vels = [uI; vI; wI]
        return vels
    end

    # function outputs normal and tangential velocity, pressure, and friction coefficient
    # around a cylinder at M points with projection and interpolation
    function surface_values(M, U, V, W, Pnon, grid, eps, ν)

        # angle around circle starting due south
        θ = LinRange(-π/2 , 3*π/2-(2*π/M), M)
        x₀ = hcat(R*cos.(θ), R*sin.(θ), .5*ones(M,1))
        Vⁿ = zeros(M, 1)
        Vᵗ¹ = zeros(M, 1)
        p_sfc = zeros(M, 1)
        dVt = zeros(M, 1)

        # at each location around circle
        for m = 1:M

            # surface normal vector
            nm = normalDist(x₀[m,:])
            # location eps distance outside cylinder in fluid
            xI = x₀[m,:] + eps*nm

            # velocity vector at both points
            Vm = interp_values(x₀[m,:], U, V, W)
            Vm_ext = interp_values(xI, U, V, W)

            # new projected velocity vector: V_rot = [Vt1, Vt2, Vn]
            matrix = projection_matrix(nm)
            V_rot = matrix*Vm
            V_rot_ext = matrix*Vm_ext

            # normal and tangential velocities
            Vⁿ[m] = V_rot[3]
            Vᵗ¹[m] = V_rot[1]
            Vᵗ¹_ext = V_rot_ext[1]

            # surface normal derivative dVt/dn
            dVt[m] = (Vᵗ¹_ext - Vᵗ¹[m])/eps

            # pressure interpolated near the surface
            p_sfc[m] = interpolate(Pnon, xI[1], xI[2], xI[3])
        end

        # friction coefficient
        Cf = dVt.* (2*ν)

        return Vⁿ,Vᵗ¹, Cf, p_sfc, θ
    end

    # get data
    filepath = experiment_name * ".jld2"

    u_timeseries = FieldTimeSeries(filepath,  "u")
    v_timeseries = FieldTimeSeries(filepath,  "v")
    w_timeseries = FieldTimeSeries(filepath,  "w")
    p_timeseries = FieldTimeSeries(filepath,  "pNHS")
    m_timeseries = FieldTimeSeries(filepath,  "mass")

    # final iteration for steady state values
    final_iter = size(u_timeseries)[end]
    U = u_timeseries[final_iter]
    V = v_timeseries[final_iter]
    W = w_timeseries[final_iter]
    P = p_timeseries[final_iter]

    center_grid = p_timeseries.grid
    xm, ym, zm = nodes(m_timeseries)

    Vⁿ,Vᵗ, Cf, p_sfc, θ  = surface_values(140, U, V, W, P, center_grid, center_grid.Δxᶜᵃᵃ, nu)

    # function creating mesh of (x,y)
    function meshgrid(vx::AbstractVector{T}, vy::AbstractVector{T}) where T
        m, n = length(vy), length(vx)
        vx = reshape(vx, 1, n)
        vy = reshape(vy, m, 1)
        (repeat(vx, m, 1), repeat(vy, 1, n))
    end

    @inline inside_cylinder(x, y, z) = (x.^2 .+ y.^2) .<= R # immersed solid

    # calculates the volume integrated concentration of tracer
    function volume_integrated_concentration(x, y, mass_series)
        X, Y = meshgrid(x, y)
        # only including concentration values outside the cylinder in the fluid
        fluid_concent = mass_series .* (1 .- inside_cylinder(X', Y', 1))
        # summing over the concentration values and multiplying by ΔxΔy at each iteration
        vol_int_concent = (sum(fluid_concent[:,:,1,:], dims = 1:2))[1,1,:] * (y[2]-y[1]) * (x[2]-x[1])
    end

    concentration_m = volume_integrated_concentration(xm, ym, m_timeseries)

    # taking the percent change in concentration at each time step in ratio to the initial amount
    percent_leakage = ((concentration_m[1] .- concentration_m) ./ concentration_m[1]) .* 100

    @info "Plotting Normal and Tangential Velocity..."

    Vnt = plot(θ*180/π, Vⁿ, xlabel="θ", ylabel="Velocity", seriestype=:scatter,
        guidefontsize = 14, titlefont=14, legendfont = 8, tickfont = 8,
        xticks = [-90,0,90,180,275], xlims = (-90, 275), markersize = 3,
        legend = true, size = (700,500), label = "Vⁿ")
    plot!(θ*180/π, Vᵗ, label= "Vᵗ", seriestype=:scatter, markersize = 3)

    @info "Plotting Pressure Coefficient..."

    Pcoef = plot(θ*180/π, 2. .* p_sfc, xlabel="θ", ylabel= "Cp", seriestype=:scatter,
        guidefontsize = 14, titlefont=14, legendfont = 10, tickfont = 8,
        xticks = [-90,0,90,180,275], xlims = (-90, 275), markersize = 3,
        legend = false, size = (700,500))

    @info "Plotting Friction Coefficient..."

    Fcoef = plot(θ*180/π, Cf, xlabel="θ", ylabel= "Cf", seriestype=:scatter,
        guidefontsize = 14, titlefont=14,legendfont = 10, tickfont = 8,
        xticks = [-90,0,90,180,275], xlims = (-90, 275), markersize = 3,
        legend = false, size = (700,500))

    @info "Plotting Concentration Leakage"
    mConcent = plot(m_timeseries.times, percent_leakage, xlabel="t", ylabel= "% Leakage",
        seriestype=:scatter, guidefontsize = 14, titlefont=14, legendfont = 10,
        tickfont = 8, markersize = 3, legend = false, size = (700,500))

    Vnt_title = @sprintf("||Vⁿ||₂  = %.3f, ||Vᵗ||₂  = %.3f", norm(Vⁿ,2), norm(Vᵗ,2))
    psfc_title = "Pressure Coefficient"
    cf_title = "Friction Coefficient"
    trac_title = "% Concentration Leakage"

    analysisplot = plot(Vnt, Pcoef, Fcoef, mConcent,
             title = [Vnt_title psfc_title cf_title trac_title],
             layout = (2, 2), left_margin=5mm, bottom_margin=10mm)

    Plots.savefig(analysisplot, "analysis_" * experiment_name * ".png")
end

advection = CenteredSecondOrder()
experiment_name = run_cylinder_steadystate(Nh = 350, advection = advection, radius = R, stop_time = 100, ν = nu)
visualize_cylinder_steadystate(experiment_name)
analyze_cylinder_steadystate(experiment_name)


text:
cd("/home/astupidbear/.julia/v0.4/JuLearn/examples/mlp/classification/")
using JuLearn,ArrayFire,Plots
@Eval lr=1e-1 momentum=0. batchsize=100 λ=0. ▽max=0.1
include("dataProvider.jl")
include("predict.jl")

inputBase,outputBase=dataProvider()
net=chain([:d,:h,:h,:p,:l],[2 2;2 10;10 2;2 2;2 2],:classification)
@time err=backprop(net,Array(inputBase),Array(outputBase),nIter=100)
p=plot(err);savefig(p,"error.png")
run(`xdg-open error.png`)


text:
using DelayEmbeddings, SparseArrays
include("GroupSlices.jl")

export
    TransferOperator, # the probabilities estimator
    InvariantMeasure, invariantmeasure,
    transfermatrix

"""
    TransferOperator(ϵ::RectangularBinning) <: BinningProbabilitiesEstimator

A probability estimator based on binning data into rectangular boxes dictated by
the binning scheme `ϵ`, then approxmating the transfer (Perron-Frobenius) operator
over the bins, then taking the invariant measure associated with that transfer operator
as the bin probabilities. Assumes that the input data are sequential (time-ordered).

This implementation follows the grid estimator approach in Diego et al. (2019)[^Diego2019].

## Description

The transfer operator ``P^{N}``is computed as an `N`-by-`N` matrix of transition
probabilities between the states defined by the partition elements, where `N` is the
number of boxes in the partition that is visited by the orbit/points.

If  ``\\{x_t^{(D)} \\}_{n=1}^L`` are the ``L`` different ``D``-dimensional points over
which the transfer operator is approximated, ``\\{ C_{k=1}^N \\}`` are the ``N`` different
partition elements (as dictated by `ϵ`) that gets visited by the points, and
 ``\\phi(x_t) = x_{t+1}``, then

```math
P_{ij} = \\dfrac
{\\#\\{ x_n | \\phi(x_n) \\in C_j \\cap x_n \\in C_i \\}}
{\\#\\{ x_m | x_m \\in C_i \\}},
```

where ``\\#`` denotes the cardinal. The element ``P_{ij}`` thus indicates how many points
that are initially in box ``C_i`` end up in box ``C_j`` when the points in ``C_i`` are
projected one step forward in time. Thus, the row ``P_{ik}^N`` where
``k \\in \\{1, 2, \\ldots, N \\}`` gives the probability
of jumping from the state defined by box ``C_i`` to any of the other ``N`` states. It
follows that ``\\sum_{k=1}^{N} P_{ik} = 1`` for all ``i``. Thus, ``P^N`` is a row/right
stochastic matrix.

### Invariant measure estimation from transfer operator

The left invariant distribution ``\\mathbf{\\rho}^N`` is a row vector, where
``\\mathbf{\\rho}^N P^{N} = \\mathbf{\\rho}^N``. Hence, ``\\mathbf{\\rho}^N`` is a row
eigenvector of the transfer matrix ``P^{N}`` associated with eigenvalue 1. The distribution
``\\mathbf{\\rho}^N`` approximates the invariant density of the system subject to the
partition `ϵ`, and can be taken as a probability distribution over the partition elements.

In practice, the invariant measure ``\\mathbf{\\rho}^N`` is computed using
[`invariantmeasure`](@ref), which also approximates the transfer matrix. The invariant distribution
is initialized as a length-`N` random distribution which is then applied to ``P^{N}``.
The resulting length-`N` distribution is then applied to ``P^{N}`` again. This process
repeats until the difference between the distributions over consecutive iterations is
below some threshold.

## Probability and entropy estimation

- `probabilities(x::AbstractDataset, est::TransferOperator{RectangularBinning})` estimates
    probabilities for the bins defined by the provided binning (`est.ϵ`)
- `genentropy(x::AbstractDataset, est::TransferOperator{RectangularBinning})` does the same,
    but computes generalized entropy using the probabilities.

See also: [`RectangularBinning`](@ref), [`invariantmeasure`](@ref).

[^Diego2019]: Diego, D., Haaga, K. A., & Hannisdal, B. (2019). Transfer entropy computation using the Perron-Frobenius operator. Physical Review E, 99(4), 042212.
"""
struct TransferOperator{R} <: BinningProbabilitiesEstimator
    ϵ::R
end

# If x is not sorted, we need to look at all pairwise comparisons
function inds_in_terms_of_unique(x)
    U = unique(x)
    N = length(x)
    Nu = length(U)
    inds = zeros(Int, N)

    for j = 1:N
        xⱼ = view(x, j)
        for i = 1:Nu
            # using views doesn't allocate
            @inbounds if xⱼ == view(U, i)
                inds[j] = i
            end
        end
    end

    return inds
end

# Taking advantage of the fact that x is sorted reduces runtime by 1.5 orders of magnitude
# for datasets of >100 000+ points
function inds_in_terms_of_unique_sorted(x) # assumes sorted
    @assert issorted(x)
    U = unique(x)
    N, Nu = length(x), length(U)
    prev = view(x, 1)
    inds = zeros(Int, N)
    uidx = 1
    @inbounds for j = 1:N
        xⱼ = view(x, j)
        # if the current value has changed, then we know that the corresponding index
        # for the unique point must be incremented by 1
        if xⱼ != prev
            prev = xⱼ
            uidx += 1
        end
        inds[j] = uidx
    end

    return inds
end

function inds_in_terms_of_unique(x, sorted::Bool)
    if sorted
        return inds_in_terms_of_unique_sorted(x)
    else
        return inds_in_terms_of_unique(x)
    end
end

inds_in_terms_of_unique(x::AbstractDataset) = inds_in_terms_of_unique(x.data)

"""
    TransferOperatorApproximationRectangular(to, ϵ::RectangularBinning, mini, edgelengths,
        bins, sort_idxs)

The `N`-by-`N` matrix `to` is an approximation to the transfer operator, subject to the
partition `ϵ`, computed over some set of sequentially ordered points.

For convenience, `mini` and `edgelengths` provide the minima and box edge lengths along
each coordinate axis, as determined by applying `ϵ` to the points. The coordinates of
the (leftmost, if axis is ordered low-high) box corners are given in `bins`.

Only bins actually visited by the points are considered, and `bins` give the coordinates
of these bins. The element `bins[i]` correspond to the `i`-th state of the system, which
corresponds to the `i`-th column/row of the transfer operator `to`.

`sort_idxs` contains the indices that would sort the input points. `visitors` is a
vector of vectors, where `visitors[i]` contains the indices of the (sorted)
points that visits `bins[i]`.
See also: [`RectangularBinning`](@ref).
"""
struct TransferOperatorApproximationRectangular{T<:Real}
    transfermatrix::AbstractArray{T, 2}
    ϵ::RectangularBinning
    mini
    edgelengths
    bins
    sort_idxs::Vector{Int}
    visitors::Vector{Vector{Int}}
end
# TODO: The above is type unstable!

"""
    transferoperator(pts::AbstractDataset{D, T}, ϵ::RectangularBinning) → TransferOperatorApproximationRectangular

Estimate the transfer operator given a set of sequentially ordered points subject to a
rectangular partition given by `ϵ`.

## Example

```julia
using DynamicalSystems, Plots, Entropy
D = 4
ds = Systems.lorenz96(D; F = 32.0)
N, dt = 20000, 0.1
orbit = trajectory(ds, N*dt; dt = dt, Ttr = 10.0)

# Estimate transfer operator over some coarse graining of the orbit.
transferoperator(orbit, RectangularBinning(10))
```

See also: [`RectangularBinning`](@ref).
"""
function transferoperator(pts::AbstractDataset{D, T}, ϵ::RectangularBinning;
        boundary_condition = :circular) where {D, T<:Real}

    L = length(pts)
    mini, edgelengths = Entropies.minima_edgelengths(pts, ϵ)

    # The L points visits a total of L bins, which are the following bins:
    visited_bins = Entropies.encode_as_bin(pts, mini, edgelengths)
    sort_idxs = sortperm(visited_bins)

    # TODO: fix re-indexing after sorting. Sorting is much faster, so we want to do so.
    #sort!(visited_bins)

    # There are N=length(unique(visited_bins)) unique bins.
    # Which of the unqiue bins does each of the L points visit?
    visits_whichbin = inds_in_terms_of_unique(visited_bins, false) # set to true when sorting is fixed

    # `visitors` lists the indices of the points visiting each of the N unique bins.
    slices = GroupSlices.groupslices(visited_bins)
    visitors = GroupSlices.groupinds(slices)

    # first_visited_by == [x[1] for x in visitors]
    first_visited_by = GroupSlices.firstinds(slices)
    L = length(first_visited_by)

    I = Int32[]
    J = Int32[]
    P = Float64[]

    # Preallocate target index for the case where there is only
    # one point of the orbit visiting a bin.
    target_bin_j::Int = 0
    n_visitsᵢ::Int = 0

    if boundary_condition == :circular
        #warn("Using circular boundary condition")
        append!(visits_whichbin, [1])
    elseif boundary_condition == :random
        #warn("Using random circular boundary condition")
        append!(visits_whichbin, [rand(1:length(visits_whichbin))])
    else
        error("Boundary condition $(boundary_condition) not implemented")
    end

    # Loop over the visited bins bᵢ
    for i in 1:L
        # How many times is this bin visited?
        n_visitsᵢ = length(visitors[i])

        # If both conditions below are true, then there is just one
        # point visiting the i-th bin. If there is only one visiting point and
        # it happens to be the last, we skip it, because we don't know its
        # image.
        if n_visitsᵢ == 1 && !(i == visits_whichbin[end])
            # To which bin does the single point visiting bᵢ jump if we
            # shift it one time step ahead along its orbit?
            target_bin_j = visits_whichbin[visitors[i][1] + 1][1]

            # We now know that exactly one point (the i-th) does the
            # transition from i to the target j.
            push!(I, i)
            push!(J, target_bin_j)
            push!(P, 1.0)
        end

        # If more than one point of the orbit visits the i-th bin, we
        # identify the visiting points and track which bins bⱼ they end up
        # in after the forward linear map of the points.
        if n_visitsᵢ > 1
            timeindices_visiting_pts = visitors[i]

            # TODO: Introduce circular boundary condition. Simply excluding
            # might lead to a cascade of loosing points.

            # If bᵢ is the bin visited by the last point in the orbit, then
            # the last entry of `visiting_pts` will be the time index of the
            # last point of the orbit. In the next time step, that point will
            # have nowhere to go along its orbit (precisely because it is the
            # last data point). Thus, we exclude it.
            if i == visits_whichbin[end]
                #warn("Removing last point")
                n_visitsᵢ = length(timeindices_visiting_pts) - 1
                timeindices_visiting_pts = timeindices_visiting_pts[1:(end - 1)]
            end

            # To which boxes do each of the visitors to bᵢ jump in the next
            # time step?
            target_bins = visits_whichbin[timeindices_visiting_pts .+ 1]
            unique_target_bins = unique(target_bins)

            # Count how many points jump from the i-th bin to each of
            # the unique target bins, and use that to calculate the transition
            # probability from bᵢ to bⱼ.
            #for j in 1:length(unique_target_bins)
            for (j, bᵤ) in enumerate(unique(target_bins))
                n_transitions_i_to_j = sum(target_bins .== bᵤ)

                push!(I, i)
                push!(J, bᵤ)
                push!(P, n_transitions_i_to_j / n_visitsᵢ)
            end
        end
    end

    # Transfer operator is just the normalized transition probabilities between the boxes.
    TO = sparse(I, J, P)

    # Compute the coordinates of the visited bins. bins[i] corresponds to the i-th
    # row/column of the transfer operator
    unique!(visited_bins)
    bins = [β .* edgelengths .+ mini for β in visited_bins]

    TransferOperatorApproximationRectangular(TO, ϵ, mini, edgelengths, bins,
        sort_idxs, visitors)
end

"""
    InvariantMeasure(to, ρ)

Minimal return struct for [`invariantmeasure`](@ref) that contains the estimated invariant
measure `ρ`, as well as the transfer operator `to` from which it is computed (including
bin information).

See also: [`invariantmeasure`](@ref).
"""
struct InvariantMeasure{T}
    to::T
    ρ::Probabilities
end

function invariantmeasure(iv::InvariantMeasure)
    return iv.ρ, iv.to.bins
end

import LinearAlgebra: norm
"""
    invariantmeasure(x::AbstractDataset, ϵ::RectangularBinning) → iv::InvariantMeasure

Estimate an invariant measure over the points in `x` based on binning the data into
rectangular boxes dictated by the binning scheme `ϵ`, then approximate the transfer
(Perron-Frobenius) operator over the bins. From the approximation to the transfer operator,
compute an invariant distribution over the bins. Assumes that the input data are sequential.

Details on the estimation procedure is found the [`TransferOperator`](@ref) docstring.

## Example

```julia
using DynamicalSystems, Plots, Entropies
D = 4
ds = Systems.lorenz96(D; F = 32.0)
N, dt = 20000, 0.1
orbit = trajectory(ds, N*dt; dt = dt, Ttr = 10.0)

# Estimate the invariant measure over some coarse graining of the orbit.
iv = invariantmeasure(orbit, RectangularBinning(15))

# Get the probabilities and bins
invariantmeasure(iv)
```

## Probabilities and bin information

    invariantmeasure(iv::InvariantMeasure) → (ρ::Probabilities, bins::Vector{<:SVector})

From a pre-computed invariant measure, return the probabilities and associated bins.
The element `ρ[i]` is the probability of visitation to the box `bins[i]`. Analogous to
[`binhist`](@ref).

!!! hint "Transfer operator approach vs. naive histogram approach"

    Why bother with the transfer operator instead of using regular histograms to obtain
    probabilities?

    In fact, the naive histogram approach and the
    transfer operator approach are equivalent in the limit of long enough time series
    (as ``n \\to \\intfy``), which is guaranteed by the ergodic theorem. There is a crucial
    difference, however:

    The naive histogram approach only gives the long-term probabilities that
    orbits visit a certain region of the state space. The transfer operator encodes that
    information too, but comes with the added benefit of knowing the *transition
    probabilities* between states (see [`transfermatrix`](@ref)).

See also: [`InvariantMeasure`](@ref).
"""
function invariantmeasure(to::TransferOperatorApproximationRectangular;
        N::Int = 200, tolerance::Float64 = 1e-8, delta::Float64 = 1e-8)

    TO = to.transfermatrix
    #=
    # Start with a random distribution `Ρ` (big rho). Normalise it so that it
    # sums to 1 and forms a true probability distribution over the partition elements.
    =#
    Ρ = rand(Float64, 1, size(to.transfermatrix, 1))
    Ρ = Ρ ./ sum(Ρ, dims = 2)

    #=
    # Start estimating the invariant distribution. We could either do this by
    # finding the left-eigenvector of M, or by repeated application of M on Ρ
    # until the distribution converges. Here, we use the latter approach,
    # meaning that we iterate until Ρ doesn't change substantially between
    # iterations.
    =#
    distribution = Ρ * to.transfermatrix

    distance = norm(distribution - Ρ) / norm(Ρ)

    check = floor(Int, 1 / delta)
    check_pts = floor.(Int, transpose(collect(1:N)) ./ check) .* transpose(collect(1:N))
    check_pts = check_pts[check_pts .> 0]
    num_checkpts = size(check_pts, 1)
    check_pts_counter = 1

    counter = 1
    while counter <= N && distance >= tolerance
        counter += 1
        Ρ = distribution

        # Apply the Markov matrix to the current state of the distribution
        distribution = Ρ * to.transfermatrix

        if (check_pts_counter <= num_checkpts &&
           counter == check_pts[check_pts_counter])

            check_pts_counter += 1
            colsum_distribution = sum(distribution, dims = 2)[1]
            if abs(colsum_distribution - 1) > delta
                distribution = distribution ./ colsum_distribution
            end
        end

        distance = norm(distribution - Ρ) / norm(Ρ)
    end
    distribution = dropdims(distribution, dims = 1)

    # Do the last normalisation and check
    colsum_distribution = sum(distribution)

    if abs(colsum_distribution - 1) > delta
        distribution = distribution ./ colsum_distribution
    end

    # Find partition elements with strictly positive measure.
    δ = tolerance/size(to.transfermatrix, 1)
    inds_nonzero = findall(distribution .> δ)

    # Extract the elements of the invariant measure corresponding to these indices
    return InvariantMeasure(to, Probabilities(distribution))
end

function invariantmeasure(x::AbstractDataset, ϵ::RectangularBinning)
    to = transferoperator(x, ϵ)
    invariantmeasure(to)
end

function probabilities(x::AbstractDataset, est::TransferOperator)
    to = transferoperator(x, est.ϵ)
    iv = invariantmeasure(to)
    return iv.ρ
end

"""
    transfermatrix(iv::InvariantMeasure) → (M::AbstractArray{<:Real, 2}, bins::Vector{<:SVector})

Return the transfer matrix/operator and corresponding bins. Here, `bins[i]` corresponds
to the i-th row/column of the transfer matrix. Thus, the entry `M[i, j]` is the
probability of jumping from the state defined by `bins[i]` to the state defined by
`bins[j]`.

See also: [`TransferOperator`](@ref).
"""
function transfermatrix(iv::InvariantMeasure)
    return iv.to.transfermatrix, iv.to.bins
end