car-topology 0.55.0

//! The execution-grounded scorer — reads `vec(A)`, regresses measured tokens.
//!
//! The proxy has two jobs, and the incumbent design in the literature fails
//! both of them.
//!
//! ## 1. It must read the adjacency
//!
//! Every published learned designer scores candidates with a message-passing
//! network over a graph whose nodes carry agent-*profile* embeddings and whose
//! edges are the candidate adjacency `A`. On a **homogeneous team** — where
//! `x_1 = … = x_N = x`, the default configuration of the published benchmarks —
//! every aggregator (mean, sum, attention) folds identical features, so every
//! node state, and therefore the pooled score, is independent of `A`. Every
//! candidate receives the same score. The ranking does not exist.
//!
//! This is not a subtle failure. It means hundreds of guided diffusion steps
//! reproduce a constant, and the only thing left separating candidates is the
//! structural-cost head. [`crate::diagnostics::team_homogeneity`] detects the
//! regime before you build on it.
//!
//! ## 2. It must regress measured tokens, not `|E|`
//!
//! The structural cost head is almost always the edge count. On the paper's
//! records `|E|` correlates with measured tokens at `r ≈ -0.4`: sparse
//! communication yields longer completions, so minimizing `|E|` *maximizes* the
//! bill it was introduced to reduce. The chain topology is the clean control —
//! sparsest connected, most token-expensive on every math-format benchmark.
//! [`crate::diagnostics::edge_count_token_correlation`] measures the sign on
//! your own records.
//!
//! So this proxy reads the flattened adjacency directly and is regressed on
//! `(u, τ̃)` — measured utility and per-task normalized token cost.
//!
//! ## The substitution, stated plainly
//!
//! The paper's proxy is `MLP([vec(A); c]) → (û, ĉ)` (Eq. 8) trained under
//! squared error on measured targets (Eq. 9), plus a VQGraph-style
//! structure-token auxiliary head (Eq. 10) as a regularizer. This
//! implementation keeps Eqs. (8) and (9) — same features, same targets — but
//! fits them in closed form rather than by SGD on an MLP.
//!
//! Which closed form is [`ProxyConditioning`], and the choice matters more than
//! it looks:
//!
//! - [`ProxyConditioning::Global`] is one ridge regression over
//!   `[1; vec(A); c]`. It is the literal linear reading of Eq. (8) — and it
//!   cannot do the job on its own, because a linear model in `A` and `c`
//!   separately has no term that couples them. "The chain is cheap for coding
//!   queries and dear for math queries" is exactly such a coupling, and it is
//!   the premise of per-query topology design. A global fit will learn the
//!   average and rank the same way for every query, quietly overriding a
//!   correct choice by the prior.
//! - [`ProxyConditioning::Local`] (the default) is **locally-weighted ridge**:
//!   at selection time the records are weighted by cosine similarity between
//!   their query and this one, and a small ridge over `[1; vec(A)]` alone is
//!   solved against those weights. Query conditioning moves from the feature
//!   vector into the weighting, which is what restores the coupling an MLP gets
//!   from its hidden layer.
//!
//! The local solve happens **once per query, not once per candidate** — it
//! depends only on `c` — so scoring the whole candidate set is still the
//! paper's single batched pass, plus one `F × F` Cholesky where
//! `F = 1 + N(N-1)`. For a four-agent team that is a 13×13 factorization.
//!
//! Two costs of the substitution, stated rather than buried:
//!
//! - A local proxy **retains its training records**; it is not a fixed-size
//!   model. Memory is `O(|D| · (N² + d))`, which for the paper's 300 records
//!   is trivial and for a million would not be.
//! - Eq. (10) is dropped, not approximated. It regularizes a *shared hidden
//!   representation*, and neither ridge variant has one — there is no layer for
//!   an auxiliary head to attach to. Its role (stabilizing the fit where `|D|`
//!   covers the topology space thinly) is served here by the ridge penalty and,
//!   in local mode, by the kernel bandwidth, which are weaker but honest
//!   substitutes.

use serde::{Deserialize, Serialize};

use crate::error::TopologyError;
use crate::record::RecordSet;
use crate::topology::Topology;

