text:
export sum2, jacobian, ε⁻¹_bar!, ε⁻¹_bar, ∂ω²∂k_adj, Mₖᵀ_plus_Mₖ, ∂²ω²∂k², herm
export ∇ₖmag_m_n, ∇HMₖH, ∇M̂, ∇solve_k, ∇solve_k!, solve_adj!, neff_ng_gvd, ∂ε⁻¹_∂ω, ∂nng⁻¹_∂ω
export ∇ₖmag_mn
### ForwardDiff Comoplex number support
# ref: https://github.com/JuliaLang/julia/pull/36030
# https://github.com/JuliaDiff/ForwardDiff.jl/pull/455
# Base.float(d::ForwardDiff.Dual{T}) where T = ForwardDiff.Dual{T}(float(d.value), d.partials)
# Base.prevfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = ForwardDiff.Dual{T}(prevfloat(float(d.value)), d.partials)
# Base.nextfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = ForwardDiff.Dual{T}(nextfloat(float(d.value)), d.partials)
# function Base.ldexp(x::T, e::Integer) where T<:ForwardDiff.Dual
#     if e >=0
#         x * (1<<e)
#     else
#         x / (1<<-e)
#     end
# end

### ForwardDiff FFT support
# ref: https://github.com/JuliaDiff/ForwardDiff.jl/pull/495/files
# https://discourse.julialang.org/t/forwarddiff-and-zygote-cannot-automatically-differentiate-ad-function-from-c-n-to-r-that-uses-fft/52440/18
ForwardDiff.value(x::Complex{<:ForwardDiff.Dual}) =
    Complex(x.re.value, x.im.value)

ForwardDiff.partials(x::Complex{<:ForwardDiff.Dual}, n::Int) =
    Complex(ForwardDiff.partials(x.re, n), ForwardDiff.partials(x.im, n))

ForwardDiff.npartials(x::Complex{<:ForwardDiff.Dual{T,V,N}}) where {T,V,N} = N
ForwardDiff.npartials(::Type{<:Complex{<:ForwardDiff.Dual{T,V,N}}}) where {T,V,N} = N

# AbstractFFTs.complexfloat(x::AbstractArray{<:ForwardDiff.Dual}) = float.(x .+ 0im)
AbstractFFTs.complexfloat(x::AbstractArray{<:ForwardDiff.Dual}) = AbstractFFTs.complexfloat.(x)
AbstractFFTs.complexfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = convert(ForwardDiff.Dual{T,float(V),N}, d) + 0im

AbstractFFTs.realfloat(x::AbstractArray{<:ForwardDiff.Dual}) = AbstractFFTs.realfloat.(x)
AbstractFFTs.realfloat(d::ForwardDiff.Dual{T,V,N}) where {T,V,N} = convert(ForwardDiff.Dual{T,float(V),N}, d)

for plan in [:plan_fft, :plan_ifft, :plan_bfft]
    @eval begin

        AbstractFFTs.$plan(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x)) =
            AbstractFFTs.$plan(ForwardDiff.value.(x) .+ 0im, region)

        AbstractFFTs.$plan(x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, region=1:ndims(x)) =
            AbstractFFTs.$plan(ForwardDiff.value.(x), region)

    end
end

# rfft only accepts real arrays
AbstractFFTs.plan_rfft(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x)) =
    AbstractFFTs.plan_rfft(ForwardDiff.value.(x), region)

for plan in [:plan_irfft, :plan_brfft]  # these take an extra argument, only when complex?
    @eval begin

        AbstractFFTs.$plan(x::AbstractArray{<:ForwardDiff.Dual}, region=1:ndims(x)) =
            AbstractFFTs.$plan(ForwardDiff.value.(x) .+ 0im, region)

        AbstractFFTs.$plan(x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}, d::Integer, region=1:ndims(x)) =
            AbstractFFTs.$plan(ForwardDiff.value.(x), d, region)

    end
end

for P in [:Plan, :ScaledPlan]  # need ScaledPlan to avoid ambiguities
    @eval begin

        Base.:*(p::AbstractFFTs.$P, x::AbstractArray{<:ForwardDiff.Dual}) =
            _apply_plan(p, x)

        Base.:*(p::AbstractFFTs.$P, x::AbstractArray{<:Complex{<:ForwardDiff.Dual}}) =
            _apply_plan(p, x)

    end
end

function _apply_plan(p::AbstractFFTs.Plan, x::AbstractArray)
    xtil = p * ForwardDiff.value.(x)
    dxtils = ntuple(ForwardDiff.npartials(eltype(x))) do n
        p * ForwardDiff.partials.(x, n)
    end
    map(xtil, dxtils...) do val, parts...
        Complex(
            ForwardDiff.Dual(real(val), map(real, parts)),
            ForwardDiff.Dual(imag(val), map(imag, parts)),
        )
    end
end

# used with the ForwardDiff+FFTW code above, this Zygote.extract method
# enables Zygote.hessian to work on real->real functions that internally use
# FFTs (and thus complex numbers)
import Zygote: extract
function Zygote.extract(xs::AbstractArray{<:Complex{<:ForwardDiff.Dual{T,V,N}}}) where {T,V,N}
  J = similar(xs, complex(V), N, length(xs))
  for i = 1:length(xs), j = 1:N
    J[j, i] = xs[i].re.partials.values[j] + im * xs[i].im.partials.values[j]
  end
  x0 = ForwardDiff.value.(xs)
  return x0, J
end

####
# Example code for defining custom ForwardDiff rules, copied from YingboMa's gist:
# https://gist.github.com/YingboMa/c22dcf8239a62e01b27ac679dfe5d4c5
# using ForwardDiff
# goo((x, y, z),) = [x^2*z, x*y*z, abs(z)-y]
# foo((x, y, z),) = [x^2*z, x*y*z, abs(z)-y]
# function foo(u::Vector{ForwardDiff.Dual{T,V,P}}) where {T,V,P}
#     # unpack: AoS -> SoA
#     vs = ForwardDiff.value.(u)
#     # you can play with the dimension here, sometimes it makes sense to transpose
#     ps = mapreduce(ForwardDiff.partials, hcat, u)
#     # get f(vs)
#     val = foo(vs)
#     # get J(f, vs) * ps (cheating). Write your custom rule here
#     jvp = ForwardDiff.jacobian(goo, vs) * ps
#     # pack: SoA -> AoS
#     return map(val, eachrow(jvp)) do v, p
#         ForwardDiff.Dual{T}(v, p...) # T is the tag
#     end
# end
# ForwardDiff.gradient(u->sum(cumsum(foo(u))), [1, 2, 3]) == ForwardDiff.gradient(u->sum(cumsum(goo(u))), [1, 2, 3])
####

# AD rules for StaticArrays Constructors
rrule(T::Type{<:SArray}, xs::Number...) = ( T(xs...), dv -> (NoTangent(), dv...) )
rrule(T::Type{<:SArray}, x::AbstractArray) = ( T(x), dv -> (NoTangent(), dv) )
rrule(T::Type{<:SMatrix}, xs::Number...) = ( T(xs...), dv -> (NoTangent(), dv...) )
rrule(T::Type{<:SMatrix}, x::AbstractMatrix) = ( T(x), dv -> (NoTangent(), dv) )
rrule(T::Type{<:SVector}, xs::Number...) = ( T(xs...), dv -> (NoTangent(), dv...) )
rrule(T::Type{<:SVector}, x::AbstractVector) = ( T(x), dv -> (NoTangent(), dv) )
rrule(T::Type{<:HybridArray}, x::AbstractArray) = ( T(x), dv -> (NoTangent(), dv) )

# AD rules for reinterpreting back and forth between N-D arrays of SVectors and (N+1)-D arrays
function rrule(::typeof(reinterpret),reshape,type::Type{T1},A::AbstractArray{SVector{N1,T2},N2}) where {T1,N1,T2,N2}
	# return ( reinterpret(reshape,T1,A), Δ->( NoTangent(), ZeroTangent(), ZeroTangent(), reinterpret( reshape,SVector{N1,T1}, Δ ) ) )
	function reinterpret_reshape_SV_pullback(Δ)
		return (NoTangent(), ZeroTangent(), ZeroTangent(), reinterpret(reshape,SVector{N1,eltype(Δ)},Δ))
	end
	( reinterpret(reshape,T1,A), reinterpret_reshape_SV_pullback )
end
function rrule(::typeof(reinterpret),reshape,type::Type{<:SVector{N1,T1}},A::AbstractArray{T1}) where {T1,N1}
	return ( reinterpret(reshape,type,A), Δ->( NoTangent(), ZeroTangent(), ZeroTangent(), reinterpret( reshape, eltype(A), Δ ) ) )
end

# need adjoint for constructor:
# Base.ReinterpretArray{Float64, 2, SVector{3, Float64}, Matrix{SVector{3, Float64}}, false}. Gradient is of type FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}
import Base: ReinterpretArray
function rrule(T::Type{ReinterpretArray{T1, N1, SVector{N2, T2}, Array{SVector{N3, T3},N4}, IsReshaped}}, x::AbstractArray)  where {T1, N1, N2, T2, N3, T3, N4, IsReshaped}
	function ReinterpretArray_SV_pullback(Δ)
		if IsReshaped
			Δr = reinterpret(reshape,SVector{N2,eltype(Δ)},Δ)
		else
			Δr = reinterpret(SVector{N2,eltype(Δ)},Δ)
		end
		return (NoTangent(), Δr)
	end
	( T(x), ReinterpretArray_SV_pullback )
end

# AD rules for reinterpreting back and forth between N-D arrays of SMatrices and (N+2)-D arrays
function rrule(::typeof(reinterpret),reshape,type::Type{T1},A::AbstractArray{SMatrix{N1,N2,T2,N3},N4}) where {T1<:Real,T2,N1,N2,N3,N4}
	# @show A
	# @show eltype(A)
	# @show type
	# @show size(reinterpret(reshape,T1,A))
	# @show N1*N2
	# function f_pb(Δ)
	# 	@show eltype(Δ)
	# 	@show size(Δ)
	# 	# @show Δ
	# 	@show typeof(Δ)
	# 	return ( NoTangent(), ZeroTangent(), ZeroTangent(), reinterpret( reshape,SMatrix{N1,N2,T1,N3}, Δ ) )
	# end
	# return ( reinterpret(reshape,T1,A), Δ->f_pb(Δ) )
	return ( reinterpret(reshape,T1,A), Δ->( NoTangent(), ZeroTangent(), ZeroTangent(), reinterpret( reshape,SMatrix{N1,N2,T1,N3}, real(Δ) ) ) )
end

function rrule(::typeof(reinterpret),reshape,type::Type{<:SMatrix{N1,N2,T1,N3}},A::AbstractArray{T1}) where {T1,T2,N1,N2,N3}
	# @show type
	# @show eltype(A)
	return ( reinterpret(reshape,type,A), Δ->( NoTangent(), ZeroTangent(), ZeroTangent(), reinterpret( reshape, eltype(A), Δ ) ) )
end

# adjoint for constructor Base.ReinterpretArray{SMatrix{3, 3, Float64, 9}, 1, Float64, Vector{Float64}, false}.
# Gradient is of type FillArrays.Fill{FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, 1, Tuple{Base.OneTo{Int64}}}

function rrule(::typeof(reinterpret), ::typeof(reshape), ::Type{R}, A::AbstractArray{T}) where {N1, N2, T, R <: SMatrix{N1,N2,T}}
    function pullback(Ā)
        ∂A = mapreduce(v -> v isa R ? v : zero(R), vcat, Ā; init = similar(A, 0))
        return (NoTangent(), DoesNotExist(), DoesNotExist(), reshape(∂A, size(A)))
    end
    return (reinterpret(reshape, R, A), pullback)
end

# function rrule(T::Type{::R3, x::R2) where {T1, N1, N2, T2, R1<:SMatrix{N1,N2,T1}, R2<:AbstractArray{T2}, R3<:ReinterpretArray{R1}}}
# 	function ReinterpretArray_SM_pullback(Δ)
# 		Δr = reshape(reinterpret(T1,collect(Δ)),size(x))
# 		# if IsReshaped
# 		# 	Δr = reshape(reinterpret(T4,collect(Δ)),size(x))
# 		# else
# 		# 	Δr = reshape(reinterpret(T4,collect(Δ)),size(x))
# 		# end
# 		return (NoTangent(), Δr)
# 	end
# 	( T(x), ReinterpretArray_SM_pullback )
# end

# AD rules for fast norms of types SVector{2,T} and SVector{2,3}

function _norm2_back_SV2r(x::SVector{2,T}, y, Δy) where T<:Real
    ∂x = Vector{T}(undef,2)
    ∂x .= x .* (real(Δy) * pinv(y))
    return reinterpret(SVector{2,T},∂x)[1]
end

function _norm2_back_SV3r(x::SVector{3,T}, y, Δy) where T<:Real
    ∂x = Vector{T}(undef,3)
    ∂x .= x .* (real(Δy) * pinv(y))
    return reinterpret(SVector{3,T},∂x)[1]
end

function _norm2_back_SV2r(x::SVector{2,T}, y, Δy) where T<:Complex
    ∂x = Vector{T}(undef,2)
    ∂x .= conj.(x) .* (real(Δy) * pinv(y))
    return reinterpret(SVector{2,T},∂x)[1]
end

function _norm2_back_SV3r(x::SVector{3,T}, y, Δy) where T<:Complex
    ∂x = Vector{T}(undef,3)
    ∂x .= conj.(x) .* (real(Δy) * pinv(y))
    return reinterpret(SVector{3,T},∂x)[1]
end

function rrule(::typeof(norm), x::SVector{3,T}) where T<:Real
	y = LinearAlgebra.norm(x)
	function norm_pb(Δy)
		∂x = Thunk() do
			_norm2_back_SV3r(x, y, Δy)
		end
		return ( NoTangent(), ∂x )
	end
	norm_pb(::ZeroTangent) = (NoTangent(), ZeroTangent())
    return y, norm_pb
end

function rrule(::typeof(norm), x::SVector{2,T}) where T<:Real
	y = LinearAlgebra.norm(x)
	function norm_pb(Δy)
		∂x = Thunk() do
			_norm2_back_SV2r(x, y, Δy)
		end
		return ( NoTangent(), ∂x )
	end
	norm_pb(::ZeroTangent) = (NoTangent(), ZeroTangent())
    return y, norm_pb
end

function rrule(::typeof(norm), x::SVector{3,T}) where T<:Complex
	y = LinearAlgebra.norm(x)
	function norm_pb(Δy)
		∂x = Thunk() do
			_norm2_back_SV3c(x, y, Δy)
		end
		return ( NoTangent(), ∂x )
	end
	norm_pb(::ZeroTangent) = (NoTangent(), ZeroTangent())
    return y, norm_pb
end

function rrule(::typeof(norm), x::SVector{2,T}) where T<:Complex
	y = LinearAlgebra.norm(x)
	function norm_pb(Δy)
		∂x = Thunk() do
			_norm2_back_SV2c(x, y, Δy)
		end
		return ( NoTangent(), ∂x )
	end
	norm_pb(::ZeroTangent) = (NoTangent(), ZeroTangent())
    return y, norm_pb
end

@non_differentiable KDTree(::Any)
@non_differentiable g⃗(::Any)
@non_differentiable _fftaxes(::Any)
# Examples of how to assert type stability for broadcasting custom types (see https://github.com/FluxML/Zygote.jl/issues/318 )
# Base.similar(bc::Base.Broadcast.Broadcasted{Base.Broadcast.ArrayStyle{V}}, ::Type{T}) where {T<:Real, V<:Real3Vector} = Real3Vector(Vector{T}(undef,3))
# Base.similar(bc::Base.Broadcast.Broadcasted{Base.Broadcast.ArrayStyle{V}}, ::Type{T}) where {T, V<:Real3Vector} = Array{T}(undef, size(bc))

@adjoint enumerate(xs) = enumerate(xs), diys -> (map(last, diys),)
_ndims(::Base.HasShape{d}) where {d} = d
_ndims(x) = Base.IteratorSize(x) isa Base.HasShape ? _ndims(Base.IteratorSize(x)) : 1
@adjoint function Iterators.product(xs...)
                    d = 1
                    Iterators.product(xs...), dy -> ntuple(length(xs)) do n
                        nd = _ndims(xs[n])
                        dims = ntuple(i -> i<d ? i : i+nd, ndims(dy)-nd)
                        d += nd
                        func = sum(y->y[n], dy; dims=dims)
                        ax = axes(xs[n])
                        reshape(func, ax)
                    end
                end

function sum2(op,arr)
    return sum(op,arr)
end

function sum2adj( Δ, op, arr )
    n = length(arr)
    g = x->Δ*Zygote.gradient(op,x)[1]
    return ( nothing, map(g,arr))
end

@adjoint function sum2(op,arr)
    return sum2(op,arr),Δ->sum2adj(Δ,op,arr)
end

# now-removed Zygote trick to improve stability of `norm` pullback
# found referenced here: https://github.com/JuliaDiff/ChainRules.jl/issues/338
function Zygote._pullback(cx::Zygote.AContext, ::typeof(norm), x::AbstractArray, p::Real = 2)
  fallback = (x, p) -> sum(abs.(x).^p .+ eps(0f0)) ^ (one(eltype(x)) / p) # avoid d(sqrt(x))/dx == Inf at 0
  Zygote._pullback(cx, fallback, x, p)
end

# """
# jacobian(f,x) : stolen from https://github.com/FluxML/Zygote.jl/pull/747/files
#
# Construct the Jacobian of `f` where `x` is a real-valued array
# and `f(x)` is also a real-valued array.
# """
# function jacobian(f,x)
#     y,back  = Zygote.pullback(f,x)
#     k  = length(y)
#     n  = length(x)
#     J  = Matrix{eltype(y)}(undef,k,n)
#     e_i = fill!(similar(y), 0)
#     @inbounds for i = 1:k
#         e_i[i] = oneunit(eltype(x))
#         J[i,:] = back(e_i)[1]
#         e_i[i] = zero(eltype(x))
#     end
#     (J,)
# end

