text:
using OrdinaryDiffEq, RecursiveArrayTools, Test, StaticArrays

f = function (u,p,t)
  - u + sin(-t)
end

prob = ODEProblem(f,1.0,(0.0,-10.0))

condition= function (u,t,integrator) # Event when event_f(u,t,k) == 0
  - t - 2.95
end

affect! = function (integrator)
  integrator.u = integrator.u + 2
end

callback = ContinuousCallback(condition,affect!)

sol = solve(prob,Tsit5(),callback=callback)

f = function (du,u,p,t)
  du[1] = - u[1] + sin(t)
end

prob = ODEProblem(f,[1.0],(0.0,10.0))

condtion= function (u,t,integrator) # Event when event_f(u,t,k) == 0
  t - 2.95
end

affect! = function (integrator)
  integrator.u = integrator.u .+ 2
end

callback = ContinuousCallback(condtion,affect!)

sol = solve(prob,Tsit5(),callback=callback,abstol=1e-8,reltol=1e-6)

f = function (du,u,p,t)
  du[1] = u[2]
  du[2] = -9.81
end

condtion= function (u,t,integrator) # Event when event_f(u,t,k) == 0
  u[1]
end

affect! = nothing
affect_neg! = function (integrator)
  integrator.u[2] = -integrator.u[2]
end

callback = ContinuousCallback(condtion,affect!,affect_neg!,interp_points=100)

u0 = [50.0,0.0]
tspan = (0.0,15.0)
prob = ODEProblem(f,u0,tspan)

sol = solve(prob,Tsit5(),callback=callback,adaptive=false,dt=1/4)

condtion_single = function (u,t,integrator) # Event when event_f(u,t,k) == 0
  u
end

affect! = nothing
affect_neg! = function (integrator)
  integrator.u[2] = -integrator.u[2]
end

callback_single = ContinuousCallback(condtion_single,affect!,affect_neg!,interp_points=100,idxs=1)

u0 = [50.0,0.0]
tspan = (0.0,15.0)
prob = ODEProblem(f,u0,tspan)

sol = solve(prob,Tsit5(),callback=callback_single,adaptive=false,dt=1/4)

#plot(sol,denseplot=true)

sol = solve(prob,Vern6(),callback=callback)
#plot(sol,denseplot=true)
sol = solve(prob,BS3(),callback=callback)

sol33 = solve(prob,Vern7(),callback=callback)

bounced = ODEProblem(f,sol[8],(0.0,1.0))
sol_bounced = solve(bounced,Vern6(),callback=callback,dt=sol.t[9]-sol.t[8])
#plot(sol_bounced,denseplot=true)
sol_bounced(0.04) # Complete density
@test maximum(maximum.(map((i)->sol.k[9][i]-sol_bounced.k[2][i],1:length(sol.k[9])))) == 0

sol2= solve(prob,Vern6(),callback=callback,adaptive=false,dt=1/2^4)
#plot(sol2)

sol2= solve(prob,Vern6())

sol3= solve(prob,Vern6(),saveat=[.5])

## Saving callback

condtion = function (u,t,integrator)
  true
end
affect! = function (integrator) end

save_positions = (true,false)
saving_callback = DiscreteCallback(condtion,affect!,save_positions=save_positions)

sol4 = solve(prob,Tsit5(),callback=saving_callback)

@test sol2(3) ≈ sol(3)

affect! = function (integrator)
  u_modified!(integrator,false)
end
saving_callback2 = DiscreteCallback(condtion,affect!,save_positions=save_positions)
sol4 = solve(prob,Tsit5(),callback=saving_callback2)

cbs = CallbackSet(saving_callback,saving_callback2)
sol4_extra = solve(prob,Tsit5(),callback=cbs)

@test length(sol4_extra) == 2length(sol4) - 1

condtion= function (u,t,integrator)
  u[1]
end

affect! = function (integrator)
  terminate!(integrator)
end

terminate_callback = ContinuousCallback(condtion,affect!)

tspan2 = (0.0,Inf)
prob2 = ODEProblem(f,u0,tspan2)

sol5 = solve(prob2,Tsit5(),callback=terminate_callback)

@test sol5[end][1] < 3e-12
@test sol5.t[end] ≈ sqrt(50*2/9.81)

affect2! = function (integrator)
  if integrator.t >= 3.5
    terminate!(integrator)
  else
    integrator.u[2] = -integrator.u[2]
  end
end
terminate_callback2 = ContinuousCallback(condtion,nothing,affect2!,interp_points=100)

sol5 = solve(prob2,Vern7(),callback=terminate_callback2)

@test sol5[end][1] < 1.3e-10
@test_broken sol5.t[end] ≈ 3*sqrt(50*2/9.81)

condtion= function (u,t,integrator) # Event when event_f(u,t,k) == 0
  t-4
end

affect! = function (integrator)
  terminate!(integrator)
end

terminate_callback3 = ContinuousCallback(condtion,affect!,interp_points=1000)

bounce_then_exit = CallbackSet(callback,terminate_callback3)

sol6 = solve(prob2,Vern7(),callback=bounce_then_exit)

@test sol6[end][1] > 0
@test sol6[end][1] < 100
@test sol6.t[end] ≈ 4

# More ODE event tests, cf. #201, #199, #198, #197
function test_callback_inplace(alg; kwargs...)
    f = (du, u, p, t) -> @. du = u
    cb = ContinuousCallback((u,t,int) -> u[1] - exp(1), terminate!)
    prob = ODEProblem(f, [1.0], (0.0, 2.0), callback=cb)
    sol = solve(prob, alg; kwargs...)
    sol.u[end][1] ≈ exp(1)
end

function test_callback_outofplace(alg; kwargs...)
    f = (u, p, t) -> copy(u)
    cb = ContinuousCallback((u,t,int) -> u[1] - exp(1), terminate!)
    prob = ODEProblem(f, [1.0], (0.0, 2.0), callback=cb)
    sol = solve(prob, alg; kwargs...)
    sol.u[end][1] ≈ exp(1)
end

function test_callback_scalar(alg; kwargs...)
    f = (u, p, t) -> u
    cb = ContinuousCallback((u,t,int) -> u - exp(1), terminate!)
    prob = ODEProblem(f, 1.0, (0.0, 2.0), callback=cb)
    sol = solve(prob, alg; kwargs...)
    sol.u[end] ≈ exp(1)
end

function test_callback_svector(alg; kwargs...)
    f = (u, p, t) -> u
    cb = ContinuousCallback((u,t,int) -> u[1] - exp(1), terminate!)
    prob = ODEProblem(f, SVector(1.0), (0.0, 2.0), callback=cb)
    sol = solve(prob, alg; kwargs...)
    sol.u[end][1] ≈ exp(1)
end

function test_callback_mvector(alg; kwargs...)
    f = (u, p, t) -> copy(u)
    cb = ContinuousCallback((u,t,int) -> u[1] - exp(1), terminate!)
    prob = ODEProblem(f, MVector(1.0), (0.0, 2.0), callback=cb)
    sol = solve(prob, alg; kwargs...)
    sol.u[end][1] ≈ exp(1)
end

println("inplace")
@test test_callback_inplace(BS3())
@test test_callback_inplace(BS5())
@test test_callback_inplace(SSPRK432())
@test test_callback_inplace(SSPRK932())
@test test_callback_inplace(OwrenZen3())
@test test_callback_inplace(OwrenZen4())
@test test_callback_inplace(OwrenZen5())
@test test_callback_inplace(DP5())
@test test_callback_inplace(DP8())
@test test_callback_inplace(Feagin10())
@test test_callback_inplace(Feagin12())
@test test_callback_inplace(Feagin14())
@test test_callback_inplace(TanYam7())
@test test_callback_inplace(Tsit5())
@test test_callback_inplace(TsitPap8())
@test test_callback_inplace(Vern6())
@test test_callback_inplace(Vern7())
@test test_callback_inplace(Vern8())
@test test_callback_inplace(Vern9())
@test test_callback_inplace(Rosenbrock23())
@test test_callback_inplace(Rosenbrock32())
@test test_callback_inplace(AB3(); dt=0.1)
@test test_callback_inplace(ABM32(); dt=0.1)
@test test_callback_inplace(AB4(); dt=0.1)
@test test_callback_inplace(ABM43(); dt=0.1)
@test test_callback_inplace(AB5(); dt=0.1)
@test test_callback_inplace(ABM54(); dt=0.1)
@test test_callback_inplace(VCAB3())
@test test_callback_inplace(VCAB4())
@test test_callback_inplace(VCAB5())
@test test_callback_inplace(VCABM3())
@test test_callback_inplace(VCABM4())
@test test_callback_inplace(VCABM5())

println("outofplace")
@test test_callback_outofplace(BS3())
@test test_callback_outofplace(BS5())
@test test_callback_outofplace(SSPRK432())
@test test_callback_outofplace(SSPRK932())
@test test_callback_outofplace(OwrenZen3())
@test test_callback_outofplace(OwrenZen4())
@test test_callback_outofplace(OwrenZen5())
@test test_callback_outofplace(DP5())
@test test_callback_outofplace(DP8())
@test test_callback_outofplace(Feagin10())
@test test_callback_outofplace(Feagin12())
@test test_callback_outofplace(Feagin14())
@test test_callback_outofplace(TanYam7())
@test test_callback_outofplace(Tsit5())
@test test_callback_outofplace(TsitPap8())
@test test_callback_outofplace(Vern6())
@test test_callback_outofplace(Vern7())
@test test_callback_outofplace(Vern8())
@test test_callback_outofplace(Vern9())
@test test_callback_outofplace(Rosenbrock23())
@test test_callback_outofplace(Rosenbrock32())
@test test_callback_outofplace(AB3(); dt=0.1)
@test test_callback_outofplace(ABM32(); dt=0.1)
@test test_callback_outofplace(AB4(); dt=0.1)
@test test_callback_outofplace(ABM43(); dt=0.1)
@test test_callback_outofplace(AB5(); dt=0.1)
@test test_callback_outofplace(ABM54(); dt=0.1)
@test test_callback_outofplace(VCAB3())
@test test_callback_outofplace(VCAB4())
@test test_callback_outofplace(VCAB5())
@test test_callback_outofplace(VCABM3())
@test test_callback_outofplace(VCABM4())
@test test_callback_outofplace(VCABM5())

println("scalar")
@test test_callback_scalar(BS3())
@test test_callback_scalar(BS5())
@test test_callback_scalar(SSPRK432())
@test test_callback_scalar(SSPRK932())
@test test_callback_scalar(OwrenZen3())
@test test_callback_scalar(OwrenZen4())
@test test_callback_scalar(OwrenZen5())
@test test_callback_scalar(DP5())
@test test_callback_scalar(DP8())
@test test_callback_scalar(Feagin10())
@test test_callback_scalar(Feagin12())
@test test_callback_scalar(Feagin14())
@test test_callback_scalar(TanYam7())
@test test_callback_scalar(Tsit5())
@test test_callback_scalar(TsitPap8())
@test test_callback_scalar(Vern6())
@test test_callback_scalar(Vern7())
@test test_callback_scalar(Vern8())
@test test_callback_scalar(Vern9())
@test test_callback_scalar(Rosenbrock23())
@test test_callback_scalar(Rosenbrock32())
@test test_callback_scalar(AB3(); dt=0.1)
@test test_callback_scalar(ABM32(); dt=0.1)
@test test_callback_scalar(AB4(); dt=0.1)
@test test_callback_scalar(ABM43(); dt=0.1)
@test test_callback_scalar(AB5(); dt=0.1)
@test test_callback_scalar(ABM54(); dt=0.1)
@test test_callback_scalar(VCAB3())
@test test_callback_scalar(VCAB4())
@test test_callback_scalar(VCAB5())
@test test_callback_scalar(VCABM3())
@test test_callback_scalar(VCABM4())
@test test_callback_scalar(VCABM5())

println("svector")
@test test_callback_svector(BS3())
@test test_callback_svector(BS5())
@test test_callback_svector(SSPRK432())
@test test_callback_svector(SSPRK932())
@test test_callback_svector(OwrenZen3())
@test test_callback_svector(OwrenZen4())
@test test_callback_svector(OwrenZen5())
@test test_callback_svector(DP5())
@test test_callback_svector(DP8())
@test test_callback_svector(Feagin10())
@test test_callback_svector(Feagin12())
@test test_callback_svector(Feagin14())
@test test_callback_svector(TanYam7())
@test test_callback_svector(Tsit5())
@test test_callback_svector(TsitPap8())
@test test_callback_svector(Vern6())
@test test_callback_svector(Vern7())
@test test_callback_svector(Vern8())
@test test_callback_svector(Vern9())
@test test_callback_svector(Rosenbrock23())
@test test_callback_svector(Rosenbrock32())
@test test_callback_svector(AB3(); dt=0.1)
@test test_callback_svector(ABM32(); dt=0.1)
@test test_callback_svector(AB4(); dt=0.1)
@test test_callback_svector(ABM43(); dt=0.1)
@test test_callback_svector(AB5(); dt=0.1)
@test test_callback_svector(ABM54(); dt=0.1)
@test test_callback_svector(VCAB3())
@test test_callback_svector(VCAB4())
@test test_callback_svector(VCAB5())
@test test_callback_svector(VCABM3())
@test test_callback_svector(VCABM4())
@test test_callback_svector(VCABM5())

println("mvector")
@test test_callback_mvector(BS3())
@test test_callback_mvector(BS5())
@test test_callback_mvector(SSPRK432())
@test test_callback_mvector(SSPRK932())
@test test_callback_mvector(OwrenZen3())
@test test_callback_mvector(OwrenZen4())
@test test_callback_mvector(OwrenZen5())
@test test_callback_mvector(DP5())
@test test_callback_mvector(DP8())
@test test_callback_mvector(Feagin10())
@test test_callback_mvector(Feagin12())
@test test_callback_mvector(Feagin14())
@test test_callback_mvector(TanYam7())
@test test_callback_mvector(Tsit5())
@test test_callback_mvector(TsitPap8())
@test test_callback_mvector(Vern6())
@test test_callback_mvector(Vern7())
@test test_callback_mvector(Vern8())
@test test_callback_mvector(Vern9())
@test test_callback_mvector(Rosenbrock23())
@test test_callback_mvector(Rosenbrock32())
@test test_callback_mvector(AB3(); dt=0.1)
@test test_callback_mvector(ABM32(); dt=0.1)
@test test_callback_mvector(AB4(); dt=0.1)
@test test_callback_mvector(ABM43(); dt=0.1)
@test test_callback_mvector(AB5(); dt=0.1)
@test test_callback_mvector(ABM54(); dt=0.1)
@test test_callback_mvector(VCAB3())
@test test_callback_mvector(VCAB4())
@test test_callback_mvector(VCAB5())
@test test_callback_mvector(VCABM3())
@test test_callback_mvector(VCABM4())
@test test_callback_mvector(VCABM5())

# Test ContinuousCallback hits values on the steps
t_event = 100.0
f_simple(u,p,t) = 1.00001*u
event_triggered = false
condition_simple(u,t,integrator) = t_event-t
function affect_simple!(integrator)
  global event_triggered
  event_triggered = true
end
cb = ContinuousCallback(condition_simple, nothing, affect_simple!)
prob = ODEProblem(f_simple, [1.0], (0.0, 2.0*t_event))
sol = solve(prob,Tsit5(),callback=cb, adaptive = false, dt = 10.0)
@test event_triggered


text:
@testset "Quadratic Forms" begin
    ndim = 10
    # A = randn(Float,ndim,ndim)
    A = VectorDataUtils.randposdefmatrix(Float,ndim)
    x = randn(Float,ndim)
    y = randn(Float,ndim)
    q = x'*A*x
    d = (x-y)'*A*(x-y)
    @test VectorDataUtils.xAx(x,A) ≈ q
    @test VectorDataUtils.sqmahaldist(x,y,A) ≈ d
    @test VectorDataUtils.mahaldist(x,y,A) ≈ sqrt(d)