/// How the proxy conditions its prediction on the query.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[serde(tag = "mode", rename_all = "snake_case")]
pub enum ProxyConditioning {
    /// One global ridge over `[1; vec(A); c]`.
    ///
    /// Cheapest, smallest, and structurally unable to rank two topologies
    /// differently for two different queries — the coupling term does not
    /// exist in a linear model. Correct only where the best topology is a
    /// property of the workload rather than of the query, which is the case
    /// the paper's Q2 says holds *within* a benchmark and not across settings.
    Global,
    /// Locally-weighted ridge over `[1; vec(A)]`, records weighted by cosine
    /// similarity to the query at `temperature`.
    ///
    /// Higher temperature makes the fit more local (fewer records dominate the
    /// weighting); lower blends the whole set toward the global fit. 8 matches
    /// [`crate::PredictorConfig::temperature`] so the prior and the reranker
    /// agree on what "a nearby query" means.
    Local { temperature: f32 },
}

impl Default for ProxyConditioning {
    fn default() -> Self {
        ProxyConditioning::Local { temperature: 8.0 }
    }
}

/// How the proxy is fitted.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ProxyConfig {
    /// Ridge penalty `α`. With few records over many features the normal
    /// equations are singular without it; it is not a tuning knob so much as
    /// what makes the solve exist.
    pub ridge: f32,
    /// Global or locally-weighted. See [`ProxyConditioning`].
    pub conditioning: ProxyConditioning,
}

impl Default for ProxyConfig {
    fn default() -> Self {
        Self {
            ridge: 1.0,
            conditioning: ProxyConditioning::default(),
        }
    }
}

/// What the proxy predicts for one candidate.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct ProxyScore {
    /// `û` — predicted utility.
    pub utility: f32,
    /// `ĉ` — predicted per-task normalized token cost.
    pub cost: f32,
}

impl ProxyScore {
    /// `û − λ·ĉ` (Eq. 11) — the objective candidates are ranked on.
    pub fn objective(&self, cost_weight: f32) -> f32 {
        self.utility - cost_weight * self.cost
    }
}

/// The fitted state, which differs by conditioning mode.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
enum Fitted {
    /// Solved once at fit time; scoring is two dot products.
    Global {
        utility_weights: Vec<f32>,
        cost_weights: Vec<f32>,
    },
    /// Solved per query at score time against kernel weights.
    Local {
        /// `[1; vec(A)]` per record.
        design: Vec<Vec<f32>>,
        /// L2-normalized query per record, and whether it had a direction.
        directions: Vec<Vec<f32>>,
        has_direction: Vec<bool>,
        utility_targets: Vec<f32>,
        cost_targets: Vec<f32>,
        ridge: f32,
        temperature: f32,
    },
}

/// A fitted `(A, c) → (û, ĉ)` regression over measured utility and measured,
/// per-task normalized token cost.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "ExecutionProxyWire")]
pub struct ExecutionProxy {
    fitted: Fitted,
    n: usize,
    query_dim: usize,
}

/// On-disk shape, validated on the way in — see [`ExecutionProxy::validate`].
#[derive(Deserialize)]
struct ExecutionProxyWire {
    fitted: Fitted,
    n: usize,
    query_dim: usize,
}

impl TryFrom<ExecutionProxyWire> for ExecutionProxy {
    type Error = TopologyError;

    fn try_from(wire: ExecutionProxyWire) -> Result<Self, Self::Error> {
        let proxy = ExecutionProxy {
            fitted: wire.fitted,
            n: wire.n,
            query_dim: wire.query_dim,
        };
        proxy.validate()?;
        Ok(proxy)
    }
}

/// The proxy's two heads, solved for one query. Scoring every candidate against
/// these is the paper's single batched pass.
#[derive(Debug, Clone, PartialEq)]
pub struct ConditionedProxy {
    utility_weights: Vec<f32>,
    cost_weights: Vec<f32>,
    /// Whether the weights expect the query appended to the topology features.
    includes_query: bool,
    n: usize,
}

impl ConditionedProxy {
    /// Score one candidate against the already-solved heads.
    pub fn score(&self, topology: &Topology, query: &[f32]) -> Result<ProxyScore, TopologyError> {
        if topology.n() != self.n {
            return Err(TopologyError::SizeMismatch {
                expected: self.n,
                found: topology.n(),
            });
        }
        let x = features(topology, query, self.includes_query);
        Ok(ProxyScore {
            utility: dot(&x, &self.utility_weights),
            cost: dot(&x, &self.cost_weights),
        })
    }
}

