//! PyO3 bindings for StatsPAI's HDFE inner kernel.
//!
//! Phase 1 contents:
//! - ``group_demean``      — single-FE, in-place column demean (legacy entry point).
//! - ``demean_2d``         — K-way alternating-projection demean with Irons-Tuck
//!                           acceleration; in-place on a Fortran-order matrix,
//!                           parallel over columns via Rayon.
//! - ``singleton_mask``    — iterative K-way singleton-row detection; returns
//!                           a boolean keep-mask.
//!
//! Phase A additions (v0.3.0, weighted variants for IRLS-internal demean):
//! - ``demean_2d_weighted`` — same as ``demean_2d`` but takes per-observation
//!                            weights and a caller-precomputed wsum
//!                            (``Σ_{i ∈ g} weights[i]``); used by the IRLS
//!                            inner loop in ``sp.fast.fepois``.
//!
//! All functions take pre-factorised int64 codes and float64 counts. The
//! Python wrapper at ``statspai.fast.demean`` packs DataFrames / mixed
//! dtypes via ``pd.factorize`` before calling here.
//!
//! All wheels are optional — the Python side falls back to the NumPy /
//! Numba kernel gracefully when the compiled extension is missing.

mod demean;
mod singletons;
mod sort_perm;
mod cholesky;
mod irls;
mod separation;
mod cluster;

use numpy::{
    PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray2, PyReadwriteArray1,
    PyReadwriteArray2, PyUntypedArrayMethods,
};
use numpy::ndarray::Array2;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use pyo3::wrap_pyfunction;

/// In-place group demean: `y[i] -= mean(y[codes == codes[i]])`.
///
/// Legacy single-FE entry point. Kept for backward compatibility with
/// ``statspai.panel.hdfe_rust.group_demean_rust``; new callers should
/// prefer the more general ``demean_2d``.
#[pyfunction]
fn group_demean(
    codes: PyReadonlyArray1<i64>,
    mut y: PyReadwriteArray1<f64>,
    mut sums: PyReadwriteArray1<f64>,
    counts: PyReadonlyArray1<i64>,
) -> PyResult<()> {
    let codes = codes.as_slice()?;
    let y_slice = y.as_slice_mut()?;
    let sums_slice = sums.as_slice_mut()?;
    let counts_slice = counts.as_slice()?;

    for s in sums_slice.iter_mut() {
        *s = 0.0;
    }
    for i in 0..y_slice.len() {
        let g = codes[i] as usize;
        sums_slice[g] += y_slice[i];
    }
    for g in 0..sums_slice.len() {
        let c = counts_slice[g];
        if c > 0 {
            sums_slice[g] /= c as f64;
        }
    }
    for i in 0..y_slice.len() {
        let g = codes[i] as usize;
        y_slice[i] -= sums_slice[g];
    }
    Ok(())
}

/// Helper: extract a list of int64 array views from a PyList.
fn py_list_to_i64_views<'py>(
    list: &Bound<'py, PyList>,
) -> PyResult<Vec<PyReadonlyArray1<'py, i64>>> {
    let mut out = Vec::with_capacity(list.len());
    for item in list.iter() {
        let arr: PyReadonlyArray1<i64> = item.extract()?;
        out.push(arr);
    }
    Ok(out)
}

/// Helper: extract a list of float64 array views from a PyList.
fn py_list_to_f64_views<'py>(
    list: &Bound<'py, PyList>,
) -> PyResult<Vec<PyReadonlyArray1<'py, f64>>> {
    let mut out = Vec::with_capacity(list.len());
    for item in list.iter() {
        let arr: PyReadonlyArray1<f64> = item.extract()?;
        out.push(arr);
    }
    Ok(out)
}

