text:
### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ f1c9d31f-c565-4e96-8801-22352b014397
import Pkg

# ╔═╡ f6daaa32-fe10-4a20-b30e-ba31b63c6b90
Pkg.add("Distributions")

# ╔═╡ 338200b2-c10c-4e7d-8a24-bfd32e6757e5
Pkg.add("DataFrames")

# ╔═╡ 525ba450-c635-4bb9-a2f3-aeb12d5c30d3
Pkg.add("CSV")

# ╔═╡ 64adb0f4-b695-4be1-9b02-c42562951b39
Pkg.add("Plots")

# ╔═╡ 55405eea-1c12-4060-81f5-cbcc082a992c
Pkg.add("IterTools")

# ╔═╡ 7e392b1a-31a8-4364-ac93-742221501f1c
Pkg.add("PlutoUI")

# ╔═╡ b708f1d2-e8dd-40ee-92fd-cb4ae02a6a9a
Pkg.add("StatsBase")

# ╔═╡ 42c0a736-b762-40e3-a4fb-46b6303f0cff
Pkg.add("FFTW")

# ╔═╡ 34a840a1-7050-4719-a81e-c4545551ec7d
using Random, Distributions, StatsBase

# ╔═╡ 83920316-034e-4ca1-94fa-5aea4d4d117b
using LinearAlgebra

# ╔═╡ 8287cc85-37d3-4b53-ab65-0aef5d996a72
using Plots

# ╔═╡ 1196f179-17ff-42d9-a582-1b8498b7c372
using IterTools

# ╔═╡ 2d8d5528-d64d-4026-b5a6-93b66e17e609
using PlutoUI

# ╔═╡ c3929637-0224-4f95-a8ce-b1a159084014
using DataFrames

# ╔═╡ 2f28edec-8461-4eda-a95f-05f4db5b3a81
using CSV

# ╔═╡ 258b3c88-b1ca-49b2-b414-8ce8db178724
using FFTW

# ╔═╡ 0a435f6c-9f0a-11eb-114f-fb8c733dafbb
md"# Brute Force Simulation"

# ╔═╡ 5c181bf0-f92e-40ac-9193-62c48927f60f
md"""
## Simulation Configuration
"""

# ╔═╡ 0fe833e8-a359-41e2-ba4a-b8313a222182
md"""
### Classical field parameters

$(@bind bigN NumberField(1:10000;default=100))

Number of emitting atoms

$(@bind temp NumberField(293:3000;default=400))

Temperature (K)

$(@bind tres NumberField(10:500))

Time resolution (ps)

$(@bind tmax NumberField(1.0:0.1:50.0;default=10))

Maximum time (ns)

"""

# ╔═╡ a8b59a40-e281-4e59-97ae-c523be21e9eb
md"""
$(@bind shift Slider(1:30,show_value=true ))

Line separation (GHz)
"""

# ╔═╡ c0e353c3-53e7-4419-b8cf-878f6fbc761d
md"""
### Photon count parameters

$(@bind nbar NumberField(1.0:1000.0;default=10))

Expected counts in time tmax

$(@bind ntot NumberField(1:1000000;default=1000))

Total photon counts overall (approximate)
"""

# ╔═╡ 659c5d8a-6722-4c74-834a-f79bc4312e8e

md"""
Make separate instances of fields: $(@bind makeInstances CheckBox())
"""

# ╔═╡ 402b4a2a-d22f-4ee7-ae3d-ca5fe7fda110
md"""
Number of classical simulations: $(convert(Integer,ceil(ntot/nbar)))
"""

# ╔═╡ 874462ef-e735-4f19-bbf9-ee2d2401eb15
md"""
Set seed $\rightarrow$ $(@bind setSeed CheckBox())
"""

# ╔═╡ 095ebc52-38d8-46d7-a5af-d2de0719ea92
if setSeed
	md"""
	$(@bind seed NumberField(1:1000000;default=42))
	"""
end

# ╔═╡ 89b17676-6496-4d37-929d-3f626e329cdd
md"""
#### Run Simulation $\rightarrow$ [$(@bind sim CheckBox())]
"""

# ╔═╡ 30f9f516-568b-4e73-8aa2-4f5bb48f9ab4
md"""
## Plots
"""

# ╔═╡ 2f0b2664-8969-4712-b26d-fbaa06b36ae3
if sim
md"""
### Plots from simulated photon counts
"""
end

# ╔═╡ c48c9d04-ca21-4f2e-96c5-3603a729900c
md"""
**Window:**
"""

# ╔═╡ d512f123-088b-4a65-965e-af512147249d
if sim
	md"""
	### Plots using classical field
	"""
end

# ╔═╡ 13e70af4-146d-43e7-9fd6-ec8f901185d4
md"""
## Calculations
"""

# ╔═╡ 807e056a-3c89-46f9-8dc2-5dde04023ce2
md"""
### Classical calculations
"""

# ╔═╡ 93a1cad0-6b86-40a0-b122-19b854f51174
# define beamsplitter
begin
	bs = [1,1]
	bs = bs/norm(bs)
end;

# ╔═╡ 282a7d64-8d67-4a10-9aa9-d64be7c5dd54
md"""
## Functions
"""

# ╔═╡ c84cf75f-17da-4d54-9e97-9bf69968732c
md"""
### Electric field parameters and generation
"""

# ╔═╡ 26e08020-6775-431f-b33a-0efaad94189a
begin
	"""
		eFieldParams(Em::Vector,ωm::Vector,ω0::Number,σ::Number,rng::MersenneTwister)

	Static parameters for the electric field
	"""
	struct eFieldParams
		Em::Vector # field magnitudes
		ωm::Vector # central frequencies of lines
		ω0::Number # central frequency for composite line
		σ::Number # Doppler broadening
		rng::MersenneTwister
		function eFieldParams(Em::Vector,ωm::Vector,ω0::Number,σ::Number,rng::MersenneTwister)
			@assert length(Em) == length(ωm) "Length of field magnitude and frequency vectors must match"
			new(
				convert(Vector{Complex},Em),
				convert(Vector{Real},ωm),
				convert(Real,ω0),
				convert(Real,σ),
				rng
			)
		end
	end

	"""
		eFieldParams(Em::Vector,ωm::Vector,σ::Number,seed::Integer=-1)

	Static parameters for the electric field
	"""
	function eFieldParams(Em::Vector,ωm::Vector,σ::Number,seed::Integer = -1)
		if seed == -1
			rng = MersenneTwister()
		else
			rng = MersenneTwister(seed)
		end
		return eFieldParams(Em,ωm,ωm[1],σ,rng)
	end
end

# ╔═╡ e51b2eac-2de8-4f30-bc11-26c320bd6afb
"""
	function seed!(params::eFieldParams,seed::Integer)

Reset the seed for the random number generator in params
"""
function seed!(params::eFieldParams,seed::Integer)
	Random.seed!(params.rng,seed)
end

# ╔═╡ 904ec626-6678-4f19-8247-e404fe9d278b
begin
	"""
	Container holding frequencies and phases for one realization of the electric field.
	"""
	struct eFieldInstance
		ωn::Vector
		ϕmn::Matrix
		params::eFieldParams

		function eFieldInstance(ωn::Vector,ϕmn::Matrix,params::eFieldParams)
			@assert (size(ϕmn)[1] == length(params.ωm) && size(ϕmn)[2] == length(ωn)) "Phase array must have shape (m,n)"
			new(ωn,ϕmn,params)
		end
	end

	"""
		eFieldInstance(n::Integer,params::eFieldParams)

	Generate frequencies and phases for the electric field from n emitting atoms and return an eFieldInstance object.
	"""
	function eFieldInstance(n::Integer,params::eFieldParams)
		ωDist = Normal(params.ω0,params.σ)
		ωn    = rand(params.rng,ωDist,n)
		ϕmn   = 2*π*rand(params.rng,Float64,(length(params.ωm),n))
		return eFieldInstance(ωn,ϕmn,params)
	end
end

# ╔═╡ 8175ac47-b261-4942-a70d-f127f22aee74
"""
	function electricField(t::Real,params::eFieldParams)

Returns the electric field value at time t
"""
function electricField(t::Real,instance::eFieldInstance)
	Δm = instance.params.ωm .- instance.params.ω0
	# generate frequencies
	ωmn = transpose(instance.ωn) .+ Δm
	# add the phase
	exponentmn = -im*(t*ωmn+instance.ϕmn)
	# put them in the exponent
	enm = exp.(exponentmn)
	# multiply by the field magnitude
	fieldnm = instance.params.Em .* enm
	# add it all together
	return sum(ivec(fieldnm))
end

# ╔═╡ 58b00c01-9dca-4220-97a6-5519d7672bd5
"""
	function σTemp(ω0::Real,temp::Number)

Returns the standard deviation due to Doppler broadening of frequency ω0 at temp
"""
function σTemp(ω0::Real,temp::Number)
	kbOverMhC2 = 9.178e-14;
	return sqrt(kbOverMhC2*temp)*ω0
end

# ╔═╡ 2f81c1e8-b9bd-4b41-8d9f-f54f964518cf
md"""
### Classical Intensity
"""

# ╔═╡ 5b656847-12d5-442d-8c2a-bbdbe3de0ab7
begin
	function intensity(eFieldT::Number)
		return real(eFieldT*conj(eFieldT))
	end

	function intensity(t::Real,instance::eFieldInstance)
		eFieldT = electricField(t,instance)
		return intensity(eFieldT)
	end
end

# ╔═╡ b409aae3-7ea5-4278-b56a-463718b40426
md"""
### Classical Correlations
"""

# ╔═╡ a8a6e7d0-d871-40e4-ac37-9685cd9e001b
md"""
- Our calculated version

$g^{(2)}(\tau) =1-\frac{1}{N}\frac{\sum_{m=1}^M|\mathcal{E}|_m^4}{\left(\sum_{m=1}^M|\mathcal{E}_m|^2\right)^2}+\left|\frac{\sum_{m=1}^M|\mathcal{E}_m|^2e^{-i\Delta_m\tau}}{\sum_{m=1}^M|\mathcal{E}|_m^2}\right|^2\frac{\langle S(\tau)\rangle_\omega}{N^2}$
"""

# ╔═╡ 4bf8740c-b840-4e3f-9413-27490cc9f79b
begin
	function stauAvg(τ::Number,params::eFieldParams,n::Integer)
		return stauAvg(τ,params.σ,n)
	end

	function stauAvg(τ::Number,σ::Number,n::Integer)
		term1 = n
		term2 = n*(n-1)
		term2 *= exp(-σ^2*τ^2)
		return term1 + term2
	end
end

# ╔═╡ 9df57654-c599-4313-87be-3e1beeb82448
begin
	function stauVar(τ::Number,params::eFieldParams,n::Integer)
		return stauVar(τ,params.σ,n)
	end

	function stauVar(τ::Number,σ::Number,n::Integer)
		στ2 = σ^2*τ^2

		prod1 = 8*n*(n-1)
		prod1 *= exp(-2*στ2)

		prod2 = n-1+cosh(στ2)

		prod3 = sinh(στ2/2)^2

		return prod1*prod2*prod3
	end
end

# ╔═╡ fd228d30-48e2-4b8f-96a4-1c846743686a
begin
	function stau(τ::Number,instance::eFieldInstance)
		return stau(τ,instance.ωn)
	end

	function stau(τ::Number,ωn::Vector)
		terms = exp.(-im*τ*ωn)
		sumterms = sum(terms)
		return real(sumterms*conj(sumterms))
	end
end

# ╔═╡ 637c1989-e452-4eea-9a11-3d94f3e2f6f5
function g2Calc(τ::Number,n::Integer,params::eFieldParams)
	em2 = real.(params.Em .* conj.(params.Em))
	em4 = em2 .* em2
	sumEm2 = sum(em2)
	sumEm4 = sum(em4)

	g2τ = 1.0

	term2 = -sumEm4/(n*sumEm2^2)

	g2τ += term2

	Δm = params.ωm .- params.ω0
	term3 = sum(em2 .* exp.(-im*τ*Δm))/sumEm2
	term3 *= conj(term3)
	term3 = real(term3)
	term3 *= stauAvg(τ,params,n)/n^2

	return g2τ+term3