### Zygote StructArrays rules from https://github.com/cossio/ZygoteStructArrays.jl
# @adjoint function (::Type{SA})(t::Tuple) where {SA<:StructArray}
#     sa = SA(t)
#     back(Δ::NamedTuple) = (values(Δ),)
#     function back(Δ::AbstractArray{<:NamedTuple})
#         nt = (; (p => [getproperty(dx, p) for dx in Δ] for p in propertynames(sa))...)
#         return back(nt)
#     end
#     return sa, back
# end
#
# @adjoint function (::Type{SA})(t::NamedTuple) where {SA<:StructArray}
#     sa = SA(t)
#     back(Δ::NamedTuple) = (NamedTuple{propertynames(sa)}(Δ),)
#     function back(Δ::AbstractArray)
#         back((; (p => [getproperty(dx, p) for dx in Δ] for p in propertynames(sa))...))
#     end
#     return sa, back
# end
#
# @adjoint function (::Type{SA})(a::A) where {T,SA<:StructArray,A<:AbstractArray{T}}
#     sa = SA(a)
#     function back(Δsa)
#         Δa = [(; (p => Δsa[p][i] for p in propertynames(Δsa))...) for i in eachindex(a)]
#         return (Δa,)
#     end
#     return sa, back
# end
#
# # Must special-case for Complex (#1)
# @adjoint function (::Type{SA})(a::A) where {T<:Complex,SA<:StructArray,A<:AbstractArray{T}}
#     sa = SA(a)
#     function back(Δsa) # dsa -> da
#         Δa = [Complex(Δsa.re[i], Δsa.im[i]) for i in eachindex(a)]
#         (Δa,)
#     end
#     return sa, back
# end
#
# @adjoint function literal_getproperty(sa::StructArray, ::Val{key}) where {key}
#     key::Symbol
#     result = getproperty(sa, key)
#     function back(Δ::AbstractArray)
#         nt = (; (k => zero(v) for (k,v) in pairs(fieldarrays(sa)))...)
#         return (Base.setindex(nt, Δ, key), nothing)
#     end
#     return result, back
# end
#
# @adjoint Base.getindex(sa::StructArray, i...) = sa[i...], Δ -> ∇getindex(sa,i,Δ)
# @adjoint Base.view(sa::StructArray, i...) = view(sa, i...), Δ -> ∇getindex(sa,i,Δ)
# function ∇getindex(sa::StructArray, i, Δ::NamedTuple)
#     dsa = (; (k => ∇getindex(v,i,Δ[k]) for (k,v) in pairs(fieldarrays(sa)))...)
#     di = map(_ -> nothing, i)
#     return (dsa, map(_ -> nothing, i)...)
# end
# # based on
# # https://github.com/FluxML/Zygote.jl/blob/64c02dccc698292c548c334a15ce2100a11403e2/src/lib/array.jl#L41
# ∇getindex(a::AbstractArray, i, Δ::Nothing) = nothing
# function ∇getindex(a::AbstractArray, i, Δ)
#     if i isa NTuple{<:Any, Integer}
#         da = Zygote._zero(a, typeof(Δ))
#         da[i...] = Δ
#     else
#         da = Zygote._zero(a, eltype(Δ))
#         dav = view(da, i...)
#         dav .= Zygote.accum.(dav, Zygote._droplike(Δ, dav))
#     end
#     return da
# end
#
# @adjoint function (::Type{NT})(t::Tuple) where {K,NT<:NamedTuple{K}}
#     nt = NT(t)
#     back(Δ::NamedTuple) = (values(NT(Δ)),)
#     return nt, back
# end

# # https://github.com/FluxML/Zygote.jl/issues/680
# @adjoint function (T::Type{<:Complex})(re, im)
# 	back(Δ::Complex) = (nothing, real(Δ), imag(Δ))
# 	back(Δ::NamedTuple) = (nothing, Δ.re, Δ.im)
# 	T(re, im), back
# end

#### AD Rules for Iterative eigensolves of Helmholtz Operator

"""
Inversion/Conjugate-transposition equalities of Maxwell operator components
----------------------------------------------------------------------------

FFT operators:
--------------
	If 𝓕⁻¹ === `bfft` (symmetric, unnormalized case)

	(1a)	(𝓕)' 	= 	𝓕⁻¹

	(2a)	(𝓕⁻¹)' = 	𝓕

	If 𝓕⁻¹ === `ifft` (asymmetric, normalized case)

	(1a)	(𝓕)' 	= 	𝓕⁻¹ * N		( N := Nx * Ny * Nz )

	(2a)	(𝓕⁻¹)' = 	𝓕	 / N

Combined Curl+Basis-Change operators:
--------------------------------------
	(3)	( [(k⃗+g⃗)×]cₜ )' 	= 	-[(k⃗+g⃗)×]ₜc

	(4)	( [(k⃗+g⃗)×]ₜc )' 	= 	-[(k⃗+g⃗)×]cₜ

Combined Cross+Basis-Change operators:
--------------------------------------
	(3)	( [(ẑ)×]cₜ )' 	= 	[(ẑ)×]ₜc

	(4)	( [(ẑ)×]ₜc )' 	= 	[(ẑ)×]cₜ

--------------
"""

"""
Calculate k̄ contribution from M̄ₖ, where M̄ₖ is backpropagated from ⟨H|M̂ₖ|H⟩

Consider Mₖ as composed of three parts:

	(1) Mₖ	= 	[(k⃗+g⃗)×]cₜ  ⋅  [ 𝓕  nn̂g⁻¹ 𝓕⁻¹ ]   ⋅  [ẑ×]ₜc
				------------	-----------------	  -------
  			= 		 A				    B                C

where the "cₜ" and "ₜc" labels on the first and third components denote
Cartesian-to-Transverse and Transverse-to-Cartesian coordinate transformations,
respectively.

From Giles, we know that if D = A B C, then

	(2)	Ā 	=	D̄ Cᵀ Bᵀ

	(3)	B̄ 	=	Aᵀ D̄ Cᵀ

	(4) C̄	=	Bᵀ Aᵀ D̄		(to the bone)

We also know that M̄ₖ corresponding to the gradient of ⟨H|M̂ₖ|H⟩ will be

	(5) M̄ₖ	=	|H*⟩⟨H|

where `*` denote complex conjugation.
Equations (2)-(5) give us formulae for the gradients back-propagated to the three
parameterized operators composing Mₖ

	(6) 	[k+g ×]̄		 = 	 |H⟩⟨H|	 ⋅  [[ẑ×]ₜc]ᵀ  ⋅  [ 𝓕  nn̂g⁻¹ 𝓕⁻¹ ]ᵀ

							= 	|H⟩⟨ ( [ 𝓕  nn̂g⁻¹ 𝓕⁻¹ ]  ⋅  [ẑ×]ₜc ⋅ H ) |

	(7)	[ 𝓕 nn̂g⁻¹ 𝓕⁻¹ ]̄	   = 	 [[k+g ×]cₜ]ᵀ  ⋅  |H⟩⟨H| ⋅  [[ẑ×]ₜc]ᵀ

						 	= 	-[k+g ×]ₜc  ⋅  |H⟩⟨H| ⋅  [ẑ×]cₜ

							=	-| [k+g ×]ₜc  ⋅ H ⟩⟨ [ẑ×]ₜc ⋅ H |

	(8)  ⇒ 	[ nn̂g⁻¹ ]̄	 	  =   -| 𝓕 ⋅ [k+g ×]ₜc  ⋅ H ⟩⟨ 𝓕 ⋅ [ẑ×]ₜc ⋅ H |

	(9)			[ẑ ×]̄	 	  =   [ 𝓕  nn̂g⁻¹ 𝓕⁻¹ ]ᵀ ⋅ [[k+g×]cₜ]ᵀ ⋅ |H⟩⟨H|

							= 	-| [ 𝓕  nn̂g⁻¹ 𝓕⁻¹ ]  ⋅  [k+g×]cₜ ⋅ H ⟩⟨H|

where `[ẑ ×]` operators are still parameterized by k⃗ because they involve
m⃗ & n⃗ orthonormal polarization basis vectors, which are determined by k⃗+g⃗
and thus k⃗-dependent.

Our [(k⃗+g⃗)×]cₜ and [ẑ ×]ₜc operators act locally in reciprocal space with the
following local structures

	(10) [ ( k⃗+g⃗[ix,iy,iz] ) × ]ₜc  =		[	-n⃗₁	m⃗₁
											  -n⃗₂	  m⃗₂
											  -n⃗₃	  m⃗₃	]

									=	  [	   m⃗     n⃗  	]  ⋅  [  0   -1
																	1 	 0	]

	(11) [ ( k⃗+g⃗[ix,iy,iz] ) × ]cₜ  =			 [	   n⃗₁	  n⃗₂	  n⃗₃
										  			-m⃗₁   -m⃗₂	  -m⃗₃	  ]

									=	[  0   -1 		⋅	[   m⃗ᵀ
										   1 	0  ]			n⃗ᵀ		]

										=	-(	[ ( k⃗+g⃗[ix,iy,iz] ) × ]ₜc	)ᵀ

	(12) [ ẑ × ]ₜc	 	 =	[	 -m⃗₂	-n⃗₂
								 m⃗₁	n⃗₁
								 0	    0	 ]

						=	[	0 	-1	  0		 		[	m⃗₁	 n⃗₁
								1 	 0	  0	 		⋅		m⃗₂	 n⃗₂
								0 	 0	  0	  ]				m⃗₃	 n⃗₃	]

	(13) [ ẑ × ]cₜ	 	 =	   [	 -m⃗₂	m⃗₁	 	0
									-n⃗₂	n⃗₁		0		]

						=	  [   m⃗ᵀ				[	0	 1	 0
								  n⃗ᵀ	]		⋅		-1	  0	  0
								  						0	 0	 0	]

						=	  (  [ ẑ × ]ₜc  )ᵀ
"""

function update_k_pb(M̂::HelmholtzMap{T},k⃗::SVector{3,T}) where T<:Real
	(mag, m, n), mag_m_n_pb = Zygote.pullback(k⃗) do x
		mag_m_n(x,dropgrad(M̂.g⃗))
	end
	M̂.mag = mag
	M̂.inv_mag = [inv(mm) for mm in mag]
	M̂.m⃗ = m #HybridArray{Tuple{Dynamic(),Dynamic(),Dynamic(),3},T}(Array(m.parent))
	M̂.n⃗ = n #HybridArray{Tuple{Dynamic(),Dynamic(),Dynamic(),3},T}(Array(n.parent))
	M̂.m = HybridArray{Tuple{3,Dynamic(),Dynamic(),Dynamic()},T}(reinterpret(reshape,T,M̂.m⃗))
	M̂.n = HybridArray{Tuple{3,Dynamic(),Dynamic(),Dynamic()},T}(reinterpret(reshape,T,M̂.n⃗))
	M̂.k⃗ = k⃗
	return (mag, m, n), mag_m_n_pb
end

update_k_pb(M̂::HelmholtzMap{T},kz::T) where T<:Real = update_k_pb(M̂,SVector{3,T}(0.,0.,kz))

# 3D
function ε⁻¹_bar!(eī, d⃗, λ⃗d, Nx, Ny, Nz)
	# # capture 3x3 block diagonal elements of outer product -| λ⃗d X d⃗ |
	# # into (3,3,Nx,Ny,Nz) array. This is the gradient of ε⁻¹ tensor field
	eīf = flat(eī)
	@avx for iz=1:Nz,iy=1:Ny,ix=1:Nx
		q = (Nz * (iz-1) + Ny * (iy-1) + ix) # (Ny * (iy-1) + i)
		for a=1:3 # loop over diagonal elements: {11, 22, 33}
			eīf[a,a,ix,iy,iz] = real( -λ⃗d[3*q-2+a-1] * conj(d⃗[3*q-2+a-1]) )
		end
		for a2=1:2 # loop over first off diagonal
			eīf[a2,a2+1,ix,iy,iz] = real( -conj(λ⃗d[3*q-2+a2]) * d⃗[3*q-2+a2-1] - λ⃗d[3*q-2+a2-1] * conj(d⃗[3*q-2+a2]) )
		end
		# a = 1, set 1,3 and 3,1, second off-diagonal
		eīf[1,3,ix,iy,iz] = real( -conj(λ⃗d[3*q]) * d⃗[3*q-2] - λ⃗d[3*q-2] * conj(d⃗[3*q]) )
	end
	return eī
end

# 2D
function ε⁻¹_bar!(eīf, d⃗, λ⃗d, Nx, Ny)
	# # capture 3x3 block diagonal elements of outer product -| λ⃗d X d⃗ |
	# # into (3,3,Nx,Ny,Nz) array. This is the gradient of ε⁻¹ tensor field
	# eīf = flat(eī)
	@avx for iy=1:Ny,ix=1:Nx
		q = (Ny * (iy-1) + ix) # (Ny * (iy-1) + i)
		for a=1:3 # loop over diagonal elements: {11, 22, 33}
			eīf[a,a,ix,iy] = real( -λ⃗d[3*q-2+a-1] * conj(d⃗[3*q-2+a-1]) )
		end
		for a2=1:2 # loop over first off diagonal
			eīf[a2,a2+1,ix,iy] = real( -conj(λ⃗d[3*q-2+a2]) * d⃗[3*q-2+a2-1] - λ⃗d[3*q-2+a2-1] * conj(d⃗[3*q-2+a2]) )
			eīf[a2+1,a2,ix,iy] = eīf[a2,a2+1,ix,iy]
		end
		# a = 1, set 1,3 and 3,1, second off-diagonal
		eīf[1,3,ix,iy] = real( -conj(λ⃗d[3*q]) * d⃗[3*q-2] - λ⃗d[3*q-2] * conj(d⃗[3*q]) )
		eīf[3,1,ix,iy] = eīf[1,3,ix,iy]
	end
	return eī # inv( (eps' + eps) / 2)

	# eīM = Matrix.(eī)
	# for iy=1:Ny,ix=1:Nx
	# 	q = (Ny * (iy-1) + ix) # (Ny * (iy-1) + i)
	# 	for a=1:3 # loop over diagonal elements: {11, 22, 33}
	# 		eīM[ix,iy][a,a] = real( -λ⃗d[3*q-2+a-1] * conj(d⃗[3*q-2+a-1]) )
	# 	end
	# 	for a2=1:2 # loop over first off diagonal
	# 		eīM[ix,iy][a2,a2+1] = real( -conj(λ⃗d[3*q-2+a2]) * d⃗[3*q-2+a2-1] - λ⃗d[3*q-2+a2-1] * conj(d⃗[3*q-2+a2]) )
	# 	end
	# 	# a = 1, set 1,3 and 3,1, second off-diagonal
	# 	eīM[ix,iy][1,3] = real( -conj(λ⃗d[3*q]) * d⃗[3*q-2] - λ⃗d[3*q-2] * conj(d⃗[3*q]) )
	# end
	# ēM = inv.(eīM)
	# eīMH = inv.( ( ēM .+ ēM' ) ./ 2 )
	# eī .= SMatrix{3,3}.( eīMH  ) # SMatrix{3,3}.(eīM)
	# return eī
end

# 3D
function ε⁻¹_bar(d⃗::AbstractVector{Complex{T}}, λ⃗d, Nx, Ny, Nz) where T<:Real
	# # capture 3x3 block diagonal elements of outer product -| λ⃗d X d⃗ |
	# # into (3,3,Nx,Ny,Nz) array. This is the gradient of ε⁻¹ tensor field

	# eīf = flat(eī)
	eīf = Buffer(Array{Float64,1}([2., 2.]),3,3,Nx,Ny,Nz) # bufferfrom(zero(T),3,3,Nx,Ny,Nz)
	# eīf = bufferfrom(zero(eltype(real(d⃗)),3,3,Nx,Ny,Nz))
	@avx for iz=1:Nz,iy=1:Ny,ix=1:Nx
		q = (Nz * (iz-1) + Ny * (iy-1) + ix) # (Ny * (iy-1) + i)
		for a=1:3 # loop over diagonal elements: {11, 22, 33}
			eīf[a,a,ix,iy,iz] = real( -λ⃗d[3*q-2+a-1] * conj(d⃗[3*q-2+a-1]) )
		end
		for a2=1:2 # loop over first off diagonal
			eīf[a2,a2+1,ix,iy,iz] = real( -conj(λ⃗d[3*q-2+a2]) * d⃗[3*q-2+a2-1] - λ⃗d[3*q-2+a2-1] * conj(d⃗[3*q-2+a2]) )
		end
		# a = 1, set 1,3 and 3,1, second off-diagonal
		eīf[1,3,ix,iy,iz] = real( -conj(λ⃗d[3*q]) * d⃗[3*q-2] - λ⃗d[3*q-2] * conj(d⃗[3*q]) )
	end
	# eī = unflat(copy(eīf))
	# eī = reinterpret(reshape,SMatrix{3,3,T,9},reshape(copy(eīf),9,Nx,Ny,Nz))
	eī = copy(eīf)
	return eī
end

# 2D
function ε⁻¹_bar(d⃗::AbstractVector{Complex{T}}, λ⃗d, Nx, Ny) where T<:Real
	# # capture 3x3 block diagonal elements of outer product -| λ⃗d X d⃗ |
	# # into (3,3,Nx,Ny,Nz) array. This is the gradient of ε⁻¹ tensor field

	# eīf = flat(eī)
	eīf = Buffer(Array{Float64,1}([2., 2.]),3,3,Nx,Ny) # bufferfrom(zero(T),3,3,Nx,Ny)
	# eīf = bufferfrom(zero(eltype(real(d⃗)),3,3,Nx,Ny))
	@avx for iy=1:Ny,ix=1:Nx
		q = (Ny * (iy-1) + ix) # (Ny * (iy-1) + i)
		for a=1:3 # loop over diagonal elements: {11, 22, 33}
			eīf[a,a,ix,iy] = real( -λ⃗d[3*q-2+a-1] * conj(d⃗[3*q-2+a-1]) )
		end
		for a2=1:2 # loop over first off diagonal
			eīf[a2,a2+1,ix,iy] = real( -conj(λ⃗d[3*q-2+a2]) * d⃗[3*q-2+a2-1] - λ⃗d[3*q-2+a2-1] * conj(d⃗[3*q-2+a2]) )
			eīf[a2+1,a2,ix,iy] = eīf[a2,a2+1,ix,iy]
		end
		# a = 1, set 1,3 and 3,1, second off-diagonal
		eīf[1,3,ix,iy] = real( -conj(λ⃗d[3*q]) * d⃗[3*q-2] - λ⃗d[3*q-2] * conj(d⃗[3*q]) )
		eīf[3,1,ix,iy] = eīf[1,3,ix,iy]
	end
	# eī = reinterpret(reshape,SMatrix{3,3,T,9},reshape(copy(eīf),9,Nx,Ny))
	eī = copy(eīf)
	return eī # inv( (eps' + eps) / 2)
end

function solve_adj!(ms::ModeSolver,H̄,eigind::Int)
	ms.adj_itr = bicgstabl_iterator!(
		ms.adj_itr.x,	# recycle previous soln as initial guess
		ms.M̂ - real(ms.ω²)*I, # A
		H̄ - ms.H⃗ * dot(ms.H⃗,H̄), # b,
		3;	# l = number of GMRES iterations per CG iteration
		Pl = ms.P̂) # left preconditioner
	for (iteration, item) = enumerate(ms.adj_itr) end # iterate until convergence or until (iters > max_iters || mvps > max_mvps)
	copyto!(ms.λ⃗,ms.adj_itr.x) # copy soln. to ms.λ⃗ where other contributions/corrections can be accumulated
	# λ₀, ch = bicgstabl(
	# 	ms.adj_itr.x,	# recycle previous soln as initial guess
	# 	ms.M̂ - real(ms.ω²[eigind])*I, # A
	# 	H̄[:,eigind] - ms.H⃗[:,eigind] * dot(ms.H⃗[:,eigind],H̄[:,eigind]), # b,
	# 	3;	# l = number of GMRES iterations per CG iteration
	# 	Pl = ms.P̂, # left preconditioner
	# 	reltol = 1e-10,
	# 	log=true,
	# 	)
	# copyto!(ms.λ⃗,λ₀) # copy soln. to ms.λ⃗ where other contributions/corrections can be accumulated
	# println("\t\tAdjoint Problem for kz = $( ms.M̂.k⃗[3] ) ###########")
	# println("\t\t\tadj converged?: $ch")
	# println("\t\t\titrs, mvps: $(ch.iters), $(ch.mvps)")
	# uplot(ch;name="log10( adj. prob. res. )")
	return ms.λ⃗