/// K-way alternating-projection demean of a Fortran-order (n, p) matrix in
/// place. Returns a list of dicts (one per column) with iter / converged /
/// max_dx fields.
///
/// Parameters
/// ----------
/// x : 2-D float64 ndarray, shape (n, p), Fortran-contiguous
///     The matrix to residualise. Each column is a contiguous slice we can
///     split for parallel processing. Pass ``np.asfortranarray(X)`` on the
///     Python side to materialise.
/// fe_codes : list[ndarray[int64, shape (n,)]]
///     One code array per FE dimension (K total).
/// counts : list[ndarray[float64, shape (G_k,)]]
///     Per-group sizes for each FE dimension. Float so weighted variants
///     can drop in later.
/// max_iter : int
///     Cap on AP iterations per column.
/// tol_abs, tol_rel : float
///     Stop when ``max|dx| <= tol_abs + tol_rel * base_scale``.
/// accelerate : bool
/// accel_period : int
#[pyfunction]
#[pyo3(signature = (x, fe_codes, counts, max_iter, tol_abs, tol_rel, accelerate, accel_period))]
#[allow(clippy::too_many_arguments)]
fn demean_2d<'py>(
    py: Python<'py>,
    mut x: PyReadwriteArray2<'py, f64>,
    fe_codes: &Bound<'py, PyList>,
    counts: &Bound<'py, PyList>,
    max_iter: u32,
    tol_abs: f64,
    tol_rel: f64,
    accelerate: bool,
    accel_period: u32,
) -> PyResult<Bound<'py, PyList>> {
    if fe_codes.len() != counts.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "len(fe_codes)={} but len(counts)={}",
            fe_codes.len(),
            counts.len()
        )));
    }

    let code_views = py_list_to_i64_views(fe_codes)?;
    let count_views = py_list_to_f64_views(counts)?;

    let arr = x.as_array();
    let shape = arr.shape();
    if shape.len() != 2 {
        return Err(pyo3::exceptions::PyValueError::new_err("x must be 2-D"));
    }
    let n = shape[0];
    let p = shape[1];

    for v in &code_views {
        if v.as_slice()?.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "fe_codes entry has length {} but n={}",
                v.as_slice()?.len(),
                n
            )));
        }
    }

    if !x.is_fortran_contiguous() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x must be Fortran-contiguous; pass np.asfortranarray(X)",
        ));
    }
    let mat = x.as_slice_mut()?;
    let codes_slices: Vec<&[i64]> =
        code_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let counts_slices: Vec<&[f64]> =
        count_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let counts_lens: Vec<usize> = counts_slices.iter().map(|s| s.len()).collect();

    let infos = py.allow_threads(|| {
        demean::demean_matrix_fortran_inplace(
            mat,
            n,
            p,
            &codes_slices,
            &counts_slices,
            &counts_lens,
            max_iter,
            tol_abs,
            tol_rel,
            accelerate,
            accel_period,
        )
    });

    let out = PyList::empty_bound(py);
    for info in &infos {
        let d = PyDict::new_bound(py);
        d.set_item("iters", info.iters)?;
        d.set_item("converged", info.converged)?;
        d.set_item("max_dx", info.max_dx)?;
        out.append(d)?;
    }
    Ok(out)
}

