gam-sae 0.3.151

use super::*;

/// One undirected candidate edge between graph anchors.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct GraphEdge {
    pub a: usize,
    pub b: usize,
}

impl GraphEdge {
    pub fn new(a: usize, b: usize) -> Result<Self, String> {
        if a == b {
            return Err(format!("GraphEdge cannot join vertex {a} to itself"));
        }
        Ok(if a < b {
            Self { a, b }
        } else {
            Self { a: b, b: a }
        })
    }
}

/// Per-edge rank charge for a learned graph edge.
///
/// This is the tiered spine's BIC/Laplace currency applied to one edge:
/// `0.5 * d_eff * ln(n_eff)`. For a graph atom edge, `d_eff` is the number of
/// fiber channels coupled by that edge.
pub fn graph_edge_rank_charge(n_eff: f64, fiber_rank: usize) -> f64 {
    0.5 * fiber_rank as f64 * n_eff.max(2.0).ln()
}

/// Pairwise mutual information in NATS of two atoms' binary co-activation
/// indicators, from the `2×2` table implied by a [`crate::atom_codes::CoactivationStats`]
/// (its `n_obs`, `n_a`, `n_b`, `n_joint` counts). The `0·log 0` cells are dropped;
/// the value is non-negative up to floating-point noise. This is the co-fire
/// coupling strength the graph-atom enrollment reads.
pub fn coactivation_mi_nats(stats: &crate::atom_codes::CoactivationStats) -> f64 {
    let n = stats.n_obs as f64;
    if n <= 0.0 {
        return 0.0;
    }
    let p1x = stats.n_a as f64 / n;
    let px1 = stats.n_b as f64 / n;
    let p11 = stats.n_joint as f64 / n;
    let p10 = (p1x - p11).max(0.0);
    let p01 = (px1 - p11).max(0.0);
    let p00 = (1.0 - p11 - p10 - p01).max(0.0);
    let cell = |p: f64, pa: f64, pb: f64| -> f64 {
        if p > 0.0 && pa > 0.0 && pb > 0.0 {
            p * (p / (pa * pb)).ln()
        } else {
            0.0
        }
    };
    (cell(p11, p1x, px1)
        + cell(p10, p1x, 1.0 - px1)
        + cell(p01, 1.0 - p1x, px1)
        + cell(p00, 1.0 - p1x, 1.0 - px1))
    .max(0.0)
}

/// Per-edge enrollment evidence for one discovered co-fire pair: the tuple
/// `(precision, delta_loss)` the graph-atom ARD survival rule consumes.
///
/// The edge **precision** (the graph-Laplacian weight) is the pair's mutual
/// information in nats — its co-fire coupling strength. The **delta-loss** (the
/// REML currency [`LearnedGraphAtom::from_reml_candidate_edges`] compares to the
/// one-edge [`graph_edge_rank_charge`]) is `n_eff · MI_nats`: the edge's
/// likelihood-ratio evidence, i.e. the `G²/2` deviance of its `2×2`
/// co-activation table. A pair with no statistical dependence carries zero
/// evidence and cannot survive enrollment.
pub fn coactivation_edge_evidence(
    stats: &crate::atom_codes::CoactivationStats,
    n_eff: f64,
) -> (f64, f64) {
    let mi = coactivation_mi_nats(stats);
    (mi, n_eff.max(0.0) * mi)
}

/// Exact read-out of the surviving graph topology.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GraphTopologyReadout {
    pub vertices: usize,
    pub surviving_edges: usize,
    pub b0: usize,
    pub b1: usize,
}

/// Named-shape compression certified after the graph has been learned.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GraphCompressionKind {
    Circle,
    Interval,
    FiniteSet,
    /// A contractible bounded surface (`χ = 1`, orientable, `b₁ = 0`, `b₂ = 0`):
    /// the topological type of a sampled sheet or disk. This is what a swiss roll
    /// glues to — the fold unrolls into one flat chart with no handle and no
    /// closed 2-cycle (#2280 acceptance: "swiss roll → sheet").
    Disk,
    Cylinder,
    /// The non-orientable bounded surface with one boundary loop (`χ = 0`,
    /// non-orientable, `b₁ = 1`, `b₂ = 0`): a cylinder's orientation cocycle with
    /// a single sign reversal around the loop. Recognized by the orientability
    /// certificate — the "half-twist as a discrete sign" (#2280 acceptance:
    /// "Möbius band → holonomy sign detected").
    MobiusStrip,
    Torus,
    Sphere,
    ProjectivePlane,
    KleinBottle,
    Graph,
}

/// MDL read-out for whether the learned edge set earns a standard name.
#[derive(Debug, Clone, PartialEq)]
pub struct GraphCompressionReport {
    pub kind: GraphCompressionKind,
    pub name: &'static str,
    pub generic_edge_bits: f64,
    pub named_bits: f64,
    pub bits_saved: f64,
}

/// The structure-search birth currency for a graph atom.
#[derive(Debug, Clone, PartialEq)]
pub struct GraphStructureSelection {
    pub selected: bool,
    pub total_edge_delta_loss: f64,
    pub total_edge_charge: f64,
    pub margin: f64,
    pub topology: GraphTopologyReadout,
    pub occupancy: OccupancyLaw,
    pub compression: GraphCompressionReport,
}

impl GraphCompressionReport {
    pub fn certified(
        kind: GraphCompressionKind,
        name: &'static str,
        generic_edge_bits: f64,
        named_bits: f64,
    ) -> Self {
        Self {
            kind,
            name,
            generic_edge_bits,
            named_bits,
            bits_saved: generic_edge_bits - named_bits,
        }
    }

    pub fn unnamed(generic_edge_bits: f64) -> Self {
        Self {
            kind: GraphCompressionKind::Graph,
            name: "structure without a standard name",
            generic_edge_bits,
            named_bits: generic_edge_bits,
            bits_saved: 0.0,
        }
    }

    pub fn earns_standard_name(&self) -> bool {
        self.kind != GraphCompressionKind::Graph && self.bits_saved > 0.0
    }
}

