official_code/src/sparse_precision.jl

# ------------------------------------------------------------------------------
# Sparse Precision Matrix Construction via KL-Optimal Cholesky
# ------------------------------------------------------------------------------

using FunctionalGPs
using GaussianMarkovRandomFields
using LinearAlgebra
using SparseArrays
using Distributions: Normal
using Random: randperm

using GaussianMarkovRandomFields:
    reverse_maximin_ordering,
    sparsity_pattern_from_ordering,
    sparse_approximate_cholesky!,
    sparse_approximate_cholesky,
    form_supernodes,
    PermutedMatrix,
    GMRF,
    InformationVector,
    ExponentialFamily,
    gaussian_approximation

using Kronecker: KroneckerProduct

export SparseGMRFApproximation
export sparse_precision, sparse_gmrf
export prescribe_indices
export condition_precision
export calibrate_output_scale

# ------------------------------------------------------------------------------
# Fast Block-Kronecker Matrix for Efficient Submatrix Extraction
# ------------------------------------------------------------------------------

"""
    FastBlockMatrix{T}

Optimized wrapper for BlockMatrix of KroneckerProducts that provides fast submatrix
extraction by materializing the small Kronecker factors.

For 2D tensor-product problems, this provides ~40-50x speedup in sparse Cholesky
while using ~100x less memory than full materialization.

Non-Kronecker blocks (e.g., in 1D problems) fall back to lazy evaluation.
"""
struct FastBlockMatrix{T} <: AbstractMatrix{T}
    # Materialized Kronecker factors (typed for fast access)
    A_factors::Matrix{Matrix{T}}
    B_factors::Matrix{Matrix{T}}
    B_dims::Matrix{Tuple{Int,Int}}
    # Original blocks for non-Kronecker fallback
    orig_blocks::Matrix
    is_kronecker::Matrix{Bool}
    block_offsets::Vector{Int}
    n::Int
end

Base.size(F::FastBlockMatrix) = (F.n, F.n)

"""
    FastBlockMatrix(K)

Create a FastBlockMatrix from a BlockMatrix. Automatically detects KroneckerProduct
blocks and materializes their factors for fast access.
"""
function FastBlockMatrix(K)
    orig_blocks = K.blocks
    n_blocks_r, n_blocks_c = size(orig_blocks)
    T = Float64

    A_factors = Matrix{Matrix{T}}(undef, n_blocks_r, n_blocks_c)
    B_factors = Matrix{Matrix{T}}(undef, n_blocks_r, n_blocks_c)
    B_dims = Matrix{Tuple{Int,Int}}(undef, n_blocks_r, n_blocks_c)
    is_kronecker = Matrix{Bool}(undef, n_blocks_r, n_blocks_c)

    for i in 1:n_blocks_r, j in 1:n_blocks_c
        block = orig_blocks[i, j]
        if block isa KroneckerProduct
            # Materialize Kronecker factors for fast access
            A_factors[i,j] = Matrix(block.A)
            B_factors[i,j] = Matrix(block.B)
            B_dims[i,j] = size(block.B)
            is_kronecker[i,j] = true
        else
            # Placeholder for type stability; use orig_blocks for actual access
            A_factors[i,j] = Matrix{T}(undef, 0, 0)
            B_factors[i,j] = Matrix{T}(undef, 0, 0)
            B_dims[i,j] = (0, 0)
            is_kronecker[i,j] = false
        end
    end

    block_sizes = [size(orig_blocks[i,1], 1) for i in 1:n_blocks_r]
    block_offsets = [0; cumsum(block_sizes)]
    n = block_offsets[end]

    return FastBlockMatrix{T}(A_factors, B_factors, B_dims, orig_blocks,
                               is_kronecker, block_offsets, n)
end