end

function solve_adj!(λ⃗,M̂::HelmholtzMap,H̄,ω²,H⃗,eigind::Int;log=false)
	# log=true
	res = bicgstabl!(
		λ⃗, # ms.adj_itr.x,	# recycle previous soln as initial guess
		M̂ - real(ω²)*I, # A
		H̄ - H⃗ * dot(H⃗,H̄), # b,
		2;	# l = number of GMRES iterations per CG iteration
		# Pl = HelmholtzPreconditioner(M̂), # left preconditioner
		log,
		abstol=1e-8,
		max_mv_products=500
		)
	if log
		copyto!(λ⃗,res[1])
		ch = res[2]
	else
		copyto!(λ⃗,res)
	end
	# println("#########  Adjoint Problem for kz = $( M̂.k⃗[3] ) ###########")
	# uplot(ch;name="log10( adj. prob. res. )")
	# println("\t\t\tadj converged?: $ch")
	# println("\t\t\titrs, mvps: $(ch.iters), $(ch.mvps)")
	return λ⃗
end

using LinearMaps: ⊗
export eig_adjt, linsolve
using IterativeSolvers: gmres
function linsolve(Â, b⃗; x⃗₀=nothing, P̂=IterativeSolvers.Identity())
	# x⃗ = isnothing(x⃗₀) ? randn(eltype(b⃗),first(size(b⃗))) : copy(x⃗₀)
	# x⃗ = isnothing(x⃗₀) ? zero(b⃗) : copy(x⃗₀)

	# return bicgstabl!(x⃗, Â, b⃗, 2; Pl=P̂, max_mv_products=5000)
	# return bicgstabl!(x⃗, Â, b⃗, 2; Pl=P̂, max_mv_products=3000)
	# bicgstabl(Â, b⃗, 3; Pl=P̂, max_mv_products=3000)
	# cg(Â, b⃗; Pl=P̂, maxiter=3000)
	# bicgstabl(Â, b⃗, 2; Pl=P̂, max_mv_products=10000)
	gmres(Â, b⃗; Pl=P̂, maxiter=1000)
end

function rrule(::typeof(linsolve), Â, b⃗;
		x⃗₀=nothing, P̂=IterativeSolvers.Identity())
	x⃗ = linsolve(Â, b⃗; x⃗₀, P̂)
	function linsolve_pullback(x̄)
		λ⃗ = linsolve(Â', vec(x̄))
		Ā = (-λ⃗) ⊗ x⃗'
		return (NoTangent(), Ā, λ⃗)
	end
	return (x⃗, linsolve_pullback)
end

"""
	eig_adjt(A, α, x⃗, ᾱ, x̄; λ⃗₀, P̂)

Compute the adjoint vector `λ⃗` for a single eigenvalue/eigenvector pair (`α`,`x⃗`) of `Â` and
sensitivities (`ᾱ`,`x̄`). It is assumed (but not checked) that ``Â ⋅ x⃗ = α x⃗``. `λ⃗` is the
sum of two components,

	``λ⃗ = λ⃗ₐ + λ⃗ₓ``

where ``λ⃗ₐ = ᾱ x⃗`` and ``λ⃗ₓ`` correspond to `ᾱ` and `x̄`, respectively. When `x̄` is non-zero
``λ⃗ₓ`` is computed by iteratively solving

	``(Â - αÎ) ⋅ λ⃗ₓ = x̄ - (x⃗ ⋅ x̄)``

An inital guess can be supplied for `λ⃗ₓ` via the keyword argument `λ⃗₀`, otherwise a random
vector is used. A preconditioner `P̂` can also be supplied to improve convergeance.
"""
function eig_adjt(Â, α, x⃗, ᾱ, x̄; λ⃗₀=nothing, P̂=IterativeSolvers.Identity())
	if iszero(x̄)
		λ⃗ = iszero(ᾱ)	? zero(x⃗) : ᾱ * x⃗
 	else
		λ⃗ₓ₀ = linsolve(
			Â - α*I,
		 	x̄ - x⃗ * dot(x⃗,x̄);
			P̂,
			# x⃗₀=λ⃗₀,
		)
		# λ⃗ -= x⃗ * dot(x⃗,λ⃗)	# re-orthogonalize λ⃗ w.r.t. x⃗, correcting for round-off err.
		# λ⃗ += ᾱ * x⃗
		λ⃗ₓ = λ⃗ₓ₀  - x⃗ * dot(x⃗,λ⃗ₓ₀)	# re-orthogonalize λ⃗ w.r.t. x⃗, correcting for round-off err.
		λ⃗ = λ⃗ₓ + ᾱ * x⃗
	end

	return λ⃗
end

function rrule(::typeof(eig_adjt), Â, α, x⃗, ᾱ, x̄; λ⃗₀=nothing, P̂=IterativeSolvers.Identity())
	# primal
	if iszero(x̄)
		λ⃗ = iszero(ᾱ)	? zero(x⃗) : ᾱ * x⃗
 	else
		λ⃗ₓ₀ = linsolve(
			Â - α*I,
		 	x̄ - x⃗ * dot(x⃗,x̄);
			P̂,
			# x⃗₀=λ⃗₀,
		)
		λ⃗ₓ = λ⃗ₓ₀ - x⃗ * dot(x⃗,λ⃗ₓ₀)	# re-orthogonalize λ⃗ w.r.t. x⃗, correcting for round-off err.
		λ⃗ = λ⃗ₓ + ᾱ * x⃗
	end

	# define pullback
	# function eig_adjt_pullback(lm̄)
	# 	if lm̄ isa AbstractZero
	# 		return (NoTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent())
	# 	else
	# 		if iszero(x̄)
	# 			if iszero(ᾱ)
	# 				return (NoTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent())
	# 			else
	# 				x⃗_bar = ᾱ * lm̄
	# 				ᾱ_bar = dot(lm̄,x⃗)
	# 				return (NoTangent(), ZeroTangent(), ZeroTangent(), x⃗_bar, ᾱ_bar, ZeroTangent())
	# 			end
	# 	 	else	# λ⃗ₓ exists, must solve 2ⁿᵈ-order adjoint problem
	# 			lm̄_perp = lm̄ - x⃗ * dot(x⃗,lm̄)
	# 			ξ⃗ = linsolve(
	# 				Â - α*I,
	# 			 	lm̄_perp;
	# 				P̂,
	# 				# x⃗₀=λ⃗ₓ,
	# 			)
	# 			# println("")
	# 			# println("ξ⃗ err: $(sum(abs2,(Â - α*I)*ξ⃗-(lm̄ - x⃗ * dot(x⃗,lm̄))))")
	#
	# 			xv_dot_ξ⃗ = dot(x⃗,ξ⃗)
	# 			ξ⃗perp = ξ⃗ - x⃗ * xv_dot_ξ⃗
	#
	# 			# println("sum(abs2,ξ⃗): $(sum(abs2,ξ⃗)))")
	# 			# println("sum(abs2,ξ⃗perp): $(sum(abs2,ξ⃗perp)))")
	# 			# println("sum(abs2,lm̄): $(sum(abs2,lm̄)))")
	# 			# println("sum(abs2,lm̄_perp): $(sum(abs2,lm̄_perp)))")
	#
	# 			A_bar = (-ξ⃗perp) ⊗ λ⃗ₓ' # * (2*length(x⃗)))
	# 			α_bar = dot(ξ⃗,λ⃗ₓ)
	# 			x⃗_bar = dot(x̄,x⃗) * ξ⃗perp + ᾱ * lm̄ + dot(ξ⃗,x⃗) * x̄
	# 			ᾱ_bar = dot(lm̄,x⃗)
	# 			x̄_bar = ξ⃗perp
	#
	# 			# A_bar = -ξ⃗ ⊗ λ⃗ₓ'
	# 			# α_bar = dot(ξ⃗,λ⃗ₓ)
	# 			# x⃗_bar = dot(x̄,x⃗) * ξ⃗ + dot(ξ⃗,x⃗) * x̄  + ᾱ * lm̄
	# 			# ᾱ_bar = dot(lm̄,x⃗)
	# 			# x̄_bar = ξ⃗ - (x⃗ * dot(x⃗,ξ⃗))
	# 		end
	# 		# return (NoTangent(), A_bar, α_bar, x⃗_bar, ᾱ_bar, x̄_bar)
	# 		return (NoTangent(), A_bar, α_bar, x⃗_bar, ᾱ_bar, x̄_bar)
	# 		# return (NoTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent())
	# 	end
	# end

	# define pullback
	function eig_adjt_pullback(lm̄)
		if lm̄ isa AbstractZero
			return (NoTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent())
		else
			if iszero(x̄)
				if iszero(ᾱ)
					return (NoTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent(), ZeroTangent())
				else
					x⃗_bar = ᾱ * lm̄
					ᾱ_bar = dot(lm̄,x⃗)
					return (NoTangent(), ZeroTangent(), ZeroTangent(), x⃗_bar, ᾱ_bar, ZeroTangent())
				end
		 	else	# λ⃗ₓ exists, must solve 2ⁿᵈ-order adjoint problem
				ξ⃗ = linsolve(
					Â - α*I,
				 	lm̄ - x⃗ * dot(x⃗,lm̄);
					P̂,
				)
				# println("")
				# println("ξ⃗ err: $(sum(abs2,(Â - α*I)*ξ⃗-(lm̄ - x⃗ * dot(x⃗,lm̄))))")
				A_bar = -ξ⃗ ⊗ λ⃗ₓ'
				α_bar = dot(ξ⃗,λ⃗ₓ)
				x⃗_bar = dot(x̄,x⃗) * ξ⃗ + dot(ξ⃗,x⃗) * x̄  + ᾱ * lm̄
				ᾱ_bar = dot(lm̄,x⃗)
				x̄_bar = ξ⃗ - (x⃗ * dot(x⃗,ξ⃗))
			end
			return (NoTangent(), A_bar, α_bar, x⃗_bar, ᾱ_bar, x̄_bar)
		end
	end

	# return primal and pullback
	return (λ⃗, eig_adjt_pullback)
end

function ∇solve_ω²(ΔΩ,Ω,k,ε⁻¹,grid)
	@show ω̄sq, H̄ = ΔΩ
	@show ω², H⃗ = Ω
	M̂ = HelmholtzMap(k,ε⁻¹,grid)
	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	Nranges = eachindex(grid) #(1:NN for NN in Ns) # 1:Nx, 1:Ny, 1:Nz for 3D, 1:Nx, 1:Ny for 2D
	H = reshape(H⃗,(2,Ns...))
	g⃗s = g⃗(dropgrad(grid))
	(mag, m⃗, n⃗), mag_m_n_pb = Zygote.pullback(x->mag_m_n(x,g⃗s),k)
	λd = similar(M̂.d)
	λẽ = similar(M̂.d)
	ẽ = similar(M̂.d)
	ε⁻¹_bar = similar(ε⁻¹)
	if typeof(ω̄sq)==ZeroTangent()
		ω̄sq = 0.
	end
	if typeof(H̄) != ZeroTangent()
		λ⃗ = solve_adj!(M̂,H̄,ω²,H⃗,eigind) 												# overwrite ms.λ⃗ with soln to (M̂ + ω²I) λ⃗ = H̄ - dot(H⃗,H̄)*H⃗
		λ⃗ -= (ω̄sq + dot(H⃗,λ⃗)) * H⃗
	else
		λ⃗ = -ω̄sq * H⃗
	end
	λ = reshape(λ⃗,(2,Ns...))
	d = _H2d!(M̂.d, H * M̂.Ninv, M̂) # =  M̂.𝓕 * kx_tc( H , mn2, mag )  * M̂.Ninv
	λd = _H2d!(λd,λ,M̂) # M̂.𝓕 * kx_tc( reshape(λ⃗,(2,M̂.Nx,M̂.Ny,M̂.Nz)) , mn2, mag )
	ε⁻¹_bar!(ε⁻¹_bar, vec(M̂.d), vec(λd), Ns...)
	# back-propagate gradients w.r.t. `(k⃗+g⃗)×` operator to k via (m⃗,n⃗) pol. basis and |k⃗+g⃗|
	λd *=  M̂.Ninv
	λẽ .= reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(λẽ , λd  , M̂ ) )
	ẽ .= reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(M̂.e,M̂.d, M̂) )
	kx̄_m⃗ = real.( λẽ .* conj.(view(H,2,Nranges...)) .+ ẽ .* conj.(view(λ,2,Nranges...)) )
	kx̄_n⃗ =  -real.( λẽ .* conj.(view(H,1,Nranges...)) .+ ẽ .* conj.(view(λ,1,Nranges...)) )
	māg .= dot.(n⃗, kx̄_n⃗) + dot.(m⃗, kx̄_m⃗)
	k̄ = -mag_m_n_pb(( māg, kx̄_m⃗.*mag, kx̄_n⃗.*mag ))[1] # m̄ = kx̄_m⃗ .* mag, n̄ = kx̄_n⃗ .* mag, #NB: not sure why this is needs to be negated, inputs match original version
	# if !(typeof(k)<:SVector)
	# 	k̄_kx = k̄_kx[3]
	# end
	return (NoTangent(), ZeroTangent(), k̄ , ε⁻¹_bar)
end

function rrule(::typeof(solve_ω²), k::Union{T,SVector{3,T}},shapes::Vector{<:Shape},grid::Grid{ND};
		nev=1,eigind=1,maxiter=3000,tol=1e-8,log=false) where {ND,T<:Real}
	# println("using new rrule")
	ms = @ignore(ModeSolver(k, shapes, grid)) # ; nev, eigind, maxiter, tol, log))
	ε⁻¹ = εₛ⁻¹(shapes;ms=dropgrad(ms))
	ω²H⃗ = solve_ω²(ms,k,ε⁻¹; nev, eigind, maxiter, tol, log)
    solve_ω²_pullback(ΔΩ) = ∇solve_ω²(ΔΩ,ω²H⃗,k,ε⁻¹,grid)
    return (ω²H⃗, solve_ω²_pullback)
end

function rrule(::typeof(solve_ω²), ms::ModeSolver{ND,T},k::Union{T,SVector{3,T}},ε⁻¹::AbstractArray{<:SMatrix{3,3},ND};
		nev=1,eigind=1,maxiter=3000,tol=1e-8,log=false) where {ND,T<:Real}
	ω²,H⃗ = solve_ω²(ms,k,ε⁻¹; nev, eigind, maxiter, tol, log)
	(mag, m⃗, n⃗), mag_m_n_pb = Zygote.pullback(k) do x
		mag_m_n(x,dropgrad(ms.M̂.g⃗))
	end
    function solve_ω²_pullback(ΔΩ)
		ω̄sq, H̄ = ΔΩ
		Ns = size(ms.grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
		Nranges = eachindex(ms.grid) #(1:NN for NN in Ns) # 1:Nx, 1:Ny, 1:Nz for 3D, 1:Nx, 1:Ny for 2D
		H = reshape(H⃗,(2,Ns...))
		# mn2 = vcat(reshape(ms.M̂.m,(1,3,Ns...)),reshape(ms.M̂.n,(1,3,Ns...)))
		if typeof(ω̄sq)==ZeroTangent()
			ω̄sq = 0.
		end
		if typeof(H̄) != ZeroTangent()
			solve_adj!(ms,H̄,eigind) 												# overwrite ms.λ⃗ with soln to (M̂ + ω²I) λ⃗ = H̄ - dot(H⃗,H̄)*H⃗
			ms.λ⃗ -= (ω̄sq[eigind] + dot(H⃗,ms.λ⃗)) * H⃗
		else
			ms.λ⃗ = -ω̄sq[eigind] * H⃗
		end
		λ = reshape(ms.λ⃗,(2,Ns...))
		d = _H2d!(ms.M̂.d, H * ms.M̂.Ninv, ms) # =  ms.M̂.𝓕 * kx_tc( H , mn2, mag )  * ms.M̂.Ninv
		λd = _H2d!(ms.λd,λ,ms) # ms.M̂.𝓕 * kx_tc( reshape(ms.λ⃗,(2,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)) , mn2, mag )
		ε⁻¹_bar!(ms.ε⁻¹_bar, vec(ms.M̂.d), vec(ms.λd), Ns...)
		# back-propagate gradients w.r.t. `(k⃗+g⃗)×` operator to k via (m⃗,n⃗) pol. basis and |k⃗+g⃗|
		ms.λd *=  ms.M̂.Ninv
		λẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(ms.λẽ , ms.λd  ,ms ) )
		ẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(ms.M̂.e,ms.M̂.d,ms) )
		ms.kx̄_m⃗ .= real.( λẽ .* conj.(view(H,2,Nranges...)) .+ ẽ .* conj.(view(λ,2,Nranges...)) )
		ms.kx̄_n⃗ .=  -real.( λẽ .* conj.(view(H,1,Nranges...)) .+ ẽ .* conj.(view(λ,1,Nranges...)) )
		ms.māg .= dot.(n⃗, ms.kx̄_n⃗) + dot.(m⃗, ms.kx̄_m⃗)
		k̄ = -mag_m_n_pb(( ms.māg, ms.kx̄_m⃗.*mag, ms.kx̄_n⃗.*mag ))[1] # m̄ = kx̄_m⃗ .* mag, n̄ = kx̄_n⃗ .* mag, #NB: not sure why this is needs to be negated, inputs match original version
		# if !(typeof(k)<:SVector)
		# 	k̄_kx = k̄_kx[3]
		# end
		return (NoTangent(), ZeroTangent(), k̄ , ms.ε⁻¹_bar)
    end
    return ((ω², H⃗), solve_ω²_pullback)
end

"""
function mapping |H⟩ ⤇ ( (∂M/∂k)ᵀ + ∂M/∂k )|H⟩
"""
function Mₖᵀ_plus_Mₖ(H⃗::AbstractVector{Complex{T}},k,ε⁻¹,grid) where T<:Real
	Ns = size(grid)
	# g⃗s = g⃗(grid)
	# mag,m⃗,n⃗ = mag_m_n(k,g⃗s)
	mag,m⃗,n⃗ = mag_m_n(k,grid)
	H = reshape(H⃗,(2,Ns...))
	mn = vcat(reshape(flat(m⃗),1,3,Ns...),reshape(flat(n⃗),1,3,Ns...))
	X = -kx_ct( ifft( ε⁻¹_dot( fft( zx_tc(H,mn), (2:3) ), real(ε⁻¹)), (2:3) ), mn, mag )
	Y = zx_ct( ifft( ε⁻¹_dot( fft( kx_tc(H,mn,mag), (2:3) ), real(ε⁻¹)), (2:3)), mn )
	vec(X + Y)
end

"""
solve the adjoint sensitivity problem corresponding to ∂ω²∂k = <H|∂M/∂k|H>
"""
function ∂ω²∂k_adj(M̂::HelmholtzMap,ω²,H⃗,H̄;eigind=1,log=false)
	res = bicgstabl(
		M̂ - real(ω²)*I, # A
		H̄ - H⃗ * dot(H⃗,H̄), # b,
		3;	# l = number of GMRES iterations per CG iteration
		# Pl = HelmholtzPreconditioner(M̂), # left preconditioner
		log,
		)