impl ExecutionProxy {
    /// Fit both heads against measured `(u, τ̃)`.
    pub fn fit(records: &RecordSet, config: &ProxyConfig) -> Result<Self, TopologyError> {
        if !config.ridge.is_finite() || config.ridge < 0.0 {
            return Err(TopologyError::BadConfig {
                field: "ridge",
                expected: "finite and non-negative",
                found: format!("{}", config.ridge),
            });
        }

        let n = records.team_size();
        let query_dim = records.query_dim();

        let include_query = matches!(config.conditioning, ProxyConditioning::Global);
        let mut design = Vec::with_capacity(records.len());
        let mut utility_targets = Vec::with_capacity(records.len());
        let mut cost_targets = Vec::with_capacity(records.len());
        for (index, record) in records.records().iter().enumerate() {
            design.push(features(&record.topology, &record.query, include_query));
            utility_targets.push(record.utility);
            cost_targets.push(records.normalized_cost(index));
        }

        let fitted = match config.conditioning {
            ProxyConditioning::Global => {
                let dim = feature_dim(n, query_dim, true);
                let weights = vec![1.0f64; design.len()];
                let gram = gram_matrix(&design, dim, &weights, config.ridge as f64);
                let utility_weights = solve(&gram, &rhs(&design, &utility_targets, &weights, dim))?;
                let cost_weights = solve(&gram, &rhs(&design, &cost_targets, &weights, dim))?;
                Fitted::Global {
                    utility_weights: utility_weights.into_iter().map(|w| w as f32).collect(),
                    cost_weights: cost_weights.into_iter().map(|w| w as f32).collect(),
                }
            }
            ProxyConditioning::Local { temperature } => {
                if !temperature.is_finite() {
                    return Err(TopologyError::BadConfig {
                        field: "temperature",
                        expected: "finite",
                        found: format!("{temperature}"),
                    });
                }
                let mut directions = Vec::with_capacity(records.len());
                let mut has_direction = Vec::with_capacity(records.len());
                for record in records.records() {
                    let (d, has) = normalize(&record.query);
                    directions.push(d);
                    has_direction.push(has);
                }
                Fitted::Local {
                    design,
                    directions,
                    has_direction,
                    utility_targets,
                    cost_targets,
                    ridge: config.ridge,
                    temperature,
                }
            }
        };

        Ok(Self {
            fitted,
            n,
            query_dim,
        })
    }

