gam-sae 0.3.155

//! Unified support-sparse assignment state for the SAE fit-path unification
//! (design: gam#2232, Increment 1).
//!
//! [`SaeAssignmentState`] is the ONE per-row routing state the unified engine
//! carries: per row it stores the active atom set `S_i ⊆ [K]` as
//!
//!   * `indices[i]`      — the active atom indices (`u32`), `|S_i|` of them;
//!   * `gate_params[i]`  — one routing scalar per active atom (the dense
//!     `logits[i, k]` value; the realized gate is derived from it through the
//!     [`AssignmentMode`], so this is the *fundamental* stored parameter, and it
//!     is exactly the `f64` the `SaeTopKCurvedBudget` "gate values" slot
//!     budgets, one per active atom — see the layout contract below);
//!   * `coords[i]`       — the on-manifold coordinates of ONLY the active atoms,
//!     flattened in support order (`Σ_{k∈S_i} d_k` scalars per row).
//!
//! The dense [`SaeAssignment`] (`assignment.rs`) is the FULL-SUPPORT
//! materialization of this state: `S_i = [0, K)` for every row, so the
//! per-`(row, atom)` coordinate storage is the transpose of the dense
//! `Vec<LatentCoordValues>` per-atom blocks, and `gate_params[i]` is the dense
//! `logits` row. `SaeAssignmentState::materialize_dense` reconstructs that dense
//! layout bit-for-bit, and `SaeAssignment::as_state` is its inverse.
//!
//! # Layout contract vs. the `SaeTopKCurvedBudget` ledger
//!
//! `manifold/streaming_plan.rs` charges the honest support-sparse assignment
//! state as
//!
//! ```text
//!   active_state_bytes = N · k_active · (2 + d_max) · 8
//! ```
//!
//! (`sae_topk_curved_budget_from_budget`, streaming_plan.rs:482-485). The three
//! `(2 + d_max)` terms are exactly this state's three per-row arrays: `1` for
//! `indices`, `1` for `gate_params`, and `d_max` for `coords` — every cell an
//! 8-byte word (`SAE_BYTES_PER_F64`; the `u32` index cell is budgeted as a full
//! 8-byte slot, matching the ledger's uniform-word accounting). For a uniform
//! `k_active = s`, `d_k = d_max` shape this state therefore occupies
//! `SaeAssignmentState::active_state_bytes`
//! `= N · s · (2 + d_max) · 8 = active_state_bytes`, verified by
//! `sparse_topk_state_memory_shape_matches_budget_formula`.

use gam_problem::LatentRetractionRegistry;
use gam_terms::latent::{LatentIdMode, LatentManifold};
use ndarray::Array1;

use crate::assignment::{AssignmentMode, SaeAssignment};

/// Per-atom coordinate metadata needed to reconstruct the dense
/// [`LatentCoordValues`] block bit-for-bit on materialization.
#[derive(Debug, Clone)]
struct AtomCoordMeta {
    latent_dim: usize,
    manifold: LatentManifold,
    retraction: LatentRetractionRegistry,
}

/// Coordinate geometry for one atom in a support-sparse assignment.
///
/// Unlike the former `d_max` constructor, this is indexed by atom and therefore
/// preserves mixed intrinsic dimensions and topologies without padding inactive
/// coordinates. The retraction and stable identity travel with the coordinate
/// block so a later full-support specialization remains an exact inverse of the
/// dense representation.
#[derive(Debug, Clone)]
pub struct SaeAssignmentAtomSpec {
    pub latent_dim: usize,
    pub id_mode: LatentIdMode,
    pub manifold: LatentManifold,
    pub retraction: LatentRetractionRegistry,
    pub latent_id: u64,
}

impl SaeAssignmentAtomSpec {
    /// Euclidean atom metadata for uniform-dimension callers.
    #[must_use]
    pub fn euclidean(latent_dim: usize) -> Self {
        Self {
            latent_dim,
            id_mode: LatentIdMode::None,
            manifold: LatentManifold::Euclidean,
            retraction: LatentRetractionRegistry::all_euclidean(),
            latent_id: 0,
        }
    }

}