/// K-way **weighted** alternating-projection demean of a Fortran-order
/// (n, p) matrix in place. Returns a list of dicts (one per column)
/// with ``iters`` / ``converged`` / ``max_dx`` fields, mirroring
/// ``demean_2d``.
///
/// Parameters
/// ----------
/// x : 2-D float64 ndarray, shape (n, p), Fortran-contiguous
///     The matrix to residualise (in place).
/// fe_codes : list[ndarray[int64, shape (n,)]]
///     One code array per FE dimension (K total).
/// wsum : list[ndarray[float64, shape (G_k,)]]
///     Per-group **weighted** sum ``Σ_{i ∈ g} weights[i]``. Caller
///     precomputes via ``np.bincount(codes, weights=weights, minlength=G)``.
/// weights : ndarray[float64, shape (n,)]
///     Per-observation weights. Caller is responsible for non-negativity
///     and finiteness — no re-validation here on the hot path.
/// max_iter, tol_abs, tol_rel, accelerate, accel_period
///     Same semantics as ``demean_2d``.
#[pyfunction]
#[pyo3(signature = (x, fe_codes, wsum, weights, max_iter, tol_abs, tol_rel, accelerate, accel_period))]
#[allow(clippy::too_many_arguments)]
fn demean_2d_weighted<'py>(
    py: Python<'py>,
    mut x: PyReadwriteArray2<'py, f64>,
    fe_codes: &Bound<'py, PyList>,
    wsum: &Bound<'py, PyList>,
    weights: PyReadonlyArray1<'py, f64>,
    max_iter: u32,
    tol_abs: f64,
    tol_rel: f64,
    accelerate: bool,
    accel_period: u32,
) -> PyResult<Bound<'py, PyList>> {
    if fe_codes.len() != wsum.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "len(fe_codes)={} but len(wsum)={}",
            fe_codes.len(),
            wsum.len()
        )));
    }

    let code_views = py_list_to_i64_views(fe_codes)?;
    let wsum_views = py_list_to_f64_views(wsum)?;
    let weights_view = weights.as_slice()?;

    let arr = x.as_array();
    let shape = arr.shape();
    if shape.len() != 2 {
        return Err(pyo3::exceptions::PyValueError::new_err("x must be 2-D"));
    }
    let n = shape[0];
    let p = shape[1];

    if weights_view.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "weights length {} but n={}",
            weights_view.len(),
            n
        )));
    }

    for v in &code_views {
        if v.as_slice()?.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "fe_codes entry has length {} but n={}",
                v.as_slice()?.len(),
                n
            )));
        }
    }

    if !x.is_fortran_contiguous() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x must be Fortran-contiguous; pass np.asfortranarray(X)",
        ));
    }
    let mat = x.as_slice_mut()?;
    let codes_slices: Vec<&[i64]> =
        code_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let wsum_slices: Vec<&[f64]> =
        wsum_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let wsum_lens: Vec<usize> = wsum_slices.iter().map(|s| s.len()).collect();

    let infos = py.allow_threads(|| {
        demean::weighted_demean_matrix_fortran_inplace(
            mat,
            n,
            p,
            &codes_slices,
            weights_view,
            &wsum_slices,
            &wsum_lens,
            max_iter,
            tol_abs,
            tol_rel,
            accelerate,
            accel_period,
        )
    });

    let out = PyList::empty_bound(py);
    for info in &infos {
        let d = PyDict::new_bound(py);
        d.set_item("iters", info.iters)?;
        d.set_item("converged", info.converged)?;
        d.set_item("max_dx", info.max_dx)?;
        out.append(d)?;
    }
    Ok(out)
}