function Base.getindex(F::FastBlockMatrix{T}, I::AbstractVector{Int}, J::AbstractVector{Int}) where T
    n_I, n_J = length(I), length(J)
    result = Matrix{T}(undef, n_I, n_J)

    # Decode indices to block and local indices
    bi_I = Vector{Int}(undef, n_I)
    li_I = Vector{Int}(undef, n_I)
    @inbounds for k in 1:n_I
        bi_I[k] = max(1, searchsortedlast(F.block_offsets, I[k]-1))
        li_I[k] = I[k] - F.block_offsets[bi_I[k]]
    end

    bj_J = Vector{Int}(undef, n_J)
    lj_J = Vector{Int}(undef, n_J)
    @inbounds for k in 1:n_J
        bj_J[k] = max(1, searchsortedlast(F.block_offsets, J[k]-1))
        lj_J[k] = J[k] - F.block_offsets[bj_J[k]]
    end

    @inbounds for ki in 1:n_I
        bi = bi_I[ki]
        li = li_I[ki]
        for kj in 1:n_J
            bj = bj_J[kj]
            lj = lj_J[kj]

            if F.is_kronecker[bi, bj]
                # Fast path: use materialized Kronecker factors
                A_mat = F.A_factors[bi, bj]
                B_mat = F.B_factors[bi, bj]
                n_B_i, n_B_j = F.B_dims[bi, bj]
                li_A, li_B = divrem(li - 1, n_B_i)
                lj_A, lj_B = divrem(lj - 1, n_B_j)
                result[ki, kj] = A_mat[li_A+1, lj_A+1] * B_mat[li_B+1, lj_B+1]
            else
                # Fallback: use original lazy block
                result[ki, kj] = F.orig_blocks[bi, bj][li, lj]
            end
        end
    end
    return result
end

Base.getindex(F::FastBlockMatrix, i::Int, j::Int) = F[[i], [j]][1, 1]

"""
    FastPermutedMatrix{T}

PermutedMatrix wrapper for FastBlockMatrix that returns Symmetric views
for square submatrix extractions (required by sparse Cholesky).
"""
struct FastPermutedMatrix{T} <: AbstractMatrix{T}
    F::FastBlockMatrix{T}
    P::Vector{Int}
end

Base.size(M::FastPermutedMatrix) = size(M.F)

function Base.getindex(M::FastPermutedMatrix{T}, I::AbstractVector{Int}, J::AbstractVector{Int}) where T
    raw = M.F[M.P[I], M.P[J]]
    # Return Symmetric for square extractions (sparse Cholesky requirement)
    return length(I) == length(J) ? Symmetric(raw) : raw
end

Base.getindex(M::FastPermutedMatrix, i::Int, j::Int) = M.F[M.P[i], M.P[j]]

"""
    is_block_matrix(K) -> Bool

Check if K is a BlockMatrix (has .blocks field).
"""
is_block_matrix(K) = hasproperty(K, :blocks)

# ------------------------------------------------------------------------------
# Ordering Strategies
# ------------------------------------------------------------------------------

"""
    OrderingStrategy

Strategy for ordering functional blocks in sparse Cholesky.

- `:integrals_coarsest`: Integrals rightmost (coarsest), then evaluations, then derivatives (finest)
- `:evaluations_coarsest`: Evaluations rightmost (coarsest)
- `:derivatives_coarsest`: Derivatives rightmost (coarsest)
- `:natural`: Use the order provided by the user
"""
const OrderingStrategy = Symbol

# Priority maps: lower number = finer scale = leftmost in Cholesky
# Note: OTHER always has the same priority as DERIVATIVE (safest default)
const ORDERING_PRIORITIES = Dict(
    :integrals_coarsest => Dict(
        DERIVATIVE => 1,
        OTHER => 1,
        EVALUATION => 2,
        FACE_INTEGRAL => 3,
        INTEGRAL => 4,
    ),
    :evaluations_coarsest => Dict(
        DERIVATIVE => 1,
        OTHER => 1,
        FACE_INTEGRAL => 2,
        INTEGRAL => 3,
        EVALUATION => 4,
    ),
    :derivatives_coarsest => Dict(
        INTEGRAL => 1,
        FACE_INTEGRAL => 2,
        EVALUATION => 3,
        DERIVATIVE => 4,
        OTHER => 4,
    )
)

