antecedent-estimate 0.5.2

//! Propensity-score nearest-neighbor matching.
//!
//! Analytic standard errors follow Abadie–Imbens (2006) with donor-usage counts
//! `Kᵢ` (matching with replacement). A linear within-arm regression bias
//! adjustment (Abadie–Imbens) is applied on the match feature(s).
//!
//! **Bootstrap caution:** the nonparametric bootstrap is invalid for nearest-neighbor
//! matching with a fixed number of matches (Abadie–Imbens 2008). Prefer the analytic
//! SE; bootstrap replicates (when enabled) are retained only for diagnostics and must
//! not be treated as valid confidence-interval input for NN matching.
//!
//! SPDX-License-Identifier: MIT OR Apache-2.0

#![allow(clippy::many_single_char_names, clippy::too_many_lines, clippy::type_complexity)]

use antecedent_core::{
    AssumptionSet, AverageEffectQuery, ExecutionContext, PopulationRegistry, TargetPopulation,
};
use antecedent_data::TabularData;
use antecedent_expr::IdentifiedEstimand;
use antecedent_stats::{FaerBackend, GlmOptions, MatchingDistance, fit_propensity};

use super::prepare::{
    PreparedPropensityProblem, PropensityEstimationWorkspace, PropensityModel, clamp_scores,
    clip_of, default_propensity_overlap, gather, gather_optional_multiway,
    gather_optional_row_labels, gather_rowmajor, prepare_propensity_problem_with_registry,
    restrict_to_rows, split_by_treatment, trim_of, trim_retained_rows,
};
use crate::adjustment::EffectEstimate;
use crate::error::EstimationError;
use crate::overlap::{IpwTarget, OverlapPolicy};
use crate::se::{AnalyticSeKind, influence_se_kind};
use crate::util::{BootstrapSeResult, bootstrap_se, sample_std, stats_err};

/// Scale on which the propensity-score matching distance (and [`PropensityMatching::caliper`])
/// is computed.
///
/// The field-standard convention (Rosenbaum & Rubin 1985; Austin 2011, "Optimal caliper
/// widths for propensity-score matching") is that a caliper is applied to the **logit** of
/// the propensity score, `logit(e) = ln(e / (1 - e))` — a caliper of 0.2 means "0.2 standard
/// deviations of the logit propensity," not 0.2 of raw probability. The raw-probability scale
/// compresses badly near 0 and 1 (exactly where matching quality matters most for units with
/// extreme propensities), so a caller who follows the textbook 0.2 rule of thumb but matches
/// on raw probability gets behavior very different from the convention they intended.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum CaliperScale {
    /// Match on `logit(e) = ln(e / (1 - e))`. This is what the 0.2 rule of thumb means.
    #[default]
    Logit,
    /// Match on the raw (clipped) propensity probability, `e ∈ [clip, 1 - clip]`.
    Raw,
}

/// Propensity-score nearest-neighbor matching (Absolute distance, optional caliper).
///
/// Positivity is mandatory: [`OverlapPolicy::ExplicitOverride`] is refused. Supports
/// ATT/ATC/ATE via `TargetPopulation`.
///
/// Analytic SEs use Abadie–Imbens (2006) donor-reuse variance; see module docs for the
/// bootstrap caveat (Abadie–Imbens 2008).
///
/// **Caliper scale:** [`Self::caliper`] is interpreted on [`Self::caliper_scale`], which
/// defaults to [`CaliperScale::Logit`] — the field-standard convention (Rosenbaum & Rubin
/// 1985; Austin 2011). A caliper of 0.2 under this default means 0.2 on the logit scale, per
/// the literature's rule of thumb; see [`CaliperScale`] for why the raw-probability scale is
/// not equivalent. This is a behavior change for callers who previously relied on the
/// undocumented raw-probability matching distance; set `caliper_scale: CaliperScale::Raw`
/// explicitly to keep the old behavior.
#[derive(Clone, Debug)]
pub struct PropensityMatching {
    /// Dense linear-algebra backend used for the logistic IRLS fit.
    pub backend: FaerBackend,
    /// Bootstrap replicates (0 = skip bootstrap). Invalid for NN matching CIs — see module docs.
    pub bootstrap_replicates: u32,
    /// Overlap policy; must be [`OverlapPolicy::RequireDiagnostics`].
    pub overlap: OverlapPolicy,
    /// GLM fitting options for the propensity model.
    pub glm_options: GlmOptions,
    /// Optional maximum matching distance for an accepted match, interpreted on
    /// [`Self::caliper_scale`].
    pub caliper: Option<f64>,
    /// Scale on which matching distance and `caliper` are computed. Defaults to
    /// [`CaliperScale::Logit`] — see the type-level docs.
    pub caliper_scale: CaliperScale,
    /// Analytic SE kind (Abadie–Imbens / hetero / cluster).
    pub se_kind: AnalyticSeKind,
    /// Optional cluster ids aligned to prepared complete-case rows.
    pub cluster_ids: Option<Vec<u32>>,
    /// Optional bindings for named predicates / custom target distributions.
    pub population_registry: Option<PopulationRegistry>,
    /// Multiway cluster ids (one `Vec<u32>` per clustering dimension).
    pub multiway_ids: Option<Vec<Vec<u32>>>,
    /// Optional panel time labels for panel HAC.
    pub panel_times: Option<Vec<i64>>,
}