/// Support-sparse per-row assignment state (see module docs). Internal type: the
/// unified engine's ONE routing state, of which the dense [`SaeAssignment`] is
/// the full-support specialization.
#[derive(Debug, Clone)]
pub struct SaeAssignmentState {
    n_obs: usize,
    k_atoms: usize,
    /// Active atom indices per row (`indices[i]` sorted-ascending for a
    /// full-support state; TopK constructors pass the routed support).
    indices: Vec<Vec<u32>>,
    /// Routing scalar per active atom per row — the dense `logits[i, k]` value.
    gate_params: Vec<Vec<f64>>,
    /// Active-atom coordinates per row, flattened in support order: for row `i`
    /// the concatenation over `j` of the `d_{indices[i][j]}` coordinate scalars.
    coords: Vec<Vec<f64>>,
    /// Per-atom coordinate metadata (length `K`) for dense reconstruction.
    atom_coord_meta: Vec<AtomCoordMeta>,
    mode: AssignmentMode,
}

impl SaeAssignmentState {

    /// Construct the canonical hard-TopK state for heterogeneous atoms.
    ///
    /// Every input row contains exactly `support_k` distinct atom indices. Its
    /// coordinate row is the concatenation of those atoms' unpadded coordinate
    /// blocks in the same order. Construction sorts each support by atom index
    /// and moves the corresponding gate/coordinate blocks with it, so logically
    /// equivalent routings have one deterministic representation.
    #[must_use = "state build error must be handled"]
    pub fn from_topk_support_heterogeneous(
        n_obs: usize,
        k_atoms: usize,
        support_k: usize,
        atom_specs: Vec<SaeAssignmentAtomSpec>,
        mut indices: Vec<Vec<u32>>,
        mut gate_params: Vec<Vec<f64>>,
        mut coords: Vec<Vec<f64>>,
    ) -> Result<Self, String> {
        if support_k == 0 || support_k > k_atoms {
            return Err(format!(
                "SaeAssignmentState::from_topk_support_heterogeneous: support_k must satisfy 1 <= s <= K={k_atoms}; got {support_k}"
            ));
        }
        if atom_specs.len() != k_atoms {
            return Err(format!(
                "SaeAssignmentState::from_topk_support_heterogeneous: atom_specs length {} must equal K={k_atoms}",
                atom_specs.len()
            ));
        }
        for (atom, spec) in atom_specs.iter().enumerate() {
            if spec.latent_dim == 0 {
                return Err(format!(
                    "SaeAssignmentState::from_topk_support_heterogeneous: atom {atom} latent_dim must be positive"
                ));
            }
            let manifold_dim = spec.manifold.ambient_dim(spec.latent_dim);
            if manifold_dim != spec.latent_dim {
                return Err(format!(
                    "SaeAssignmentState::from_topk_support_heterogeneous: atom {atom} manifold ambient dimension {manifold_dim} != latent_dim {}",
                    spec.latent_dim
                ));
            }
            spec.retraction.validate_dim(
                spec.latent_dim,
                "SaeAssignmentState::from_topk_support_heterogeneous",
            )?;
        }
        if indices.len() != n_obs || gate_params.len() != n_obs || coords.len() != n_obs {
            return Err(format!(
                "SaeAssignmentState::from_topk_support_heterogeneous: per-row arrays must all have length N={n_obs}; \
                 got indices={}, gate_params={}, coords={}",
                indices.len(),
                gate_params.len(),
                coords.len()
            ));
        }
        for i in 0..n_obs {
            if indices[i].len() > support_k
                || indices[i].is_empty()
                || gate_params[i].len() != indices[i].len()
            {
                return Err(format!(
                    "SaeAssignmentState::from_topk_support_heterogeneous: row {i} widths must be indices={support_k}, gate_params={support_k}; got {}, {}",
                    indices[i].len(),
                    gate_params[i].len(),
                ));
            }
            if gate_params[i].iter().any(|value| !value.is_finite()) {
                return Err(format!(
                    "SaeAssignmentState::from_topk_support_heterogeneous: row {i} contains a non-finite gate parameter"
                ));
            }
            if coords[i].iter().any(|value| !value.is_finite()) {
                return Err(format!(
                    "SaeAssignmentState::from_topk_support_heterogeneous: row {i} contains a non-finite coordinate"
                ));
            }

            let mut coord_cursor = 0usize;
            let mut slots = Vec::with_capacity(indices[i].len());
            for slot in 0..indices[i].len() {
                let atom = indices[i][slot] as usize;
                if atom >= k_atoms {
                    return Err(format!(
                        "SaeAssignmentState::from_topk_support_heterogeneous: row {i} atom index {atom} out of range K={k_atoms}"
                    ));
                }
                let d = atom_specs[atom].latent_dim;
                let end = coord_cursor.saturating_add(d);
                if end > coords[i].len() {
                    return Err(format!(
                        "SaeAssignmentState::from_topk_support_heterogeneous: row {i} coordinate width {} is too short for its declared support",
                        coords[i].len()
                    ));
                }
                slots.push((
                    atom as u32,
                    gate_params[i][slot],
                    coords[i][coord_cursor..end].to_vec(),
                ));
                coord_cursor = end;
            }
            if coord_cursor != coords[i].len() {
                return Err(format!(
                    "SaeAssignmentState::from_topk_support_heterogeneous: row {i} coordinate width {} != support-implied width {coord_cursor}",
                    coords[i].len()
                ));
            }
            slots.sort_by_key(|slot| slot.0);
            if slots.windows(2).any(|pair| pair[0].0 == pair[1].0) {
                return Err(format!(
                    "SaeAssignmentState::from_topk_support_heterogeneous: row {i} support contains a duplicate atom"
                ));
            }
            indices[i] = slots.iter().map(|slot| slot.0).collect();
            gate_params[i] = slots.iter().map(|slot| slot.1).collect();
            coords[i] = slots.into_iter().flat_map(|slot| slot.2).collect();
        }
        let atom_coord_meta = atom_specs
            .into_iter()
            .map(|spec| AtomCoordMeta {
                latent_dim: spec.latent_dim,
                manifold: spec.manifold,
                retraction: spec.retraction,
            })
            .collect();
        Ok(Self {
            n_obs,
            k_atoms,
            indices,
            gate_params,
            coords,
            atom_coord_meta,
            mode: AssignmentMode::top_k_support(support_k),
        })
    }