end

"""
pull back sensitivity w.r.t. ∂ω²∂k = 2⟨H|∂M/∂k|H⟩ to corresponding
k̄ (scalar) and nn̄g⁻¹ (tensor field) sensitivities
"""
function ∇HMₖH(k::Real,H⃗::AbstractArray{Complex{T}},nng⁻¹::AbstractArray{T2,N2},grid::Grid{ND};eigind=1) where {T<:Real,ND,T2<:Real,N2}
	# Setup
	local 	zxtc_to_mn = SMatrix{3,3}(	[	0 	-1	  0
											1 	 0	  0
											0 	 0	  0	  ]	)

	local 	kxtc_to_mn = SMatrix{2,2}(	[	0 	-1
											1 	 0	  ]	)

	g⃗s, Ninv, Ns, 𝓕, 𝓕⁻¹ = Zygote.ignore() do
		Ninv 		= 		1. / N(grid)
		Ns			=		size(grid)
		g⃗s = g⃗(grid)
		d0 = randn(Complex{T}, (3,Ns...))
		𝓕	 =	plan_fft(d0,_fftaxes(grid),flags=FFTW.PATIENT) # planned out-of-place FFT operator 𝓕
		𝓕⁻¹ =	plan_bfft(d0,_fftaxes(grid),flags=FFTW.PATIENT) # planned out-of-place iFFT operator 𝓕⁻¹
		return (g⃗s,Ninv,Ns,𝓕,𝓕⁻¹)
	end
	mag, m⃗, n⃗  = mag_m_n(k,g⃗s)
	H = reshape(H⃗,(2,Ns...))
	Hsv = reinterpret(reshape, SVector{2,Complex{T}}, H )

	#TODO: Banish this quadruply re(shaped,interpreted) m,n,mns format back to hell
	# mns = mapreduce(x->reshape(flat(x),(1,3,size(x)...)),vcat,(m⃗,n⃗))
	m = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,m⃗)))
	n = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,n⃗)))
	mns = vcat(reshape(m,(1,3,Ns...)),reshape(n,(1,3,Ns...)))

	### calculate k̄ contribution from M̄ₖ ( from ⟨H|M̂ₖ|H⟩ )
	Ā₁		=	conj.(Hsv)
	Ā₂ = reinterpret(
		reshape,
		SVector{3,Complex{T}},
		# reshape(
		# 	𝓕⁻¹ * nngsp * 𝓕 * zxtcsp * vec(H),
		# 	(3,size(gr)...),
		# 	),
		𝓕⁻¹ * ε⁻¹_dot(  𝓕 * zx_tc(H * Ninv,mns) , real(nng⁻¹)),
		)
	Ā 	= 	Ā₁  .*  transpose.( Ā₂ )
	m̄n̄_Ā = transpose.( (kxtc_to_mn,) .* real.(Ā) )
	m̄_Ā = 		view.( m̄n̄_Ā, (1:3,), (1,) )
	n̄_Ā = 		view.( m̄n̄_Ā, (1:3,), (2,) )
	māg_Ā = dot.(n⃗, n̄_Ā) + dot.(m⃗, m̄_Ā)

	# # diagnostic for nngī accuracy
	# B̄₁_old = reinterpret(
	# 	reshape,
	# 	SVector{3,Complex{T}},
	# 	# 𝓕  *  kxtcsp	 *	vec(H),
	# 	𝓕 * kx_tc( conj.(H) ,mns,mag),
	# 	)
	# B̄₂_old = reinterpret(
	# 	reshape,
	# 	SVector{3,Complex{T}},
	# 	# 𝓕  *  zxtcsp	 *	vec(H),
	# 	𝓕 * zx_tc( H * Ninv ,mns),
	# 	)
	# B̄_old 	= 	 SMatrix{3,3,Float64,9}.(real.(Hermitian.(  B̄₁_old  .*  transpose.( B̄₂_old )  )) )
	# B̄_oldf = copy(flat(B̄_old))
	# println("sum(B̄_oldf): $(sum(B̄_oldf))")
	# println("maximum(B̄_oldf): $(maximum(B̄_oldf))")
	# # end diagnostic for nngī accuracy

	B̄₁ = 𝓕 * kx_tc( conj.(H) ,mns,mag)
	B̄₂ = 𝓕 * zx_tc( H * Ninv ,mns)
	@tullio B̄[a,b,i,j] := B̄₁[a,i,j] * B̄₂[b,i,j] + B̄₁[b,i,j] * B̄₂[a,i,j]   #/2 + real(B̄₁[b,i,j] * B̄₂[a,i,j])/2

	# # diagnostic for nngī accuracy
	#
	# # println("sum(B̄): $(sum(real(B̄)))")
	# # println("maximum(B̄): $(maximum(real(B̄)))")
	# B̄_herm = real(B̄)/2
	# println("sum(B̄_herm): $(sum(B̄_herm))")
	# println("maximum(B̄_herm): $(maximum(B̄_herm))")
	# # end diagnostic for nngī accuracy

	C̄₁ = reinterpret(
		reshape,
		SVector{3,Complex{T}},
		# reshape(
		# 	𝓕⁻¹ * nngsp * 𝓕 * kxtcsp * -vec(H),
		# 	(3,size(gr)...),
		# 	),
		𝓕⁻¹ * ε⁻¹_dot(  𝓕 * -kx_tc( H * Ninv, mns, mag) , nng⁻¹),
		)
	C̄₂ =   conj.(Hsv)
	C̄ 	= 	C̄₁  .*  transpose.( C̄₂ )
	m̄n̄_C̄ = 			 (zxtc_to_mn,) .* real.(C̄)
	m̄_C̄ = 		view.( m̄n̄_C̄, (1:3,), (1,) )
	n̄_C̄ = 		view.( m̄n̄_C̄, (1:3,), (2,) )

	# Accumulate gradients and pull back
	nngī 	=  real(B̄)/2 #( B̄ .+ transpose.(B̄) ) ./ 2
	k̄	 	= ∇ₖmag_m_n(
						māg_Ā, 				# māg total
						m̄_Ā.*mag .+ m̄_C̄, 	  # m̄  total
						n̄_Ā.*mag .+ n̄_C̄,	  # n̄  total
						mag,
						m⃗,
						n⃗,
					)
	# H̄ = Mₖᵀ_plus_Mₖ(H⃗,k,nng⁻¹,grid)
	# Y = zx_ct( ifft( ε⁻¹_dot( fft( kx_tc(H,mns,mag), (2:3) ), nng⁻¹), (2:3)), mns )
	# X = -kx_ct( ifft( ε⁻¹_dot( fft( zx_tc(H,mns), (2:3) ), nng⁻¹), (2:3) ), mns, mag )

	# nngif = real(flat(nng⁻¹))
	X = -kx_ct( 𝓕⁻¹ * ε⁻¹_dot( 𝓕 * zx_tc(H,mns)		, nng⁻¹), mns, mag )
	Y =  zx_ct( 𝓕⁻¹ * ε⁻¹_dot( 𝓕 * kx_tc(H,mns,mag)	, nng⁻¹), mns )
	H̄ = vec(X + Y) * Ninv
	return k̄, H̄, nngī
	# return k̄, H̄, reinterpret(SMatrix{3,3,Float64,9},reshape( nngī ,9*128,128))
end

"""
pull back sensitivity w.r.t. ng_z = ∂kz/∂ω = ⟨E|(nng+ε)|E⟩ / ∫dA 2Re( conj(E) × H)⋅ẑ ) to corresponding
k̄ (scalar) and nn̄g⁻¹ (tensor field) sensitivities
"""
function ∇ng_z(k::Real,H⃗::AbstractArray{Complex{T}},nng⁻¹::AbstractArray{T2,N2},grid::Grid{ND};eigind=1) where {T<:Real,ND,T2<:Real,N2}
	# Setup
	local 	kxtc_to_mn = SMatrix{3,3}(	[	0 	-1	  0
											1 	 0	  0
											0 	 0	  0	  ]	)

	local 	kxtc_to_mn = SMatrix{2,2}(	[	0 	-1
											1 	 0	  ]	)

	g⃗s, Ninv, Ns, 𝓕, 𝓕⁻¹ = Zygote.ignore() do
		Ninv 		= 		1. / N(grid)
		Ns			=		size(grid)
		g⃗s = g⃗(grid)
		d0 = randn(Complex{T}, (3,Ns...))
		𝓕	 =	plan_fft(d0,_fftaxes(grid),flags=FFTW.PATIENT) # planned out-of-place FFT operator 𝓕
		𝓕⁻¹ =	plan_bfft(d0,_fftaxes(grid),flags=FFTW.PATIENT) # planned out-of-place iFFT operator 𝓕⁻¹
		return (g⃗s,Ninv,Ns,𝓕,𝓕⁻¹)
	end
	mag, m⃗, n⃗  = mag_m_n(k,g⃗s)
	H = reshape(H⃗,(2,Ns...))
	Hsv = reinterpret(reshape, SVector{2,Complex{T}}, H )

	#TODO: Banish this quadruply re(shaped,interpreted) m,n,mns format back to hell
	# mns = mapreduce(x->reshape(flat(x),(1,3,size(x)...)),vcat,(m⃗,n⃗))
	m = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,m⃗)))
	n = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,n⃗)))
	mns = vcat(reshape(m,(1,3,Ns...)),reshape(n,(1,3,Ns...)))

	### calculate k̄ contribution from M̄ₖ ( from ⟨H|M̂ₖ|H⟩ )
	Ā₁		=	conj.(Hsv)
	Ā₂ = reinterpret(
		reshape,
		SVector{3,Complex{T}},
		# reshape(
		# 	𝓕⁻¹ * nngsp * 𝓕 * zxtcsp * vec(H),
		# 	(3,size(gr)...),
		# 	),
		𝓕⁻¹ * ε⁻¹_dot(  𝓕 * zx_tc(H * Ninv,mns) , real(nng⁻¹)),
		)
	Ā 	= 	Ā₁  .*  transpose.( Ā₂ )
	m̄n̄_Ā = transpose.( (kxtc_to_mn,) .* real.(Ā) )
	m̄_Ā = 		view.( m̄n̄_Ā, (1:3,), (1,) )
	n̄_Ā = 		view.( m̄n̄_Ā, (1:3,), (2,) )
	māg_Ā = dot.(n⃗, n̄_Ā) + dot.(m⃗, m̄_Ā)

	# # diagnostic for nngī accuracy
	# B̄₁_old = reinterpret(
	# 	reshape,
	# 	SVector{3,Complex{T}},
	# 	# 𝓕  *  kxtcsp	 *	vec(H),
	# 	𝓕 * kx_tc( conj.(H) ,mns,mag),
	# 	)
	# B̄₂_old = reinterpret(
	# 	reshape,
	# 	SVector{3,Complex{T}},
	# 	# 𝓕  *  zxtcsp	 *	vec(H),
	# 	𝓕 * zx_tc( H * Ninv ,mns),
	# 	)
	# B̄_old 	= 	 SMatrix{3,3,Float64,9}.(real.(Hermitian.(  B̄₁_old  .*  transpose.( B̄₂_old )  )) )
	# B̄_oldf = copy(flat(B̄_old))
	# println("sum(B̄_oldf): $(sum(B̄_oldf))")
	# println("maximum(B̄_oldf): $(maximum(B̄_oldf))")
	# # end diagnostic for nngī accuracy

	B̄₁ = 𝓕 * kx_tc( conj.(H) ,mns,mag)
	B̄₂ = 𝓕 * zx_tc( H * Ninv ,mns)
	@tullio B̄[a,b,i,j] := B̄₁[a,i,j] * B̄₂[b,i,j] + B̄₁[b,i,j] * B̄₂[a,i,j]   #/2 + real(B̄₁[b,i,j] * B̄₂[a,i,j])/2

	# # diagnostic for nngī accuracy
	#
	# # println("sum(B̄): $(sum(real(B̄)))")
	# # println("maximum(B̄): $(maximum(real(B̄)))")
	# B̄_herm = real(B̄)/2
	# println("sum(B̄_herm): $(sum(B̄_herm))")
	# println("maximum(B̄_herm): $(maximum(B̄_herm))")
	# # end diagnostic for nngī accuracy

	C̄₁ = reinterpret(
		reshape,
		SVector{3,Complex{T}},
		# reshape(
		# 	𝓕⁻¹ * nngsp * 𝓕 * kxtcsp * -vec(H),
		# 	(3,size(gr)...),
		# 	),
		𝓕⁻¹ * ε⁻¹_dot(  𝓕 * -kx_tc( H * Ninv, mns, mag) , nng⁻¹),
		)
	C̄₂ =   conj.(Hsv)
	C̄ 	= 	C̄₁  .*  transpose.( C̄₂ )
	m̄n̄_C̄ = 			 (zxtc_to_mn,) .* real.(C̄)
	m̄_C̄ = 		view.( m̄n̄_C̄, (1:3,), (1,) )
	n̄_C̄ = 		view.( m̄n̄_C̄, (1:3,), (2,) )

	# Accumulate gradients and pull back
	nngī 	=  real(B̄)/2 #( B̄ .+ transpose.(B̄) ) ./ 2
	k̄	 	= ∇ₖmag_m_n(
						māg_Ā, 				# māg total
						m̄_Ā.*mag .+ m̄_C̄, 	  # m̄  total
						n̄_Ā.*mag .+ n̄_C̄,	  # n̄  total
						mag,
						m⃗,
						n⃗,
					)
	# H̄ = Mₖᵀ_plus_Mₖ(H⃗,k,nng⁻¹,grid)
	# Y = zx_ct( ifft( ε⁻¹_dot( fft( kx_tc(H,mns,mag), (2:3) ), nng⁻¹), (2:3)), mns )
	# X = -kx_ct( ifft( ε⁻¹_dot( fft( zx_tc(H,mns), (2:3) ), nng⁻¹), (2:3) ), mns, mag )

	# nngif = real(flat(nng⁻¹))
	X = -kx_ct( 𝓕⁻¹ * ε⁻¹_dot( 𝓕 * zx_tc(H,mns)		, nng⁻¹), mns, mag )
	Y =  zx_ct( 𝓕⁻¹ * ε⁻¹_dot( 𝓕 * kx_tc(H,mns,mag)	, nng⁻¹), mns )
	H̄ = vec(X + Y) * Ninv
	return k̄, H̄, nngī
	# return k̄, H̄, reinterpret(SMatrix{3,3,Float64,9},reshape( nngī ,9*128,128))
end

# nng = inv.(nnginv)
# ε = inv.(ε⁻¹)
# ∂ε∂ω_man = (2/ω) * (nng .- ε)
# ∂ei∂ω_man = copy(flat(-(ε⁻¹.^2) .* ∂ε∂ω_man ))
# ∂ε⁻¹_∂ω(ε⁻¹,nng⁻¹,ω) = -(2.0/ω) * ε⁻¹.^2 .* (  inv.(nng⁻¹) .- inv.(ε⁻¹) ) #(2.0/ω) * ε⁻¹ .* (  ε⁻¹ .* inv.(nng⁻¹) - I )

# function ∂nng∂ω_man_LN(om)
# 	 ng = ng_MgO_LiNbO₃(inv(om))[1,1]
# 	 n = sqrt(ε_MgO_LiNbO₃(inv(om))[1,1])
# 	 gvd = gvd_MgO_LiNbO₃(inv(om))[1,1]  #/ (2π)
# 	 # om = 1/om
# 	 om*(ng^2 - n*ng) + n * gvd
# end

# previously working
# ∂ε⁻¹_∂ω(ε⁻¹,nng⁻¹,ω) = -(2.0/ω) * (  ε⁻¹.^2 .* inv.(nng⁻¹) .- ε⁻¹ )
# ∂nng⁻¹_∂ω(ε⁻¹,nng⁻¹,ngvd,ω) = -(nng⁻¹.^2 ) .* ( ω*(ε⁻¹.*inv.(nng⁻¹).^2 .- inv.(nng⁻¹)) .+ ngvd) # (1.0/ω) * (nng⁻¹ .- ε⁻¹ ) .- (  ngvd .* (nng⁻¹).^2  )

"""
	∂ε⁻¹_∂ω(ε⁻¹,nng,ω) computes:
  ∂ε⁻¹_∂ω(ε⁻¹,nng⁻¹,ω) = -(2.0/ω) * (  ε⁻¹.^2 .* inv.(nng⁻¹) .- ε⁻¹ )
"""
function ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω)
	deps_dom = inv(ω) * (nng - ε)
	dei_dom = -1.0 * _dot(ε⁻¹,deps_dom,ε⁻¹)  #-(2.0/om) * ( _dot(ei,ei,nng) - ei )
end

"""
	∂nng⁻¹_∂ω(ε⁻¹,nng,nng⁻¹,ngvd,ω) computes:
  -(nng⁻¹.^2 ) .* ( ω*(ε⁻¹.*inv.(nng⁻¹).^2 .- inv.(nng⁻¹)) .+ ngvd)
"""
function ∂nng⁻¹_∂ω(ε⁻¹,nng,nng⁻¹,ngvd,ω)
	dnngi_dom = -1*_dot(nng⁻¹, ngvd, nng⁻¹)
end

"""
	herm(A) converts backpropagating gradients of nested Hermitian arrays/tensors to be Hermitian
"""
function herm(A::AbstractArray{TA,N}) where {TA<:AbstractMatrix,N}
	(real.(A) .+ transpose.(real.(A)) ) ./ 2
end

function herm(A::AbstractArray{T,4}) where {T<:Number}
	0.5 * real(A + conj(permutedims(A,(2,1,3,4))))
end

function herm(A::AbstractArray{T,5}) where {T<:Number}
	0.5 * real(A + conj(permutedims(A,(2,1,3,4,5))))
end

