export GainObject
export assign_pumpstr!

# GainObject has
# - Shape
# - scale factor to be multiplied to the true/false array.
#
# Later, to put an object with a spatially changing pump strength, I could define a new
# GainObject.  Different types of GainObjects will need to define different assignment
# function.  Eventually, I will need to write something like MaxwellFDFD's assign function
# such that it internally uses the individual shape's assign function.  But let's not think
# about that right now.
#
# This may require defining GainObject as an abstract type later, and defining its concrete
# types.

# See Julia issue #26321 for using S<:Shape{K}.
mutable struct GainObject{K,Kₑ,S<:Shape{K},F}  # Kₑ is number of Cartesian components of E-field
    shape::S
    D₀fun::F  # population inversion in absence of lasing; function of pump parameter d
    GainObject{K,Kₑ,S,F}(shape, D₀fun) where {K,Kₑ,S<:Shape{K},F} = new(shape, D₀fun)
end

GainObject{K,Kₑ}(shape::S, D₀fun::F) where {K,Kₑ,S<:Shape{K},F} = GainObject{K,Kₑ,S,F}(shape, D₀fun)
GainObject{K,Kₑ}(shape::Shape{K}) where {K,Kₑ} = GainObject{K,Kₑ}(shape, d::Real->d)

add_gainobj!(gobj_vec::AbsVec{<:GainObject}, shapes::Shape...; D₀fun::Any=identity) = add_gainobj!(gobj_vec, [shapes...]; D₀fun)
function add_gainobj!(gobj_vec::AbsVec{<:GainObject{K,Kₑ}}, shapes::AbsVec{<:Shape{K}}; D₀fun::Any=identity) where {K,Kₑ} # identity: d::Real->d
    for shape = shapes  # shapes is tuple
        gobj = GainObject{K,Kₑ}(shape, D₀fun)
        push!(gobj_vec, gobj)
    end

    return nothing
end

# Assign the pump strength specified by gain objects to an array.
# The implementation depends on MaxwellFDFD, so this function is moved from SALTBase to
# MaxwellSALT.
function assign_pumpstr!(D₀::AbsVecReal,  # output vector to write on
                         gobj_vec::AbsVec{<:GainObject{K,Kₑ}},  # vector of gain objects; later objects overwrite earlier objects
                         d::Real,  # pump parameter at which D₀ is evaluated
                         gt₀::SVec{K,GridType},  # grid type of E-field voxel corners; generated by ft2gt.(EE, boundft)
                         l::Tuple2{NTuple{K,AbsVecReal}},  # l[PRIM][k], l[DUAL][k]: primal, dual vertex locations in k-direction
                         ise˔shp::Bool,  # true if E-field dimension is orthogonal to shape dimension; false otherwise
                         order_cmpfirst::Bool  #  if true, D₀ is indexed D₀[w,x,y,z] for K = 3; otherwise indexed D₀[x,y,z,w]
                         ) where {K,Kₑ}
    N = SVec(length.(l[nPR]))  # SInt{K}
    D₀ .= 0  # initialize (pump strength is nonzero only inside gain media)
    D₀array = order_cmpfirst ? reshape(D₀, Kₑ, N.data...) : reshape(D₀, N.data..., Kₑ)  # related: MaxwellBase/field.jl/field_arr2vec().
    for nw = 1:Kₑ
        gt_cmp = ise˔shp ? gt₀ : gt_w(nw, gt₀)
        lcmp = t_ind(l, gt_cmp)
        D₀cmp = order_cmpfirst ? selectdim(D₀array, 1, nw) :  # D₀array[nw,:,:,:] for K = 3 (note 1 ≤ nw ≤ Kₑ)
                                 selectdim(D₀array, K+1, nw)  # D₀array[:,:,:,nw] for K = 3 (note 1 ≤ nw ≤ Kₑ)
        for gobj = gobj_vec
            shape = gobj.shape
            D₀fun = gobj.D₀fun
            D₀val = D₀fun(d)
            assign_val_shape!(D₀cmp, D₀val, shape, lcmp)  # exported from MaxwellBase
        end
    end

    return nothing
end