end

ext: jl


text:
using SimpleValueGraphs.AbstractTuples

@testset "AbstractTuples" begin

tuples = [(), (1, ), (1, 2)]
named_tuples = [NamedTuple{}(), (a=1,), (a=1, b='2')]
all_tuples = [tuples; named_tuples]
not_abstract_tuples = [Int[], [1, 2], 1, Tuple, Tuple{Int, Int}, NamedTuple]

@testset "$t isa AbstractTuple" for t in all_tuples
    @test t isa AbstractTuple
end

@testset "$t isa AbstractTuple{$T}" for t in tuples, T in (typeof(t),)
    @test t isa AbstractTuple{T}
end

@testset "$t isa AbstractTuple{$T}" for
        t in named_tuples,
        T in (typeof(Tuple(t)),)

    @test t isa AbstractTuple{T}
end

@testset "$T <: AbstractTuple" for T in map(typeof, tuples)
    @test T <: AbstractTuple
end

@testset "$T <: AbstractTuple{$T}" for T in map(typeof, tuples)
    @test T <: AbstractTuple{T}
end

@testset "$T <: AbstractTuple{$TT}" for
        t in named_tuples,
        T in (typeof(t),),
        TT =  (typeof(Tuple(t)),)

    @test T <: AbstractTuple
end

@testset "!($t isa AbstractTuple" for t in not_abstract_tuples
    @test !(t isa AbstractTuple)
end

@test (a=1, b=2) isa NamedNTuple
@test (a=1, b=2) isa NamedNTuple{2}
@test (a=1, b=2) isa NamedNTuple{2, Int}
@test !((a=1, b=2) isa NamedNTuple{3})
@test !((a=1, b="2") isa NamedNTuple{2})
@test !((a=1, b=2) isa NamedNTuple{2, String})
@test !((1, 2) isa NamedNTuple)

@test (1, 2) isa AbstractNTuple
@test (a=1, b=2) isa AbstractNTuple
@test (1, 2) isa AbstractNTuple{2}
@test (a=1, b=2) isa AbstractNTuple{2}
@test (1, 2) isa AbstractNTuple{2, Int}
@test (a=1, b=2) isa AbstractNTuple{2, Int}
@test !((1, 2) isa AbstractNTuple{3})
@test !((a=1, b=2) isa AbstractNTuple{3})
@test !((1, "2") isa AbstractNTuple{2, Int})
@test !((a=1, b="2") isa AbstractNTuple{2, Int})
@test !((1, 2) isa AbstractNTuple{2, String})
@test !((a=1, b=2) isa AbstractNTuple{2, String})

end


text:
export solve_ω, _solve_Δω², solve_k, solve_n, solve, ng, k_guess, solve_nω, solve_ω², replan_ffts!, filter_eigs

"""
################################################################################
#																			   #
#	Routines to shield expensive initialization calculations from memory-	   #
#						intensive reverse-mode AD				  			   #
#																			   #
################################################################################
"""

import Base: show

function show(res::IterativeSolvers.LOBPCGResults;ind=1,color=:blue)
	itr_nums = map(x->getfield(x,:iteration)[ind],res.trace)
	residuals = map(x->getfield(x,:residual_norms)[ind],res.trace)
	log10res = log10.(residuals)
	log10res_min, log10res_max = extrema(log10res)
	resplt = UnicodePlots.lineplot(itr_nums,
		log10res;
		color,
		ylim=[floor(log10res_min),ceil(log10res_max)],
		xlim=[extrema(itr_nums)...],
		xlabel="iteration #",
		ylabel="log10(residuals)",
		title="LOBPCG Convergence (eig #: $ind)",
		)
	annotate!(resplt,:r,1,"converged: $(res.converged)")
	annotate!(resplt,:r,2,"iterations: $(res.iterations)   ($(res.maxiter) max)")
	annotate!(resplt,:r,3,"eigenvalue: $(res.λ[ind])")
	annotate!(resplt,:r,4,"residual_norm: $(res.residual_norms[ind])   (tol: $(res.tolerance))")
end

"""
################################################################################
#																			   #
#						solve_ω² methods: (ε⁻¹, k) --> (H, ω²)				   #
#																			   #
################################################################################
"""

# add try/catch with
# res = DFTK.LOBPCG(ms.M̂,rand(ComplexF64,size(ms.M̂)[1],1),I,ms.P̂,1e-8,3000)

# struct Eigensolver end
# struct IS_LOBPCG <: Eigensolver end
# struct DFTK_LOBPCG <: Eigensolver end

"""
f_filter takes in a ModeSolver and an eigenvalue/vector pair `αX` and outputs boolean,
ex. f_filter = (ms,αX)->sum(abs2,𝓟x(ms.grid)*αX[2])>0.9
where the modesolver `ms` is passed for access to any auxilary information
"""
function filter_eigs(ms::ModeSolver{ND,T},f_filter::Function)::Tuple{Vector{T},Matrix{Complex{T}}} where {ND,T<:Real}
	ω²H_filt = filter(ω²H->f_filter(ms,ω²H), [(real(ms.ω²[i]),ms.H⃗[:,i]) for i=1:length(ms.ω²)] )
	return copy(getindex.(ω²H_filt,1)), copy(hcat(getindex.(ω²H_filt,2)...)) # ω²_filt, H_filt
	# return getindex.(ω²H_filt,1), hcat(getindex.(ω²H_filt,2)...) # ω²_filt, H_filt
end

# function _solve_ω²(ms::ModeSolver{ND,T},::;nev=1,eigind=1,maxiter=100,tol=1.6e-8

function solve_ω²(ms::ModeSolver{ND,T};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing)::Tuple{Vector{T},Matrix{Complex{T}}} where {ND,T<:Real}

		# evals,evecs,convinfo = eigsolve(x->ms.M̂*x,ms.H⃗,size(ms.H⃗,2),:SR;maxiter,tol,krylovdim=50)
		evals,evecs,convinfo = eigsolve(x->ms.M̂*x,ms.H⃗,size(ms.H⃗,2),:SR;maxiter,tol,krylovdim=50) #,verbosity=2)
		copyto!(ms.H⃗,hcat(evecs...)[1:size(ms.H⃗,2)])
		copyto!(ms.ω²,evals[1:size(ms.H⃗,2)])

		# res = lobpcg!(ms.eigs_itr; log,not_zeros=false,maxiter,tol)

		# res = LOBPCG(ms.M̂,ms.H⃗,I,ms.P̂,tol,maxiter)
		# copyto!(ms.H⃗,res.X)
		# copyto!(ms.ω²,res.λ)

	if isnothing(f_filter)
		return   (copy(real(ms.ω²)), copy(ms.H⃗))
	else
		return filter_eigs(ms, f_filter)
	end
end

function solve_ω²(ms::ModeSolver{ND,T},k::TK,ε⁻¹::AbstractArray{SMatrix{3,3,T,9},ND};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,TK<:Union{T,SVector{3,T}}}
		# nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where T<:Real
	@ignore(update_k!(ms,k))
	@ignore(update_ε⁻¹(ms,ε⁻¹))
	solve_ω²(ms; nev, eigind, maxiter, tol, log, f_filter)
end

function solve_ω²(ms::ModeSolver{ND,T},k::TK;nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,TK<:Union{T,SVector{3,T}}}
		# nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where T<:Real
	@ignore(update_k!(ms,k))
	solve_ω²(ms; nev, eigind, maxiter, tol, log, f_filter)
end

"""
################################################################################
#																			   #
#						solve_k methods: (ε⁻¹, ω) --> (H, k)				   #
#																			   #
################################################################################
"""

"""
modified solve_ω version for Newton solver, which wants (x -> f(x), f(x)/f'(x)) as input to solve f(x) = 0
"""
function _solve_Δω²(ms::ModeSolver{ND,T},k::TK,ωₜ::T;nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,TK}
	# println("k: $(k)")
	ω²,H⃗ = solve_ω²(ms,k; nev, eigind, maxiter, tol, log, f_filter)
	Δω² = ω²[eigind] - ωₜ^2
	# ms.∂ω²∂k[eigind] = 2 * HMₖH(H⃗[:,eigind],ms.M̂.ε⁻¹,ms.M̂.mag,ms.M̂.m,ms.M̂.n) # = 2ω*ωₖ; ωₖ = ∂ω/∂kz = group velocity = c / ng; c = 1 here
	∂ω²∂k = 2 * HMₖH(H⃗[:,eigind],ms.M̂.ε⁻¹,ms.M̂.mag,ms.M̂.mn)
	ms.∂ω²∂k[eigind] = ∂ω²∂k
	# println("Δω²: $(Δω²)")
	# println("∂ω²∂k: $(∂ω²∂k)")
    return Δω² , Δω² / ∂ω²∂k #Δω² / copy(ms.∂ω²∂k[eigind])
end

function solve_k(ms::ModeSolver{ND,T},ω::T;nev=1,eigind=1,maxiter=100,tol=1e-8,atol=tol,maxevals=60,log=false,f_filter=nothing)::Tuple{T,Vector{Complex{T}}} where {ND,T<:Real} #
	if iszero(ms.M̂.k⃗[3])
		ms.M̂.k⃗ = SVector(0., 0., ω*ñₘₐₓ(ms.M̂.ε⁻¹))
	end
    kz = Roots.find_zero(x -> _solve_Δω²(ms,x,ω;nev,eigind,maxiter,tol,f_filter), ms.M̂.k⃗[3], Roots.Newton(); atol,maxevals,) #  verbose=true,)
	if isnothing(f_filter)
		Hv = ms.H⃗[:,eigind]
	else
		Hv = filter_eigs(ms, f_filter)[2][:,eigind]
	end
    return ( copy(kz), copy(Hv) ) # maybe copy(ds.H⃗) instead?
end

function solve_k(ms::ModeSolver{ND,T},ω::T,ε⁻¹::AbstractArray{T};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing)::Tuple{T,Vector{Complex{T}}} where {ND,T<:Real} #
	Zygote.@ignore(update_ε⁻¹(ms,ε⁻¹))
	solve_k(ms, ω; nev, eigind, maxiter, tol, log, f_filter)
end

function solve_k(ω::T,p::AbstractVector,geom_fn::F,grid::Grid{ND};kguess=nothing,Hguess=nothing,nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing)::Tuple{T,Vector{Complex{T}}} where {ND,T<:Real,F<:Function}
	ε⁻¹ = smooth(ω,p,:fεs,true,geom_fn,grid)
	ms = ignore() do
		kguess = isnothing(kguess) ? k_guess(ω,ε⁻¹) : kguess
		ms = ModeSolver(kguess, ε⁻¹, grid; nev, maxiter, tol)
		if !isnothing(Hguess)
			ms.H⃗ = reshape(Hguess,size(ms.H⃗))
		end
		return ms
	end
	solve_k(ms, ω; nev, eigind, maxiter, tol, log, f_filter)
end

function solve_k(ω::AbstractVector{T},p::AbstractVector,geom_fn::F,grid::Grid{ND};kguess=nothing,Hguess=nothing,nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,F<:Function}
	ε⁻¹ = smooth(ω,p,:fεs,true,geom_fn,grid)
	# ms = @ignore(ModeSolver(k_guess(first(ω),first(ε⁻¹)), first(ε⁻¹), grid; nev, maxiter, tol))
	ms = ignore() do
		kguess = isnothing(kguess) ? k_guess(ω,ε⁻¹) : kguess
		ms = ModeSolver(kguess, ε⁻¹, grid; nev, maxiter, tol)
		if !isnothing(Hguess)
			ms.H⃗ = Hguess
		end
		return ms
	end
	nω = length(ω)
	k = Buffer(ω,nω)
	Hv = Buffer([1.0 + 3.0im, 2.1+4.0im],(size(ms.M̂)[1],nω))
	for ωind=1:nω
		@ignore(update_ε⁻¹(ms,ε⁻¹[ωind]))
		kHv = solve_k(ms,ω[ωind]; nev, eigind, maxiter, tol, log, f_filter)
		k[ωind] = kHv[1]
		Hv[:,ωind] = kHv[2]
	end
	return copy(k), copy(Hv)
end

"""
################################################################################
#																			   #
#						`solve` methods: (ε⁻¹, ω) --> (E, n, ng, gvd)		   #
#																			   #
################################################################################
"""

function ∇ₖmag_mn2(māg,mn̄s,mag,mns)
	m = view(mns,1,:,:,:)
	n = view(mns,2,:,:,:)
	@tullio kp̂g_over_mag[i,ix,iy] := m[mod(i-2),ix,iy] * n[mod(i-1),ix,iy] / mag[ix,iy] - m[mod(i-1),ix,iy] * n[mod(i-2),ix,iy] / mag[ix,iy] (i in 1:3)
	kp̂g_over_mag_x_dk̂ = _cross(kp̂g_over_mag,dk̂)
	@tullio k̄_mag := māg[ix,iy] * mag[ix,iy] * kp̂g_over_mag[j,ix,iy] * dk̂[j]
	@tullio k̄_mn := -conj(mn̄s)[imn,i,ix,iy] * mns[imn,mod(i-2),ix,iy] * kp̂g_over_mag_x_dk̂[mod(i-1),ix,iy] + conj(mn̄s)[imn,i,ix,iy] * mns[imn,mod(i-1),ix,iy] * kp̂g_over_mag_x_dk̂[mod(i-2),ix,iy] (i in 1:3)
	k̄_magmn = k̄_mag + k̄_mn
	return k̄_magmn
end