"""
solve the adjoint sensitivity problem corresponding to ∂ω²∂k = <H|∂M/∂k|H>
"""
# function ∂²ω²∂k²(ω,ε⁻¹,nng⁻¹,k,Hv,grid::Grid{ND,T};eigind=1,log=true) where {ND,T<:Real}
function ∂²ω²∂k²(ω,p,geom_fn,k,Hv,grid::Grid{ND,T};eigind=1,log=true) where {ND,T<:Real}

	# nng⁻¹, nnginv_pb = Zygote.pullback(nngₛ⁻¹,ω,geom,grid)
	# ε⁻¹, epsi_pb = Zygote.pullback(εₛ⁻¹,ω,geom,grid)

	# ngvd = ngvdₛ(ω,geom,grid)
	# nng⁻¹, nnginv_pb = Zygote.pullback(x->nngₛ⁻¹(x,geom,grid),ω)
	# ε⁻¹, epsi_pb = Zygote.pullback(x->εₛ⁻¹(x,geom,grid),ω)

	# nng⁻¹, nnginv_pb = Zygote._pullback(Zygote.Context(),x->nngₛ⁻¹(x,dropgrad(geom),dropgrad(grid)),ω)
	# ε⁻¹, epsi_pb = Zygote._pullback(Zygote.Context(),x->εₛ⁻¹(x,dropgrad(geom),dropgrad(grid)),ω)

	# ε, nng, ngvd = εₛ_nngₛ_ngvdₛ(ω,geom,grid)
	# nng⁻¹ = inv.(nng)
	# ε⁻¹ = inv.(ε)

	mag,m⃗,n⃗ = mag_m_n(k,grid)

	ei,ei_pb = Zygote.pullback(ω) do ω
		ε⁻¹,nng,nng⁻¹ = deepcopy(smooth(ω,p,(:fεs,:fnn̂gs,:fnn̂gs),[true,false,true],geom_fn,grid));
		return ε⁻¹
	end

	nngi,nngi_pb = Zygote.pullback(ω) do ω
		ε⁻¹,nng,nng⁻¹ = deepcopy(smooth(ω,p,(:fεs,:fnn̂gs,:fnn̂gs),[true,false,true],geom_fn,grid));
		return nng⁻¹
	end

	ngvd0,ngvd_pb = Zygote.pullback(ω) do ω
		# ngvd,nng2,nngi2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs,:fnn̂gs),[false,false,true],geom_fn,grid,volfrac_smoothing));
		ngvd,nng2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs),[false,false],geom_fn,grid,volfrac_smoothing));
		return ngvd
	end
	nng20,nng2_pb = Zygote.pullback(ω) do ω
		# ngvd,nng2,nngi2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs,:fnn̂gs),[false,false,true],geom_fn,grid,volfrac_smoothing));
		ngvd,nng2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs),[false,false],geom_fn,grid,volfrac_smoothing));
		return nng2
	end

	nngi2,nngi2_pb = Zygote.pullback(ω) do ω
		# ngvd,nng2,nngi2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs,:fnn̂gs),[false,false,true],geom_fn,grid,volfrac_smoothing));
		ngvd,nngi2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs),[false,true],geom_fn,grid,volfrac_smoothing));
		return nngi2
	end

	ε,ε⁻¹,nng,nng⁻¹ = deepcopy(smooth(ω,p,(:fεs,:fεs,:fnn̂gs,:fnn̂gs),[false,true,false,true],geom_fn,grid));
	ngvd,nng2 = deepcopy(smooth(ω,p,(:fnĝvds,:fnn̂gs),[false,false],geom_fn,grid,volfrac_smoothing));

	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	mag,m⃗,n⃗ = mag_m_n(k,grid)
	∂ω²∂k_nd = 2 * HMₖH(Hv,ε⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗)))
	k̄, H̄, nngī  = ∇HMₖH(k,Hv,nng⁻¹,grid; eigind)
	( _, _, om̄₁, eī₁ ) = ∇solve_k(	  (k̄,H̄),
									 	(k,Hv),
									  	∂ω²∂k_nd,
									   	ω,
									    ε⁻¹,
										grid; eigind)

	println("")
	println("\n manual calc.:")
	om̄₂ = dot(herm(nngī), ∂nng⁻¹_∂ω(ε⁻¹,nng,nng⁻¹,ngvd,ω))
	om̄₃ = dot(herm(eī₁), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	println("om̄₁: $(om̄₁)")
	println("om̄₂: $(om̄₂)")
	println("om̄₃: $(om̄₃)")
	om̄ = om̄₁ + om̄₂ + om̄₃
	println("om̄: $(om̄)")

	# k̄_nd, H̄_nd, nngī_nd  = ∇HMₖH(k,Hv,ε⁻¹,grid; eigind)
	# om̄₂_nd = dot(herm(nngī_nd), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω)) #dot(herm(nngī_nd), ∂nng⁻¹_∂ω(ε⁻¹,nng,nng⁻¹,ngvd,ω))
	# om̄₃ = dot(herm(eī₁), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	# println("om̄₁: $(om̄₁)")
	# println("om̄₂_nd: $(om̄₂_nd)")
	# println("om̄₃: $(om̄₃)")
	# om̄_nd = om̄₁ + om̄₂_nd + om̄₃
	# println("om̄_nd: $(om̄_nd)")

	nngis = copy(reinterpret(reshape,SMatrix{3,3,Float64,9},reshape(nng⁻¹,(9,128,128))))
	eis = copy(reinterpret(reshape,SMatrix{3,3,Float64,9},reshape(ε⁻¹,(9,128,128))))
	ngvds = copy(reinterpret(reshape,SMatrix{3,3,Float64,9},reshape(ngvd,(9,128,128))))
	dnngi_dom_s = -(nngis.^2 ) .* ( ω*(eis.*inv.(nngis).^2 .- inv.(nngis)) .+ ngvds)
	dnngi_dom_sr = copy(flat(dnngi_dom_s))
	om̄₂2 = dot(herm(nngī), dnngi_dom_sr)
	println("om̄₂2: $(om̄₂2)")

	#######

	# calculate and print neff = k/ω, ng = ∂k/∂ω, gvd = ∂²k/∂ω²
	Hₜ = reshape(Hv,(2,Ns...))
	mns = vcat(reshape(flat(m⃗),1,3,Ns...),reshape(flat(n⃗),1,3,Ns...))
	EE = 1im * ε⁻¹_dot( fft( kx_tc( Hₜ,mns,mag), (2:1+ND) ), ε⁻¹)
	HH = fft(tc(kx_ct( ifft( EE, (2:1+ND) ), mns,mag), mns),(2:1+ND) ) / ω
	EEs = copy(reinterpret(reshape,SVector{3,ComplexF64},EE))
	HHs = copy(reinterpret(reshape,SVector{3,ComplexF64},HH))
	Sz = dot.(cross.(conj.(EEs),HHs),(SVector(0.,0.,1.),))
	PP = 2*real(sum(Sz))
	WW = dot(EE,_dot((ε+nng),EE))
	ng = WW / PP

	∂ω²∂k_disp = 2 * HMₖH(Hv,nng⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗)))
	neff = k / ω
	# ng = 2 * ω / ∂ω²∂k_disp # HMₖH(H⃗,nng⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗))) # ng = ∂k/∂ω
	gvd = 2 / ∂ω²∂k_disp - ω * 4 / ∂ω²∂k_disp^2 * om̄ #( ng / ω ) * ( 1. - ( ng * om̄ ) )
	# println("∂ω²∂k_disp: $(∂ω²∂k_disp)")
	println("neff: $(neff)")
	println("ng: $(ng)")
	println("gvd: $(gvd)")

	println("")
	println("calc. with pullbacks:")
	# nngī2 = copy(reinterpret(SMatrix{3,3,T,9},copy(reshape( nngī , 9*Ns[1], Ns[2:end]...))))
	# nngī_herm = (real.(nngī2) .+ transpose.(real.(nngī2)) ) ./ 2
	# eī_herm = (real.(eī₁) .+ transpose.(real.(eī₁)) ) ./ 2
	om̄₂_pb = nngi_pb(herm(nngī))[1] #nngī2)
	om̄₃_pb = ei_pb(herm(eī₁))[1] #eī₁)
	println("om̄₁: $(om̄₁)")
	println("om̄₂_pb: $(om̄₂_pb)")
	println("om̄₃_pb: $(om̄₃_pb)")
	om̄_pb = om̄₁ + om̄₂_pb + om̄₃_pb
	println("om̄_pb: $(om̄_pb)")
	gvd_pb = 2 / ∂ω²∂k_disp - ω * 4 / ∂ω²∂k_disp^2 * om̄_pb #( ng / ω ) * ( 1. - ( ng * om̄ ) )
	println("gvd_pb: $(gvd_pb)")
	println("")
	return om̄
end

function neff_ng_gvd(ω,geom,k,H⃗,grid::Grid{ND,T};eigind=1,log=true) where {ND,T<:Real}
	ε, nng, ngvd = εₛ_nngₛ_ngvdₛ(ω,geom,grid)
	nng⁻¹ = inv.(nng)
	ε⁻¹ = inv.(ε)
	# calculate om̄ = ∂²ω²/∂k²
	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	mag,m⃗,n⃗ = mag_m_n(k,grid)
	∂ω²∂k_nd = 2 * HMₖH(H⃗,ε⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗)))
	k̄, H̄, nngī  = ∇HMₖH(k,H⃗,nng⁻¹,grid; eigind)
	( _, _, om̄₁, eī₁ ) = ∇solve_k(	  (k̄,H̄),
									 	(k,H⃗),
									  	∂ω²∂k_nd,
									   	ω,
									    ε⁻¹,
										grid; eigind)

    # nngī2 = copy(reinterpret(SMatrix{3,3,T,9},copy(reshape( nngī , 9*Ns[1], Ns[2:end]...))))
	# nngī_herm = (real.(nngī2) .+ transpose.(real.(nngī)) ) ./ 2
	# eī_herm = (real.(eī₁) .+ transpose.(real.(eī₁)) ) ./ 2
	# om̄₂ = dot(nngī_herm, ∂nng⁻¹_∂ω(ε⁻¹,nng,nng⁻¹,ngvd,ω))
	# om̄₃ = dot(eī_herm, ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	om̄₂ = dot(herm(nngī), ∂nng⁻¹_∂ω(ε⁻¹,nng,nng⁻¹,ngvd,ω))
	om̄₃ = dot(herm(eī₁), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	om̄ = om̄₁ + om̄₂ + om̄₃
	# calculate and return neff = k/ω, ng = ∂k/∂ω, gvd = ∂²k/∂ω²
	∂ω²∂k_disp = 2 * HMₖH(H⃗,nng⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗)))
	neff = k / ω
	ng = 2 * ω / ∂ω²∂k_disp # HMₖH(H⃗,nng⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗))) # ng = ∂k/∂ω
	gvd = 2 / ∂ω²∂k_disp - ω * 4 / ∂ω²∂k_disp^2 * om̄ #( ng / ω ) * ( 1. - ( ng * om̄ ) )
	return neff, ng, gvd
end

function neff_ng_gvd(ω,ε,ε⁻¹,nng,nng⁻¹,ngvd,k,Hv,grid::Grid{ND,T};eigind=1,log=true) where {ND,T<:Real}
	# ε, nng, ngvd = εₛ_nngₛ_ngvdₛ(ω,geom,grid)
	# nng⁻¹ = inv.(nng)
	# ε⁻¹ = inv.(ε)
	# calculate om̄ = ∂²ω²/∂k²
	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	mag,m⃗,n⃗ = mag_m_n(k,grid)
	m = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},Float64}(reinterpret(reshape,Float64,m⃗)))
	n = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},Float64}(reinterpret(reshape,Float64,n⃗)))
	∂ω²∂k_nd = 2 * HMₖH(Hv,ε⁻¹,mag,m,n)
	k̄, H̄, nngī  = ∇HMₖH(k,Hv,nng⁻¹,grid; eigind)
	( _, _, om̄₁, eī₁ ) = ∇solve_k(	  (k̄,H̄),
									 	(k,Hv),
									  	∂ω²∂k_nd,
									   	ω,
									    ε⁻¹,
										grid; eigind)
	# nngī2 = copy(reinterpret(SMatrix{3,3,T,9},copy(reshape( nngī , 9*Ns[1], Ns[2:end]...))))
	# nngī_herm = (real.(nngī2) .+ transpose.(real.(nngī)) ) ./ 2
	# eī_herm = (real.(eī₁) .+ transpose.(real.(eī₁)) ) ./ 2
	om̄₂ = dot(herm(nngī), ∂nng⁻¹_∂ω(ε⁻¹,nng,nng⁻¹,ngvd,ω))
	om̄₃ = dot(herm(eī₁), ∂ε⁻¹_∂ω(ε,ε⁻¹,nng,ω))
	om̄ = om̄₁ + om̄₂ + om̄₃
	# calculate and return neff = k/ω, ng = ∂k/∂ω, gvd = ∂²k/∂ω²
	∂ω²∂k_disp = 2 * HMₖH(Hv,nng⁻¹,mag,m,n)
	neff = k / ω
	# ng = 2 * ω / ∂ω²∂k_disp # HMₖH(H⃗,nng⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗))) # ng = ∂k/∂ω
	gvd = 2 / ∂ω²∂k_disp - ω * 4 / ∂ω²∂k_disp^2 * om̄ #( ng / ω ) * ( 1. - ( ng * om̄ ) )

	Hₜ = reshape(Hv,(2,Ns...))
	mns = vcat(reshape(flat(m⃗),1,3,Ns...),reshape(flat(n⃗),1,3,Ns...))
	EE = 1im * ε⁻¹_dot( fft( kx_tc( Hₜ,mns,mag), (2:1+ND) ), ε⁻¹)
	HH = inv(ω) * fft(tc(kx_ct( ifft( EE, (2:1+ND) ), mns,mag), mns),(2:1+ND) )
	EEs = copy(reinterpret(reshape,SVector{3,Complex{T}},EE))
	HHs = copy(reinterpret(reshape,SVector{3,Complex{T}},HH))
	# Sz = dot.(cross.(conj.(EEs),HHs),(SVector(0.,0.,1.),))
	Sz = getindex.(cross.(conj.(EEs),HHs),(3,))
	PP = 2*sum(Sz)
	# PP = 2*real( mapreduce((a,b)->dot(cross(conj(a),b),SVector(0.,0.,1.)),+,zip(EEs,HHs)))
	WW = dot(EE,_dot((ε+nng),EE))
	ng = real( WW / PP )

	return neff, ng, gvd
end

# function ∂²ω²∂k²(ω,ε⁻¹,nng⁻¹,k,H⃗,grid::Grid{ND,T};eigind=1,log=true) where {ND,T<:Real}
# 	ω² = ω^2
# 	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
# 	mag,m⃗,n⃗ = mag_m_n(k,grid)
# 	∂ω²∂k_nd = 2 * HMₖH(H⃗,ε⁻¹,real(mag),real(flat(m⃗)),real(flat(n⃗)))
# 	k̄, H̄, nngī  = ∇HMₖH(k,H⃗,nng⁻¹,grid; eigind)
# 	( _, _, om̄₁, eī₁ ) = ∇solve_k(	  (k̄,H̄),
# 									 	(k,H⃗),
# 									  	∂ω²∂k_nd,
# 									   	ω,
# 									    ε⁻¹,
# 										grid; eigind)
# 	nngī2 = copy(reinterpret(SMatrix{3,3,T,9},copy(reshape( nngī , 9*Ns[1], Ns[2:end]...))))
# 	nngī_herm = (nngī2 .+ adjoint.(nngī2) ) / 2
# 	eī_herm = (eī₁ .+ adjoint.(eī₁) ) / 2
# 	return om̄₁, eī_herm, nngī_herm
# end

function ∇M̂(k,ε⁻¹,λ⃗,H⃗,grid::Grid{ND,T}) where {ND,T<:Real}
	Nranges, Ninv, Ns, 𝓕, 𝓕⁻¹ = Zygote.ignore() do
		Ninv 		= 		1. / N(grid)
		Ns			=		size(grid)
		# g⃗s = g⃗(grid)
		Nranges		=		eachindex(grid)
		d0 = randn(Complex{T}, (3,Ns...))
		𝓕	 =	plan_fft(d0,_fftaxes(grid),flags=FFTW.PATIENT) # planned out-of-place FFT operator 𝓕
		𝓕⁻¹ =	plan_bfft(d0,_fftaxes(grid),flags=FFTW.PATIENT) # planned out-of-place iFFT operator 𝓕⁻¹
		return (Nranges,Ninv,Ns,𝓕,𝓕⁻¹)
	end
	mag, m⃗, n⃗  = mag_m_n(k,grid)
	# mns = vcat(reshape(flat(m⃗),(1,3,Ns...)),reshape(flat(n⃗),(1,3,Ns...)))
	m = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,m⃗)))
	n = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,n⃗)))
	mns = vcat(reshape(m,(1,3,Ns...)),reshape(n,(1,3,Ns...)))
	H = reshape(H⃗,(2,Ns...))
	λ = reshape(λ⃗,(2,Ns...))
	d 	= 	𝓕 * kx_tc( H , mns, mag ) * Ninv
	λd 	= 	𝓕 * kx_tc( λ , mns, mag )
	eī	 = 	 ε⁻¹_bar(vec(d), vec(λd), Ns...)
	eif = ε⁻¹ #flat(ε⁻¹)
	# eif = reshape(reinterpret(reshape,T,ε⁻¹),3,3,Ns...) #flat(ε⁻¹)
	# eif = reshape(reinterpret(T,ε⁻¹),3,3,Ns...)
	λẽ  =   𝓕⁻¹ * ε⁻¹_dot(λd * Ninv, real(eif)) #flat(ε⁻¹)) # _d2ẽ!(λẽ , λd  ,M̂ )
	ẽ 	 =   𝓕⁻¹ * ε⁻¹_dot(d        , real(eif)) #flat(ε⁻¹)) # _d2ẽ!(M̂.e,M̂.d,M̂)
	λẽ_sv  = reinterpret(reshape, SVector{3,Complex{T}}, λẽ )
	ẽ_sv 	= reinterpret(reshape, SVector{3,Complex{T}}, ẽ )
	m̄_kx = real.( λẽ_sv .* conj.(view(H,2,Nranges...)) .+ ẽ_sv .* conj.(view(λ,2,Nranges...)) )	#NB: m̄_kx and n̄_kx would actually
	n̄_kx =  -real.( λẽ_sv .* conj.(view(H,1,Nranges...)) .+ ẽ_sv .* conj.(view(λ,1,Nranges...)) )	# be these quantities mulitplied by mag, I do that later because māg is calc'd with m̄/mag & n̄/mag
	māg_kx = dot.(n⃗, n̄_kx) + dot.(m⃗, m̄_kx)
	k̄		= ∇ₖmag_m_n(
				māg_kx, 		# māg total
				m̄_kx.*mag, 	# m̄  total
				n̄_kx.*mag,	  	# n̄  total
				mag, m⃗, n⃗,
			)
	return k̄, eī
end

function rrule(::typeof(HelmholtzMap), kz::T, ε⁻¹, grid::Grid; shift=0.) where {T<:Real}
	function HelmholtzMap_pullback(M̄)
		if M̄ isa AbstractZero
			k̄	= ZeroTangent()
			eī = ZeroTangent()
		else
			λ⃗ = -M̄.maps[1].lmap
			H⃗ = M̄.maps[2].lmap'
			k̄, eī = ∇M̂(kz,ε⁻¹,λ⃗,H⃗,grid)
		end

		return (NoTangent(), k̄, eī, ZeroTangent())
	end
	return HelmholtzMap(kz, ε⁻¹, grid; shift), HelmholtzMap_pullback
end