impl Default for PropensityMatching {
    fn default() -> Self {
        Self::new()
    }
}

impl PropensityMatching {
    /// Defaults: no caliper, 200 bootstrap replicates, clip = 0.01, no trim.
    #[must_use]
    pub fn new() -> Self {
        Self {
            backend: FaerBackend,
            bootstrap_replicates: 200,
            overlap: default_propensity_overlap(),
            glm_options: GlmOptions::default(),
            caliper: None,
            caliper_scale: CaliperScale::Logit,
            se_kind: AnalyticSeKind::Homoskedastic,
            cluster_ids: None,
            population_registry: None,
            multiway_ids: None,
            panel_times: None,
        }
    }

    /// Set the dense linear-algebra backend used for the logistic IRLS fit.
    #[must_use]
    pub const fn with_backend(mut self, backend: FaerBackend) -> Self {
        self.backend = backend;
        self
    }

    /// Set the number of bootstrap replicates (0 = skip bootstrap).
    ///
    /// Defaults to 200, but the nonparametric bootstrap is invalid for nearest-neighbor
    /// matching CIs (Abadie–Imbens 2008) — see the module docs. Prefer the analytic SE.
    #[must_use]
    pub const fn with_bootstrap_replicates(mut self, replicates: u32) -> Self {
        self.bootstrap_replicates = replicates;
        self
    }

    /// Set the overlap policy. Positivity is mandatory here:
    /// [`OverlapPolicy::ExplicitOverride`] is refused by `prepare`.
    #[must_use]
    pub const fn with_overlap(mut self, overlap: OverlapPolicy) -> Self {
        self.overlap = overlap;
        self
    }

    /// Set the GLM fitting options for the propensity model.
    #[must_use]
    pub const fn with_glm_options(mut self, glm_options: GlmOptions) -> Self {
        self.glm_options = glm_options;
        self
    }

    /// Set the maximum matching distance for an accepted match, interpreted on
    /// [`Self::caliper_scale`] (default [`CaliperScale::Logit`] — see the type-level docs for
    /// why 0.2 on that scale is the literature's rule of thumb).
    ///
    /// Defaults to `None` (no caliper): every query row is matched to its nearest donor
    /// regardless of distance.
    #[must_use]
    pub const fn with_caliper(mut self, caliper: f64) -> Self {
        self.caliper = Some(caliper);
        self
    }

    /// Set the scale on which matching distance and [`Self::caliper`] are computed.
    #[must_use]
    pub const fn with_caliper_scale(mut self, caliper_scale: CaliperScale) -> Self {
        self.caliper_scale = caliper_scale;
        self
    }

    /// Set the analytic SE kind (Abadie–Imbens / hetero / cluster).
    #[must_use]
    pub const fn with_se_kind(mut self, se_kind: AnalyticSeKind) -> Self {
        self.se_kind = se_kind;
        self
    }

    /// Set cluster ids aligned to prepared complete-case rows.
    #[must_use]
    pub fn with_cluster_ids(mut self, cluster_ids: Vec<u32>) -> Self {
        self.cluster_ids = Some(cluster_ids);
        self
    }

    /// Set bindings for named predicates / custom target distributions.
    #[must_use]
    pub fn with_population_registry(mut self, registry: PopulationRegistry) -> Self {
        self.population_registry = Some(registry);
        self
    }

