wifi-densepose-nn 0.3.2

//! ADR-146 — RF encoder multi-task heads + uncertainty quantification.
//!
//! Extends ADR-024 (AETHER contrastive embedding) with seven task-specific head
//! branches over a shared RF embedding, per-head uncertainty, a
//! calibration-robustness loss tying invariance to the ADR-135 `calibration_id`,
//! and a `ContrastiveBatcher` sampling contract. The tensor ABI is **pure-Rust
//! `f32`** (no backend-specific tensor type at this boundary) so inference is
//! deterministic and witnessable (ADR-136 §2.5) and a head can be toggled by the
//! ADR-145 ablation matrix.

/// Shared RF embedding dimension (ADR-146 / ADR-024 AETHER).
pub const EMBEDDING_DIM: usize = 256;

/// A 256-d shared RF embedding (pure-Rust f32 ABI).
#[derive(Debug, Clone, PartialEq)]
pub struct RfEmbedding(pub Vec<f32>);

impl RfEmbedding {
    /// Wrap a vector, asserting it is [`EMBEDDING_DIM`] long.
    #[must_use]
    pub fn new(v: Vec<f32>) -> Self {
        debug_assert_eq!(v.len(), EMBEDDING_DIM, "embedding must be {EMBEDDING_DIM}-d");
        Self(v)
    }

    /// Squared L2 distance to another embedding.
    #[must_use]
    pub fn sq_dist(&self, other: &RfEmbedding) -> f32 {
        self.0.iter().zip(&other.0).map(|(a, b)| (a - b).powi(2)).sum()
    }
}

/// The seven task heads over the shared encoder (ADR-146 §2.1).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TaskKind {
    /// 17-keypoint pose.
    Pose,
    /// Binary presence.
    Presence,
    /// Person count.
    Count,
    /// Activity class.
    Activity,
    /// Vital signs (HR/BR).
    Vitals,
    /// Gait signature.
    Gait,
    /// Identity embedding (AETHER re-ID).
    IdentityEmbedding,
}

impl TaskKind {
    /// All seven heads.
    pub const ALL: [TaskKind; 7] = [
        TaskKind::Pose,
        TaskKind::Presence,
        TaskKind::Count,
        TaskKind::Activity,
        TaskKind::Vitals,
        TaskKind::Gait,
        TaskKind::IdentityEmbedding,
    ];
}

/// One head's output: task values plus a scalar predictive uncertainty
/// (ADR-146 §2.2). `uncertainty` mirrors the spirit of the ADR-136
/// `QualityScored` trait — lower is more confident.
#[derive(Debug, Clone, PartialEq)]
pub struct HeadOutput {
    /// Which head produced this.
    pub task: TaskKind,
    /// Raw output activations.
    pub values: Vec<f32>,
    /// Predictive uncertainty in [0, ∞); softplus of a learned log-variance.
    pub uncertainty: f32,
}

impl HeadOutput {
    /// Confidence in [0, 1] derived from uncertainty (`1 / (1 + uncertainty)`),
    /// matching the ADR-136 `QualityScored::quality_score` contract shape.
    #[must_use]
    pub fn confidence(&self) -> f32 {
        1.0 / (1.0 + self.uncertainty)
    }
}

/// A linear task head: `out = W·emb + b`, plus a separate scalar log-variance
/// projection `lv = wᵥ·emb + bᵥ` whose softplus is the predictive uncertainty.
#[derive(Debug, Clone)]
pub struct LinearHead {
    task: TaskKind,
    /// Row-major `[out_dim × EMBEDDING_DIM]` weights.
    w: Vec<f32>,
    b: Vec<f32>,
    out_dim: usize,
    /// Uncertainty (log-variance) projection over the embedding.
    var_w: Vec<f32>,
    var_b: f32,
}