/// Sort-aware weighted demean of a Fortran-order (n, p) matrix in place.
/// Caller has applied the primary-FE sort permutation π to ``x`` (rows),
/// ``weights``, and the secondary FE codes; this function does NOT
/// permute. Result is in π-order; caller applies π⁻¹ on return.
///
/// Parameters
/// ----------
/// x : 2-D float64 ndarray, shape (n, p), Fortran-contiguous, in π order.
/// primary_starts : ndarray[int64, shape (G1+1,)]
///     Group-start offsets for the primary FE (caller computes once via
///     ``primary_fe_sort_perm`` + cumulative count).
/// primary_wsum : ndarray[float64, shape (G1,)]
///     Weighted group sums for the primary FE, computed in π order
///     (i.e., ``np.bincount(codes_perm, weights=weights_perm)``).
/// secondary_codes : list[ndarray[int64, shape (n,)]]
///     K-1 arrays, one per non-primary FE; codes are under π.
/// secondary_wsum : list[ndarray[float64, shape (G_k,)]]
///     Weighted group sums for non-primary FEs.
/// weights_sorted : ndarray[float64, shape (n,)]
///     Per-obs weights in π order.
/// max_iter, tol_abs, tol_rel, accelerate, accel_period :
///     Same semantics as ``demean_2d_weighted``.
#[pyfunction]
#[pyo3(signature = (x, primary_starts, primary_wsum, secondary_codes, secondary_wsum, weights_sorted, max_iter, tol_abs, tol_rel, accelerate, accel_period))]
#[allow(clippy::too_many_arguments)]
fn demean_2d_weighted_sorted<'py>(
    py: Python<'py>,
    mut x: PyReadwriteArray2<'py, f64>,
    primary_starts: PyReadonlyArray1<'py, i64>,
    primary_wsum: PyReadonlyArray1<'py, f64>,
    secondary_codes: &Bound<'py, PyList>,
    secondary_wsum: &Bound<'py, PyList>,
    weights_sorted: PyReadonlyArray1<'py, f64>,
    max_iter: u32,
    tol_abs: f64,
    tol_rel: f64,
    accelerate: bool,
    accel_period: u32,
) -> PyResult<Bound<'py, PyList>> {
    if secondary_codes.len() != secondary_wsum.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "len(secondary_codes)={} but len(secondary_wsum)={}",
            secondary_codes.len(),
            secondary_wsum.len()
        )));
    }

    let primary_starts_view = primary_starts.as_slice()?;
    let primary_wsum_view = primary_wsum.as_slice()?;
    let weights_view = weights_sorted.as_slice()?;
    let sec_code_views = py_list_to_i64_views(secondary_codes)?;
    let sec_wsum_views = py_list_to_f64_views(secondary_wsum)?;

    let arr = x.as_array();
    let shape = arr.shape();
    if shape.len() != 2 {
        return Err(pyo3::exceptions::PyValueError::new_err("x must be 2-D"));
    }
    let n = shape[0];
    let p = shape[1];

    if weights_view.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "weights_sorted length {} but n={}",
            weights_view.len(),
            n
        )));
    }
    if primary_starts_view.len() != primary_wsum_view.len() + 1 {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "primary_starts length {} but expected {} (= primary_wsum.len + 1)",
            primary_starts_view.len(),
            primary_wsum_view.len() + 1
        )));
    }
    for v in &sec_code_views {
        if v.as_slice()?.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "secondary_codes entry has length {} but n={}",
                v.as_slice()?.len(),
                n
            )));
        }
    }

    if !x.is_fortran_contiguous() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x must be Fortran-contiguous; pass np.asfortranarray(X)",
        ));
    }
    let mat = x.as_slice_mut()?;

    // Convert primary_starts (passed as i64 from Python) to usize.
    let primary_starts_usize: Vec<usize> =
        primary_starts_view.iter().map(|&v| v as usize).collect();

    let sec_codes_slices: Vec<&[i64]> =
        sec_code_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let sec_wsum_slices: Vec<&[f64]> =
        sec_wsum_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let sec_lens: Vec<usize> = sec_wsum_slices.iter().map(|s| s.len()).collect();

    let infos = py.allow_threads(|| {
        demean::weighted_demean_matrix_fortran_inplace_sorted(
            mat,
            n,
            p,
            &primary_starts_usize,
            primary_wsum_view,
            &sec_codes_slices,
            &sec_wsum_slices,
            &sec_lens,
            weights_view,
            max_iter,
            tol_abs,
            tol_rel,
            accelerate,
            accel_period,
        )
    });

    let out = PyList::empty_bound(py);
    for info in &infos {
        let d = PyDict::new_bound(py);
        d.set_item("iters", info.iters)?;
        d.set_item("converged", info.converged)?;
        d.set_item("max_dx", info.max_dx)?;
        out.append(d)?;
    }
    Ok(out)
}