function solve(ω::T,p::AbstractVector,geom_fn::F,grid::Grid{ND};kguess=nothing,Hguess=nothing,dk̂=SVector(0.0,0.0,1.0),nev=1,eigind=1,maxiter=500,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,F<:Function} # output type ::Tuple{T,T,T,Vector{Complex{T}}}
	ε,ε⁻¹,nng,nng⁻¹ = deepcopy(smooth(ω,p,(:fεs,:fεs,:fnn̂gs,:fnn̂gs),[false,true,false,true],geom_fn,grid));
	# ε = smooth(ω,p,:fεs,false,geom_fn,grid)
	# ε⁻¹ = smooth(ω,p,:fεs,true,geom_fn,grid)
	# nng = smooth(ω,p,:fnn̂gs,false,geom_fn,grid)
	# nng⁻¹ = smooth(ω,p,:fnn̂gs,true,geom_fn,grid)

	ngvd,nng2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs),[false,false],geom_fn,grid,volfrac_smoothing));
	ms = ignore() do
		kguess = isnothing(kguess) ? k_guess(ω,ε⁻¹) : kguess
		ms = ModeSolver(kguess, ε⁻¹, grid; nev, maxiter, tol)
		if !isnothing(Hguess)
			ms.H⃗ = Hguess
		end
		return ms
	end
	k, Hv = solve_k(ms,ω,ε⁻¹;nev,eigind,maxiter,tol,log, f_filter) #ω²_tol)
	neff::T = k/ω

	# calculate effective group index `ng`
	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	Ninv 		= 		1. / N(grid)
	mag,m⃗,n⃗ = mag_m_n(k,grid)
	m = flat(m⃗)
	n = flat(n⃗)
	mns = copy(vcat(reshape(m,1,3,Ns...),reshape(n,1,3,Ns...)))
    Hₜ = reshape(Hv,(2,Ns...))
	D = 1im * fft( kx_tc( Hₜ,mns,mag), _fftaxes(grid) )
	E = ε⁻¹_dot( D, ε⁻¹)
	# E = 1im * ε⁻¹_dot( fft( kx_tc( Hₜ,mns,mag), (2:1+ND) ), ε⁻¹)
	# H = inv(ω) * fft(tc(kx_ct( ifft( E, (2:1+ND) ), mns,mag), mns),(2:1+ND) )
	H = fft( tc(Hₜ,mns), (2:3) ) * (-1im * ω)
	P = 2*real(_sum_cross_z(conj(E),H))    # integrated Poyting flux parallel to ẑ: P = ∫dA S⃗⋅ẑ
	# W = dot(E,_dot((ε+nng),E))             # energy density per unit length
	W = real(dot(E,_dot(nng,E))) + (N(grid)* (ω^2))     # energy density per unit length
	ng = real( W / P )

	# calculate GVD = ∂(ng) / ∂ω = (∂²k)/(∂ω²)
	W̄ = inv(P)
	om̄₁₁ = 2*ω * N(grid) * W̄
	nnḡ = _outer(E,E) * W̄
	# H̄ = (-2*ng*W̄) * _cross(repeat([0.,0.,1.],outer=(1,Ns...)), E)
	# Ē = 2W̄*( _dot(nng,E) - ng * _cross(H,repeat([0.,0.,1.],outer=(1,Ns...))) )
	H̄ = (-2*ng*W̄) * _cross(dk̂, E)
	Ē = 2W̄*( _dot(nng,E) - ng * _cross(H,dk̂) )
	om̄₁₂ = dot(H,H̄) / ω
	om̄₁ = om̄₁₁ + om̄₁₂
	# eī₁ = _outer(Ē,D) ####################################
	𝓕⁻¹_ε⁻¹_Ē = bfft(ε⁻¹_dot( Ē, ε⁻¹),(2:3))
	𝓕⁻¹_H̄ = bfft( H̄ ,(2:3))
	H̄ₜ = 1im*( kx_ct(𝓕⁻¹_ε⁻¹_Ē,mns,mag) + ω*ct(𝓕⁻¹_H̄,mns) )
	local one_mone = [1.0im, -1.0im]
	@tullio 𝓕⁻¹_ε⁻¹_Ē_xHₜᵀ[i,j,ix,iy] := one_mone[i] * reverse(Hₜ;dims=1)[i,ix,iy] * conj(𝓕⁻¹_ε⁻¹_Ē)[j,ix,iy] nograd=one_mone
	∂ω²∂k_nd = 2 * HMₖH(Hv,ε⁻¹,mag,m,n)

	##### grad solve k
	# solve_adj!(λ⃗,M̂,H̄,ω^2,H⃗,eigind)
	M̂2 = HelmholtzMap(k,ε⁻¹,grid)
	λ⃗	= eig_adjt(
		M̂2,								 # Â
		ω^2, 							# α
		Hv, 					 		 # x⃗
		0.0, 							# ᾱ
		vec(H̄ₜ);								 # x̄
		# λ⃗₀,
		P̂	= HelmholtzPreconditioner(M̂2),
	)
	### k̄ₕ, eīₕ = ∇M̂(k,ε⁻¹,λ⃗,H⃗,grid)
	λ = reshape(λ⃗,(2,Ns...))
	λd 	= 	fft(kx_tc( λ , mns, mag ),_fftaxes(grid))
	# eīₕ	 = 	 ε⁻¹_bar(vec(D * (Ninv * -1.0im)), vec(λd), Ns...) ##########################
	λẽ  =   bfft(ε⁻¹_dot(λd , ε⁻¹),_fftaxes(grid))
	ẽ 	 =   bfft(E * -1.0im,_fftaxes(grid))
	@tullio mn̄s_kx0[i,j,ix,iy] := -1.0im * one_mone[i] * reverse(conj(Hₜ);dims=1)[i,ix,iy] * (Ninv*λẽ)[j,ix,iy] + -1.0im * one_mone[i] * reverse(conj(λ);dims=1)[i,ix,iy] * (Ninv*ẽ)[j,ix,iy]  nograd=one_mone
	# @tullio mn̄s_kx0[i,j,ix,iy] := -1.0im * one_mone[i] * reverse(conj(Hₜ);dims=1)[i,ix,iy] * λẽ[j,ix,iy] + -1.0im * one_mone[i] * reverse(conj(λ);dims=1)[i,ix,iy] * ẽ[j,ix,iy]  nograd=one_mone
	# @tullio mn̄s_kx0[i,j,ix,iy] := -1.0im * one_mone[i] * reverse(conj(Hₜ);dims=1)[i,ix,iy] * λẽ[j,ix,iy] + -1.0im * one_mone[i] * reverse(conj(λ);dims=1)[i,ix,iy] * ẽ[j,ix,iy]  nograd=one_mone
	@tullio mn̄s[i,j,ix,iy] := mag[ix,iy] * (mn̄s_kx0-conj(𝓕⁻¹_ε⁻¹_Ē_xHₜᵀ))[i,j,ix,iy]  + 1im*ω*conj(Hₜ)[i,ix,iy]*𝓕⁻¹_H̄[j,ix,iy]
	@tullio māg[ix,iy] := mns[a,b,ix,iy] * (mn̄s_kx0-conj(𝓕⁻¹_ε⁻¹_Ē_xHₜᵀ))[a,b,ix,iy]
	# k̄ = ∇ₖmag_mn(māg,mn̄s,mag,mns)
	@tullio kp̂g_over_mag[i,ix,iy] := m[mod(i-2),ix,iy] * n[mod(i-1),ix,iy] / mag[ix,iy] - m[mod(i-1),ix,iy] * n[mod(i-2),ix,iy] / mag[ix,iy] (i in 1:3)
	kp̂g_over_mag_x_dk̂ = _cross(kp̂g_over_mag,dk̂)
	@tullio k̄_mag := māg[ix,iy] * mag[ix,iy] * kp̂g_over_mag[j,ix,iy] * dk̂[j]
	@tullio k̄_mn := -conj(mn̄s)[imn,i,ix,iy] * mns[imn,mod(i-2),ix,iy] * kp̂g_over_mag_x_dk̂[mod(i-1),ix,iy] + conj(mn̄s)[imn,i,ix,iy] * mns[imn,mod(i-1),ix,iy] * kp̂g_over_mag_x_dk̂[mod(i-2),ix,iy] (i in 1:3)
	k̄ = k̄_mag + k̄_mn
	### \ k̄ₕ, eīₕ = ∇M̂(k,ε⁻¹,λ⃗,H⃗,grid)

	# combine k̄ₕ with k̄, scale by ( 2ω / ∂ω²∂k ) and calculate ω̄ and eīₖ
	λₖ  = ( k̄ / ∂ω²∂k_nd ) * Hₜ #reshape(λ⃗ₖ, (2,Ns...))
	λdₖ	=	fft(kx_tc( λₖ , mns, mag ),_fftaxes(grid))
	# eīₖ = ε⁻¹_bar(vec(D* (Ninv * -1.0im)), vec(λdₖ), Ns...) ####################################
	om̄₂  =  2ω * k̄ / ∂ω²∂k_nd
	##### \grad solve k
	om̄₃ = dot(herm(nnḡ), ngvd)
	om̄₄ = dot( herm(_outer(Ē+(λd+λdₖ)*(Ninv * -1.0im),D) ), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	# @show om̄₄_old = dot( ( eīₖ + eīₕ + eī₁ ), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	gvd = real( om̄₁ + om̄₂ + om̄₃ + om̄₄ )

	return ( neff, ng, gvd, E )
end

function solve_old(ω::T,p::AbstractVector,geom_fn::F,grid::Grid{ND};kguess=nothing,Hguess=nothing,dk̂=SVector(0.0,0.0,1.0),nev=1,eigind=1,maxiter=500,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,F<:Function} # output type ::Tuple{T,T,T,Vector{Complex{T}}}
	# ε⁻¹,nng⁻¹ = smooth(ω,p,(:fεs,:fnn̂gs),[true,true],geom_fn,grid)
	# ε⁻¹,nng,nng⁻¹ = deepcopy(smooth(ω,p,(:fεs,:fnn̂gs,:fnn̂gs),[true,false,true],geom_fn,grid));
	ε,ε⁻¹,nng,nng⁻¹ = deepcopy(smooth(ω,p,(:fεs,:fεs,:fnn̂gs,:fnn̂gs),[false,true,false,true],geom_fn,grid));
	# ngvd = copy(smooth((ω,),p,:fnĝvds,false,SMatrix{3,3,T,9}(0.,0.,0.,0.,0.,0.,0.,0.,0.),geom_fn,grid,volfrac_smoothing));
	ngvd,nng2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs),[false,false],geom_fn,grid,volfrac_smoothing));

	# ε⁻¹,nng⁻¹,ngvd = smooth(ω,p,(:fεs,:fnn̂gs,:fnĝvds),[true,true,false],geom_fn,grid)
	# ms = @ignore(ModeSolver(k_guess(ω,ε⁻¹), ε⁻¹, grid; nev, maxiter, tol))
	ms = ignore() do
		kguess = isnothing(kguess) ? k_guess(ω,ε⁻¹) : kguess
		ms = ModeSolver(kguess, ε⁻¹, grid; nev, maxiter, tol)
		if !isnothing(Hguess)
			ms.H⃗ = Hguess
		end
		return ms
	end
	# update_ε⁻¹(ms,ε⁻¹)
	# k, Hv = solve_k(ω,p,geom_fn,grid; nev, eigind, maxiter, tol, log, f_filter)
	k, Hv = solve_k(ms,ω,ε⁻¹;nev,eigind,maxiter,tol,log, f_filter) #ω²_tol)
	# (mag,m⃗,n⃗) = mag_m_n(k,dropgrad(ms.M̂.g⃗))
	# ng::T = ω / HMₖH(Hv,real(nng⁻¹),real(mag),real(reinterpret(reshape,T,m⃗)),real(reinterpret(reshape,T,n⃗))) #  material disp. included
	neff::T = k/ω

	# calculate effective group index `ng`
	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	mag,m⃗,n⃗ = mag_m_n(k,grid)
	mns = vcat(reshape(flat(m⃗),1,3,Ns...),reshape(flat(n⃗),1,3,Ns...))
    Hₜ = reshape(Hv,(2,Ns...))
	D = 1im * fft( kx_tc( Hₜ,mns,mag), (2:1+ND) )
	E = ε⁻¹_dot( D, ε⁻¹)
	# E = 1im * ε⁻¹_dot( fft( kx_tc( Hₜ,mns,mag), (2:1+ND) ), ε⁻¹)
	# H = inv(ω) * fft(tc(kx_ct( ifft( E, (2:1+ND) ), mns,mag), mns),(2:1+ND) )
    H = fft( tc(Hₜ,mns), (2:3) ) * (-1im * ω)
	P = 2*real(_sum_cross_z(conj(E),H))    # integrated Poyting flux parallel to ẑ: P = ∫dA S⃗⋅ẑ
	# W = dot(E,_dot((ε+nng),E))             # energy density per unit length
    W = real(dot(E,_dot(nng,E))) + (N(grid)* (ω^2))     # energy density per unit length
	ng = real( W / P )

	# calculate GVD = ∂(ng) / ∂ω = (∂²k)/(∂ω²)
	W̄ = inv(P)
	om̄₁₁ = 2*ω * N(grid) * W̄
	nnḡ = _outer(E,E) * W̄
	# H̄ = (-2*ng*W̄) * _cross(repeat([0.,0.,1.],outer=(1,Ns...)), E)
	# Ē = 2W̄*( _dot(nng,E) - ng * _cross(H,repeat([0.,0.,1.],outer=(1,Ns...))) )
	H̄ = (-2*ng*W̄) * _cross(dk̂, E)
	Ē = 2W̄*( _dot(nng,E) - ng * _cross(H,dk̂) )
	om̄₁₂ = dot(H,H̄) / ω
	om̄₁ = om̄₁₁ + om̄₁₂
	eī₁ = _outer(Ē,D)
	𝓕⁻¹_ε⁻¹_Ē = bfft(ε⁻¹_dot( Ē, ε⁻¹),(2:3))
	𝓕⁻¹_H̄ = bfft( H̄ ,(2:3))
	H̄ₜ = 1im*( kx_ct(𝓕⁻¹_ε⁻¹_Ē,mns,mag) + ω*ct(𝓕⁻¹_H̄,mns) )
	# 𝓕⁻¹_ε⁻¹_Ē_xHₜᵀ = 1im *_outer(_dot(repeat([0.0+0.0im 1.0+0.0im ;-1.0+0.0im 0.0+0.0im ],outer=(1,1,Ns...)), Hₜ), 𝓕⁻¹_ε⁻¹_Ē )
	local one_mone = [1.0im, -1.0im]
	@tullio 𝓕⁻¹_ε⁻¹_Ē_xHₜᵀ[i,j,ix,iy] := one_mone[i] * reverse(Hₜ;dims=1)[i,ix,iy] * conj(𝓕⁻¹_ε⁻¹_Ē)[j,ix,iy] nograd=one_mone
	@tullio māg[ix,iy] := mns[a,b,ix,iy] * -conj(𝓕⁻¹_ε⁻¹_Ē_xHₜᵀ)[a,b,ix,iy]
	mn̄s = -conj( 1im*ω*_outer(Hₜ,𝓕⁻¹_H̄) + _mult(𝓕⁻¹_ε⁻¹_Ē_xHₜᵀ,mag))
	# m̄ = reinterpret(reshape,SVector{3,T},view(mn̄s,1,:,:,:))
	# n̄ = reinterpret(reshape,SVector{3,T},view(mn̄s,2,:,:,:))
	m̄ = reinterpret(reshape,SVector{3,eltype(mn̄s)},view(mn̄s,1,:,:,:))
	n̄ = reinterpret(reshape,SVector{3,eltype(mn̄s)},view(mn̄s,2,:,:,:))
	k̄ = ∇ₖmag_m_n(māg,m̄,n̄,mag,m⃗,n⃗;dk̂)
	∂ω²∂k_nd = 2 * HMₖH(Hv,ε⁻¹,mag,flat(m⃗),flat(n⃗))
	( _, _, om̄₂, eī₂ ) = ∇solve_k(
		(k̄,vec(H̄ₜ)),
		(k,Hv),
		∂ω²∂k_nd,
		ω,
		ε⁻¹,
		grid;
		eigind,
	)
	om̄₃ = dot(herm(nnḡ), ngvd)
	om̄₄ = dot(herm(eī₁+eī₂), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	gvd = real( om̄₁ + om̄₂ + om̄₃ + om̄₄ )

	return ( neff, ng, gvd, E )
end

function solve(ω::AbstractVector{T},p::AbstractVector,geom_fn::F,grid::Grid{ND};nev=1,eigind=1,maxiter=500,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,F<:Function} # output type ::Tuple{T,T,T,Vector{Complex{T}}}
	nω = length(ω)
	neff = Buffer(ω,nω)
	ng = Buffer(ω,nω)
	gvd = Buffer(ω,nω)
	E = Buffer([1.0 + 3.0im, 2.1+4.0im],(3,size(grid)...,nω))
	ε⁻¹_nng⁻¹_ngvd = smooth(ω,p,(:fεs,:fnn̂gs,:fnĝvds),[true,true,false],geom_fn,grid)
	# create ModeSolver and solve for first frequency
	om = first(ω)
	ms = @ignore(ModeSolver(k_guess(om,ε⁻¹_nng⁻¹_ngvd[1,1]), ε⁻¹_nng⁻¹_ngvd[1,1], grid; nev, maxiter, tol))
	k,Hv = solve_k(ms,om; nev, eigind, maxiter, tol, log, f_filter)
	(mag,m⃗,n⃗) = mag_m_n(k,dropgrad(ms.M̂.g⃗))
	neff[1] = k/om
	ng[1] = om / HMₖH(Hv,ε⁻¹_nng⁻¹_ngvd[2,1],real(mag),real(reinterpret(reshape,T,m⃗)),real(reinterpret(reshape,T,n⃗))) #  material disp. included
	gvd[1] = 0.0
	E[:,:,:,1] = E⃗(k,reshape(Hv,(2,size(grid)...)),om,ε⁻¹_nng⁻¹_ngvd[1,1],ε⁻¹_nng⁻¹_ngvd[2,1],grid; normalized=true, nnginv=true)
	# solve at all remaining frequencies by updating the ModeSolver
	for ωind=2:nω
		om = ω[ωind]
		@ignore(update_ε⁻¹(ms,ε⁻¹_nng⁻¹_ngvd[1,ωind]))
		k,Hv = solve_k(ms,om; nev, eigind, maxiter, tol, log, f_filter)
		(mag,m⃗,n⃗) = mag_m_n(k,dropgrad(ms.M̂.g⃗))
		neff[ωind] = k/om
		ng[ωind] = om / HMₖH(Hv,ε⁻¹_nng⁻¹_ngvd[2,ωind],real(mag),real(reinterpret(reshape,T,m⃗)),real(reinterpret(reshape,T,n⃗))) #  material disp. included
		gvd[ωind] = 0.0
		E[:,:,:,ωind] = E⃗(k,reshape(Hv,(2,size(grid)...)),om,ε⁻¹_nng⁻¹_ngvd[1,ωind],ε⁻¹_nng⁻¹_ngvd[2,ωind],grid; normalized=true, nnginv=true)
	end
	return copy(neff), copy(ng), copy(gvd), copy(E)
end

# function solve_n(ms::ModeSolver{ND,T},ω::T,ε⁻¹::AbstractArray{<:SMatrix{3,3},ND},nnginv::AbstractArray{<:SMatrix{3,3},ND};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol,f_filter=nothing) where {ND,T<:Real}
# 	@ignore(update_ε⁻¹(ms,ε⁻¹))
# 	k, H⃗ = solve_k(ms,ω,ε⁻¹;nev,eigind,maxiter,tol,log, f_filter) #ω²_tol)
# 	(mag,m⃗,n⃗) = mag_m_n(k,dropgrad(ms.M̂.g⃗))
# 	ng = ω / HMₖH(H⃗,nnginv,real(mag),real(reinterpret(reshape,T,m⃗)),real(reinterpret(reshape,T,n⃗))) # new, material disp. included
# 	gvd = ∂²ω²∂k²(ω,geom,k,Hv,grid)
# 	neff, ng, gvd = neff_ng_gvd(ω,geom,k,Hv,ms.grid; eigind)
# 	E = E⃗(k,reshape(H⃗,(2,size(ms.grid)...)),ω,geom,ms.grid; svecs=false, normalized=true)
# 	return ( k/ω, ng, E )
# end

# function solve(ms::ModeSolver{ND,T},ω::T,geom::Geometry;nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real}
# 	ε⁻¹ = εₛ⁻¹(ω,geom;ms) # make_εₛ⁻¹(ω,shapes,dropgrad(ms))
# 	nnginv = nngₛ⁻¹(ω,geom;ms)
# 	# solve_n(ms, ω,ε⁻¹,nnginv; nev, eigind, maxiter, tol, log, f_filter)
# 	# update_ε⁻¹(ms,ε⁻¹)
# 	k, Hv = solve_k(ms,ω,ε⁻¹;nev,eigind,maxiter,tol,log, f_filter) #ω²_tol)
# 	(mag,m⃗,n⃗) = mag_m_n(k,dropgrad(ms.M̂.g⃗))
# 	# ng::T = ω / HMₖH(Hv,nnginv,real(mag),real(reinterpret(reshape,T,m⃗)),real(reinterpret(reshape,T,n⃗))) # new, material disp. included
# 	# neff::T = k/ω
# 	# gvd = 0.0
# 	neff, ng, gvd = neff_ng_gvd(ω,geom,k,Hv,ms.grid; eigind)
# 	E = E⃗(k,reshape(Hv,(2,size(ms.grid)...)),ω,geom,ms.grid; normalized=true)
# 	return ( neff, ng, gvd, E )
# end

function solve_n(ω::T,geom::Geometry,gr::Grid{ND};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real}
	ms = @ignore(ModeSolver(kguess(ω,geom), geom, gr;nev))
	solve_n(dropgrad(ms), ω, geom; nev, eigind, maxiter, tol, log, f_filter)
end

# prev method could also be:
# function solve_n(ω::Real,geom::Vector{<:Shape},gr::Grid{ND,T};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol, k₀=kguess(ω,geom)) where {ND,T<:Real}
# 	ms::ModeSolver{ND,T} = @ignore( ModeSolver(k₀, geom, gr) );
# 	solve_n(ms,ω,geom;nev,eigind,maxiter,tol,log)
# end

# function solve_n(ms::ModeSolver{ND,T},ω::Vector{T};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol) where {ND,T<:Real}
# 	nω = length(ω)
# 	n = Buffer(ω,nω)
# 	ng = Buffer(ω,nω)
# 	@inbounds for ωind=1:nω
# 		@inbounds nng = solve_n(ms,ω[ωind]; nev, eigind, maxiter, tol, log)
# 		@inbounds n[ωind] = nng[1]
# 		@inbounds ng[ωind] = nng[2]
# 	end
# 	return ( copy(n), copy(ng) )
# end
#
# function solve_n(ms::ModeSolver{ND,T},ω::Vector{T},ε⁻¹::AbstractArray{<:SMatrix{3,3},ND};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol) where {ND,T<:Real}
# 	@ignore(update_ε⁻¹(ms,ε⁻¹))
# 	nω = length(ω)
# 	n = Buffer(ω,nω)
# 	ng = Buffer(ω,nω)
# 	@inbounds for ωind=1:nω
# 		@inbounds nng = solve_n(ms,ω[ωind],ε⁻¹; nev, eigind, maxiter, tol, log)
# 		@inbounds n[ωind] = nng[1]
# 		@inbounds ng[ωind] .= nng[2]
# 	end
# 	return ( copy(n), copy(ng) )
# end

function replan_ffts!(ms::ModeSolver{3,T}) where T<:Real
	ms.M̂.𝓕! = plan_fft!(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)),(2:4),flags=FFTW.PATIENT);
	ms.M̂.𝓕⁻¹! = plan_bfft!(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)),(2:4),flags=FFTW.PATIENT);
	ms.M̂.𝓕 = plan_fft(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)),(2:4),flags=FFTW.PATIENT);
	ms.M̂.𝓕⁻¹ = plan_bfft(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)),(2:4),flags=FFTW.PATIENT);
