# Below, there are two ways to define a function with the same name as the one defined in
# another package (e.g., MaxwellFDFD): extending the function defined in the other package,
# or calling the function defined in the other package by qualifying it with the package
# name within the definition of the function in the present package.
#
# For example, when defining a function that delegates to a function fun_XXX!() in
# MaxwellFDFD, we would want to use the same function name.  This can be done as either
#     MaxwellFDFD.fun_XXX!(mdl::Model, ...) = fun_XXX!(mdl.mx, ...)
# or
#     fun_XXX!(mdl::Model, ...) = MaxwellFDFD.fun_XXX!(mdl.mx, ...)
#
# Both definitions works.  If we reexport MaxwellFDFD from the present package, we do not
# even need to export the function, because the function name is already exported by
# MaxwellFDFD.  (If MaxwellFDFD were not reexported, we would need to export fun_XXX!
# explicitly in this file.)
#
# Which one is the right approach?
#
# If we are extending a function that is already exported by the present package (e.g.,
# using reexport), I think the former is better, because it explicitly specifies that we are
# extending the capabilities of an already exported function.
#
# On the other hand, if we are "extending" a function that is exported by another package
# (e.g., SALTBase) that is not reexported by the present package, then we need to export the
# function anyway to make it available to the users.  Then, we are actually not extending
# the capability of some function that is already available to the users.  Therefore, in
# that case I think using the latter makes more sense.
#
# Currently, I am reexporting MaxwellFDFD from the present package, but not SALTBase.  (The
# users may want to solve a driven Maxwell's equations, so I reexport MaxwellFDFD, but they
# do not need to know the internals of SALTBase, so I do not reexport SALTBase.)  Therefore,
# when redefining functions already defined in MaxwellFDFD and SALTBase below, I extend
# MaxwellFDFD's functions (the former approach), but call SALTBase's functions with
# qualification (the latter approach).

export GainProfile
export get_εcold, get_dblcurl, add_gainobj!
export init_salt!, simulate!

# - Do not export Model; quality it with the package name MaxwellSALT, because I would have
# similar types in other packages such as MaxwellFDFD and MaxwellGuide.
# - It might be possible to define mx for other Maxwell solvers (e.g., COMSOL).  Then MX
# doesn't have to be <:MaxwellFDFD.Model.  Let's consider the possibility of this
# generalization later.
mutable struct Model{K,Kₑ,VF,VC,GP<:GainProfile{VF},MX<:MaxwellFDFD.Model{K,Kₑ},MXS<:AbstractMaxwellSolver}
    # Model components
    mx::MX  # MaxwellFDFD's model
    gp::GP  # gain profile
    gobj_vec::Vector{GainObject{K,Kₑ}}  # gain objects

    # Solver
    mxsolver::MXS

    # Initial guess
    ωguess::VecComplexF
    Ψguess::Vector{VC}

    # Solution placeholders
    nlsol::NonlasingSol{VC}
    nlvar::NonlasingVar{MXS,VC}
    lsol::LasingSol{VF,VC}
    lvar::LasingVar{MXS,VF,VC}

    function Model{K,Kₑ,VF,VC,GP,MX}(mx::MaxwellFDFD.Model, gp::GainProfile) where {
            K,Kₑ,VF,VC,GP<:GainProfile{VF},MX<:MaxwellFDFD.Model{K,Kₑ}}
        Ps = create_paramops(mx)
        Cs = create_curls(mx)
        CC = create_A(EE, 0.0, Ps, Cs)  # ω = 0.0

        Me2c, _ = create_Mcs(mx)

        # Construct Mc2e.
        grid = mx.grid
        boundft = mx.boundft
        e⁻ⁱᵏᴸ = create_e⁻ⁱᵏᴸ(mx)

        N = grid.N
        isbloch = grid.isbloch

        # Arguments of create_mean:
        # - isfwd is set to average the fields at the voxel corners to get the fields at the
        # solution field locations (but with polarization not meant for that locations).
        # (For boundft[w] = ft, voxel-corner field is forward-averaged along w.)
        # - ∆l and ∆l′⁻¹ are not supplied; use unweighted arithmetic averaging.
        isfwdₑ = boundft.==EE
        Mₑ = create_mean(isfwdₑ, N, isbloch, e⁻ⁱᵏᴸ; mx.order_cmpfirst)

        permute = SVec(ntuple(identity,Val(Kₑ)))
        Mc2e = Vector{typeof(CC)}(undef, Kₑ-1)
        for k = 1:Kₑ-1
            permute = circshift(permute, 1)  # currently return MVector; see StaticArrays.jl/issues/745
            P = create_πcmp(N, permute; mx.order_cmpfirst)
            Mc2e[k] = Mₑ * P
        end

        mxsolver = DirectMaxwellSolver(CC, Me2c, tuple(Mc2e...))

        gobj_vec = GainObject{K,Kₑ}[]

        return new{K,Kₑ,VF,VC,GP,MX,typeof(mxsolver)}(mx, gp, gobj_vec, mxsolver)
    end