    pub fn n_obs(&self) -> usize {
        self.n_obs
    }

    pub fn k_atoms(&self) -> usize {
        self.k_atoms
    }

    pub fn mode(&self) -> AssignmentMode {
        self.mode
    }

    /// Active atom indices `S_i` for `row`.
    pub fn support_indices(&self, row: usize) -> &[u32] {
        &self.indices[row]
    }

    /// Routing scalars (dense `logits` values) on `row`'s active support.
    pub fn gate_params(&self, row: usize) -> &[f64] {
        &self.gate_params[row]
    }

    /// The whole flattened active-atom coordinate block for `row`
    /// (`Σ_{k∈S_i} d_k` scalars, support order).
    pub fn coords_row(&self, row: usize) -> &[f64] {
        &self.coords[row]
    }

    /// Per-atom intrinsic dimension, with no `d_max` padding.
    pub fn atom_coord_dim(&self, atom: usize) -> usize {
        self.atom_coord_meta[atom].latent_dim
    }

    /// Effective per-axis periodicity for one atom, including a retraction
    /// override attached to an otherwise Euclidean coordinate block.
    ///
    /// `LatentManifold::Euclidean` is dimension-generic and therefore reports a
    /// single topology axis. The assignment state owns the concrete atom
    /// dimension, so every sparse-coordinate consumer must obtain its topology
    /// here rather than index the raw manifold descriptor directly.
    pub fn atom_axis_periods(&self, atom: usize) -> Vec<Option<f64>> {
        let meta = &self.atom_coord_meta[atom];
        let periods = if meta.manifold.is_euclidean() {
            meta.retraction.axis_periods(meta.latent_dim)
        } else {
            meta.manifold.axis_periods()
        };
        assert_eq!(
            periods.len(),
            meta.latent_dim,
            "SaeAssignmentState atom {atom} axis-period count {} != latent dimension {}",
            periods.len(),
            meta.latent_dim,
        );
        periods
    }

    /// Coordinate block for one active support slot.
    pub fn coords_for_slot(&self, row: usize, slot: usize) -> &[f64] {
        let start: usize = self.indices[row][..slot]
            .iter()
            .map(|&atom| self.atom_coord_meta[atom as usize].latent_dim)
            .sum();
        let atom = self.indices[row][slot] as usize;
        &self.coords[row][start..start + self.atom_coord_meta[atom].latent_dim]
    }