# function rrule(T::Type{<:LinearMaps.LinearCombination{Complex{T1}}},As::Tuple{<:HelmholtzMap,<:LinearMaps.UniformScalingMap}) where T1<:Real
# 	function LinComb_Helmholtz_USM_pullback(M̄)
# 		# return (NoTangent(), (M̄, M̄))
# 		return (NoTangent(), Composite{Tuple{LinearMap,LinearMap}}(M̄, M̄))
# 	end
# 	return LinearMaps.LinearCombination{Complex{T1}}(As), LinComb_Helmholtz_USM_pullback
# end
#
# function rrule(T::Type{<:LinearMaps.UniformScalingMap},α::T1,N::Int) where T1
# 	function USM_pullback(M̄)
# 		# ᾱ = dot(M̄.maps[1].lmap/N, M̄.maps[2].lmap')
# 		ᾱ = mean( M̄.maps[1].lmap .* M̄.maps[2].lmap' )
# 		return (NoTangent(), ᾱ, ZeroTangent())
# 	end
# 	return LinearMaps.UniformScalingMap(α,N), USM_pullback
# end
#
# function rrule(T::Type{<:LinearMaps.UniformScalingMap},α::T1,N::Int,N2::Int) where T1
# 	function USM_pullback(M̄)
# 		# ᾱ = dot(M̄.maps[1].lmap/N, M̄.maps[2].lmap')
# 		ᾱ = mean( M̄.maps[1].lmap .* M̄.maps[2].lmap' )
# 		return (NoTangent(), ᾱ, ZeroTangent(), ZeroTangent())
# 	end
# 	return LinearMaps.UniformScalingMap(α,N,N2), USM_pullback
# end
#
# function rrule(T::Type{<:LinearMaps.UniformScalingMap},α::T1,Ns::Tuple{<:Int,<:Int}) where T1
# 	function USM_pullback(M̄)
# 		# ᾱ = dot(M̄.maps[1].lmap/first(Ns), M̄.maps[2].lmap')
# 		ᾱ = mean( M̄.maps[1].lmap .* M̄.maps[2].lmap' )
# 		return (NoTangent(), ᾱ, DoesNotExist())
# 	end
# 	return LinearMaps.UniformScalingMap(α,Ns), USM_pullback
# end

function ∇solve_k(ΔΩ, Ω, ∂ω²∂k, ω, ε⁻¹, grid::Grid{ND,T}; eigind=1, λ⃗₀=nothing) where {ND,T<:Real}
	k̄ₖ, H̄ = ΔΩ
	k, Hv = Ω
	# Ninv, Ns, 𝓕 = Zygote.ignore() do
	# 	Ninv 		= 		1. / N(grid)
	# 	Ns			=		size(grid)
	# 	d0 = randn(Complex{T}, (3,Ns...))
	# 	𝓕	 =	plan_fft(d0,_fftaxes(grid),flags=FFTW.PATIENT) # planned out-of-place FFT operator 𝓕
	# 	return (Ninv,Ns,𝓕)
	# end
	Ninv 		= 		1. / N(grid)
	Ns			=		size(grid)
	M̂ = HelmholtzMap(k,ε⁻¹,dropgrad(grid))
	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	Nranges = eachindex(grid) #(1:NN for NN in Ns) # 1:Nx, 1:Ny, 1:Nz for 3D, 1:Nx, 1:Ny for 2D
	g⃗s = g⃗(dropgrad(grid))
	mag,m⃗,n⃗ = mag_m_n(k,grid)
	# mn = vcat(reshape(flat(m⃗),(1,3,Ns...)),reshape(flat(n⃗),(1,3,Ns...)))
	m = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,m⃗)))
	n = real(HybridArray{Tuple{3,Dynamic(),Dynamic()},T}(reinterpret(reshape,T,n⃗)))
	mns = vcat(reshape(m,(1,3,Ns...)),reshape(n,(1,3,Ns...)))
	if !iszero(H̄)
		# solve_adj!(λ⃗,M̂,H̄,ω^2,Hv,eigind)
		λ⃗	= eig_adjt(
				M̂,								 # Â
				ω^2, 							# α
				Hv, 					 		 # x⃗
				0.0, 							# ᾱ
				H̄;								 # x̄
				λ⃗₀,
				# P̂	= HelmholtzPreconditioner(M̂),
			)
		k̄ₕ, eīₕ = ∇M̂(k,ε⁻¹,λ⃗,Hv,grid)
	else
		eīₕ 	= zero(ε⁻¹) #fill(SMatrix{3,3}(0.,0.,0.,0.,0.,0.,0.,0.,0.),size(ε⁻¹))
		k̄ₕ 	= 0.0
	end
	# combine k̄ₕ with k̄, scale by ( 2ω / ∂ω²∂k ) and calculate ω̄ and eīₖ
	# println("")
	# println("k̄ₖ = $(k̄ₖ)")
	# println("k̄ₕ = $(k̄ₕ)")
	# println("k̄ₖ + k̄ₕ = $(k̄ₖ+k̄ₕ)")
	λ⃗ₖ	 = ( (k̄ₖ + k̄ₕ ) / ∂ω²∂k ) * Hv
	H 	= reshape(Hv,(2,Ns...))
	λₖ  = reshape(λ⃗ₖ, (2,Ns...))
	# d	= 	𝓕 * kx_tc( H  , mns, mag ) * Ninv
	# λdₖ	=	𝓕 * kx_tc( λₖ , mns, mag )
	d	= 	fft(kx_tc( H  , mns, mag ),_fftaxes(grid)) * Ninv
	λdₖ	=	fft(kx_tc( λₖ , mns, mag ),_fftaxes(grid))
	eīₖ = ε⁻¹_bar(vec(d), vec(λdₖ), Ns...)
	ω̄  =  2ω * (k̄ₖ + k̄ₕ ) / ∂ω²∂k
	# if !(typeof(k)<:SVector)
	# 	k̄_kx = k̄_kx[3]
	# end
	# ms.ω̄  = 2ω * ( k̄_kx  / ms.∂ω²∂k ) # = 2ω * ω²̄
	return (NoTangent(), ZeroTangent(), ω̄  , eīₖ + eīₕ)
end

function ∇solve_k!(ΔΩ, Ω, ∂ω²∂k, ω::T, ε⁻¹, grid; eigind=1) where T<:Real
	k̄, H̄ = ΔΩ
	# println("k̄ = $(k̄)")
	k, Hv = Ω
	M̂ = HelmholtzMap(k,ε⁻¹,dropgrad(grid))
	Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	Nranges = eachindex(grid) #(1:NN for NN in Ns) # 1:Nx, 1:Ny, 1:Nz for 3D, 1:Nx, 1:Ny for 2D
	H = reshape(Hv,(2,Ns...))
	g⃗s = g⃗(dropgrad(grid))
	mag = M̂.mag
	n⃗ 	 = M̂.n⃗
	m⃗ 	= M̂.m⃗
	λd = similar(M̂.d)
	λẽ = similar(M̂.d)
	ẽ = similar(M̂.d)
	eīₕ = similar(ε⁻¹)
	eīₖ = similar(ε⁻¹)
	λ⃗ = similar(Hv)
	λ = reshape(λ⃗,(2,Ns...))
	if iszero(k̄)
		k̄ = 0.
	end
	if !iszero(H̄) #typeof(H̄) != ZeroTangent()
		solve_adj!(λ⃗,M̂,H̄,ω^2,Hv,eigind)
		λ⃗ -= dot(Hv,λ⃗) * Hv
		d = _H2d!(M̂.d, H * M̂.Ninv, M̂) # =  M̂.𝓕 * kx_tc( H , mn2, mag )  * M̂.Ninv
		λd = _H2d!(λd,λ,M̂) # M̂.𝓕 * kx_tc( reshape(λ⃗,(2,M̂.Nx,M̂.Ny,M̂.Nz)) , mn2, mag )
		ε⁻¹_bar!(eīₕ, vec(M̂.d), vec(λd), Ns...)
		# back-propagate gradients w.r.t. `(k⃗+g⃗)×` operator to k via (m⃗,n⃗) pol. basis and |k⃗+g⃗|
		λd *=  M̂.Ninv
		λẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(λẽ , λd  ,M̂ ) )
		ẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(M̂.e,M̂.d,M̂) )
		m̄_kx = real.( λẽ .* conj.(view(H,2,Nranges...)) .+ ẽ .* conj.(view(λ,2,Nranges...)) )	#NB: m̄_kx and n̄_kx would actually
		n̄_kx =  -real.( λẽ .* conj.(view(H,1,Nranges...)) .+ ẽ .* conj.(view(λ,1,Nranges...)) )	# be these quantities mulitplied by mag, I do that later because māg is calc'd with m̄/mag & n̄/mag
		māg_kx = dot.(n⃗, n̄_kx) + dot.(m⃗, m̄_kx)
		k̄ₕ	= -∇ₖmag_m_n(māg_kx, 		# māg total
						m̄_kx.*mag, 	# m̄  total
						n̄_kx.*mag,	  	# n̄  total
						mag, m⃗, n⃗ )
	else
		eīₕ = fill(SMatrix{3,3}(0.,0.,0.,0.,0.,0.,0.,0.,0.),size(ε⁻¹))
		k̄ₕ = 0.0
	end
	# combine k̄ₕ with k̄, scale by ( 2ω / ∂ω²∂k ) and calculate ω̄ and eīₖ
	copyto!(λ⃗, ( (k̄ + k̄ₕ ) / ∂ω²∂k ) * Hv )
	λ = reshape(λ⃗,(2,Ns...))
	d = _H2d!(M̂.d, H * M̂.Ninv, M̂) # =  M̂.𝓕 * kx_tc( H , mn2, mag )  * M̂.Ninv
	λd = _H2d!(λd,λ,M̂) # M̂.𝓕 * kx_tc( reshape(λ⃗,(2,M̂.Nx,M̂.Ny,M̂.Nz)) , mn2, mag )
	ε⁻¹_bar!(eīₖ, vec(M̂.d), vec(λd),Ns...)
	ω̄  =  2ω * (k̄ + k̄ₕ ) / ∂ω²∂k #[eigind]
	# if !(typeof(k)<:SVector)
	# 	k̄_kx = k̄_kx[3]
	# end
	# ms.ω̄  = 2ω * ( k̄_kx  / ms.∂ω²∂k[eigind] ) # = 2ω * ω²̄
	return (NoTangent(), ZeroTangent(), ω̄  , eīₖ + eīₕ)
end

# function rrule(::typeof(solve_k), ms::ModeSolver{ND,T},ω::T,ε⁻¹::AbstractArray{<:SMatrix{3,3},ND};
# 		nev=1,eigind=1,maxiter=3000,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real}
# 	kH⃗ = solve_k(ms,ω,ε⁻¹; nev, eigind, maxiter, tol, log, f_filter)
#     solve_k_pullback(ΔΩ) = let kH⃗=kH⃗, ∂ω²∂k=ms.∂ω²∂k[eigind], ω=ω, ε⁻¹=ε⁻¹, grid=ms.grid, eigind=eigind
# 		∇solve_k!(ΔΩ,kH⃗,∂ω²∂k,ω,ε⁻¹,grid;eigind)
# 	end
#     return (kH⃗, solve_k_pullback)
# end

# 	println("#########  ∂ω²/∂k Adjoint Problem for kz = $( M̂.k⃗[3] ) ###########")
# 	uplot(ch;name="log10( adj. prob. res. )")
# 	println("\t\t\tadj converged?: $ch")
# 	println("\t\t\titrs, mvps: $(ch.iters), $(ch.mvps)")

# function rrule(::typeof(solve_k), ω::T,geom::Vector{<:Shape},gr::Grid{ND};
# 		nev=1,eigind=1,maxiter=3000,tol=1e-8,log=false,ω²_tol=tol) where {ND,T<:Real}
#
# 	es = vcat(εs(geom,( 1. / ω )),[εᵥ,])		# dielectric tensors for each material, vacuum permittivity tensor appended
# 	eis = inv.(es)
#
# 	Srvol,proc_sinds,mat_inds = ignore() do
# 		xyz = x⃗(gr)			# (Nx × Ny × Nz) 3-Array of (x,y,z) vectors at pixel/voxel centers
# 		xyzc = x⃗c(gr)
# 		ms = ModeSolver(kguess(ω,geom), geom, gr))
# 		corner_sinds!(ms.corner_sinds,geom,xyz,xyzc))
# 		proc_sinds!(ms.sinds_proc,ms.corner_sinds))
# 		Srvol(x) = let psinds=ms.sinds_proc, xyz=xyz, vxlmin=vxl_min(xyzc), vxlmax=vxl_max(xyzc)
# 			S_rvol(sinds_proc,xyz,vxlmin,vxlmax,x)
# 		end
# 		eism(om,x) =
# 		(Srvol, ms.sinds_proc)
# 	end
# 	# Srvol = S_rvol(proc_sinds,xyz,vxl_min(xyzc),vxl_max(xyzc),shapes)
# 	ε⁻¹ = εₛ⁻¹(ω,geom;ms=dropgrad(ms))
# 	kH⃗ = solve_k(ms,ω,ε⁻¹; nev, eigind, maxiter, tol, log)
#     solve_k_pullback(ΔΩ) = let kH⃗=kH⃗, ∂ω²∂k=ms.∂ω²∂k, ω=ω, ε⁻¹=ε⁻¹, grid=ms.grid, eigind=eigind
# 		∇solve_k(ΔΩ,kH⃗,∂ω²∂k,ω,ε⁻¹,grid;eigind)
# 	end
#     return (kH⃗, solve_k_pullback)
# end

function rrule(::typeof(solve_k),ω::T,p::AbstractVector,geom_fn::F,grid::Grid{ND};
	nev=1,eigind=1,maxiter=300,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,F<:Function}
	ε⁻¹, ε⁻¹_pb = Zygote.pullback(ω,p) do ω,p
		# smooth(ω,p,:fεs,true,geom_fn,grid)
		smooth(ω,p,(:fεs,:fnn̂gs),[true,false,],geom_fn,grid)[1]
	end
	ms = ModeSolver(k_guess(ω,ε⁻¹), ε⁻¹, grid; nev, maxiter, tol)
	k, Hv = solve_k(ms, ω; nev, eigind, maxiter, tol, log, f_filter)
	g⃗ = copy(ms.M̂.g⃗)
	(mag, m⃗, n⃗), mag_m_n_pb = Zygote.pullback(k) do x
		mag_m_n(x,dropgrad(g⃗))
	end

	# Ns = copy(size(grid)) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	# Nranges = copy(eachindex(grid)) #(1:NN for NN in Ns) # 1:Nx, 1:Ny, 1:Nz for 3D, 1:Nx, 1:Ny for 2D
	# println("\tsolve_k:")
	# println("\t\tω² (target): $(ω^2)")
	# println("\t\tω² (soln): $(ms.ω²[eigind])")
	# println("\t\tΔω² (soln): $(real(ω^2 - ms.ω²[eigind]))")
	# println("\t\tk: $k")
	# println("\t\t∂ω²∂k: $∂ω²∂k")
	∂ω²∂k = copy(ms.∂ω²∂k[eigind])
	omsq_soln = copy(ms.ω²[eigind])
	# ε⁻¹_copy = copy(ε⁻¹)
	k = copy(k)
	Hv = copy(Hv)
    function solve_k_pullback(ΔΩ)
		k̄, H̄ = ΔΩ
		# println("\tsolve_k_pullback:")
		# println("k̄ (bar): $k̄")
		ms = ModeSolver(k, ε⁻¹, grid; nev, maxiter, tol)
		update_k!(ms,k)
		update_ε⁻¹(ms,ε⁻¹) #ε⁻¹)
		ms.ω²[eigind] = omsq_soln # ω^2
		ms.∂ω²∂k[eigind] = ∂ω²∂k
		copyto!(ms.H⃗[:,eigind], Hv)
		# replan_ffts!(ms)	# added  to check if this enables pmaps to work without crashing
		λ⃗ = similar(Hv)
		λd =  similar(ms.M̂.d)
		λẽ = similar(ms.M̂.d)
		Ns = size(grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
		Nranges = eachindex(grid)
		# ε⁻¹_bar = similar(ε⁻¹)
		# ∂ω²∂k = ms.∂ω²∂k[eigind] # copy(ms.∂ω²∂k[eigind])
		# Ns = size(ms.grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
		# Nranges = eachindex(ms.grid)

		H = reshape(Hv,(2,Ns...))
	    if iszero(k̄) # typeof(k̄)==ZeroTangent()
			k̄ = 0.
		end
		if !iszero(H̄) # if typeof(H̄) != ZeroTangent()
			# solve_adj!(ms,H̄,eigind) 												# overwrite ms.λ⃗ with soln to (M̂ + ω²I) λ⃗ = H̄ - dot(Hv,H̄)*Hv
			solve_adj!(λ⃗,ms.M̂,H̄,omsq_soln,Hv,eigind;log=false)
			# solve_adj!(ms,H̄,ω^2,Hv,eigind)
			λ⃗ -= dot(Hv,λ⃗) * Hv
			λ = reshape(λ⃗,(2,Ns...))
			d = _H2d!(ms.M̂.d, H * ms.M̂.Ninv, ms) # =  ms.M̂.𝓕 * kx_tc( H , mn2, mag )  * ms.M̂.Ninv
			λd = _H2d!(λd,λ,ms) # ms.M̂.𝓕 * kx_tc( reshape(λ⃗,(2,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)) , mn2, mag )
			# eīₕ = ε⁻¹_bar!(ε⁻¹_bar, vec(ms.M̂.d), vec(λd), Ns...)
			eīₕ = ε⁻¹_bar(vec(ms.M̂.d), vec(λd), Ns...)
			# eīₕ = copy(ε⁻¹_bar)
			# back-propagate gradients w.r.t. `(k⃗+g⃗)×` operator to k via (m⃗,n⃗) pol. basis and |k⃗+g⃗|
			λd *=  ms.M̂.Ninv
			λẽ_sv = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(λẽ , λd  ,ms ) )
			ẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(ms.M̂.e,ms.M̂.d,ms) )
			kx̄_m⃗ = real.( λẽ_sv .* conj.(view(H,2,Nranges...)) .+ ẽ .* conj.(view(λ,2,Nranges...)) )
			kx̄_n⃗ =  -real.( λẽ_sv .* conj.(view(H,1,Nranges...)) .+ ẽ .* conj.(view(λ,1,Nranges...)) )
			māg = dot.(n⃗, kx̄_n⃗) + dot.(m⃗, kx̄_m⃗)
			k̄ₕ = -mag_m_n_pb(( māg, kx̄_m⃗.*mag, kx̄_n⃗.*mag ))[1] # m̄ = kx̄_m⃗ .* mag, n̄ = kx̄_n⃗ .* mag, #NB: not sure why this is needs to be negated, inputs match original version
		else
			eīₕ = zero(ε⁻¹)#fill(SMatrix{3,3}(0.,0.,0.,0.,0.,0.,0.,0.,0.),size(ε⁻¹))
			k̄ₕ = 0.0
		end
		# combine k̄ₕ with k̄, scale by ( 2ω / ∂ω²∂k ) and calculate ω̄ and eīₖ
		copyto!(λ⃗, ( (k̄ + k̄ₕ ) / ∂ω²∂k ) * Hv )
		λ = reshape(λ⃗,(2,Ns...))
		d = _H2d!(ms.M̂.d, H * ms.M̂.Ninv, ms) # =  ms.M̂.𝓕 * kx_tc( H , mn2, mag )  * ms.M̂.Ninv
		λd = _H2d!(λd,λ,ms) # ms.M̂.𝓕 * kx_tc( reshape(λ⃗,(2,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)) , mn2, mag )
		# ε⁻¹_bar!(ε⁻¹_bar, vec(ms.M̂.d), vec(λd), Ns...)
		# eīₖ = copy(ε⁻¹_bar)
		eīₖ = ε⁻¹_bar(vec(ms.M̂.d), vec(λd), Ns...)
		# ε⁻¹_bar = eīₖ + eīₕ
		eibar = eīₖ + eīₕ
		ω̄_ε⁻¹, p̄ = ε⁻¹_pb(eibar)
		ω̄  =  ( 2ω * (k̄ + k̄ₕ ) / ∂ω²∂k ) + ω̄_ε⁻¹ #2ω * k̄ₖ / ms.∂ω²∂k[eigind]
		# if !(typeof(k)<:SVector)
		# 	k̄_kx = k̄_kx[3]
		# end
		# ms.ω̄  = 2ω * ( k̄_kx  / ms.∂ω²∂k[eigind] ) # = 2ω * ω²̄

		return (NoTangent(),  ω̄  , p̄, ZeroTangent(), ZeroTangent())
    end
    return ((k, Hv), solve_k_pullback)
