functionals.jl

using KernelFunctions: Kernel, ZeroKernel, KernelSum, ScaledKernel

export LinearLagrangeIntegrals, Q1Projection, mass_matrix

# Default quadrature degree assumed for applicands without a polynomial degree
# (e.g. non-compact kernels, general mean functions).
const _DEFAULT_QUAD_DEGREE = 3

# ──────────────────────────────────────────────────────────────────────────────
# 1D linear (P1) Lagrange L² projection — the tensor-product building block
# ──────────────────────────────────────────────────────────────────────────────
"""
    LinearLagrangeIntegrals(grid::AbstractRange)

Linear functional that computes the **1D linear (P1) Lagrange L² loads** on the
**interior** nodes of the **uniform** `grid` (an `AbstractRange`; a descending
range is flipped automatically).

The two boundary nodes are held at zero (homogeneous Dirichlet) and dropped, so a
grid of `N + 2` nodes yields `N` loads — one per interior node, whose hat `φ_i`
spans its two **full** neighbouring elements. Applied to a function `f`, it
returns the vector of **hat-weighted loads**

```
ℒ(f)_i = b_i = ∫ φ_i f dx,   i = 2, …, N + 1.
```

The (tridiagonal, SPD) interior P1 mass matrix `M_{ij} = ∫ φ_i φ_j dx` (size
`N × N`) is computed and stored at construction (see [`mass_matrix`](@ref)) but
**not** applied: the orthogonal L² projection coefficients are `M \\ ℒ(f)`,
formed downstream when wanted.

# Quadrature as a linear combination of evaluation functionals

Each element's integrand `φ_i · f` has degree `deg f + 1`, so the Gauss rule —
**dispatched on what the functional is applied to** — is exact to that degree.
For a [`CompactPolynomialKernel`](@ref) of polynomial degree `d` it uses
`N = clamp(⌈(d + 2)/2⌉, 1, 5)` nodes per element (exact to degree `d + 1`);
everything else assumes degree `$(_DEFAULT_QUAD_DEGREE)`. Only the hard-coded
`N ∈ {1, 2, 3, 4, 5}` schemes are provided.

On a uniform grid every element has width `h = step(grid)`, so the Gauss nodes of
all elements at rule node `τ` are a single shifted subrange `grid[k:k+N-1] .+ τ·h`.
Each interior load is the right-hat (`τ`) contribution from its left element plus
the left-hat (`1-τ`) contribution from its right element, so the whole functional
is one sum over the (few) rule nodes of **scalar-scaled vector
[`EvaluationFunctional`](@ref)s**, which reuses the ordinary evaluation machinery
and produces no bespoke cross-covariance type.

See also [`Q1Projection`](@ref) for the tensor-product (multi-dimensional) version.
"""
struct LinearLagrangeIntegrals{R <: AbstractRange} <: AbstractLinearFunctional
    grid::R
    mass_matrix::Matrix{Float64}

    function LinearLagrangeIntegrals(grid::AbstractRange)
        length(grid) >= 3 ||
            throw(ArgumentError("the grid needs at least three nodes (one interior)"))
        g = step(grid) < 0 ? reverse(grid) : grid
        step(g) > 0 || throw(ArgumentError("the grid must have a nonzero step"))
        return new{typeof(g)}(g, _mass_matrix(g))
    end
end

"""
    mass_matrix(ℒ::LinearLagrangeIntegrals) -> Matrix

The interior P1 mass matrix `M_{ij} = ∫ φ_i φ_j dx` (`N × N`, boundary nodes
excluded). `ℒ` returns the loads `b = (∫ φ_i f)_i`; the orthogonal L² projection
coefficients are `mass_matrix(ℒ) \\ b`.
"""
mass_matrix(ℒ::LinearLagrangeIntegrals) = ℒ.mass_matrix

FunctionalGPs.output_shape(ℒ::LinearLagrangeIntegrals) = (length(ℒ.grid) - 2,)