/// A shape mismatch when building a [`LinearHead`] from supplied weights.
///
/// Returned by [`LinearHead::try_new`] so a caller loading weights from an
/// **untrusted / deserialized** source can validate the tensor shapes without
/// the panic that [`LinearHead::new`] raises on a programmer-supplied mismatch
/// (ADR-155 M2 §3: a pure-Rust input guard ahead of the construction contract).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RfHeadError {
    /// `w.len()` was not `out_dim * EMBEDDING_DIM`.
    WeightShape {
        /// Expected length (`out_dim * EMBEDDING_DIM`).
        expected: usize,
        /// Actual `w.len()`.
        got: usize,
    },
    /// `b.len()` was not `out_dim`.
    BiasShape {
        /// Expected length (`out_dim`).
        expected: usize,
        /// Actual `b.len()`.
        got: usize,
    },
    /// `var_w.len()` was not `EMBEDDING_DIM`.
    VarWeightShape {
        /// Expected length (`EMBEDDING_DIM`).
        expected: usize,
        /// Actual `var_w.len()`.
        got: usize,
    },
}

impl std::fmt::Display for RfHeadError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::WeightShape { expected, got } => {
                write!(f, "weight shape mismatch: expected {expected}, got {got}")
            }
            Self::BiasShape { expected, got } => {
                write!(f, "bias shape mismatch: expected {expected}, got {got}")
            }
            Self::VarWeightShape { expected, got } => {
                write!(f, "var weight shape mismatch: expected {expected}, got {got}")
            }
        }
    }
}

impl std::error::Error for RfHeadError {}

impl LinearHead {
    /// Build a head with given weights. `w.len()` must be `out_dim * EMBEDDING_DIM`.
    ///
    /// # Panics
    ///
    /// Panics on a shape mismatch (`w`/`b`/`var_w`). This is a construction-time
    /// API contract on *programmer-supplied* vectors. For weights from an
    /// untrusted / deserialized source, prefer [`LinearHead::try_new`], which
    /// returns a typed [`RfHeadError`] instead of panicking.
    #[must_use]
    pub fn new(task: TaskKind, out_dim: usize, w: Vec<f32>, b: Vec<f32>, var_w: Vec<f32>, var_b: f32) -> Self {
        assert_eq!(w.len(), out_dim * EMBEDDING_DIM, "weight shape mismatch");
        assert_eq!(b.len(), out_dim, "bias shape mismatch");
        assert_eq!(var_w.len(), EMBEDDING_DIM, "var weight shape mismatch");
        Self { task, w, b, out_dim, var_w, var_b }
    }

    /// Fallible constructor: validate the weight shapes and return a typed
    /// [`RfHeadError`] on mismatch instead of panicking (ADR-155 M2 §3).
    ///
    /// Use this when `w` / `b` / `var_w` originate from a checkpoint or any
    /// untrusted source. On success the produced head is byte-for-byte identical
    /// to [`LinearHead::new`] with the same arguments.
    ///
    /// # Errors
    ///
    /// Returns [`RfHeadError`] when any of:
    /// - `w.len() != out_dim * EMBEDDING_DIM`
    /// - `b.len() != out_dim`
    /// - `var_w.len() != EMBEDDING_DIM`
    pub fn try_new(
        task: TaskKind,
        out_dim: usize,
        w: Vec<f32>,
        b: Vec<f32>,
        var_w: Vec<f32>,
        var_b: f32,
    ) -> Result<Self, RfHeadError> {
        let expected_w = out_dim * EMBEDDING_DIM;
        if w.len() != expected_w {
            return Err(RfHeadError::WeightShape { expected: expected_w, got: w.len() });
        }
        if b.len() != out_dim {
            return Err(RfHeadError::BiasShape { expected: out_dim, got: b.len() });
        }
        if var_w.len() != EMBEDDING_DIM {
            return Err(RfHeadError::VarWeightShape { expected: EMBEDDING_DIM, got: var_w.len() });
        }
        Ok(Self { task, w, b, out_dim, var_w, var_b })
    }