end

# ╔═╡ f2e2b9fd-9d21-410c-89d8-2286b163bb2a
# these parameters are needed for all simulations
begin
	# Balmer-α lines specified here
	# ωM = 2*π*[456811.0, 456812.0]
	ωM = 2*π*[456811.0, 456811.0+shift]
	# magnitude of each line
	mag = convert(Vector{ComplexF64},ones(length(ωM)))
	# Doppler broadening
	σDopp = σTemp(ωM[1],temp)
	# Store parameters
	eParams = eFieldParams(mag,ωM,σDopp)

	if setSeed
		seed!(eParams,seed)
	end
	# generate times in tres ps intervals up to 2*tmax
	times = collect(0:tres*1e-3:2*tmax);
	# limit the window to tmax to avoid correlation cutoff
	window = convert(Integer,floor(length(times)/2));
	# τ is just the times up to our window
	τ = times[1:window];

	g2τCalc = map(tau->g2Calc(tau,bigN,eParams),τ)
end;

# ╔═╡ 7ac53f6b-15ee-440d-acae-005dc3acce51
md"""
f0 (GHz) = $(ωM[1]/(2*π))

Δf (GHz) = $(round((ωM[2]-ωM[1])/(2*π),digits=2))

σ (GHz) = $(round(sqrt(9.178e-14*temp)*ωM[1]/(2*π),digits=2))
"""

# ╔═╡ a3cd4171-5bc5-4988-a5ef-428dd6fa31ab
if sim
	md"""
	Size: $(@bind windowSize Slider(1:window,default=window))
	"""
end

# ╔═╡ e2fab25a-ccde-44e8-a3bf-e402fc32d2ba
if sim
	md"""
	Position: $(@bind windowPos Slider(0:(window-windowSize)))
	"""
end

# ╔═╡ b61954a4-a0ba-4f88-90af-0baa98d1325c
if sim
	windowRange = 1+windowPos:windowSize+windowPos
end;

# ╔═╡ 8b607526-47b0-4a8a-aa16-2aee351b590f
if sim
	md"""
	**Window:**

	Window size (ns) = $(tres*windowSize/1000)

	Window position (ns) = $(windowPos*tres/1000)
	"""
end

# ╔═╡ 17d19e10-604f-4a5f-9ed9-73bca2f68d4b
md"""
### Fourier transform related functions for classical case
"""

# ╔═╡ 664ede52-c0c5-4591-b94b-e11abd2de1ff
function fftPositiveFreq(ft,f)
	@assert length(ft) == length(f) "Fourier transform and frequency vectors much have matching length"
	fftRange = f .>= 0
	return (ft[fftRange],f[fftRange])
end

# ╔═╡ 3c7feb92-1223-4428-a86c-0b46960df59f
function g2ClassicalFFT(g2τ::Vector{S},cuts::Tuple{T,T}) where {S<:Real, T<:Integer}
	@assert length(τ) == length(g2τ) "Time and coherence arrays must have the same length"
	g2τPrep = g2τ[1+cuts[1]:end-cuts[2]]
	g2τPrepMean = mean(g2τPrep)
	g2τPrep = g2τPrep .- g2τPrepMean
	g2τFFT = fft(g2τPrep)
end

# ╔═╡ fe3043d0-5417-4a0a-90bc-03f82186edd9
function g2ClassicalFreq(tres::Real,τ::Vector,cuts::Tuple{T,T}) where {T<:Integer}
	τlen = length(τ) - sum(cuts)
	freqFFT = fftfreq(τlen,1/tres)
end

# ╔═╡ 9c2a61f6-dc60-4d35-bfaf-45aad4de2905
md"""
### General correlation functions
"""

# ╔═╡ b237ca1e-1719-420c-820d-67b8e18c7a57
"""
	function correlate(u::Vector{T},v::Vector{T},offset::Integer,window::Integer = -1) where {T<:Number}

Calculates correlation between vectors u and v with given offset. Specify averaging window to limit range of correlation. If the window extends beyond the end of one vector, it treats out-of-bounds indices as zero.
"""
function correlate(u::Vector{T},v::Vector{T},offset::Integer,window::Integer = -1) where {T<:Number}

	@assert offset <= length(u) "Offset out of bounds"
	@assert window <= length(u) && window <= length(v) "Window must be smaller than input vector lengths"

	if window == -1
		window = length(u)
	end

	v1 = view(u,1:window)
	v2 = view(v,1+offset:min(window+offset,length(v)))
	if window+offset > length(v)
		v2 = vcat(v2,zeros(window+offset-length(v)))
	end

	return dot(v1,v2)/window
end

# ╔═╡ 3620f708-3eac-4cd8-90db-72ac383b5917
"""
	function autocorrelate(u::Vector{T},offset::Integer, window::Integer = -1) where {T<:Number}

Calculates correlation of vector u with itself.
"""
function autocorrelate(u::Vector{T},offset::Integer, window::Integer = -1) where {T<:Number}
	correlate(u,u,offset,window)
end

# ╔═╡ 361cec13-e752-4166-8cd3-d74cb532dd56
md"""
### Photon counts from intensity
"""

# ╔═╡ 48fd09eb-f5a0-4bc4-bdbf-3648366978a1
"""
	function γIntensity(intensity::Vector,nbar::Real)

Calculates the photon count rate in each bin of an intensity histogram
"""
function γIntensity(intensity::Vector,nbar::Real)
	nintensity = intensity/sum(intensity)
	return nbar*nintensity
end

# ╔═╡ 933ae343-2226-4ce0-a9e7-9bc581b07de2
md"""
### Photon correlations
"""

# ╔═╡ c2b7737e-faab-4ed4-9b03-0e1286c635cd
"""
	function autocorrTimes(τ::Vector,γCounts::Vector)

Returns an array of τ values for which the γCounts autocorrelation is non-zero.
"""
function autocorrTimes(τ::Vector,γCounts::Vector)
	return τ[map(i->autocorrelate(γCounts,i,length(τ)) > 0 ? true : false,collect(0:length(τ)-1))]
end

# ╔═╡ 43c5fdbe-823d-4b1e-8801-e41e467ffeef
"""
	function corrTimes(τ::Vector,γCounts::Vector)

Returns an array of τ values for which the γCounts autocorrelation is non-zero.
"""
function corrTimes(τ::Vector,γCounts1::Vector,γCounts2::Vector)
	@assert length(γCounts1) == length(γCounts2) "Count vectors must have the same length"
	return τ[map(i->correlate(γCounts1,γCounts2,i,length(τ)) > 0 ? true : false,collect(0:length(τ)-1))]
end

# ╔═╡ 0abe72e2-92c4-4752-856f-b5975ab4c5e3
"""
	function countDeltaTimes(τ::Vector,γCounts::Vector)

Returns an array of τ values for which the γCounts autocorrelation is non-zero.
"""
function countTimes(times::Vector,γCounts::Vector)
	out = Vector{Real}(undef,0)
	for (i,counts) in enumerate(γCounts)
		if counts != 0
			countTimes = times[i]*ones(counts)
			out = vcat(out,countTimes)
		end
	end
	return out
end

# ╔═╡ 727d4c4c-d374-4b3b-8f4b-8e73b27edc38
md"""
### Count generators
"""

# ╔═╡ 31b92f9f-6be0-48cb-9a7e-bfc5edb79bd2
"""
	function poissonCount(nbar::Real)

Returns Poisson distributed counts for average count rate nbar
"""
function poissonCount(nbar::Real)
	p = exp(-nbar)
	s = p
	r = rand()
	count = 0
	while r > s
		count += 1
		p *= nbar/count
		s += p
	end
	return count
end

# ╔═╡ 18f7bc41-751f-4283-8a90-af3ad217b58b
# This is the main calculation cell for the simulation
if sim
	# low cut on τ (indices from start)
	τstart =100
	# high cut on τ (indices from end)
	τend = 0

	# generate frequency bins of fourier transform
	allfreqs = g2ClassicalFreq(tres*1e-3,τ,(τstart,τend))

	# array to accumulate g2(τ) Fourier transform
	g2τ1FFTsum = zeros(ComplexF64,length(τ)-τstart)

	# array to accumulate all photon counts
	totCounts = zeros(Integer,length(times))

	# concatenation of all photon correlation times
	allCorrTimes = Vector{Float64}[]
	γCorrTimes = Vector{typeof(allCorrTimes)}[]

	### Calculate initial instance of electric field ###

	# instantiate electric field
	eInstance1 = eFieldInstance(bigN,eParams)

	# calculate electric field vs time
	eFieldT1 = map(t->electricField(t,eInstance1),times)

	# apply beam splitter
	eFieldT1Beam = bs[1] * eFieldT1

	# calculate intensity
	intensity1Beam = intensity.(eFieldT1Beam)

	# calculate classical g2τ
	g2τ1Norm = mean(intensity1Beam)^2
	g2τ1 = map(i->autocorrelate(intensity1Beam,i,window),collect(0:window-1))/g2τ1Norm

	# calculate Fourier transform of this instance
	g2τ1FFT = g2ClassicalFFT(g2τ1,(τstart,τend))

	# accumulate Fourier transform sum
	g2τ1FFTsum .+= g2τ1FFT

	# number of instances
	trials = convert(Integer,ceil(ntot/nbar))

	# loop over instances
	# for k = 1:γtrials
	for n = 1:trials

		# calculate the average photon counts in each time bin from the beam intensity and the overall average photon count rate
		γavg1Beam = γIntensity(intensity1Beam,nbar/2)

		# generate counts for each beam
		γcounts1Beam1 = poissonCount.(γavg1Beam)
		γcounts1Beam2 = poissonCount.(γavg1Beam)

		# concatenate the correlations from this run into all others
		allCorrTimes = vcat(allCorrTimes,corrTimes(τ,γcounts1Beam1,γcounts1Beam2))

		# add counts from both beams
		totCounts .+= γcounts1Beam1
		totCounts .+= γcounts1Beam2

		# reinstantiate if desired
		if makeInstances && n != trials
			# reinstantiate electric field
			eInstance1 = eFieldInstance(bigN,eParams)

			# calculate electric field vs time
			eFieldT1 = map(t->electricField(t,eInstance1),times)

			# apply beam splitter
			eFieldT1Beam = bs[1] * eFieldT1

			# calculate intensity
			intensity1Beam = intensity.(eFieldT1Beam)

			# calculate classical g2τ
			g2τ1Norm = mean(intensity1Beam)^2
			g2τ1 = map(i->autocorrelate(intensity1Beam,i,window),collect(0:window-1))/g2τ1Norm

			# calculate Fourier transform of this instance
			g2τ1FFT = g2ClassicalFFT(g2τ1,(τstart,τend))

			# accumulate Fourier transform sum
			g2τ1FFTsum .+= g2τ1FFT
		end
	end
	# end

	g2τ1FFTsinglePos,allfreqsPos = fftPositiveFreq(g2τ1FFT,allfreqs)
	g2τ1FFTsumPos,_ = fftPositiveFreq(g2τ1FFTsum,allfreqs)

end;

# ╔═╡ 4a6763e9-46c3-470e-8712-2e4aa65493cf
if sim
	γCountsPlot = bar(times[windowRange],totCounts[windowRange],label=false)
	xlabel!(γCountsPlot,"time (ns)")
	title!(γCountsPlot,"Total photon counts = $(sum(totCounts))")

end

# ╔═╡ 5fdf6412-ab5e-434b-b5df-3c16fc98216d
# make electric field plot
if sim
	eFieldPlot = plot(times[windowRange],real.(eFieldT1Beam[windowRange]),label=false)
	xlabel!(eFieldPlot,"t (ns)")
	title!(eFieldPlot,"E(t)")
end

# ╔═╡ dde1cc54-1ae5-4c0f-8762-b1cf3e2f30e9
# make intensity plot
if sim
	intensityPlot = plot(times[windowRange],intensity1Beam[windowRange],label=false)
	xlabel!(intensityPlot,"t (ns)")
	title!(intensityPlot,"I(t)")
end