Base.show(io::IO, ℒ::LinearLagrangeIntegrals) =
    print(io, "LinearLagrangeIntegrals($(length(ℒ.grid) - 2) interior nodes)")

# ── Hard-coded Gauss–Legendre rules on [-1, 1] (N ∈ {1, 2, 3, 4, 5}) ──────────
_gauss_legendre(n::Int) = _gauss_legendre(Val(n))
_gauss_legendre(::Val{1}) = ([0.0], [2.0])
_gauss_legendre(::Val{2}) = ([-0.5773502691896257, 0.5773502691896257], [1.0, 1.0])
_gauss_legendre(::Val{3}) =
    ([-0.7745966692414834, 0.0, 0.7745966692414834], [5 / 9, 8 / 9, 5 / 9])
_gauss_legendre(::Val{4}) = (
    [-0.8611363115940526, -0.3399810435848563, 0.3399810435848563, 0.8611363115940526],
    [0.3478548451374538, 0.6521451548625461, 0.6521451548625461, 0.3478548451374538],
)
_gauss_legendre(::Val{5}) = (
    [-0.906179845938664, -0.5384693101056831, 0.0, 0.5384693101056831, 0.906179845938664],
    [0.2369268850561891, 0.4786286704993665, 0.5688888888888889, 0.4786286704993665,
        0.2369268850561891],
)
_gauss_legendre(::Val{N}) where {N} =
    throw(ArgumentError("only Gauss–Legendre schemes with N ∈ 1:5 are hard-coded, got N = $N"))

# Per-element Gauss node count, exact to degree `d + 1` (capped at the N = 5 rule).
_gauss_nodes_for_degree(d::Integer) = clamp(cld(d + 2, 2), 1, 5)

# Polynomial degree of the applicand that drives the quadrature dispatch.
_quad_degree(::Any) = _DEFAULT_QUAD_DEGREE
_quad_degree(k::CompactPolynomialKernel) = length(k.poly.coeffs) - 1

# ── Interior P1 mass matrix on a uniform grid ─────────────────────────────────
# Every element of width `h` contributes the classic (h/6)·[2 1; 1 2] block; the
# boundary nodes are then dropped (homogeneous Dirichlet).
function _mass_matrix(g::AbstractRange)
    n = length(g)
    h = step(g)
    M = zeros(n, n)
    for e in 1:(n - 1)
        M[e, e]         += h / 3
        M[e + 1, e + 1] += h / 3
        M[e, e + 1]     += h / 6
        M[e + 1, e]     += h / 6
    end
    return M[2:(end - 1), 2:(end - 1)]
end

# ── Interior loads as one loop over scaled vector evaluation functionals ──────
# Interior node `j` (global node `j+1`) integrates `φ_j · f` over its two full
# elements: element `j` (right hat, weight `τ`, Gauss points `g[1:end-2] .+ τ·h`)
# and element `j+1` (left hat, weight `1-τ`, Gauss points `g[2:end-1] .+ τ·h`).
# Both subranges have the interior length, so they add directly — no padding.
function _quadrature_functional(g::AbstractRange, d::Integer)
    h = step(g)
    ξ, w = _gauss_legendre(_gauss_nodes_for_degree(d))
    terms = map(zip(ξ, w)) do (ξᵢ, wᵢ)
        τ = (ξᵢ + 1) / 2
        c = (h / 2) * wᵢ
        return c * τ       * EvaluationFunctional(g[1:(end - 2)] .+ τ * h) +
               c * (1 - τ) * EvaluationFunctional(g[2:(end - 1)] .+ τ * h)
    end
    return reduce(+, terms)
end

# ── Applying the functional ───────────────────────────────────────────────────
# To a kernel → reuse the evaluation machinery via the combinator.
(ℒ::LinearLagrangeIntegrals)(k::Kernel; arg::Integer = 2) =
    _quadrature_functional(ℒ.grid, _quad_degree(k))(k; arg = arg)