    /// Set multiway cluster ids (one `Vec<u32>` per clustering dimension).
    #[must_use]
    pub fn with_multiway_ids(mut self, multiway_ids: Vec<Vec<u32>>) -> Self {
        self.multiway_ids = Some(multiway_ids);
        self
    }

    /// Set panel time labels for panel HAC.
    #[must_use]
    pub fn with_panel_times(mut self, panel_times: Vec<i64>) -> Self {
        self.panel_times = Some(panel_times);
        self
    }

    /// Prepare the covariate design.
    ///
    /// # Errors
    ///
    /// See [`PropensityWeighting::prepare`](crate::propensity::PropensityWeighting::prepare).
    pub fn prepare(
        &self,
        data: &TabularData,
        estimand: &IdentifiedEstimand,
        query: &AverageEffectQuery,
    ) -> Result<PreparedPropensityProblem, EstimationError> {
        prepare_propensity_problem_with_registry(
            data,
            estimand,
            query,
            self.overlap,
            self.population_registry.as_ref(),
        )
    }

    /// Fit the propensity model and compute the matched effect.
    ///
    /// # Errors
    ///
    /// Unsupported target population, empty treated/control arm, no matches within the
    /// caliper, or GLM failure.
    pub fn fit(
        &self,
        problem: &PreparedPropensityProblem,
        workspace: &mut PropensityEstimationWorkspace,
        ctx: &ExecutionContext,
        assumptions: AssumptionSet,
    ) -> Result<EffectEstimate, EstimationError> {
        let trim = trim_of(problem.overlap);
        let model = PropensityModel::fit(
            problem,
            &self.backend,
            &mut workspace.propensity,
            &self.glm_options,
        )?;
        // Trim on RAW scores (mirrors PropensityWeighting): both query and donor sets are
        // restricted to common-support rows before matching.
        let retained = trim_retained_rows(&model.fit.scores, trim)?;
        let (t_used, y_used, s_used) = restrict_to_rows(
            &problem.treatment,
            &problem.outcome,
            &model.clipped_scores,
            1,
            retained.as_deref(),
        );
        let s_used = apply_caliper_scale(s_used, self.caliper_scale);
        let tw_used: Option<Vec<f64>> = problem.target_weights.as_ref().map(|w| match &retained {
            Some(idx) => idx.iter().map(|&i| w[i]).collect(),
            None => w.to_vec(),
        });
        let clusters_used = gather_optional_row_labels(
            self.cluster_ids.as_deref(),
            problem.nrows,
            retained.as_deref(),
            "cluster_ids",
        )?;
        let times_used = gather_optional_row_labels(
            self.panel_times.as_deref(),
            problem.nrows,
            retained.as_deref(),
            "panel_times",
        )?;
        let multiway_used = gather_optional_multiway(
            self.multiway_ids.as_deref(),
            problem.nrows,
            retained.as_deref(),
        )?;
        let result = matching_contrast(
            &t_used,
            &y_used,
            &s_used,
            1,
            MatchingDistance::Absolute,
            &problem.target_population,
            self.caliper,
            workspace,
            self.se_kind,
            clusters_used.as_deref(),
            tw_used.as_deref(),
            multiway_used.as_ref(),
            times_used.as_deref(),
        )?;

        let boot = if self.bootstrap_replicates == 0 {
            None
        } else {
            Some(self.bootstrap_se(problem, trim, workspace, ctx)?)
        };

        let ipw_target = IpwTarget::from_population(&problem.target_population).ok();
        let mut overlap_report = crate::propensity::propensity_overlap_report(
            problem,
            &model.fit.scores,
            None,
            ipw_target,
        );
        overlap_report.retained_fraction *= result.retained_fraction;
        let overlap_report = Some(overlap_report);

        Ok(EffectEstimate::new(result.ate, result.se_analytic, assumptions, problem.overlap)
            .with_overlap_report(overlap_report)
            .with_retained_memory_bytes(Some(workspace.retained_memory_bytes()))
            .with_bootstrap(boot))
    }