# ╔═╡ 0591df60-f054-4136-819b-909bad833ed6
if sim
	g2τPlot = plot(τ[windowRange],g2τ1[windowRange],label="Simulated")
	plot!(τ[windowRange],g2τCalc[windowRange],label="Theoretical (noiseless)")
	xlabel!(g2τPlot,"τ (ns)")
	title!(g2τPlot,"\$g^{(2)}(\\tau)\$")
end

# ╔═╡ 94f06e2a-bf0f-4d78-b3a9-83e7865d05a1
if sim
	g2τFFTplot = plot(allfreqsPos,abs.(g2τ1FFTsinglePos),label=false)
	title!(g2τFFTplot,"\$\\hat{g}^{(2)}(\\nu)\$")
	xlabel!(g2τFFTplot,"frequency (GHz)")
	vline!(g2τFFTplot,[shift],label="Frequency shift",ls=:dash)
end

# ╔═╡ a2d886d7-f030-4d38-860b-399e501b24fc
if sim && makeInstances
	g2τFFTsumPlot = plot(allfreqsPos,abs.(g2τ1FFTsumPos),label=false)
	title!(g2τFFTsumPlot,"\$\\sum_{i=1}^{$(convert(Integer,ceil(ntot/nbar)))}\\hat{g}_i^{(2)}(\\nu)\$")
	xlabel!(g2τFFTsumPlot,"frequency (GHz)")
	vline!(g2τFFTsumPlot,[shift],label="Frequency shift",ls=:dash)

end

# ╔═╡ dd5ba54e-5049-414c-af9f-f8cb1925b662
# make photon count correlation histogram and take fourier transform
if sim

	# bin correlation times into a histogram
	corrHist = fit(Histogram,allCorrTimes,vcat(τ,τ[end]+tres*1e-3),closed=:left)

	# normalize the histogram
	normCorr = corrHist.weights/sum(corrHist.weights)

	# subtract mean from histogram to reduce constant term
	normAvgCorr = normCorr .- mean(normCorr)

	# take the fourier transform
	γfft = fft(normAvgCorr)

	# recover fourier transform frequencies
	γfreqs = fftfreq(length(γfft),1/(tres*1e-3))

	# select positive frequencies for plotting
	γfftPos,γfreqsPos = fftPositiveFreq(γfft,γfreqs)
end;

# ╔═╡ 34a9830b-8b4a-441d-a1db-20ae72565ef5
if sim
	md"""
	Size: $(@bind fwindowSize Slider(1:length(γfreqsPos),default=length(γfreqsPos)-10))
	"""
end

# ╔═╡ 92980437-333e-4205-9b53-9de074a849d4
if sim
	md"""
	Position: $(@bind fwindowPos Slider(0:(length(γfreqsPos)-fwindowSize),default=5))
	"""
end

# ╔═╡ 730295a5-5217-4a80-a327-b0d61a88165a
if sim
	fwindow = 1+fwindowPos:fwindowPos+fwindowSize
end;

# ╔═╡ 9f1f5473-1acd-4d13-a7aa-e88dcdf2420a
if sim
	γCorrFreqPlot = plot(γfreqsPos[fwindow],abs.(γfftPos)[fwindow],label = false)
	xlabel!(γCorrFreqPlot,"frequency (GHz)")
	title!(γCorrFreqPlot,"Fourier transform of photon correlations")
	vline!(γCorrFreqPlot,[shift],label="Frequency shift",ls=:dash)
end

# ╔═╡ 50c3ba83-89fd-4211-a3d1-1439e6a7cb5e
if sim
	γCorrPlot = plot(corrHist,label=false)
	xlabel!(γCorrPlot,"τ (ns)")
	title!(γCorrPlot,"Correlated counts vs correlation time")

end

# ╔═╡ 3edfcf46-9f6d-4ff0-8179-f9a7ddb1897f
"""
	function beCount(nbar::Real)

Returns Bose-Einstein distributed counts for average count rate nbar
"""
function beCount(nbar::Real)
	p = 1/(nbar+1)
	fnbar = p*nbar
	f = p*nbar
	s = p
	r = rand()
	count = 0
	while r>s
		count += 1
		p *= f
		s += p
	end
	return count
end

# ╔═╡ 72cbac0a-ae8a-4296-ba78-ab5696539a5a
md"""
### Helping functions
"""

# ╔═╡ 2328359b-8f27-4a18-8d63-0538dbdfcdff
function vectorAvg(someVector::Vector)
	return +(someVector...)/length(someVector)
end

# ╔═╡ 57b0a316-782e-4f85-8a77-8743abc601f1
md"""
## Load Prerequisites
"""

# ╔═╡ Cell order:
# ╟─0a435f6c-9f0a-11eb-114f-fb8c733dafbb
# ╟─5c181bf0-f92e-40ac-9193-62c48927f60f
# ╟─0fe833e8-a359-41e2-ba4a-b8313a222182
# ╟─a8b59a40-e281-4e59-97ae-c523be21e9eb
# ╟─7ac53f6b-15ee-440d-acae-005dc3acce51
# ╟─c0e353c3-53e7-4419-b8cf-878f6fbc761d
# ╟─659c5d8a-6722-4c74-834a-f79bc4312e8e
# ╟─402b4a2a-d22f-4ee7-ae3d-ca5fe7fda110
# ╟─874462ef-e735-4f19-bbf9-ee2d2401eb15
# ╟─095ebc52-38d8-46d7-a5af-d2de0719ea92
# ╟─89b17676-6496-4d37-929d-3f626e329cdd
# ╟─30f9f516-568b-4e73-8aa2-4f5bb48f9ab4
# ╟─2f0b2664-8969-4712-b26d-fbaa06b36ae3
# ╟─730295a5-5217-4a80-a327-b0d61a88165a
# ╟─c48c9d04-ca21-4f2e-96c5-3603a729900c
# ╠═34a9830b-8b4a-441d-a1db-20ae72565ef5
# ╠═92980437-333e-4205-9b53-9de074a849d4
# ╠═9f1f5473-1acd-4d13-a7aa-e88dcdf2420a
# ╠═50c3ba83-89fd-4211-a3d1-1439e6a7cb5e
# ╠═4a6763e9-46c3-470e-8712-2e4aa65493cf
# ╠═d512f123-088b-4a65-965e-af512147249d
# ╠═b61954a4-a0ba-4f88-90af-0baa98d1325c
# ╠═8b607526-47b0-4a8a-aa16-2aee351b590f
# ╠═a3cd4171-5bc5-4988-a5ef-428dd6fa31ab
# ╠═e2fab25a-ccde-44e8-a3bf-e402fc32d2ba
# ╠═5fdf6412-ab5e-434b-b5df-3c16fc98216d
# ╠═dde1cc54-1ae5-4c0f-8762-b1cf3e2f30e9
# ╠═0591df60-f054-4136-819b-909bad833ed6
# ╠═94f06e2a-bf0f-4d78-b3a9-83e7865d05a1
# ╠═a2d886d7-f030-4d38-860b-399e501b24fc
# ╟─13e70af4-146d-43e7-9fd6-ec8f901185d4
# ╟─807e056a-3c89-46f9-8dc2-5dde04023ce2
# ╠═f2e2b9fd-9d21-410c-89d8-2286b163bb2a
# ╠═93a1cad0-6b86-40a0-b122-19b854f51174
# ╠═18f7bc41-751f-4283-8a90-af3ad217b58b
# ╠═dd5ba54e-5049-414c-af9f-f8cb1925b662
# ╟─282a7d64-8d67-4a10-9aa9-d64be7c5dd54
# ╟─c84cf75f-17da-4d54-9e97-9bf69968732c
# ╠═26e08020-6775-431f-b33a-0efaad94189a
# ╠═e51b2eac-2de8-4f30-bc11-26c320bd6afb
# ╠═904ec626-6678-4f19-8247-e404fe9d278b
# ╠═8175ac47-b261-4942-a70d-f127f22aee74
# ╠═58b00c01-9dca-4220-97a6-5519d7672bd5
# ╟─2f81c1e8-b9bd-4b41-8d9f-f54f964518cf
# ╠═5b656847-12d5-442d-8c2a-bbdbe3de0ab7
# ╟─b409aae3-7ea5-4278-b56a-463718b40426
# ╟─a8a6e7d0-d871-40e4-ac37-9685cd9e001b
# ╠═4bf8740c-b840-4e3f-9413-27490cc9f79b
# ╠═9df57654-c599-4313-87be-3e1beeb82448
# ╠═fd228d30-48e2-4b8f-96a4-1c846743686a
# ╠═637c1989-e452-4eea-9a11-3d94f3e2f6f5
# ╟─17d19e10-604f-4a5f-9ed9-73bca2f68d4b
# ╠═664ede52-c0c5-4591-b94b-e11abd2de1ff
# ╠═3c7feb92-1223-4428-a86c-0b46960df59f
# ╠═fe3043d0-5417-4a0a-90bc-03f82186edd9
# ╟─9c2a61f6-dc60-4d35-bfaf-45aad4de2905
# ╠═b237ca1e-1719-420c-820d-67b8e18c7a57
# ╠═3620f708-3eac-4cd8-90db-72ac383b5917
# ╟─361cec13-e752-4166-8cd3-d74cb532dd56
# ╠═48fd09eb-f5a0-4bc4-bdbf-3648366978a1
# ╟─933ae343-2226-4ce0-a9e7-9bc581b07de2
# ╠═c2b7737e-faab-4ed4-9b03-0e1286c635cd
# ╠═43c5fdbe-823d-4b1e-8801-e41e467ffeef
# ╠═0abe72e2-92c4-4752-856f-b5975ab4c5e3
# ╟─727d4c4c-d374-4b3b-8f4b-8e73b27edc38
# ╠═31b92f9f-6be0-48cb-9a7e-bfc5edb79bd2
# ╠═3edfcf46-9f6d-4ff0-8179-f9a7ddb1897f
# ╟─72cbac0a-ae8a-4296-ba78-ab5696539a5a
# ╠═2328359b-8f27-4a18-8d63-0538dbdfcdff
# ╟─57b0a316-782e-4f85-8a77-8743abc601f1
# ╠═f1c9d31f-c565-4e96-8801-22352b014397
# ╠═f6daaa32-fe10-4a20-b30e-ba31b63c6b90
# ╠═338200b2-c10c-4e7d-8a24-bfd32e6757e5
# ╠═525ba450-c635-4bb9-a2f3-aeb12d5c30d3
# ╠═64adb0f4-b695-4be1-9b02-c42562951b39
# ╠═55405eea-1c12-4060-81f5-cbcc082a992c
# ╠═7e392b1a-31a8-4364-ac93-742221501f1c
# ╠═b708f1d2-e8dd-40ee-92fd-cb4ae02a6a9a
# ╠═42c0a736-b762-40e3-a4fb-46b6303f0cff
# ╠═34a840a1-7050-4719-a81e-c4545551ec7d
# ╠═83920316-034e-4ca1-94fa-5aea4d4d117b
# ╠═8287cc85-37d3-4b53-ab65-0aef5d996a72
# ╠═1196f179-17ff-42d9-a582-1b8498b7c372
# ╠═2d8d5528-d64d-4026-b5a6-93b66e17e609
# ╠═c3929637-0224-4f95-a8ce-b1a159084014
# ╠═2f28edec-8461-4eda-a95f-05f4db5b3a81
# ╠═258b3c88-b1ca-49b2-b414-8ce8db178724

max_stars_repo_name: greenspaceexplorer/SpeckleSim

max_issues_repo_name: greenspaceexplorer/SpeckleSim

max_forks_repo_name: greenspaceexplorer/SpeckleSim


text:
#=
 Example 2: Find the global distribution of a "water mass" defined by an oceanographically-relevant surface region.
 Steps: (a) define the water mass 1). by a pre-defined surface
            dyed with passive tracer concentration of 1,
        (b) propagate the dye with the matrix A, with the result
            being the fraction of water originating from the
            surface region.
 See Section 2b of Gebbie & Huybers 2010, esp. eqs. (15)-(17).
=#
using Revise
using TMI, BenchmarkTools, PyPlot, PyCall

#TMIversion = "modern_180x90x33_GH10_GH12"
TMIversion = "modern_90x45x33_GH10_GH12"
A, Alu, γ, TMIfile, L, B = config_from_nc(TMIversion)

