antecedent-stats 0.4.0

//! CI request / result types and workspace.
//!
//! SPDX-License-Identifier: MIT OR Apache-2.0

use antecedent_core::ExecutionContext;
use antecedent_kernels::ParCorrWorkspace;

use crate::error::StatsError;

/// Reusable kNN index + permutation plan for CMI.
#[derive(Clone, Debug, Default)]
pub struct KnnDependenceWorkspace {
    /// Built neighbor index generation (bumps only on rebuild).
    pub index_generation: u64,
    /// Number of times a new [`MatchingIndex`] was constructed.
    pub index_builds: u32,
    /// Last built feature dim.
    pub last_dim: usize,
    /// Last n.
    pub last_n: usize,
    /// Fingerprint of the (x, y, z) inputs behind the cached index (pointer, length,
    /// and sampled-content hash), so different pairs in one batch rebuild correctly.
    pub last_fingerprint: u64,
    /// Cached joint features (row-major `n * dim`).
    pub features: Vec<f64>,
    /// Cached nearest-neighbor index over [`Self::features`].
    pub index: Option<crate::matching::MatchingIndex>,
    /// Permutation plan (row indexes) reused for null shuffles.
    pub perm: Vec<usize>,
    /// Distance scratch (per-query kth distances).
    pub distances: Vec<f64>,
}

/// Significance method for a CI statistic.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub enum SignificanceMethod {
    /// Analytic Fisher-z / Student-t for partial correlation.
    Analytic,
    /// Block-shuffle null distribution.
    BlockShuffle {
        /// Number of null replicates.
        replicates: u32,
        /// Block length for shuffling.
        block_size: usize,
    },
}

/// Confidence interval method for a CI statistic.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum ConfidenceMethod {
    /// No interval.
    None,
    /// Analytic Fisher-z interval at the given level in `(0, 1)`.
    Analytic {
        /// Confidence level (e.g. `0.95`).
        level: f64,
    },
}

impl Default for ConfidenceMethod {
    fn default() -> Self {
        Self::Analytic { level: 0.95 }
    }
}

/// How many null replicates a nonparametric CI test should draw.
///
/// [`SignificanceMethod::Analytic`] has no closed-form null for distance / MI proxies;
/// those tests use a documented default of 49. [`SignificanceMethod::BlockShuffle`]
/// honors the caller's `replicates`.
#[must_use]
pub fn nonparametric_permutation_count(significance: SignificanceMethod) -> usize {
    match significance {
        SignificanceMethod::Analytic => 49,
        SignificanceMethod::BlockShuffle { replicates, .. } => replicates.max(1) as usize,
    }
}

/// The block length the caller asked for, or `1` when no blocking was requested.
///
/// `1` means "permute exchangeably": [`SignificanceMethod::Analytic`] has no blocking concept,
/// and `BlockShuffle { block_size: 0 | 1, .. }` imposes none.
#[must_use]
pub fn requested_block_size(significance: SignificanceMethod) -> usize {
    match significance {
        SignificanceMethod::Analytic => 1,
        SignificanceMethod::BlockShuffle { block_size, .. } => block_size.max(1),
    }
}

/// Reject a block-preserving request from a test that can only permute exchangeably.
///
/// [`SignificanceMethod::BlockShuffle`]'s `block_size` exists so the null keeps the serial
/// dependence of the observed series; an exchangeable permutation under-disperses that null
/// and inflates Type I error for exactly the autocorrelated data the parameter is meant to
/// protect. Tests whose null generation cannot honour `block_size` must say so rather than
/// accept the argument and discard it, which silently returns an anticonservative p-value.
///
/// `block_size <= 1` imposes no blocking, so it is always accepted.
///
/// Call this per query, not per batch: for a test that builds its null by stratifying on Z,
/// whether `block_size` can be honoured depends on whether *that query's* conditioning set is
/// empty, which varies across a batch.
///
/// # Errors
///
/// [`StatsError::Unsupported`] when `block_size > 1`.
pub fn reject_unsupported_block_size(
    significance: SignificanceMethod,
    test_name: &'static str,
) -> Result<(), crate::error::StatsError> {
    if let SignificanceMethod::BlockShuffle { block_size, .. } = significance {
        if block_size > 1 {
            let _ = test_name;
            return Err(crate::error::StatsError::Unsupported {
                message: "block-preserving permutation (block_size > 1) is not supported for \
                          this CI test when the conditioning set is non-empty: its null is a \
                          within-Z-stratum exchange over strata whose members are scattered \
                          across time, which is structurally incompatible with a contiguous-block \
                          permutation and would be invalid for autocorrelated data if silently \
                          substituted. This is a structural limit, not an unimplemented feature. \
                          Options: query with an empty conditioning set, where blocking is \
                          supported because the null is a single time-ordered stratum; use \
                          ParCorr / WeightedParCorr / MultivariateParCorr / Gpdc, which \
                          residualize on Z and so support block-preserving nulls with \
                          conditioning; or set block_size = 1 to accept an exchangeable null.",
            });
        }
    }
    Ok(())
}