/// A canonical learned graph atom: anchors with a learned subset of a derived kNN
/// candidate edge set.
///
/// The smoothness penalty is `beta^T (L_W kron I_r) beta`, where `L_W` is the
/// weighted graph Laplacian of the surviving graph and `r` is the fiber rank.
/// Edge survival is read from the same rank-charge discipline used by tiered
/// births/deaths: an edge survives only when the REML loss increase from
/// removing it is greater than its one-edge charge. Betti read-out is exact on
/// surviving edges; named shapes are secondary MDL compressions of the graph.
#[derive(Debug, Clone)]
pub struct LearnedGraphAtom {
    anchor_embeddings: Array2<f64>,
    candidate_edges: Vec<GraphEdge>,
    edge_precisions: Vec<f64>,
    edge_delta_loss: Vec<f64>,
    surviving_edges: Vec<bool>,
    n_eff: f64,
    occupancy: OccupancyLaw,
}

impl LearnedGraphAtom {
    /// Derived k for the anchor-kNN candidate graph. The graph atom is expected to
    /// learn edge survival by ARD; k only supplies a sparse local superset.
    pub fn derived_knn_k(anchors: usize) -> usize {
        anchors.saturating_sub(1).min(2)
    }

    /// Build the derived undirected kNN candidate edge set from anchor embeddings.
    pub fn knn_candidate_edges(
        anchor_embeddings: ArrayView2<'_, f64>,
    ) -> Result<Vec<GraphEdge>, String> {
        validate_anchor_embeddings(anchor_embeddings)?;
        let anchors = anchor_embeddings.nrows();
        let k = Self::derived_knn_k(anchors);
        if k == 0 {
            return Ok(Vec::new());
        }
        let mut edges = Vec::new();
        for a in 0..anchors {
            let mut distances = Vec::<(f64, usize)>::with_capacity(anchors.saturating_sub(1));
            for b in 0..anchors {
                if a == b {
                    continue;
                }
                let mut dist2 = 0.0_f64;
                for c in 0..anchor_embeddings.ncols() {
                    let d = anchor_embeddings[[a, c]] - anchor_embeddings[[b, c]];
                    dist2 += d * d;
                }
                distances.push((dist2, b));
            }
            distances.sort_by(|left, right| {
                left.0
                    .partial_cmp(&right.0)
                    .unwrap_or(std::cmp::Ordering::Equal)
                    .then_with(|| left.1.cmp(&right.1))
            });
            for &(_, b) in distances.iter().take(k) {
                let edge = GraphEdge::new(a, b)?;
                if !edges.contains(&edge) {
                    edges.push(edge);
                }
            }
        }
        edges.sort_by_key(|edge| (edge.a, edge.b));
        Ok(edges)
    }

    /// Build a learned graph atom from the derived kNN candidate edge set.
    pub fn from_reml_knn_edges(
        anchor_embeddings: ArrayView2<'_, f64>,
        row_coordinates: &[f64],
        n_eff: f64,
        edge_precisions: &[f64],
        edge_delta_loss: &[f64],
    ) -> Result<Self, String> {
        let candidate_edges = Self::knn_candidate_edges(anchor_embeddings)?;
        Self::from_reml_candidate_edges(
            anchor_embeddings,
            row_coordinates,
            n_eff,
            &candidate_edges,
            edge_precisions,
            edge_delta_loss,
        )
    }

    /// Build a graph atom from a caller-supplied candidate edge set.
    ///
    /// Production births should pass the kNN set from [`Self::knn_candidate_edges`].
    /// Tests and certified imports use this form to make the candidate superset
    /// explicit while preserving the same ARD survival rule.
    pub fn from_reml_candidate_edges(
        anchor_embeddings: ArrayView2<'_, f64>,
        row_coordinates: &[f64],
        n_eff: f64,
        candidate_edges: &[GraphEdge],
        edge_precisions: &[f64],
        edge_delta_loss: &[f64],
    ) -> Result<Self, String> {
        validate_anchor_embeddings(anchor_embeddings)?;
        let anchors = anchor_embeddings.nrows();
        let fiber_rank = anchor_embeddings.ncols();
        if candidate_edges.is_empty() {
            return Err("LearnedGraphAtom requires at least one candidate edge".to_string());
        }
        if edge_precisions.len() != candidate_edges.len() {
            return Err(format!(
                "LearnedGraphAtom edge_precisions length {} must equal candidate edges {}",
                edge_precisions.len(),
                candidate_edges.len()
            ));
        }
        if edge_delta_loss.len() != candidate_edges.len() {
            return Err(format!(
                "LearnedGraphAtom edge_delta_loss length {} must equal candidate edges {}",
                edge_delta_loss.len(),
                candidate_edges.len()
            ));
        }
        if !(n_eff.is_finite() && n_eff > 0.0) {
            return Err(format!(
                "LearnedGraphAtom n_eff must be finite and positive; got {n_eff}"
            ));
        }
        let mut normalized = Vec::<GraphEdge>::with_capacity(candidate_edges.len());
        for (idx, edge) in candidate_edges.iter().enumerate() {
            if edge.a >= anchors || edge.b >= anchors || edge.a == edge.b {
                return Err(format!(
                    "LearnedGraphAtom candidate edge {idx} = ({}, {}) is invalid for {anchors} anchors",
                    edge.a, edge.b
                ));
            }
            let edge = GraphEdge::new(edge.a, edge.b)?;
            if normalized.contains(&edge) {
                return Err(format!(
                    "LearnedGraphAtom candidate edge {idx} duplicates ({}, {})",
                    edge.a, edge.b
                ));
            }
            normalized.push(edge);
        }
        for (edge, &precision) in edge_precisions.iter().enumerate() {
            if !(precision.is_finite() && precision >= 0.0) {
                return Err(format!(
                    "LearnedGraphAtom edge {edge} precision must be finite and nonnegative; got {precision}"
                ));
            }
        }
        for (edge, &delta) in edge_delta_loss.iter().enumerate() {
            if !delta.is_finite() {
                return Err(format!(
                    "LearnedGraphAtom edge {edge} deletion loss must be finite; got {delta}"
                ));
            }
        }

        let charge = graph_edge_rank_charge(n_eff, fiber_rank);
        let surviving_edges = edge_precisions
            .iter()
            .zip(edge_delta_loss.iter())
            .map(|(&precision, &delta)| precision > 0.0 && delta > charge)
            .collect();

        Ok(Self {
            anchor_embeddings: anchor_embeddings.to_owned(),
            candidate_edges: normalized,
            edge_precisions: edge_precisions.to_vec(),
            edge_delta_loss: edge_delta_loss.to_vec(),
            surviving_edges,
            n_eff,
            occupancy: classify_occupancy(row_coordinates),
        })
    }