"""
    create_block_ordering(X_blocks, categories, names, layout, strategy)

Create a global ordering for sparse Cholesky from multiple functional blocks.

# Arguments
- `X_blocks`: Vector of coordinate matrices, one per functional
- `categories`: Vector of `FunctionalCategory` for each block
- `names`: Vector of Symbol names for each block
- `layout`: Layout mapping names to index ranges
- `strategy`: Ordering strategy (`:integrals_coarsest`, `:evaluations_coarsest`, etc.)

# Returns
- `P`: Global permutation vector
- `ℓ`: Maximin lengthscales for each index
- `X`: Global coordinate matrix
"""
function create_block_ordering(X_blocks::Vector, categories::Vector{FunctionalCategory},
                               names::Vector{Symbol}, layout::Layout,
                               strategy::OrderingStrategy)
    n_blocks = length(X_blocks)

    # Determine block order based on strategy
    if strategy == :natural
        block_order = collect(1:n_blocks)
    else
        priorities = get(ORDERING_PRIORITIES, strategy, ORDERING_PRIORITIES[:integrals_coarsest])
        block_order = sortperm([priorities[cat] for cat in categories])
    end

    # Compute maximin ordering for each block
    P_blocks = Vector{Vector{Int}}(undef, n_blocks)
    ℓ_blocks = Vector{Vector{Float64}}(undef, n_blocks)

    for i in 1:n_blocks
        P_local, ℓ_local = reverse_maximin_ordering(X_blocks[i])
        P_blocks[i] = P_local
        ℓ_blocks[i] = ℓ_local
    end

    # Find coarsest lengthscale (from the coarsest block = last in order)
    # Use the lengthscale of the first point in the permutation (the coarsest point)
    coarsest_block = block_order[end]
    P_coarsest = P_blocks[coarsest_block]
    ℓ_coarse = ℓ_blocks[coarsest_block][P_coarsest[1]]

    # Total state dimension
    n_total = sum(size(X, 2) for X in X_blocks)

    # Build global ordering following block_order
    # ℓ_global must be indexed by GLOBAL STATE INDEX (not permutation order)
    # because sparsity_pattern_from_ordering does ℓ[P[j]] to get lengthscale
    P_global = Int[]
    ℓ_global = zeros(n_total)

    for block_idx in block_order
        name = names[block_idx]
        P_local = P_blocks[block_idx]
        ℓ_local = ℓ_blocks[block_idx]

        # Map local indices to global indices
        global_indices = indices(layout, name)
        P_global_block = global_indices[P_local]

        # Set lengthscales at the GLOBAL indices (not appended in permutation order!)
        # For non-coarsest blocks, use the coarsest lengthscale
        if block_idx != coarsest_block
            for global_idx in global_indices
                ℓ_global[global_idx] = ℓ_coarse
            end
        else
            for (local_idx, global_idx) in enumerate(global_indices)
                ℓ_global[global_idx] = ℓ_local[local_idx]
            end
        end

        append!(P_global, P_global_block)
    end

    # Build global coordinate matrix (in original order, not block_order)
    X_global = hcat(X_blocks...)

    return P_global, ℓ_global, X_global
end

# ------------------------------------------------------------------------------
# Main API
# ------------------------------------------------------------------------------

"""
    SparseGMRFApproximation

Result of sparse Cholesky approximation to a GP defined by linear functionals.

# Fields
- `Q`: Sparse precision matrix (in original ordering)
- `layout`: Layout for named access to state vector
- `functionals`: Original linear functionals
- `info`: Named tuple with sparsity statistics
"""
struct SparseGMRFApproximation{TQ<:AbstractMatrix, TL<:Layout}
    Q::TQ
    layout::TL
    functionals::Vector
    info::NamedTuple
end

function Base.show(io::IO, approx::SparseGMRFApproximation)
    method = approx.info.supernodal ? "supernodal" : "simplicial"
    print(io, "SparseGMRFApproximation(n=$(approx.info.n), " *
              "nnz=$(approx.info.nnz), fill=$(round(approx.info.fill_pct, digits=1))%, $method)")
end