end

function Model(mx::MX, gp::GP; Vtype::Type=Vector) where {K,Kₑ,GP<:GainProfile,MX<:MaxwellFDFD.Model{K,Kₑ}}
    return Model{K,Kₑ,Vtype{Float},Vtype{ComplexF},GP,MX}(mx, gp)
end


# Operators needed for calculating the initial guess solution
get_dblcurl(mdl::Model) = mdl.mxsolver.CC

function get_εcold(mdl::Model)
    Pε, _ = create_paramops(mdl.mx)

    return diag(Pε)
end

function init_salt!(mdl::Model{K,Kₑ}, ωguess::AbsVecNumber, Ψguess::AbsMatComplexF) where {K,Kₑ}
    mxsolver = mdl.mxsolver
    gp = mdl.gp

    grid = mdl.mx.grid
    N = Kₑ * prod(grid.N)
    M = length(ωguess)

    mdl.nlsol = NonlasingSol(ωguess, Ψguess)
    mdl.nlvar = NonlasingVar(mxsolver, N, M)
    mdl.lsol = LasingSol(N, M)
    mdl.lvar = LasingVar(mxsolver, N, length(gp), M)

    return nothing
end

#= Gain objects =#
add_gainobj!(mdl::Model, shapes::Shape...; kwargs...) = add_gainobj!(mdl.gobj_vec, shapes...; kwargs...)
add_gainobj!(mdl::Model, shapes::AbsVec{<:Shape}; kwargs...) = add_gainobj!(mdl.gobj_vec, shapes; kwargs...)

function simulate!(mdl::Model{K,Kₑ,VF},
                   dvec::AbsVecReal;  # trajectory of pump strength parameter to follow
                   outωaψ::NTuple{3,Bool}=(true,true,false),  # true to output ω, a, ψ
                   doutvec::AbsVecReal=dvec,  # output results when dvec[i] ∈ doutvec
                   kwargs...  # keyword arguments of SALTBase > simulation.jl/simulate!()
                   ) where {K,Kₑ,VF}
    mx = mdl.mx
    grid, boundft, ise˔shp, order_cmpfirst = getproperty.(Ref(mx), (:grid, :boundft, :ise˔shp, :order_cmpfirst))

    gt₀ = ft2gt.(EE, boundft)
    setD₀!(D₀::VF, d::Real) = assign_pumpstr!(D₀, mdl.gobj_vec, d, gt₀, grid.l, ise˔shp, order_cmpfirst)

    gp, lsol, lvar, nlsol, nlvar = getproperty.(Ref(mdl), (:gp, :lsol, :lvar, :nlsol, :nlvar))

    # Return, as functions of d, the followings:
    # eigenfrequencies, modal amplitudes, normalized mode profiles, numbers of AA steps, times taken for convergence of AA
    ωout, aout, ψout, nAA, tAA = SALTBase.simulate!(lsol, lvar, nlsol, nlvar, gp, get_εcold(mdl),
                                                    dvec, setD₀!; outωaψ, doutvec, kwargs...)

    return ωout, aout, ψout, nAA, tAA
end