    /// ENROLL a discovered co-fire graph structure as a first-class graph atom
    /// (#985 / E1), rather than only detecting it.
    ///
    /// The **vertices are the atoms** and the **candidate edges are the co-firing
    /// pairs** returned by
    /// [`crate::atom_codes::SparseAtomCodes::coactive_pair_stats`] whose symmetric
    /// code dependence ([`crate::atom_codes::CoactivationStats::dependence`])
    /// clears `dependence_floor` — the sparse candidate superset, the co-fire
    /// analogue of the anchor-kNN edge set. Each surviving candidate edge's
    /// precision and REML deletion loss are supplied by
    /// [`coactivation_edge_evidence`]. The SAME rank-charge ARD survival rule as
    /// every other graph atom then keeps only the edges whose co-fire evidence
    /// (`n_eff·MI`) outweighs the one-edge charge, so enrollment *proposes* the
    /// co-fire graph and the REML currency *disposes*: the structure is enrolled
    /// into the dictionary as a `LearnedGraphAtom`, and its Betti numbers are read
    /// back from the surviving edges via [`Self::topology_readout`].
    ///
    /// `anchor_embeddings` is `n_atoms × r` (one fiber-`r` vertex embedding per
    /// atom); `row_coordinates` classifies the occupancy law as for any graph
    /// atom. Errors when a co-fire pair indexes a non-existent atom or no pair
    /// clears the dependence floor (nothing to enroll).
    pub fn enroll_from_coactivation(
        anchor_embeddings: ArrayView2<'_, f64>,
        row_coordinates: &[f64],
        n_eff: f64,
        coactive_pairs: &[(usize, usize, crate::atom_codes::CoactivationStats)],
        dependence_floor: f64,
    ) -> Result<Self, String> {
        validate_anchor_embeddings(anchor_embeddings)?;
        let anchors = anchor_embeddings.nrows();
        let mut candidate_edges = Vec::new();
        let mut edge_precisions = Vec::new();
        let mut edge_delta_loss = Vec::new();
        for (a, b, stats) in coactive_pairs {
            if stats.dependence() < dependence_floor {
                continue;
            }
            if *a >= anchors || *b >= anchors || a == b {
                return Err(format!(
                    "enroll_from_coactivation: co-fire pair ({a}, {b}) is invalid for {anchors} atoms"
                ));
            }
            let edge = GraphEdge::new(*a, *b)?;
            if candidate_edges.contains(&edge) {
                continue;
            }
            let (precision, delta) = coactivation_edge_evidence(stats, n_eff);
            candidate_edges.push(edge);
            edge_precisions.push(precision);
            edge_delta_loss.push(delta);
        }
        if candidate_edges.is_empty() {
            return Err(format!(
                "enroll_from_coactivation: no co-fire pair cleared the dependence floor {dependence_floor}"
            ));
        }
        Self::from_reml_candidate_edges(
            anchor_embeddings,
            row_coordinates,
            n_eff,
            &candidate_edges,
            &edge_precisions,
            &edge_delta_loss,
        )
    }

    pub fn anchors(&self) -> usize {
        self.anchor_embeddings.nrows()
    }

    pub fn fiber_rank(&self) -> usize {
        self.anchor_embeddings.ncols()
    }

    pub fn n_eff(&self) -> f64 {
        self.n_eff
    }

    pub fn one_edge_charge(&self) -> f64 {
        graph_edge_rank_charge(self.n_eff, self.fiber_rank())
    }

    pub fn summed_edge_charge(&self) -> f64 {
        self.one_edge_charge() * self.topology_readout().surviving_edges as f64
    }

    pub fn occupancy(&self) -> OccupancyLaw {
        self.occupancy
    }

    pub fn candidate_edges(&self) -> &[GraphEdge] {
        &self.candidate_edges
    }

    pub fn edge_precisions(&self) -> &[f64] {
        &self.edge_precisions
    }

    pub fn edge_delta_loss(&self) -> &[f64] {
        &self.edge_delta_loss
    }

    pub fn surviving_edges(&self) -> &[bool] {
        &self.surviving_edges
    }

    /// Vertex degrees in the surviving graph.
    pub fn surviving_degrees(&self) -> Vec<usize> {
        let mut degrees = vec![0usize; self.anchors()];
        for (idx, edge) in self.candidate_edges.iter().enumerate() {
            if self.surviving_edges[idx] {
                degrees[edge.a] += 1;
                degrees[edge.b] += 1;
            }
        }
        degrees
    }

    /// Weighted graph Laplacian `L_W` over all non-retired edges.
    pub fn surviving_laplacian(&self) -> Array2<f64> {
        self.weighted_laplacian_from_mask(&self.surviving_edges)
    }

    /// Weighted graph Laplacian `L_W` before edge retirement.
    pub fn full_laplacian(&self) -> Array2<f64> {
        let all_edges = vec![true; self.candidate_edges.len()];
        self.weighted_laplacian_from_mask(&all_edges)
    }

    /// Smoothness value `beta^T (L_W kron I_r) beta =
    /// sum_e w_e ||beta_i-beta_j||^2` over the surviving edge set.
    pub fn surviving_smoothness_value(&self) -> f64 {
        self.smoothness_value_from_mask(&self.surviving_edges)
    }

    /// Exact `O(E)` topology read-out from the surviving edge set. No persistence
    /// and no topology menu are involved.
    pub fn topology_readout(&self) -> GraphTopologyReadout {
        let vertices = self.anchors();
        let mut parent: Vec<usize> = (0..vertices).collect();
        let mut surviving_edges = 0usize;

        for (idx, edge) in self.candidate_edges.iter().enumerate() {
            if self.surviving_edges[idx] {
                surviving_edges += 1;
                graph_union(&mut parent, edge.a, edge.b);
            }
        }

        let mut roots = Vec::with_capacity(vertices);
        for vertex in 0..vertices {
            let root = graph_find(&mut parent, vertex);
            if !roots.contains(&root) {
                roots.push(root);
            }
        }
        let b0 = roots.len();
        let b1 = surviving_edges + b0 - vertices;
        GraphTopologyReadout {
            vertices,
            surviving_edges,
            b0,
            b1,
        }
    }