    fn bootstrap_se(
        &self,
        problem: &PreparedPropensityProblem,
        trim: Option<f64>,
        workspace: &mut PropensityEstimationWorkspace,
        ctx: &ExecutionContext,
    ) -> Result<BootstrapSeResult, EstimationError> {
        let clip = clip_of(problem.overlap);
        let n = problem.nrows;
        let ncols = problem.design_ncols;
        let mut x_boot = vec![0.0; n * ncols];
        let mut t_boot = vec![0.0; n];
        let mut y_boot = vec![0.0; n];
        bootstrap_se(self.bootstrap_replicates, ctx, 0x51E7_u64, n, |idx| {
            crate::util::gather_bootstrap_vector(&mut t_boot, &problem.treatment, idx);
            crate::util::gather_bootstrap_vector(&mut y_boot, &problem.outcome, idx);
            crate::util::gather_bootstrap_design(
                &mut x_boot,
                &problem.design_matrix,
                n,
                ncols,
                idx,
            );
            let Ok(fit) = fit_propensity(
                &x_boot,
                n,
                ncols,
                &t_boot,
                &self.backend,
                &mut workspace.propensity,
                &self.glm_options,
            ) else {
                return Ok(None);
            };
            let raw = fit.scores;
            let mut scores = raw.clone();
            if let Some(c) = clip {
                clamp_scores(&mut scores, c);
            }
            let Ok(retained) = trim_retained_rows(&raw, trim) else {
                return Ok(None);
            };
            let (t_used, y_used, s_used) =
                restrict_to_rows(&t_boot, &y_boot, &scores, 1, retained.as_deref());
            let s_used = apply_caliper_scale(s_used, self.caliper_scale);
            match matching_contrast(
                &t_used,
                &y_used,
                &s_used,
                1,
                MatchingDistance::Absolute,
                &problem.target_population,
                self.caliper,
                workspace,
                AnalyticSeKind::Homoskedastic,
                None,
                None,
                None,
                None,
            ) {
                Ok(m) => Ok(Some(m.ate)),
                Err(_) => Ok(None),
            }
        })
    }
}

/// Transform clipped propensity scores onto `scale` for use as the `Absolute`-distance
/// matching feature (and, transitively, the caliper comparison).
///
/// Scores passed in here are always the already-clipped `[clip, 1 - clip]` scores (default
/// `[0.01, 0.99]`), so `logit(e) = ln(e / (1 - e))` is always finite — no additional epsilon
/// guard against `e == 0` or `e == 1` is needed.
fn apply_caliper_scale(mut scores: Vec<f64>, scale: CaliperScale) -> Vec<f64> {
    if let CaliperScale::Logit = scale {
        for s in &mut scores {
            *s = (*s / (1.0 - *s)).ln();
        }
    }
    scores
}

/// Match each `query` row to its nearest `donor` row; returns bias-corrected
/// `query_y[q] − donor_y[matched]` and the local donor indices used (for `Kᵢ`).
///
/// Reuses [`PropensityEstimationWorkspace`]'s cached [`MatchingIndex`] when donor geometry
/// is unchanged.
pub(crate) fn match_diffs(
    donor_features: &[f64],
    donor_outcome: &[f64],
    dim: usize,
    distance: MatchingDistance,
    query_features: &[f64],
    query_outcome: &[f64],
    caliper: Option<f64>,
    workspace: &mut PropensityEstimationWorkspace,
) -> Result<(Vec<f64>, Vec<usize>, Vec<usize>), EstimationError> {
    let n_donors = donor_outcome.len();
    if n_donors == 0 {
        return Err(EstimationError::data_msg("matching requires at least one donor row"));
    }
    workspace.ensure_matching_index(donor_features, dim, distance)?;
    let n_queries = query_outcome.len();
    let mut donor_rows = std::mem::take(&mut workspace.matching_donor_rows);
    let mut distances = std::mem::take(&mut workspace.matching_distances);
    donor_rows.clear();
    donor_rows.resize(n_queries, 0);
    distances.clear();
    distances.resize(n_queries, 0.0);
    {
        let index = workspace.matching_index.as_ref().expect("ensured");
        index
            .match_all(query_features, n_queries, caliper, &mut donor_rows, &mut distances)
            .map_err(stats_err)?;
    }
    let mut diffs = Vec::with_capacity(n_queries);
    let mut used_donors = Vec::with_capacity(n_queries);
    let mut used_queries = Vec::with_capacity(n_queries);
    let mu_donor = fit_linear_mean(donor_features, donor_outcome, dim);
    for q in 0..n_queries {
        let d = donor_rows[q];
        if d != usize::MAX {
            let raw = query_outcome[q] - donor_outcome[d];
            let bias = match &mu_donor {
                Some(beta) => {
                    let mq = predict_linear(beta, query_features, dim, q);
                    let md = predict_linear(beta, donor_features, dim, d);
                    mq - md
                }
                None => 0.0,
            };
            diffs.push(raw - bias);
            used_donors.push(d);
            used_queries.push(q);
        }
    }
    workspace.matching_donor_rows = donor_rows;
    workspace.matching_distances = distances;
    Ok((diffs, used_donors, used_queries))
}