# choose water mass (i.e., surface patch) of interest
# Enhancement: provide list of choices
list = ("GLOBAL","ANT","SUBANT",
            "NATL","NPAC","TROP","ARC",
            "MED","ROSS","WED","LAB","GIN",
            "ADEL","SUBANTATL","SUBANTPAC","SUBANTIND",
            "TROPATL","TROPPAC","TROPIND")

region = list[2]

# do numerical analysis
g = watermassdistribution(TMIversion,Alu,region,γ);

# plot a section at 330 east longitude (i.e., 30 west)
lon_section = 330 # only works if exact
lims = 0:5:100
tlabel = region * " water-mass fraction [%]"
sectionplot(100g,lon_section,lims,titlelabel = tlabel)

ext: jl


text:
# Balsa-Canto, E., Alonso, A. A., Banga, J. R.,
# An iterative identification procedure for dynamic modeling of biochemical networks
# https://doi.org/10.1186/1752-0509-4-11
# The model is out of reach at the moment but can be analyzed by SIAN (https://github.com/pogudingleb/SIAN)
using Logging

using StructuralIdentifiability

logger = Logging.SimpleLogger(stdout, Logging.Info)
global_logger(logger)

ode = @ODEmodel(
    x1'(t) = k_prod  - k_deg * x1(t) - k1 * x1(t) * u(t),
    x2'(t) = -k3 * x2(t) - k_deg * x2(t) - a2 * x2(t) * x10(t) + t1 * x4(t) - a3 * x2(t) * x13(t) + t2 * x5(t) + (k1 * x1(t) - k2 * x2(t) * x8(t)) * u(t),
    x3'(t) = k3 * x2(t) - k_deg * x3(t) + k2 * x2(t) * x8(t) * u(t),
    x4'(t) = a2 * x2(t) * x10(t) - t1 * x4(t),
    x5'(t) = a3 *  x2(t) * x13(t) - t2 *  x5(t),
    x6'(t) = c6a * x13(t) - a1 * x6(t) * x10(t) + t2 * x5(t) - i1 * x6(t),
    x7'(t) = i1 * kv * x6(t) - a1 * x11(t) *  x7(t),
    x8'(t) = c4 * x9(t) - c5 * x8(t),
    x9'(t) = c2 - c1 * x7(t) - c3 * x9(t),
    x10'(t) = -a2 * x2(t) * x10(t) - a1 * x10(t) * x6(t) + c4a * x12(t) - c5a * x10(t) - i1a * x10(t) + e1a * x11(t),
    x11'(t) = -a1 * x11(t) * x7(t) + i1a * kv * x10(t) - e1a * kv * x11(t),
    x12'(t) = c2a + c1a * x7(t) - c3a * x12(t),
    x13'(t) = a1 * x10(t) * x6(t) - c6a * x13(t) - a3 * x2(t) * x13(t) + e2a * x14(t),
    x14'(t) = a1 * x11(t) * x7(t) - e2a * kv * x14(t),
    x15'(t) = c2c + c1c * x7(t) - c3c * x15(t),
    y1(t) = x7(t),
    y2(t) = x10(t) + x13(t),
    y3(t) = x9(t),
    y4(t) = x1(t) + x2(t) + x3(t),
    y5(t) = x2(t),
    y6(t) = x12(t)
)

ode = set_parameter_values(ode, Dict(
    a1 => QQ(1, 2),
    a2 => QQ(1, 5),
    a3 => QQ(1),
    c1a => QQ(5, 10^(7)),
    c2a => QQ(0),
    c5a => QQ(1, 10^(4)),
    c6a => QQ(2, 10^(5)),
    c1 => QQ(5, 10^(7)),
    c2 => QQ(0),
    c3 => QQ(4, 10^(4)),
    c4 => QQ(1, 2),
    kv => QQ(5),
    e1a => QQ(5, 10^(4)),
    c1c => QQ(5, 10^(7)),
    c2c => QQ(0),
    c3c => QQ(4, 10^(4))
))

@time io_equations = find_ioequations(ode)


text:
# Julia wrapper for header: types.h
# Automatically generated using Clang.jl

# Julia wrapper for header: cutensor.h
# Automatically generated using Clang.jl

function cutensorCreate(handle)
    @check ccall((:cutensorCreate, @libcutensor), cutensorStatus_t,
                 (Ptr{cutensorHandle_t},),
                 handle)
end

function cutensorDestroy(handle)
    @check ccall((:cutensorDestroy, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t,),
                 handle)
end

function cutensorCreateTensorDescriptor(desc, numModes, extent, stride, dataType, unaryOp,
                                        vectorWidth, vectorModeIndex)
    @check ccall((:cutensorCreateTensorDescriptor, @libcutensor), cutensorStatus_t,
                 (Ptr{cutensorTensorDescriptor_t}, UInt32, Ptr{Int64}, Ptr{Int64},
                  cudaDataType_t, cutensorOperator_t, UInt32, UInt32),
                 desc, numModes, extent, stride, dataType, unaryOp, vectorWidth,
                 vectorModeIndex)
end

function cutensorDestroyTensorDescriptor(desc)
    @check ccall((:cutensorDestroyTensorDescriptor, @libcutensor), cutensorStatus_t,
                 (cutensorTensorDescriptor_t,),
                 desc)
end

function cutensorElementwiseTrinary(handle, alpha, A, descA, modeA, beta, B, descB, modeB,
                                    gamma, C, descC, modeC, D, descD, modeD, opAB, opABC,
                                    typeCompute, stream)
    @check ccall((:cutensorElementwiseTrinary, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  cutensorOperator_t, cutensorOperator_t, cudaDataType_t, CUstream),
                 handle, alpha, A, descA, modeA, beta, B, descB, modeB, gamma, C, descC,
                 modeC, D, descD, modeD, opAB, opABC, typeCompute, stream)
end

function cutensorElementwiseBinary(handle, alpha, A, descA, modeA, gamma, C, descC, modeC,
                                   D, descD, modeD, opAC, typeCompute, stream)
    @check ccall((:cutensorElementwiseBinary, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  cutensorOperator_t, cudaDataType_t, CUstream),
                 handle, alpha, A, descA, modeA, gamma, C, descC, modeC, D, descD, modeD,
                 opAC, typeCompute, stream)
end

function cutensorPermutation(handle, alpha, A, descA, modeA, B, descB, modeB, typeCompute,
                             stream)
    @check ccall((:cutensorPermutation, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  cudaDataType_t, CUstream),
                 handle, alpha, A, descA, modeA, B, descB, modeB, typeCompute, stream)
end

function cutensorContraction(handle, alpha, A, descA, modeA, B, descB, modeB, beta, C,
                             descC, modeC, D, descD, modeD, opOut, typeCompute, algo,
                             workspace, workspaceSize, stream)
    @check ccall((:cutensorContraction, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32}, cutensorOperator_t,
                  cudaDataType_t, cutensorAlgo_t, CuPtr{Cvoid}, UInt64, CUstream),
                 handle, alpha, A, descA, modeA, B, descB, modeB, beta, C, descC, modeC, D,
                 descD, modeD, opOut, typeCompute, algo, workspace, workspaceSize, stream)
end

function cutensorContractionGetWorkspace(handle, A, descA, modeA, B, descB, modeB, C,
                                         descC, modeC, D, descD, modeD, opOut, typeCompute,
                                         algo, pref, workspaceSize)
    @check ccall((:cutensorContractionGetWorkspace, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32}, CuPtr{Cvoid},
                  cutensorTensorDescriptor_t, Ptr{Int32}, CuPtr{Cvoid},
                  cutensorTensorDescriptor_t, Ptr{Int32}, cutensorOperator_t,
                  cudaDataType_t, cutensorAlgo_t, cutensorWorksizePreference_t, Ptr{UInt64}),
                 handle, A, descA, modeA, B, descB, modeB, C, descC, modeC, D, descD,
                 modeD, opOut, typeCompute, algo, pref, workspaceSize)
end

function cutensorContractionMaxAlgos(maxNumAlgos)
    @check ccall((:cutensorContractionMaxAlgos, @libcutensor), cutensorStatus_t,
                 (Ptr{Int32},),
                 maxNumAlgos)
end

function cutensorReduction(handle, alpha, A, descA, modeA, beta, C, descC, modeC, D, descD,
                           modeD, opReduce, typeCompute, workspace, workspaceSize, stream)
    @check ccall((:cutensorReduction, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, Ptr{Cvoid}, CuPtr{Cvoid}, cutensorTensorDescriptor_t,
                  Ptr{Int32}, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  cutensorOperator_t, cudaDataType_t, CuPtr{Cvoid}, UInt64, CUstream),
                 handle, alpha, A, descA, modeA, beta, C, descC, modeC, D, descD, modeD,
                 opReduce, typeCompute, workspace, workspaceSize, stream)
end

function cutensorReductionGetWorkspace(handle, A, descA_, modeA, C, descC_, modeC, D,
                                       descD_, modeD, opReduce, typeCompute, workspaceSize)
    @check ccall((:cutensorReductionGetWorkspace, @libcutensor), cutensorStatus_t,
                 (cutensorHandle_t, CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32},
                  CuPtr{Cvoid}, cutensorTensorDescriptor_t, Ptr{Int32}, CuPtr{Cvoid},
                  cutensorTensorDescriptor_t, Ptr{Int32}, cutensorOperator_t,
                  cudaDataType_t, Ptr{UInt64}),
                 handle, A, descA_, modeA, C, descC_, modeC, D, descD_, modeD, opReduce,
                 typeCompute, workspaceSize)
end

function cutensorGetErrorString(error)
    ccall((:cutensorGetErrorString, @libcutensor), Cstring,
          (cutensorStatus_t,),
          error)
end

function cutensorGetVersion()
    ccall((:cutensorGetVersion, @libcutensor), Csize_t, ())
end

function cutensorGetCudartVersion()
    ccall((:cutensorGetCudartVersion, @libcutensor), Csize_t, ())
end


text:
"""
    NUTS(n_iters::Int, n_adapt::Int, delta::Float64)

No-U-Turn Sampler (NUTS) sampler.

Usage:

```julia
NUTS(1000, 200, 0.6j_max)
```

Example:

```julia
# Define a simple Normal model with unknown mean and variance.
@model gdemo(x) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0, sqrt(s))
  x[1] ~ Normal(m, sqrt(s))
  x[2] ~ Normal(m, sqrt(s))
  return s, m
end

sample(gdemo([1.j_max, 2]), NUTS(1000, 200, 0.6j_max))
```
"""
mutable struct NUTS{T} <: Hamiltonian
  n_iters   ::  Int       # number of samples
  n_adapt   ::  Int       # number of samples with adaption for epsilon
  delta     ::  Float64   # target accept rate
  space     ::  Set{T}    # sampling space, emtpy means all
  gid       ::  Int       # group ID
end
function NUTS(n_adapt::Int, delta::Float64, space...)
  NUTS(1, n_adapt, delta, isa(space, Symbol) ? Set([space]) : Set(space), 0)
end
function NUTS(n_iters::Int, n_adapt::Int, delta::Float64, space...)
  NUTS(n_iters, n_adapt, delta, isa(space, Symbol) ? Set([space]) : Set(space), 0)
end
function NUTS(n_iters::Int, delta::Float64)
  n_adapt_default = Int(round(n_iters / 2))
  NUTS(n_iters, n_adapt_default > 1000 ? 1000 : n_adapt_default, delta, Set(), 0)
end
NUTS(alg::NUTS, new_gid::Int) = NUTS(alg.n_iters, alg.n_adapt, alg.delta, alg.space, new_gid)

