export DirectMaxwellSolver

# To do:
#
# We can create a similar type that reduces the number of factorization.  During the Newton
# iteration for solving the nonlinear equation, the matrix changes only diagonally.
# Furthermore, this diagonal change is parametrized by ω.  Therefore, we can perform
# factorization only once at a single ω, and use the idea of fast frequency sweep for other
# ω's.  Because γ(ω) has a pole, we may need to use the Padé approximation instead of the
# simple power series expansion, for which γ(ω) will generate infinitely many terms because
# γ(ω) is infinitely differentiable.
#
# A simiar idea cannot be applied to solving the lasing equation, where the diagonal change
# involves a change in the hole-burning term and is not parametrized by the scalar ω.
# However, we may construct a new parametrized matrix for each hole-burning term, by fixing
# the hole burning term and changing only ω.  The downside of this is that the coefficients
# of the series expansion of the solution needs to be calculated for every ω.  (For the
# nonlasing equation, this needed to be calculated only once for varying ω during the Newton
# iteration...  Is this right?  This is correct when both A and b are parametrized by ω, but
# in our case b is probably not, so this may not be correct.)  However, if the major
# performance improvement comes from the reduced number of factorizations rather than the
# reduced number of linear solves, this scheme should still have a huge benefit.
#
# However, all these ideas become irrelevant for 3D problems, where linear systems will be
# solved iteratively and therefore no factorization will be performed.  Hence, implementing
# the above ideas is not a priority.

mutable struct DirectMaxwellSolver{VR<:AbsVecFloat,VC<:AbsVecComplexF,MC<:AbsMatComplexF,FC<:Factorization{ComplexF},Kₑ₋₁} <: AbstractMaxwellSolver
    CC::MC  # double curl operator ∇ × μ⁻¹ ∇ ×
    A::MC  # Maxwell operator
    Me2c::MC  # operator interpolating edge Ew for w = 1, ..., Kₑ at voxel corners
    Mc2e::NTuple{Kₑ₋₁,MC}  # operator interpolating voxel-corner Ew for w = 1, ..., Kₑ at Ev for v that is circularly shifted 1, ..., Kₑ
    vtemp::VC  # temporary storage for complex vector
    wtemp::VC  # temporary storage for complex vector
    fact::FC  # matrix factors of A

    function DirectMaxwellSolver{VR,VC,MC,FC,Kₑ₋₁}(CC::AbsMatNumber, A::AbsMatNumber,
                                                 Me2c::AbsMatNumber, Mc2e::NTuple{Kₑ₋₁,AbsMatNumber},
                                                 vtemp::AbsVecNumber, wtemp::AbsVecNumber
                                                 ) where {VR<:AbsVecFloat,VC<:AbsVecComplexF,MC<:AbsMatComplexF,FC<:Factorization{ComplexF},Kₑ₋₁}
        isequal(size(CC)...) || @error "CC should be square, but size(CC) = $(size(CC))."

        size(CC)==size(A)==size(Me2c) || @error "size.((CC,A,Me2c)) = $(size.((CC,A,Me2c))) should be all same."

        for k = 1:Kₑ₋₁
            size(Mc2e[k])==size(CC) || @error "size(Mc2e[$k]) = $(size(Mc2e[k])) and size(CC) = $(size(CC)) should be the same."
        end

        size(CC,1)==length(vtemp)==length(wtemp) ||
            @error "size.((CC,vtemp,wtemp),1) = $(size.((CC,vtemp,wtemp),1)) should be all same."

        return new(CC, A, Me2c, Mc2e, vtemp, wtemp)
    end
end

# Below, CC is shared between modes, but A is not (it is overwritten for new modes).
# Therefore we need to keep a separate copy of A.
function DirectMaxwellSolver(CC::MC, Me2c::MC, Mc2e::Tuple{Vararg{MC}}; Vtype::Type=Vector) where {MC<:AbsMatComplexF}
    Kₑ₋₁ = length(Mc2e)
    FC = Base.return_types(lu, tuple(MC))[1];
    n = size(CC, 2)

    return DirectMaxwellSolver{Vtype{Float},Vtype{ComplexF},MC,FC,Kₑ₋₁}(CC, similar(CC), Me2c, Mc2e, VecComplexF(undef,n), VecComplexF(undef,n))
end

Base.similar(mxsolver::DirectMaxwellSolver) = DirectMaxwellSolver(mxsolver.CC, mxsolver.Me2c, mxsolver.Mc2e)
Base.size(mxsolver::DirectMaxwellSolver) = size(mxsolver.CC)

# Factorize the matrix A and cache it.
function SALTBase.init_mxsolver!(mxsolver::DirectMaxwellSolver, ω::Number, ε::AbsVecNumber)
    mA = size(mxsolver.A, 1)
    nε = length(ε)
    mA==nε || @error "size(mxsolver.A,1) = $mA and length(ε) = $nε should be same."

    mxsolver.A .= mxsolver.CC  # initialize; works for sparse matrices with same nonzero entry pattern
    for i = 1:nε
        mxsolver.A[i,i] -= ω^2 * ε[i]
    end

    t = @elapsed mxsolver.fact = lu(mxsolver.A)
    # @info "time factorization: $t"

    return nothing
end

function SALTBase.mxsolve!(x::VC, mxsolver::DirectMaxwellSolver{VR,VC}, b::VC) where {VR,VC}
    t = @elapsed ldiv!(x, mxsolver.fact, b)
    # @info "time mxsolve!: $t"

    return nothing
end

function SALTBase.mxsolve_transpose!(x::VC, mxsolver::DirectMaxwellSolver{VR,VC}, b::VC) where {VR,VC}
    t = @elapsed ldiv!(x, transpose(mxsolver.fact), b)
    # @info "time mxsolve_transpose!: $t"

    return nothing
end

function SALTBase.mxapply!(b::VC, mxsolver::DirectMaxwellSolver{VR,VC}, x::VC) where {VR,VC}
    mul!(b, mxsolver.A, x)

    return nothing
end

function SALTBase.abs2_interp!(abs2ψ::VR, mxsolver::DirectMaxwellSolver{VR,VC}, ψ::VC) where {VR,VC}
    # Calculate the amplitude squares of Ex, Ey, Ez at their Yee grid locations.
    abs2ψ .= abs2.(ψ)

    # Interpolate Ex, Ey, Ez at the Yee grid voxel corner locations.
    ψc = mxsolver.vtemp
    mul!(ψc, mxsolver.Me2c, ψ)

    # Calculate the amplitude squares of Ex, Ey, Ez at the Yee grid's Ez, Ex, Ey locations.
    ψe = mxsolver.wtemp
    for M = mxsolver.Mc2e
        mul!(ψe, M, ψc)
        abs2ψ .+= abs2.(ψe)
    end
end

# In IterativeMaxwellData, the only things I need to store would be the things that I need
# for applying the Maxwell operation or averaging operation matrix-free, such as ε, ω, and
# grid.  The functions themselves that apply the matrix-free Maxwell operation or
# matrix-free averaging operation to a column vector should be implemented in MaxwellFDFD.
# mxsolve!, etc that take IterativeMaxwellData simply need to call those functions with the
# arguments it stores.