pub(crate) struct MatchedEstimate {
    pub(crate) ate: f64,
    pub(crate) se_analytic: f64,
    pub(crate) retained_fraction: f64,
}

/// ATT/ATC/ATE via nearest-neighbor matching on `features` (dim columns, row-major).
///
/// ATT matches treated→nearest control; ATC matches control→nearest treated (sign-flipped);
/// ATE pools both directions' per-unit imputed effects (Abadie–Imbens style).
#[allow(clippy::too_many_arguments)]
pub(crate) fn matching_contrast(
    treatment: &[f64],
    outcome: &[f64],
    features: &[f64],
    dim: usize,
    distance: MatchingDistance,
    target: &TargetPopulation,
    caliper: Option<f64>,
    workspace: &mut PropensityEstimationWorkspace,
    se_kind: AnalyticSeKind,
    cluster_ids: Option<&[u32]>,
    target_weights: Option<&[f64]>,
    multiway_ids: Option<&Vec<Vec<u32>>>,
    panel_times: Option<&[i64]>,
) -> Result<MatchedEstimate, EstimationError> {
    if let Some(ids) = cluster_ids {
        if ids.len() != treatment.len() {
            return Err(EstimationError::data_msg("matching cluster_ids length != treatment rows"));
        }
    }
    if let Some(times) = panel_times {
        if times.len() != treatment.len() {
            return Err(EstimationError::data_msg("matching panel_times length != treatment rows"));
        }
    }
    if let Some(dims) = multiway_ids {
        for (i, d) in dims.iter().enumerate() {
            if d.len() != treatment.len() {
                return Err(EstimationError::data_msg(format!(
                    "matching multiway_ids[{i}] length {} != treatment rows",
                    d.len()
                )));
            }
        }
    }
    let (treated_idx, control_idx) = split_by_treatment(treatment);
    if treated_idx.is_empty() || control_idx.is_empty() {
        return Err(EstimationError::data_msg("matching requires both treated and control rows"));
    }
    let treated_feat = gather_rowmajor(features, dim, &treated_idx);
    let control_feat = gather_rowmajor(features, dim, &control_idx);
    let treated_y = gather(outcome, &treated_idx);
    let control_y = gather(outcome, &control_idx);

    let (per_unit_effects, donor_usage, n_donors, effect_rows): (
        Vec<f64>,
        Vec<usize>,
        usize,
        Vec<usize>,
    ) = match target {
        TargetPopulation::Treated => {
            let (diffs, donors, q_local) = match_diffs(
                &control_feat,
                &control_y,
                dim,
                distance,
                &treated_feat,
                &treated_y,
                caliper,
                workspace,
            )?;
            let rows: Vec<usize> = q_local.iter().map(|&q| treated_idx[q]).collect();
            (diffs, donors, control_y.len(), rows)
        }
        TargetPopulation::Untreated => {
            let (diffs, donors, q_local) = match_diffs(
                &treated_feat,
                &treated_y,
                dim,
                distance,
                &control_feat,
                &control_y,
                caliper,
                workspace,
            )?;
            let flipped: Vec<f64> = diffs.into_iter().map(|d| -d).collect();
            let rows: Vec<usize> = q_local.iter().map(|&q| control_idx[q]).collect();
            (flipped, donors, treated_y.len(), rows)
        }
        TargetPopulation::AllObserved
        | TargetPopulation::Predicate(_)
        | TargetPopulation::CustomDistribution(_) => {
            let (att_diffs, att_donors, att_q) = match_diffs(
                &control_feat,
                &control_y,
                dim,
                distance,
                &treated_feat,
                &treated_y,
                caliper,
                workspace,
            )?;
            let (atc_raw, atc_donors, atc_q) = match_diffs(
                &treated_feat,
                &treated_y,
                dim,
                distance,
                &control_feat,
                &control_y,
                caliper,
                workspace,
            )?;
            let atc_diffs: Vec<f64> = atc_raw.into_iter().map(|d| -d).collect();
            let n_control = control_y.len();
            let mut effects = att_diffs;
            effects.extend(atc_diffs);
            let mut donors = att_donors;
            donors.extend(atc_donors.into_iter().map(|d| d + n_control));
            let mut rows: Vec<usize> = att_q.iter().map(|&q| treated_idx[q]).collect();
            rows.extend(atc_q.iter().map(|&q| control_idx[q]));
            (effects, donors, n_control + treated_y.len(), rows)
        }
        _ => {
            return Err(EstimationError::unsupported(
                "matching estimators support AllObserved, Treated, Untreated, Predicate, or CustomDistribution",
            ));
        }
    };
    if per_unit_effects.is_empty() {
        return Err(EstimationError::data_msg("no matched units within caliper"));
    }
    let n_eligible = match target {
        TargetPopulation::Treated => treated_idx.len(),
        TargetPopulation::Untreated => control_idx.len(),
        _ => treated_idx.len() + control_idx.len(),
    };
    let retained_fraction = per_unit_effects.len() as f64 / n_eligible.max(1) as f64;
    let ate = if matches!(target, TargetPopulation::CustomDistribution(_)) {
        let Some(tw) = target_weights else {
            return Err(EstimationError::unsupported(
                "CustomDistribution requires PopulationRegistry weights on the prepared problem",
            ));
        };
        let mut num = 0.0;
        let mut den = 0.0;
        for (i, &eff) in per_unit_effects.iter().enumerate() {
            let w = tw.get(effect_rows[i]).copied().unwrap_or(0.0);
            num += w * eff;
            den += w;
        }
        if den <= 0.0 {
            return Err(EstimationError::data_msg(
                "CustomDistribution weights left no mass on matched units",
            ));
        }
        num / den
    } else {
        per_unit_effects.iter().sum::<f64>() / per_unit_effects.len() as f64
    };
    let se_analytic = match se_kind {
        AnalyticSeKind::Homoskedastic => {
            abadie_imbens_se(&per_unit_effects, &donor_usage, n_donors)
        }
        AnalyticSeKind::Hc0 | AnalyticSeKind::Hc1 | AnalyticSeKind::Hc2 | AnalyticSeKind::Hc3 => {
            return Err(EstimationError::unsupported(
                "matching does not implement HC0–HC3 sandwich SEs; use Homoskedastic (Abadie–Imbens) or Cluster",
            ));
        }
        AnalyticSeKind::Cluster
        | AnalyticSeKind::Multiway
        | AnalyticSeKind::NeweyWest { .. }
        | AnalyticSeKind::PanelClusterHac { .. } => {
            let mut k = vec![0usize; n_donors.max(1)];
            for &d in &donor_usage {
                if d < k.len() {
                    k[d] += 1;
                }
            }
            let mut psi = Vec::with_capacity(per_unit_effects.len());
            for (i, &d) in donor_usage.iter().enumerate() {
                let kd = k.get(d).copied().unwrap_or(0) as f64;
                psi.push((per_unit_effects[i] - ate) * (1.0 + kd));
            }
            influence_se_kind(
                se_kind,
                &psi,
                treatment.len(),
                cluster_ids,
                multiway_ids.map(Vec::as_slice),
                panel_times,
                Some(&effect_rows),
            )?
        }
    };
    Ok(MatchedEstimate { ate, se_analytic, retained_fraction })
}