/// Native Rust Poisson IRLS for ``sp.fast.fepois``.
///
/// Single PyO3 entry point that runs the entire IRLS state machine
/// in Rust — no per-iter FFI round-trips. The Python side parses the
/// formula and runs the singleton/separation pre-passes; this function
/// runs the IRLS body and returns the final β plus the demeaned X̃ and
/// working weights so the caller can compute the requested vcov in
/// Python.
///
/// Parameters
/// ----------
/// y : ndarray[float64, shape (n,)]
///     Outcome vector (after pre-passes).
/// x : ndarray[float64, shape (n, p)], Fortran-contiguous
///     Regressor matrix. The function does NOT mutate `x`.
/// fe_codes : list[ndarray[int64, shape (n,)]]
///     K dense FE code arrays.
/// g_per_fe : ndarray[int64, shape (K,)]
///     Cardinality of each FE.
/// obs_weights : ndarray[float64, shape (n,)]
///     Per-observation weights. Pass an all-1 array for unweighted MLE.
/// config : dict
///     IRLS knobs: maxiter / tol / fe_tol / fe_maxiter / eta_clip /
///     accel_period / max_halvings. Missing keys fall back to
///     `FePoisIRLSConfig::default()`.
///
/// Returns
/// -------
/// dict with keys: ``beta`` (ndarray, shape (p,)),
/// ``x_tilde_flat`` (ndarray, shape (n*p,), F-order data),
/// ``x_tilde_n`` (int), ``x_tilde_p`` (int),
/// ``w`` (shape (n,)), ``eta`` (shape (n,)), ``mu`` (shape (n,)),
/// ``deviance`` (float), ``log_likelihood`` (float), ``iters`` (int),
/// ``converged`` (bool), ``n_halvings`` (int), ``max_inner_dx`` (float).
///
/// Notes
/// -----
/// ``x_tilde`` is returned as a flat 1-D array in Fortran (column-major)
/// order together with shape scalars ``x_tilde_n`` and ``x_tilde_p``.
/// Reconstruct on the Python side with:
/// ``x_tilde = result["x_tilde_flat"].reshape(result["x_tilde_p"],
///              result["x_tilde_n"]).T``
/// which produces a row-major (n, p) view of the F-order data.
#[pyfunction]
#[pyo3(signature = (y, x, fe_codes, g_per_fe, obs_weights, config))]
fn fepois_irls<'py>(
    py: Python<'py>,
    y: PyReadonlyArray1<'py, f64>,
    x: PyReadonlyArray2<'py, f64>,
    fe_codes: &Bound<'py, PyList>,
    g_per_fe: PyReadonlyArray1<'py, i64>,
    obs_weights: PyReadonlyArray1<'py, f64>,
    config: &Bound<'py, PyDict>,
) -> PyResult<Bound<'py, PyDict>> {
    let y_view = y.as_slice()?;
    let obs_w_view = obs_weights.as_slice()?;
    let g_view = g_per_fe.as_slice()?;
    let code_views = py_list_to_i64_views(fe_codes)?;

    let arr = x.as_array();
    let shape = arr.shape();
    if shape.len() != 2 {
        return Err(pyo3::exceptions::PyValueError::new_err("x must be 2-D"));
    }
    let n = shape[0];
    let p = shape[1];

    if y_view.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "y length {} but x has {} rows", y_view.len(), n
        )));
    }
    if obs_w_view.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "obs_weights length {} but x has {} rows", obs_w_view.len(), n
        )));
    }
    if g_view.len() != code_views.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "len(g_per_fe)={} but len(fe_codes)={}",
            g_view.len(), code_views.len()
        )));
    }
    for v in &code_views {
        if v.as_slice()?.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "fe_codes entry has length {} but n={}",
                v.as_slice()?.len(), n
            )));
        }
    }
    if !x.is_fortran_contiguous() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x must be Fortran-contiguous; pass np.asfortranarray(X)",
        ));
    }

    // Read config knobs (each optional — fall back to default).
    let mut cfg = irls::FePoisIRLSConfig::default();
    if let Ok(Some(v)) = config.get_item("maxiter") {
        cfg.maxiter = v.extract::<u32>()?;
    }
    if let Ok(Some(v)) = config.get_item("tol") {
        cfg.tol = v.extract::<f64>()?;
    }
    if let Ok(Some(v)) = config.get_item("fe_tol") {
        cfg.fe_tol = v.extract::<f64>()?;
    }
    if let Ok(Some(v)) = config.get_item("fe_maxiter") {
        cfg.fe_maxiter = v.extract::<u32>()?;
    }
    if let Ok(Some(v)) = config.get_item("eta_clip") {
        cfg.eta_clip = v.extract::<f64>()?;
    }
    if let Ok(Some(v)) = config.get_item("accel_period") {
        cfg.accel_period = v.extract::<u32>()?;
    }
    if let Ok(Some(v)) = config.get_item("max_halvings") {
        cfg.max_halvings = v.extract::<u32>()?;
    }

    // Build code slices + g_per_fe Vec.
    let codes_slices: Vec<&[i64]> =
        code_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let g_per_fe_vec: Vec<usize> = g_view.iter().map(|&g| g as usize).collect();

    // Snapshot x into an owned Vec (the Rust IRLS doesn't mutate the
    // input but needs slice access; the workspace handles π internally).
    let x_owned = x.as_slice()?.to_vec();

    let result = py.allow_threads(|| {
        let mut ws = irls::FePoisIRLSWorkspace::new(
            n, p, &codes_slices, &g_per_fe_vec,
        );
        irls::fepois_loop(y_view, &x_owned, obs_w_view, &cfg, &mut ws)
    });

    let out = PyDict::new_bound(py);
    out.set_item("beta", PyArray1::from_vec_bound(py, result.beta))?;
    // x_tilde is F-order flat Vec (length n*p). Return it as a 1-D array
    // together with shape scalars so the Python caller can reconstruct:
    //   x_tilde = result["x_tilde_flat"].reshape(p, n).T
    out.set_item("x_tilde_flat", PyArray1::from_vec_bound(py, result.x_tilde))?;
    out.set_item("x_tilde_n", n)?;
    out.set_item("x_tilde_p", p)?;
    out.set_item("w", PyArray1::from_vec_bound(py, result.w))?;
    out.set_item("eta", PyArray1::from_vec_bound(py, result.eta))?;
    out.set_item("mu", PyArray1::from_vec_bound(py, result.mu))?;
    out.set_item("deviance", result.deviance)?;
    out.set_item("log_likelihood", result.log_likelihood)?;
    out.set_item("iters", result.iters)?;
    out.set_item("converged", result.converged)?;
    out.set_item("n_halvings", result.n_halvings)?;
    out.set_item("max_inner_dx", result.max_inner_dx)?;
    Ok(out)
}