    /// A zero-initialised head (uncertainty = softplus(0) ≈ 0.693).
    #[must_use]
    pub fn zeros(task: TaskKind, out_dim: usize) -> Self {
        Self::new(
            task,
            out_dim,
            vec![0.0; out_dim * EMBEDDING_DIM],
            vec![0.0; out_dim],
            vec![0.0; EMBEDDING_DIM],
            0.0,
        )
    }

    /// Forward pass over a shared embedding.
    #[must_use]
    pub fn forward(&self, emb: &RfEmbedding) -> HeadOutput {
        let mut values = vec![0.0f32; self.out_dim];
        for o in 0..self.out_dim {
            let row = &self.w[o * EMBEDDING_DIM..(o + 1) * EMBEDDING_DIM];
            let dot: f32 = row.iter().zip(&emb.0).map(|(wi, xi)| wi * xi).sum();
            values[o] = dot + self.b[o];
        }
        let log_var: f32 = self.var_w.iter().zip(&emb.0).map(|(wi, xi)| wi * xi).sum::<f32>() + self.var_b;
        let uncertainty = softplus(log_var);
        HeadOutput { task: self.task, values, uncertainty }
    }
}

/// Input magnitude above which `softplus(x) ≈ x` to f32 precision, so the
/// `exp` is skipped to avoid overflow (ADR-155 M2 §8: de-magicked from a bare
/// `20.0`; value unchanged). At x = 20, `ln(1+e^20) − 20 ≈ 2e-9`, below f32 eps.
const SOFTPLUS_LINEAR_THRESHOLD: f32 = 20.0;

fn softplus(x: f32) -> f32 {
    // Numerically stable softplus.
    if x > SOFTPLUS_LINEAR_THRESHOLD {
        x
    } else {
        (1.0 + x.exp()).ln()
    }
}

/// Multi-task encoder: a shared embedding feeding a set of [`LinearHead`]s
/// (ADR-146 §2.1). Heads can be subset for ADR-145 ablation.
#[derive(Debug, Clone, Default)]
pub struct MultiTaskHeads {
    heads: Vec<LinearHead>,
}

impl MultiTaskHeads {
    /// Empty head set.
    #[must_use]
    pub fn new() -> Self {
        Self { heads: Vec::new() }
    }

    /// Add a head.
    pub fn push(&mut self, head: LinearHead) {
        self.heads.push(head);
    }

    /// Number of active heads.
    #[must_use]
    pub fn len(&self) -> usize {
        self.heads.len()
    }

    /// Whether no heads are configured.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.heads.is_empty()
    }

    /// Run every head on the shared embedding.
    #[must_use]
    pub fn forward(&self, emb: &RfEmbedding) -> Vec<HeadOutput> {
        self.heads.iter().map(|h| h.forward(emb)).collect()
    }

    /// Run only the heads in `enabled` (ADR-145 ablation toggle).
    #[must_use]
    pub fn forward_subset(&self, emb: &RfEmbedding, enabled: &[TaskKind]) -> Vec<HeadOutput> {
        self.heads
            .iter()
            .filter(|h| enabled.contains(&h.task))
            .map(|h| h.forward(emb))
            .collect()
    }
}

/// Calibration-robustness loss (ADR-146 §2.3): the encoder should produce the
/// same embedding for the same physical input under two different ADR-135
/// calibration baselines. Returns the mean squared embedding difference — a
/// penalty that is 0 under perfect calibration invariance.
#[must_use]
pub fn calibration_robustness_loss(under_cal_a: &RfEmbedding, under_cal_b: &RfEmbedding) -> f32 {
    under_cal_a.sq_dist(under_cal_b) / EMBEDDING_DIM as f32
}

/// Triplet contrastive loss (ADR-024 / ADR-146 §2.4): pull `anchor` toward
/// `positive` (same physical state), push from `negative` (different), with a
/// margin. `max(0, d(a,p) - d(a,n) + margin)`.
#[must_use]
pub fn triplet_loss(anchor: &RfEmbedding, positive: &RfEmbedding, negative: &RfEmbedding, margin: f32) -> f32 {
    (anchor.sq_dist(positive) - anchor.sq_dist(negative) + margin).max(0.0)
}