    /// Check the invariants [`ExecutionProxy::fit`] establishes.
    ///
    /// This is the one worth having. `dot` zips its two slices, so a weight
    /// vector that is too short does not panic — it silently scores against a
    /// truncated feature vector and returns a plausible number. A cost head
    /// that is quietly wrong defeats the entire point of regressing measured
    /// tokens, and nothing downstream could tell.
    pub fn validate(&self) -> Result<(), TopologyError> {
        if self.n < 2 {
            return Err(TopologyError::TeamTooSmall { n: self.n });
        }
        let expect = |field: &'static str, found: usize, want: usize| {
            if found == want {
                Ok(())
            } else {
                Err(TopologyError::BadConfig {
                    field,
                    expected: "one weight per feature",
                    found: format!("{found} for {want} features"),
                })
            }
        };
        match &self.fitted {
            Fitted::Global {
                utility_weights,
                cost_weights,
            } => {
                let want = feature_dim(self.n, self.query_dim, true);
                expect("utility_weights", utility_weights.len(), want)?;
                expect("cost_weights", cost_weights.len(), want)?;
            }
            Fitted::Local {
                design,
                directions,
                has_direction,
                utility_targets,
                cost_targets,
                ridge,
                temperature,
            } => {
                if design.is_empty() {
                    return Err(TopologyError::NoRecords { kind: "proxy" });
                }
                let rows = design.len();
                for (field, found) in [
                    ("directions", directions.len()),
                    ("has_direction", has_direction.len()),
                    ("utility_targets", utility_targets.len()),
                    ("cost_targets", cost_targets.len()),
                ] {
                    if found != rows {
                        return Err(TopologyError::BadConfig {
                            field,
                            expected: "one entry per design row",
                            found: format!("{found} for {rows} rows"),
                        });
                    }
                }
                let want = feature_dim(self.n, self.query_dim, false);
                for row in design {
                    expect("design row", row.len(), want)?;
                }
                for direction in directions {
                    if direction.len() != self.query_dim {
                        return Err(TopologyError::QueryDimMismatch {
                            expected: self.query_dim,
                            found: direction.len(),
                        });
                    }
                }
                if !ridge.is_finite() || *ridge < 0.0 {
                    return Err(TopologyError::BadConfig {
                        field: "ridge",
                        expected: "finite and non-negative",
                        found: format!("{ridge}"),
                    });
                }
                if !temperature.is_finite() {
                    return Err(TopologyError::BadConfig {
                        field: "temperature",
                        expected: "finite",
                        found: format!("{temperature}"),
                    });
                }
            }
        }
        Ok(())
    }

    /// Team size this proxy scores topologies over.
    pub fn team_size(&self) -> usize {
        self.n
    }

    /// Query dimension this proxy expects.
    pub fn query_dim(&self) -> usize {
        self.query_dim
    }

    /// Solve the two heads for one query.
    ///
    /// This is the whole per-query cost of the proxy: nothing below depends on
    /// the candidate, so a caller scoring `M` candidates pays this once.
    pub fn condition(&self, query: &[f32]) -> Result<ConditionedProxy, TopologyError> {
        if query.len() != self.query_dim {
            return Err(TopologyError::QueryDimMismatch {
                expected: self.query_dim,
                found: query.len(),
            });
        }
        match &self.fitted {
            Fitted::Global {
                utility_weights,
                cost_weights,
            } => Ok(ConditionedProxy {
                utility_weights: utility_weights.clone(),
                cost_weights: cost_weights.clone(),
                includes_query: true,
                n: self.n,
            }),
            Fitted::Local {
                design,
                directions,
                has_direction,
                utility_targets,
                cost_targets,
                ridge,
                temperature,
            } => {
                let (direction, query_has_direction) = normalize(query);
                let logits: Vec<f32> = directions
                    .iter()
                    .zip(has_direction.iter())
                    .map(|(d, &has)| {
                        if has && query_has_direction {
                            temperature * dot(&direction, d)
                        } else {
                            0.0
                        }
                    })
                    .collect();
                // Scaled so the weights sum to the record count rather than to
                // 1: the ridge penalty is then on the same footing as in the
                // global fit, instead of being overwhelming by a factor of |D|.
                let softmaxed = softmax(&logits);
                let weights: Vec<f64> = softmaxed
                    .iter()
                    .map(|w| *w as f64 * design.len() as f64)
                    .collect();

                let dim = feature_dim(self.n, self.query_dim, false);
                let gram = gram_matrix(design, dim, &weights, *ridge as f64);
                let utility_weights = solve(&gram, &rhs(design, utility_targets, &weights, dim))?;
                let cost_weights = solve(&gram, &rhs(design, cost_targets, &weights, dim))?;
                Ok(ConditionedProxy {
                    utility_weights: utility_weights.into_iter().map(|w| w as f32).collect(),
                    cost_weights: cost_weights.into_iter().map(|w| w as f32).collect(),
                    includes_query: false,
                    n: self.n,
                })
            }
        }
    }

    /// Score one candidate.
    ///
    /// Convenience over [`ExecutionProxy::condition`]; scoring several
    /// candidates for the same query should go through
    /// [`ExecutionProxy::score_batch`] so the per-query solve happens once.
    pub fn score(&self, topology: &Topology, query: &[f32]) -> Result<ProxyScore, TopologyError> {
        self.condition(query)?.score(topology, query)
    }

    /// Score every candidate — the paper's single batched forward pass.
    ///
    /// One per-query solve, then a dot product per candidate. There is no
    /// per-candidate graph to construct and no message passing, so the cost is
    /// flat in the number of candidates.
    pub fn score_batch(
        &self,
        candidates: &[Topology],
        query: &[f32],
    ) -> Result<Vec<ProxyScore>, TopologyError> {
        let conditioned = self.condition(query)?;
        candidates
            .iter()
            .map(|t| conditioned.score(t, query))
            .collect()
    }
}

fn feature_dim(n: usize, query_dim: usize, use_query: bool) -> usize {
    1 + n * (n - 1) + if use_query { query_dim } else { 0 }
}