end

function replan_ffts!(ms::ModeSolver{2,T}) where T<:Real
	ms.M̂.𝓕! = plan_fft!(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny)),(2:3),flags=FFTW.PATIENT);
	ms.M̂.𝓕⁻¹! = plan_bfft!(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny)),(2:3),flags=FFTW.PATIENT);
	ms.M̂.𝓕 = plan_fft(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny)),(2:3),flags=FFTW.PATIENT);
	ms.M̂.𝓕⁻¹ = plan_bfft(randn(Complex{T}, (3,ms.M̂.Nx,ms.M̂.Ny)),(2:3),flags=FFTW.PATIENT);
end

using Distributed

function _solve_n_serial(ms::ModeSolver{ND,T},ωs::Vector{T},geom::Geometry;nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol,wp=nothing,f_filter=nothing) where {ND,T<:Real}

	nω = length(ωs)
	ns = Buffer(ωs,nω)
	ngs = Buffer(ωs,nω)
	# @inbounds for ωind=1:nω
	# 	@inbounds nng = solve_n(ms, ωs[ωind], geom; nev, eigind, maxiter, tol, log)
	# 	# @inbounds nng = solve_n(ms, ωs[ωind], εₛ⁻¹(ωs[ωind],geom;ms=dropgrad(ms)); nev, eigind, maxiter, tol, log)
	# 	@inbounds ns[ωind] = nng[1]
	# 	@inbounds ngs[ωind] = nng[2]
	# end
	Srvol = S_rvol(geom;ms=dropgrad(ms))
	# ms_copies = @ignore( [ deepcopy(ms) for om in 1:length(ωs) ] )

	nω = length(ωs)
	n_buff = Buffer(ωs,nω)
	ng_buff = Buffer(ωs,nω)
	gvd_buff = Buffer(ωs,nω)
	E_buff = Buffer(ms.H⃗,(3,size(ms.grid)...,nω))
	for ωind=1:nω
		ωinv = inv(ωs[ωind])
		es = vcat(map(f->SMatrix{3,3}(f( ωinv )),geom.fεs),[εᵥ,])
		eis = inv.(es)	# corresponding list of inverse dielectric tensors for each material
		ε⁻¹_ω = εₛ⁻¹(es,eis,dropgrad(ms.sinds_proc),dropgrad(ms.minds),Srvol)
		# ei_new = εₛ⁻¹(es,eis,dropgrad(ps),dropgrad(minds),Srvol)  # new spatially smoothed ε⁻¹ tensor array
		k_ω, Hv_ω = solve_k(ms, ωs[ωind], ε⁻¹_ω; nev, eigind, maxiter, tol, log, f_filter) #ω²_tol)
		(mag,m⃗,n⃗) = mag_m_n(k_ω,dropgrad(ms.M̂.g⃗))
		### dispersive ng calculation
		# nngs_ω = vcat( nn̂g.(materials(geom), ωinv) ,[εᵥ,]) # = √.(ε̂) .* nĝ (elementwise product of index and group index tensors) for each material, vacuum permittivity tensor appended
		# nngis_ω = inv.(nngs_ω)
		# nngs_ω = vcat(map(f->SMatrix{3,3}(f( ωinv )),geom.fnn̂gs),[εᵥ,])
		# nngis_ω = inv.(nngs_ω)	# corresponding list of inverse dielectric tensors for each material
		# # nngi_ω = εₛ⁻¹(nngs,nngis,dropgrad(ps),dropgrad(minds),Srvol)  # new spatially smoothed ε⁻¹ tensor array
		# nngi_ω = εₛ⁻¹(nngs_ω,nngis_ω,dropgrad(ms.sinds_proc),dropgrad(ms.minds),Srvol)
		#
		# ng_ω = ωs[ωind] / HMₖH(Hv_ω[:,eigind],nngi_ω,real(mag),real(reinterpret(reshape,T,m⃗)),real(reinterpret(reshape,T,n⃗))) # new spatially smoothed ε⁻¹ tensor array
		# gvd_ω = ∂²ω²∂k²(ωs[ωind],geom,k_ω,Hv_ω,ms.grid; eigind)
		neff_ω, ng_ω, gvd_ω = neff_ng_gvd(ωs[ωind],geom,k_ω,Hv_ω,ms.grid; eigind)

		E_buff[1:3,axes(ms.grid)...,ωind] = E⃗(k_ω,reshape(Hv_ω,(2,size(ms.grid)...)),ωs[ωind],geom,ms.grid; normalized=true)
		n_buff[ωind] =  neff_ω #k_ω/ωs[ωind]
		ng_buff[ωind] = ng_ω
		gvd_buff[ωind] = gvd_ω
	end
	return ( copy(n_buff), copy(ng_buff), copy(gvd_buff), copy(E_buff) )
end

function _solve_n_parallel(ms::ModeSolver{ND,T},ωs::Vector{T},geom::Vector{<:Shape};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol,wp=nothing) where {ND,T<:Real}
	# pids = length(workers())<n_procs ? addprocs(n_procs) : workers()
	ω0 = ωs[Int(ceil(nω/2))]
	nng0 = solve_n(ms,ω0,shapes)
	ms_copies = [ deepcopy(ms) for om in 1:nω ]
	shapess = [deepcopy(shapes) for om in 1:nω ]
	nng = pmap(wp,ms_copies,ωs,shapess) do m,om,s
		@ignore( replan_ffts!(m) );
		solve_n(dropgrad(m),om,s)
	end
	n = [res[1] for res in nng]; ng = [res[2] for res in nng]
	return n, ng
end

function solve_n(ms::ModeSolver{ND,T},ωs::Vector{T},geom::Geometry;nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol,wp=nothing,f_filter=nothing) where {ND,T<:Real}
	_solve_n_serial(ms, ωs, geom; nev, eigind, maxiter, tol, log,f_filter=dropgrad(f_filter))
end

function solve_n(ωs::Vector{T},geom::Geometry,gr::Grid;nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol,wp=nothing,f_filter=nothing) where {ND,T<:Real}
	ms = @ignore(ModeSolver(kguess(first(ωs),geom), geom, gr))
	_solve_n_serial(ms,ωs, geom; nev, eigind, maxiter, tol, log,f_filter=dropgrad(f_filter))
end

	# ε⁻¹ = make_εₛ⁻¹(shapes,dropgrad(ms))
	# @ignore(update_ε⁻¹(ms,ε⁻¹))

	# m = @ignore( ModeSolver(1.45, shapes, 6., 4., 1., 128, 128, 1) )
	# ms_copies = @ignore( [ deepcopy(ms) for om in 1:length(ωs) ] )
	# nng = map((x,y)->solve_n(y,x,geom), ωs, ms_copies)
	# n = [res[1] for res in nng]
	# ng = [res[2] for res in nng]

	# Srvol = S_rvol(geom;ms=dropgrad(ms))
	# ms_copies = @ignore( [ deepcopy(ms) for om in 1:length(ωs) ] )

	# nω = length(ωs)
	# n_buff = Buffer(ωs,nω)
	# ng_buff = Buffer(ωs,nω)
	# for ωind=1:nω
	# 	# calculate ε⁻¹ for current ω
	# 	# m = ms_copies[ωind]
	# 	# @ignore( replan_ffts!(m) );
	# 	# es = vcat(εs(ms.geom,( 1. / ωs[ωind] )),[εᵥ,])		# dielectric tensors for each material, vacuum permittivity tensor appended
	# 	# eis = inv.(es)
	# 	# ε⁻¹_ω = εₛ⁻¹(es,eis,dropgrad(ms.sinds_proc),dropgrad(ms.minds),Srvol)
	# 	ε⁻¹_ω = εₛ⁻¹(ωs[ωind],geom;ms=dropgrad(ms))
	# 	# solve for n, ng with new ε⁻¹
	# 	nng = solve_n(ms,ωs[ωind],ε⁻¹_ω; nev, eigind, maxiter, tol, log)
	# 	# nng = solve_n(ms_copies[ωind],ω[ωind],ε⁻¹; nev, eigind, maxiter, tol, log)
	# 	n_buff[ωind] = nng[1]
	# 	ng_buff[ωind] = nng[2]
	# end
	# return ( copy(n_buff), copy(ng_buff) )

	# if isnothing(wp)			# solve_n for all ωs on current process
	# 	n_buff = Buffer(ωs,nω)
	# 	ng_buff = Buffer(ωs,nω)
	# 	for ωind=1:nω
	# 		nng = solve_n(ms,ωs[ωind],shapes; nev, eigind, maxiter, tol, log)
	# 		# nng = solve_n(ms_copies[ωind],ω[ωind],ε⁻¹; nev, eigind, maxiter, tol, log)
	# 		n_buff[ωind] = nng[1]
	# 		ng_buff[ωind] = nng[2]
	# 	end
	# 	# return ( copy(n), copy(ng) )
	# 	n = copy(n_buff); ng = copy(ng_buff);
	# else						# distribute solve_n's across worker pool wp
	# 	# pids = length(workers())<n_procs ? addprocs(n_procs) : workers()
	# 	ω0 = ωs[Int(ceil(nω/2))]
	# 	nng0 = solve_n(ms,ω0,shapes)
	# 	ms_copies = [ deepcopy(ms) for om in 1:nω ]
	# 	shapess = [deepcopy(shapes) for om in 1:nω ]
	# 	nng = pmap(wp,ms_copies,ωs,shapess) do m,om,s
	# 		@ignore( replan_ffts!(m) );
	# 		solve_n(dropgrad(m),om,s)
	# 	end
	# 	# nng = pmap(ωs) do om
	# 	# 	m = @ignore( ModeSolver(1.45, shapes, 6., 4., 1., 128, 128, 1) )
	# 	# 	solve_n(
	# 	# 		m,
	# 	# 		om,
	# 	# 		shapes,
	# 	# 		)
	# 	# end
	# 	n = [res[1] for res in nng]; ng = [res[2] for res in nng]
	# end