    /// Certified named compression of the learned graph. A positive bit saving
    /// returns the named shape; otherwise the structure remains an unnamed graph.
    pub fn certified_compression(&self) -> GraphCompressionReport {
        let readout = self.topology_readout();
        let degrees = self.surviving_degrees();
        let max_edges = readout
            .vertices
            .saturating_mul(readout.vertices.saturating_sub(1))
            / 2;
        let generic = crate::description_length::selection_bits(
            max_edges as i64,
            readout.surviving_edges as i64,
        );
        let log_vertices = (readout.vertices.max(2) as f64).log2();
        let named = if readout.b0 == 1
            && readout.b1 == 1
            && degrees.iter().all(|&d| d == 2)
            && self.surviving_edge_weights_are_uniform()
        {
            Some((GraphCompressionKind::Circle, "circle", log_vertices))
        } else if readout.b0 == 1
            && readout.b1 == 0
            && readout.vertices >= 2
            && degrees.iter().filter(|&&d| d == 1).count() == 2
            && degrees.iter().filter(|&&d| d == 2).count() == readout.vertices.saturating_sub(2)
            && self.surviving_edge_weights_are_uniform()
        {
            Some((
                GraphCompressionKind::Interval,
                "interval",
                2.0 * log_vertices,
            ))
        } else if matches!(
            self.occupancy,
            OccupancyLaw::Discrete { anchors } if anchors == readout.vertices
        ) && readout.b1 == 0
        {
            Some((GraphCompressionKind::FiniteSet, "finite_set", log_vertices))
        } else {
            None
        };
        if let Some((kind, name, named_bits)) = named {
            let report = GraphCompressionReport::certified(kind, name, generic, named_bits);
            if report.bits_saved > 0.0 {
                return report;
            }
        }
        GraphCompressionReport::unnamed(generic)
    }

    /// Birth-selection readout: graph existence is paid for by the sum of the
    /// surviving one-edge charges, not by a fixed topology menu.
    pub fn structure_selection(&self) -> GraphStructureSelection {
        let topology = self.topology_readout();
        let total_edge_delta_loss = self
            .edge_delta_loss
            .iter()
            .zip(self.surviving_edges.iter())
            .filter_map(|(&delta, &survives)| survives.then_some(delta))
            .sum::<f64>();
        let total_edge_charge = self.one_edge_charge() * topology.surviving_edges as f64;
        let margin = total_edge_delta_loss - total_edge_charge;
        GraphStructureSelection {
            selected: topology.surviving_edges > 0 && margin > 0.0,
            total_edge_delta_loss,
            total_edge_charge,
            margin,
            topology,
            occupancy: self.occupancy,
            compression: self.certified_compression(),
        }
    }

    pub fn surviving_penalty_op(
        &self,
        global_offset: usize,
        beta_dim: usize,
    ) -> Arc<dyn gam_solve::arrow_schur::BetaPenaltyOp> {
        Arc::new(IdentityRightKroneckerPenaltyOp {
            factor_a: self.surviving_laplacian(),
            p: self.fiber_rank(),
            global_offset,
            k: beta_dim,
        })
    }

    fn weighted_laplacian_from_mask(&self, active_edges: &[bool]) -> Array2<f64> {
        let anchors = self.anchors();
        let mut laplacian = Array2::<f64>::zeros((anchors, anchors));
        for (idx, edge) in self.candidate_edges.iter().enumerate() {
            if !active_edges[idx] {
                continue;
            }
            let w = self.edge_precisions[idx];
            if w == 0.0 {
                continue;
            }
            laplacian[[edge.a, edge.a]] += w;
            laplacian[[edge.b, edge.b]] += w;
            laplacian[[edge.a, edge.b]] -= w;
            laplacian[[edge.b, edge.a]] -= w;
        }
        laplacian
    }

    fn smoothness_value_from_mask(&self, active_edges: &[bool]) -> f64 {
        let fiber_rank = self.fiber_rank();
        let mut value = 0.0_f64;
        for (idx, edge) in self.candidate_edges.iter().enumerate() {
            if !active_edges[idx] {
                continue;
            }
            let w = self.edge_precisions[idx];
            if w == 0.0 {
                continue;
            }
            for channel in 0..fiber_rank {
                let diff = self.anchor_embeddings[[edge.a, channel]]
                    - self.anchor_embeddings[[edge.b, channel]];
                value += w * diff * diff;
            }
        }
        value
    }

    fn surviving_edge_weights_are_uniform(&self) -> bool {
        let mut min_weight = f64::INFINITY;
        let mut max_weight = f64::NEG_INFINITY;
        let mut count = 0usize;
        for (idx, survives) in self.surviving_edges.iter().enumerate() {
            if *survives {
                let weight = self.edge_precisions[idx];
                min_weight = min_weight.min(weight);
                max_weight = max_weight.max(weight);
                count += 1;
            }
        }
        if count == 0 {
            return false;
        }
        let scale = max_weight.abs().max(min_weight.abs()).max(1.0);
        max_weight - min_weight <= f64::EPSILON * scale * count as f64
    }
}

fn validate_anchor_embeddings(anchor_embeddings: ArrayView2<'_, f64>) -> Result<(), String> {
    let anchors = anchor_embeddings.nrows();
    let fiber_rank = anchor_embeddings.ncols();
    if anchors < 2 {
        return Err(format!(
            "LearnedGraphAtom requires at least 2 anchors; got {anchors}"
        ));
    }
    if fiber_rank == 0 {
        return Err("LearnedGraphAtom requires fiber_rank >= 1".to_string());
    }
    if anchor_embeddings.iter().any(|v| !v.is_finite()) {
        return Err("LearnedGraphAtom anchor_embeddings contain a non-finite value".to_string());
    }
    Ok(())
}

fn graph_find(parent: &mut [usize], x: usize) -> usize {
    let mut root = x;
    while parent[root] != root {
        root = parent[root];
    }
    let mut cur = x;
    while parent[cur] != root {
        let next = parent[cur];
        parent[cur] = root;
        cur = next;
    }
    root
}

fn graph_union(parent: &mut [usize], a: usize, b: usize) {
    let ra = graph_find(parent, a);
    let rb = graph_find(parent, b);
    if ra != rb {
        parent[rb] = ra;
    }
}