/// Abadie–Imbens (2006) SE for 1-NN matching with replacement (homoskedastic).
///
/// With unit-level matched effects `τ̂ᵢ` and donor reuse counts `Kⱼ`,
/// `Var = σ̂² (n + Σⱼ Kⱼ²) / n²` where `σ̂² = Var(τ̂ᵢ) / 2` (equal-arm residual variance).
fn abadie_imbens_se(effects: &[f64], donor_local: &[usize], n_donors: usize) -> f64 {
    let n = effects.len();
    if n < 2 || donor_local.len() != n {
        return sample_std(effects) / (n as f64).sqrt();
    }
    let mut k = vec![0usize; n_donors.max(1)];
    for &d in donor_local {
        if d < k.len() {
            k[d] += 1;
        }
    }
    let mean = effects.iter().sum::<f64>() / n as f64;
    let var_tau = effects.iter().map(|e| (e - mean).powi(2)).sum::<f64>() / (n as f64 - 1.0);
    let sigma2 = (var_tau * 0.5).max(0.0);
    let sum_k2: f64 = k.iter().map(|&kj| (kj as f64).powi(2)).sum();
    let var = sigma2 * (n as f64 + sum_k2) / (n as f64).powi(2);
    var.sqrt()
}

/// Heteroskedastic Abadie–Imbens SE using demeaned pair-level variance proxies.
/// Not exposed as `Hc0`–`Hc3` (those names are sandwich estimators matching does not implement).
#[allow(dead_code)]
fn abadie_imbens_se_hetero(effects: &[f64], donor_local: &[usize], n_donors: usize) -> f64 {
    let n = effects.len();
    if n < 2 || donor_local.len() != n {
        return sample_std(effects) / (n as f64).sqrt();
    }
    let mut k = vec![0usize; n_donors.max(1)];
    for &d in donor_local {
        if d < k.len() {
            k[d] += 1;
        }
    }
    let mean = effects.iter().sum::<f64>() / n as f64;
    let mut var = 0.0;
    for (i, &d) in donor_local.iter().enumerate() {
        let centered = effects[i] - mean;
        let sigma2_i = 0.5 * centered * centered;
        let kd = k.get(d).copied().unwrap_or(0) as f64;
        var += sigma2_i * (1.0 + kd).powi(2);
    }
    (var / (n as f64).powi(2)).max(0.0).sqrt()
}