end

# function rrule(::typeof(solve_k),ω::T,p::AbstractVector,geom_fn::F,grid::Grid{ND};
# 	nev=1,eigind=1,maxiter=300,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real,F<:Function}
function rrule(::typeof(solve_k), ms::ModeSolver{ND,T},ω::T,ε⁻¹::AbstractArray{T};
		nev=1,eigind=1,maxiter=300,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real}
		# ms::ModeSolver{ND,T},ω::T,ε⁻¹::AbstractArray{T};nev=1,eigind=1,maxiter=300,tol=1e-8,log=false,f_filter=nothing
	update_ε⁻¹(ms,ε⁻¹)
	k, Hv = solve_k(ms, ω; nev, eigind, maxiter, tol, log, f_filter)
	g⃗ = copy(ms.M̂.g⃗)
	(mag, m⃗, n⃗), mag_m_n_pb = Zygote.pullback(k) do x
		mag_m_n(x,dropgrad(g⃗))
	end

	Ns = size(ms.grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
	Nranges = eachindex(ms.grid) #(1:NN for NN in Ns) # 1:Nx, 1:Ny, 1:Nz for 3D, 1:Nx, 1:Ny for 2D
	# println("\tsolve_k:")
	# println("\t\tω² (target): $(ω^2)")
	# println("\t\tω² (soln): $(ms.ω²[eigind])")
	# println("\t\tΔω² (soln): $(real(ω^2 - ms.ω²[eigind]))")
	# println("\t\tk: $k")
	# println("\t\t∂ω²∂k: $∂ω²∂k")
	∂ω²∂k = copy(ms.∂ω²∂k[eigind])
	omsq_soln = copy(ms.ω²[eigind])
	# ε⁻¹_copy = copy(ε⁻¹)
	k = copy(k)
	Hv = copy(Hv)
    function solve_k_pullback(ΔΩ)
		k̄, H̄ = ΔΩ
		# println("\tsolve_k_pullback:")
		# println("k̄ (bar): $k̄")
		update_k!(ms,k)
		update_ε⁻¹(ms,ε⁻¹) #ε⁻¹)
		ms.ω²[eigind] = omsq_soln # ω^2
		ms.∂ω²∂k[eigind] = ∂ω²∂k
		copyto!(ms.H⃗, Hv)
		replan_ffts!(ms)	# added  to check if this enables pmaps to work without crashing
		λ⃗ = similar(Hv)
		λd =  similar(ms.M̂.d)
		λẽ = similar(ms.M̂.d)
		# ε⁻¹_bar = similar(ε⁻¹)
		# ∂ω²∂k = ms.∂ω²∂k[eigind] # copy(ms.∂ω²∂k[eigind])
		# Ns = size(ms.grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
		# Nranges = eachindex(ms.grid)

		H = reshape(Hv,(2,Ns...))
	    # if typeof(k̄)==ZeroTangent()
		if isa(k̄,AbstractZero)
			k̄ = 0.
		end
		# if typeof(H̄) != ZeroTangent()
		if !isa(H̄,AbstractZero)
			# solve_adj!(ms,H̄,eigind) 												# overwrite ms.λ⃗ with soln to (M̂ + ω²I) λ⃗ = H̄ - dot(Hv,H̄)*Hv
			solve_adj!(λ⃗,ms.M̂,H̄,omsq_soln,Hv,eigind;log=false)
			# solve_adj!(ms,H̄,ω^2,Hv,eigind)
			λ⃗ -= dot(Hv,λ⃗) * Hv
			λ = reshape(λ⃗,(2,Ns...))
			d = _H2d!(ms.M̂.d, H * ms.M̂.Ninv, ms) # =  ms.M̂.𝓕 * kx_tc( H , mn2, mag )  * ms.M̂.Ninv
			λd = _H2d!(λd,λ,ms) # ms.M̂.𝓕 * kx_tc( reshape(λ⃗,(2,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)) , mn2, mag )
			# eīₕ = ε⁻¹_bar!(ε⁻¹_bar, vec(ms.M̂.d), vec(λd), Ns...)
			eīₕ = ε⁻¹_bar(vec(ms.M̂.d), vec(λd), Ns...)
			# eīₕ = copy(ε⁻¹_bar)
			# back-propagate gradients w.r.t. `(k⃗+g⃗)×` operator to k via (m⃗,n⃗) pol. basis and |k⃗+g⃗|
			λd *=  ms.M̂.Ninv
			λẽ_sv = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(λẽ , λd  ,ms ) )
			ẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(ms.M̂.e,ms.M̂.d,ms) )
			kx̄_m⃗ = real.( λẽ_sv .* conj.(view(H,2,Nranges...)) .+ ẽ .* conj.(view(λ,2,Nranges...)) )
			kx̄_n⃗ =  -real.( λẽ_sv .* conj.(view(H,1,Nranges...)) .+ ẽ .* conj.(view(λ,1,Nranges...)) )
			māg = dot.(n⃗, kx̄_n⃗) + dot.(m⃗, kx̄_m⃗)
			k̄ₕ = -mag_m_n_pb(( māg, kx̄_m⃗.*mag, kx̄_n⃗.*mag ))[1] # m̄ = kx̄_m⃗ .* mag, n̄ = kx̄_n⃗ .* mag, #NB: not sure why this is needs to be negated, inputs match original version
		else
			eīₕ = zero(ε⁻¹)#fill(SMatrix{3,3}(0.,0.,0.,0.,0.,0.,0.,0.,0.),size(ε⁻¹))
			k̄ₕ = 0.0
		end
		# combine k̄ₕ with k̄, scale by ( 2ω / ∂ω²∂k ) and calculate ω̄ and eīₖ
		copyto!(λ⃗, ( (k̄ + k̄ₕ ) / ∂ω²∂k ) * Hv )
		λ = reshape(λ⃗,(2,Ns...))
		d = _H2d!(ms.M̂.d, H * ms.M̂.Ninv, ms) # =  ms.M̂.𝓕 * kx_tc( H , mn2, mag )  * ms.M̂.Ninv
		λd = _H2d!(λd,λ,ms) # ms.M̂.𝓕 * kx_tc( reshape(λ⃗,(2,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)) , mn2, mag )
		# ε⁻¹_bar!(ε⁻¹_bar, vec(ms.M̂.d), vec(λd), Ns...)
		# eīₖ = copy(ε⁻¹_bar)
		eīₖ = ε⁻¹_bar(vec(ms.M̂.d), vec(λd), Ns...)
		# ε⁻¹_bar = eīₖ + eīₕ
		eibar = eīₖ + eīₕ
		ω̄  =  ( 2ω * (k̄ + k̄ₕ ) / ∂ω²∂k )  #2ω * k̄ₖ / ms.∂ω²∂k[eigind]
		# if !(typeof(k)<:SVector)
		# 	k̄_kx = k̄_kx[3]
		# end
		# ms.ω̄  = 2ω * ( k̄_kx  / ms.∂ω²∂k[eigind] ) # = 2ω * ω²̄

		return (NoTangent(), ZeroTangent(), ω̄  , eibar)
    end
    return ((k, Hv), solve_k_pullback)
end

# function rrule(::typeof(solve_k), ms::ModeSolver{ND,T},ω::T,ε⁻¹::AbstractArray{<:SMatrix{3,3},ND};
# 		nev=1,eigind=1,maxiter=3000,tol=1e-8,log=false,f_filter=nothing) where {ND,T<:Real}
# 	k, Hv = solve_k(ms,ω,ε⁻¹; nev, eigind, maxiter, tol, log, f_filter)
# 	# k, Hv = copy.(solve_k(ms,ω,ε⁻¹; nev, eigind, maxiter, tol, log)) # ,ω²_tol)	 # returned data are refs to fields in ms struct. copy to preserve result for (possibly delayed) pullback closure.
# 	g⃗ = copy(ms.M̂.g⃗)
# 	(mag, m⃗, n⃗), mag_m_n_pb = Zygote.pullback(k) do x
# 		mag_m_n(x,dropgrad(g⃗))
# 	end
# 	∂ω²∂k = copy(ms.∂ω²∂k[eigind])
# 	Ns = size(ms.grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
# 	Nranges = eachindex(ms.grid) #(1:NN for NN in Ns) # 1:Nx, 1:Ny, 1:Nz for 3D, 1:Nx, 1:Ny for 2D
# 	# println("\tsolve_k:")
# 	# println("\t\tω² (target): $(ω^2)")
# 	# println("\t\tω² (soln): $(ms.ω²[eigind])")
# 	# println("\t\tΔω² (soln): $(real(ω^2 - ms.ω²[eigind]))")
# 	# println("\t\tk: $k")
# 	# println("\t\t∂ω²∂k: $∂ω²∂k")
# 	omsq_soln = ms.ω²[eigind]
# 	ε⁻¹_copy = copy(ε⁻¹)
# 	k_copy = copy(k)
# 	Hv = copy(Hv)
#     function solve_k_pullback(ΔΩ)
# 		k̄, H̄ = ΔΩ
# 		# println("\tsolve_k_pullback:")
# 		# println("k̄ (bar): $k̄")
# 		update_k!(ms,k_copy)
# 		update_ε⁻¹(ms,ε⁻¹_copy) #ε⁻¹)
# 		ms.ω²[eigind] = omsq_soln # ω^2
# 		ms.∂ω²∂k[eigind] = ∂ω²∂k
# 		copyto!(ms.H⃗, Hv)
# 		replan_ffts!(ms)	# added  to check if this enables pmaps to work without crashing
# 		# ∂ω²∂k = ms.∂ω²∂k[eigind] # copy(ms.∂ω²∂k[eigind])
# 		# Ns = size(ms.grid) # (Nx,Ny,Nz) for 3D or (Nx,Ny) for 2D
# 		# Nranges = eachindex(ms.grid)
#
# 		H = reshape(Hv,(2,Ns...))
# 	    if typeof(k̄)==ZeroTangent()
# 			k̄ = 0.
# 		end
# 		if typeof(H̄) != ZeroTangent()
# 			# solve_adj!(ms,H̄,eigind) 												# overwrite λ⃗ with soln to (M̂ + ω²I) λ⃗ = H̄ - dot(Hv,H̄)*Hv
# 			solve_adj!(ms.λ⃗,ms.M̂,H̄,omsq_soln,Hv,eigind;log=false)
# 			# solve_adj!(ms,H̄,ω^2,Hv,eigind)
# 			ms.λ⃗ -= dot(Hv,ms.λ⃗) * Hv
# 			λ = reshape(ms.λ⃗,(2,Ns...))
# 			d = _H2d!(ms.M̂.d, H * ms.M̂.Ninv, ms) # =  ms.M̂.𝓕 * kx_tc( H , mn2, mag )  * ms.M̂.Ninv
# 			λd = _H2d!(ms.λd,λ,ms) # ms.M̂.𝓕 * kx_tc( reshape(ms.λ⃗,(2,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)) , mn2, mag )
# 			ε⁻¹_bar!(ms.ε⁻¹_bar, vec(ms.M̂.d), vec(ms.λd), Ns...)
# 			eīₕ = copy(ms.ε⁻¹_bar)
# 			# back-propagate gradients w.r.t. `(k⃗+g⃗)×` operator to k via (m⃗,n⃗) pol. basis and |k⃗+g⃗|
# 			ms.λd *=  ms.M̂.Ninv
# 			λẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(ms.λẽ , ms.λd  ,ms ) )
# 			ẽ = reinterpret(reshape, SVector{3,Complex{T}}, _d2ẽ!(ms.M̂.e,ms.M̂.d,ms) )
# 			ms.kx̄_m⃗ .= real.( λẽ .* conj.(view(H,2,Nranges...)) .+ ẽ .* conj.(view(λ,2,Nranges...)) )
# 			ms.kx̄_n⃗ .=  -real.( λẽ .* conj.(view(H,1,Nranges...)) .+ ẽ .* conj.(view(λ,1,Nranges...)) )
# 			ms.māg .= dot.(n⃗, ms.kx̄_n⃗) + dot.(m⃗, ms.kx̄_m⃗)
# 			k̄ₕ = -mag_m_n_pb(( ms.māg, ms.kx̄_m⃗.*mag, ms.kx̄_n⃗.*mag ))[1] # m̄ = kx̄_m⃗ .* mag, n̄ = kx̄_n⃗ .* mag, #NB: not sure why this is needs to be negated, inputs match original version
# 		else
# 			eīₕ = fill(SMatrix{3,3}(0.,0.,0.,0.,0.,0.,0.,0.,0.),size(ε⁻¹))
# 			k̄ₕ = 0.0
# 		end
# 		# combine k̄ₕ with k̄, scale by ( 2ω / ∂ω²∂k ) and calculate ω̄ and eīₖ
# 		copyto!(ms.λ⃗, ( (k̄ + k̄ₕ ) / ∂ω²∂k ) * Hv )
# 		λ = reshape(ms.λ⃗,(2,Ns...))
# 		d = _H2d!(ms.M̂.d, H * ms.M̂.Ninv, ms) # =  ms.M̂.𝓕 * kx_tc( H , mn2, mag )  * ms.M̂.Ninv
# 		λd = _H2d!(ms.λd,λ,ms) # ms.M̂.𝓕 * kx_tc( reshape(ms.λ⃗,(2,ms.M̂.Nx,ms.M̂.Ny,ms.M̂.Nz)) , mn2, mag )
# 		ε⁻¹_bar!(ms.ε⁻¹_bar, vec(ms.M̂.d), vec(ms.λd),Ns...)
# 		eīₖ = copy(ms.ε⁻¹_bar)
# 		ω̄  =  2ω * (k̄ + k̄ₕ ) / ∂ω²∂k #2ω * k̄ₖ / ms.∂ω²∂k[eigind]
# 		ε⁻¹_bar = eīₖ + eīₕ
# 		# if !(typeof(k)<:SVector)
# 		# 	k̄_kx = k̄_kx[3]
# 		# end
# 		# ms.ω̄  = 2ω * ( k̄_kx  / ms.∂ω²∂k[eigind] ) # = 2ω * ω²̄
# 		return (NoTangent(), ZeroTangent(), ω̄  , ε⁻¹_bar)
#     end
#     return ((k, Hv), solve_k_pullback)
# end


text:
export VaxFloatF

struct VaxFloatF <: VaxFloat
    x::UInt32

    VaxFloatF(x::UInt32) = new(ltoh(x))
end

function VaxFloatF(x::T) where T <: Real
    ieeepart1 = reinterpret(UInt32, convert(Float32, x))
    e = reinterpret(Int32, ieeepart1 & IEEE_S_EXPONENT_MASK)

    if ieeepart1 & ~SIGN_BIT === zero(UInt32)
        # ±0.0 becomes 0.0
        return zero(VaxFloatF)
    elseif e === IEEE_S_EXPONENT_MASK
        # Vax types don't support ±Inf or NaN
        throw(InexactError(:VaxFloatF, VaxFloatF, x))
    else
        e >>>= VAX_F_MANTISSA_SIZE
        m = ieeepart1 & VAX_F_MANTISSA_MASK

        if e === zero(Int32)
            m <<= 1
            while m & VAX_F_HIDDEN_BIT === zero(UInt32)
                m <<= UNO
                e -= one(Int32)
            end
            m &= VAX_F_MANTISSA_MASK
        end

        e += one(Int32) + VAX_F_EXPONENT_BIAS - IEEE_S_EXPONENT_BIAS
        if e <= 0
            # Silent underflow
            return zero(VaxFloatF)
        elseif e > (2*VAX_F_EXPONENT_BIAS - 1)
            # Overflow
            throw(InexactError(:VaxFloatF, VaxFloatF, x))
        else
            vaxpart = (ieeepart1 & SIGN_BIT) | (e << VAX_F_MANTISSA_SIZE) | m
        end
    end

    vaxpart = htol(vaxpart)
    vaxpart1 = vaxpart & bmask16
    vaxpart2 = (vaxpart >>> 16) & bmask16
    vaxpart1 = (vaxpart1 << 16) | vaxpart2

    return VaxFloatF(vaxpart1)
end

Base.typemax(::Type{VaxFloatF}) = VaxFloatF(0xffff7fff)
Base.floatmax(::Type{VaxFloatF}) = VaxFloatF(0xffff7fff)
Base.typemin(::Type{VaxFloatF}) = VaxFloatF(0x00000080)
Base.floatmin(::Type{VaxFloatF}) = VaxFloatF(0x00000080)

Base.zero(::Type{VaxFloatF}) = VaxFloatF(0x00000000)
Base.one(::Type{VaxFloatF}) = VaxFloatF(0x00004080)

function Base.convert(::Type{Float32}, x::VaxFloatF)
    y = x.x
    vaxpart1 = y & bmask16
    vaxpart2 = (y >>> 16) & bmask16
    vaxpart1 = (vaxpart1 << 16) | vaxpart2

    e = reinterpret(Int32, vaxpart1 & VAX_F_EXPONENT_MASK)

    if e === zero(Int32)
        if vaxpart1 & SIGN_BIT === SIGN_BIT
            # Reserved floating-point reserved operand
            throw(InexactError(:convert, Float32, x))
        end

        # Dirty zero
        return zero(Float32)
    else
        e >>>= VAX_F_MANTISSA_SIZE

        e -= one(Int32) + VAX_F_EXPONENT_BIAS - IEEE_S_EXPONENT_BIAS
        if e > zero(Int32)
            out = vaxpart1 -
                (( UNO + VAX_F_EXPONENT_BIAS - IEEE_S_EXPONENT_BIAS ) <<
                    IEEE_S_MANTISSA_SIZE)
        else
            # out will be a subnormal
            out = (vaxpart1 & SIGN_BIT) |
                ((VAX_F_HIDDEN_BIT | (vaxpart1 & VAX_F_MANTISSA_MASK)) >>> (UNO - e))
        end
    end

    return reinterpret(Float32, out)
end

function Base.convert(::Type{T}, x::VaxFloatF) where T <: Union{Float16,Float64,BigFloat,Integer}
    return convert(T, convert(Float32, x))
end

ext: jl


text:
#
# Minimize total variation of a signal of R3 data.
#
# This example is part of Example 6.1 in the publication
#
# > R. Bergmann, R. Herzog, M. Silva Louzeiro, D. Tenbrinck, J. Vidal Núñez,
# > Fenchel Duality Theory and a Primal-Dual Algorithm on Riemannian Manifolds,
# > arXiv: [1908.02022](https://arxiv.org/abs/1908.02022)
#
using Manopt, Manifolds, LinearAlgebra

#
# Script Settings
experiment_name = "R3_Signal_TV_CP"
export_orig = true
export_primal = true
export_table = true
use_debug = true
#
# Automatic Script Settings
current_folder = @__DIR__
export_any = export_orig || export_primal || export_table
results_folder = joinpath(current_folder, "Signal_TV")
# Create folder if we have an export
(export_any && !isdir(results_folder)) && mkdir(results_folder)