// ===========================================================================
// SPECTRAL DECODE — the open-world atom's basis and out-of-sample coordinate.
//
// The learned graph atom certifies and prices topology (Betti read-out, named-
// shape MDL) but by itself cannot *reconstruct*: it has no basis `Φ` and no
// per-row coordinate. The spectral decode closes that gap without leaving the
// currency the atom already uses:
//
//   * BASIS — the leading `q` non-trivial eigenvectors of the SAME survived,
//     ARD-weighted Laplacian `L_W` that assembles the smoothness penalty. In
//     that eigenbasis the Dirichlet form IS the penalty: `Φᵀ L_W Φ = diag(λ)`
//     (a diagonalisation, not a parallel computation) — see
//     [`GraphSpectralBasis::penalty`].
//   * COORDINATE — a differentiable Nyström (geometric-harmonics) extension of
//     those eigenvectors to any out-of-sample row `z`, with an analytic jet.
//   * RACE — a decodable candidate ([`SpectralGraphRaceCandidate`]) presenting
//     the same {basis eval, penalty, jet, rank charge} interface the typed
//     atoms hand the birth topology race.
// ===========================================================================

/// Upper cap on the spectral decode dimension `q`. The eigengap rule
/// ([`select_spectral_q`]) never keeps more than this many non-trivial modes:
/// a decode coordinate is meant to be a *small* intrinsic chart (a circle is
/// `q = 2`, a torus `q = 4`), and the pricing charges every kept mode
/// ([`spectral_decode_rank_charge`]), so an unbounded `q` would both defeat the
/// compression story and make the Nyström jet needlessly wide.
pub const SPECTRAL_DECODE_MAX_Q: usize = 8;

/// Rank charge of a `q`-dimensional spectral decode under the graph atom's own
/// currency. The atom prices one graph edge at `0.5·d_eff·ln(n_eff)`
/// ([`graph_edge_rank_charge`]); a spectral decode's `d_eff` is the number of
/// decode coordinates `q` (each eigenvector is one independent latent axis the
/// reconstruction spends), so the charge is the identical BIC/Laplace form with
/// `q` in the `d_eff` slot. Reported by [`GraphSpectralBasis::rank_charge_dof`]
/// and carried on [`SpectralGraphRaceCandidate`] so the race compares the
/// spectral alternative against the typed atoms in one commensurable currency.
pub fn spectral_decode_rank_charge(n_eff: f64, q: usize) -> f64 {
    0.5 * q as f64 * n_eff.max(2.0).ln()
}

/// Eigengap rule for the spectral decode dimension `q`.
///
/// `mu` is the ascending list of *non-trivial* Laplacian eigenvalues (the
/// `b0` near-zero constant-per-component modes already stripped). We keep the
/// number of leading modes that sits just before the largest **multiplicative**
/// spectral gap: `q = argmax_{1 ≤ k < cap} μ_{k+1} / μ_k`, with
/// `cap = min(mu.len(), SPECTRAL_DECODE_MAX_Q)`.
///
/// The ratio (not the additive gap `μ_{k+1} − μ_k`) is used because it is
/// scale-free: multiplying every edge precision by a constant rescales `L_W`
/// and every eigenvalue by that constant, and the ratio is invariant, so the
/// selected `q` does not drift with the ARD precisions' overall magnitude.
///
/// For a clean cycle (a circle) the first non-trivial eigenvalue is *doubly
/// degenerate* — the `[cos θ, sin θ]` pair — so `μ_2/μ_1 ≈ 1` while
/// `μ_3/μ_2 ≈ (1−cos 4π/N)/(1−cos 2π/N) ≈ 4`; the largest gap is after the
/// pair, giving `q = 2` (the minimal embedding dimension of `S¹`). A path /
/// interval has a single dominant Fiedler mode and a large `μ_2/μ_1`, giving
/// `q = 1`.
fn select_spectral_q(mu: &[f64], q_max: usize) -> usize {
    let cap = q_max.min(mu.len());
    if cap <= 1 {
        return cap.max(1);
    }
    let tiny = f64::MIN_POSITIVE;
    let mut best_q = 1usize;
    let mut best_ratio = f64::NEG_INFINITY;
    for k in 1..cap {
        let ratio = mu[k] / mu[k - 1].max(tiny);
        if ratio > best_ratio {
            best_ratio = ratio;
            best_q = k;
        }
    }
    best_q
}

/// The eigengap-selected spectral decode basis of a learned graph atom: the
/// leading `q` non-trivial eigenvectors of the survived weighted Laplacian
/// `L_W`, evaluated at the graph vertices, together with the eigenvalues that
/// ARE the Dirichlet penalty in this basis and the Gaussian-affinity Nyström
/// data that extends it out of sample.
#[derive(Debug, Clone)]
pub struct GraphSpectralBasis {
    /// `Φ` at the graph vertices, shape `(anchors × q)`. Column `k` is the
    /// unit-norm eigenvector `v_k` of `L_W` at the `k`-th smallest non-trivial
    /// eigenvalue.
    basis_values: Array2<f64>,
    /// `λ_1 ≤ … ≤ λ_q`, the non-trivial Laplacian eigenvalues. `diag(λ)` is
    /// literally `Φᵀ L_W Φ` — the Dirichlet form of the basis columns and the
    /// decode penalty are the SAME object.
    eigenvalues: Vec<f64>,
    /// Training-vertex embeddings `x_i`, shape `(anchors × r)`; the Nyström
    /// anchors the out-of-sample kernel weights against.
    anchor_embeddings: Array2<f64>,
    /// Gaussian affinity bandwidth `ε` — the median squared length of the
    /// surviving graph edges (the same median-neighbour-distance rule the
    /// Laplacian-eigenmap seed uses, `gam_geometry::latent_seed`).
    bandwidth: f64,
    /// Effective sample mass carried from the atom, for the decode rank charge.
    n_eff: f64,
}

impl GraphSpectralBasis {
    /// Selected decode dimension `q`.
    pub fn selected_q(&self) -> usize {
        self.eigenvalues.len()
    }

    /// The kept non-trivial Laplacian eigenvalues `λ_1 ≤ … ≤ λ_q`.
    pub fn eigenvalues(&self) -> &[f64] {
        &self.eigenvalues
    }