# 	return n, ng
# end
# 	n = Buffer(ω,nω)
# 	ng = Buffer(ω,nω)
# 	for ωind=1:nω
# 		# nng = solve_n(ms,ω[ωind],ε⁻¹; nev, eigind, maxiter, tol, log)
# 		nng = solve_n(ms_copies[ωind],ω[ωind],ε⁻¹; nev, eigind, maxiter, tol, log)
# 		n[ωind] = nng[1]
# 		ng[ωind] = nng[2]
# 	end
# 	return ( copy(n), copy(ng) )
# end

# nng = vmap(ms_copies,ω) do m,om
# 	solve_n(m,om,ε⁻¹)
# end
# n = [ res[1] for res in nng ]
# ng = [ res[2] for res in nng ]
# return n, ng
# end

# function solve(ms::ModeSolver{ND,T},ω::T,ε⁻¹::AbstractArray{<:SMatrix{3,3},ND},nng⁻¹::AbstractArray{<:SMatrix{3,3},ND};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false,ω²_tol=tol) where {ND,T<:Real}
# 	@ignore(update_ε⁻¹(ms,ε⁻¹))
# 	k, H⃗ = solve_k(ms,ω,ε⁻¹;nev,eigind,maxiter,tol,log) #ω²_tol)
# 	mag,m⃗,n⃗ = mag_m_n(k,grid)
#
# 	∂ω²∂k_nondisp 	= 2 * HMₖH(H⃗[:,eigind],ε⁻¹  ,real(mag),real(flat(m⃗)),real(flat(n⃗)))
# 	∂ω²∂k_disp 		= 2 * HMₖH(H⃗[:,eigind],nng⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗)))
# 	ng = 2ω / ∂ω²∂k_disp # new, material disp. included
# 	# calculate second order dispersion
# 	k̄, H̄, nngī  = ∇HMₖH(k,H⃗,nng⁻¹,grid; eigind)
# 	( _, _, om̄₁, eī₁ ) = ∇solve_k(	  (k̄,H̄),
# 									 	(k,H⃗[:,eigind]),
# 									  	∂ω²∂k_nondisp,
# 									   	ω,
# 									    ε⁻¹,
# 										grid; eigind)
#
#
# 	ng = 2ω / ∂ω²∂k_disp # new, material disp. included
#
# 	end
#     return ( k/ω, ng )
# end

"""
################################################################################
#																			   #
#					solve_nω methods: (ε⁻¹, k) --> (n, ng)					   #
#						(mostly for debugging gradients)					   #
#																			   #
################################################################################
"""

function solve_nω(ms::ModeSolver{T},k,ε⁻¹::AbstractArray{T,5};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false) where T<:Real
    ω², H⃗ = solve_ω²(ms,k,ε⁻¹;nev,eigind,maxiter,tol,log)
	ω = sqrt(ω²)
	(mag,m⃗,n⃗) = mag_m_n(k,dropgrad(ms.M̂.g⃗))
	ng = ω / HMₖH(H⃗,ε⁻¹,real(mag),real(reinterpret(reshape,Float64,m⃗)),real(reinterpret(reshape,Float64,n⃗)))
	return ( k/ω, ng )
end

# function solve_nω(ms::ModeSolver{T},k,shapes::Vector{<:Shape};nev=1,eigind=1,maxiter=100,tol=1e-8,log=false) where T<:Real
# 	g::MaxwellGrid = make_MG(ms.M̂.Δx,ms.M̂.Δy,ms.M̂.Δz,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)
# 	# ε⁻¹ = HybridArray{Tuple{3,3,Dynamic(),Dynamic(),Dynamic()},Float64,5,5,Array{Float64,5}}( make_εₛ⁻¹(shapes,g) )
# 	ε⁻¹ = make_εₛ⁻¹(shapes,dropgrad(ms))
# 	ω², H⃗ = solve_ω²(ms,k,ε⁻¹;nev,eigind,maxiter,tol,log)
# 	ω = sqrt(ω²)
# 	(mag,m⃗,n⃗) = mag_m_n(k,dropgrad(ms.M̂.g⃗))
# 	ng = ω / HMₖH(H⃗,ε⁻¹,real(mag),real(reinterpret(reshape,T,m⃗)),real(reinterpret(reshape,T,n⃗)))
# 	return ( k/ω, ng )
# end

# function solve_nω(kz::T,shapes::Vector{<:Shape},Δx,Δy,Δz,Nx,Ny,Nz;neigs=1,eigind=1,maxiter=100,tol=1e-8) where T<:Real
# 	# g::MaxwellGrid = make_MG(Zygote.dropgrad(Δx),Zygote.dropgrad(Δy),Zygote.dropgrad(Δz),Zygote.dropgrad(Nx),Zygote.dropgrad(Ny),Zygote.dropgrad(Nz)) #Δx,Δy,Δz,Nx,Ny,Nz)  	# MaxwellGrid(Δx,Δy,Δz,Nx,Ny,Nz)
# 	g::MaxwellGrid = make_MG(Δx,Δy,Δz,Nx,Ny,Nz)  	# MaxwellGrid(Δx,Δy,Δz,Nx,Ny,Nz)
# 	# ds::MaxwellData = make_MD(kz,g) # MaxwellData(kz,g)
# 	# kpg_mag,kpg_mn = calc_kpg(kz,Zygote.dropgrad(Δx),Zygote.dropgrad(Δy),Zygote.dropgrad(Δz),Zygote.dropgrad(Nx),Zygote.dropgrad(Ny),Zygote.dropgrad(Nz))
# 	# mag,mn = calc_kpg(kz,Δx,Δy,Δz,Nx,Ny,Nz)
# 	mag,mn = calc_kpg(kz,g.g⃗)
# 	ε⁻¹::Array{Float64,5} = make_εₛ⁻¹(shapes,g)
# 	H,ω² = solve_ω²(kz,ε⁻¹,Δx,Δy,Δz;neigs,eigind,maxiter,tol)
# 	# println("ω² = $ω²")
# 	ω = sqrt(ω²)
# 	Ha = reshape(H,(2,Nx,Ny,Nz))
# 	# ng = -ω / real( dot(Ha, kx_c2t( ifft( ε⁻¹_dot( fft( zx_t2c(Ha,ds.mn), (2:4) ), ε⁻¹), (2:4)),ds.mn,ds.mag) ) )
# 	ng = ω / HMₖH(Ha,ε⁻¹,mag,mn[:,1,:,:,:],mn[:,2,:,:,:])
# 	# ng = ω / real( dot(H, -vec( kx_c2t( ifft( ε⁻¹_dot( fft( zx_t2c(Ha,mn), (2:4) ), ε⁻¹), (2:4)),mn,mag) ) ) )
# 	# ng = -ω / real( dot(Ha, kx_c2t( ifft( ε⁻¹_dot( fft( zx_t2c(Ha,Zygote.@showgrad(mn)), (2:4) ), ε⁻¹), (2:4)), Zygote.@showgrad(mn),Zygote.@showgrad(mag)) ) )
# 	( kz/ω, ng )
# end

"""
################################################################################
#																			   #
#							   Plotting methods					   			   #
#																			   #
################################################################################
"""
function uplot(ms::ModeSolver;xlim=[0.5,1.8])
	ls_mats = uplot(ms.materials;xlim)
end

uplot(ch::IterativeSolvers.ConvergenceHistory; kwargs...) = lineplot(log10.(ch.data[:resnorm]); name="log10(resnorm)", kwargs...)


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct IoK8sApiCoreV1DownwardAPIVolumeFile <: SwaggerModel
    fieldRef::Any # spec type: Union{ Nothing, IoK8sApiCoreV1ObjectFieldSelector } # spec name: fieldRef
    mode::Any # spec type: Union{ Nothing, Int32 } # spec name: mode
    path::Any # spec type: Union{ Nothing, String } # spec name: path
    resourceFieldRef::Any # spec type: Union{ Nothing, IoK8sApiCoreV1ResourceFieldSelector } # spec name: resourceFieldRef

    function IoK8sApiCoreV1DownwardAPIVolumeFile(;fieldRef=nothing, mode=nothing, path=nothing, resourceFieldRef=nothing)
        o = new()
        validate_property(IoK8sApiCoreV1DownwardAPIVolumeFile, Symbol("fieldRef"), fieldRef)
        setfield!(o, Symbol("fieldRef"), fieldRef)
        validate_property(IoK8sApiCoreV1DownwardAPIVolumeFile, Symbol("mode"), mode)
        setfield!(o, Symbol("mode"), mode)
        validate_property(IoK8sApiCoreV1DownwardAPIVolumeFile, Symbol("path"), path)
        setfield!(o, Symbol("path"), path)
        validate_property(IoK8sApiCoreV1DownwardAPIVolumeFile, Symbol("resourceFieldRef"), resourceFieldRef)
        setfield!(o, Symbol("resourceFieldRef"), resourceFieldRef)
        o
    end
end # type IoK8sApiCoreV1DownwardAPIVolumeFile

const _property_map_IoK8sApiCoreV1DownwardAPIVolumeFile = Dict{Symbol,Symbol}(Symbol("fieldRef")=>Symbol("fieldRef"), Symbol("mode")=>Symbol("mode"), Symbol("path")=>Symbol("path"), Symbol("resourceFieldRef")=>Symbol("resourceFieldRef"))
const _property_types_IoK8sApiCoreV1DownwardAPIVolumeFile = Dict{Symbol,String}(Symbol("fieldRef")=>"IoK8sApiCoreV1ObjectFieldSelector", Symbol("mode")=>"Int32", Symbol("path")=>"String", Symbol("resourceFieldRef")=>"IoK8sApiCoreV1ResourceFieldSelector")
Base.propertynames(::Type{ IoK8sApiCoreV1DownwardAPIVolumeFile }) = collect(keys(_property_map_IoK8sApiCoreV1DownwardAPIVolumeFile))
Swagger.property_type(::Type{ IoK8sApiCoreV1DownwardAPIVolumeFile }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_IoK8sApiCoreV1DownwardAPIVolumeFile[name]))}
Swagger.field_name(::Type{ IoK8sApiCoreV1DownwardAPIVolumeFile }, property_name::Symbol) =  _property_map_IoK8sApiCoreV1DownwardAPIVolumeFile[property_name]

function check_required(o::IoK8sApiCoreV1DownwardAPIVolumeFile)
    (getproperty(o, Symbol("path")) === nothing) && (return false)
    true
end

function validate_property(::Type{ IoK8sApiCoreV1DownwardAPIVolumeFile }, name::Symbol, val)
end


text:
using BenchPerf
using BenchmarkTools
using Printf
using ToyStencils
@assert !ToyStencils.DEBUG_ENABLED

try
    using Base: @something
catch
    using Compat: @something
end

np = 14
n = 2^np
prob = ToyStencils.example_prob(; nx = n, nt = n)

@show sizeof(prob) / 2^20

print_percent(x) = print_percent(stdout, x)
function print_percent(io::IO, x)
    if x === nothing
        print(io, "??.?%")
    else
        @printf(io, "%.2f%%", x)
    end
end

print_label(label; kwargs...) = print_label(stdout, label; kwargs...)
function print_label(io::IO, label; prefix = ' ')
    print(io, prefix)
    printstyled(io, label; color = :light_black)
    print(io, ": ")
end

# trials = map([13, 9]) do p
trials = map(np:-1:max(1, np - 5)) do p
    print_label("smallsize"; prefix = "")
    @printf("2^%-2d", p)
    bench = @benchmarkable(
        ToyStencils.stencil!(prob; smallsize = 2^$p, basesize = Inf, timewindow = Inf)
    )
    t = BenchPerf.run(
        bench;
        event = [
            "branch-misses",
            "branches",
            "L1-dcache-load-misses",
            "L1-dcache-loads",
            "L1-icache-load-misses",
            "L1-icache-loads",
            "LLC-load-misses",
            "LLC-loads",
            "dTLB-load-misses",
            "dTLB-loads",
            "iTLB-load-misses",
            "iTLB-loads",
            "stalled-cycles-frontend",
            "stalled-cycles-backend",
            "cycles",
            "instructions",
        ],
    )
    print("  ")
    show(IOContext(stdout, :typeinfo => typeof(t.benchmark)), t.benchmark)
    print(" ")
    print_label("L1d miss")
    print_percent(t.perf.percent.l1_dcache_load_misses)
    print_label("L1i miss")
    print_percent(t.perf.percent.l1_icache_load_misses)
    print_label("br miss")
    print_percent(t.perf.percent.branch_misses)
    print_label("dTLB miss")
    print_percent(t.perf.percent.dtlb_load_misses)
    print_label("iTLB miss")
    print_percent(t.perf.percent.itlb_load_misses)
    print_label("stl.fr.")
    print_percent(t.perf.percent.stalled_cycles_frontend)
    print_label("stl.bk.")
    print_percent(t.perf.percent.stalled_cycles_backend)
    print_label("insn/cy")
    @printf("%.2f", something(t.perf.instructions_per_cycle, NaN))
    print_label("stl./insn")
    @printf("%.2f", something(t.perf.stalled_cycles_per_instructions, NaN))
    println()
    return 2^p => t
end

println()
for (n, t) in trials
    p = Int(log2(n))
    print_label("smallsize"; prefix = "")
    @printf("2^%-2d", p)
    print_label("memory")
    print(BenchmarkTools.memory(t.benchmark))
    print_label("#allocs")
    print(BenchmarkTools.allocs(t.benchmark))
    println()
end

if lowercase(get(ENV, "VERBOSE", "false")) == "true"
    for (n, t) in trials
        @show n
        display(t)
        println()
    end
end


text:
struct Visitante
    visitante

    function Visitante(ast::Node)

    	function mivisitante(x::Array{Node,1},rule::Rule)
    		for c in x
		     mivisitante(c,rule)
			end
		end

		function mivisitante(x::Node,rules::Rule)
			t = typeof(x)
			rules.enter(x)
			for f in fieldnames(t)
				subarbol= getfield(x, f)
				if( (typeof(subarbol )<: Node) & !(isa(subarbol, Diana.Name) ) )
					mivisitante(subarbol,rules)
				elseif (typeof(subarbol) <: Array{Node,1})
				    mivisitante(subarbol,rules)
				end
			end
			rules.leave(x)
		end

	    function visitante(rule::Rule)
			mivisitante(ast,rule)
	    end
      new(visitante)
    end
end

#=
struct VisitanteParallel
    visitante

    function VisitanteParallel(ast::Node)

    	function mivisitante(x::Array{Node,1},rule::Rule)
    		for c in x  # paralelizacion
		     mivisitante(c,rule)
			end
		end

		function mivisitante(x::Node,rules::Rule)
			t = typeof(x)
			rules.enter(x)
			for f in fieldnames(t)  # paralelizacion
				subarbol= getfield(x, f)
				if( (typeof(subarbol )<: Node) & !(isa(subarbol, Diana.Name) ) )
					mivisitante(subarbol,rules)
				elseif (typeof(subarbol) <: Array{Node,1})
				    mivisitante(subarbol,rules)
				end
			end
			rules.leave(x)
		end

	    function visitante(rule::Rule)
			mivisitante(ast,rule)
	    end
      new(visitante)
    end
end=#


text:
using Base: promote_op, @propagate_inbounds

function Base.show(io::IO, ::MIME"text/plain", k::Kernel)
    print(io, "Kernel{$(axes(k))} with window function $(k.f)")
    #print(code_lowered(k.f, (AbstractArray{Any},))[1])
end

Base.axes(::Type{<:Kernel{X}}) where X = X
Base.axes(k::Kernel) = axes(typeof(k))
Base.ndims(k::Kernel) = length(axes(k))
Base.size(k::Kernel) = length.(axes(k))
Base.keys(k::Kernel) = CartesianIndices(axes(k))

@inline (k::Kernel)(w::Window...) = k.f(w...)

# FIXME: revise the following questionable interface

"""
    eltype(k::Kernel, a::AbstractArray)

Infer the return type of `k` when applied to an interior window of `a`.
"""
Base.eltype(k::Kernel, a::AbstractArray...) = promote_op(k.f, map(ai -> wintype(k, ai), a)...)