    /// Replace one compact coordinate row and project every heterogeneous atom
    /// block onto its declared manifold. This is the exact snapshot-restore
    /// operation required by support-sparse line searches: applying a negated
    /// step is not an inverse retraction at interval boundaries or on curved
    /// manifolds, so rollback must restore the accepted point itself.
    /// Move the per-row coordinate storage out for a disjoint-rows parallel
    /// sweep: rows own non-overlapping blocks, so a parallel solver may take
    /// the storage, hand each row its own block mutably, and put it back with
    /// [`Self::restore_coords`]. While taken, coordinate reads must go through
    /// the caller-held blocks.
    pub fn take_coords(&mut self) -> Vec<Vec<f64>> {
        std::mem::take(&mut self.coords)
    }

    /// Restore storage taken by [`Self::take_coords`].
    pub fn restore_coords(&mut self, coords: Vec<Vec<f64>>) -> Result<(), String> {
        if coords.len() != self.n_obs {
            return Err(format!(
                "SaeAssignmentState::restore_coords: {} rows != N={}",
                coords.len(),
                self.n_obs
            ));
        }
        self.coords = coords;
        Ok(())
    }

    /// `Self::apply_row_coord_step` against a caller-held coordinate block —
    /// the identical per-atom retraction.
    /// Project a compact row step onto each atom's tangent space at its current
    /// coordinates.
    ///
    /// [`Self::retract_row_coords`] travels only the tangent component of a
    /// step, because the manifold exponential map discards any part of the
    /// vector that leaves the tangent space. A caller that certifies descent
    /// against the un-projected step therefore certifies a direction it will
    /// not travel, and for a dim-2 sphere -- whose tangent space is a line --
    /// that silently drops one of the two coordinates being optimised. Callers
    /// project the step through here first so the certificate and the motion
    /// describe the same vector.
    ///
    /// For Euclidean atoms the tangent space is the whole ambient space and
    /// `project_tangent` is the identity, so this is a no-op on flat charts.
    pub fn project_row_tangent(
        &self,
        row: usize,
        coords: &[f64],
        delta: &mut [f64],
    ) -> Result<(), String> {
        if row >= self.n_obs {
            return Err(format!(
                "SaeAssignmentState::project_row_tangent: row {row} out of range N={}",
                self.n_obs
            ));
        }
        if delta.len() != coords.len() {
            return Err(format!(
                "SaeAssignmentState::project_row_tangent: row {row} delta width {} != compact coordinate width {}",
                delta.len(),
                coords.len()
            ));
        }
        let mut cursor = 0usize;
        for &atom in &self.indices[row] {
            let meta = &self.atom_coord_meta[atom as usize];
            let end = cursor + meta.latent_dim;
            let point = Array1::from_vec(coords[cursor..end].to_vec());
            let vector = Array1::from_vec(delta[cursor..end].to_vec());
            let projected = meta
                .manifold
                .project_to_tangent(point.view(), vector.view());
            delta[cursor..end].copy_from_slice(
                projected
                    .as_slice()
                    .expect("tangent projection is contiguous"),
            );
            cursor = end;
        }
        Ok(())
    }

    pub fn retract_row_coords(
        &self,
        row: usize,
        coords: &mut [f64],
        delta: &[f64],
    ) -> Result<(), String> {
        if delta.len() != coords.len() {
            return Err(format!(
                "SaeAssignmentState::retract_row_coords: row {row} delta width {} != compact coordinate width {}",
                delta.len(),
                coords.len()
            ));
        }
        let mut cursor = 0usize;
        for slot in 0..self.indices[row].len() {
            let atom = self.indices[row][slot] as usize;
            let meta = &self.atom_coord_meta[atom];
            let end = cursor + meta.latent_dim;
            let mut current = Array1::from_vec(coords[cursor..end].to_vec());
            let step = Array1::from_vec(delta[cursor..end].to_vec());
            if meta.retraction.is_all_euclidean() {
                current = meta.manifold.retract(current.view(), step.view());
            } else {
                meta.retraction
                    .retract(&mut current.view_mut(), step.view());
            }
            coords[cursor..end]
                .copy_from_slice(current.as_slice().expect("retraction is contiguous"));
            cursor = end;
        }
        Ok(())
    }