#
# Example Settings
signal_section_size = 15
α = 5.0
σ = 0.5
τ = 0.5
θ = 1.0
γ = 0.0
max_iterations = 500
noise_level = 0.0
noise_type = :Gaussian

pixelM = Euclidean(3)
base = [1.0, 0.0, 0.0]
X = π / 4 * [0.0, 1.0, 0.0]
# Generate a signal with two sections - same signal as in S2_Signal_TV
p1 = exp(Sphere(2), base, X)
p2 = exp(Sphere(2), base, -X)
f = vcat(fill(p1, signal_section_size), fill(p2, signal_section_size))
#
# Compute exact minimizer
include("Signal_TV_commons.jl")
jump_height = distance(pixelM, f[signal_section_size], f[signal_section_size + 1])
δ = min(2 / (size(f, 1) * jump_height) * α, 1 / 2)
x_hat = shortest_geodesic(M, f, reverse(f), δ)
include("Ck.jl")

# add noise
if noise_level > 0
    f = [exp(pixelM, p, random_tangent(pixelM, p, noise_type, noise_level)) for p in f]
end
if export_orig
    orig_file = joinpath(results_folder, experiment_name * "-original.asy")
    asymptote_export_S2_data(orig_file; data=f)
    render_asymptote(orig_file)
end
#
# Initial values
m = fill(base, size(f))
n = Λ(m)
x0 = deepcopy(f)
ξ0 = ProductRepr(zero_tangent_vector(M, m), zero_tangent_vector(M, m))

storage = StoreOptionsAction((:x, :n, :ξbar))

@time o = ChambollePock(
    M,
    N,
    cost,
    x0,
    ξ0,
    m,
    n,
    prox_F,
    prox_G_dual,
    DΛ,
    adjoint_DΛ;
    primal_stepsize=σ,
    dual_stepsize=τ,
    relaxation=θ,
    acceleration=γ,
    relax=:dual,
    variant=:linearized,
    debug=if use_debug
        [
            :Iteration,
            " ",
            DebugPrimalChange(),
            " | ",
            DebugCk(storage),
            " | ",
            :Cost,
            "\n",
            100,
            :Stop,
        ]
    else
        missing
    end,
    record=if export_table
        [:Iteration, RecordPrimalChange(x0), RecordDualChange((ξ0, n)), :Cost, RecordCk()]
    else
        missing
    end,
    stopping_criterion=StopAfterIteration(max_iterations),
    return_options=true,
)
y = get_solver_result(o)
if has_record(o)
    r = get_record(o)
    Ck_values = [s[5] for s in r]
    println("The Ck Estimate lies between $(minimum(Ck_values)) and $(maximum(Ck_values))")
end
println("Distance from result to minimizer: ", distance(M, x_hat, y), "\n")

if export_primal
    orig_file = joinpath(results_folder, experiment_name * "-result.asy")
    asymptote_export_S2_data(orig_file; data=y)
    render_asymptote(orig_file)
end


text:
module Foo
"""
```jldoctest
julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8

julia> reshape(Int64[1, 2, 3, 4], (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 1
 2
 3
 4
```
```jldoctest
julia> reshape(Int64[1, 2, 3, 4], (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 1
 2
 3
 4

julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest
julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8

julia> reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest
julia> begin
           reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
       end
4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest
reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))

# output

4×1×1 Array{Int64,3}:
[:, :, 1] =
 2
 4
 6
 8
```
```jldoctest; filter = r"foo"
julia> println("  foobar")
  foobar
```
```jldoctest
julia> 1 + 2
3

julia> 3 + 4
7
```
"""
foo() = 1

    """
    ```jldoctest
    julia> begin
               reshape(Int64[1, 2, 3, 4] * 2, (4,1,1))
           end
    4×1×1 Array{Int64,3}:
    [:, :, 1] =
     2
     4
     6
     8
    ```
    ```jldoctest
    julia> println(); println("foo")

    foo
    ```
    """
    foo(x) = 1

end # module


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

using Markdown
using InteractiveUtils

# ╔═╡ 70b737a2-a9b2-11eb-1f6f-a330281aaaeb
using Omega

# ╔═╡ 305de14b-4732-4abf-ac08-c542b1bf0203
using CSV

# ╔═╡ f8ac57f8-680a-4246-b2a2-92714c38c167
using DataFrames

# ╔═╡ 4d5a4967-d8c6-4bf9-8c16-c1486d39b334
using Plots

# ╔═╡ 0d6a1f13-1b0c-46c0-8885-23530fd9a02c
using Query

# ╔═╡ e8195cc0-45c8-4edd-b3bc-425f0f902645
using VegaLite

# ╔═╡ 601b0f0d-5dab-43c2-9440-5254c802a3dc
using Distributions

# ╔═╡ 46fcaa08-ba11-41f4-aa65-9cc5ac9f44ec
md"# Inferring Degrees of Racial Descrimination through Bayesian Inference from Police Stopping Data"

# ╔═╡ 3f4880ba-65f2-4dd0-8eed-2d4f811fb8b0
md""" ## The Data
The data here is of police stops in San Diego
"""

# ╔═╡ 71b12af6-1886-4af9-bd3b-a7f06d89f83b
datapath = joinpath(dirname(pathof(Omega)), "..", "examples")

# ╔═╡ 847d6099-c655-4a92-a39a-ec7ae9982848
pd = DataFrame(CSV.File(joinpath(datapath, "sandiegopolicestops.csv")))

# ╔═╡ 471dea65-1eb1-415f-89bc-0612d7d88ed2
x = @from i in pd begin
    @where i.race_category == "Black"
    @select {i.block, i.land_mark}
    @collect DataFrame
end;

# ╔═╡ 1a4e1fb2-9c59-4cca-b925-ba06947f82ce
md"""# Stopping Test

The basic idea behind the threshold test is the following.

Suppose there is a disparity between the the rate at which members of minority groups are being stopped compared to non-minority groups.

Alone, this evidence is not sufficient to prove bias against the minority groups because there may be other causal factors that lead to the same conclusion.  For example, ethnic minorities tend to live in cities, which have a higher density of police offers.

One indication that disparities in stopping are indeed due to discrimination are if there are disparities in the rates at which the stops were justified.
A justified stop for example is one where an illegal weapon was found on the participant.

However the presence of a disparity is still insufficient to indicate bias.
This can be demonstrated through a counterexample: it's possible for a police offer to adopt a completely random (and hence indifferent to race) criteria by which they choose to stop someone, but there still to be differences in the justified stopping rates.
"""

# ╔═╡ e8d79dfc-13b3-4ae1-bc19-0d4eadd60cdc
md"""## A Bayesian Model

We will model an encounter between a police officer and a civillian.  The basic setup is as follows:

- A police-officer encounters a civillian
- Based on the civillian's race and the location, the officer infers a belief $p$ which represents the probability that the civillian is carrying a weapon
- The police-offer actually executes a search if $p$ exceeds some threshold $t_r$ which is race specific.

We use Bayesian inference to infer the thresholds for different races.  Any major disparity of $t_r$ between races indicates discrimination.
"""

# ╔═╡ 71ab4e3a-3888-480e-aa7e-82cc771f6c55
md"First we extract the race categories from the dataset"

# ╔═╡ cd2e8516-68f8-4c72-baf6-11cbbfd721ec
#race_categories = (@from i in pd begin
#    @select i.race_category
#	end) |> @unique() |> collect

# ╔═╡ 71b4adbc-4f84-4d9a-a158-4a3667881a4a
race_categories = ["Black", "White"]

# ╔═╡ ad52aaa5-fbb5-4747-9554-e536b81f1736
md"Let's add a little helper function to get the race index from the string"

# ╔═╡ 4a0aa2ef-654c-4bec-8d88-3ac5af3f0a94
race_id(race) = findfirst(x -> x == race, race_categories);

# ╔═╡ 8a174224-9e12-43aa-96be-fb0e3e951c2a
n_races = length(race_categories)

# ╔═╡ c92c75a7-2a72-4520-a03b-39116e7106ac
md"The following function `risk` constructs the parameters for the risk distribution as a function of race, as indicated by its id `i` (which spans from 1 to the total number of races `n_races`)"

# ╔═╡ a6678b03-f493-447e-83db-b63a03028744
risk(i) =
  (ϕ_r = (i, 1, 1) ~ Normal(0, 1),
   ϕ_d = (i, 2, 1) ~ Normal(0, 1),
   λ_r = (i, 3, 1) ~ Normal(0, 1),
   λ_d = (i, 4, 1) ~ Normal(0, 1))

# ╔═╡ ae577e52-8ced-46b9-a1ea-ee1098428de4
md"For each race category we assume there is an independent risk distribution"

# ╔═╡ 825b4d05-78f7-4619-a9c8-54f016dfd7f4
risks = map(risk, 1:n_races);

# ╔═╡ de8c1f27-6dc3-477f-b5e1-9e89d608bf92
md"In addition, for each race category there is a threshold value $t_r$, above which, the model suggests that a person will be searched"

# ╔═╡ d59e8d13-cce1-4ef2-b14c-27604fd42f3f
thresholds = [(i, 2) ~ Beta(1, 2) for i = 1:n_races];

# ╔═╡ 09fd183f-8677-4922-b86e-bf786b667bbe
invlogit(x) = (e = exp(x); e / (e + 1))

# ╔═╡ 5a63da51-65fe-4c8d-8821-78a4644c509a
md"Finally we can construct the stopping model.  Models in Omega are normal Julia functions which take as input ω (which we can think of as an random number generator object, or if we are more mathematically inclined as an element in probability space."

# ╔═╡ 1b73ba1c-3e7b-4594-a8fb-9e1835ebf3cc
function stop(id, ω)
	race_ = ω |> ((id, 3) ~ DiscreteUniform(1, n_races))

	# Get risk parameters
	riskp = risks[race_]

	# Construct beta distribution priors
	ϕ_rd = invlogit(riskp.ϕ_r(ω) + riskp.ϕ_d(ω))
	λ_rd = exp(riskp.λ_r(ω) + riskp.λ_d(ω))

	# Draw $p$ -- belief weapon-carrying from Beta
	p = ω |> ((id, 4) ~ Beta(ϕ_rd, λ_rd))

	# Draw threshold (as function of race)
	thresh = thresholds[race_](ω)

	# Officer stops if p exceeds threshold
	race_, Int(p > thresh)
end

# ╔═╡ 5afbcc6c-0012-4d10-b1ec-b194899396ba
n = 10

# ╔═╡ 3438de0c-692d-4376-b45e-85de49d95129
md"`stop` models a single data point.  To model multiple data points we next construct `data_size` (conditionally) independent draws from this model"

# ╔═╡ a251f256-3c95-46ac-be85-26aa08ca0eda
stopmany = Mv(1:n, stop)

# ╔═╡ 6bc7da07-7a60-4916-9d2f-3bf4d7838fcf
md"Here's a single sample from `stopmany`"

# ╔═╡ ecb05030-bb6e-40ea-aae8-718eb85a08b7
randsample(stopmany)

# ╔═╡ 92308112-eb7f-42bb-b303-e37dbbc4c0d4
md"We can show the unconditional distribution of the stops vs non-stops";

# ╔═╡ e659f421-19d5-43fe-a9bf-c845494efbe6
histogram(randsample(stopmany));

# ╔═╡ 599a49f2-0b86-4661-8bc8-a8473fe68241
md"""### Conditioning

Finally we can condition the model with data to infer the posterior distribution over $t$.

First, let's construct the actual proposition we want to condition on
"""

# ╔═╡ 1be486c7-cd84-4843-a388-aaadb39881ea
black_threshold = thresholds[race_id("Black")]

# ╔═╡ be0f16df-608f-43fa-accb-5af34c5608cb
d1 = @from i in pd begin
     @where i.search_conducted == 0
	 @where i.race_category in race_categories
     @select (race_id(i.race_category), i.search_conducted)
	 @collect
end;

# ╔═╡ 1b72954e-0a25-40a2-9638-5d6f75ae77bf
d2 = @from i in pd begin
     @where i.search_conducted == 1
     @where i.race_category in race_categories
     @select (race_id(i.race_category), i.search_conducted)
	 @collect
end;

# ╔═╡ 2d70f306-abda-4805-ab8d-0cde05decdc7
data_slice =
Tuple{Int, Int64}[d1[1:div(n,2)]; d2[1:div(n,2)]]

# ╔═╡ 2a8a54d0-5460-4679-96b7-fbde5c82cf68
data_condition(ω) = stopmany(ω) == data_slice;

# ╔═╡ c52be3dd-6123-418d-837f-b111ce908fdd
md"Construct the posterior distribution"

# ╔═╡ abaa6f63-c89b-4b54-afc3-27316352e9bc
black_threshold_posterior = black_threshold |ᶜ data_condition;

# ╔═╡ 0d8bef6f-92cb-449c-a77a-fd6650a9e586
md"Sample from the posterior distribution"

# ╔═╡ 7aa560da-08b9-4a61-96f0-b24ddb7a6da0
samples = randsample(black_threshold_posterior, 10000; alg = RejectionSample)

# ╔═╡ 4d10e6c7-0824-4258-822c-624e132cd309
histogram(samples)

# ╔═╡ Cell order:
# ╟─46fcaa08-ba11-41f4-aa65-9cc5ac9f44ec
# ╠═70b737a2-a9b2-11eb-1f6f-a330281aaaeb
# ╠═305de14b-4732-4abf-ac08-c542b1bf0203
# ╠═f8ac57f8-680a-4246-b2a2-92714c38c167
# ╠═4d5a4967-d8c6-4bf9-8c16-c1486d39b334
# ╠═0d6a1f13-1b0c-46c0-8885-23530fd9a02c
# ╠═e8195cc0-45c8-4edd-b3bc-425f0f902645
# ╠═601b0f0d-5dab-43c2-9440-5254c802a3dc
# ╟─3f4880ba-65f2-4dd0-8eed-2d4f811fb8b0
# ╟─71b12af6-1886-4af9-bd3b-a7f06d89f83b
# ╠═847d6099-c655-4a92-a39a-ec7ae9982848
# ╟─471dea65-1eb1-415f-89bc-0612d7d88ed2
# ╟─1a4e1fb2-9c59-4cca-b925-ba06947f82ce
# ╟─e8d79dfc-13b3-4ae1-bc19-0d4eadd60cdc
# ╟─71ab4e3a-3888-480e-aa7e-82cc771f6c55
# ╟─cd2e8516-68f8-4c72-baf6-11cbbfd721ec
# ╟─71b4adbc-4f84-4d9a-a158-4a3667881a4a
# ╟─ad52aaa5-fbb5-4747-9554-e536b81f1736
# ╠═4a0aa2ef-654c-4bec-8d88-3ac5af3f0a94
# ╠═8a174224-9e12-43aa-96be-fb0e3e951c2a
# ╟─c92c75a7-2a72-4520-a03b-39116e7106ac
# ╠═a6678b03-f493-447e-83db-b63a03028744
# ╟─ae577e52-8ced-46b9-a1ea-ee1098428de4
# ╠═825b4d05-78f7-4619-a9c8-54f016dfd7f4
# ╟─de8c1f27-6dc3-477f-b5e1-9e89d608bf92
# ╠═d59e8d13-cce1-4ef2-b14c-27604fd42f3f
# ╠═09fd183f-8677-4922-b86e-bf786b667bbe
# ╟─5a63da51-65fe-4c8d-8821-78a4644c509a
# ╠═1b73ba1c-3e7b-4594-a8fb-9e1835ebf3cc
# ╠═5afbcc6c-0012-4d10-b1ec-b194899396ba
# ╟─3438de0c-692d-4376-b45e-85de49d95129
# ╠═a251f256-3c95-46ac-be85-26aa08ca0eda
# ╟─6bc7da07-7a60-4916-9d2f-3bf4d7838fcf
# ╠═ecb05030-bb6e-40ea-aae8-718eb85a08b7
# ╟─92308112-eb7f-42bb-b303-e37dbbc4c0d4
# ╟─e659f421-19d5-43fe-a9bf-c845494efbe6
# ╟─599a49f2-0b86-4661-8bc8-a8473fe68241
# ╠═1be486c7-cd84-4843-a388-aaadb39881ea
# ╟─be0f16df-608f-43fa-accb-5af34c5608cb
# ╟─1b72954e-0a25-40a2-9638-5d6f75ae77bf
# ╟─2d70f306-abda-4805-ab8d-0cde05decdc7
# ╠═2a8a54d0-5460-4679-96b7-fbde5c82cf68
# ╟─c52be3dd-6123-418d-837f-b111ce908fdd
# ╠═abaa6f63-c89b-4b54-afc3-27316352e9bc
# ╟─0d8bef6f-92cb-449c-a77a-fd6650a9e586
# ╠═7aa560da-08b9-4a61-96f0-b24ddb7a6da0
# ╠═4d10e6c7-0824-4258-822c-624e132cd309

max_stars_repo_name: zenna/expect

max_issues_repo_name: zenna/expect

max_forks_repo_name: zenna/expect


text:
@testset "quantile and pearson residuals" begin
    n = 10^3
    lags = 1
    normality_quantile_and_pearson_residuals(Normal, n, lags)
    normality_quantile_and_pearson_residuals(LogNormal, n, lags)
    normality_quantile_and_pearson_residuals(Gamma, n, lags)
end


text:
######################################################################
# Packets used in Pia peer-to-peer communication.
# -----
# Licensed under Apache License 2.0
module Packets
import ExtraFun: Ident

# ===== Packet Prefixes =====
# Always write fixed enum values for backward compatibility! Never delete packet IDs either, even when deprecated.
# This way, order of packet IDs does not matter and we are free to rearrange and insert packets as needed.

export PacketClass
@enum PacketClass begin
    ConnectivityClass = 0x01
    CryptographyClass = 0x02
    DiscoveryClass    = 0x03
    DataClass         = 0x04
    UtilityClass      = 0x10
end

@enum ConnectivityPrefix begin
    HandshakeStart    = 0x01
    HandshakeResponse = 0x02
    HandshakeConfirm  = 0x03
    Ping              = 0x04
    Pong              = 0x05
    WantPromote       = 0x06
    AcceptPromote     = 0x07
    Demote            = 0x08
    Goodbye           = 0xFF
end

@enum DiscoveryPrefix begin
    QueryPeer = 0x01
end

@enum UtilityPrefix begin
    Wakeup = 0x01
end

@enum DataPrefix begin
    AppData = 0x01
end

classid(::ConnectivityPrefix) = ConnectivityClass
classid(::DiscoveryPrefix)    = DiscoveryClass
classid(::UtilityPrefix)      = UtilityClass
classid(::DataPrefix)         = DataClass
packetid(cls, pktid::UInt8) = packetid(Ident{cls}(), pktid)
packetid(::Ident{ConnectivityClass}, pktid::UInt8) = ConnectivityPrefix(pktid)
packetid(::Ident{DiscoveryClass},    pktid::UInt8) = DiscoveryPrefix(pktid)
packetid(::Ident{UtilityClass},      pktid::UInt8) = UtilityPrefix(pktid)
packetid(::Ident{DataClass},         pktid::UInt8) = DataPrefix(pktid)

end # module Packets