# Applied to a parametric kernel's feature map `ϕ` (e.g. via `ℒ(::ParametricKernel)`
# or the tensor-product factors): route the loads through the Gauss quadrature — a
# sum of `EvaluationFunctional`s — so the existing `feature_matrix` machinery
# integrates `ϕ` at the Gauss nodes, rather than calling `ℒ(ϕ)` directly.
feature_matrix(ϕ, ℒ::LinearLagrangeIntegrals) =
    feature_matrix(ϕ, _quadrature_functional(ℒ.grid, _quad_degree(ϕ)))

# Disambiguate the wrapper kernels that FunctionalGPs/base.jl handles generically.
(ℒ::LinearLagrangeIntegrals)(k::ScaledKernel; arg::Integer = 2) =
    k.σ² * ℒ(k.kernel; arg = arg)
(ℒ::LinearLagrangeIntegrals)(k::LinearlyScaledKernel; arg::Integer = 2) =
    k.scalar * ℒ(k.kernel; arg = arg)
(ℒ::LinearLagrangeIntegrals)(k::KernelSum; arg::Integer = 2) =
    mapreduce(kᵢ -> ℒ(kᵢ; arg = arg), +, k.kernels)
(ℒ::LinearLagrangeIntegrals)(::ZeroKernel; arg::Integer = 2) =
    FunctionalGPs.ZeroPVCrosscov(output_shape(ℒ), Int(arg))
# A `ParametricKernel` is handled through its feature map by the generic
# `(::AbstractLinearFunctional)(::ParametricKernel)` in `functional_gps.jl`. This more
# specific method disambiguates it against the `(::LinearLagrangeIntegrals)(::Kernel)`
# quadrature method above, which matters when a parametric factor (e.g. a constant-in-z
# indicator kernel) appears inside a tensor-product kernel evaluated by `Q1Projection`.
(ℒ::LinearLagrangeIntegrals)(k::ParametricKernel; arg::Integer = 2) =
    ParametricPVCrosscov{arg}(k.feature_map, k.inprod, ℒ, feature_matrix(k.feature_map, ℒ))

# To a mean function → loads of the mean.
(ℒ::LinearLagrangeIntegrals)(m::AbstractGPs.MeanFunction) =
    _quadrature_functional(ℒ.grid, _quad_degree(m))(m)
(ℒ::LinearLagrangeIntegrals)(::AbstractGPs.ZeroMean{T}) where {T} =
    zeros(T, output_shape(ℒ))
# Every interior hat spans two full elements of width `h`, so `∫φ_i = h` and a
# constant `c` has loads `c·h` at every interior node — no quadrature needed.
(ℒ::LinearLagrangeIntegrals)(m::AbstractGPs.ConstMean) =
    fill(m.c * step(ℒ.grid), length(ℒ.grid) - 2)

"""
    Q1Projection(grid₁, grid₂, …)

L² projection onto **Q1 (multilinear) Lagrange elements** on the tensor-product
grid `grid₁ × grid₂ × …`, built as the tensor product (`⊗`) of
[`LinearLagrangeIntegrals`](@ref) factors. With a single grid it returns the
1D projection itself.

Each `gridᵈ` is a uniform `AbstractRange`. Applied to a tensor-product kernel
`k₁ ⊗ k₂ ⊗ …`, the existing `TensorProductFunctional` machinery evaluates each 1D
factor independently.

```julia
ℒ = Q1Projection(0:0.25:1, 0:0.5:1)
μ, Σ = mean_and_cov(GP(WendlandKernel(1, 1) ⊗ WendlandKernel(1, 1)), ℒ)
```
"""
function Q1Projection(grids::AbstractRange...)
    isempty(grids) && throw(ArgumentError("at least one grid range is required"))
    return reduce(⊗, map(LinearLagrangeIntegrals, grids))
end