    /// The decode basis `Φ` at the graph vertices, shape `(anchors × q)`.
    pub fn vertex_basis(&self) -> ArrayView2<'_, f64> {
        self.basis_values.view()
    }

    /// Nyström Gaussian bandwidth `ε`.
    pub fn bandwidth(&self) -> f64 {
        self.bandwidth
    }

    /// The decode penalty in the spectral basis: `diag(λ)`, shape `(q × q)`.
    ///
    /// This is not a second computation of the roughness — it is exactly the
    /// atom's Dirichlet form `Φᵀ L_W Φ` read in the eigenbasis, where `L_W` is
    /// the *same* survived weighted Laplacian [`LearnedGraphAtom::surviving_laplacian`]
    /// whose Kronecker lift `L_W ⊗ I_r` is the atom's smoothness penalty
    /// [`LearnedGraphAtom::surviving_penalty_op`]. Diagonalising `L_W` on its
    /// own eigenvectors returns `diag(λ)`, so the penalty a decode consumer
    /// reads here and the penalty the graph atom prices are one operator.
    pub fn penalty(&self) -> Array2<f64> {
        let q = self.selected_q();
        let mut penalty = Array2::<f64>::zeros((q, q));
        for k in 0..q {
            penalty[[k, k]] = self.eigenvalues[k];
        }
        penalty
    }

    /// The decode rank charge `0.5·q·ln(n_eff)` — see
    /// [`spectral_decode_rank_charge`].
    pub fn rank_charge_dof(&self) -> f64 {
        spectral_decode_rank_charge(self.n_eff, self.selected_q())
    }

    /// Out-of-sample decode coordinate `φ(z) ∈ ℝ^q` for a single query row `z`
    /// (`r` ambient features) plus its analytic jet `∂φ/∂z` (`q × r`).
    pub fn nystrom_coordinate(&self, z: &[f64]) -> Result<(Vec<f64>, Array2<f64>), String> {
        let r = self.anchor_embeddings.ncols();
        if z.len() != r {
            return Err(format!(
                "GraphSpectralBasis::nystrom_coordinate: query has {} features but graph anchors have {r}",
                z.len()
            ));
        }
        let query = ArrayView2::from_shape((1, r), z)
            .map_err(|e| format!("GraphSpectralBasis::nystrom_coordinate: bad query shape: {e}"))?;
        let (phi, jet) = self.nystrom_coordinates(query)?;
        let q = self.selected_q();
        let coord = phi.row(0).to_vec();
        let mut jac = Array2::<f64>::zeros((q, r));
        for k in 0..q {
            for c in 0..r {
                jac[[k, c]] = jet[[0, k, c]];
            }
        }
        Ok((coord, jac))
    }

    /// Batched Nyström extension of the decode basis to arbitrary out-of-sample
    /// rows `points` (`n × r`).
    ///
    /// The coordinate is the affinity-weighted (Nadaraya–Watson / geometric-
    /// harmonics) average of the training eigenvectors,
    ///
    /// ```text
    ///   φ_k(z) = N_k(z) / S(z),
    ///   N_k(z) = Σ_i w_i(z) Φ[i, k],   S(z) = Σ_i w_i(z),
    ///   w_i(z) = exp(−‖z − x_i‖² / ε),
    /// ```
    ///
    /// with `ε` the graph's edge-length bandwidth [`Self::bandwidth`]. At a
    /// training vertex the Gaussian mass concentrates on that vertex, so
    /// `φ_k(x_j) ≈ Φ[j, k]`, and between vertices it interpolates smoothly — the
    /// standard Nyström / geometric-harmonics extension of a graph eigenmap.
    ///
    /// The jet is analytic (no finite differences). With
    /// `∂w_i/∂z_c = w_i · (−2 (z_c − x_{i,c}) / ε)`,
    ///
    /// ```text
    ///   ∂φ_k/∂z_c = ( (∂N_k/∂z_c)·S − N_k·(∂S/∂z_c) ) / S²,
    ///   ∂N_k/∂z_c = Σ_i (∂w_i/∂z_c) Φ[i, k],   ∂S/∂z_c = Σ_i ∂w_i/∂z_c.
    /// ```
    ///
    /// Returns `(φ, ∂φ/∂z)` with shapes `(n × q)` and `(n × q × r)`.
    pub fn nystrom_coordinates(
        &self,
        points: ArrayView2<'_, f64>,
    ) -> Result<(Array2<f64>, Array3<f64>), String> {
        let anchors = self.anchor_embeddings.nrows();
        let r = self.anchor_embeddings.ncols();
        let q = self.selected_q();
        if points.ncols() != r {
            return Err(format!(
                "GraphSpectralBasis::nystrom_coordinates: query has {} features but graph anchors have {r}",
                points.ncols()
            ));
        }
        if points.iter().any(|v| !v.is_finite()) {
            return Err(
                "GraphSpectralBasis::nystrom_coordinates: query contains a non-finite value".into(),
            );
        }
        let n = points.nrows();
        let eps = self.bandwidth;
        if !(eps > 0.0 && eps.is_finite()) {
            return Err(format!(
                "GraphSpectralBasis::nystrom_coordinates: non-positive bandwidth {eps}"
            ));
        }
        let mut phi = Array2::<f64>::zeros((n, q));
        let mut jet = Array3::<f64>::zeros((n, q, r));
        let mut w = vec![0.0_f64; anchors];
        let mut n_k = vec![0.0_f64; q];
        let mut dw = vec![0.0_f64; anchors];
        for row in 0..n {
            // Affinities and the normaliser S(z).
            let mut s = 0.0_f64;
            for i in 0..anchors {
                let mut d2 = 0.0_f64;
                for c in 0..r {
                    let d = points[[row, c]] - self.anchor_embeddings[[i, c]];
                    d2 += d * d;
                }
                let wi = (-d2 / eps).exp();
                w[i] = wi;
                s += wi;
            }
            if !(s > 0.0 && s.is_finite()) {
                return Err(
                    "GraphSpectralBasis::nystrom_coordinates: query point has vanishing affinity \
                     to every anchor (bandwidth underflow)"
                        .into(),
                );
            }
            // Numerators N_k and the decode coordinate φ_k = N_k / S.
            for k in 0..q {
                let mut acc = 0.0_f64;
                for i in 0..anchors {
                    acc += w[i] * self.basis_values[[i, k]];
                }
                n_k[k] = acc;
                phi[[row, k]] = acc / s;
            }
            // Analytic jet, one ambient channel c at a time.
            for c in 0..r {
                let mut ds_c = 0.0_f64;
                for i in 0..anchors {
                    let dwi =
                        w[i] * (-2.0 * (points[[row, c]] - self.anchor_embeddings[[i, c]]) / eps);
                    dw[i] = dwi;
                    ds_c += dwi;
                }
                for k in 0..q {
                    let mut dn_kc = 0.0_f64;
                    for i in 0..anchors {
                        dn_kc += dw[i] * self.basis_values[[i, k]];
                    }
                    jet[[row, k, c]] = (dn_kc * s - n_k[k] * ds_c) / (s * s);
                }
            }
        }
        Ok((phi, jet))
    }

    /// The Nyström extension as a first-class [`SaeBasisEvaluator`], so the
    /// decode presents the exact {basis values, jet} interface the typed atoms'
    /// evaluators do. Its input coordinates are the ambient row embeddings `z`
    /// (`r` features), its output the `q`-dimensional decode coordinate.
    pub fn evaluator(&self) -> Arc<dyn SaeBasisEvaluator> {
        Arc::new(NystromSpectralEvaluator {
            basis: self.clone(),
        })
    }
}