/// `[1; vec(A)]`, with `c` appended when the fit uses query features.
fn features(topology: &Topology, query: &[f32], use_query: bool) -> Vec<f32> {
    let mut x = Vec::with_capacity(1 + topology.flat().len() + query.len());
    x.push(1.0);
    x.extend(topology.flat().iter().map(|&e| if e { 1.0 } else { 0.0 }));
    if use_query {
        x.extend_from_slice(query);
    }
    x
}

fn dot(x: &[f32], w: &[f32]) -> f32 {
    x.iter().zip(w.iter()).map(|(a, b)| a * b).sum()
}

/// L2-normalize, reporting whether the vector had any magnitude to normalize.
fn normalize(v: &[f32]) -> (Vec<f32>, bool) {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > f32::EPSILON {
        (v.iter().map(|x| x / norm).collect(), true)
    } else {
        (vec![0.0; v.len()], false)
    }
}

/// Numerically stable softmax.
fn softmax(logits: &[f32]) -> Vec<f32> {
    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = logits.iter().map(|l| (l - max).exp()).collect();
    let total: f32 = exps.iter().sum();
    if total > 0.0 {
        exps.into_iter().map(|e| e / total).collect()
    } else {
        vec![1.0 / logits.len() as f32; logits.len()]
    }
}

/// `XᵀWX + αI`, in row-major `dim × dim`.
fn gram_matrix(design: &[Vec<f32>], dim: usize, weights: &[f64], ridge: f64) -> Vec<f64> {
    let mut gram = vec![0f64; dim * dim];
    for (row, &w) in design.iter().zip(weights.iter()) {
        if w == 0.0 {
            continue;
        }
        for i in 0..dim {
            let xi = row[i] as f64 * w;
            if xi == 0.0 {
                continue;
            }
            for j in 0..dim {
                gram[i * dim + j] += xi * row[j] as f64;
            }
        }
    }
    for i in 0..dim {
        // The bias term is left unpenalized: shrinking it toward zero biases
        // every prediction toward the origin rather than toward the mean, which
        // on a utility head in [0,1] is a systematic underestimate.
        if i > 0 {
            gram[i * dim + i] += ridge;
        }
    }
    gram
}

/// `XᵀWy`.
fn rhs(design: &[Vec<f32>], targets: &[f32], weights: &[f64], dim: usize) -> Vec<f64> {
    let mut b = vec![0f64; dim];
    for ((row, &y), &w) in design.iter().zip(targets.iter()).zip(weights.iter()) {
        if w == 0.0 {
            continue;
        }
        let wy = y as f64 * w;
        for i in 0..dim {
            b[i] += row[i] as f64 * wy;
        }
    }
    b
}

/// Solve `A w = b` for a symmetric positive-semidefinite `A` by Cholesky with
/// escalating jitter.
///
/// The unpenalized bias column can leave the Gram matrix singular when the
/// ridge alone does not cover it (a degenerate record set — one record, or one
/// repeated topology). Rather than fail, the solve retries with progressively
/// larger jitter; a caller with such a set gets a heavily shrunk but finite
/// model, which is the correct answer for data that constrains nothing.
fn solve(gram: &[f64], b: &[f64]) -> Result<Vec<f64>, TopologyError> {
    let dim = b.len();
    let mut jitter = 0f64;
    for attempt in 0..8 {
        let mut a = gram.to_vec();
        if jitter > 0.0 {
            for i in 0..dim {
                a[i * dim + i] += jitter;
            }
        }
        if let Some(w) = cholesky_solve(&a, b, dim) {
            return Ok(w);
        }
        jitter = if attempt == 0 { 1e-8 } else { jitter * 100.0 };
    }
    Err(TopologyError::BadConfig {
        field: "ridge",
        expected: "large enough to make the normal equations solvable",
        found: "singular even with jitter".into(),
    })
}