function step(model::Function, spl::Sampler{<:NUTS}, vi::VarInfo, is_first::Bool)
  if is_first
    if ~haskey(spl.info, :wum)
      if spl.alg.gid != 0 link!(vi, spl) end      # X -> R

      init_warm_up_params(vi, spl)

      θ = vi[spl]
      ϵ = find_good_eps(model, vi, spl)           # heuristically find optimal ϵ
      vi[spl] = θ

      if spl.alg.gid != 0 invlink!(vi, spl) end   # R -> X

      push!(spl.info[:wum][:ϵ], ϵ)
      update_da_μ(spl.info[:wum], ϵ)
    end

    push!(spl.info[:accept_his], true)

    vi
  else
    # Set parameters
    ϵ = spl.info[:wum][:ϵ][end]; @debug "current ϵ: $ϵ"

    spl.info[:lf_num] = 0   # reset current lf num counter

    @debug "X-> R..."
    if spl.alg.gid != 0
      link!(vi, spl)
      runmodel!(model, vi, spl)
    end

    grad_func = gen_grad_func(vi, spl, model)
    lj_func = gen_lj_func(vi, spl, model)
    rev_func = gen_rev_func(vi, spl)
    log_func = gen_log_func(spl)

    θ = vi[spl]
    lj = vi.logp
    stds = spl.info[:wum][:stds]

    θ_new, da_stat = _nuts_step(θ, ϵ, lj_func, grad_func, stds)

    if PROGRESS[] && spl.alg.gid == 0
      stds_str = string(spl.info[:wum][:stds])
      stds_str = length(stds_str) >= 32 ? stds_str[1:30]*"..." : stds_str
      haskey(spl.info, :progress) && ProgressMeter.update!(
                                       spl.info[:progress],
                                       spl.info[:progress].counter; showvalues = [(:ϵ, ϵ), (:pre_cond, stds_str)])
    end

    push!(spl.info[:accept_his], true)
    vi[spl][:] = θ_new[:]
    setlogp!(vi, lj_func(θ_new))

    # Adapt step-size and pre-cond
    # TODO: figure out whether or not the condition below is needed
    # if τ_valid > 0
      adapt!(spl.info[:wum], da_stat, vi[spl], adapt_M = true, adapt_ϵ = true)
    # end

    @debug "R -> X..."
    spl.alg.gid != 0 && invlink!(vi, spl)

    vi
  end
end

"""
  function _build_tree(θ::T, r::AbstractVector, logu::AbstractFloat, v::Int, j::Int, ϵ::AbstractFloat,
                       H0::AbstractFloat,lj_func::Function, grad_func::Function, stds::AbstractVector;
                       Δ_max::AbstractFloat=1000) where {T<:Union{Vector,SubArray}}

Recursively build balanced tree.

Ref: Algorithm 6 on http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf

Arguments:

- `θ`         : model parameter
- `r`         : momentum variable
- `logu`      : slice variable (in log scale)
- `v`         : direction ∈ {-1, 1}
- `j`         : depth of tree
- `ϵ`         : leapfrog step size
- `H0`        : initial H
- `lj_func`   : function for log-joint
- `grad_func` : function for the gradient of log-joint
- `stds`      : pre-conditioning matrix
- `Δ_max`     : threshold for exploeration error tolerance
"""
function _build_tree(θ::T, r::AbstractVector, logu::AbstractFloat, v::Int, j::Int, ϵ::AbstractFloat,
                     H0::AbstractFloat,lj_func::Function, grad_func::Function, stds::AbstractVector;
                     Δ_max::AbstractFloat=1000.0) where {T<:Union{AbstractVector,SubArray}}
    if j == 0
      # Base case - take one leapfrog step in the direction v.
      θ′, r′, τ_valid = _leapfrog(θ, r, 1, v * ϵ, grad_func)
      # Use old H to save computation
      H′ = τ_valid == 0 ? Inf : _find_H(θ′, r′, lj_func, stds)
      n′ = (logu <= -H′) ? 1 : 0
      s′ = (logu < Δ_max + -H′) ? 1 : 0
      α′ = exp(min(0, -H′ - (-H0)))

      return θ′, r′, θ′, r′, θ′, n′, s′, α′, 1
    else
      # Recursion - build the left and right subtrees.
      θm, rm, θp, rp, θ′, n′, s′, α′, n′α = _build_tree(θ, r, logu, v, j - 1, ϵ, H0, lj_func, grad_func, stds)

      if s′ == 1
        if v == -1
          θm, rm, _, _, θ′′, n′′, s′′, α′′, n′′α = _build_tree(θm, rm, logu, v, j - 1, ϵ, H0, lj_func, grad_func, stds)
        else
          _, _, θp, rp, θ′′, n′′, s′′, α′′, n′′α = _build_tree(θp, rp, logu, v, j - 1, ϵ, H0, lj_func, grad_func, stds)
        end
        if rand() < n′′ / (n′ + n′′)
          θ′ = θ′′
        end
        α′ = α′ + α′′
        n′α = n′α + n′′α
        s′ = s′′ * (dot(θp - θm, rm) >= 0 ? 1 : 0) * (dot(θp - θm, rp) >= 0 ? 1 : 0)
        n′ = n′ + n′′
      end

      θm, rm, θp, rp, θ′, n′, s′, α′, n′α
    end
  end

"""
  function _nuts_step(θ::T, ϵ::AbstractFloat, lj_func::Function, grad_func::Function, stds::AbstractVector;
                      j_max::Int=j_max) where {T<:Union{AbstractVector,SubArray}}

Perform one NUTS step.

Ref: Algorithm 6 on http://www.stat.columbia.edu/~gelman/research/published/nuts.pdf

Arguments:

- `θ`         : model parameter
- `ϵ`         : leapfrog step size
- `lj_func`   : function for log-joint
- `grad_func` : function for the gradient of log-joint
- `stds`      : pre-conditioning matrix
- `j_max`     : maximum expanding of doubling tree
"""
function _nuts_step(θ::T, ϵ::AbstractFloat, lj_func::Function, grad_func::Function, stds::AbstractVector;
                    j_max::Int=5) where {T<:Union{AbstractVector,SubArray}}

  d = length(θ)
  r0 = randn(d)
  H0 = _find_H(θ, r0, lj_func, stds)
  logu = log(rand()) + -H0

  θm = θ; θp = θ; rm = r0; rp = r0; j = 0; θ_new = θ; n = 1; s = 1
  local da_stat

  while s == 1 && j <= j_max

    v = rand([-1, 1])
    if v == -1
        θm, rm, _, _, θ′, n′, s′, α, nα = _build_tree(θm, rm, logu, v, j, ϵ, H0, lj_func, grad_func, stds)
    else
        _, _, θp, rp, θ′, n′, s′, α, nα = _build_tree(θp, rp, logu, v, j, ϵ, H0, lj_func, grad_func, stds)
    end

    if s′ == 1
        if rand() < min(1, n′ / n)
            θ_new = θ′
        end
    end

    n = n + n′
    s = s′ * (dot(θp - θm, rm) >= 0 ? 1 : 0) * (dot(θp - θm, rp) >= 0 ? 1 : 0)
    j = j + 1

    da_stat = α / nα

  end

  return θ_new, da_stat

end


text:
#' # Simple pipeline

#' Loading a single conditon from subject where epoch data was extracted via BESA®. In This
#' example we use some more plausible data than the test data

using Plots
gr()
using Measures
using MegTools
test = load_cont_epochs("condEpoch.mat")

#' The data from condition "2" (defined in BESA®) is averaged across all the trials
test_average = average_across_trials(test["2"])

#' We select data that we're interested in: in this case the pre-defined auditoryN1m channels
#' and take a look at the averaged waveforms from the selected channels.
test_average_auditory,auditory_left,auditory_right = select_channels(
    test_average,
    paradigm="auditoryN1m",
)

plot(test_average_auditory, legend=false, size=(1000,500))

#' Applying filters to smoothen out the waveform
filtered = highlow_butterworth_filter(test_average_auditory, 1000)
plot(test["2"].time, filtered(channels=:MEG1621), label="Filtered", size=(1000,500), linewidth=3, color=:royalblue);
plot!(test["2"].time, test_average_auditory(channels=:MEG1621), label="Raw", color=:lightblue)

#' Follow up with baseline correction of the data
baseline_corrected = baseline_correction(filtered)
plot(test["2"].time, filtered(channels=:MEG1621), label="Only Filtered", size=(1000,500), linewidth=3, color=:royalblue);
plot!(test["2"].time, baseline_corrected(channels=:MEG1621), label="Baseline Corrected", size=(1000,500), linewidth=3, color=:crimson)

#' And finally finding peaks in both left and right sets of channels
a,_,b,_,c,d = find_peaks(baseline_corrected, auditory_left, auditory_right)

# Takinga look at all the channels with marked peak values
channel_plots = plot(layout=(9,3), size = (1000,3000), margin=5mm, legend=false, ticks=[])
for channel = 1:26
    plot!(channel_plots, test_average_auditory[:,channel], color=:lightblue, subplot=channel, label = "Raw Signal")
    plot!(channel_plots, filtered[:,channel], subplot=channel, width=2, color=:royalblue,label = "filtered Signal")
    plot!(channel_plots, baseline_corrected[:,channel], subplot=channel, color=:crimson, width=2, label="Baseline corrected")
    if test_average_auditory.channels[channel] == c
        plot!(reverse(findmax(baseline_corrected[:,channel])), seriestype=:scatter, subplot=channel, color=:green, label="Peak left hem", markersize=8)
    elseif test_average_auditory.channels[channel] == d
        plot!(reverse(findmin(baseline_corrected[:,channel])), seriestype=:scatter, subplot=channel, color=:red, label="Peak right hem", markersize=8)
    end

end

channel_plots

max_stars_repo_name: ElectronicTeaCup/MegTools

max_issues_repo_name: ElectronicTeaCup/MegTools

max_forks_repo_name: ElectronicTeaCup/MegTools


text:
# Nest
#
# If the line exceeds the print width it will be nested.
#
# This is done by replacing `PLACEHOLDER` nodes with `NEWLINE`
# nodes and updating the FST's indent.
#
# `extra_margin` provides additional width to consider from
# the top-level node.
#
# Example:
#
#     LHS op RHS
#
# the length of " op" will be considered when nesting LHS

# unnest, converts newlines to whitespace
unnest!(fst::FST, ind::Int) = fst[ind] = Whitespace(fst[ind].len)
function unnest!(fst::FST, nl_inds::Vector{Int})
    for (i, ind) in enumerate(nl_inds)
        unnest!(fst, ind)
        i == length(nl_inds) || continue
        ind2 = ind - 1
        if fst[ind2].typ === TRAILINGCOMMA
            fst[ind2].val = ""
            fst[ind2].len = 0
        elseif fst[ind-1].typ === TRAILINGSEMICOLON
            fst[ind2].val = ";"
            fst[ind2].len = 1
        end
    end
end

function skip_indent(fst::FST)
    if fst.typ === CSTParser.LITERAL && fst.val == ""
        return true
    elseif fst.typ === NEWLINE || fst.typ === NOTCODE
        return true
    end
    false
end

function walk(f, fst::FST, s::State)
    f(fst, s)
    is_leaf(fst) && return
    for (i, n) in enumerate(fst.nodes)
        if n.typ === NEWLINE && i < length(fst.nodes)
            if is_closer(fst[i+1])
                s.line_offset = fst[i+1].indent
            elseif !skip_indent(fst[i+1])
                s.line_offset = fst.indent
            end
        else
            walk(f, n, s)
        end
    end
end

function reset_line_offset!(fst::FST, s::State)
    is_leaf(fst) || return
    s.line_offset += length(fst)
end

function add_indent!(fst::FST, s::State, indent)
    indent == 0 && return
    lo = s.line_offset
    f = (fst::FST, s::State) -> fst.indent += indent
    walk(f, fst, s)
    s.line_offset = lo
end

function dedent!(fst::FST, s::State)
    if is_leaf(fst)
        s.line_offset += length(fst)
        if is_closer(fst) || fst.typ === NOTCODE
            fst.indent -= s.indent_size
        end
        return
    end
    fst.typ === CSTParser.ConditionalOpCall && return
    fst.typ === CSTParser.StringH && return
    fst.typ === CSTParser.BinaryOpCall && return

    # dedent
    fst.indent -= s.indent_size
    fst.force_nest && return

    nl_inds = findall(n -> n.typ === NEWLINE && !n.force_nest, fst.nodes)
    length(nl_inds) > 0 || return
    margin = s.line_offset + fst.extra_margin + length(fst)
    # @info "" fst.typ margin s.line_offset fst.extra_margin length(fst) length(nl_inds)
    margin <= s.margin || return
    unnest!(fst, nl_inds)