/// Iterative K-way singleton detection. Returns a uint8 keep-mask
/// (1 = keep, 0 = drop).
#[pyfunction]
fn singleton_mask<'py>(
    py: Python<'py>,
    fe_codes: &Bound<'py, PyList>,
) -> PyResult<Bound<'py, PyArray1<u8>>> {
    let code_views = py_list_to_i64_views(fe_codes)?;
    if code_views.is_empty() {
        return Ok(PyArray1::<u8>::zeros_bound(py, 0, false));
    }
    let codes_slices: Vec<&[i64]> =
        code_views.iter().map(|v| v.as_slice().unwrap()).collect();

    let keep = py.allow_threads(|| singletons::detect_singletons(&codes_slices));
    let as_u8: Vec<u8> = keep.into_iter().map(|b| if b { 1 } else { 0 }).collect();
    Ok(PyArray1::from_vec_bound(py, as_u8))
}

/// Iterative Poisson-separation detection. Returns a uint8 keep-mask
/// (1 = keep, 0 = drop). Drops rows whose FE group has a zero y-sum;
/// iterates until fixed point. The Python wrapper at
/// ``statspai.fast.fepois._drop_separation_dispatcher`` falls back to
/// a pure-NumPy implementation when this entry point is missing.
#[pyfunction]
fn separation_mask<'py>(
    py: Python<'py>,
    y: PyReadonlyArray1<'py, f64>,
    fe_codes: &Bound<'py, PyList>,
    g_per_fe: PyReadonlyArray1<'py, i64>,
) -> PyResult<Bound<'py, PyArray1<u8>>> {
    let y_view = y.as_slice()?;
    let g_view = g_per_fe.as_slice()?;
    let code_views = py_list_to_i64_views(fe_codes)?;

    if code_views.is_empty() {
        // No FE → no separation possible; keep every row.
        let all_keep: Vec<u8> = vec![1u8; y_view.len()];
        return Ok(PyArray1::from_vec_bound(py, all_keep));
    }
    if g_view.len() != code_views.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "len(g_per_fe)={} but len(fe_codes)={}",
            g_view.len(),
            code_views.len()
        )));
    }
    let n = y_view.len();
    for v in &code_views {
        if v.as_slice()?.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "fe_codes entry has length {} but n={}",
                v.as_slice()?.len(),
                n
            )));
        }
    }

    let codes_slices: Vec<&[i64]> =
        code_views.iter().map(|v| v.as_slice().unwrap()).collect();
    let g_per_fe_vec: Vec<usize> = g_view.iter().map(|&g| g as usize).collect();

    let keep = py.allow_threads(|| {
        separation::separation_mask(y_view, &codes_slices, &g_per_fe_vec)
    });
    let as_u8: Vec<u8> = keep.into_iter().map(|b| if b { 1 } else { 0 }).collect();
    Ok(PyArray1::from_vec_bound(py, as_u8))
}