/// Cholesky factor-and-solve. Returns `None` when the matrix is not positive
/// definite at this jitter level.
fn cholesky_solve(a: &[f64], b: &[f64], dim: usize) -> Option<Vec<f64>> {
    let mut l = vec![0f64; dim * dim];
    for i in 0..dim {
        for j in 0..=i {
            let mut sum = a[i * dim + j];
            for k in 0..j {
                sum -= l[i * dim + k] * l[j * dim + k];
            }
            if i == j {
                if sum <= 0.0 || !sum.is_finite() {
                    return None;
                }
                l[i * dim + j] = sum.sqrt();
            } else {
                l[i * dim + j] = sum / l[j * dim + j];
            }
        }
    }

    // Forward substitution: L y = b.
    let mut y = vec![0f64; dim];
    for i in 0..dim {
        let mut sum = b[i];
        for k in 0..i {
            sum -= l[i * dim + k] * y[k];
        }
        y[i] = sum / l[i * dim + i];
    }
    // Back substitution: Lᵀ w = y.
    let mut w = vec![0f64; dim];
    for i in (0..dim).rev() {
        let mut sum = y[i];
        for k in (i + 1)..dim {
            sum -= l[k * dim + i] * w[k];
        }
        w[i] = sum / l[i * dim + i];
    }
    if w.iter().all(|v| v.is_finite()) {
        Some(w)
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::record::{ExecutionRecord, DEFAULT_COST_WEIGHT};

    /// Records in which the chain topology is the sparsest AND the most
    /// expensive — the paper's positive control for the inverted surrogate.
    fn inverted_cost_records() -> RecordSet {
        let n = 4;
        let chain = Topology::chain(n).unwrap();
        let complete = Topology::complete(n).unwrap();
        let star = Topology::star(n, 0).unwrap();
        let mut out = Vec::new();
        for i in 0..8 {
            let q = vec![i as f32 / 8.0, 1.0 - i as f32 / 8.0];
            let task = format!("t{i}");
            out.push(ExecutionRecord::new(
                &task,
                q.clone(),
                chain.clone(),
                1.0,
                3000,
            ));
            out.push(ExecutionRecord::new(
                &task,
                q.clone(),
                star.clone(),
                1.0,
                1500,
            ));
            out.push(ExecutionRecord::new(&task, q, complete.clone(), 1.0, 900));
        }
        RecordSet::new(out).unwrap()
    }

    /// Records where the cheap topology DEPENDS on the query family: complete
    /// is cheap for "math" queries, chain for "code" queries. A global linear
    /// fit has no term that can express this; a local one does.
    fn query_dependent_records() -> RecordSet {
        let n = 4;
        let chain = Topology::chain(n).unwrap();
        let complete = Topology::complete(n).unwrap();
        let mut out = Vec::new();
        for i in 0..6 {
            let drift = i as f32 * 0.01;
            let math = format!("math{i}");
            let math_q = vec![1.0, 0.0, drift];
            out.push(ExecutionRecord::new(
                &math,
                math_q.clone(),
                complete.clone(),
                1.0,
                600,
            ));
            out.push(ExecutionRecord::new(
                &math,
                math_q,
                chain.clone(),
                1.0,
                2400,
            ));

            let code = format!("code{i}");
            let code_q = vec![0.0, 1.0, drift];
            out.push(ExecutionRecord::new(
                &code,
                code_q.clone(),
                chain.clone(),
                1.0,
                600,
            ));
            out.push(ExecutionRecord::new(
                &code,
                code_q,
                complete.clone(),
                1.0,
                2400,
            ));
        }
        RecordSet::new(out).unwrap()
    }

    fn global_config(ridge: f32) -> ProxyConfig {
        ProxyConfig {
            ridge,
            conditioning: ProxyConditioning::Global,
        }
    }

    #[test]
    fn the_proxy_learns_the_measured_cost_ordering_not_the_edge_count_one() {
        let records = inverted_cost_records();
        let proxy = ExecutionProxy::fit(&records, &ProxyConfig::default()).unwrap();
        let q = vec![0.5, 0.5];
        let n = 4;
        let chain = proxy.score(&Topology::chain(n).unwrap(), &q).unwrap();
        let complete = proxy.score(&Topology::complete(n).unwrap(), &q).unwrap();

        // Chain has the FEWEST edges and the HIGHEST measured cost. An
        // edge-count head would rank it cheapest; the proxy ranks it dearest.
        assert!(
            Topology::chain(n).unwrap().edge_count() < Topology::complete(n).unwrap().edge_count()
        );
        assert!(
            chain.cost > complete.cost,
            "chain {} should score dearer than complete {}",
            chain.cost,
            complete.cost
        );
        assert!(complete.objective(DEFAULT_COST_WEIGHT) > chain.objective(DEFAULT_COST_WEIGHT));
    }

    #[test]
    fn the_proxy_separates_candidates_on_a_homogeneous_team() {
        // The whole point: unlike a profile-node message-passing scorer, this
        // gives different scores to different adjacencies with no profile
        // information at all.
        let records = inverted_cost_records();
        for config in [ProxyConfig::default(), global_config(1.0)] {
            let proxy = ExecutionProxy::fit(&records, &config).unwrap();
            let scores = proxy
                .score_batch(
                    &[
                        Topology::chain(4).unwrap(),
                        Topology::star(4, 0).unwrap(),
                        Topology::complete(4).unwrap(),
                    ],
                    &[0.5, 0.5],
                )
                .unwrap();
            assert!(scores[0].cost != scores[1].cost, "{config:?}");
            assert!(scores[1].cost != scores[2].cost, "{config:?}");
        }
    }

    #[test]
    fn local_conditioning_ranks_the_same_topology_differently_per_query() {
        let records = query_dependent_records();
        let proxy = ExecutionProxy::fit(&records, &ProxyConfig::default()).unwrap();
        let n = 4;
        let chain = Topology::chain(n).unwrap();
        let complete = Topology::complete(n).unwrap();

        let math = proxy
            .score_batch(&[chain.clone(), complete.clone()], &[1.0, 0.0, 0.0])
            .unwrap();
        assert!(
            math[1].cost < math[0].cost,
            "math query should price complete below chain: {math:?}"
        );

        let code = proxy
            .score_batch(&[chain, complete], &[0.0, 1.0, 0.0])
            .unwrap();
        assert!(
            code[0].cost < code[1].cost,
            "code query should price chain below complete: {code:?}"
        );
    }

    #[test]
    fn global_conditioning_cannot_and_reports_the_average() {
        // Documents the limitation rather than hiding it: on the same records,
        // a global linear fit gives the SAME ranking for both query families.
        let records = query_dependent_records();
        let proxy = ExecutionProxy::fit(&records, &global_config(0.01)).unwrap();
        let n = 4;
        let candidates = [Topology::chain(n).unwrap(), Topology::complete(n).unwrap()];
        let math = proxy.score_batch(&candidates, &[1.0, 0.0, 0.0]).unwrap();
        let code = proxy.score_batch(&candidates, &[0.0, 1.0, 0.0]).unwrap();
        assert_eq!(
            math[0].cost < math[1].cost,
            code[0].cost < code[1].cost,
            "a global linear fit has no query/topology coupling term"
        );
    }

    #[test]
    fn predictions_track_the_targets_they_were_fitted_on() {
        let records = inverted_cost_records();
        for config in [
            ProxyConfig {
                ridge: 0.01,
                conditioning: ProxyConditioning::Local { temperature: 8.0 },
            },
            global_config(0.01),
        ] {
            let proxy = ExecutionProxy::fit(&records, &config).unwrap();
            for (index, record) in records.records().iter().enumerate() {
                let score = proxy.score(&record.topology, &record.query).unwrap();
                let target = records.normalized_cost(index);
                assert!(
                    (score.cost - target).abs() < 0.15,
                    "{config:?} predicted {} for target {target}",
                    score.cost
                );
            }
        }
    }

    #[test]
    fn utility_head_learns_a_failing_topology() {
        let n = 4;
        let mut out = Vec::new();
        for i in 0..8 {
            let q = vec![0.5, 0.5];
            let task = format!("t{i}");
            out.push(ExecutionRecord::new(
                &task,
                q.clone(),
                Topology::chain(n).unwrap(),
                0.0,
                1000,
            ));
            out.push(ExecutionRecord::new(
                &task,
                q,
                Topology::complete(n).unwrap(),
                1.0,
                1000,
            ));
        }
        let records = RecordSet::new(out).unwrap();
        let proxy = ExecutionProxy::fit(
            &records,
            &ProxyConfig {
                ridge: 0.01,
                conditioning: ProxyConditioning::Local { temperature: 8.0 },
            },
        )
        .unwrap();
        let q = vec![0.5, 0.5];
        let chain = proxy.score(&Topology::chain(n).unwrap(), &q).unwrap();
        let complete = proxy.score(&Topology::complete(n).unwrap(), &q).unwrap();
        assert!(complete.utility > chain.utility + 0.5);
    }

    #[test]
    fn conditioning_once_matches_scoring_each_candidate() {
        let records = query_dependent_records();
        let proxy = ExecutionProxy::fit(&records, &ProxyConfig::default()).unwrap();
        let q = [0.3, 0.7, 0.02];
        let candidates = [
            Topology::chain(4).unwrap(),
            Topology::star(4, 0).unwrap(),
            Topology::complete(4).unwrap(),
        ];
        let batched = proxy.score_batch(&candidates, &q).unwrap();
        let conditioned = proxy.condition(&q).unwrap();
        for (candidate, expected) in candidates.iter().zip(batched.iter()) {
            assert_eq!(&conditioned.score(candidate, &q).unwrap(), expected);
            assert_eq!(&proxy.score(candidate, &q).unwrap(), expected);
        }
    }

    #[test]
    fn a_degenerate_single_record_set_still_fits() {
        // One record: the normal equations are rank-1 over ~14 features. The
        // jitter escalation must produce a finite model, not an error.
        let records = RecordSet::new(vec![ExecutionRecord::new(
            "t",
            vec![1.0, 0.0],
            Topology::chain(4).unwrap(),
            1.0,
            100,
        )])
        .unwrap();
        for config in [ProxyConfig::default(), global_config(1.0)] {
            let proxy = ExecutionProxy::fit(&records, &config).unwrap();
            let score = proxy
                .score(&Topology::complete(4).unwrap(), &[1.0, 0.0])
                .unwrap();
            assert!(
                score.utility.is_finite() && score.cost.is_finite(),
                "{config:?}"
            );
        }
    }

    #[test]
    fn a_zero_query_scores_finitely_under_local_conditioning() {
        let records = query_dependent_records();
        let proxy = ExecutionProxy::fit(&records, &ProxyConfig::default()).unwrap();
        let score = proxy
            .score(&Topology::chain(4).unwrap(), &[0.0, 0.0, 0.0])
            .unwrap();
        assert!(score.utility.is_finite() && score.cost.is_finite());
    }

    #[test]
    fn the_global_fit_carries_query_features_and_the_local_one_does_not() {
        let records = inverted_cost_records();
        let global = ExecutionProxy::fit(&records, &global_config(1.0)).unwrap();
        match &global.fitted {
            Fitted::Global {
                utility_weights, ..
            } => {
                assert_eq!(utility_weights.len(), 1 + 4 * 3 + 2);
            }
            other => panic!("expected a global fit, got {other:?}"),
        }
        let local = ExecutionProxy::fit(&records, &ProxyConfig::default()).unwrap();
        let conditioned = local.condition(&[0.5, 0.5]).unwrap();
        assert_eq!(conditioned.utility_weights.len(), 1 + 4 * 3);
        assert!(!conditioned.includes_query);
    }

    #[test]
    fn shape_mismatches_are_rejected() {
        let records = inverted_cost_records();
        let proxy = ExecutionProxy::fit(&records, &ProxyConfig::default()).unwrap();
        assert!(matches!(
            proxy.score(&Topology::chain(5).unwrap(), &[0.5, 0.5]),
            Err(TopologyError::SizeMismatch { .. })
        ));
        assert!(matches!(
            proxy.score(&Topology::chain(4).unwrap(), &[0.5]),
            Err(TopologyError::QueryDimMismatch { .. })
        ));
    }

    #[test]
    fn a_negative_ridge_is_rejected() {
        let records = inverted_cost_records();
        assert!(matches!(
            ExecutionProxy::fit(&records, &global_config(-1.0)),
            Err(TopologyError::BadConfig { field: "ridge", .. })
        ));
    }

    #[test]
    fn a_non_finite_temperature_is_rejected() {
        let records = inverted_cost_records();
        assert!(matches!(
            ExecutionProxy::fit(
                &records,
                &ProxyConfig {
                    ridge: 1.0,
                    conditioning: ProxyConditioning::Local {
                        temperature: f32::NAN
                    },
                }
            ),
            Err(TopologyError::BadConfig {
                field: "temperature",
                ..
            })
        ));
    }

    #[test]
    fn fitting_is_deterministic() {
        let records = inverted_cost_records();
        for config in [ProxyConfig::default(), global_config(1.0)] {
            let a = ExecutionProxy::fit(&records, &config).unwrap();
            let b = ExecutionProxy::fit(&records, &config).unwrap();
            assert_eq!(a, b);
            assert_eq!(
                a.score(&Topology::chain(4).unwrap(), &[0.5, 0.5]).unwrap(),
                b.score(&Topology::chain(4).unwrap(), &[0.5, 0.5]).unwrap()
            );
        }
    }
}