end

function nest!(
    ds::DefaultStyle,
    nodes::Vector{FST},
    s::State,
    indent::Int;
    extra_margin = 0,
)
    style = getstyle(ds)
    for (i, n) in enumerate(nodes)
        if n.typ === NEWLINE && nodes[i+1].typ === CSTParser.Block
            s.line_offset = nodes[i+1].indent
        elseif n.typ === NOTCODE && nodes[i+1].typ === CSTParser.Block
            s.line_offset = nodes[i+1].indent
        elseif n.typ === NEWLINE
            s.line_offset = indent
        else
            n.extra_margin = extra_margin
            nest!(style, n, s)
        end
    end
end
nest!(
    style::S,
    nodes::Vector{FST},
    s::State,
    indent::Int;
    extra_margin = 0,
) where {S<:AbstractStyle} =
    nest!(DefaultStyle(style), nodes, s, indent, extra_margin = extra_margin)

function nest!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    if is_leaf(fst)
        s.line_offset += length(fst)
        return
    end

    if fst.typ === CSTParser.Import
        n_import!(style, fst, s)
    elseif fst.typ === CSTParser.Export
        n_export!(style, fst, s)
    elseif fst.typ === CSTParser.Using
        n_using!(style, fst, s)
    elseif fst.typ === CSTParser.WhereOpCall
        n_whereopcall!(style, fst, s)
    elseif fst.typ === CSTParser.ConditionalOpCall
        n_conditionalopcall!(style, fst, s)
    elseif fst.typ === CSTParser.BinaryOpCall
        n_binaryopcall!(style, fst, s)
    elseif fst.typ === CSTParser.Curly
        n_curly!(style, fst, s)
    elseif fst.typ === CSTParser.Call
        n_call!(style, fst, s)
    elseif fst.typ === CSTParser.MacroCall
        n_macrocall!(style, fst, s)
    elseif fst.typ === CSTParser.Ref
        n_ref!(style, fst, s)
    elseif fst.typ === CSTParser.TypedVcat
        n_typedvcat!(style, fst, s)
    elseif fst.typ === CSTParser.TupleH
        n_tupleh!(style, fst, s)
    elseif fst.typ === CSTParser.Vect
        n_vect!(style, fst, s)
    elseif fst.typ === CSTParser.Vcat
        n_vcat!(style, fst, s)
    elseif fst.typ === CSTParser.Parameters
        n_parameters!(style, fst, s)
    elseif fst.typ === CSTParser.Braces
        n_braces!(style, fst, s)
    elseif fst.typ === CSTParser.InvisBrackets
        n_invisbrackets!(style, fst, s)
    elseif fst.typ === CSTParser.Comprehension
        n_comprehension!(style, fst, s)
    elseif fst.typ === CSTParser.Do
        n_do!(style, fst, s)
    elseif fst.typ === CSTParser.Generator
        n_generator!(style, fst, s)
    elseif fst.typ === CSTParser.Filter
        n_filter!(style, fst, s)
    elseif fst.typ === CSTParser.Block
        n_block!(style, fst, s)
    elseif fst.typ === CSTParser.ChainOpCall
        n_chainopcall!(style, fst, s)
    elseif fst.typ === CSTParser.Comparison
        n_comparison!(style, fst, s)
    elseif fst.typ === CSTParser.For
        n_for!(style, fst, s)
    elseif fst.typ === CSTParser.Let
        n_let!(style, fst, s)
    elseif fst.typ === CSTParser.UnaryOpCall && fst[2].typ === CSTParser.OPERATOR
        n_unaryopcall!(style, fst, s)
    else
        nest!(style, fst.nodes, s, fst.indent, extra_margin = fst.extra_margin)
    end
end
nest!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    nest!(DefaultStyle(style), fst, s)

function n_unaryopcall!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    fst[1].extra_margin = fst.extra_margin + length(fst[2])
    nest!(style, fst[1], s)
    nest!(style, fst[2], s)
end
n_unaryopcall!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_unaryopcall!(DefaultStyle(style), fst, s)

function n_do!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    extra_margin = sum(length.(fst[2:3]))
    # make sure there are nodes after "do"
    if fst[4].typ === WHITESPACE
        extra_margin += length(fst[4])
        extra_margin += length(fst[5])
    end
    fst[1].extra_margin = fst.extra_margin + extra_margin
    nest!(style, fst[1], s)
    nest!(style, fst[2:end], s, fst.indent, extra_margin = fst.extra_margin)
end
n_do!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_do!(DefaultStyle(style), fst, s)

# Import,Using,Export
function n_import!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    line_margin = s.line_offset + length(fst) + fst.extra_margin
    idx = findfirst(n -> n.typ === PLACEHOLDER, fst.nodes)
    if idx !== nothing && (line_margin > s.margin || fst.force_nest)
        fst.indent += s.indent_size

        if !fst.force_nest
            if fst.indent + sum(length.(fst[idx+1:end])) <= s.margin
                fst[idx] = Newline(length = fst[idx].len)
                walk(reset_line_offset!, fst, s)
                return
            end
        end

        for (i, n) in enumerate(fst.nodes)
            if n.typ === NEWLINE
                s.line_offset = fst.indent
            elseif n.typ === PLACEHOLDER
                fst[i] = Newline(length = n.len)
                s.line_offset = fst.indent
            else
                nest!(style, n, s)
            end
        end
    else
        nest!(style, fst.nodes, s, fst.indent, extra_margin = fst.extra_margin)
    end
end
n_import!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_import!(DefaultStyle(style), fst, s)

@inline n_export!(ds::DefaultStyle, fst::FST, s::State) = n_import!(ds, fst, s)
n_export!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_export!(DefaultStyle(style), fst, s)

@inline n_using!(ds::DefaultStyle, fst::FST, s::State) = n_import!(ds, fst, s)
n_using!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_using!(DefaultStyle(style), fst, s)

function n_tupleh!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    line_margin = s.line_offset + length(fst) + fst.extra_margin
    idx = findlast(n -> n.typ === PLACEHOLDER, fst.nodes)
    # @info "ENTERING" idx fst.typ s.line_offset length(fst) fst.extra_margin
    opener = length(fst.nodes) > 0 ? is_opener(fst[1]) : false
    if idx !== nothing && (line_margin > s.margin || fst.force_nest)
        line_offset = s.line_offset
        if opener
            fst[end].indent = fst.indent
        end
        if fst.typ !== CSTParser.Parameters && !(fst.typ === CSTParser.TupleH && !opener)
            fst.indent += s.indent_size
        end

        # @info "DURING" fst.typ fst.indent s.line_offset
        for (i, n) in enumerate(fst.nodes)
            if n.typ === NEWLINE
                s.line_offset = fst.indent
            elseif n.typ === PLACEHOLDER
                fst[i] = Newline(length = n.len)
                s.line_offset = fst.indent
            elseif n.typ === TRAILINGCOMMA
                n.val = ","
                n.len = 1
                nest!(style, n, s)
            elseif n.typ === TRAILINGSEMICOLON
                n.val = ""
                n.len = 0
                nest!(style, n, s)
            elseif opener && (i == 1 || i == length(fst.nodes))
                nest!(style, n, s)
            else
                diff = fst.indent - fst[i].indent
                add_indent!(n, s, diff)
                n.extra_margin = 1
                nest!(style, n, s)
            end
        end

        if opener
            s.line_offset = fst[end].indent + 1
        end
        # @info "EXITING" fst.typ s.line_offset fst.indent fst[end].indent
    else
        extra_margin = fst.extra_margin
        opener && (extra_margin += 1)
        nest!(style, fst.nodes, s, fst.indent, extra_margin = extra_margin)
    end
end
n_tupleh!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_tupleh!(DefaultStyle(style), fst, s)

@inline n_vect!(ds::DefaultStyle, fst::FST, s::State) = n_tupleh!(ds, fst, s)
n_vect!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_vect!(DefaultStyle(style), fst, s)

@inline n_vcat!(ds::DefaultStyle, fst::FST, s::State) = n_tupleh!(ds, fst, s)
n_vcat!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_vcat!(DefaultStyle(style), fst, s)

@inline n_braces!(ds::DefaultStyle, fst::FST, s::State) = n_tupleh!(ds, fst, s)
n_braces!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_braces!(DefaultStyle(style), fst, s)

@inline n_parameters!(ds::DefaultStyle, fst::FST, s::State) = n_tupleh!(ds, fst, s)
n_parameters!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_parameters!(DefaultStyle(style), fst, s)

@inline n_invisbrackets!(ds::DefaultStyle, fst::FST, s::State) = n_tupleh!(ds, fst, s)
n_invisbrackets!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_invisbrackets!(DefaultStyle(style), fst, s)

@inline n_comprehension!(ds::DefaultStyle, fst::FST, s::State) = n_tupleh!(ds, fst, s)
n_comprehension!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_comprehension!(DefaultStyle(style), fst, s)

function n_call!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    line_margin = s.line_offset + length(fst) + fst.extra_margin
    idx = findlast(n -> n.typ === PLACEHOLDER, fst.nodes)
    # @info "ENTERING" fst.typ s.line_offset length(fst) fst.extra_margin s.margin fst[1].val
    if idx !== nothing && (line_margin > s.margin || fst.force_nest)
        line_offset = s.line_offset
        fst[end].indent = fst.indent
        fst.indent += s.indent_size

        # @info "DURING" fst.typ fst.indent s.line_offset
        for (i, n) in enumerate(fst.nodes)
            if n.typ === NEWLINE
                s.line_offset = fst.indent
            elseif n.typ === PLACEHOLDER
                fst[i] = Newline(length = n.len)
                s.line_offset = fst.indent
            elseif n.typ === TRAILINGCOMMA
                n.val = ","
                n.len = 1
                nest!(style, n, s)
            elseif n.typ === TRAILINGSEMICOLON
                n.val = ""
                n.len = 0
                nest!(style, n, s)
            elseif i == 1 || i == length(fst.nodes)
                n.typ === CSTParser.Parameters && (n.force_nest = true)
                n.extra_margin = 1
                nest!(style, n, s)
            else
                diff = fst.indent - fst[i].indent
                add_indent!(n, s, diff)
                n.typ === CSTParser.Parameters && (n.force_nest = true)
                n.extra_margin = 1
                nest!(style, n, s)
            end
        end

        s.line_offset = fst[end].indent + 1
    else
        extra_margin = fst.extra_margin
        is_closer(fst[end]) && (extra_margin += 1)
        nest!(style, fst.nodes, s, fst.indent, extra_margin = extra_margin)
    end
end
n_call!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_call!(DefaultStyle(style), fst, s)

@inline n_curly!(ds::DefaultStyle, fst::FST, s::State) = n_call!(ds, fst, s)
n_curly!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_curly!(DefaultStyle(style), fst, s)

@inline n_macrocall!(ds::DefaultStyle, fst::FST, s::State) = n_call!(ds, fst, s)
n_macrocall!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_macrocall!(DefaultStyle(style), fst, s)

@inline n_ref!(ds::DefaultStyle, fst::FST, s::State) = n_call!(ds, fst, s)
n_ref!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_ref!(DefaultStyle(style), fst, s)

@inline n_typedvcat!(ds::DefaultStyle, fst::FST, s::State) = n_call!(ds, fst, s)
n_typedvcat!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_typedvcat!(DefaultStyle(style), fst, s)