"""
    sparse_precision(named_functionals, kernel; ρ=2.0, λ=1.5, ordering=:integrals_coarsest)

Build a sparse precision matrix from named linear functionals using KL-optimal Cholesky approximation.

# Arguments
- `named_functionals`: Vector of `name => functional` pairs, e.g.,
  ```julia
  [:f => EvaluationFunctional(X),
   :f_dx => L_eval ∘ PartialDerivative((1,)),
   :f_int => VectorizedLebesgueIntegral(intervals)]
  ```
- `kernel`: GP kernel (e.g., `HalfIntegerMaternKernel`)

# Keyword Arguments
- `ρ=2.0`: Sparsity threshold. Larger values give denser but more accurate approximations.
- `λ=1.5`: Supernodal clustering threshold. Controls how columns are grouped for
  efficient factorization. Set to `nothing` to use simplicial (column-by-column)
  factorization instead. Supernodal is ~2-3x more efficient (more nnz per compute).
- `ordering=:integrals_coarsest`: Block ordering strategy. Options:
  - `:integrals_coarsest`: Integrals as coarsest (recommended for FVM)
  - `:evaluations_coarsest`: Evaluations as coarsest
  - `:derivatives_coarsest`: Derivatives as coarsest
  - `:natural`: Use the order provided

# Returns
- `SparseGMRFApproximation` containing the sparse precision matrix and metadata

# Example
```julia
# Define functionals
endpoints = range(0, 1, length=51)
intervals = intervals_from_endpoints(collect(endpoints))
k = HalfIntegerMaternKernel(2, [0.1])

L_eval = EvaluationFunctional(endpoints)
L_deriv = L_eval ∘ PartialDerivative((1,))
L_int = VectorizedLebesgueIntegral(intervals)

# Build sparse precision (supernodal by default)
approx = sparse_precision([
    :f => L_eval,
    :f_dx => L_deriv,
    :f_int => L_int
], k; ρ=2.0)

# Use simplicial instead
approx_simp = sparse_precision([...], k; ρ=2.0, λ=nothing)

# Create GMRF
x = GMRF(zeros(approx.info.n), approx.Q)
```
"""
function sparse_precision(named_functionals::Vector{<:Pair{Symbol}}, kernel;
                          ρ::Real=2.0, λ::Union{Real,Nothing}=1.5,
                          ordering::OrderingStrategy=:integrals_coarsest)
    # Extract names and functionals
    names = Symbol[first(p) for p in named_functionals]
    functionals = [last(p) for p in named_functionals]
    n_blocks = length(functionals)

    # Classify functionals
    categories = [functional_category(L) for L in functionals]

    # Extract coordinates
    X_blocks = [get_coordinates(L) for L in functionals]

    # Get output sizes and build layout
    sizes = [output_length(L) for L in functionals]
    state_layout = layout(NamedTuple{Tuple(names)}(Tuple(sizes)))
    n_total = sum(sizes)

    # Build stacked functional and covariance matrix
    L_stack = StackedLinearFunctional(functionals...)
    K = L_stack(L_stack(kernel))
    K_sym = Symmetric(K)

    # Create ordering
    P, ℓ, X = create_block_ordering(X_blocks, categories, names, state_layout, ordering)

    # Create sparsity pattern
    S = sparsity_pattern_from_ordering(X, P, ℓ, Float64(ρ))

    # Apply sparse Cholesky (supernodal by default, simplicial if λ=nothing)
    if λ === nothing
        # Simplicial (column-by-column) factorization
        K_P = PermutedMatrix(K_sym, P)
        sparse_approximate_cholesky!(K_P, S)
        L = S
    else
        # Supernodal factorization with FastBlockMatrix optimization
        # For BlockMatrix inputs (2D tensor products), this provides ~40-50x speedup
        if is_block_matrix(K)
            K_P = FastPermutedMatrix(FastBlockMatrix(K), P)
        else
            K_P = PermutedMatrix(K_sym, P)
        end
        sc = form_supernodes(S, P, ℓ; λ=Float64(λ))
        L = sparse_approximate_cholesky(K_P, sc)
    end

    # Build sparse precision (unpermuted to original ordering)
    Q_perm = L * L'
    P_inv = invperm(P)
    Q_sparse = Q_perm[P_inv, P_inv]

    # Compute sparsity statistics
    nnz_Q = nnz(Q_sparse)
    nnz_L = nnz(L)
    dense_nnz = n_total * (n_total + 1) ÷ 2
    fill_pct = 100.0 * nnz_Q / (2 * dense_nnz)

    info = (
        ρ = Float64(ρ),
        λ = λ === nothing ? nothing : Float64(λ),
        supernodal = λ !== nothing,
        n = n_total,
        nnz = nnz_Q,
        nnz_L = nnz_L,
        fill_pct = fill_pct,
        ordering = ordering,
        categories = categories
    )

    return SparseGMRFApproximation(Q_sparse, state_layout, functionals, info)
end