/// OLS of `y` on `[1, x]` (row-major `x` with `dim` columns). Returns `[intercept, β…]`.
fn fit_linear_mean(features: &[f64], y: &[f64], dim: usize) -> Option<Vec<f64>> {
    let n = y.len();
    if n < dim + 1 || dim == 0 {
        // Fall back to intercept-only mean when underdetermined.
        if n == 0 {
            return None;
        }
        return Some(vec![y.iter().sum::<f64>() / n as f64]);
    }
    let p = dim + 1;
    let mut xtx = vec![0.0; p * p];
    let mut xty = vec![0.0; p];
    for i in 0..n {
        let mut row = vec![1.0; p];
        for d in 0..dim {
            row[d + 1] = features[i * dim + d];
        }
        for a in 0..p {
            xty[a] += row[a] * y[i];
            for b in 0..p {
                xtx[a * p + b] += row[a] * row[b];
            }
        }
    }
    solve_linear_system(&mut xtx, &mut xty, p)
}

fn predict_linear(beta: &[f64], features: &[f64], dim: usize, row: usize) -> f64 {
    if beta.len() == 1 {
        return beta[0];
    }
    let mut y = beta[0];
    for d in 0..dim.min(beta.len().saturating_sub(1)) {
        y += beta[d + 1] * features[row * dim + d];
    }
    y
}