    /// [`Self::set_row_coords`] against a caller-held coordinate block — the
    /// identical validation and per-atom manifold projection.
    pub fn project_row_coords(
        &self,
        row: usize,
        values: &[f64],
        coords: &mut [f64],
    ) -> Result<(), String> {
        if row >= self.n_obs {
            return Err(format!(
                "SaeAssignmentState::project_row_coords: row {row} out of range N={}",
                self.n_obs
            ));
        }
        if values.len() != coords.len() {
            return Err(format!(
                "SaeAssignmentState::project_row_coords: row {row} value width {} != compact coordinate width {}",
                values.len(),
                coords.len()
            ));
        }
        if values.iter().any(|value| !value.is_finite()) {
            return Err(format!(
                "SaeAssignmentState::project_row_coords: row {row} contains a non-finite coordinate"
            ));
        }
        let mut cursor = 0usize;
        for &atom in &self.indices[row] {
            let meta = &self.atom_coord_meta[atom as usize];
            let end = cursor + meta.latent_dim;
            let candidate = Array1::from_vec(values[cursor..end].to_vec());
            let projected = meta.manifold.project_point(candidate.view());
            coords[cursor..end]
                .copy_from_slice(projected.as_slice().expect("projection is contiguous"));
            cursor = end;
        }
        Ok(())
    }

    /// #2502 occupancy-earned topology: retype one atom's coordinate block to
    /// dimension-generic Euclidean. The caller owns the coordinate remap; this
    /// method changes ONLY the topology metadata every sparse consumer reads
    /// through [`Self::atom_axis_periods`] and the retraction registry.
    pub fn convert_atom_to_euclidean(&mut self, atom: usize) -> Result<(), String> {
        if atom >= self.k_atoms {
            return Err(format!(
                "SaeAssignmentState::convert_atom_to_euclidean: atom {atom} out of range K={}",
                self.k_atoms
            ));
        }
        let meta = &mut self.atom_coord_meta[atom];
        meta.manifold = LatentManifold::Euclidean;
        meta.retraction = LatentRetractionRegistry::all_euclidean();
        Ok(())
    }

    /// Overwrite ONE support slot's coordinate block, projecting through the
    /// atom's manifold exactly as [`Self::set_row_coords`] does for a row.
    pub fn set_slot_coords(
        &mut self,
        row: usize,
        slot: usize,
        values: &[f64],
    ) -> Result<(), String> {
        if row >= self.n_obs || slot >= self.indices[row].len() {
            return Err(format!(
                "SaeAssignmentState::set_slot_coords: row {row} slot {slot} out of range"
            ));
        }
        let atom = self.indices[row][slot] as usize;
        let width = self.atom_coord_meta[atom].latent_dim;
        if values.len() != width {
            return Err(format!(
                "SaeAssignmentState::set_slot_coords: value width {} != atom width {width}",
                values.len()
            ));
        }
        if values.iter().any(|value| !value.is_finite()) {
            return Err(format!(
                "SaeAssignmentState::set_slot_coords: row {row} slot {slot} non-finite coordinate"
            ));
        }
        let start: usize = self.indices[row][..slot]
            .iter()
            .map(|&prior| self.atom_coord_meta[prior as usize].latent_dim)
            .sum();
        let candidate = Array1::from_vec(values.to_vec());
        let projected = self.atom_coord_meta[atom]
            .manifold
            .project_point(candidate.view());
        self.coords[row][start..start + width]
            .copy_from_slice(projected.as_slice().expect("projection is contiguous"));
        Ok(())
    }

    pub fn set_row_coords(&mut self, row: usize, values: &[f64]) -> Result<(), String> {
        if row >= self.n_obs {
            return Err(format!(
                "SaeAssignmentState::set_row_coords: row {row} out of range N={}",
                self.n_obs
            ));
        }
        if values.len() != self.coords[row].len() {
            return Err(format!(
                "SaeAssignmentState::set_row_coords: row {row} value width {} != compact coordinate width {}",
                values.len(),
                self.coords[row].len()
            ));
        }
        if values.iter().any(|value| !value.is_finite()) {
            return Err(format!(
                "SaeAssignmentState::set_row_coords: row {row} contains a non-finite coordinate"
            ));
        }
        let mut cursor = 0usize;
        for &atom in &self.indices[row] {
            let meta = &self.atom_coord_meta[atom as usize];
            let end = cursor + meta.latent_dim;
            let candidate = Array1::from_vec(values[cursor..end].to_vec());
            let projected = meta.manifold.project_point(candidate.view());
            self.coords[row][cursor..end]
                .copy_from_slice(projected.as_slice().expect("projection is contiguous"));
            cursor = end;
        }
        Ok(())
    }

    // -- Layout-contract cell accounting (see module docs) -------------------

}

impl SaeAssignment {

}