/// [`SaeBasisEvaluator`] adapter over a [`GraphSpectralBasis`]: `evaluate`
/// returns the Nyström decode coordinate `φ(z)` and its analytic first jet
/// `∂φ/∂z` at each queried ambient row `z`. It declares no analytic second /
/// third jet (`None`): the spectral decode's roughness is the graph Dirichlet
/// form `diag(λ)` carried directly on [`GraphSpectralBasis::penalty`], not a
/// second-jet curvature Gram, so no consumer needs a Nyström Hessian and the
/// honest capability declaration is absence.
#[derive(Debug, Clone)]
pub struct NystromSpectralEvaluator {
    basis: GraphSpectralBasis,
}

impl SaeBasisEvaluator for NystromSpectralEvaluator {
    fn evaluate(&self, coords: ArrayView2<'_, f64>) -> Result<(Array2<f64>, Array3<f64>), String> {
        self.basis.nystrom_coordinates(coords)
    }

    fn second_jet_dyn(&self, coords: ArrayView2<'_, f64>) -> Option<Result<Array4<f64>, String>> {
        // A mismatched query width is a caller bug, not a missing capability:
        // surface it as an error exactly like `nystrom_coordinates` would,
        // instead of a silent "no jet" that sends the caller down a fallback.
        let r = self.basis.anchor_embeddings.ncols();
        if coords.ncols() != r {
            return Some(Err(format!(
                "NystromSpectralEvaluator::second_jet_dyn: query has {} features but graph anchors have {r}",
                coords.ncols()
            )));
        }
        None
    }

    fn third_jet_dyn(
        &self,
        coords: ArrayView2<'_, f64>,
    ) -> Option<Result<ndarray::Array5<f64>, String>> {
        let r = self.basis.anchor_embeddings.ncols();
        if coords.ncols() != r {
            return Some(Err(format!(
                "NystromSpectralEvaluator::third_jet_dyn: query has {} features but graph anchors have {r}",
                coords.ncols()
            )));
        }
        None
    }
}

/// A spectral-graph decode candidate shaped for the birth topology race.
///
/// This mirrors the private `TopologyRaceFit` shape the typed candidates carry
/// in [`crate::structure_harvest`] — evaluator, basis kind, latent manifold,
/// decode design `Φ`, jet `∂Φ`, penalized decoder `B`, roughness penalty — but
/// for the open-world graph atom, whose penalty is the graph Dirichlet form
/// `diag(λ)` rather than a basis second-jet Gram.
///
/// # Where this plugs into the race
///
/// The single call site is
/// [`crate::structure_harvest::topology_candidates_for_dim`] (crate-internal,
/// concurrently edited elsewhere, so it is NOT touched here). After that
/// function builds the typed `TopologyCandidateSpec`s for the born atom's `d_k`,
/// a spectral candidate is appended when the born atom already carries a
/// [`LearnedGraphAtom`], by calling
/// [`LearnedGraphAtom::spectral_race_candidate`] with the birth target and the
/// per-row ambient embeddings; its `{phi, jet, penalty, rank_charge_dof}` feed
/// the SAME `TopologyAutoFitEvidence` inputs `fit_topology_candidate` produces —
/// with `penalty` supplied directly instead of re-derived from a second jet —
/// and its `evaluator` seeds the born atom's out-of-sample decode. Everything
/// up to that append (basis, penalty, decoder, jet, charge, evaluator) is
/// implemented here; only the one `specs.push(...)` line lives across the seam.
pub struct SpectralGraphRaceCandidate {
    /// Basis-kind tag; a precomputed decode basis with no closed-form typed
    /// evaluator family.
    pub basis_kind: SaeAtomBasisKind,
    /// The flat `q`-coordinate decode chart the atom carries.
    pub manifold: LatentManifold,
    /// Decode dimension `q`.
    pub latent_dim: usize,
    /// The ambient row embeddings `z` (`n × r`) the Nyström evaluator reads.
    pub row_coords: Array2<f64>,
    /// Decode design `Φ(z)` at the rows (`n × q`).
    pub phi: Array2<f64>,
    /// Decode design jet `∂Φ/∂z` (`n × q × r`).
    pub jet: Array3<f64>,
    /// Dirichlet-penalized decoder `B` (`q × p`) at the REML-optimal `λ̂` on the
    /// spectral penalty — fit through the SAME closed-form entry point the typed
    /// candidates use, so the decode is priced commensurably.
    pub decoder: Array2<f64>,
    /// The decode penalty `diag(λ)` (`q × q`) — the graph Dirichlet form.
    pub penalty: Array2<f64>,
    /// The Nyström out-of-sample decode map.
    pub evaluator: Arc<dyn SaeBasisEvaluator>,
    /// The decode rank charge `0.5·q·ln(n_eff)`.
    pub rank_charge_dof: f64,
}