/// Gaussian elimination with partial pivoting; returns solution in `b`, or `None` if singular.
fn solve_linear_system(a: &mut [f64], b: &mut [f64], p: usize) -> Option<Vec<f64>> {
    for col in 0..p {
        let mut pivot = col;
        let mut best = a[col * p + col].abs();
        for r in (col + 1)..p {
            let v = a[r * p + col].abs();
            if v > best {
                best = v;
                pivot = r;
            }
        }
        if best < 1e-14 {
            return None;
        }
        if pivot != col {
            for c in 0..p {
                a.swap(col * p + c, pivot * p + c);
            }
            b.swap(col, pivot);
        }
        let diag = a[col * p + col];
        for r in (col + 1)..p {
            let f = a[r * p + col] / diag;
            for c in col..p {
                a[r * p + c] -= f * a[col * p + c];
            }
            b[r] -= f * b[col];
        }
    }
    let mut x = vec![0.0; p];
    for i in (0..p).rev() {
        let mut s = b[i];
        for j in (i + 1)..p {
            s -= a[i * p + j] * x[j];
        }
        x[i] = s / a[i * p + i];
    }
    Some(x)
}

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

    #[test]
    fn abadie_imbens_se_grows_with_donor_reuse() {
        let effects = [1.0, 1.2, 0.8, 1.1];
        // Four queries, two unique donors reused twice each.
        let donors_reuse = vec![0usize, 0, 1, 1];
        let donors_unique = vec![0usize, 1, 2, 3];
        let se_reuse = abadie_imbens_se(&effects, &donors_reuse, 2);
        let se_unique = abadie_imbens_se(&effects, &donors_unique, 4);
        assert!(se_reuse > se_unique, "reuse={se_reuse} unique={se_unique}");
    }

    #[test]
    fn cluster_se_grows_with_donor_reuse() {
        let effects = [1.0, 1.2, 0.8, 1.1];
        let ate = effects.iter().sum::<f64>() / effects.len() as f64;
        let donors_reuse = vec![0usize, 0, 1, 1];
        let donors_unique = vec![0usize, 1, 2, 3];
        let groups = vec![0u32, 0, 1, 1];
        let se = |donors: &[usize], n_donors: usize| {
            let mut k = vec![0usize; n_donors.max(1)];
            for &d in donors {
                if d < k.len() {
                    k[d] += 1;
                }
            }
            let psi: Vec<f64> = effects
                .iter()
                .enumerate()
                .map(|(i, &e)| {
                    let kd = k.get(donors[i]).copied().unwrap_or(0) as f64;
                    (e - ate) * (1.0 + kd)
                })
                .collect();
            crate::se::cluster_influence_se(&psi, &groups).unwrap()
        };
        let se_reuse = se(&donors_reuse, 2);
        let se_unique = se(&donors_unique, 4);
        assert!(se_reuse > se_unique, "cluster reuse={se_reuse} unique={se_unique}");
    }

    /// Regression test for the caliper-scale defect: with a fixed numeric caliper of 0.2 and
    /// propensity scores near the extremes of `[0, 1]`, `CaliperScale::Raw` and
    /// `CaliperScale::Logit` must select materially different match sets. This is the
    /// discriminating case the literature's 0.2 rule of thumb is about — raw-probability
    /// caliper 0.2 is far too permissive near 0/1, while the logit-scale caliper (what 0.2
    /// actually means per Rosenbaum & Rubin 1985 / Austin 2011) correctly excludes those
    /// pairs.
    #[test]
    fn caliper_scale_logit_vs_raw_diverge_near_extremes() {
        // Donor at each extreme plus one in the middle; queries just inside each extreme.
        let donor_probs = vec![0.02, 0.5, 0.98];
        let donor_y = vec![0.0, 0.0, 0.0];
        let query_probs = vec![0.08, 0.5, 0.92];
        let query_y = vec![0.0, 0.0, 0.0];
        let caliper = Some(0.2);

        let raw_donors = apply_caliper_scale(donor_probs.clone(), CaliperScale::Raw);
        let raw_queries = apply_caliper_scale(query_probs.clone(), CaliperScale::Raw);
        let mut ws_raw = PropensityEstimationWorkspace::default();
        let (raw_diffs, _, _) = match_diffs(
            &raw_donors,
            &donor_y,
            1,
            MatchingDistance::Absolute,
            &raw_queries,
            &query_y,
            caliper,
            &mut ws_raw,
        )
        .unwrap();

        let logit_donors = apply_caliper_scale(donor_probs, CaliperScale::Logit);
        let logit_queries = apply_caliper_scale(query_probs, CaliperScale::Logit);
        let mut ws_logit = PropensityEstimationWorkspace::default();
        let (logit_diffs, _, _) = match_diffs(
            &logit_donors,
            &donor_y,
            1,
            MatchingDistance::Absolute,
            &logit_queries,
            &query_y,
            caliper,
            &mut ws_logit,
        )
        .unwrap();

        // Raw-probability caliper 0.2 is loose everywhere: all three queries land within 0.2
        // of a donor (|0.08-0.02|=0.06, |0.5-0.5|=0, |0.92-0.98|=0.06).
        assert_eq!(
            raw_diffs.len(),
            3,
            "raw-scale caliper=0.2 should admit all 3 near-extreme queries, got {}",
            raw_diffs.len()
        );
        // Logit caliper 0.2 correctly rejects the extreme pairs: logit(0.08) to logit(0.02) is
        // ~1.45 apart (>> 0.2), and symmetrically at the top; only the mid-range query (whose
        // donor and query coincide, logit distance 0) survives.
        assert_eq!(
            logit_diffs.len(),
            1,
            "logit-scale caliper=0.2 should admit only the mid-range query, got {}",
            logit_diffs.len()
        );
    }
}