/// A contrastive training triplet over the shared embedding space.
#[derive(Debug, Clone)]
pub struct Triplet {
    /// Anchor sample index.
    pub anchor: usize,
    /// Positive (same state, different environment) index.
    pub positive: usize,
    /// Negative (different state) index.
    pub negative: usize,
}

/// Formalised contrastive pair/triplet sampler (ADR-146 §2.4): positives are the
/// *same physical state across different environments* (cross-room invariance,
/// ADR-027 MERIDIAN); negatives are *different states*.
#[derive(Debug, Clone)]
pub struct ContrastiveBatcher {
    /// `state_of[i]` = the physical-state label of sample `i`.
    state_of: Vec<u32>,
    /// `env_of[i]` = the environment/room label of sample `i`.
    env_of: Vec<u32>,
}

impl ContrastiveBatcher {
    /// Build from per-sample (state, environment) labels.
    #[must_use]
    pub fn new(state_of: Vec<u32>, env_of: Vec<u32>) -> Self {
        assert_eq!(state_of.len(), env_of.len(), "label vectors must align");
        Self { state_of, env_of }
    }

    /// Deterministically enumerate triplets: for each anchor, the first sample
    /// with the *same state but a different environment* is the positive, and
    /// the first sample with a *different state* is the negative. Anchors with
    /// no valid positive or negative are skipped. Determinism (lowest-index
    /// choice) keeps the batch witnessable (ADR-136 §2.5).
    #[must_use]
    pub fn triplets(&self) -> Vec<Triplet> {
        let n = self.state_of.len();
        let mut out = Vec::new();
        for a in 0..n {
            let positive = (0..n).find(|&p| {
                p != a && self.state_of[p] == self.state_of[a] && self.env_of[p] != self.env_of[a]
            });
            let negative = (0..n).find(|&q| self.state_of[q] != self.state_of[a]);
            if let (Some(positive), Some(negative)) = (positive, negative) {
                out.push(Triplet { anchor: a, positive, negative });
            }
        }
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn emb(fill: f32) -> RfEmbedding {
        RfEmbedding::new(vec![fill; EMBEDDING_DIM])
    }

    /// ADR-155 M2 §8: the de-magicked softplus linear-threshold must equal the
    /// prior inline `20.0` literal exactly (operating-value guard).
    #[test]
    fn softplus_threshold_unchanged_from_literal() {
        assert_eq!(SOFTPLUS_LINEAR_THRESHOLD, 20.0_f32);
    }

    /// ADR-155 M2 §3: `try_new` accepts correctly-shaped weights and produces a
    /// head byte-identical to `new`, but returns a typed error on a mismatched
    /// (e.g. corrupt-checkpoint) shape instead of panicking.
    #[test]
    fn try_new_accepts_valid_and_rejects_each_bad_shape() {
        let out_dim = 2;
        let w = vec![0.0; out_dim * EMBEDDING_DIM];
        let b = vec![0.0; out_dim];
        let var_w = vec![0.0; EMBEDDING_DIM];

        // Valid: try_new == new (forward identical on a probe embedding).
        let head = LinearHead::try_new(TaskKind::Presence, out_dim, w.clone(), b.clone(), var_w.clone(), 0.0)
            .expect("valid shapes must construct");
        let reference = LinearHead::new(TaskKind::Presence, out_dim, w.clone(), b.clone(), var_w.clone(), 0.0);
        assert_eq!(head.forward(&emb(0.5)).values, reference.forward(&emb(0.5)).values);

        // Bad weight length.
        assert_eq!(
            LinearHead::try_new(TaskKind::Presence, out_dim, vec![0.0; 3], b.clone(), var_w.clone(), 0.0)
                .unwrap_err(),
            RfHeadError::WeightShape { expected: out_dim * EMBEDDING_DIM, got: 3 }
        );
        // Bad bias length.
        assert_eq!(
            LinearHead::try_new(TaskKind::Presence, out_dim, w.clone(), vec![0.0; 1], var_w.clone(), 0.0)
                .unwrap_err(),
            RfHeadError::BiasShape { expected: out_dim, got: 1 }
        );
        // Bad var-weight length.
        assert_eq!(
            LinearHead::try_new(TaskKind::Presence, out_dim, w, b, vec![0.0; 5], 0.0).unwrap_err(),
            RfHeadError::VarWeightShape { expected: EMBEDDING_DIM, got: 5 }
        );
    }

    #[test]
    fn head_forward_produces_values_and_finite_uncertainty() {
        let head = LinearHead::zeros(TaskKind::Presence, 2);
        let out = head.forward(&emb(1.0));
        assert_eq!(out.values, vec![0.0, 0.0]); // zero weights
        assert!(out.uncertainty.is_finite() && out.uncertainty > 0.0);
        assert!((out.confidence() - 1.0 / (1.0 + out.uncertainty)).abs() < 1e-6);
    }

    #[test]
    fn uncertainty_responds_to_log_variance_weights() {
        // var_w all 1 → log_var = sum(emb) = 256 → softplus ≈ 256 (clamped path).
        let head = LinearHead::new(
            TaskKind::Vitals,
            1,
            vec![0.0; EMBEDDING_DIM],
            vec![0.0],
            vec![1.0; EMBEDDING_DIM],
            0.0,
        );
        let out = head.forward(&emb(1.0));
        assert!(out.uncertainty > 100.0, "high log-var → high uncertainty");
        assert!(out.confidence() < 0.02);
    }

    #[test]
    fn calibration_robustness_loss_zero_for_identical() {
        assert_eq!(calibration_robustness_loss(&emb(0.5), &emb(0.5)), 0.0);
        assert!(calibration_robustness_loss(&emb(0.0), &emb(1.0)) > 0.0);
    }

    #[test]
    fn triplet_loss_properties() {
        let a = emb(0.0);
        let p = emb(0.1); // close
        let n = emb(5.0); // far
        // d(a,p) << d(a,n) → loss should be 0 with a modest margin.
        assert_eq!(triplet_loss(&a, &p, &n, 0.5), 0.0);
        // Swap: positive far, negative close → positive loss.
        assert!(triplet_loss(&a, &n, &p, 0.5) > 0.0);
    }

    #[test]
    fn multitask_subset_ablation() {
        let mut heads = MultiTaskHeads::new();
        heads.push(LinearHead::zeros(TaskKind::Presence, 1));
        heads.push(LinearHead::zeros(TaskKind::Pose, 51));
        heads.push(LinearHead::zeros(TaskKind::Vitals, 2));
        assert_eq!(heads.forward(&emb(1.0)).len(), 3);
        // Ablate to just presence + vitals.
        let sub = heads.forward_subset(&emb(1.0), &[TaskKind::Presence, TaskKind::Vitals]);
        assert_eq!(sub.len(), 2);
        assert!(sub.iter().all(|o| o.task != TaskKind::Pose));
    }

    #[test]
    fn contrastive_batcher_samples_cross_env_positives() {
        // samples: 0=(stateA,room0) 1=(stateA,room1) 2=(stateB,room0)
        let b = ContrastiveBatcher::new(vec![0, 0, 1], vec![0, 1, 0]);
        let trips = b.triplets();
        // Anchor 0: positive=1 (same state, diff room), negative=2 (diff state).
        let t0 = trips.iter().find(|t| t.anchor == 0).unwrap();
        assert_eq!(t0.positive, 1);
        assert_eq!(t0.negative, 2);
        // Anchor 2 (stateB) has no same-state-diff-env positive → skipped.
        assert!(trips.iter().all(|t| t.anchor != 2));
        // Deterministic.
        assert_eq!(b.triplets().len(), trips.len());
    }

    #[test]
    fn seven_task_heads() {
        assert_eq!(TaskKind::ALL.len(), 7);
    }
}