/// Confidence level for analytic intervals, if requested.
#[must_use]
pub fn analytic_confidence_level(confidence: ConfidenceMethod) -> Option<f64> {
    match confidence {
        ConfidenceMethod::None => None,
        ConfidenceMethod::Analytic { level } => Some(level),
    }
}

/// Preparation plan for a CI session.
#[derive(Clone, Debug)]
pub struct CiPreparationPlan {
    /// Significance method applied to subsequent queries.
    pub significance: SignificanceMethod,
    /// Confidence method applied when analytic intervals are available.
    pub confidence: ConfidenceMethod,
}

impl Default for CiPreparationPlan {
    fn default() -> Self {
        Self { significance: SignificanceMethod::Analytic, confidence: ConfidenceMethod::default() }
    }
}

/// Prepared CI state after [`ConditionalIndependenceTest::prepare`].
#[derive(Clone, Debug)]
pub struct PreparedCiTest {
    /// Row count observed at prepare time.
    pub n: usize,
    /// Column count observed at prepare time.
    pub ncols: usize,
    /// Plan used for preparation.
    pub plan: CiPreparationPlan,
}

impl PreparedCiTest {
    /// Ensure a batch request matches this prepare (row/column counts).
    ///
    /// # Errors
    ///
    /// Shape mismatch vs prepare-time `n` / `ncols`.
    pub fn ensure_compatible(&self, request: &CiBatchRequest<'_>) -> Result<(), StatsError> {
        let n = request.nrows()?;
        if n != self.n {
            return Err(StatsError::Shape { message: "CI batch row count differs from prepare()" });
        }
        if request.columns.len() != self.ncols {
            return Err(StatsError::Shape {
                message: "CI batch column count differs from prepare()",
            });
        }
        Ok(())
    }

    /// Copy of `request` with significance/confidence taken from this prepare plan.
    #[must_use]
    pub fn bind_request<'a>(&self, request: &CiBatchRequest<'a>) -> CiBatchRequest<'a> {
        CiBatchRequest {
            columns: request.columns,
            queries: request.queries,
            z_flat: request.z_flat,
            significance: self.plan.significance,
            confidence: self.plan.confidence,
        }
    }
}

/// One CI query over column indexes into a shared matrix.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct CiQuery {
    /// X column index.
    pub x: usize,
    /// Y column index.
    pub y: usize,
    /// Start into flat conditioning indexes.
    pub z_start: usize,
    /// Conditioning arity.
    pub z_len: usize,
}

/// Batch of CI queries (deterministic output order).
#[derive(Clone, Debug)]
pub struct CiBatchRequest<'a> {
    /// Column-major / list of equal-length float columns.
    pub columns: &'a [&'a [f64]],
    /// Queries.
    pub queries: &'a [CiQuery],
    /// Flat conditioning column indexes.
    pub z_flat: &'a [usize],
    /// Significance.
    pub significance: SignificanceMethod,
    /// Confidence intervals (when the test supports analytic intervals).
    pub confidence: ConfidenceMethod,
}

/// One CI result.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct CiResult {
    /// Test statistic (partial correlation for partial-correlation CI).
    pub statistic: f64,
    /// Two-sided p-value.
    pub p_value: f64,
    /// Residual degrees of freedom (analytic path).
    pub df: f64,
    /// Optional analytic confidence interval `(lower, upper)` for the statistic.
    pub ci: Option<(f64, f64)>,
}

/// Batch results aligned with request queries.
#[derive(Clone, Debug, Default)]
pub struct CiBatchResult {
    /// Per-query results.
    pub results: Vec<CiResult>,
}