/// Cluster-robust sandwich "meat" matrix.
///
/// ``X`` and ``residuals`` must be **pre-sorted by cluster** (the
/// Python wrapper at ``statspai.core._numba_kernels.cluster_meat`` does
/// this with ``np.argsort(cluster_ids, kind='mergesort')``).
/// ``cluster_starts[g] .. cluster_ends[g]`` is the row range of cluster
/// ``g`` (half-open).
///
/// Parallelism: clusters are summed independently with Rayon over a
/// thread-local k×k upper-tri accumulator and reduced.
///
/// Parameters
/// ----------
/// x : 2-D float64 ndarray, shape (n, k), C-contiguous
///     Pre-sorted design matrix.
/// residuals : 1-D float64 ndarray, shape (n,)
///     Pre-sorted residuals.
/// cluster_starts : 1-D int64 ndarray, shape (G,)
/// cluster_ends   : 1-D int64 ndarray, shape (G,)
///     Half-open row ranges per cluster.
///
/// Returns
/// -------
/// (k, k) float64 ndarray (fully symmetric).
#[pyfunction]
fn cluster_meat<'py>(
    py: Python<'py>,
    x: PyReadonlyArray2<'py, f64>,
    residuals: PyReadonlyArray1<'py, f64>,
    cluster_starts: PyReadonlyArray1<'py, i64>,
    cluster_ends: PyReadonlyArray1<'py, i64>,
) -> PyResult<Bound<'py, PyArray2<f64>>> {
    if !x.is_c_contiguous() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "x must be C-contiguous; pass np.ascontiguousarray(X)",
        ));
    }
    let arr = x.as_array();
    let shape = arr.shape();
    if shape.len() != 2 {
        return Err(pyo3::exceptions::PyValueError::new_err("x must be 2-D"));
    }
    let n = shape[0];
    let k = shape[1];

    let r = residuals.as_slice()?;
    if r.len() != n {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "len(residuals)={} but n={}",
            r.len(),
            n
        )));
    }
    let starts_i = cluster_starts.as_slice()?;
    let ends_i = cluster_ends.as_slice()?;
    if starts_i.len() != ends_i.len() {
        return Err(pyo3::exceptions::PyValueError::new_err(format!(
            "len(cluster_starts)={} but len(cluster_ends)={}",
            starts_i.len(),
            ends_i.len()
        )));
    }
    // Validate cluster bounds before going into ``allow_threads``: any
    // out-of-range index would otherwise panic inside the parallel
    // section, which is harder to surface as a Python exception.
    for (g, (&s, &e)) in starts_i.iter().zip(ends_i.iter()).enumerate() {
        if s < 0 || e < 0 {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "cluster {g}: negative start/end ({s},{e})"
            )));
        }
        let (su, eu) = (s as usize, e as usize);
        if eu < su || eu > n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "cluster {g}: invalid range [{su},{eu}) with n={n}"
            )));
        }
    }
    let starts: Vec<usize> = starts_i.iter().map(|&v| v as usize).collect();
    let ends: Vec<usize> = ends_i.iter().map(|&v| v as usize).collect();

    let x_slice = arr.as_slice().ok_or_else(|| {
        pyo3::exceptions::PyValueError::new_err(
            "internal: x.as_slice() failed despite C-contiguous check",
        )
    })?;

    let meat_flat = py.allow_threads(|| {
        cluster::cluster_meat_sorted(x_slice, n, k, r, &starts, &ends)
    });

    let arr2 = Array2::from_shape_vec((k, k), meat_flat).map_err(|e| {
        pyo3::exceptions::PyValueError::new_err(format!("from_shape_vec: {e}"))
    })?;
    Ok(PyArray2::from_owned_array_bound(py, arr2))
}

/// Python module definition.
#[pymodule]
fn statspai_hdfe(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(group_demean, m)?)?;
    m.add_function(wrap_pyfunction!(demean_2d, m)?)?;
    m.add_function(wrap_pyfunction!(demean_2d_weighted, m)?)?;
    m.add_function(wrap_pyfunction!(demean_2d_weighted_sorted, m)?)?;
    m.add_function(wrap_pyfunction!(fepois_irls, m)?)?;
    m.add_function(wrap_pyfunction!(singleton_mask, m)?)?;
    m.add_function(wrap_pyfunction!(separation_mask, m)?)?;
    m.add_function(wrap_pyfunction!(cluster_meat, m)?)?;  // Phase 2
    m.add("__version__", "0.7.0-alpha.1")?;               // BUMPED 0.6.0 → 0.7.0-alpha.1
    Ok(())
}