impl LearnedGraphAtom {
    /// Median squared length of the surviving graph edges — the Nyström Gaussian
    /// bandwidth `ε`. Mirrors the median-neighbour-distance bandwidth the
    /// Laplacian-eigenmap seed uses (`gam_geometry::latent_seed`): a scale-free,
    /// data-driven choice that keeps the out-of-sample affinities matched to the
    /// graph's own locality.
    fn surviving_edge_bandwidth(&self) -> Result<f64, String> {
        let fiber_rank = self.fiber_rank();
        let mut lengths = Vec::new();
        for (idx, edge) in self.candidate_edges.iter().enumerate() {
            if !self.surviving_edges[idx] {
                continue;
            }
            let mut d2 = 0.0_f64;
            for c in 0..fiber_rank {
                let d = self.anchor_embeddings[[edge.a, c]] - self.anchor_embeddings[[edge.b, c]];
                d2 += d * d;
            }
            lengths.push(d2);
        }
        if lengths.is_empty() {
            return Err(
                "LearnedGraphAtom::spectral_decode_basis: no surviving edge to set the Nyström \
                 bandwidth"
                    .into(),
            );
        }
        lengths.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        let median = lengths[lengths.len() / 2];
        if !(median > 0.0 && median.is_finite()) {
            return Err(format!(
                "LearnedGraphAtom::spectral_decode_basis: degenerate surviving-edge bandwidth {median}"
            ));
        }
        Ok(median)
    }

    /// Build the eigengap-selected [`GraphSpectralBasis`] from the survived,
    /// ARD-weighted Laplacian `L_W`.
    ///
    /// Diagonalises the SAME `L_W = ` [`Self::surviving_laplacian`] the atom's
    /// smoothness penalty lifts, strips the `b0` near-zero constant-per-component
    /// null modes, keeps the leading `q` non-trivial eigenvectors chosen by
    /// [`select_spectral_q`], and reports their eigenvalues as the decode penalty
    /// (`diag(λ)`). Errors when the survived graph has no non-trivial mode (every
    /// vertex isolated, or a single constant component) — such a graph has no
    /// continuous coordinate to decode.
    pub fn spectral_decode_basis(&self) -> Result<GraphSpectralBasis, String> {
        let anchors = self.anchors();
        let laplacian = self.surviving_laplacian();
        let (evals, evecs) = laplacian.eigh(Side::Lower).map_err(|e| {
            format!(
                "LearnedGraphAtom::spectral_decode_basis: Laplacian eigendecomposition failed: {e}"
            )
        })?;
        // Ascending eigenvalue order (faer does not guarantee it).
        let mut order: Vec<usize> = (0..evals.len()).collect();
        order.sort_by(|&a, &b| {
            evals[a]
                .partial_cmp(&evals[b])
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        let lambda_max = order.last().map(|&i| evals[i]).unwrap_or(0.0).max(0.0);
        // Trivial (null) modes: one constant per connected component, at ~0
        // eigenvalue. Strip them with a relative threshold.
        let zero_tol = (lambda_max * 1e-8).max(1e-12);
        let nontrivial: Vec<usize> = order
            .iter()
            .copied()
            .filter(|&i| evals[i] > zero_tol)
            .collect();
        if nontrivial.is_empty() {
            return Err(
                "LearnedGraphAtom::spectral_decode_basis: survived graph has no non-trivial \
                 Laplacian mode to decode (all vertices isolated or a single constant component)"
                    .into(),
            );
        }
        let sorted_mu: Vec<f64> = nontrivial.iter().map(|&i| evals[i]).collect();
        let q = select_spectral_q(&sorted_mu, SPECTRAL_DECODE_MAX_Q);
        let mut basis_values = Array2::<f64>::zeros((anchors, q));
        for (col, &eig_idx) in nontrivial.iter().take(q).enumerate() {
            for v in 0..anchors {
                basis_values[[v, col]] = evecs[[v, eig_idx]];
            }
        }
        let eigenvalues: Vec<f64> = sorted_mu.iter().take(q).copied().collect();
        let bandwidth = self.surviving_edge_bandwidth()?;
        Ok(GraphSpectralBasis {
            basis_values,
            eigenvalues,
            anchor_embeddings: self.anchor_embeddings.clone(),
            bandwidth,
            n_eff: self.n_eff,
        })
    }

    /// Realise the spectral decode as a birth-race candidate: the eigengap decode
    /// basis, the Nyström design `Φ(z)` + jet at the birth rows, the graph
    /// Dirichlet penalty `diag(λ)`, and a REML-optimal Dirichlet-penalized
    /// decoder fit against `target` (`n × p`) through the same closed-form entry
    /// point the typed candidates use. See [`SpectralGraphRaceCandidate`] for the
    /// exact call site this plugs into.
    pub fn spectral_race_candidate(
        &self,
        target: ArrayView2<'_, f64>,
        row_embeddings: ArrayView2<'_, f64>,
    ) -> Result<SpectralGraphRaceCandidate, String> {
        let basis = self.spectral_decode_basis()?;
        let (phi, jet) = basis.nystrom_coordinates(row_embeddings)?;
        let n = target.nrows();
        if phi.nrows() != n {
            return Err(format!(
                "LearnedGraphAtom::spectral_race_candidate: {n} targets but {} row embeddings",
                phi.nrows()
            ));
        }
        let penalty = basis.penalty();
        let reml = gam_solve::gaussian_reml::gaussian_reml_multi_closed_form(
            phi.view(),
            target,
            penalty.view(),
            None,
            None,
        )
        .map_err(|e| {
            format!("LearnedGraphAtom::spectral_race_candidate: REML decode fit: {e:?}")
        })?;
        Ok(SpectralGraphRaceCandidate {
            basis_kind: SaeAtomBasisKind::Precomputed("spectral_graph".to_string()),
            manifold: LatentManifold::Euclidean,
            // The candidate's coordinates ARE the ambient row embeddings (n × r),
            // and the Nyström jet differentiates w.r.t. those r embedding
            // channels — so the declared latent dimension is r, matching
            // `row_coords`/`jet`, NOT the spectral basis width q (which lives in
            // `phi`/`decoder`/`rank_charge_dof`). Declaring q here handed
            // consumers a `(n, q, r)` jet labelled `(n, q, q)`.
            latent_dim: row_embeddings.ncols(),
            row_coords: row_embeddings.to_owned(),
            phi,
            jet,
            decoder: reml.coefficients,
            penalty,
            evaluator: basis.evaluator(),
            rank_charge_dof: basis.rank_charge_dof(),
        })
    }
}