"""
    sparse_gmrf(named_functionals, kernel; ρ=2.0, λ=1.5, ordering=:integrals_coarsest)

Create a GMRF with sparse precision from named linear functionals.

Convenience wrapper around `sparse_precision` that directly returns a GMRF.

# Returns
- `gmrf`: GMRF with zero mean and sparse precision
- `layout`: Layout for named access to state vector
- `approx`: Full `SparseGMRFApproximation` with metadata
"""
function sparse_gmrf(named_functionals::Vector{<:Pair{Symbol}}, kernel;
                     ρ::Real=2.0, λ::Union{Real,Nothing}=1.5,
                     ordering::OrderingStrategy=:integrals_coarsest)
    approx = sparse_precision(named_functionals, kernel; ρ=ρ, λ=λ, ordering=ordering)
    gmrf = GMRF(zeros(approx.info.n), approx.Q)
    return gmrf, approx.layout, approx
end

# ------------------------------------------------------------------------------
# Convenience: prescribe_indices (conditioning helper)
# ------------------------------------------------------------------------------

"""
    prescribe_indices(x::GMRF, indices, values; noise_std=1e-3)

Condition a GMRF on observations at specific indices.

# Arguments
- `x`: GMRF to condition
- `indices`: Vector of indices where observations are made
- `values`: Observation values
- `noise_std`: Observation noise standard deviation

# Returns
- Conditioned GMRF
"""
function prescribe_indices(x::GMRF, indices, values; noise_std::Real=1e-3)
    obs_model = ExponentialFamily(Normal, indices=indices)
    obs_lik = obs_model(values; σ=Float64(noise_std))
    return gaussian_approximation(x, obs_lik)
end

# ------------------------------------------------------------------------------
# Output Scale Calibration
# ------------------------------------------------------------------------------

"""
    calibrate_output_scale(Q₀, A_fvm, μ_BC; n_samples=10)

Estimate output scale σ via quasi-MLE from FVM innovations.

After conditioning on ICs/BCs to get `μ_BC`, the FVM constraints have nonzero
innovations. This function estimates σ² such that the standardized innovations
have unit variance (well-calibrated uncertainty).

# Arguments
- `Q₀`: Unit-scale sparse precision matrix
- `A_fvm`: FVM constraint matrix (n_constraints × n_state)
- `μ_BC`: Mean after BC conditioning (length n_state)
- `n_samples`: Number of random constraints to sample (default 10)

# Returns
- `σ`: Estimated output scale (square root of σ²)

# Reference
Based on CAPOS (Bosch, Hennig, Tronarp 2021) quasi-MLE approach for
probabilistic ODE solvers.

# Example
```julia
# After building sparse precision with unit output scale
approx = sparse_precision(functionals, kernel; ρ=2.0)

# Condition on boundary conditions
x_BC = prescribe_indices(x0, bc_indices, bc_values)
μ_BC = mean(x_BC)

# Estimate output scale from FVM innovations
σ = calibrate_output_scale(approx.Q, A_fvm, μ_BC; n_samples=10)

# Apply calibrated scale
Q_calibrated = approx.Q / σ^2
```
"""
function calibrate_output_scale(Q₀, A_fvm, μ_BC; n_samples::Int=10)
    n_constraints = size(A_fvm, 1)
    n_samples = min(n_samples, n_constraints)

    # Compute Cholesky factorization once
    F = cholesky(Q₀)

    # Sample random constraint indices
    sample_idx = randperm(n_constraints)[1:n_samples]

    σ²_sum = 0.0
    for i in sample_idx
        aᵢ = Vector(A_fvm[i, :])

        # Innovation: how far is BC-conditioned mean from satisfying constraint?
        rᵢ = -dot(aᵢ, μ_BC)

        # Prior variance of constraint: sᵢ = aᵢ' Σ₀ aᵢ = aᵢ' Q₀⁻¹ aᵢ
        # Solve Q₀ x = aᵢ to get x = Σ₀ aᵢ, then sᵢ = aᵢ' x
        x = F \ aᵢ
        sᵢ = dot(aᵢ, x)

        # Accumulate normalized squared innovation
        if sᵢ > 0
            σ²_sum += rᵢ^2 / sᵢ
        end
    end

    σ² = σ²_sum / n_samples

    # Guard against σ = 0 (can happen if all innovations are zero)
    σ = sqrt(max(σ², 1e-10))

    return σ
end