/// Conditional independence test.
///
/// Numeric kernels live in `antecedent-stats`; discovery owns the algorithm surface and
/// re-exports this trait.
pub trait ConditionalIndependenceTest {
    /// Prepare once for a data view / plan (sample planning, caches).
    ///
    /// # Errors
    ///
    /// Shape failures.
    fn prepare(
        &self,
        columns: &[&[f64]],
        plan: &CiPreparationPlan,
        _ctx: &ExecutionContext,
    ) -> Result<PreparedCiTest, StatsError> {
        if columns.is_empty() {
            return Err(StatsError::Shape { message: "no columns" });
        }
        let n = columns[0].len();
        for col in columns {
            if col.len() != n {
                return Err(StatsError::Shape { message: "column length mismatch" });
            }
        }
        Ok(PreparedCiTest { n, ncols: columns.len(), plan: plan.clone() })
    }

    /// Single query (convenience over [`Self::test_batch`]).
    ///
    /// # Errors
    ///
    /// Shape / numerical failures.
    fn test(
        &self,
        prepared: &PreparedCiTest,
        columns: &[&[f64]],
        query: CiQuery,
        z_flat: &[usize],
        workspace: &mut CiWorkspace,
        ctx: &ExecutionContext,
    ) -> Result<CiResult, StatsError> {
        let req = CiBatchRequest {
            columns,
            queries: std::slice::from_ref(&query),
            z_flat,
            significance: prepared.plan.significance,
            confidence: prepared.plan.confidence,
        };
        let out = self.test_batch(prepared, &req, workspace, ctx)?;
        out.results
            .into_iter()
            .next()
            .ok_or(StatsError::Shape { message: "CI test returned no results" })
    }

    /// Evaluate a batch against a prior [`Self::prepare`].
    ///
    /// Implementations must call [`PreparedCiTest::ensure_compatible`] and should
    /// honor `prepared.plan` for significance / confidence.
    ///
    /// # Errors
    ///
    /// Shape / numerical failures.
    fn test_batch(
        &self,
        prepared: &PreparedCiTest,
        request: &CiBatchRequest<'_>,
        workspace: &mut CiWorkspace,
        ctx: &ExecutionContext,
    ) -> Result<CiBatchResult, StatsError>;

    /// Ad-hoc batch: prepare from the request plan, then [`Self::test_batch`].
    ///
    /// Prefer an explicit prepare-once session for discovery / repeated queries.
    ///
    /// # Errors
    ///
    /// Shape / numerical failures.
    fn test_batch_adhoc(
        &self,
        request: &CiBatchRequest<'_>,
        workspace: &mut CiWorkspace,
        ctx: &ExecutionContext,
    ) -> Result<CiBatchResult, StatsError> {
        let plan = CiPreparationPlan {
            significance: request.significance,
            confidence: request.confidence,
        };
        let prepared = self.prepare(request.columns, &plan, ctx)?;
        self.test_batch(&prepared, request, workspace, ctx)
    }
}

/// Conditional independence test.
///
/// Prefer this name; [`ConditionalIndependenceTest`] is the same trait.
pub use ConditionalIndependenceTest as ConditionalIndependence;

impl CiBatchRequest<'_> {
    /// Validate non-empty equal-length columns; returns `n`.
    ///
    /// # Errors
    ///
    /// Empty column list or length mismatch.
    pub fn nrows(&self) -> Result<usize, StatsError> {
        if self.columns.is_empty() {
            return Err(StatsError::Shape { message: "no columns" });
        }
        let n = self.columns[0].len();
        for col in self.columns {
            if col.len() != n {
                return Err(StatsError::Shape { message: "column length mismatch" });
            }
        }
        Ok(n)
    }
}

/// Shared scratch for CI batches.
#[derive(Clone, Debug, Default)]
pub struct CiWorkspace {
    /// Partial-correlation residualization workspace.
    pub parcorr: ParCorrWorkspace,
    /// Temporary statistic buffer.
    pub stats: Vec<Option<f64>>,
    /// Block-shuffle column scratch.
    pub shuffled: Vec<f64>,
    /// Block starts for shuffle / reusable permutation plan.
    pub block_perm: Vec<usize>,
    /// Contingency X level-code scratch (G² / discrete CI).
    pub contingency_x_codes: Vec<u32>,
    /// Contingency Y level-code scratch (G² / discrete CI).
    pub contingency_y_codes: Vec<u32>,
    /// kNN CMI index / permutation reuse state.
    pub knn: KnnDependenceWorkspace,
}

impl CiWorkspace {
    /// Prepare for `n_queries` results.
    pub fn prepare_queries(&mut self, n_queries: usize) {
        if self.stats.len() < n_queries {
            self.stats.resize(n_queries, None);
        }
    }
}