function n_whereopcall!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    line_margin = s.line_offset + length(fst) + fst.extra_margin
    # @info "" s.line_offset fst.typ line_margin fst.extra_margin length(fst) fst.force_nest line_margin > s.margin
    if line_margin > s.margin || fst.force_nest
        line_offset = s.line_offset
        # after "A where "
        idx = findfirst(n -> n.typ === PLACEHOLDER, fst.nodes)
        Blen = sum(length.(fst[idx+1:end]))

        # A, +7 is the length of " where "
        fst[1].extra_margin = Blen + 7 + fst.extra_margin
        nest!(style, fst[1], s)

        for (i, n) in enumerate(fst[2:idx-1])
            nest!(style, n, s)
        end

        # "B"

        has_braces = is_closer(fst[end])
        if has_braces
            fst[end].indent = fst.indent
        end

        over = (s.line_offset + Blen + fst.extra_margin > s.margin) || fst.force_nest
        fst.indent += s.indent_size

        for (i, n) in enumerate(fst[idx+1:end])
            if n.typ === NEWLINE
                s.line_offset = fst.indent
            elseif is_opener(n)
                if fst.indent - s.line_offset > 1
                    fst.indent = s.line_offset + 1
                    fst[end].indent = s.line_offset
                end
                nest!(style, n, s)
            elseif n.typ === PLACEHOLDER && over
                fst[i+idx] = Newline(length = n.len)
                s.line_offset = fst.indent
            elseif n.typ === TRAILINGCOMMA && over
                n.val = ","
                n.len = 1
                nest!(style, n, s)
            elseif has_braces
                n.extra_margin = 1 + fst.extra_margin
                nest!(style, n, s)
            else
                n.extra_margin = fst.extra_margin
                nest!(style, n, s)
            end
        end

        s.line_offset = line_offset
        walk(reset_line_offset!, fst, s)
    else
        nest!(style, fst.nodes, s, fst.indent, extra_margin = fst.extra_margin)
    end
end
n_whereopcall!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_whereopcall!(DefaultStyle(style), fst, s)

function n_conditionalopcall!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    line_margin = s.line_offset + length(fst) + fst.extra_margin
    # @info "ENTERING" fst.typ s.line_offset line_margin fst.force_nest
    if line_margin > s.margin || fst.force_nest
        line_offset = s.line_offset
        fst.indent = s.line_offset
        phs = reverse(findall(n -> n.typ === PLACEHOLDER, fst.nodes))
        for (i, idx) in enumerate(phs)
            if i == 1
                fst[idx] = Newline(length = fst[idx].len)
            else
                nidx = phs[i-1]
                l1 = sum(length.(fst[1:idx-1]))
                l2 = sum(length.(fst[idx:nidx-1]))
                width = line_offset + l1 + l2
                if fst.force_nest || width > s.margin
                    fst[idx] = Newline(length = fst[idx].len)
                end
            end
        end

        for (i, n) in enumerate(fst.nodes)
            if i == length(fst.nodes)
                n.extra_margin = fst.extra_margin
                nest!(style, n, s)
            elseif fst[i+1].typ === WHITESPACE
                n.extra_margin = length(fst[i+1]) + length(fst[i+2])
                nest!(style, n, s)
            elseif n.typ === NEWLINE
                s.line_offset = fst.indent
            else
                nest!(style, n, s)
            end
        end

        s.line_offset = line_offset
        for (i, n) in enumerate(fst.nodes)
            if n.typ === NEWLINE && !is_comment(fst[i+1]) && !is_comment(fst[i-1])
                # +1 for newline to whitespace conversion
                width = s.line_offset + 1
                if i == length(fst.nodes) - 1
                    width += sum(length.(fst[i+1:end])) + fst.extra_margin
                else
                    width += sum(length.(fst[i+1:i+3]))
                end
                # @debug "" s.line_offset l  s.margin
                if width <= s.margin
                    fst[i] = Whitespace(1)
                else
                    s.line_offset = fst.indent
                end
            end
            s.line_offset += length(fst[i])
        end

        s.line_offset = line_offset
        walk(reset_line_offset!, fst, s)
    else
        nest!(style, fst.nodes, s, fst.indent, extra_margin = fst.extra_margin)
    end
end
n_conditionalopcall!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_conditionalopcall!(DefaultStyle(style), fst, s)

no_unnest(fst::FST) = fst.typ === CSTParser.BinaryOpCall && contains_comment(fst)

# lhs op rhs
#
# nest in order of
#
# lhs op
# lhs
function n_binaryopcall!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    # If there's no placeholder the binary call is not nestable
    idxs = findall(n -> n.typ === PLACEHOLDER, fst.nodes)
    line_margin = s.line_offset + length(fst) + fst.extra_margin
    rhs = fst[end]
    rhs.typ === CSTParser.Block && (rhs = rhs[1])
    # @info "ENTERING" fst.typ fst.extra_margin s.line_offset length(fst) idxs fst.ref[][2]
    if length(idxs) == 2 && (line_margin > s.margin || fst.force_nest || rhs.force_nest)
        line_offset = s.line_offset
        i1 = idxs[1]
        i2 = idxs[2]
        fst[i1] = Newline(length = fst[i1].len)
        cst = fst.ref[]

        has_eq = CSTParser.defines_function(cst) || nest_assignment(cst)
        has_arrow = cst[2].kind === Tokens.PAIR_ARROW || cst[2].kind === Tokens.ANON_FUNC
        indent_nest = has_eq || has_arrow

        if indent_nest
            s.line_offset = fst.indent + s.indent_size
            fst[i2] = Whitespace(s.indent_size)
            add_indent!(fst[end], s, s.indent_size)
        else
            fst.indent = s.line_offset
        end

        # rhs
        fst[end].extra_margin = fst.extra_margin
        nest!(style, fst[end], s)

        # "lhs op" rhs
        s.line_offset = line_offset

        # extra margin for " op"
        fst[1].extra_margin = length(fst[2]) + length(fst[3])
        nest!(style, fst[1], s)
        for n in fst[2:i1]
            nest!(style, n, s)
        end

        # Undo nest if possible
        if !fst.force_nest && !no_unnest(rhs)
            cst = rhs.ref[]
            line_margin = s.line_offset

            if (
                rhs.typ === CSTParser.BinaryOpCall &&
                (!(is_lazy_op(cst) && !indent_nest) && cst[2].kind !== Tokens.IN)
            ) || rhs.typ === CSTParser.UnaryOpCall ||
               rhs.typ === CSTParser.ChainOpCall || rhs.typ === CSTParser.Comparison
                line_margin += length(fst[end])
            elseif rhs.typ === CSTParser.Do && is_iterable(rhs[1])
                rw, _ = length_to(fst, [NEWLINE], start = i2 + 1)
                line_margin += rw
            elseif is_block(cst)
                idx = findfirst(n -> n.typ === NEWLINE, rhs.nodes)
                if idx === nothing
                    line_margin += length(fst[end])
                else
                    line_margin += sum(length.(rhs[1:idx-1]))
                end
            else
                rw, _ = length_to(fst, [NEWLINE], start = i2 + 1)
                line_margin += rw
            end

            # @info "" rhs.typ indent_nest s.line_offset line_margin fst.extra_margin length(fst[end])
            if line_margin + fst.extra_margin <= s.margin
                fst[i1] = Whitespace(1)
                if indent_nest
                    fst[i2] = Whitespace(0)
                    walk(dedent!, rhs, s)
                end
            end
        end

        # @info "before reset" line_offset s.line_offset
        s.line_offset = line_offset
        walk(reset_line_offset!, fst, s)
        # @info "after reset" line_offset s.line_offset
    else
        # Handles the case of a function def defined
        # as "foo(a)::R where {A,B} = body".
        #
        # In this case instead of it being parsed as:
        #
        # CSTParser.BinaryOpCall
        #  - CSTParser.WhereOpCall
        #  - OP
        #  - rhs
        #
        # It's parsed as:
        #
        # CSTParser.BinaryOpCall
        #  - CSTParser.BinaryOpCall
        #   - lhs
        #   - OP
        #   - CSTParser.WhereOpCall
        #    - R
        #    - ...
        #  - OP
        #  - rhs
        #
        # The result being extra width is trickier to deal with.
        idx = findfirst(n -> n.typ === CSTParser.WhereOpCall, fst.nodes)
        return_width = 0
        if idx !== nothing && idx > 1
            return_width = length(fst[idx]) + length(fst[2])
        elseif idx === nothing
            return_width, _ = length_to(fst, [NEWLINE], start = 2)
        end

        # @info "" idx map(n -> n.typ, fst.nodes)
        # @info "" s.line_offset length(fst[1]) return_width fst.extra_margin

        for (i, n) in enumerate(fst.nodes)
            if n.typ === NEWLINE
                s.line_offset = fst.indent
            elseif i == 1
                n.extra_margin = return_width + fst.extra_margin
                nest!(style, n, s)
            elseif i == idx
                n.extra_margin = fst.extra_margin
                nest!(style, n, s)
            else
                n.extra_margin = fst.extra_margin
                nest!(style, n, s)
            end
        end
    end
end
n_binaryopcall!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_binaryopcall!(DefaultStyle(style), fst, s)

function n_for!(ds::DefaultStyle, fst::FST, s::State)
    style = getstyle(ds)
    block_idx = findfirst(n -> !is_leaf(n), fst.nodes)
    if block_idx === nothing
        nest!(style, fst.nodes, s, fst.indent, extra_margin = fst.extra_margin)
        return
    end
    ph_idx = findfirst(n -> n.typ === PLACEHOLDER, fst[block_idx].nodes)
    nest!(style, fst.nodes, s, fst.indent, extra_margin = fst.extra_margin)

    # return if the argument block was nested
    ph_idx !== nothing && fst[3][ph_idx].typ === NEWLINE && return

    idx = 5
    n = fst[idx]
    if n.typ === NOTCODE && n.startline == n.endline
        res = get(s.doc.comments, n.startline, (0, ""))
        res == (0, "") && deleteat!(fst.nodes, idx - 1)
    end
end
n_for!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_for!(DefaultStyle(style), fst, s)

@inline n_let!(ds::DefaultStyle, fst::FST, s::State) = n_for!(ds, fst, s)
n_let!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_let!(DefaultStyle(style), fst, s)

function n_block!(ds::DefaultStyle, fst::FST, s::State; custom_indent = 0)
    style = getstyle(ds)
    line_margin = s.line_offset + length(fst) + fst.extra_margin
    idx = findfirst(n -> n.typ === PLACEHOLDER, fst.nodes)
    # @info "ENTERING" idx fst.typ s.line_offset length(fst) fst.extra_margin
    if idx !== nothing && (line_margin > s.margin || fst.force_nest)
        line_offset = s.line_offset
        custom_indent > 0 && (fst.indent = custom_indent)
        # @info "" fst.indent

        # @info "DURING" fst.indent s.line_offset fst.typ
        for (i, n) in enumerate(fst.nodes)
            if n.typ === NEWLINE
                s.line_offset = fst.indent
            elseif n.typ === PLACEHOLDER
                fst[i] = Newline(length = n.len)
                s.line_offset = fst.indent
            elseif n.typ === TRAILINGCOMMA
                n.val = ","
                n.len = 1
                nest!(style, n, s)
            elseif n.typ === TRAILINGSEMICOLON
                n.val = ""
                n.len = 0
                nest!(style, n, s)
            else
                diff = fst.indent - fst[i].indent
                add_indent!(n, s, diff)
                n.extra_margin = 1
                nest!(style, n, s)
            end
        end

        # @info "EXITING" fst.typ s.line_offset fst.indent fst[end].indent
    else
        nest!(style, fst.nodes, s, fst.indent, extra_margin = fst.extra_margin)
    end
end
n_block!(style::S, fst::FST, s::State; custom_indent = 0) where {S<:AbstractStyle} =
    n_block!(DefaultStyle(style), fst, s, custom_indent = custom_indent)

@inline n_generator!(ds::DefaultStyle, fst::FST, s::State) =
    n_block!(ds, fst, s, custom_indent = fst.indent)
n_generator!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_generator!(DefaultStyle(style), fst, s)

@inline n_filter!(ds::DefaultStyle, fst::FST, s::State) =
    n_block!(ds, fst, s, custom_indent = fst.indent)
n_filter!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_filter!(DefaultStyle(style), fst, s)

@inline n_comparison!(ds::DefaultStyle, fst::FST, s::State) = n_block!(ds, fst, s)
n_comparison!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_comparison!(DefaultStyle(style), fst, s)

@inline n_chainopcall!(ds::DefaultStyle, fst::FST, s::State) = n_block!(ds, fst, s)
n_chainopcall!(style::S, fst::FST, s::State) where {S<:AbstractStyle} =
    n_chainopcall!(DefaultStyle(style), fst, s)


text:
# # Gaussian Process with Heteroscedastic likelihoods
#
# ### Loading necessary packages
using AugmentedGaussianProcesses
using Distributions
using LinearAlgebra
using Plots
using Random
default(; lw=3.0, msw=0.0)
# using CairoMakie