wintype(k::Kernel, a::AbstractArray) =
    Window{eltype(a),ndims(a),axes(k),typeof(k),typeof(a)}
wintype(k::Kernel, a::ExtensionArray) =
    Window{eltype(a),ndims(a),axes(k),typeof(k),
    Array{Base.promote_type(eltype(a),eltype_extension(a)),ndims(a)}}

"""
    axes(k::Kernel, a::AbstractArray)

Return axes along which `k` can be applied to a window of `a`.
"""
@inline function Base.axes(k::Kernel, a::AbstractArray)
    ndims(a) == ndims(k) ||
        throw(DimensionMismatch("$(ndims(a)) vs $(ndims(k))"))

    extension(a) != ExtensionNone() && return axes(a)

    return map(axes(a), axes(k)) do ax, kx
        first(ax) - first(kx) : last(ax) - last(kx)
    end
end

"""
    size(k::Kernel, a::AbstractArray)

Return size of the cartesian region over which `k` can be applied to a window
of `a`.
"""
@inline Base.size(k::Kernel, a::AbstractArray) = length.(axes(k, a))

"""
    @kernel(wf)

Create a `Kernel` from the supplied anonymous window function `wf` by inferring
the kernel axes automatically.

For example the following two lines create functionally equivalent kernels:
```julia
@kernel w -> w[1] - w[0]
Kernel{(0:1,)}(w -> w[1] - w[0])
```

This macro currently expects an anonymous function and is intended purely as a
shorthand in the frequent case that your kernel axes can easily be inferred
from constant indices.
"""
macro kernel(wf)
    function walk(f, expr)
        f(expr) || return
        isa(expr, Expr) || return
        foreach(x -> walk(f, x), expr.args)
    end
    enclose(a) = first(a) : last(a)
    enclose(a, b) = min(first(a), first(b)) : max(last(a), last(b))

    wf.head == :-> ||
        throw(ArgumentError("anonymous function expected"))

    if wf.args[1] isa Symbol
        wfargs = [wf.args[1]]
    elseif wf.args[1] isa Expr && wf.args[1].head == :tuple
        wfargs = wf.args[1].args
    else
        throw(ArgumentError(error("unexpected function arguments")))
    end
    wf.args[2] = :(Base.@_inline_meta; $(wf.args[2]))
    wfbody = wf.args[2]

    d = nothing
    ax = nothing
    walk(wfbody) do x
        if x isa Expr && x.head == :ref && x.args[1] in wfargs
            all(x -> isa(x, Int), x.args[2:end]) ||
                throw(ArgumentError("""
                    encountered non-integer index, consider using constant
                    integer indices or the non-macro syntax instead
                """))

            if isnothing(d)
                d = length(x.args) - 1
                ax = enclose.(x.args[2:end])
            else
                d == length(x.args) - 1 ||
                    throw(ArgumentError("window index dimensions don't match"))
                ax = enclose.(ax, x.args[2:end])
            end
            y = copy(x)
            x.head = :macrocall
            x.args = [Symbol("@inbounds"), LineNumberNode(0), y]
            return false
        end
        return true
    end
    (isnothing(d) || isnothing(ax)) &&
        error("could not determine kernel axes")
    return :( $Kernel{($(ax...),)}($(esc(wf))) )
end


text:
"""
    SurfaceFluxes

Surface fluxes parameters, shared by multiple models
including the atmosphere, ocean.
"""
module SurfaceFluxes

include("UniversalFunctions.jl")

end # module SurfaceFluxes

max_stars_repo_name: climate-machine/CLIMAParameters

max_issues_repo_name: climate-machine/CLIMAParameters

max_forks_repo_name: climate-machine/CLIMAParameters


text:
# call stack:
# run/replay
#   > calls apply(GreensIterator, group, dqmc, model, sweep)
#       > calls measure(LatticeIterator, measurement, dqmc, model, sweep, GreensResults...)
#           > calls apply!(LatticeIterator, measurement, dqmc, model, GreensResults...)
#               > calls kernel(measurement, dqmc, model, MaskResult..., GreensResults...)

# Does:
# creates groups, runs simulation
#   > resolves GreensIterator (calclates G, Gkk, Gkl or whatever)
#       > resolves sweep (skipping, frequency based measurements) and commits
#           > resolves Mask (to Greens indices)
#               > calculate element from Wicks theorem

# TODO
# try replacing global function rather than creatin new local ones

struct DQMCMeasurement{GI, LI, F <: Function, OT, T} <: AbstractMeasurement
    greens_iterator::GI
    lattice_iterator::LI
    kernel::F
    observable::OT
    temp::T
end

missing_kernel(args...) = error("kernel couldn't be loaded.")

function DQMCMeasurement(
        m::DQMCMeasurement;
        greens_iterator = m.greens_iterator, lattice_iterator = m.lattice_iterator,
        kernel = m.kernel,
        observable = m.observable, temp = m.temp,
        capacity = nothing
    )
    if capacity === nothing
        DQMCMeasurement(greens_iterator, lattice_iterator, kernel, observable, temp)
    else
        binner = rebuild(observable, capacity)
        DQMCMeasurement(greens_iterator, lattice_iterator, kernel, binner, temp)
    end
end
rebuild(B::LogBinner, capacity) = LogBinner(B, capacity=capacity)
rebuild(B::T, capacity) where T = T(B, capacity=capacity)

function Measurement(
        dqmc, _model, greens_iterator, lattice_iterator, kernel;
        capacity = _default_capacity(dqmc), eltype = geltype(dqmc),
        temp = _measurement_buffer(dqmc, _model, lattice_iterator, eltype),
        obs = LogBinner(
            _binner_zero_element(dqmc, _model, lattice_iterator, eltype),
            capacity=capacity
        ),
        # let
        #     shape = _get_final_shape(dqmc, _model, lattice_iterator)
        #     _zero = shape === nothing ? zero(eltype) : zeros(eltype, shape)
        #     LogBinner(_zero, capacity=capacity)
        # end
    )
    DQMCMeasurement(greens_iterator, lattice_iterator, kernel, obs, temp)
end

################################################################################
### DQMCMeasurement utilities
################################################################################

function Base.show(io::IO, ::MIME"text/plain", m::DQMCMeasurement)
    max = applicable(capacity, m.observable) ? capacity(m.observable) : Inf
    current = length(m.observable)
    GI = m.greens_iterator
    LI = if m.lattice_iterator isa Type
        m.lattice_iterator
    else
        typeof(m.lattice_iterator)
    end
    print(io, "[$current/$max] DQMCMeasurement($GI, $LI)")
end

MonteCarloObservable.mean(m::DQMCMeasurement) = mean(m.observable)
MonteCarloObservable.var(m::DQMCMeasurement) = var(m.observable)
MonteCarloObservable.std_error(m::DQMCMeasurement) = std_error(m.observable)
MonteCarloObservable.tau(m::DQMCMeasurement) = tau(m.observable)
Base.length(m::DQMCMeasurement) = length(m.observable)
Base.isempty(m::DQMCMeasurement) = isempty(m.observable)
Base.empty!(m::DQMCMeasurement) = empty!(m.observable)

"""
    _default_capacity(mc::DQMC)

Returns a default capacity based in the number of sweeps and the measure rate.
"""
_default_capacity(mc::DQMC) = 2 * ceil(Int, mc.parameters.sweeps / mc.parameters.measure_rate)

################################################################################
### Buffers
################################################################################

# function _measurement_buffer(mc::DQMC, m::Model, li, eltype)
#     shape = _get_temp_shape(dqmc, _model, lattice_iterator)
#     shape === nothing ? nothing : Array{eltype}(undef, shape)
# end

# General (Float64)
_measurement_buffer(mc, model, li, eltype) = _simple_buffer(mc, li, eltype)
_measurement_buffer(mc, model, ::Sum, eltype) = zeros(eltype, 1)
function _measurement_buffer(mc, model, ::LatticeIterationWrapper{LI}, eltype) where {LI}
    _measurement_buffer(mc, model, LI, eltype)
end

# StructArrays / ComplexF64
# This unwraps wrapper to see if any of the wrapped types matches what we're
# looking for. Currently looks for T1 in T2
function is_wrapped_type_of(T1, T2)
    if T1 == T2
        return true
    else
        for t2 in T2.types
            is_subtype_of(T1, t2) && return true
        end
    end
    false
end

function maybe_structarray(mc, A)
    is_wrapped_type_of(CMat64, gmattype(mc)) ? StructArray(A) : A
end

_measurement_buffer(mc, model, li, ::ComplexF64) = maybe_structarray(mc, _simple_buffer(mc, li, eltype))
_measurement_buffer(mc, model, ::Sum, ::ComplexF64) = maybe_structarray(mc, zeros(ComplexF64, 1))
function _measurement_buffer(mc, model, ::LatticeIterationWrapper{LI}, ::ComplexF64) where {LI}
    maybe_structarray(mc, _measurement_buffer(mc, model, LI, eltype))
end

# can be determined from type
_simple_buffer(mc, t::Type, T) = _simple_buffer(mc, t(), T)
_simple_buffer(mc, ::Type{Nothing}, T) = nothing
_simple_buffer(mc, ::Nothing, T) = nothing
_simple_buffer(mc, ::EachSite, T) = zeros(T, length(lattice(mc)))
_simple_buffer(mc, ::EachSiteAndFlavor, T) = zeros(T, nflavors(field(mc)) * length(lattice(mc)))
_simple_buffer(mc, ::EachSitePair, T) = zeros(T, length(lattice(mc)), length(lattice(mc)))

# determined from type data
_simple_buffer(mc, li::DeferredLatticeIteratorTemplate, T) = zeros(T, ndirections(mc, li)) # TODO

# zero element for binner
_binner_zero_element(mc, model, li, eltype) = _simple_buffer(mc, li, eltype)
_binner_zero_element(mc, model, ::Type{Nothing}, eltype) = zero(eltype)
_binner_zero_element(mc, model, ::Nothing, eltype) = zero(eltype)
_binner_zero_element(mc, model, ::Sum, eltype) = zero(eltype)
# _binner_zero_element(mc, model, ::SuperfluidDensity, eltype) = zero(eltype)
function _binner_zero_element(mc, model, li::ApplySymmetries{LI, N}, eltype) where {LI, N}
    if LI <: EachLocalQuadByDistance || LI <: EachLocalQuadBySyncedDistance
        return zeros(eltype, first(ndirection(LI(mc, model))), N) # TODO
    else
        throw(MethodError(_binner_zero_element, (mc, model, li, eltype)))
    end
end

# Once we change greens iterators to follow the "template -> iter" structure
# this will be compat only
# Some type piracy to make things easier
Base.Nothing(::DQMC, ::Model) = nothing

# To identify the requirement of equal-time Greens functions
struct Greens <: AbstractGreensIterator end
Greens(::DQMC, ::Model)= Greens()

struct GreensAt <: AbstractUnequalTimeGreensIterator
    k::Int
    l::Int
end
GreensAt(l::Integer) = GreensAt(l, l)
GreensAt(k::Integer, l::Integer) = GreensAt(k, l)
# GreensAt{k, l}(::DQMC, ::Model) where {k, l} = GreensAt(k, l)
Base.:(==)(a::GreensAt, b::GreensAt) = (a.k == b.k) && (a.l == b.l)

requires(::AbstractMeasurement) = (Nothing, Nothing)
requires(m::DQMCMeasurement) = (m.greens_iterator, m.lattice_iterator)

@bm function generate_groups(mc, model, measurements)
    empty!(mc.lattice_iterator_cache)

    # get unique requirements
    requirements = requires.(measurements)
    GIs = unique(first.(requirements))

    # init requirements
    # lattice_iterators = map(T -> T(mc, model), LIs)
    if any(x -> (x isa Type ? x : typeof(x)) <: AbstractUnequalTimeGreensIterator, GIs)
        initialize_stack(mc, mc.ut_stack)
    end

    # Group measurements with the same greens iterator together
    # lattice_iterators
    # greens_iterator => [
    #     (lattice_iterator, measurement),
    #     (lattice_iterator, measurement),
    #     ...
    # ]
    output = map(enumerate(GIs)) do (i, G)
        ms = filter(m -> G == requires(m)[1], measurements)
        group = map(ms) do m
            LI = requires(m)[2]
            (LI === nothing ? nothing : LI(mc, model), m)
        end
        (G isa Type ? G(mc, model) : G) => group
    end

    if length(measurements) != mapreduce(x -> length(x[2]), +, output, init = 0)
        for (G, group) in output
            println(G)
            for (li, m) in group
                println("\t", typeof(li), typeof(m.kernel))
            end
        end
        N = length(measurements)
        M = mapreduce(x -> length(x[2]), +, output, init = 0)
        error("Oh no. We lost some measurements. $N -> $M")
    end

    return output
end

lattice_iterator(m::DQMCMeasurement, mc, model) = m.lattice_iterator(mc, model)

# TODO
# Saving the kernel function as a symbol is kinda risky because the function
# definition is not guaranteed to be the same
# though maybe that's a good thing - changing function definitions doesn't
# break stuff this way
function _save(file::JLDFile, m::DQMCMeasurement, key::String)
    write(file, "$key/VERSION", 1)
    write(file, "$key/tag", "DQMCMeasurement")
    write(file, "$key/GI", m.greens_iterator)
    write(file, "$key/LI", m.lattice_iterator)
    # maybe add module for eval?
    write(file, "$key/kernel", Symbol(m.kernel))
    write(file, "$key/obs", m.observable)
    write(file, "$key/temp", m.temp)
end

function _load(data, ::Val{:DQMCMeasurement})
    temp = haskey(data, "temp") ? data["temp"] : data["output"]

    kernel = try
        eval(data["kernel"])
    catch e
        @warn "Failed to load kernel in module MonteCarlo." exception=e
        missing_kernel
    end
    DQMCMeasurement(data["GI"], data["LI"], kernel, data["obs"], temp)
end

to_tag(::Type{<: DQMCMeasurement}) = Val(:DQMCMeasurement)

################################################################################
### Greens function related
################################################################################