# ## Model generated data
# The heteroscedastic noise mean that the variance of the likelihood
# directly depends on the input.
# To model this with Gaussian process, we define a GP `f` for the mean and another GP `g` for the variance
# ``y \sim f + \epsilon``
# where ``\epsilon \sim \mathcal{N}(0, (\lambda \sigma(g))^{-1})``
# We create a toy dataset with X ∈ [-10, 10] and sample `f`, `g` and `y` given this same generative model
rng = MersenneTwister(42)
N = 200
x = (sort(rand(rng, N)) .- 0.5) * 20.0
x_test = range(-10, 10; length=500)
kernel = 5.0 * SqExponentialKernel() ∘ ScaleTransform(1.0) # Kernel function
K = kernelmatrix(kernel, x) + 1e-5I # The kernel matrix
f = rand(rng, MvNormal(K)); # We draw a random sample from the GP prior

# We add a prior mean on `g` so that the variance does not become too big
μ₀ = -3.0
g = rand(rng, MvNormal(μ₀ * ones(N), K))
λ = 3.0 # The maximum possible precision
σ = inv.(sqrt.(λ * AGP.logistic.(g))) # We use the following transform to obtain the std. deviation
y = f + σ .* randn(N); # We finally sample the ouput
# We can visualize the data:
n_sig = 2 # Number of std. dev. around the mean
plot(x, f; ribbon=n_sig * σ, lab="p(y|f,g)") # Mean and std. dev. of y
scatter!(x, y; alpha=0.5, msw=0.0, lab="y") # Observation samples

# ## Model creation and training
# We will now use the augmented model to infer both `f` and `g`
model = VGP(
    x,
    y,
    kernel,
    HeteroscedasticLikelihood(λ),
    AnalyticVI();
    optimiser=true, # We optimise both the mean parameters and kernel hyperparameters
    mean=μ₀,
    verbose=1,
)

# Model training, we train for around 100 iterations to wait for the convergence of the hyperparameters
train!(model, 100);

# ## Predictions
# We can now look at the predictions and compare them with out original model
(f_m, g_m), (f_σ, g_σ) = predict_f(model, x_test; cov=true)
y_m, y_σ = proba_y(model, x_test);
# Let's first look at the differece between the latent `f` and `g`
plot(x, [f, g]; label=["f" "g"])
plot!(
    x_test,
    [f_m, g_m];
    ribbon=[n_sig * f_σ n_sig * g_σ],
    lab=["f_pred" "g_pred"],
    legend=true,
)
# But it's more interesting to compare the predictive probability of `y` directly:
plot(x, f; ribbon=n_sig * σ, lab="p(y|f,g)")
plot!(x_test, y_m; ribbon=n_sig * sqrt.(y_σ), lab="p(y|f,g) pred")
scatter!(x, y; lab="y", alpha=0.2)
# Or to explore the heteroscedasticity itself, we can look at the residuals
scatter(x, (f - y) .^ 2; yaxis=:log, lab="residuals", alpha=0.2)
plot!(x, σ .^ 2; lab="true σ²(x)")
plot!(x_test, y_σ; lab="predicted σ²(x)")


text:
function flux_negative_function(flux_positive_function, scheme)
    func(scheme, vector_in::Array{<:Real, 1}) = flux_positive_function(scheme,
                                                inverted_vector(vector_in))
    return func
end # flux_negative_function


text:
module PolynomialZeros

using Polynomials
import PolynomialRoots

using PolynomialFactors
import Roots: bisection64

using AbstractAlgebra
const AA = AbstractAlgebra

using Printf

export poly_roots
export Over

# export agcd, multroot # qualify these

include("utils.jl")
include("over.jl")
include("special_cases.jl")

include("agcd/agcd.jl")        # AGCD.agcd
include("agcd/multroot.jl")    # MultRoot.multroot
include("amvw/AMVW.jl")        # AMVW.poly_roots
include("RealRoots/real_roots.jl")  # RealRoots.real_roots(p::Poly)

using .AGCD
using .MultRoot

"""

`poly_roots(f, domain)`: Find zeros of the polynomial `f` within the specified domain.

* `f` can be an instance of `Poly` type (from `Polynomials.jl`) or a callable object which can be so converted. Will throw a `DomainError` if the object `f` can not be converted to `Poly{T}`.

* `domain` is one of

    - `over.C` (the default) for solving over complex values
      (`Complex{Float64}`). Use `Over.CC{T}` to specfy a type `T<:
      AbstractFloat` other than `Float64`. The default method is from
      `PolynomialRoots`. Pass the argument `method=:roots` to use the
      `roots` function from `Polynomials.jl`. Pass the argument
      `method=:amvw` to use an algorithm by Aurentz, Mach, Vandebril,
      and Watkins. For a degree n polynomial over C, all n roots
      should be returned (including multiplicities).

    - `over.R` for solving over the real line (`Float64`). Use
      `Over.RR{T}` to specify a `T <: Integer` other than
      `Float64`. The algorithm assumes the polynomial is square free
      (none of its factors are squares over R). This is important for
      floating point coefficients. Pass the argument
      `square_free=false` to have an *approximate* gcd used to create
      a square-free version. Only unique real roots are returned (no
      multiplicities).

    - `over.Q` for solving over the rational numbers
      (`Rational{Int}`). Use `Over.RR{T}` to specify a `T <: Integer`
      other than `Int`. Only unique rational roots are returned (no
      multiplicities).

    - `over.Z` for solving over the integers (`Int`). Use `Over.ZZ{T}`
      to specify a `T` other than `Int`. Only unique integer roots are
      returned (no multiplicities).

    - `over.Zp{p}` for solving over the finite field `ZZ_p`, p a
      prime.  Only unique roots are returned (no multiplicities).

Returns an array of zeros, possibly empty. May throw error if polynomial type is inappropriate for specified domain. For an unspecified domain, the domain may reflect the element type of `f(0)`.

Examples:

```julia
using Polynomials, PolynomialZeros
x = variable(); p = x^5 - x - 1
poly_roots(p, Over.C)  # 5
poly_roots(p, Over.R)  # 1
poly_roots(p, Over.Q)  # empty

p = x^3 - 1
poly_roots(p, Over.C)  # 3
poly_roots(p, Over.R)  # 1
poly_roots(p, Over.Q)  # 1
poly_roots(p, Over.Z)  # 1
poly_roots(p, Over.Zp{7})  # 3
```

"""
function poly_roots(f; method=:PolynomialRoots)
    poly_roots(f, Over.C; method=method) # default
end

function poly_roots(f, ::Type{Over.C}; method=:PolynomialRoots)
    T = promote_type(Float64, e_type(f))
    poly_roots(f, Over.CC{T}, method=method)
end

function poly_roots(f, U::Type{Over.CC{T}}; method=:PolynomialRoots) where {T<:AbstractFloat}

    ps = poly_coeffs(T, f)
    fn = special_case(ps, U)

    if fn == identity
        if method == :PolynomialRoots && T <: Union{Float64, BigFloat}
            PolynomialRoots.roots(ps, polish=true)
        elseif method == :Roots && T <: Float64
            convert(Vector{Complex{T}}, Polynomials.roots(Poly(ps)))
        elseif (method == :AMVW || method == :amvw)
            AMVW.poly_roots(ps)
        else # default to AMVW
            AMVW.poly_roots(ps)
        end
    else
        fn(ps, U)
    end

end

function poly_roots(f, ::Type{Over.R};square_free=false)
    T = promote_type(Float64, e_type(f))
    poly_roots(f, Over.RR{T}, square_free=square_free)
end
function poly_roots(f, U::Type{Over.RR{T}}; square_free::Bool=false) where {T <: Real}

    ps = convert(Vector{T}, poly_coeffs(f))
    fn = special_case(ps, U)

    if fn == identity
        RealRoots.real_roots(as_poly(T, f), square_free=square_free)
    else
        fn(ps, U)
    end

end
# should I do an alias? Best to add if requestd, keeping name space light for now
## realroots(f; kwargs...) = poly_roots(f, Over.R; kwargs...) #real_roots?

_rational_T(::Type{Rational{T}}) where {T}= T
_rational_T(::Rational{T}) where {T} =T
function poly_roots(f, ::Type{Over.Q})
    T = promote_type(Int, e_type(f))

    if T <: Integer
        return poly_roots(f, Over.QQ{T})
    elseif T <: Rational
        S = promote_type(_rational_T(T), Int)
        p = as_poly(T, f)
        return poly_roots(f, Over.QQ{S})
    else
        throw(ArgumentError("Use Over.Q for polynomials with integer or rational coefficients"))
    end
end

# for v0.6 -> v0.7 transition on filter with Dicts

function poly_roots(f, U::Type{Over.QQ{T}}) where {T <: Integer}
    p = as_poly(Rational{T}, f)
    fn = special_case(poly_coeffs(p), U)
    if fn == identity
        ps = poly_coeffs(p)
        l = lcm(denominator.(ps))
        qs = numerator.(l * ps)
        d = PolynomialFactors.poly_factor(qs)
        d = filter(kv -> AA.degree(kv[1]) == 1, d) ## is this broken
        rts = Rational{T}[]
        for k in keys(d)
            p0 = k(0); p1 = k(1) - p0
            push!(rts, -p0//p1)
        end
        rts
    else
        fn(p.a, U)
    end

end

## ----

function poly_roots(f, ::Type{Over.Z})
    T = eltype(as_poly(f)(0))
    poly_roots(f, Over.ZZ{T})
end

function poly_roots(f, U::Type{Over.ZZ{T}}) where {T <: Integer}

    p = as_poly(T, f)

    fn = special_case(poly_coeffs(p), U)
    if fn == identity
        d = PolynomialFactors.poly_factor(poly_coeffs(p))
        d = filter(kv -> AA.degree(kv[1]) == 1, d)
        rts = T[]
        for (k,v) in d
            x1, x0= k(1)-k(0), k(0)
            if iszero(rem(x0, x1))
                push!(rts, div(x0, x1))
            end
        end
        rts
    else
        fn(p.a, U)
    end
end

function poly_roots(f, U::Type{Over.Zp{q}}) where {q}
    # error if q is not prime?

    p = as_poly(e_type(f), f)
    paa = PolynomialFactors.as_poly(p.a)

    fn = special_case(poly_coeffs(p), U)
    if fn == identity
        fs = PolynomialFactors.factormod(paa,q)
        ls = filter(kv -> AA.degree(kv[1]) == 1, fs)
        rts = eltype(p.a)[]
        for (k,v) in ls
            x0 = k(0)
            x1 = k(1) - k(0)
            u = x0 * inv(x1)
            push!(rts, PolynomialFactors.value(u))
        end
        rts
    else
        fn(p.a, U)
    end

end

end # module


text:
struct Filter <: AbstractPipeline
    table::Observable{SimpleTable}
    filters::Observable{Filters}
    value::Observable{SimpleTable}
end

function get_vertices(f::Filter)
    isoriginal = all(f.filters[].options) do option
        return option.value.isoriginal[]
    end
    if isoriginal
        return Vertex[]
    else
        names = collect(Tables.columnnames(f.table[]))
        return [Vertex(:Filter, names, names)]
    end
end

get_vertex_names(::Filter) = [:Filter]

function Filter(table::Observable{SimpleTable})
    filters = Observable(Filters(table[]))
    value = Observable(selected_data(filters[], table[]))
    return Filter(table, filters, value)
end

selected_data(f::Filter) = selected_data(f.filters[], f.table[])

function jsrender(session::Session, f::Filter)

    tryon(session, f.table) do table
        f.filters[] = Filters(table)
        f.value[] = selected_data(f)
    end

    filter_button = Button("Filter", class=buttonclass(true))
    clear_button = Button("Clear", class=buttonclass(false))
    ui = scrollable_component(
        f.filters,
        DOM.div(class="mt-12 pl-4", filter_button, clear_button)
    )

    tryon(session, filter_button.value) do _
        f.value[] = selected_data(f)
    end
    tryon(session, clear_button.value) do _
        f.filters[] = Filters(f.table[])
        f.value[] = selected_data(f)
    end

    return jsrender(session, ui)
end