@bm function apply!(::Nothing, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    for (lattice_iterator, measurement) in combined
        # Clear temp if necessary
        prepare!(lattice_iterator, model, measurement)
        # Write measurement to ouput
        measure!(lattice_iterator, measurement, mc, model, sweep)
        # Finalize computation (temp) and commit
        finish!(lattice_iterator, model, measurement)
    end

    nothing
end

@bm function apply!(::Greens, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    G = greens!(mc)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
        measure!(lattice_iterator, measurement, mc, model, sweep, G)
        finish!(lattice_iterator, model, measurement)
    end
    nothing
end

@bm function apply!(g::GreensAt, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    G = greens!(mc, g.k, g.l)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
        measure!(lattice_iterator, measurement, mc, model, sweep, G)
        finish!(lattice_iterator, model, measurement)
    end
    nothing
end

@bm function apply!(iter::CombinedGreensIterator, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
    end

    G00 = greens!(mc)
    for (G0l, Gl0, Gll) in init(mc, iter)
        for (lattice_iterator, measurement) in combined
            measure!(lattice_iterator, measurement, mc, model, sweep, (G00, G0l, Gl0, Gll))
        end
    end

    for (lattice_iterator, measurement) in combined
        finish!(lattice_iterator, model, measurement, mc.parameters.delta_tau)
    end
    nothing
end

@bm function apply!(iter::AbstractGreensIterator, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
    end

    for packed_greens in init(mc, iter)
        for (lattice_iterator, measurement) in combined
            measure!(lattice_iterator, measurement, mc, model, sweep, packed_greens)
        end
    end

    for (lattice_iterator, measurement) in combined
        finish!(lattice_iterator, model, measurement, 1.0)
    end
    nothing
end

################################################################################
### measure!
################################################################################

@bm function measure!(lattice_iterator, measurement, mc::DQMC, model, sweep, packed_greens)
    # ignore sweep
    apply!(measurement.temp, lattice_iterator, measurement, mc, model, packed_greens)
    nothing
end

# Lattice irrelevant
@bm function measure!(::Nothing, measurement, mc::DQMC, model, sweep, packed_greens)
    flv = Val(max(nflavors(field(mc)), nflavors(model)))
    push!(measurement.observable, measurement.kernel(mc, model, packed_greens, flv))
    nothing
end

################################################################################
### apply Lattice Iterators
################################################################################

# Call kernel for each site (linear index)
@bm function apply!(temp::Array, iter::DirectLatticeIterator, measurement, mc::DQMC, model, packed_greens)
    flv = Val(max(nflavors(field(mc)), nflavors(model)))
    for i in iter
        temp[i] += measurement.kernel(mc, model, i, packed_greens, flv)
    end
    nothing
end

# Call kernel for each pair (src, trg) (Nsties² total)
@bm function apply!(temp::Array, iter::EachSitePair, measurement, mc::DQMC, model, packed_greens)
    flv = Val(max(nflavors(field(mc)), nflavors(model)))
    for (i, j) in iter
        temp[i, j] += measurement.kernel(mc, model, (i, j), packed_greens, flv)
    end
    nothing
end

# Call kernel for each pair (site, site) (i.e. on-site)
@bm function apply!(temp::Array, iter::_OnSite, measurement, mc::DQMC, model, packed_greens)
    flv = Val(max(nflavors(field(mc)), nflavors(model)))
    for (i, j) in iter
        temp[i] += measurement.kernel(mc, model, (i, j), packed_greens, flv)
    end
    nothing
end

@bm function apply!(temp::Array, iter::DeferredLatticeIterator, measurement, mc::DQMC, model, packed_greens)
    flv = Val(max(nflavors(field(mc)), nflavors(model)))
    @inbounds for idxs in iter
        temp[first(idxs)] += measurement.kernel(mc, model, idxs[2:end], packed_greens, flv)
    end
    nothing
end

# Sums
@bm function apply!(temp::Array, iter::_Sum{<: DirectLatticeIterator}, measurement, mc::DQMC, model, packed_greens)
    flv = Val(max(nflavors(field(mc)), nflavors(model)))
    @inbounds for idxs in iter
        temp[1] += measurement.kernel(mc, model, idxs, packed_greens, flv)
    end
    nothing
end
@bm function apply!(temp::Array, iter::_Sum{<: DeferredLatticeIterator}, measurement, mc::DQMC, model, packed_greens)
    flv = Val(max(nflavors(field(mc)), nflavors(model)))
    @inbounds for idxs in iter
        temp[1] += measurement.kernel(mc, model, idxs[2:end], packed_greens, flv)
    end
    nothing
end

@inline function apply!(temp::Array, s::LatticeIterationWrapper, m, mc, model, pg)
    apply!(temp, s.iter, m, mc, model, pg)
end

################################################################################
### LatticeIterator preparation and finalization
################################################################################

# If LatticeIterator is Nothing, then things should be handled in measure!
@inline prepare!(::Nothing, model, m) = nothing
@inline prepare!(::AbstractLatticeIterator, model, m) = m.temp .= zero(eltype(m.temp))
@inline prepare!(s::_Sum, args...) = prepare!(s.iter, args...)

@inline finish!(::Nothing, args...) = nothing # handled in measure!
@inline function finish!(li, model, m, factor=1.0)
    finalize_temp!(li, model, m, factor)
    commit!(li, m)
end

@inline function finalize_temp!(::AbstractLatticeIterator, model, m, factor)
    m.temp .*= factor
end
@inline function finalize_temp!(::DeferredLatticeIterator, model, m, factor)
    m.temp .*= factor / length(lattice(model))
end
@inline function finalize_temp!(s::LatticeIterationWrapper, model, m , factor)
    finalize_temp!(s.iter, model, m, factor)
end

@inline commit!(::AbstractLatticeIterator, m) = push!(m.observable, m.temp)
@inline commit!(::_Sum, m) = push!(m.observable, m.temp[1])

function commit!(s::_ApplySymmetries{<: EachLocalQuadByDistance}, m)
    final = zeros(eltype(m.temp), size(m.temp, 1), length(s.symmetries))
    # This calculates
    # ∑_{a a'} O(Δr, a, a') f_ζ(a) f_ζ(a')
    # where a, a' are typically nearest neighbor directions and
    # f_ζ(a) is the weight in direction a for a symmetry ζ
    for (i, sym) in enumerate(s.symmetries)
        for k in 1:length(sym), l in 1:length(sym)
            @. final[:, i] += m.temp[:, k, l] * sym[k] * sym[l]
        end
    end
    push!(m.observable, final)
end
function commit!(s::_ApplySymmetries{<: EachLocalQuadBySyncedDistance}, m)
    final = zeros(eltype(m.temp), size(m.temp, 1), length(s.symmetries))
    # Same as above but with a = a'
    for (i, sym) in enumerate(s.symmetries)
        for k in 1:length(sym)
            @. final[:, i] += m.temp[:, k] * sym[k] * sym[k]
        end
    end
    push!(m.observable, final)
end


text:
using Lazy
import Lazy: cycle, range, drop, take
using FactCheck

facts("Lists") do
    @fact list(1, 2, 3)[2] --> 2
    @fact prepend(1, list(2,3,4)) --> 1:list(2, 3, 4)
    @fact seq([1, 2, 3]) --> list(1, 2, 3)
    @pending seq(1:3) --> list(1, 2, 3)
    @fact constantly(1)[50] --> 1
    testfn() = 1
    @fact repeatedly(testfn)[50] --> 1
    @fact cycle([1, 2, 3])[50] --> 2
    @fact iterate(x->x^2, 2)[3] --> 16
    @fact range(1, 5)[3] --> 3
    @fact range(1, 5)[10] --> nothing
    @fact range(1, 5)[-1] --> 1
    @fact list(1, 2, 3) * list(4, 5, 6) --> list(1, 2, 3, 4, 5, 6)
    @fact first(list(1, 2, 3)) --> 1
    @fact tail(list(1, 2, 3)) --> list(2, 3)
    @fact flatten(list(1,2,list(3,4))) --> list(1, 2, 3, 4)
    @fact list(1,2,list(3,4))[3] --> list(3, 4)
end

facts("Fibs") do
    fibs = @lazy 0:1:(fibs + drop(1, fibs));
    @fact fibs[20] --> 4181
    @fact take(5, fibs) --> list(0, 1, 1, 2, 3)
end

facts("Primes") do
    isprime(n) =
        @>> primes begin
            take_while(x -> x<=sqrt(n))
            map(x -> n % x == 0)
            any; !
        end
    primes = filter(isprime, range(2));
end

facts("Even squares") do
    esquares = @>> range() map(x->x^2) filter(iseven);
    @fact take(5, esquares) --> list(4, 16, 36, 64, 100)
end

facts("Threading macros") do
    temp = @> [2 3] sum
    @fact temp --> 5
    # Reverse from after index 2
    temp = @>> 2 reverse([1, 2, 3, 4, 5])
    @fact temp --> [1, 5, 4, 3, 2]
    temp = @as x 2 begin
        x^2
        x + 2
    end
    @fact temp --> 6
end

FactCheck.exitstatus()


text:
using Distributions
struct OrthoNNDist <: DiscreteMultivariateDistribution
	x0::Vector{Float64}
	oc::Array{Float64,2}
	x1s::Array
	prob::Float64
    #return a new uniform distribution with all vectors in x1s orthogonal to oc
	function OrthoNNDist(x0::Vector{Float64}, oc::Array{Float64,2})
		x1s = []
		for i = 1:size(oc)[2]
			x1 = x0 + oc[:, i]
			if nonneg(x1)
				push!(x1s, x1)
			end
			x1 = x0 - oc[:, i]
			if nonneg(x1)
				push!(x1s, x1)
			end
		end
		new(x0, oc, x1s, 1.0/length(x1s))
	end
end

Base.length(d::OrthoNNDist) = length(d.x0)

Distributions.rand(d::OrthoNNDist) = rand(d.x1s)

Distributions.pdf(d::OrthoNNDist, x::Vector) = x in d.x1s ? d.prob : 0.0
Distributions.pdf(d::OrthoNNDist) = fill(d.prob, size(d.x1s))
Distributions.logpdf(d::OrthoNNDist, x::Vector) = log(pdf(d, x))


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    AbstractErrorEstimator

A method for estimating error of geostatistical solvers.
"""
abstract type AbstractErrorEstimator end

"""
    error(solver, problem, eestimator)

Estimate error of `solver` in a given `problem` with
`eestimator` error estimation method.
"""
function error end

# ----------------
# IMPLEMENTATIONS
# ----------------
include("errors/leave_ball_out.jl")
include("errors/cross_validation.jl")
include("errors/block_cross_validation.jl")
include("errors/weighted_cross_validation.jl")
include("errors/density_ratio_validation.jl")


text:
struct IBMDPSolver <: Solver
    ipft::IPFTSolver
    updater::Union{Function, Updater}
    ifm::InformationMeasure
    discount_information_gain::Bool
end

function POMDPs.solve(solver::IBMDPSolver, pomdp::POMDP)
    if isa(solver.updater, Function)
        updater = solver.updater(pomdp)
    else
        updater = solver.updater
    end
    belief_mdp = InformationRewardBeliefMDP(deepcopy(pomdp),
                                     updater.resample,
                                     solver.ifm,
                                     solver.discount_information_gain,
                                     updater.max_frac_replaced
                                    )
    return solve(solver.ipft, belief_mdp)
end

function POMDPs.solve(qs::QMDPSolver, bmdp::InformationRewardBeliefMDP)
    return solve(qs, bmdp.pomdp)
end


text:
using Stheno
using Random

export to_ColVecs,
    unpack_gp,
    unpack_gpar,
    get_time_mask,
    get_output_mask,
    parse_initial_gp_params,
    parse_initial_gpar_params
"""
Transform the input space into a ColVecs for using GPAR
"""
function to_ColVecs(inputs)
    # Inputs is an array of 1-D vectors
    concatted = hcat(inputs...)
    # Transpose so that we work with columns
    tranposed = transpose(concatted)
    return ColVecs(tranposed)
end

"""
Functions to unpack positive hyperparameters. Used for optimization purposes.
"""
function unpack_gp(params)
    # Unpack the parameters for a GP with EQ kernel
    l = exp(params[1]) + 1e-3
    process_var = exp(params[2]) + 1e-3
    noise_sigma = exp(params[3]) + 1e-3

    return l, process_var, noise_sigma
end

function unpack_gpar(params)
    # Unpack the parameters for a GPAR with EQ kernel
    time_l = exp(params[1]) + 1e-3
    time_var = exp(params[2]) + 1e-3
    out_l = exp(params[3]) + 1e-3
    out_var = exp(params[4]) + 1e-3

    noise_sigma = exp(params[5]) + 1e-3

    return time_l, time_var, out_l, out_var, noise_sigma
end

"""
The masks are meant to be used toghether when constructing a GPAR kernel so
that k((t, y), (t', y')) = k̲₁(t, t') + k₂(y, y')

Example:

xs = to_ColVecs(
[[-0.08, 0.39, 1.31],
 [-1.60, 0.58, -2.34],
 [-1.71, -0.16, 1.26]])
xs_no_first = to_ColVecs(
[[-1.60, 0.58, -2.34],
[-1.71, -0.16, 1.26]])
xs_only_first = to_ColVecs([[-0.08, 0.39, 1.31]])

ys = to_ColVecs(
[[1.82, 1.06, 0.54],
 [0.14, 1.02, -0.38],
 [-1.45, -0.007, 0.006]])
ys_no_first = to_ColVecs(
[[0.14, 1.02, -0.38],
[-1.45, -0.007, 0.006]])
ys_only_first = to_ColVecs([[1.82, 1.06, 0.54]])
vanilla_kernel = EQ()

time_mask = get_time_mask(3)  # 3 is the number of vals in an input vector
time_kernel = stretch(EQ(), time_mask)

# These two should be the same
time_K = pairwise(time_kernel, xs, ys)
time_true_K = pairwise(vanilla_kernel, xs_only_first, ys_only_first)
# Compare the two results

out_mask = get_output_mask(3)
out_kernel = stretch(EQ(), out_mask)

# These two should be the same
out_K = pairwise(out_kernel, xs, ys)
out_true_K = pairwise(vanilla_kernel, xs_no_first, ys_no_first)
"""

"""
Return mask used to only select time features from inputs.
"""
function get_time_mask(input_length)
    time_mask = zeros(input_length)
    time_mask[1] = 1  # only select the first (time) element
    return time_mask
end

"""
Return mask used to only select previous output features from inputs.
"""
function get_output_mask(input_length)
    if input_length <= 1
        throw(DomainError(
            input_length,
            "Input length must be integer greater than 1",
        ))
    end
    out_mask = zeros(input_length - 1, input_length)
    for row = 1:(input_length-1)
        out_mask[row, row+1] = 1  # select this feature
    end
    return out_mask
end

"""
Helper function for doing parameter checking and creation.
"""
function _parse_param(i_log_param)
    if isnothing(i_log_param)
        i_log_param = rand(1)[1]
    end

    return i_log_param
end

"""
Functionality that takes in initial parameters (when defined), created them when
not defined (by sampling a normal), and accumulates them in a params array.
Used for GPs (3 params).
"""
function parse_initial_gp_params(i_log_l, i_log_process_var, i_log_noise_sigma)
    return [
        _parse_param(i_log_l),
        _parse_param(i_log_process_var),
        _parse_param(i_log_noise_sigma),
    ]
end

"""
Functionality that takes in initial parameters (when defined), creates them when
not defined (by sampling a normal), and accumulates them in a params array.
Used for GPARs (5 params).
"""
function parse_initial_gpar_params(
    i_log_time_l,
    i_log_time_var,
    i_log_out_l,
    i_log_out_var,
    i_log_noise_sigma,
)

    return [
        _parse_param(i_log_time_l),
        _parse_param(i_log_time_var),
        _parse_param(i_log_out_l),
        _parse_param(i_log_out_var),
        _parse_param(i_log_noise_sigma),
    ]
end

max_stars_repo_name: willtebbutt/GPAR-at-scale

max_issues_repo_name: willtebbutt/GPAR-at-scale

max_forks_repo_name: willtebbutt/GPAR-at-scale


text:
export ZabbixItem, get_items

struct ZabbixItem
  id::Int
  name::String
end
ZabbixItem(dict) =
  ZabbixItem(
    parse(Int, dict["itemid"]),
    dict["name"]
  )

showres(i::ZabbixItem; offset=0) =
  println(repeat(" ", offset), "ZabbixItem: ", i.id, " => ", i.name)

function get_items(z, host::ZabbixHost)
  items = ZabbixItem[]
  params = Dict("hostids"=>host.id)
  _items = @checkerr make_request(z, "item.get", params)
  for i in _items
    push!(items, ZabbixItem(i))
  end
  items
end


text:
# Kernel Scalar & Vector Operation  ========================================================

function kernel(κ::Kernel{T}, x::T, y::T) where {T<:AbstractFloat}
    kappa(κ, base_evaluate(basefunction(κ), x, y))
end

function kernel(
        κ::Kernel{T},
        x::AbstractArray{T},
        y::AbstractArray{T}
    ) where {T<:AbstractFloat}
    kappa(κ, base_evaluate(basefunction(κ), x, y))
end

# Kernel Matrix Calculation ================================================================

function kappamatrix!(κ::Kernel{T}, P::AbstractMatrix{T}) where {T<:AbstractFloat}
    for i in eachindex(P)
        @inbounds P[i] = kappa(κ, P[i])
    end
    P
end

function symmetric_kappamatrix!(
        κ::Kernel{T},
        P::AbstractMatrix{T},
        symmetrize::Bool
    ) where {T<:AbstractFloat}
    if !((n = size(P,1)) == size(P,2))
        throw(DimensionMismatch("Pairwise matrix must be square."))
    end
    for j = 1:n, i = (1:j)
        @inbounds P[i,j] = kappa(κ, P[i,j])
    end
    symmetrize ? LinearAlgebra.copytri!(P, 'U') : P
end

"""
    kernelmatrix!(P::Matrix, σ::Orientation, κ::Kernel, X::Matrix, symmetrize::Bool)

In-place version of `kernelmatrix` where pre-allocated matrix `K` will be overwritten
with the kernel matrix.
"""
function kernelmatrix!(
        σ::Orientation,
        P::Matrix{T},
        κ::Kernel{T},
        X::AbstractMatrix{T},
        symmetrize::Bool
    ) where {T<:AbstractFloat}
    basematrix!(σ, P, basefunction(κ), X, false)
    symmetric_kappamatrix!(κ, P, symmetrize)
end

"""
    kernelmatrix!(K::Matrix, σ::Orientation, κ::Kernel, X::Matrix, Y::Matrix)

In-place version of `kernelmatrix` where pre-allocated matrix `K` will be overwritten with
the kernel matrix.
"""
function kernelmatrix!(
        σ::Orientation,
        P::Matrix{T},
        κ::Kernel{T},
        X::AbstractMatrix{T},
        Y::AbstractMatrix{T}
    ) where {T<:AbstractFloat}
    basematrix!(σ, P, basefunction(κ), X, Y)
    kappamatrix!(κ, P)
end

function kernelmatrix(
        σ::Orientation,
        κ::Kernel{T},
        X::AbstractMatrix{T},
        symmetrize::Bool = true
    ) where {T<:AbstractFloat}
    symmetric_kappamatrix!(κ, basematrix(σ, basefunction(κ), X, false), symmetrize)
end

function kernelmatrix(
        σ::Orientation,
        κ::Kernel{T},
        X::AbstractMatrix{T},
        Y::AbstractMatrix{T}
    ) where {T<:AbstractFloat}
    kappamatrix!(κ, basematrix(σ, basefunction(κ), X, Y))
end

# Convenience Methods ======================================================================

"""
    kernel(κ::Kernel, x, y)

Apply the kernel `κ` to ``x`` and ``y`` where ``x`` and ``y`` are vectors or scalars of
some subtype of ``Real``.
"""
function kernel(κ::Kernel{T}, x::Real, y::Real) where {T}
    kernel(κ, T(x), T(y))
end

function kernel(
        κ::Kernel{T},
        x::AbstractArray{T1},
        y::AbstractArray{T2}
    ) where {T,T1,T2}
    kernel(κ, convert(AbstractArray{T}, x), convert(AbstractArray{T}, y))
end

"""
    kernelmatrix([σ::Orientation,] κ::Kernel, X::Matrix [, symmetrize::Bool])

Calculate the kernel matrix of `X` with respect to kernel `κ`.
"""
function kernelmatrix(
        σ::Orientation,
        κ::Kernel{T},
        X::AbstractMatrix{T1},
        symmetrize::Bool = true
    ) where {T,T1}
    U = convert(AbstractMatrix{T}, X)
    kernelmatrix(σ, κ, U, symmetrize)
end

function kernelmatrix(
        κ::Kernel,
        X::AbstractMatrix,
        symmetrize::Bool = true
    )
    kernelmatrix(Val(:row), κ, X, symmetrize)
end

"""
    kernelmatrix([σ::Orientation,] κ::Kernel, X::Matrix, Y::Matrix)

Calculate the base matrix of `X` and `Y` with respect to kernel `κ`.
"""
function kernelmatrix(
        σ::Orientation,
        κ::Kernel{T},
        X::AbstractMatrix{T1},
        Y::AbstractMatrix{T2}
    ) where {T,T1,T2}
    U = convert(AbstractMatrix{T}, X)
    V = convert(AbstractMatrix{T}, Y)
    kernelmatrix(σ, κ, U, V)
end

function kernelmatrix(
        κ::Kernel,
        X::AbstractMatrix,
        Y::AbstractMatrix
    )
    kernelmatrix(Val(:row), κ, X, Y)
end

# Kernel Centering =========================================================================

@doc raw"""
    centerkernelmatrix(K::Matrix)

Centers the (rectangular) kernel matrix `K` with respect to the implicit Kernel Hilbert
Space according to the following formula:

```math
[\mathbf{K}]_{ij}
= \langle\phi(\mathbf{x}_i) -\mathbf{\mu}_{\phi\mathbf{x}}, \phi(\mathbf{y}_j)
- \mathbf{\mu}_{\phi\mathbf{y}} \rangle
```
Where ``\mathbf{\mu}_{\phi\mathbf{x}}`` and ``\mathbf{\mu}_{\phi\mathbf{x}}`` are given by:

```math
\mathbf{\mu}_{\phi\mathbf{x}}
= \frac{1}{n} \sum_{i=1}^n \phi(\mathbf{x}_i)
\qquad \qquad
\mathbf{\mu}_{\phi\mathbf{y}}
= \frac{1}{m} \sum_{i=1}^m \phi(\mathbf{y}_i)
```
"""
function centerkernelmatrix!(K::Matrix{T}) where {T<:AbstractFloat}
    μx = Statistics.mean(K, dims = 2)
    μy = Statistics.mean(K, dims = 1)
    μ  = Statistics.mean(K)

    K .+= μ .- μx .- μy

    return K
end
centerkernelmatrix(K::Matrix) = centerkernelmatrix!(copy(K))


text:
__precompile__()

module RandomCorrelationMatrices

    import LinearAlgebra: cholesky, diag, norm
    import Distributions

    """
        randcormatrix(d, η)

    Generate a d-by-d random correlation matrix using the
    method described in:

        Lewandowski, Kurowicka, Joe
          "Generating random correlation matrices based on vines
            and extended onion method"
        Journal of Multivariate Analysis 100 (2009)
        doi:10.1016/j.jmva.2009.04.008
    """
    function randcormatrix(d, η)
        β = η + (d - 2) / 2
        u = rand(Distributions.Beta(β, β))
        r₁₂ = 2 * u - 1
        r = [   1 r₁₂ ;
              r₁₂   1 ]
        # In published paper this index is n = 2:d-1. n is never mentioned
        # again, but a mysterious k is, so we'll thus use k. A similar approach
        # was taken here: http://stats.stackexchange.com/a/125017/58921
        for k in 2:d - 1
            β = β - 1 / 2
            y = rand(Distributions.Beta(k / 2, β))
            u = randn(k)
            u /= norm(u)
            w = sqrt(y) * u
            A = cholesky(r).L
            z = A * w
            r = [ r  z ;
                  z' 1 ]
        end
        return r
    end

    """
        randcovmatrix(d, η, σ)

    Use randcormatrix with desired vector of standard deviations σ
    for each term to generate a random covariance matrix.
    """
    function randcovmatrix(d, η, σ)
        length(σ) != d && throw(DimensionMismatch("length(σ) doesn't match d"))
        r = randcormatrix(d, η)
        Σ = zeros(d, d)
        @inbounds for i in 1:d, j in 1:d
            Σ[i,j] = r[i,j] * σ[i] * σ[j]
        end
        return Σ
    end

end  # module


text:
using RecyclePermanent
using RNGPool
using JLD2
using BenchmarkTools

include("../models/bpModel.jl")
include("compare.jl")

@load "data/bpData.jld" ys θ
model = makeBPModel(θ, ys)

N = 40*model.n

trials = 1000

setRNGs(12345)

n = 10
N = 40*n

rpio = RPIO{model.Tx}(N, n, 1)

# simpleOutput = testEstimator(() -> simpleEstimate(model, rpio), trials)
biasedOutput = testEstimator(() -> biasedEstimateDiscrete(model, rpio), trials)
rpOutput = testEstimator(() -> rpEstimateDiscrete(model, rpio), trials)

ϵ = 1e-5
kmio = KMIO{model.Tx}(N, n)
kmOutput = testEstimator(() -> kmEstimate(model, kmio, ϵ), trials)

@btime rpEstimateDiscrete($model, $rpio)


text:
#=
module CompGraphs
=================
The CompGraph type in this module is intended to hold the computational graph of a
single composite function with vector-valued inputs and outputs. This graph
expresses the composite function as a recipe of elemental operations. Computational
graphs are required for implementing certain numerical methods, such as:

- the standard reverse/adjoint mode of automatic differentiation (AD)

- the "branch-locking" method for efficient reverse-AD-like generalized
  differentiation by Khan (2018)
  https://doi.org/10.1080/10556788.2017.1341506

- the "cone-squashing" method for generalized differentiation
  by Khan and Barton (2012, 2013)
  https://doi.org/10.1145/2491491.2491493

- the "reverse McCormick" convex relaxations incorporating constraint propagation,
  by Wechsung et al. (2015)
  https://doi.org/10.1007/s10898-015-0303-6

"load_function!" in this module uses operator overloading to construct the
computational graph of a finite composition of supported operations. Each
node and the graph overall can also hold additional user-specified data,
intended for use in methods like the reverse AD mode.

Written by Kamil Khan on February 10, 2022
=#
module CompGraphs

using Printf

export CompGraph, GraphNode

export load_function!, is_function_loaded

# structs; T and P are for smuggling in any sort of application-dependent data,
# but could easily be Any or Nothing for simplicity.
struct GraphNode{P}
    operation::Symbol
    parentIndices::Vector{Int}       # identify operands of "operation"
    constValue::Union{Float64, Int}  # only used when "operation" = :const or :^
    data::P                          # hold extra node-specific data
end

struct CompGraph{T, P}
    nodeList::Vector{GraphNode{P}}
    domainDim::Int              # domain dimension of graphed function
    rangeDim::Int               # range dimension of graphed function
    data::T                     # hold extra data not specific to any one node
end

struct GraphBuilder{P}
    index::Int
    graph::CompGraph
end

# constructors
GraphNode{P}(op::Symbol, i::Vector{Int}, p::P) where P = GraphNode{P}(op, i, 0.0, p)

# the following constructor requires a constructor T() with no arguments
function CompGraph{T, P}(n::Int, m::Int) where {T, P}
    return CompGraph{T, P}(GraphNode{P}[], n, m, T())
end

# A GraphNode.operation can be any Symbol from the following lists
unaryOpList = [:-, :inv, :exp, :log, :sin, :cos, :abs]
binaryOpList = [:+, :-, :*, :/, :^, :max, :min, :hypot]
customOpList = [:input, :output, :const]   # input, output, and Float64 constant

# print graph or individual nodes
opStringDict = Dict(
    :- => "neg",
    :inv => "inv",
    :exp => "exp",
    :log => "log",
    :sin => "sin",
    :cos => "cos",
    :abs => "abs",
    :+ => " + ",
    :- => " - ",
    :* => " * ",
    :/ => " / ",
    :^ => " ^ ",
    :max => "max",
    :min => "min",
    :hypot => "hyp",
    :input => "inp",
    :output => "out",
    :const => "con",
)

function Base.show(io::IO, node::GraphNode)
    parents = node.parentIndices
    nParents = length(parents)

    if (node.operation == :^) && (length(parents) == 1)
        opString = @sprintf " ^%1d" node.constValue
    else
        opString = opStringDict[node.operation]
    end

    if nParents <= 2
        oneParent(i::Int) = (nParents < i) ? "   " : @sprintf "%-3d" parents[i]
        parentString = oneParent(1) * "  " * oneParent(2)
    else
        parentString = string(parents)
    end

    if node.operation == :const
        dataString = @sprintf "const: % .3e" node.constValue
    else
        dataString = string(node.data)
    end

    return print(io, opString, " | ", parentString, " | ", dataString)
end

function Base.show(io::IO, graph::CompGraph)
    return begin
        println(" index | op  | parents  | data")
        println(" ------------------------------")
        for (i, node) in enumerate(graph.nodeList)
            @printf "   %3d | " i
            println(node)
        end
    end
end

# load in a function using operator overloading, and store its computational graph
function load_function!(
    f::Function,
    graph::CompGraph{T, P},
    initP::P
) where {T, P}

    empty!(graph.nodeList)

    # push new nodes for function inputs
    xGB = [GraphBuilder{P}(i, graph) for i=1:(graph.domainDim)]
    for xComp in xGB
        inputData = deepcopy(initP)
        inputNode = GraphNode{P}(:input, Int[], 0.0, inputData)
        push!(graph.nodeList, inputNode)
    end

    # push new nodes for all intermediate operations, using operator overloading
    yGB = f(xGB)
    if !(yGB isa Vector)
        yGB = [yGB]
    end

    # push new nodes for function outputs
    for yComp in yGB
        outputData = deepcopy(initP)
        outputNode = GraphNode{P}(:output, [yComp.index], 0.0, outputData)
        push!(graph.nodeList, outputNode)
    end
end

is_function_loaded(graph::CompGraph) = !isempty(graph.nodeList)

## let GraphBuilder construct a nodeList by operator overloading

# overload unary operations in unaryOpList
macro define_GraphBuilder_unary_op_rule(op)
    opEval = eval(op)
    return quote
        function Base.$opEval(u::GraphBuilder{P}) where P
            parentGraph = u.graph
            prevNodes = parentGraph.nodeList

            newNodeData = deepcopy(prevNodes[u.index].data)
            newNode = GraphNode{P}($op, [u.index], 0.0, newNodeData)
            push!(prevNodes, newNode)

            return GraphBuilder{P}(length(prevNodes), parentGraph)
        end
    end
end

for i in 1:length(unaryOpList)
    @eval @define_GraphBuilder_unary_op_rule $unaryOpList[$i]
end

# overload binary operations in binaryOpList.
# Uses nontrivial "promote" rules to push Float64 constants to the parent graph
function Base.promote(uA::GraphBuilder{P}, uB::Float64) where P
    parentGraph = uA.graph
    prevNodes = parentGraph.nodeList

    newNodeData = deepcopy(prevNodes[uA.index].data)
    newNode = GraphNode{P}(:const, Int[], uB, newNodeData)
    push!(prevNodes, newNode)

    return (uA, GraphBuilder{P}(length(prevNodes), parentGraph))
end
function Base.promote(uA::Float64, uB::GraphBuilder{P}) where P
    return reverse(promote(uB, uA))
end
Base.promote(uA::GraphBuilder{P}, uB::GraphBuilder{P}) where P = (uA, uB)

macro define_GraphBuilder_binary_op_rule(op)
    opEval = eval(op)
    return quote
        function Base.$opEval(uA::GraphBuilder{P}, uB::GraphBuilder{P}) where P
            parentGraph = uA.graph
            prevNodes = parentGraph.nodeList

            newNodeData = deepcopy(prevNodes[uA.index].data)
            newNode = GraphNode{P}($op, [uA.index, uB.index], 0.0, newNodeData)
            push!(prevNodes, newNode)

            return GraphBuilder{P}(length(prevNodes), parentGraph)
        end

        function Base.$opEval(uA::GraphBuilder{P}, uB::Float64) where P
            return $opEval(promote(uA, uB)...)
        end

        function Base.$opEval(uA::Float64, uB::GraphBuilder{P}) where P
            return $opEval(promote(uA, uB)...)
        end
    end
end

for i in 1:length(binaryOpList)
    @eval @define_GraphBuilder_binary_op_rule $binaryOpList[$i]
end

function Base.:^(uA::GraphBuilder{P}, uB::Int) where P
    parentGraph = uA.graph
    prevNodes = parentGraph.nodeList

    newNodeData = deepcopy(prevNodes[uA.index].data)
    newNode = GraphNode{P}(:^, [uA.index], uB, newNodeData)
    push!(prevNodes, newNode)

    return GraphBuilder{P}(length(prevNodes), parentGraph)
end

end # module

max_stars_repo_name: kamilkhanlab/computational-graph-tools

max_issues_repo_name: kamilkhanlab/computational-graph-tools

max_forks_repo_name: kamilkhanlab/computational-graph-tools