libqe coverage - 65.16%

Files
Source

/** @file accumulation.hpp
 *  @brief Accumulation primitives shared between the rENA stanza-window model
 *         and the tma ground/response/tensor model.
 *
 *  Both models reduce to the same core connection_matrix() operation; they
 *  differ in how ground and response vectors are assembled from the raw data.
 */
#ifndef LIBQE_ACCUMULATION_HPP
#define LIBQE_ACCUMULATION_HPP

#include <armadillo>
#include <functional>
#include <vector>
#include "adjacency.hpp"

namespace qe {

/// @name Core adjacency math
/// @{

/** @brief Compute the connection matrix for one ground + response event pair.
 *
 *  @param[in] ground          Row vector of accumulated ground codes (length p).
 *  @param[in] response        Row vector of the focal (response) codes (length p).
 *  @param[in] response_weight Scalar weight applied to the response self-connection
 *                             term.  Defaults to 1.0.
 *  @param[in] ordered         When @c true, produces a directed (asymmetric) matrix:
 *                             @c ground→response cross-product plus
 *                             @c 0.5 * response_weight * response self-connection
 *                             (diagonal zeroed).
 *                             When @c false, produces a symmetric undirected matrix:
 *                             @c (response_weight * r⊗r) + g⊗r + r⊗g.
 *
 *  @returns A p × p connection matrix.
 *
 *  @note Equivalent to @c calculate_adjacency_matrix() in @c tma/code.cpp.
 */
inline arma::mat connection_matrix(
    arma::rowvec ground, arma::rowvec response,
    double response_weight = 1.0, bool ordered = true
) {
    arma::mat resp_self = response.t() * response;
    arma::mat g_by_r   = ground.t() * response;

    if (ordered) {
        resp_self.diag().zeros();
        return g_by_r + (0.5 * response_weight * resp_self);
    }
    return (response_weight * resp_self) + g_by_r + (response.t() * ground);
}

/// @}

/// @name Traditional stanza-window accumulation (rENA model)
/// @{

/** @brief Accumulate connection vectors for every row in one conversation using
 *         a stanza window.
 *
 *  For each focal row @c k in @p codes the function assembles a ground context
 *  from the surrounding window and calls connection_matrix().
 *
 *  @par Ordering semantics
 *  - @b Undirected (@p ordered = @c false, default — rENA stanza model): the
 *    window spans [@c k - window_back, @c k + window_forward]; the upper-triangle
 *    outer-product (code_connections) is computed and back/forward-reference
 *    corrections are subtracted.  Returns a matrix of shape
 *    @c n_rows × choose_two(n_codes).
 *  - @b Directed (@p ordered = @c true): focal row @c k is the response; the
 *    sum of prior rows [@c earliest, @c k−1] is the ground.
 *    connection_matrix() is called with @p ordered = @c true and the result is
 *    vectorised.  @p window_forward is ignored (future rows cannot be causal
 *    ground context).  Returns a matrix of shape @c n_rows × n_codes².
 *
 *  @param[in] codes          Code matrix for one conversation (n_rows × n_codes).
 *  @param[in] window_back    Number of prior rows included in the window
 *                            (1 = current row only; INT_MAX = all prior rows).
 *  @param[in] window_forward Number of future rows included in the window
 *                            (ignored when @p ordered is @c true).
 *  @param[in] binary         When @c true, clamp all positive entries to 1.
 *  @param[in] ordered        When @c true, use directed accumulation; when
 *                            @c false, use undirected stanza-window accumulation.
 *
 *  @returns Connection matrix (n_rows × connection_vector_length) where
 *           connection_vector_length is choose_two(n_codes) for undirected or
 *           n_codes² for directed.
 *
 *  @note Equivalent to @c ref_window_df() in @c rENA/ena.cpp (undirected case).
 */
inline arma::mat accumulate_stanza(
    arma::mat codes,
    int window_back    = 1,
    int window_forward = 0,
    bool binary        = true,
    bool ordered       = false
) {
    int n_rows  = codes.n_rows;
    int n_codes = codes.n_cols;
    const int INT_MAX_VAL = std::numeric_limits<int>::max();

    // Shared helper: earliest row index for focal row k
    auto get_earliest = [&](int row) -> int {
        if (window_back == INT_MAX_VAL || window_back == 0)
            return (window_back == 0) ? row : 0;
        return std::max(0, row - (window_back - 1));
    };

    if (ordered) {
        arma::mat out(n_rows, n_codes * n_codes, arma::fill::zeros);
        for (int row = 0; row < n_rows; row++) {
            int earliest = get_earliest(row);
            arma::rowvec response = codes.row(row);
            arma::rowvec ground(n_codes, arma::fill::zeros);
            if (row > earliest)
                ground = arma::sum(codes.rows(earliest, row - 1));
            out.row(row) = arma::vectorise(
                connection_matrix(ground, response, 1.0, true)).t();
        }
        if (binary) out.elem(arma::find(out > 0)).ones();
        return out;
    }

    // Undirected: existing rENA stanza-window logic
    int n_tri = choose_two(n_codes);
    arma::mat out(n_rows, n_tri, arma::fill::zeros);

    for (int row = 0; row < n_rows; row++) {
        int earliest = get_earliest(row);

        int latest = row;
        if (window_forward == INT_MAX_VAL || row + window_forward >= n_rows) {
            latest = n_rows - 1;
        } else if (window_forward > 0) {
            latest = std::min(n_rows - 1, row + window_forward);
        }

        arma::mat window_rows = codes.rows(earliest, latest);
        arma::mat summed      = arma::sum(window_rows);
        arma::rowvec to_ut    = code_connections(summed);

        // Back-reference correction
        int win_rows  = latest - earliest + 1;
        if (win_rows > 0 && window_back > 1 && row - 1 >= 0) {
            int head_rows = win_rows - 1 - window_forward;
            if (head_rows > 0) {
                arma::mat refs    = window_rows.head_rows(head_rows);
                arma::mat ref_sum = arma::sum(refs);
                to_ut -= code_connections(ref_sum);
            }
        }

        // Forward-reference correction
        if (window_forward > 0 && latest <= n_rows - 1) {
            int tail_rows = latest - row;
            if (tail_rows > 0) {
                arma::mat refs    = window_rows.tail_rows(tail_rows);
                arma::mat ref_sum = arma::sum(refs);
                to_ut -= code_connections(ref_sum);
            }
        }

        out.row(row) = to_ut;
    }

    if (binary) out.elem(arma::find(out > 0)).ones();
    return out;
}

/// @}

/// @name Ground/response accumulation (tma model)
/// @{

/** @brief Accumulate connections for a single unit from its context rows.
 *
 *  For every response row belonging to this unit, the function computes a
 *  weighted ground vector from all preceding context rows and calls
 *  connection_matrix().
 *
 *  @par Ordering semantics
 *  When @p ordered is @c false (default), the result is folded into the
 *  upper-triangle representation (length choose_two(p)).  When @p ordered is
 *  @c true, the full directed p² vector is returned.
 *
 *  @param[in] codes      Full context matrix visible to this unit (n × p).
 *  @param[in] unit_rows  0-based indices of the rows that belong to this unit.
 *  @param[in] decay_fn   Callback with signature
 *                        @c arma::vec(arma::vec distances) that maps a vector
 *                        of distances (0 = current row, 1 = one step back, …)
 *                        to a vector of scalar weights; default behaviour is a
 *                        simple rectangular window.
 *  @param[in] ordered    When @c true, return a directed flat vector (length p²);
 *                        when @c false, return an undirected upper-triangle flat
 *                        vector (length choose_two(p)).
 *
 *  @returns Flat connection vector of length choose_two(p) (undirected) or p²
 *           (directed).
 *
 *  @note Equivalent to @c accumulate_network() in @c tma/code.cpp, with the R
 *        function callback replaced by @c std::function.
 */
inline arma::rowvec accumulate_unit(
    const arma::mat& codes,
    const std::vector<int>& unit_rows,
    std::function<arma::vec(arma::vec)> decay_fn,
    bool ordered = false
) {
    int code_cnt = codes.n_cols;
    arma::mat g_w_vec(code_cnt, code_cnt, arma::fill::zeros);

    for (int unit_row : unit_rows) {
        if (unit_row == 0) continue;

        arma::rowvec response = codes.row(unit_row);

        // Ground rows: everything from 0 up to (and including) unit_row
        arma::uvec ground_indices = arma::regspace<arma::uvec>(0, 1, unit_row);
        arma::vec  distances(ground_indices.n_elem);
        for (arma::uword k = 0; k < ground_indices.n_elem; k++)
            distances[k] = static_cast<double>(unit_row - ground_indices[k]);

        arma::vec decay_weights = decay_fn(distances);

        arma::mat ground_codes   = codes.rows(ground_indices);
        arma::mat weighted       = ground_codes.each_col() % decay_weights;
        arma::rowvec g_summed    = arma::sum(weighted);

        // Exclude the response row's own contribution from the ground
        arma::rowvec g_no_resp   = g_summed - weighted.tail_rows(1).row(0);

        g_w_vec += connection_matrix(g_no_resp, response, 1.0, ordered);
    }

    if (!ordered) return fold_directed_network(arma::vectorise(g_w_vec));
    return arma::vectorise(g_w_vec).t();
}

/// @}

/// @name Extended ground/response accumulation — returns per-row connection data
/// @{

/** @brief Result of accumulate_unit_with_rows().
 *
 *  Bundles the unit-level flat connection vector together with the per-response-row
 *  full p² connection matrices needed by tma's @c accumulate_network().
 */
struct UnitNetworks {
    arma::rowvec networks;      ///< Flat connection vector (choose_two(p) or p²).
    arma::mat    row_networks;  ///< Per-response-row full p² matrix (n_unit_rows × p²).
};

/** @brief Like accumulate_unit() but also returns the per-response-row connection
 *         matrix needed by tma's @c accumulate_network().
 *
 *  @param[in] codes      Full context matrix visible to this unit (n × p).
 *  @param[in] unit_rows  0-based indices of the rows that belong to this unit.
 *  @param[in] decay_fn   Callback with signature
 *                        @c arma::vec(int unit_row, arma::uvec ground_indices)
 *                        that returns a weight vector of length
 *                        @c ground_indices.n_elem.  The two-argument form lets
 *                        callers (e.g. the tma Rcpp wrapper) set R environment
 *                        variables before invoking the actual R decay function,
 *                        without any R-specific code leaking into libqe.
 *  @param[in] ordered    When @c true, return directed flat vectors (length p²);
 *                        when @c false, return undirected upper-triangle flat
 *                        vectors (length choose_two(p)).
 *
 *  @returns A UnitNetworks struct containing the aggregated connection vector
 *           and the per-response-row connection matrices.
 *
 *  @note Equivalent to @c accumulate_network() in @c tma/code.cpp (extended form).
 */
inline UnitNetworks accumulate_unit_with_rows(
    const arma::mat&                              codes,
    const std::vector<int>&                       unit_rows,
    std::function<arma::vec(int, arma::uvec)>     decay_fn,
    bool ordered = false
) {
    int code_cnt    = codes.n_cols;
    int n_unit_rows = unit_rows.size();

    arma::mat g_w_mat(code_cnt, code_cnt, arma::fill::zeros);
    arma::mat row_networks(n_unit_rows, code_cnt * code_cnt, arma::fill::zeros);

    for (int i = 0; i < n_unit_rows; ++i) {
        int unit_row = unit_rows[i];

        arma::uvec ground_indices = arma::regspace<arma::uvec>(0, 1, unit_row);
        arma::vec  decay_weights  = decay_fn(unit_row, ground_indices);

        arma::mat    ground_codes = codes.rows(ground_indices);
        arma::mat    weighted     = ground_codes.each_col() % decay_weights;
        arma::rowvec g_summed     = arma::sum(weighted);
        arma::rowvec g_no_resp    = g_summed - weighted.tail_rows(1).row(0);

        arma::rowvec response = codes.row(unit_row);
        arma::mat    conn     = connection_matrix(g_no_resp, response, 1.0, ordered);
        g_w_mat += conn;
        row_networks.row(i) = arma::vectorise(conn).t();
    }

    UnitNetworks result;
    if (!ordered)
        result.networks = fold_directed_network(arma::vectorise(g_w_mat));
    else
        result.networks = arma::vectorise(g_w_mat).t();
    result.row_networks = row_networks;
    return result;
}

/// @}

/// @name Tensor-based multi-modal accumulation (tma model)
/// @{

/** @brief Compute the linear index into a column-major multi-dimensional array.
 *
 *  @param[in] indices  Per-dimension index values (0-based).
 *  @param[in] dims     Size of each dimension.
 *
 *  @returns The scalar column-major flat index.
 *
 *  @note Equivalent to @c flat_index() in @c tma/code.cpp.
 */
inline int flat_index(const std::vector<int>& indices,
                               const std::vector<int>& dims) {
    if (indices.size() != dims.size())
        throw std::invalid_argument("Number of indices must match number of dimensions.");
    size_t linear = 0, stride = 1;
    for (size_t v = 0; v < indices.size(); ++v) {
        linear += static_cast<size_t>(indices[v]) * stride;
        stride *= static_cast<size_t>(dims[v]);
    }
    return static_cast<int>(linear);
}

/** @brief Result of apply_tensor_unit().
 *
 *  Bundles the unit-level aggregated connection vector together with the
 *  per-response-row connection matrices used by the tma tensor accumulation.
 */
struct TensorNetworks {
    arma::rowvec connection_counts;      ///< Unit-level flat vector (p²).
    arma::mat    row_connection_counts;  ///< Per-response-row connection matrix (n_unit_rows × p²).
    /// Per-response-row in-window ground rows (0-based context indices), visited
    /// order (ri-1 .. 0). Populated only when apply_tensor_unit() is called with
    /// return_members = true; empty otherwise. Feeds the webtool Data View
    /// window-span hover (tma design dataview-flexible-window.md).
    std::vector<std::vector<int>>    row_window_members;
    /// Per-response-row resolved window size for each member above (parallel to
    /// row_window_members). Populated only when return_members = true.
    std::vector<std::vector<double>> row_window_wins;
};

/** @brief Pure-C++ port of the tma @c apply_tensor() inner logic.
 *
 *  For each response row in @p unit_rows, the function looks up window and
 *  weight values from the tensor, collects all ground rows that fall within the
 *  response row's window, and calls connection_matrix() to accumulate the
 *  connection counts.
 *
 *  @param[in] tensor          Flat column-major array containing window and
 *                             weight values indexed by the factor dimensions.
 *  @param[in] dims            Sizes of each dimension of @p tensor.
 *  @param[in] dims_sender     Tensor axis indices corresponding to sender factors.
 *  @param[in] dims_receiver   Tensor axis indices corresponding to receiver
 *                             factors; these axes are overridden to the response
 *                             row's values when looking up window sizes for
 *                             ground rows.
 *  @param[in] dims_mode       Tensor axis indices corresponding to mode factors.
 *  @param[in] context_lookup  Integer matrix (n_context_rows × n_factors, 0-based)
 *                             mapping each context row to its factor level indices.
 *  @param[in] unit_rows       0-based response-row indices for this unit.
 *  @param[in] codes           Full context code matrix (n_context_rows × n_codes).
 *  @param[in] times           Timestamp associated with each context row
 *                             (length n_context_rows).
 *  @param[in] ordered         When @c true, produce a directed full p² result;
 *                             when @c false, produce an undirected upper-triangle
 *                             result.
 *
 *  @returns A TensorNetworks struct containing the unit-level connection counts
 *           and the per-response-row connection counts.
 *
 *  @note Equivalent to the inner loop of @c apply_tensor() in @c tma/code.cpp.
 */
inline TensorNetworks apply_tensor_unit(
    const arma::vec&        tensor,
    const std::vector<int>& dims,
    const std::vector<int>& dims_sender,
    const std::vector<int>& dims_receiver,
    const std::vector<int>& dims_mode,
    const arma::imat&       context_lookup,
    const std::vector<int>& unit_rows,
    const arma::mat&        codes,
    const arma::vec&        times,
    bool ordered = true,
    bool return_members = false
) {
    const int  WINDOW_DIM  = 1;
    const int  WEIGHT_DIM  = 0;
    const bool IS_DEFAULT  = (dims.size() == 1 && dims[0] == 2);

    int code_cnt    = codes.n_cols;
    int n_unit_rows = static_cast<int>(unit_rows.size());
    int ctx_cols    = static_cast<int>(context_lookup.n_cols);

    arma::mat g_w_mat(code_cnt, code_cnt, arma::fill::zeros);
    arma::mat row_conn(n_unit_rows, code_cnt * code_cnt, arma::fill::zeros);

    std::vector<std::vector<int>>    members_all;
    std::vector<std::vector<double>> wins_all;
    if (return_members) { members_all.reserve(n_unit_rows); wins_all.reserve(n_unit_rows); }

    int response_win    = 0;
    int response_weight = 0;
    if (IS_DEFAULT) {
        response_win    = static_cast<int>(tensor[1]);
        response_weight = static_cast<int>(tensor[0]);
    }

    for (int i = 0; i < n_unit_rows; ++i) {
        int    ri            = unit_rows[i];
        double response_time = times[ri];

        std::vector<int> resp_ctx(ctx_cols + 1);
        for (int j = 0; j < ctx_cols; ++j) resp_ctx[j] = context_lookup(ri, j);
        resp_ctx[ctx_cols] = WINDOW_DIM;

        if (!IS_DEFAULT)
            response_win = static_cast<int>(tensor[flat_index(resp_ctx, dims)]);

        std::vector<int>    gri_v;
        std::vector<double> grw_v;
        std::vector<double> grwin_v;
        arma::rowvec g_ws(code_cnt, arma::fill::zeros);

        if (ri > 0) {
            for (int gr = ri - 1; gr >= 0; --gr) {
                std::vector<int> row_v(ctx_cols + 1);
                for (int j = 0; j < ctx_cols; ++j) row_v[j] = context_lookup(gr, j);
                row_v[ctx_cols] = WINDOW_DIM;
                for (int dim : dims_receiver) row_v[dim] = resp_ctx[dim];

                double row_win = static_cast<double>(response_win);
                if (!IS_DEFAULT)
                    row_win = tensor[flat_index(row_v, dims)];

                if (times[gr] + row_win > response_time) {
                    gri_v.push_back(gr);
                    if (return_members) grwin_v.push_back(row_win);
                    row_v[ctx_cols] = WEIGHT_DIM;
                    double row_wgt = static_cast<double>(response_weight);
                    if (!IS_DEFAULT)
                        row_wgt = tensor[flat_index(row_v, dims)];
                    grw_v.push_back(row_wgt);
                }
            }

            if (!gri_v.empty()) {
                arma::uvec gri_u(gri_v.size());
                for (size_t k = 0; k < gri_v.size(); ++k) gri_u[k] = gri_v[k];
                arma::mat    gc  = codes.rows(gri_u);
                arma::colvec wts = arma::conv_to<arma::colvec>::from(grw_v);
                g_ws = arma::sum(gc.each_col() % wts, 0);
            }
        }

        if (!IS_DEFAULT) {
            resp_ctx[ctx_cols] = WEIGHT_DIM;
            response_weight = static_cast<int>(tensor[flat_index(resp_ctx, dims)]);
        }

        arma::rowvec row_vec = codes.row(ri);
        arma::mat resp = connection_matrix(
            g_ws, row_vec, static_cast<double>(response_weight), ordered);
        g_w_mat += resp;
        row_conn.row(i) = arma::vectorise(resp).t();

        if (return_members) {
            members_all.push_back(std::move(gri_v));
            wins_all.push_back(std::move(grwin_v));
        }
    }

    TensorNetworks result;
    result.connection_counts     = arma::vectorise(g_w_mat).t();
    result.row_connection_counts = row_conn;
    if (return_members) {
        result.row_window_members = std::move(members_all);
        result.row_window_wins    = std::move(wins_all);
    }
    return result;
}

/** @brief Aggregate per-response-row tensor connections into a unit vector.
 *
 *  Mirrors tma's R aggregation of @c apply_tensor_unit()'s
 *  @c row_connection_counts (tma/R/accum_multidim_c.R): the @c ordered branch is
 *  a plain column sum over the directed p² space, while the unordered branch
 *  folds each response row to the upper triangle (@c as.unordered) and, when
 *  @p binary is @c true, clamps each folded row to presence before summing
 *  (@c colSums.ena.matrix(binary = TRUE)).
 *
 *  This is the aggregation step tma performs in R; libqe's @c apply_tensor_unit
 *  intentionally returns the raw per-row matrix so callers can compose it.
 *  Kept as a standalone kernel function so the fold/binarize/sum semantics live
 *  in one place shared by every binding, instead of being re-implemented in each
 *  wrapper layer.
 *
 *  @param[in] row_conn  Per-response-row directed connection matrix
 *                       (n_response_rows × p²), i.e.
 *                       @c TensorNetworks::row_connection_counts.
 *  @param[in] n_codes   Number of codes @c p.
 *  @param[in] ordered   When @c true, return the directed p² column sums (no fold,
 *                       no binarization).  When @c false, fold each row to
 *                       @c choose_two(p) and sum.
 *  @param[in] binary    Unordered only: when @c true, binarize each folded
 *                       response row (presence, 0/1) before summing — the rENA /
 *                       tma co-occurrence convention.
 *
 *  @returns Flat unit connection vector of length p² (ordered) or
 *           @c choose_two(p) (unordered).
 */
inline arma::rowvec aggregate_row_connections(
    const arma::mat& row_conn,
    int  n_codes,
    bool ordered = false,
    bool binary  = true
) {
    if (ordered) return arma::sum(row_conn, 0);

    int n_rows = static_cast<int>(row_conn.n_rows);
    int n_tri  = choose_two(n_codes);
    arma::rowvec out(n_tri, arma::fill::zeros);

    for (int r = 0; r < n_rows; ++r) {
        arma::rowvec folded =
            fold_directed_network(arma::vectorise(row_conn.row(r)));
        if (binary) folded.elem(arma::find(folded > 0)).ones();
        out += folded;
    }
    return out;
}

/// @}

/// @name Per-row co-occurrence and rolling window (rENA accumulation primitives)
/// @{

/** @brief Compute the upper-triangle co-occurrence vector for each row.
 *
 *  For each row, calls code_connections() and optionally binarizes the result.
 *
 *  @param[in] codes   Code matrix (n_rows × n_codes).
 *  @param[in] binary  When @c true, clamp all positive entries to 1.
 *
 *  @returns Matrix of shape n_rows × choose_two(n_codes).
 *
 *  @note Equivalent to @c rows_to_co_occurrences() in @c rENA/ena.cpp.
 */
inline arma::mat row_connections(
    const arma::mat& codes,
    bool binary = true
) {
    int n_rows = codes.n_rows;
    int n_tri  = choose_two(codes.n_cols);
    arma::mat out(n_rows, n_tri, arma::fill::zeros);
    for (int row = 0; row < n_rows; ++row)
        out.row(row) = code_connections(codes.row(row));
    if (binary) out.elem(arma::find(out > 0)).ones();
    return out;
}

/** @brief Rolling backward window sum of the raw code matrix.
 *
 *  For each focal row @c k, sums rows [@c max(0, k - window_size + 1), @c k].
 *  Returns a matrix of the same shape as @p codes — no upper-triangle transform
 *  is applied.  A @p window_size of 0 or less is treated as 1 (current row only).
 *
 *  @param[in] codes        Code matrix (n_rows × n_codes).
 *  @param[in] window_size  Number of rows to include in the backward window
 *                          (including the focal row).  Values <= 0 are clamped
 *                          to 1.
 *
 *  @returns Matrix of the same shape as @p codes containing the windowed sums.
 *
 *  @note Equivalent to @c ref_window_lag() in @c rENA/ena.cpp.
 */
inline arma::mat rolling_window_sum(
    const arma::mat& codes,
    int window_size = 1
) {
    int n_rows = codes.n_rows;
    arma::mat out(n_rows, codes.n_cols, arma::fill::zeros);
    for (int row = 0; row < n_rows; ++row) {
        int earliest = (window_size > 0)
            ? std::max(0, row - (window_size - 1))
            : row;  // guard: window_size <= 0 -> single row
        out.row(row) = arma::sum(codes.rows(earliest, row));
    }
    return out;
}

/// @}

} // namespace qe

#endif // LIBQE_ACCUMULATION_HPP

/**
 * @file adjacency.hpp
 * @brief Combinatorics and vector/matrix utilities for ENA connection networks.
 *
 * Pure C++ / Armadillo — no Rcpp dependency.
 * Include @c <RcppArmadillo.h> before this header when building inside an R package.
 */
#ifndef LIBQE_ADJACENCY_HPP
#define LIBQE_ADJACENCY_HPP

#include <armadillo>
#include <string>
#include <vector>
#include <cmath>

namespace qe {

// ---------------------------------------------------------------------------
/// @name Combinatorics
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Number of unordered pairs from @p n items (n choose 2).
 *
 * @param n Number of codes/nodes.
 * @returns @c (n * (n - 1)) / 2.
 */
inline int choose_two(int n) {
    return (n * (n - 1)) / 2;
}

/**
 * @brief Upper-triangle index pairs for a symmetric matrix of side @p len.
 *
 * @param len Side length of the square matrix (number of codes).
 * @param row Controls what is returned:
 *   - @c -1 (default): 2 × k matrix of [row_idx; col_idx].
 *   - @c  0: row indices only (1 × k).
 *   - @c  1: column indices only (1 × k).
 * @returns An @c arma::umat of shape 2×k or 1×k depending on @p row.
 * @note Equivalent to @c triIndices() in rENA/ena.cpp and tma/code.cpp.
 */
inline arma::umat connection_indices(int len, int row = -1) {
    int vS = choose_two(len);
    int s  = 0;

    arma::umat vR(2, vS, arma::fill::zeros);
    arma::umat vRone(1, vS, arma::fill::zeros);

    for (int i = 2; i <= len; i++) {
        for (int j = 0; j < i - 1; j++) {
            vR(0, s) = j;
            vR(1, s) = i - 1;
            if (row == 0)      vRone[s] = j;
            else if (row == 1) vRone[s] = i - 1;
            s++;
        }
    }
    return (row == -1) ? vR : vRone;
}

/// @}

// ---------------------------------------------------------------------------
/// @name Vector ↔ upper-triangle conversions
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Compute pairwise products and return as a flat upper-triangle vector.
 *
 * For each pair @c (j, i) with @c j < i, outputs @c v[j] * v[i].
 * The result has length @c choose_two(v.size()).
 *
 * @param v A 1 × p row vector (or p-element matrix).
 * @returns A row vector of length @c choose_two(p).
 * @note Equivalent to @c vector_to_ut() in rENA/ena.cpp.
 */
inline arma::rowvec code_connections(arma::mat v) {
    int vL = v.size();
    int vS = choose_two(vL);
    arma::rowvec out(vS, arma::fill::zeros);
    int s = 0;
    for (int i = 2; i <= vL; i++) {
        for (int j = 0; j < i - 1; j++) {
            out[s] = v[j] * v[i - 1];
            s++;
        }
    }
    return out;
}

/**
 * @brief Fold a directed n² vector into an undirected upper-triangle vector.
 *
 * Symmetric elements are summed: A→B + B→A for every pair.
 *
 * @param v Flat directed vector of length n², column-major (n = sqrt(v.size())).
 * @returns Upper-triangle row vector of length @c choose_two(n).
 * @note Equivalent to @c vector_to_summed_uppertri() in tma/code.cpp.
 */
inline arma::rowvec fold_directed_network(arma::vec v) {
    int n    = static_cast<int>(std::round(std::sqrt(static_cast<double>(v.size()))));
    int tris = choose_two(n);

    arma::mat m(v);
    m.reshape(n, n);
    m = m + m.t();

    arma::colvec flat     = m.as_col();
    arma::uvec   up_inds  = arma::trimatu_ind(arma::size(m), 1);
    return arma::conv_to<arma::rowvec>::from(flat.elem(up_inds));
}

/**
 * @brief Flatten an adjacency matrix to a vector.
 *
 * @param x Square adjacency matrix (n × n).
 * @param full If @c true (default), returns the full n² vector (directed).
 *             If @c false, returns the upper-triangle only (undirected).
 * @returns Row vector of length n² or @c choose_two(n).
 * @note Equivalent to @c adjacency_matrix_to_vector() in tma/code.cpp.
 */
inline arma::rowvec network_to_vector(arma::mat x, bool full = true) {
    if (full) return arma::vectorise(x).t();
    arma::mat combined  = arma::trimatu(x, 1) + arma::trimatl(x, -1);
    arma::uvec up_inds  = arma::trimatu_ind(arma::size(combined), 1);
    return arma::vectorise(combined(up_inds)).t();
}

/// @}

// ---------------------------------------------------------------------------
/// @name String utilities
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Generate "A & B" pair names for every upper-triangle position.
 *
 * For code names @c {"A", "B", "C"} returns @c {"A & B", "A & C", "B & C"}.
 *
 * @param v Ordered list of code names (length p).
 * @returns Vector of @c choose_two(p) label strings.
 * @note Equivalent to @c svector_to_ut() in rENA/ena.cpp.
 */
inline std::vector<std::string> connection_names(std::vector<std::string> v) {
    int vL = v.size();
    int vS = choose_two(vL);
    std::vector<std::string> out(vS);
    int s = 0;
    for (int i = 2; i <= vL; i++) {
        for (int j = 0; j < i - 1; j++) {
            out[s] = v[j] + " & " + v[i - 1];
            s++;
        }
    }
    return out;
}

/// @}

} // namespace qe

#endif // LIBQE_ADJACENCY_HPP

/** @file ccd.hpp
 *  @brief Cross-Covariance Decay (CCD) window-size estimation.
 *
 *  Estimates the ENA moving-window size from the "discourse coherence length":
 *  the half-life decay lag of the noise-corrected Frobenius norm of pooled
 *  conversation cross-covariance matrices.
 *
 *  Reference: Shaffer, D. W. & Cai, Z. (2026). Discourse Coherence Length:
 *  Noise-Corrected Covariance Estimation of Window-Size in Epistemic Network
 *  Analysis. International Conference on Quantitative Ethnography (ICQE26).
 *
 *  This kernel is pure-numeric: it consumes conversation subsets that have
 *  already been extracted from the raw data (one code matrix per conversation,
 *  rows in sequence).  Column selection and conversation splitting belong in
 *  each language wrapper (cf. accumulation.hpp), so that the C++ core never
 *  touches data-frame semantics.
 */
#ifndef LIBQE_CCD_HPP
#define LIBQE_CCD_HPP

#include <armadillo>
#include <cmath>
#include <vector>

namespace qe {

/** @brief Result of a CCD window-size estimation.
 *
 *  The `*` vectors are all length `max_window + 1`, indexed by lag `0..max_window`.
 *  Entries for lags with no contributing conversation are NaN (weight 0).
 */
struct CCDResult {
    int       window_size;            ///< Estimated half-life lag window size (>= 1).
    int       peak_lag;               ///< Lag of the peak corrected cross-covariance norm.
    arma::vec lag;                    ///< Evaluated lags, 0..max_window.
    arma::vec frob;                   ///< Uncorrected pooled Frobenius norm per lag.
    arma::vec frob_sq_unbiased;       ///< Noise-corrected squared Frobenius norm per lag.
    arma::vec frob_unbiased_signed;   ///< sign * sqrt(|corrected sq|) per lag.
    arma::vec total_weight;           ///< Pooled overlap weight per lag.
};

/** @brief Estimate the ENA moving-window size via cross-covariance decay.
 *
 *  For each lag @c l in @c 0..max_window, pools window-matched cross-covariance
 *  matrices across all conversation subsets, subtracts a trace-variance noise
 *  floor, and locates the first lag strictly after the peak where the corrected
 *  norm falls to half of its peak value.
 *
 *  @param[in] conversations One code matrix per conversation (N_k rows × C codes,
 *                           rows in sequence).  All matrices must share the same
 *                           number of columns.
 *  @param[in] max_window    Maximum lag to evaluate.  Default 20.
 *  @param[in] min_overlap   Minimum overlapping rows (N_k - l) required for a
 *                           conversation to contribute at lag l.  Default 10.
 *
 *  @returns A @ref CCDResult.  When no conversation contributes at any positive
 *           lag (all-NaN or non-positive corrected covariance), the window size
 *           defaults to 1 and @c peak_lag to 0.
 */
inline CCDResult ccd_window(
    const std::vector<arma::mat>& conversations,
    int max_window = 20,
    int min_overlap = 10
) {
    if (max_window < 0) max_window = 0;
    const int n_lags = max_window + 1;

    CCDResult out;
    out.lag                  = arma::regspace<arma::vec>(0, max_window);
    out.frob                 = arma::vec(n_lags).fill(arma::datum::nan);
    out.frob_sq_unbiased     = arma::vec(n_lags).fill(arma::datum::nan);
    out.frob_unbiased_signed = arma::vec(n_lags).fill(arma::datum::nan);
    out.total_weight         = arma::zeros<arma::vec>(n_lags);
    out.window_size          = 1;
    out.peak_lag             = 0;

    // Determine code count from the first non-empty conversation.
    arma::uword C = 0;
    for (const arma::mat& X : conversations) {
        if (X.n_cols > 0) { C = X.n_cols; break; }
    }
    if (C == 0) return out;  // nothing to work with -> window 1

    // 1. Cross-covariance curves & noise floor across lags.
    for (int i = 0; i < n_lags; ++i) {
        const int l = i;
        arma::mat Cov_sum = arma::zeros<arma::mat>(C, C);
        double total_weight = 0.0;
        double noise_floor_sq_weighted_sum = 0.0;

        for (const arma::mat& X_k : conversations) {
            const int N_k = static_cast<int>(X_k.n_rows);
            const int weight = N_k - l;
            if (weight < min_overlap) continue;

            // A = rows [l, N_k-1], B = rows [0, N_k-1-l]; both have `weight` rows.
            const arma::mat A = X_k.rows(l, N_k - 1);
            const arma::mat B = X_k.rows(0, N_k - 1 - l);

            const arma::rowvec p_A = arma::mean(A, 0);
            const arma::rowvec p_B = arma::mean(B, 0);

            // Window-matched cross-covariance: crossprod(B, A)/w - outer(p_B, p_A).
            const arma::mat Cov_k = (B.t() * A) / weight - p_B.t() * p_A;

            // Trace-variance noise floor (n-1 denominator, matching R's var()).
            const double tr_var_A =
                arma::accu(arma::square(A.each_row() - p_A)) / (weight - 1);
            const double tr_var_B =
                arma::accu(arma::square(B.each_row() - p_B)) / (weight - 1);
            const double var_k_sq = tr_var_A * tr_var_B;

            Cov_sum += weight * Cov_k;
            total_weight += weight;
            noise_floor_sq_weighted_sum += weight * var_k_sq;
        }

        if (total_weight > 0.0) {
            const arma::mat Cov_pooled = Cov_sum / total_weight;
            const double frob_sq_pooled = arma::accu(arma::square(Cov_pooled));
            const double noise_floor_sq_pooled =
                noise_floor_sq_weighted_sum / (total_weight * total_weight);
            const double frob_sq_unbiased_pooled =
                frob_sq_pooled - noise_floor_sq_pooled;

            const double sgn = (frob_sq_unbiased_pooled > 0.0) ? 1.0
                             : (frob_sq_unbiased_pooled < 0.0) ? -1.0 : 0.0;

            out.frob[i]                 = std::sqrt(frob_sq_pooled);
            out.frob_sq_unbiased[i]     = frob_sq_unbiased_pooled;
            out.frob_unbiased_signed[i] = sgn * std::sqrt(std::abs(frob_sq_unbiased_pooled));
            out.total_weight[i]         = total_weight;
        }
    }

    // 2. Half-life decay lag detection over the signed corrected norm.
    const arma::vec& f = out.frob_unbiased_signed;

    // Peak among lags > 0 with finite, positive corrected norm.
    int    peak_lag = -1;
    double max_f    = -arma::datum::inf;
    for (int i = 1; i < n_lags; ++i) {
        if (std::isfinite(f[i]) && f[i] > 0.0 && f[i] > max_f) {
            max_f    = f[i];
            peak_lag = i;   // lag == index
        }
    }

    if (peak_lag < 0) {
        // No positive corrected covariance anywhere -> default window 1.
        out.window_size = 1;
        out.peak_lag    = 0;
        return out;
    }
    out.peak_lag = peak_lag;

    // First lag at or after the peak whose corrected norm has decayed to <= 50%.
    // NaN comparisons are false, so NaN lags are skipped (matching R's which()).
    int decay_lag = -1;
    for (int i = peak_lag; i < n_lags; ++i) {
        if (f[i] / max_f <= 0.5) { decay_lag = i; break; }
    }

    if (decay_lag >= 0) {
        out.window_size = decay_lag;
    } else {
        // No decay below half-life within max_window: fall back to the last lag
        // (at or after the peak) with a finite corrected norm.
        int last_valid = -1;
        for (int i = peak_lag; i < n_lags; ++i) {
            if (std::isfinite(f[i])) last_valid = i;
        }
        out.window_size = (last_valid >= 0) ? last_valid : 1;
    }

    return out;
}

} // namespace qe

#endif // LIBQE_CCD_HPP

/** @file door.hpp
 *  @brief Door temporal pooling and smoothing functions for trajectory modeling.
 *
 *  Provides sliding lookback window pooling (with uniform or linear weighting,
 *  sum or mean aggregation, and optional segment boundary resets) as well as
 *  Exponential Moving Average (EMA) smoothing over consecutive time steps.
 */
#ifndef LIBQE_DOOR_HPP
#define LIBQE_DOOR_HPP

#include <armadillo>
#include <vector>
#include <string>
#include <cmath>
#include <algorithm>

namespace qe {

/** @brief Apply a sliding lookback window pool over a single unit's connection matrix.
 *
 *  For each row @c i, pools rows in the window [@c max(0, i - lookback_size + 1), @c i].
 *  If @p segment_ids is non-empty, the window does not cross rows with a different segment ID.
 *
 *  @param[in] block            Matrix of connection counts for one unit (n_rows × n_features).
 *  @param[in] lookback_size    Number of rows in the lookback window (>= 1).
 *  @param[in] aggregate_mean   When @c true, averages rows; when @c false, sums rows.
 *  @param[in] weighting_linear When @c true, applies linearly increasing weights
 *                              (oldest row in window has weight 1, current row has weight k).
 *                              When @c false, all rows in window have equal weight 1.
 *  @param[in] segment_ids      Optional vector of segment/activity IDs (length n_rows).
 *                              Window pooling resets at segment transitions.
 *
 *  @returns A matrix of shape n_rows × n_features containing door-smoothed values.
 */
inline arma::mat lookback_block(
    const arma::mat& block,
    int lookback_size = 20,
    bool aggregate_mean = false,
    bool weighting_linear = false,
    const std::vector<int>& segment_ids = {}
) {
    int n_rows = static_cast<int>(block.n_rows);
    int n_cols = static_cast<int>(block.n_cols);
    arma::mat out(n_rows, n_cols, arma::fill::zeros);

    if (n_rows == 0 || n_cols == 0) return out;
    int k_max = std::max(1, lookback_size);
    bool has_segments = (!segment_ids.empty() && static_cast<int>(segment_ids.size()) >= n_rows);

    for (int i = 0; i < n_rows; ++i) {
        int start_idx = std::max(0, i - k_max + 1);

        if (has_segments) {
            int cur_seg = segment_ids[i];
            for (int s = i; s >= start_idx; --s) {
                if (segment_ids[s] != cur_seg) {
                    start_idx = s + 1;
                    break;
                }
            }
        }

        int win_len = i - start_idx + 1;
        if (win_len <= 0) continue;

        arma::rowvec acc(n_cols, arma::fill::zeros);
        arma::rowvec denom(n_cols, arma::fill::zeros);

        for (int j = start_idx; j <= i; ++j) {
            double w = 1.0;
            if (weighting_linear) {
                // Linear weight: 1 for earliest in window, win_len for current row i
                w = static_cast<double>(j - start_idx + 1);
            }
            for (int col = 0; col < n_cols; ++col) {
                double value = block(j, col);
                if (!std::isnan(value)) {
                    acc[col] += w * value;
                    denom[col] += w;
                }
            }
        }

        if (aggregate_mean) {
            for (int col = 0; col < n_cols; ++col) {
                acc[col] = (denom[col] > 0.0) ? (acc[col] / denom[col]) : arma::datum::nan;
            }
        }

        out.row(i) = acc;
    }

    return out;
}

/** @brief Apply Exponential Moving Average (EMA) smoothing over a single unit's connection matrix.
 *
 *  Computes S_0 = Y_0, S_t = alpha * Y_t + (1 - alpha) * S_{t-1}.
 *  If @p segment_ids is non-empty, S_t resets to Y_t when transitioning to a new segment.
 *
 *  @param[in] block       Matrix of connection counts for one unit (n_rows × n_features).
 *  @param[in] alpha       Smoothing parameter alpha in (0, 1].
 *  @param[in] segment_ids Optional vector of segment/activity IDs (length n_rows).
 *
 *  @returns A matrix of shape n_rows × n_features containing EMA-smoothed values.
 */
inline arma::mat ema_block(
    const arma::mat& block,
    double alpha = 0.1,
    const std::vector<int>& segment_ids = {}
) {
    int n_rows = static_cast<int>(block.n_rows);
    int n_cols = static_cast<int>(block.n_cols);
    arma::mat out(n_rows, n_cols, arma::fill::zeros);

    if (n_rows == 0 || n_cols == 0) return out;
    double a = std::max(1e-6, std::min(1.0, alpha));
    bool has_segments = (!segment_ids.empty() && static_cast<int>(segment_ids.size()) >= n_rows);

    out.row(0) = block.row(0);

    for (int i = 1; i < n_rows; ++i) {
        bool reset = has_segments && (segment_ids[i] != segment_ids[i - 1]);
        if (reset) {
            out.row(i) = block.row(i);
        } else {
            for (int col = 0; col < n_cols; ++col) {
                double cur = block(i, col);
                if (std::isnan(cur)) {
                    cur = out(i - 1, col);
                }
                out(i, col) = a * cur + (1.0 - a) * out(i - 1, col);
            }
        }
    }

    return out;
}

/** @brief Apply lookback door pooling to an entire connection matrix across multiple units.
 *
 *  @param[in] conn_counts        Complete connection matrix (all units).
 *  @param[in] unit_row_indices   Vector of row-index vectors, one per unit (0-indexed).
 *  @param[in] lookback_sizes     Vector of lookback sizes, one per unit.
 *  @param[in] aggregate_mean     When @c true, compute mean; when @c false, compute sum.
 *  @param[in] weighting_linear   When @c true, use linear weighting.
 *  @param[in] segment_ids        Optional vector of segment IDs for every row in conn_counts.
 *
 *  @returns Smoothed connection matrix with identical dimensions to @p conn_counts.
 */
inline arma::mat apply_door_lookback(
    const arma::mat& conn_counts,
    const std::vector<std::vector<int>>& unit_row_indices,
    const std::vector<int>& lookback_sizes,
    bool aggregate_mean = false,
    bool weighting_linear = false,
    const std::vector<int>& segment_ids = {}
) {
    arma::mat out = conn_counts;
    size_t n_units = unit_row_indices.size();
    bool has_segments = (!segment_ids.empty() && segment_ids.size() >= conn_counts.n_rows);

    for (size_t u = 0; u < n_units; ++u) {
        const auto& rows = unit_row_indices[u];
        if (rows.empty()) continue;

        int k = (u < lookback_sizes.size()) ? lookback_sizes[u] : 20;
        int n_u = static_cast<int>(rows.size());

        arma::mat block(n_u, conn_counts.n_cols);
        std::vector<int> block_segs(n_u, 0);

        for (int i = 0; i < n_u; ++i) {
            int r = rows[i];
            block.row(i) = conn_counts.row(r);
            if (has_segments) block_segs[i] = segment_ids[r];
        }

        arma::mat smoothed = lookback_block(block, k, aggregate_mean, weighting_linear,
                                            has_segments ? block_segs : std::vector<int>{});

        for (int i = 0; i < n_u; ++i) {
            out.row(rows[i]) = smoothed.row(i);
        }
    }

    return out;
}

/** @brief Apply EMA door smoothing to an entire connection matrix across multiple units.
 */
inline arma::mat apply_door_ema(
    const arma::mat& conn_counts,
    const std::vector<std::vector<int>>& unit_row_indices,
    double alpha = 0.1,
    const std::vector<int>& segment_ids = {}
) {
    arma::mat out = conn_counts;
    size_t n_units = unit_row_indices.size();
    bool has_segments = (!segment_ids.empty() && segment_ids.size() >= conn_counts.n_rows);

    for (size_t u = 0; u < n_units; ++u) {
        const auto& rows = unit_row_indices[u];
        if (rows.empty()) continue;

        int n_u = static_cast<int>(rows.size());
        arma::mat block(n_u, conn_counts.n_cols);
        std::vector<int> block_segs(n_u, 0);

        for (int i = 0; i < n_u; ++i) {
            int r = rows[i];
            block.row(i) = conn_counts.row(r);
            if (has_segments) block_segs[i] = segment_ids[r];
        }

        arma::mat smoothed = ema_block(block, alpha, has_segments ? block_segs : std::vector<int>{});

        for (int i = 0; i < n_u; ++i) {
            out.row(rows[i]) = smoothed.row(i);
        }
    }

    return out;
}

} // namespace qe

#endif // LIBQE_DOOR_HPP

/**
 * @file generalized_rotation.hpp
 * @brief Generalized Means Rotation (GMR) with Lasso-based covariate adjustment.
 *
 * @details
 * Mirrors rENA's `ena.rotate.by.generalized()` + `gmr()` + `get_x1_main_effect()`.
 * Authoritative reference: commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901
 * ("Fixed bugs in rotation functions") by Zhiqiang Cai.
 *
 * **Algorithm overview**
 *
 *  1. Compute x_vector: the rotation axis that best represents the target
 *     variable's contribution to the ENA point space, after controlling for
 *     covariates via Lasso (`lasso_x1_contribution`).
 *  2. Deflate V by x_vector → defA = V - V*x*x'.
 *  3. Optionally compute secondary axis x1: the normalised group-mean
 *     difference in the *original* V space (not defA) for groups 0 vs 1.
 *     x1 stays empty when the target is not categorical or there are not
 *     exactly 2 groups — there is no SVD fallback.
 *  4. Compute y_vector from a GMR call on defA, or from SVD of defA.
 *     Explicitly orthogonalise y_vector against x_vector, and against x1
 *     when x1 is available.
 *  5. Double-deflate V by x_vector and y_vector, then call `complete_rotation()`
 *     to fill remaining dimensions from the SVD of the doubly-deflated space.
 *
 * **Caller responsibilities**
 *
 *  - Build the model matrix (`X_raw`) from the data frame, including dummy
 *    variables for categorical predictors and any desired interaction terms.
 *    No formula parsing is done here.
 *  - Identify `x1_cols`: the 0-based column indices in `X_raw` that correspond to
 *    the target variable (these are unpenalized in the Lasso).
 *  - Encode categorical targets as 0-based integer codes in `x_target`.
 *  - Set `x_n_groups` to the number of distinct groups when `x_categorical=true`.
 *
 * The R wrapper (`rENA::ena.rotate.by.generalized`) handles these steps and
 * calls `libqe::generalized_means_rotation()` with the prepared matrices.
 */

#ifndef LIBQE_GENERALIZED_ROTATION_HPP
#define LIBQE_GENERALIZED_ROTATION_HPP

#include "linalg_fallback.hpp"  // qe::linalg::eig_sym (works w/o LAPACK)
#include "rotation.hpp"         // ena_svd, complete_rotation, RotationResult
#include "lasso.hpp"            // lasso_x1_contribution

#include <armadillo>
#include <stdexcept>
#include <string>
#include <vector>

namespace qe {

/// @name Between-group scatter
/// @{

/**
 * @brief Compute the between-group scatter matrix S_B.
 *
 * @details
 * Computes S_B = Σ_g  n_g * (μ_g − μ)(μ_g − μ)'  where μ is the grand mean.
 *
 * @note Mirrors rENA's `compute_SB()` (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
 *
 * @param[in] V        n×p matrix of ENA points (or covariate-adjusted contribution).
 * @param[in] labels   n×1 0-based integer group labels.
 * @param[in] n_groups Number of distinct groups.
 *
 * @returns p×p symmetric between-group scatter matrix.
 */
inline arma::mat between_group_scatter(
    const arma::mat&  V,
    const arma::uvec& labels,
    arma::uword       n_groups
) {
    const arma::uword p = V.n_cols;
    arma::rowvec mu = arma::mean(V, 0);
    arma::mat SB(p, p, arma::fill::zeros);

    for (arma::uword g = 0; g < n_groups; ++g) {
        arma::uvec idx = arma::find(labels == g);
        if (idx.is_empty()) continue;
        arma::vec diff = (arma::mean(V.rows(idx), 0) - mu).t();
        SB += static_cast<double>(idx.n_elem) * diff * diff.t();
    }
    return SB;
}

/// @}

/// @name Rotation direction
/// @{

/**
 * @brief Compute the unit rotation direction from a covariate-adjusted contribution matrix.
 *
 * @details
 * **Numeric target:**
 * Centres the target, then computes β = (t_c' t_c)^{-1} t_c' Vx (OLS slope from
 * lm(Vx ~ target)). Mirrors R: `model = lm(Vx ~ target); beta = coefficients[2,]`.
 *
 * **Categorical target:**
 * Computes the between-group scatter matrix S_B and returns its leading
 * eigenvector (the direction of maximum between-group variance).
 * Mirrors R: `sb = compute_SB(Vx, target); r = svd(sb)$v[,1]`.
 *
 * @note Mirrors rENA's `gmr()` direction computation
 *       (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
 *
 * @param[in] Vx          n×q covariate-adjusted contribution matrix.
 * @param[in] target      n×1 target variable (numeric float, or 0-based integer codes).
 * @param[in] categorical `true` if the target is categorical (uses between-group scatter);
 *                        `false` for a numeric target (uses OLS slope).
 * @param[in] n_groups    Number of distinct groups; only used when `categorical` is `true`.
 *
 * @returns Unit vector (p×1) representing the rotation direction.
 *
 * @throws std::runtime_error If `categorical` is `true` and `n_groups < 2`.
 * @throws std::runtime_error If the numeric target has zero variance after centering.
 * @warning Throws `std::runtime_error` if the resulting direction has zero norm,
 *          which can occur when the Lasso-adjusted contribution collapses to zero.
 */
inline arma::vec gmr_direction(
    const arma::mat& Vx,         // n×q covariate-adjusted contribution
    const arma::vec& target,     // n×1 (numeric float, or 0-based int codes)
    bool             categorical,
    arma::uword      n_groups = 0
) {
    arma::vec r;

    if (categorical) {
        if (n_groups < 2)
            throw std::runtime_error(
                "gmr_direction: x_n_groups must be >= 2 for categorical target");
        arma::uvec labels = arma::conv_to<arma::uvec>::from(target);
        arma::mat  SB     = between_group_scatter(Vx, labels, n_groups);

        arma::vec eigval; arma::mat eigvec;
        qe::linalg::eig_sym(eigval, eigvec, SB);  // ascending order
        r = eigvec.col(eigvec.n_cols - 1);         // largest eigenvalue → last col
    } else {
        // Centre target to get the OLS slope (mirrors R's lm() intercept handling)
        arma::vec tc   = target - arma::mean(target);
        double    tTt  = arma::dot(tc, tc);
        if (tTt < 1e-12)
            throw std::runtime_error(
                "gmr_direction: target has zero variance after centering");
        arma::rowvec beta = (tc.t() * Vx) / tTt;  // 1×q slope vector
        r = beta.t();
    }

    double norm = arma::norm(r);
    if (norm < 1e-12)
        throw std::runtime_error("gmr_direction: resulting direction has zero norm");
    return r / norm;
}

/// @}

/**
 * @brief Result of one GMR axis step.
 */
struct GMRStepResult {
    arma::vec direction;   ///< Unit vector in connection space (p×1).
};

/// @name GMR axis step
/// @{

/**
 * @brief Compute the rotation direction for one axis of the Generalized Means Rotation.
 *
 * @details
 * Applies optional row subsetting, performs Lasso-based covariate adjustment
 * (or a simple OLS projection when all columns are target columns), then
 * delegates to `gmr_direction()` to produce a unit direction vector.
 *
 * When `x1_cols.n_elem >= p` (all columns are target columns and there are no
 * penalized covariates), the function falls back to a simple OLS projection:
 * Vx = fitted values from lm(V ~ t).
 *
 * @note Mirrors rENA's `gmr()` step logic
 *       (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
 *
 * @param[in] V          n×q matrix of ENA points (full data; subset applied internally).
 * @param[in] X_model    n×p model matrix (caller-built; target cols + covariate cols).
 * @param[in] target     n×1 target variable (raw float or 0-based integer codes).
 * @param[in] x1_cols    0-based indices of target columns in `X_model` (unpenalized in Lasso).
 * @param[in] categorical `true` if the target is categorical.
 * @param[in] n_groups   Number of distinct groups (only used when `categorical` is `true`).
 * @param[in] subset     Optional 0-based row indices to use; empty means use all rows.
 * @param[in] n_lambda   Length of the Lasso lambda path (forwarded to `lasso_x1_contribution`).
 * @param[in] k_folds    Number of cross-validation folds (forwarded to `lasso_x1_contribution`).
 * @param[in] lasso_eps  lambda_min = lasso_eps * lambda_max (forwarded to `lasso_x1_contribution`).
 *
 * @returns `GMRStepResult` containing the unit direction vector.
 *
 * @throws std::runtime_error Propagated from `gmr_direction()` if the direction
 *         is degenerate or the target has zero variance.
 */
inline GMRStepResult gmr_step(
    const arma::mat&  V,
    const arma::mat&  X_model,
    const arma::vec&  target,
    const arma::uvec& x1_cols,
    bool              categorical,
    arma::uword       n_groups,
    const arma::uvec& subset,
    int               n_lambda  = 50,
    int               k_folds   = 5,
    double            lasso_eps = 0.01
) {
    const arma::uword n = V.n_rows;
    const arma::uword p = X_model.n_cols;

    // ── Select rows ──────────────────────────────────────────────────────────
    arma::uvec rows = subset.is_empty()
                    ? arma::regspace<arma::uvec>(0, n - 1)
                    : subset;

    const arma::mat V_sub  = V.rows(rows);
    const arma::mat X_sub  = X_model.rows(rows);
    const arma::vec t_sub  = target.rows(rows);

    // ── Lasso-adjusted contribution Vx ───────────────────────────────────────
    arma::mat Vx_sub;
    if (x1_cols.n_elem >= p) {
        // All columns are target columns (no penalized covariates):
        // fall back to simple OLS projection (Vx = fitted values from lm(V~t))
        arma::vec tc  = t_sub - arma::mean(t_sub);
        double    tTt = arma::dot(tc, tc);
        if (tTt > 1e-12) {
            arma::rowvec b1 = (tc.t() * V_sub) / tTt;
            Vx_sub = tc * b1;
        } else {
            Vx_sub.zeros(V_sub.n_rows, V_sub.n_cols);
        }
    } else {
        arma::vec pf(p, arma::fill::ones);
        for (arma::uword j : x1_cols) pf(j) = 0.0;

        Vx_sub = lasso_x1_contribution(
            X_sub, V_sub, x1_cols, pf, n_lambda, k_folds, lasso_eps);
    }

    // ── Rotation direction ───────────────────────────────────────────────────
    arma::vec direction = gmr_direction(Vx_sub, t_sub, categorical, n_groups);

    return {direction};
}

/// @}

/**
 * @brief Parameters for `generalized_means_rotation()`.
 *
 * @details
 * Groups all inputs needed to specify both the X and Y rotation axes, along
 * with Lasso tuning knobs. The Y axis is optional: when `has_y` is `false`
 * the second axis is taken from the leading SVD of the x-deflated space.
 */
struct GeneralizedRotationParams {
    /// @name X axis (required)
    /// @{
    arma::mat   x_model_matrix;           ///< n×p model matrix for the X axis.
    arma::vec   x_target;                 ///< n×1 target variable (raw float or 0-based integer codes).
    arma::uvec  x1_cols;                  ///< 0-based target column indices in `x_model_matrix`.
    bool        x_categorical = false;    ///< `true` if the X target is categorical.
    arma::uword x_n_groups    = 0;        ///< Number of distinct groups (used when `x_categorical` is `true`).
    arma::uvec  x_subset;                 ///< Optional row subset (0-based); empty means use all rows.
    /// @}

    /// @name Y axis (optional)
    /// @{
    /// When `has_y` is `false`, the Y axis is taken from the leading SVD of the x-deflated space.
    bool        has_y          = false;   ///< `true` if a Y-axis GMR target is provided.
    arma::mat   y_model_matrix;           ///< n×p model matrix for the Y axis.
    arma::vec   y_target;                 ///< n×1 target variable for the Y axis.
    arma::uvec  y1_cols;                  ///< 0-based target column indices in `y_model_matrix`.
    bool        y_categorical  = false;   ///< `true` if the Y target is categorical.
    arma::uword y_n_groups     = 0;       ///< Number of distinct groups for the Y axis.
    /// The Y axis always uses all rows (no subset parameter).
    /// @}

    /// @name Lasso tuning
    /// @{
    int    n_lambda  = 50;    ///< Length of the lambda path.
    int    k_folds   = 5;     ///< Number of cross-validation folds.
    double lasso_eps = 0.01;  ///< lambda_min = lasso_eps * lambda_max.
    /// @}
};

/// @name Generalized Means Rotation
/// @{

/**
 * @brief Full Generalized Means Rotation (GMR) producing a complete q×q rotation matrix.
 *
 * @details
 * C++/Armadillo implementation of `ena.rotate.by.generalized()`, matching the
 * authoritative R implementation in commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901
 * by Zhiqiang Cai.
 *
 * **Steps performed:**
 *
 *  1. **X axis via GMR** — calls `gmr_step()` on V with the X-axis parameters.
 *  2. **Deflate V** — computes defA = V − V·x·x' for subsequent Y-axis work.
 *  3. **Secondary axis x1** — when the X target is categorical with ≥ 2 groups,
 *     computes x1 = normalize(mean(V[group0]) − mean(V[group1])) using the
 *     *original* V (not defA) and the full-length target (all n rows).
 *     x1 is used only to orthogonalise y_vector; it does not deflate defA.
 *     There is no SVD fallback: x1 stays empty if the condition is not met.
 *  4. **Y axis** — either GMR on defA (when `params.has_y` is `true`) or the
 *     leading SVD direction of defA.
 *  5. **Orthogonalise y_vector** — mirrors R:
 *     `y_vector <- y_vector - (t(y_vector) %*% x_vector) * x_vector`
 *     and, when x1 is available,
 *     `y_vector <- y_vector - (t(y_vector) %*% x1) * x1`,
 *     followed by safe normalisation.
 *  6. **Complete with SVD** — double-deflates V by both x_vector and y_vector,
 *     then calls `complete_rotation()` to fill remaining dimensions.
 *
 * @note Mirrors rENA's `ena.rotate.by.generalized()`
 *       (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
 *
 * @param[in] V       n×q matrix of ENA points (line weights; caller normalises/
 *                    centres upstream).
 * @param[in] params  Fully populated `GeneralizedRotationParams` struct.
 *
 * @returns `RotationResult` with:
 *   - `rotation`     q×q full rotation matrix (column j = axis j).
 *   - `eigenvalues`  q×1 variance per SVD-filled axis (named axes carry 0).
 *   - `column_names` "GMR1", "GMR2" or "SVD2", "SVD3", …, "SVD{q}".
 *
 * @throws std::runtime_error If `gmr_direction()` produces a degenerate direction
 *         for the X or Y axis.
 * @throws std::runtime_error If y_vector has zero norm after orthogonalisation
 *         (can happen when defA is rank-deficient or perfectly aligned with x_vector).
 */
inline RotationResult generalized_means_rotation(
    const arma::mat&                  V,
    const GeneralizedRotationParams&  params
) {
    // ── Step 1: X axis via GMR ────────────────────────────────────────────
    GMRStepResult x_gmr = gmr_step(
        V,
        params.x_model_matrix,
        params.x_target,
        params.x1_cols,
        params.x_categorical,
        params.x_n_groups,
        params.x_subset,
        params.n_lambda,
        params.k_folds,
        params.lasso_eps);

    const arma::vec x_vector = x_gmr.direction;

    // ── Step 2: Deflate V by x_vector ─────────────────────────────────────
    // defA is used only for y_vector computation; x1 uses original V.
    const arma::mat defA = V - V * x_vector * x_vector.t();

    // ── Step 3: Secondary axis x1 ─────────────────────────────────────────
    // Mirrors R: x1 = normalize(mean(V[group0]) - mean(V[group1])).
    //
    // Key points matching the authoritative R:
    //   - Uses original V (not defA)
    //   - Uses the full-length target (all n rows), not just the subset
    //   - Groups are always 0 (first) and 1 (second) in 0-based encoding
    //   - NO fallback: x1 stays empty when categorical/2-group condition fails
    //   - x1 does NOT deflate defA; it is only used to orthogonalise y_vector
    arma::vec x1;

    if (params.x_categorical && params.x_n_groups >= 2) {
        arma::uvec labels = arma::conv_to<arma::uvec>::from(params.x_target);
        arma::uvec g0 = arma::find(labels == 0);
        arma::uvec g1 = arma::find(labels == 1);

        if (!g0.is_empty() && !g1.is_empty()) {
            arma::rowvec m0 = arma::mean(V.rows(g0), 0);   // original V, not defA
            arma::rowvec m1 = arma::mean(V.rows(g1), 0);
            arma::vec diff  = (m0 - m1).t();
            double   dlen   = arma::norm(diff);
            if (dlen > 1e-10) x1 = diff / dlen;
        }
    }

    // ── Step 4: Y axis ────────────────────────────────────────────────────
    // GMR on defA (deflated by x_vector only), or SVD of defA.
    arma::vec y_vector;
    std::string y_name = "SVD2";

    if (params.has_y) {
        GMRStepResult y_gmr = gmr_step(
            defA,
            params.y_model_matrix,
            params.y_target,
            params.y1_cols,
            params.y_categorical,
            params.y_n_groups,
            {},          // no row subset for y axis
            params.n_lambda,
            params.k_folds,
            params.lasso_eps);
        y_vector = y_gmr.direction;
        y_name   = "GMR2";
    } else {
        RotationResult svd_r = ena_svd(defA);
        y_vector = svd_r.rotation.col(0);
    }

    // ── Step 5: Orthogonalise y_vector ────────────────────────────────────
    // Mirrors R (authoritative commit):
    //   y_vector <- y_vector - (t(y_vector) %*% x_vector) * x_vector
    //   if (!is.null(x1)) y_vector <- y_vector - (t(y_vector) %*% x1) * x1
    //   y_vector <- safe_normalize(y_vector)
    y_vector -= arma::dot(y_vector, x_vector) * x_vector;
    if (!x1.is_empty())
        y_vector -= arma::dot(y_vector, x1) * x1;

    double yn = arma::norm(y_vector);
    if (yn < 1e-12)
        throw std::runtime_error(
            "generalized_means_rotation: y_vector has zero norm after orthogonalization");
    y_vector /= yn;

    // ── Step 6: Combine and complete with SVD ─────────────────────────────
    // Double-deflate the original data (mirrors R):
    //   defA <- A - A %*% x %*% t(x) - A %*% y %*% t(y)
    // then pass to prcomp (equivalent: SVD of defA_final).
    arma::mat defA_final = V
        - V * x_vector * x_vector.t()
        - V * y_vector * y_vector.t();

    arma::mat named_axes = arma::join_rows(x_vector, y_vector);
    return complete_rotation(defA_final, named_axes, {"GMR1", y_name});
}

/// @}

}  // namespace qe
#endif  // LIBQE_GENERALIZED_ROTATION_HPP

/**
 * @file lasso.hpp
 * @brief Coordinate-descent Lasso for multivariate response with k-fold CV.
 *
 * Mirrors `glmnet(x, y, family = "mgaussian", penalty.factor = ...)` with
 * k-fold cross-validation to select lambda, and warm-starts along the lambda
 * path for efficiency.
 *
 * The public entry point is qe::lasso_x1_contribution(), which returns an
 * n×q matrix of fitted values attributable to the "target" predictors at the
 * CV-selected lambda.  It is used by `generalized_means_rotation` to isolate
 * the target variable's contribution to the ENA point space after controlling
 * for covariates.
 *
 * **Penalty factors (`pf`):**
 * - `0.0` — predictor is always included (forced in; no shrinkage)
 * - `1.0` — standard L1 penalty
 *
 * The caller is responsible for building the model matrix (main effects and/or
 * interactions); this function accepts a numeric matrix, not a formula.
 */

#ifndef LIBQE_LASSO_HPP
#define LIBQE_LASSO_HPP

#include <armadillo>
#include <cmath>
#include <limits>
#include <stdexcept>

namespace qe {

/**
 * @brief Soft-threshold (shrinkage) operator used by coordinate descent.
 *
 * @details Computes the scalar soft-threshold function:
 *   \f$ S(z,\gamma) = \operatorname{sign}(z)\,\max(|z|-\gamma,\,0) \f$
 *
 * @param[in] z     Input value.
 * @param[in] gamma Threshold amount (must be non-negative for meaningful use).
 * @returns Soft-thresholded value; exactly zero when `|z| <= gamma`.
 */
inline double soft_threshold(double z, double gamma) noexcept {
    return z > gamma ? z - gamma : z < -gamma ? z + gamma : 0.0;
}

/**
 * @brief Single-response Lasso via cyclic coordinate descent at a fixed lambda.
 *
 * @details Fits a Lasso regression for a single response vector using cyclic
 * coordinate descent (naive updates).  The partial-residual update for
 * predictor \f$j\f$ is:
 * \f[
 *   z_j = \frac{X_j^\top r}{n} + \frac{X_j^\top X_j}{n}\,\beta_j
 * \f]
 * \f[
 *   \beta_j \leftarrow \frac{S\!\left(z_j,\;\lambda\,\mathit{pf}_j\right)}{X_j^\top X_j / n}
 * \f]
 * where \f$r = y - X\beta\f$ is the residual maintained incrementally.
 * Predictors with \f$\mathit{pf}_j = 0\f$ skip the soft-threshold and receive
 * the unconstrained OLS update instead.
 *
 * @param[in] X        n×p column-standardized predictor matrix.
 * @param[in] y        n×1 response vector (centering recommended but not required).
 * @param[in] lambda   Regularization strength, on the standardized predictor scale.
 * @param[in] pf       p×1 penalty factors (0 = unpenalized / forced in, 1 = standard L1).
 * @param[in] beta0    Warm-start coefficient vector; pass an empty `arma::vec{}`
 *                     for a cold start from zero.
 * @param[in] max_iter Maximum number of coordinate-descent passes (default 1000).
 * @param[in] tol      Convergence tolerance on the maximum coefficient change
 *                     per pass (default 1e-7).
 * @returns p×1 coefficient vector on the standardized predictor scale.
 */
inline arma::vec lasso_cd(
    const arma::mat& X,
    const arma::vec& y,
    double           lambda,
    const arma::vec& pf,
    arma::vec        beta0   = {},
    int              max_iter = 1000,
    double           tol      = 1e-7
) {
    const arma::uword p = X.n_cols;
    const double      n = static_cast<double>(X.n_rows);

    // Precompute X_j'X_j / n for each predictor
    arma::vec xTx(p);
    for (arma::uword j = 0; j < p; ++j)
        xTx(j) = arma::dot(X.col(j), X.col(j)) / n;

    arma::vec beta = beta0.is_empty()
                   ? arma::vec(p, arma::fill::zeros)
                   : beta0;
    arma::vec r = y - X * beta;   // residual, maintained incrementally

    for (int it = 0; it < max_iter; ++it) {
        double max_delta = 0.0;

        for (arma::uword j = 0; j < p; ++j) {
            if (xTx(j) < 1e-14) continue;   // skip zero-variance predictor

            // Partial correlation statistic
            double z = arma::dot(X.col(j), r) / n + xTx(j) * beta(j);

            double b_new = (pf(j) < 1e-12)
                ? z / xTx(j)                                   // forced in: OLS
                : soft_threshold(z, lambda * pf(j)) / xTx(j); // L1

            double delta = b_new - beta(j);
            if (std::abs(delta) > max_delta) max_delta = std::abs(delta);
            r     -= X.col(j) * delta;
            beta(j) = b_new;
        }

        if (max_delta < tol) break;
    }
    return beta;
}

/**
 * @brief Compute the smallest lambda that drives all penalized coefficients to zero.
 *
 * @details For column-standardized `X` and column-centered `Yc`, lambda_max is:
 * \f[
 *   \lambda_{\max} = \max_{\substack{j:\,\mathit{pf}_j > 0 \\ k}} \frac{|X_j^\top Y_{c,k}|}{n}
 * \f]
 * Above this value every penalized coefficient is exactly zero (unpenalized
 * predictors are unaffected and remain in via OLS regardless of lambda).
 *
 * @param[in] X   n×p column-standardized predictor matrix.
 * @param[in] Yc  n×q column-centered response matrix.
 * @param[in] pf  p×1 penalty factors; predictors with `pf(j) < 1e-12` are skipped.
 * @returns Scalar lambda_max value.
 */
inline double lasso_lambda_max(
    const arma::mat& X,   // column-standardized
    const arma::mat& Yc,  // column-centered responses
    const arma::vec& pf
) {
    const double n = static_cast<double>(X.n_rows);
    double lmax = 0.0;
    for (arma::uword j = 0; j < X.n_cols; ++j) {
        if (pf(j) < 1e-12) continue;   // skip unpenalized predictors
        for (arma::uword k = 0; k < Yc.n_cols; ++k) {
            double v = std::abs(arma::dot(X.col(j), Yc.col(k))) / n;
            if (v > lmax) lmax = v;
        }
    }
    return lmax;
}

/**
 * @brief Cross-validated multivariate Lasso; returns fitted values for target columns only.
 *
 * @details Fits a multivariate Lasso (one independent single-response problem per
 * response column) along a log-spaced lambda path from lambda_max down to
 * `eps * lambda_max`.  Lambda is selected by k-fold cross-validation
 * (mean squared error, averaged over folds and response columns), and the
 * model is refit on the full data at the selected lambda using warm-starts
 * along the path.
 *
 * **Preprocessing performed internally:**
 * - `X_raw` is column-standardized to zero mean and unit variance (zero-variance
 *   columns are left as-is to avoid division by zero).
 * - `Y` is column-centered to absorb the intercept.
 * - Coefficients are unstandardized before the contribution is computed.
 *
 * **Fold assignment:** row \f$i\f$ is assigned to fold \f$i \bmod k\f$
 * (deterministic, order-dependent).
 *
 * **Return value:** The n×q matrix \f$X_{\text{raw},x_1} \hat{B}_{x_1}\f$
 * represents the portion of `Y` (mean-removed) that is linearly explained by
 * the target predictors after penalizing covariate coefficients.  A zero
 * matrix is returned when the target predictors carry no signal at any lambda
 * (e.g., the target is constant after subsetting).  Constant shifts cancel
 * in the downstream between-group scatter (categorical) and slope regression
 * (numeric), so no `ymu` adjustment is applied to the return value.
 *
 * @param[in] X_raw    n×p raw model matrix built by the caller (main effects
 *                     and/or interactions); no formula parsing is done here.
 * @param[in] Y        n×q ENA point matrix (multivariate response).
 * @param[in] x1_cols  0-based column indices in `X_raw` that belong to the
 *                     target variable; these receive `pf = 0` (always included).
 * @param[in] pf       p×1 penalty factors: 0 for `x1_cols`, 1 elsewhere.
 *                     The caller constructs this; non-standard weightings are
 *                     supported.
 * @param[in] n_lambda Length of the lambda path (default 50).
 * @param[in] k_folds  Number of cross-validation folds (default 5; clamped to
 *                     `[2, n/2]`).
 * @param[in] eps      Ratio `lambda_min / lambda_max`; controls how far down
 *                     the path to search (default 0.01).
 * @returns n×q matrix of fitted values attributable to `x1_cols` at the
 *          CV-selected lambda (lambda.min), on the original `Y` scale.
 */
inline arma::mat lasso_x1_contribution(
    const arma::mat& X_raw,
    const arma::mat& Y,
    const arma::uvec& x1_cols,
    const arma::vec&  pf,
    int    n_lambda = 50,
    int    k_folds  = 5,
    double eps      = 0.01
) {
    const arma::uword n = X_raw.n_rows;
    const arma::uword p = X_raw.n_cols;
    const arma::uword q = Y.n_cols;

    // Guard: need at least 2 folds and 2 observations per fold
    k_folds = std::max(2, std::min(k_folds, static_cast<int>(n) / 2));

    // ── Standardize X (μ=0, σ=1); guard zero-variance predictors ────────────
    arma::rowvec xmu  = arma::mean(X_raw, 0);
    arma::rowvec xstd = arma::stddev(X_raw, 0, 0);
    xstd.elem(arma::find(xstd < 1e-12)).ones();   // avoid divide-by-zero
    arma::mat X = (X_raw.each_row() - xmu).each_row() / xstd;

    // ── Center Y (absorbs intercept into the standardized-scale fit) ─────────
    arma::rowvec ymu = arma::mean(Y, 0);
    arma::mat    Yc  = Y.each_row() - ymu;

    // ── Build lambda path: lambda_max → lambda_min on log scale ─────────────
    double lmax = lasso_lambda_max(X, Yc, pf);
    if (lmax < 1e-12)   // target uncorrelated with Y after standardization
        return arma::mat(n, q, arma::fill::zeros);

    arma::vec lambdas = arma::exp(
        arma::linspace(std::log(lmax),
                       std::log(eps * lmax),
                       n_lambda));

    // ── Stratified k-fold CV (fold assignment: row i → fold i % k_folds) ────
    arma::uvec fold_id(n);
    for (arma::uword i = 0; i < n; ++i)
        fold_id(i) = i % static_cast<arma::uword>(k_folds);

    arma::vec cv_err(n_lambda, arma::fill::zeros);

    for (int fold = 0; fold < k_folds; ++fold) {
        arma::uvec tr = arma::find(fold_id != static_cast<arma::uword>(fold));
        arma::uvec va = arma::find(fold_id == static_cast<arma::uword>(fold));

        const arma::mat X_tr = X.rows(tr), X_va = X.rows(va);
        const arma::mat Y_tr = Yc.rows(tr), Y_va = Yc.rows(va);
        const double n_va_q  = static_cast<double>(va.n_elem * q);

        // Warm-start along path for this fold (start cold at lambda_max)
        arma::mat bw(p, q, arma::fill::zeros);
        for (int l = 0; l < n_lambda; ++l) {
            arma::mat b(p, q);
            for (arma::uword k = 0; k < q; ++k)
                b.col(k) = lasso_cd(X_tr, Y_tr.col(k), lambdas(l), pf, bw.col(k));
            bw = b;

            arma::mat err = Y_va - X_va * b;
            cv_err(l) += arma::accu(err % err) / n_va_q;
        }
    }
    cv_err /= static_cast<double>(k_folds);

    // ── Refit on full data at lambda_min, warming along the path ─────────────
    arma::uword best = cv_err.index_min();
    arma::mat bfull(p, q, arma::fill::zeros);
    for (arma::uword l = 0; l <= best; ++l) {
        arma::mat b(p, q);
        for (arma::uword k = 0; k < q; ++k)
            b.col(k) = lasso_cd(X, Yc.col(k), lambdas(l), pf, bfull.col(k));
        bfull = b;
    }

    // ── Unstandardize: beta_orig[j] = beta_std[j] / xstd[j] ─────────────────
    for (arma::uword j = 0; j < p; ++j)
        bfull.row(j) /= xstd(j);

    // ── Return fitted contribution from x1 columns only ──────────────────────
    // This is the portion of Y (mean-removed) linearly explained by the target
    // predictors after covariate adjustment.  Constant shifts cancel in both
    // the between-group scatter (categorical) and slope regression (numeric)
    // that follow, so no ymu adjustment is needed here.
    return X_raw.cols(x1_cols) * bfull.rows(x1_cols);
}

}  // namespace qe
#endif  // LIBQE_LASSO_HPP

/** @file linalg_fallback.hpp
 *  @brief Thin LAPACK-free implementations of SVD, QR, and symmetric
 *         eigendecomposition for use in WASM / no-LAPACK builds.
 *
 *  When @c ARMA_USE_LAPACK is defined these wrappers delegate directly to the
 *  standard Armadillo routines (@c arma::svd, @c arma::qr, @c arma::eig_sym),
 *  so there is zero overhead in normal builds.  When LAPACK is not available
 *  (e.g. Emscripten WASM builds) the wrappers fall back to pure C++
 *  implementations that work for the small matrices typical in ENA
 *  (p ≤ ~100 connection dimensions).
 *
 *  All routines live in namespace @c qe::linalg and operate on @c arma::mat
 *  (double precision).
 *
 *  Fallback algorithms:
 *  - **SVD**      — one-sided Jacobi iterations on \f$A^T\f$ (Demmel & Veselić 1992)
 *  - **QR**       — modified Gram-Schmidt
 *  - **eig_sym**  — cyclic Jacobi (classical Jacobi sweep)
 */

#ifndef LIBQE_LINALG_FALLBACK_HPP
#define LIBQE_LINALG_FALLBACK_HPP

#include <armadillo>
#include <cmath>
#include <limits>

namespace qe {
namespace linalg {

/// @cond INTERNAL
// ── helpers ───────────────────────────────────────────────────────────────────

namespace detail {

/// @brief Modified Gram-Schmidt orthonormalization applied in-place.
///
/// Columns of @p Q are successively orthonormalized against all previously
/// processed columns.  A column whose norm falls below 1e-14 after projection
/// is left unchanged (not zeroed) to avoid introducing NaNs.
///
/// @param Q  Matrix whose columns are orthonormalized in-place (modified).
inline void mgs(arma::mat& Q) {
    const arma::uword p = Q.n_cols;
    for (arma::uword j = 0; j < p; ++j) {
        for (arma::uword i = 0; i < j; ++i)
            Q.col(j) -= arma::dot(Q.col(i), Q.col(j)) * Q.col(i);
        double n = arma::norm(Q.col(j));
        if (n > 1e-14) Q.col(j) /= n;
    }
}

/// @brief One cyclic Jacobi sweep on a p×p symmetric matrix.
///
/// Iterates over all super-diagonal index pairs @c (q,r) and applies a
/// two-sided Jacobi rotation that annihilates @c A(q,r).  The same rotation
/// is accumulated in @p V so that, starting with @c V = I, the full sequence
/// of sweeps yields @c V = eigenvector matrix.
///
/// @param A  Symmetric p×p matrix, updated in-place toward diagonal form.
/// @param V  Rotation accumulator (p×p); initialize to identity to recover
///           eigenvectors.
///
/// @note Algorithm follows the classical cyclic-by-rows Jacobi scheme.
///       Off-diagonal entries no larger than 1e-14 × (|a_qq| + |a_rr|) are
///       skipped for numerical stability.  The `<=` (not `<`) guard correctly
///       handles the all-zero case (a_qq = a_rr = a_qr = 0) that would
///       otherwise produce NaN via 0/0 in the theta computation.
inline void jacobi_sweep(arma::mat& A, arma::mat& V) {
    const arma::uword p = A.n_rows;
    for (arma::uword q = 0; q < p - 1; ++q) {
        for (arma::uword r = q + 1; r < p; ++r) {
            double aqq = A(q, q), arr = A(r, r), aqr = A(q, r);
            if (std::abs(aqr) <= 1e-14 * (std::abs(aqq) + std::abs(arr)))
                continue;
            double theta = (arr - aqq) / (2.0 * aqr);
            double t = (theta >= 0.0)
                ? 1.0 / (theta + std::sqrt(1.0 + theta * theta))
                : 1.0 / (theta - std::sqrt(1.0 + theta * theta));
            double c = 1.0 / std::sqrt(1.0 + t * t); ///< cosine of Jacobi rotation angle
            double s = t * c;                          ///< sine of Jacobi rotation angle
            // Apply Jacobi rotation to A (symmetric update)
            double new_qq = aqq - t * aqr;
            double new_rr = arr + t * aqr;
            A(q, q) = new_qq;  A(r, r) = new_rr;  A(q, r) = A(r, q) = 0.0;
            for (arma::uword k = 0; k < p; ++k) {
                if (k == q || k == r) continue;
                double akq = A(k, q), akr = A(k, r);
                A(k, q) = A(q, k) = c * akq - s * akr;
                A(k, r) = A(r, k) = s * akq + c * akr;
            }
            // Accumulate rotation in V
            for (arma::uword k = 0; k < p; ++k) {
                double vkq = V(k, q), vkr = V(k, r);
                V(k, q) = c * vkq - s * vkr;
                V(k, r) = s * vkq + c * vkr;
            }
        }
    }
}

} // namespace detail

/// @endcond

/// @name Eigendecomposition
/// @{

// ── eig_sym ───────────────────────────────────────────────────────────────────

/// @brief Eigendecomposition of a real symmetric matrix.
///
/// Computes eigenvalues and eigenvectors of the symmetric p×p matrix @p S and
/// returns them sorted in **ascending** eigenvalue order.
///
/// @param[out] eigval  p×1 vector of eigenvalues, ascending.
/// @param[out] eigvec  p×p matrix whose columns are the corresponding
///                     orthonormal eigenvectors.
/// @param[in]  S       Real symmetric p×p input matrix (not modified).
///
/// @warning With @c ARMA_USE_LAPACK defined this calls @c arma::eig_sym
///          directly.  Without LAPACK the fallback runs up to @c 50*p cyclic
///          Jacobi sweeps; convergence is guaranteed for well-conditioned
///          matrices but may be slow for p > 200.
///
/// @note Fallback algorithm: classical cyclic-by-rows Jacobi iteration.
///       Convergence is declared when the Frobenius norm of the strictly
///       off-diagonal part drops below 1e-13.
inline void eig_sym(arma::vec& eigval, arma::mat& eigvec, const arma::mat& S) {
#if defined(ARMA_USE_LAPACK)
    arma::eig_sym(eigval, eigvec, S);
#else
    const arma::uword p = S.n_rows;
    arma::mat A = S;                              ///< working copy of S
    eigvec.eye(p, p);                             ///< accumulate rotations; starts as identity

    for (int sweep = 0; sweep < 50 * static_cast<int>(p); ++sweep) {
        // Check off-diagonal convergence
        double off = 0.0;
        for (arma::uword q = 0; q < p - 1; ++q)
            for (arma::uword r = q + 1; r < p; ++r)
                off += A(q, r) * A(q, r);
        if (std::sqrt(off) < 1e-13) break;
        detail::jacobi_sweep(A, eigvec);
    }
    eigval = A.diag();

    // Sort ascending
    arma::uvec idx = arma::sort_index(eigval);
    eigval  = eigval(idx);
    eigvec  = eigvec.cols(idx);
#endif
}

/// @}

/// @name QR Factorization
/// @{

// ── qr_full ───────────────────────────────────────────────────────────────────

/// @brief Full (square) QR factorization of an n×p matrix.
///
/// Produces a square n×n orthogonal factor @p Q and an n×p upper-trapezoidal
/// factor @p R such that @c A = Q*R.  This matches the convention of
/// @c arma::qr(), which returns the full (non-economy) @p Q.
///
/// @param[out] Q  n×n orthogonal matrix.
/// @param[out] R  n×p upper-trapezoidal matrix.
/// @param[in]  A  n×p input matrix (not modified).
///
/// @warning With @c ARMA_USE_LAPACK defined this calls @c arma::qr directly.
///          Without LAPACK the fallback constructs the full ONB by seeding
///          columns beyond @c A.n_cols with standard basis vectors and then
///          applying modified Gram-Schmidt; numerical quality degrades for
///          nearly rank-deficient inputs.
///
/// @note The caller (@c orthogonal_svd) depends on @p Q being n×n so that
///       @c Q.cols(k, n-1) spans the null space of @c A^T.
inline void qr_full(arma::mat& Q, arma::mat& R, const arma::mat& A) {
#if defined(ARMA_USE_LAPACK)
    arma::qr(Q, R, A);
#else
    // Build a full n×n orthogonal matrix.
    // Step 1: thin QR via modified Gram-Schmidt (gives first k=A.n_cols columns).
    // Step 2: extend to a full ONB using random vectors + re-orthogonalisation.
    const arma::uword n = A.n_rows, p = A.n_cols;

    // Start with all n columns of A padded with random vectors if n > p.
    arma::mat Qfull(n, n, arma::fill::zeros);
    // Copy A's columns first
    for (arma::uword j = 0; j < p && j < n; ++j)
        Qfull.col(j) = A.col(j);
    // Fill remaining columns with standard basis vectors (or random fallback)
    for (arma::uword j = p; j < n; ++j) {
        Qfull(j, j) = 1.0;   ///< standard basis vector e_j used as seed
    }
    // Full MGS orthonormalization
    detail::mgs(Qfull);
    Q = Qfull;

    // R is n×p (upper trapezoidal); only first p rows are non-zero
    R = Q.t() * A;
    // Zero the strictly lower-triangular part of R[0:p, 0:p]
    for (arma::uword j = 0; j < p; ++j)
        for (arma::uword i = j + 1; i < n; ++i)
            R(i, j) = 0.0;
#endif
}

/// @}

/// @name Singular Value Decomposition
/// @{

// ── svd ───────────────────────────────────────────────────────────────────────

/// @brief Full SVD: @c A ≈ U * diag(s) * V', singular values descending.
///
/// Decomposes the n×p matrix @p A into left singular vectors @p U, singular
/// values @p s, and right singular vectors @p V.
///
/// @param[out] U  n×min(n,p) matrix of left singular vectors.
/// @param[out] s  min(n,p)×1 vector of singular values, descending.
/// @param[out] V  p×p matrix of right singular vectors (all p columns).
/// @param[in]  A  n×p input matrix (not modified).
///
/// @warning With @c ARMA_USE_LAPACK defined this calls @c arma::svd directly.
///          Without LAPACK the fallback computes @c A^T*A, calls @c eig_sym,
///          and derives @p U via @c U[:,j] = A*V[:,j]/s_j.  Columns with
///          @c s_j < 1e-14 are zeroed.  Accuracy is governed by the
///          conditioning of @c A^T*A (squares the condition number of @p A).
///
/// @note Fallback algorithm: eigendecomposition of \f$A^T A\f$ via cyclic
///       Jacobi (Demmel & Veselić 1992).  Suitable for the small, moderately
///       conditioned matrices typical in ENA (p ≤ ~100).
inline void svd(arma::mat& U, arma::vec& s, arma::mat& V, const arma::mat& A) {
#if defined(ARMA_USE_LAPACK)
    arma::svd(U, s, V, A);
#else
    // One-sided Jacobi on A^T:
    // We work with B = A^T (p×n). The right singular vectors of A are the
    // left singular vectors of B (= columns of U_B after Jacobi on B^T B).
    // We use the eigendecomposition of A^T * A to get V and s.
    const arma::uword n = A.n_rows, p = A.n_cols;
    const arma::uword k = std::min(n, p);

    arma::mat AtA = A.t() * A;   ///< p×p symmetric Gram matrix
    arma::vec eigval;
    arma::mat eigvec;
    eig_sym(eigval, eigvec, AtA);

    // Eigenvalues from eig_sym are ascending; we want descending.
    arma::uvec idx = arma::sort_index(eigval, "descend");
    eigval = eigval(idx);
    eigvec = eigvec.cols(idx);

    // Singular values (non-negative square roots)
    s.set_size(k);
    for (arma::uword j = 0; j < k; ++j)
        s(j) = std::sqrt(std::max(0.0, eigval(j)));

    // V = eigvec (right singular vectors of A are eigvectors of A^T*A)
    V = eigvec;

    // U = A * V / s (left singular vectors, only first k columns needed)
    U.set_size(n, k);
    for (arma::uword j = 0; j < k; ++j) {
        if (s(j) > 1e-14)
            U.col(j) = A * V.col(j) / s(j);
        else
            U.col(j).zeros();
    }
#endif
}

// ── svd_right ─────────────────────────────────────────────────────────────────

/// @brief Singular values and right singular vectors only (left vectors skipped).
///
/// A lighter-weight variant of @c svd() for callers that only need @p s and
/// @p V.  Under LAPACK the full @c arma::svd is still called internally; the
/// savings come in WASM builds where computing @p U is avoided.
///
/// @param[out] s  min(n,p)×1 vector of singular values, descending.
/// @param[out] V  p×p matrix of right singular vectors.
/// @param[in]  A  n×p input matrix (not modified).
///
/// @warning With @c ARMA_USE_LAPACK defined this still invokes the full
///          @c arma::svd; the left-vector output is discarded.
inline void svd_right(arma::vec& s, arma::mat& V, const arma::mat& A) {
#if defined(ARMA_USE_LAPACK)
    arma::mat U;
    arma::svd(U, s, V, A);
#else
    arma::mat U_unused;
    svd(U_unused, s, V, A);
#endif
}

/// @}

/// @name Linear System Solve
/// @{

// ── solve_spd ─────────────────────────────────────────────────────────────────

/// @brief Solve @c A*x = b for symmetric positive (semi-)definite @p A.
///
/// Returns the n×k solution matrix @p x given the n×n system matrix @p A and
/// the n×k right-hand side @p b.  Near-zero eigenvalues are handled gracefully,
/// making this a pseudo-inverse solve safe for rank-deficient Gram matrices.
///
/// @param[in] A    Real symmetric positive (semi-)definite n×n matrix.
/// @param[in] b    n×k right-hand side matrix.
/// @param[in] tol  Relative threshold below which eigenvalues are treated as
///                 zero.  Eigenvalues with @c |λ| < tol * max(|λ|) are
///                 inverted as zero (i.e. those directions are projected out).
///                 Default: 1e-12.
/// @returns  n×k solution matrix @p x satisfying @c A*x ≈ b.
///
/// @warning With @c ARMA_USE_LAPACK defined this delegates to
///          @c arma::solve(A, b, arma::solve_opts::fast), which uses a
///          Cholesky factorization and does **not** handle rank deficiency.
///          Without LAPACK the eigendecomposition-based pseudo-inverse is used,
///          which is safe for rank-deficient inputs but slower.
///
/// @note Without LAPACK: @c A = V * diag(λ) * V^T is computed via @c eig_sym;
///       the solution is @c x = V * diag(1/λ) * V^T * b with small λ skipped.
inline arma::mat solve_spd(const arma::mat& A, const arma::mat& b,
                            double tol = 1e-12) {
#if defined(ARMA_USE_LAPACK)
    return arma::solve(A, b, arma::solve_opts::fast);
#else
    arma::vec  eigval;
    arma::mat  eigvec;
    eig_sym(eigval, eigvec, A);                     ///< A = eigvec * diag(eigval) * eigvec^T

    const double thresh = tol * std::abs(eigval.max());
    arma::vec inv_eigval(eigval.n_elem, arma::fill::zeros);
    for (arma::uword i = 0; i < eigval.n_elem; ++i)
        if (std::abs(eigval(i)) > thresh)
            inv_eigval(i) = 1.0 / eigval(i);

    // x = eigvec * diag(inv_eigval) * eigvec^T * b
    return eigvec * arma::diagmat(inv_eigval) * eigvec.t() * b;
#endif
}

/// @}

} // namespace linalg
} // namespace qe

#endif // LIBQE_LINALG_FALLBACK_HPP

/**
 * @file modeling.hpp
 * @brief ENA modeling utilities: centering, correlation, confidence intervals,
 *        and node-position solvers for undirected and directed ENA networks.
 */
#ifndef LIBQE_MODELING_HPP
#define LIBQE_MODELING_HPP

#include <armadillo>
#include <algorithm>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <vector>
#include <libqe/linalg_fallback.hpp>

namespace qe {

// ---------------------------------------------------------------------------
/// @name Return types
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Aggregated result returned by every node-position solver.
 *
 * All matrices share the same dimensional conventions:
 * rows index units (or nodes), columns index rotated dimensions.
 */
struct NodePositions {
    arma::mat nodes;      ///< Solved node locations, n_codes × n_dims.
    arma::mat centroids;  ///< Per-unit centroids projected onto the node space, n_units × n_dims.
    arma::mat weights;    ///< Normalised half-edge weight per unit per node, n_units × n_codes.
    arma::mat points;     ///< Input rotated points echoed back unchanged, n_units × n_dims.
};

/// @}

// ---------------------------------------------------------------------------
/// @name Centering
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Subtract column means so that every dimension is centred at the origin.
 *
 * @param[in] values  Matrix of ENA unit points (n_units × n_dims).
 * @returns A copy of @p values with each column mean subtracted (centre-to-origin).
 *
 * @note Equivalent to `center_data_c()` in rENA/ena.cpp.
 */
inline arma::mat center_points(arma::mat values) {
    return values.each_row() - arma::mean(values);
}

/// @}

/// @cond INTERNAL

// ---------------------------------------------------------------------------
// Normal quantile (probit) — pure C++, no R dependency
// Rational approximation by Peter Acklam; max abs error < 1.15e-9.
// Used internally by ena_correlation() so we don't need Rcpp::qnorm.
// ---------------------------------------------------------------------------

inline double normal_quantile(double p) {
    static const double a[] = {
        -3.969683028665376e+01,  2.209460984245205e+02,
        -2.759285104469687e+02,  1.383577518672690e+02,
        -3.066479806614716e+01,  2.506628277459239e+00
    };
    static const double b[] = {
        -5.447609879822406e+01,  1.615858368580409e+02,
        -1.556989798598866e+02,  6.680131188771972e+01,
        -1.328068155288572e+01
    };
    static const double c[] = {
        -7.784894002430293e-03, -3.223964580411365e-01,
        -2.400758277161838e+00, -2.549732539343734e+00,
         4.374664141464968e+00,  2.938163982698783e+00
    };
    static const double d[] = {
         7.784695709041462e-03,  3.224671290700398e-01,
         2.445134137142996e+00,  3.754408661907416e+00
    };
    const double p_low  = 0.02425;
    const double p_high = 1.0 - p_low;
    double q, r;
    if (p < p_low) {
        q = std::sqrt(-2.0 * std::log(p));
        return (((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) /
               ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1.0);
    } else if (p <= p_high) {
        q = p - 0.5; r = q * q;
        return (((((a[0]*r+a[1])*r+a[2])*r+a[3])*r+a[4])*r+a[5])*q /
               (((((b[0]*r+b[1])*r+b[2])*r+b[3])*r+b[4])*r+1.0);
    } else {
        q = std::sqrt(-2.0 * std::log(1.0 - p));
        return -(((((c[0]*q+c[1])*q+c[2])*q+c[3])*q+c[4])*q+c[5]) /
                ((((d[0]*q+d[1])*q+d[2])*q+d[3])*q+1.0);
    }
}

// ---------------------------------------------------------------------------
// Internal helpers: regularized incomplete beta + t-distribution quantile
// ---------------------------------------------------------------------------
//
// These live in qe::detail so they don't pollute the public qe namespace.
// They are used by mean_ci() and are not part of the public API.

namespace detail {

// Lentz continued-fraction evaluation for the regularized incomplete beta.
// Ported from Numerical Recipes in C, §6.4.
inline double betacf(double a, double b, double x) {
    const int    MAXIT = 200;
    const double EPS   = std::numeric_limits<double>::epsilon();
    const double FPMIN = std::numeric_limits<double>::min() / EPS;
    double qab = a + b, qap = a + 1.0, qam = a - 1.0;
    double c = 1.0;
    double d = 1.0 - qab * x / qap;
    if (std::abs(d) < FPMIN) d = FPMIN;
    d = 1.0 / d;
    double h = d;
    for (int m = 1; m <= MAXIT; ++m) {
        const int m2 = 2 * m;
        // even step
        double aa = static_cast<double>(m) * (b - static_cast<double>(m)) * x
                    / ((qam + m2) * (a + m2));
        d = 1.0 + aa * d; if (std::abs(d) < FPMIN) d = FPMIN;
        c = 1.0 + aa / c; if (std::abs(c) < FPMIN) c = FPMIN;
        d = 1.0 / d; h *= d * c;
        // odd step
        aa = -(a + static_cast<double>(m)) * (qab + static_cast<double>(m)) * x
             / ((a + m2) * (qap + m2));
        d = 1.0 + aa * d; if (std::abs(d) < FPMIN) d = FPMIN;
        c = 1.0 + aa / c; if (std::abs(c) < FPMIN) c = FPMIN;
        d = 1.0 / d;
        const double del = d * c;
        h *= del;
        if (std::abs(del - 1.0) <= EPS) break;
    }
    return h;
}

// Regularized incomplete beta I_x(a, b).
inline double betai(double a, double b, double x) {
    if (x <= 0.0) return 0.0;
    if (x >= 1.0) return 1.0;
    const double lbab = std::lgamma(a + b) - std::lgamma(a) - std::lgamma(b);
    const double bt   = std::exp(lbab + a * std::log(x) + b * std::log(1.0 - x));
    if (x < (a + 1.0) / (a + b + 2.0))
        return bt * betacf(a, b, x) / a;
    return 1.0 - bt * betacf(b, a, 1.0 - x) / b;
}

// Quantile function (inverse CDF) for the t-distribution with `df` degrees
// of freedom.  Uses the relation:
//
//   P(T ≤ t | df)  =  1 - I_x(df/2, 1/2) / 2,   x = df / (df + t^2)
//
// Inverted via Newton–Raphson on I_x with bisection fallback; normal-quantile
// plus one Cornish–Fisher correction term provides the starting guess.
inline double t_quantile(double p, double df) {
    const double INF = std::numeric_limits<double>::infinity();
    if (p <= 0.0) return -INF;
    if (p >= 1.0) return  INF;
    if (p == 0.5) return  0.0;
    if (df <= 0.0) return INF;   // df = 0 → t-distribution undefined → treat as ∞

    // Work with pp > 0.5 for numerical stability; negate at the end if needed.
    const bool   flip = (p < 0.5);
    const double pp   = flip ? 1.0 - p : p;
    const double a    = 0.5 * df;
    const double bv   = 0.5;

    // We need x in (0,1) such that I_x(a, bv) = 2*(1-pp).
    const double target = 2.0 * (1.0 - pp);

    // Initial guess via normal quantile + single Cornish–Fisher correction.
    const double z  = normal_quantile(pp);
    const double t0 = z + (z * z * z + z) / (4.0 * df);
    double x = df / (df + t0 * t0);
    x = std::max(1e-12, std::min(1.0 - 1e-12, x));

    // Newton–Raphson with a bisection bracket.
    // f(x) = I_x(a,bv) - target;  f'(x) = x^(a-1)*(1-x)^(bv-1) / B(a,bv)
    const double lbab = std::lgamma(a + bv) - std::lgamma(a) - std::lgamma(bv);
    double xlo = 0.0, xhi = 1.0;
    for (int iter = 0; iter < 100; ++iter) {
        const double fx = betai(a, bv, x) - target;
        // betai is increasing in x, so fx < 0 → x too small → raise lower bound
        if (fx < 0.0) xlo = x; else xhi = x;

        // Derivative: may underflow near boundaries — fall back to bisection.
        const double log_fpx = lbab + (a - 1.0) * std::log(x)
                                     + (bv - 1.0) * std::log(1.0 - x);
        const double fpx = (log_fpx > -700.0) ? std::exp(log_fpx) : 0.0;

        double x_new;
        if (fpx > 0.0) {
            x_new = x - fx / fpx;
            if (x_new <= xlo || x_new >= xhi)
                x_new = 0.5 * (xlo + xhi);
        } else {
            x_new = 0.5 * (xlo + xhi);
        }

        if (std::abs(x_new - x) < 1e-13 * x) break;
        x = x_new;
    }

    const double t_val = std::sqrt(df * (1.0 - x) / x);
    return flip ? -t_val : t_val;
}

// ---------------------------------------------------------------------------
// IQR helper — matches R's quantile(type = 7) (Hyndman–Fan #7), which is
// identical to numpy's np.percentile(method="linear").
//
// For a sorted n-element vector and probability p:
//   h    = (n - 1) * p
//   lo   = floor(h),  frac = h - lo
//   Q(p) = sorted[lo] * (1 - frac) + sorted[lo+1] * frac
// ---------------------------------------------------------------------------

inline double quantile_type7(const arma::vec& sorted, double p) {
    const int n = static_cast<int>(sorted.n_elem);
    if (n == 0) return std::numeric_limits<double>::quiet_NaN();
    if (n == 1) return sorted[0];
    const double h    = (n - 1) * p;
    const int    lo   = static_cast<int>(std::floor(h));
    const double frac = h - lo;
    if (lo + 1 >= n) return sorted[n - 1];          // guard: p == 1.0
    return sorted[lo] * (1.0 - frac) + sorted[lo + 1] * frac;
}

// Interquartile range for a column vector — uses the same type-7 quantile as R.
inline double iqr(const arma::vec& col) {
    const arma::vec s = arma::sort(col);
    return quantile_type7(s, 0.75) - quantile_type7(s, 0.25);
}

} // namespace detail

/// @endcond

// ---------------------------------------------------------------------------
/// @name Correlation
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Pearson correlation with Fisher-z confidence interval between ENA
 *        unit points and their centroids.
 *
 * All unique pairs of units are formed; for each pair the per-dimension
 * difference vectors are computed, then Pearson r is obtained between the
 * @p points differences and the @p centroids differences.  The confidence
 * interval is derived via Fisher's z-transform using the pure-C++
 * `normal_quantile` function so the function works outside R without
 * `Rcpp::qnorm`.
 *
 * @param[in] points      Rotated ENA unit points, n_units × n_dims.
 * @param[in] centroids   Corresponding centroid coordinates, n_units × n_dims.
 * @param[in] conf_level  Confidence level for the interval (default 0.95).
 * @returns An n_dims × 3 matrix whose columns are [r, ci_lower, ci_upper].
 *
 * @note Equivalent to `ena_correlation()` in rENA/ena.cpp.
 */
inline arma::mat ena_correlation(arma::mat points, arma::mat centroids,
                                  double conf_level = 0.95) {
    int n = points.n_rows;
    int n_pairs = (n * (n - 1)) / 2;

    arma::umat idx1(1, n_pairs), idx2(1, n_pairs);
    int col = 0;
    for (int i = 0; i < n; i++)
        for (int j = i + 1; j < n; j++) { idx1[col] = i; idx2[col] = j; col++; }

    arma::mat pts_diff = points.rows(idx1.row(0))    - points.rows(idx2.row(0));
    arma::mat cts_diff = centroids.rows(idx1.row(0)) - centroids.rows(idx2.row(0));
    arma::mat cr       = arma::cor(pts_diff, cts_diff);

    double qq = normal_quantile((1.0 + conf_level) / 2.0);
    arma::mat out(points.n_cols, 3);
    for (arma::uword i = 0; i < points.n_cols; i++) {
        double r     = cr(i, i);
        double z     = std::atanh(r);
        double sigma = 1.0 / std::sqrt(static_cast<double>(n_pairs) - 3.0);
        out(i, 0) = r;
        out(i, 1) = std::tanh(z - sigma * qq);
        out(i, 2) = std::tanh(z + sigma * qq);
    }
    return out;
}

/// @}

// ---------------------------------------------------------------------------
/// @name Group confidence intervals
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Student-t confidence interval for the mean of a group of ENA unit points.
 *
 * For each dimension d:
 * @code
 *   mean ± t_{α/2, n-1} × (sample_sd / sqrt(n))
 * @endcode
 * where α = 1 − @p conf_level and degrees-of-freedom = n − 1.
 *
 * @param[in] points      ENA unit points for a single group, n_units × n_dims.
 * @param[in] conf_level  Confidence level for the interval (default 0.95).
 * @returns An n_dims × 3 matrix whose columns are [mean, ci_lower, ci_upper].
 *          When n == 0 all entries are NaN.
 *          When n == 1 the CI bounds are ±Inf.
 */
inline arma::mat mean_ci(const arma::mat& points, double conf_level = 0.95) {
    const int n      = static_cast<int>(points.n_rows);
    const int n_dims = static_cast<int>(points.n_cols);

    arma::mat out(n_dims, 3);
    out.fill(std::numeric_limits<double>::quiet_NaN());

    if (n == 0) return out;

    const double df     = static_cast<double>(n - 1);
    const double t_crit = detail::t_quantile(1.0 - (1.0 - conf_level) / 2.0, df);
    const double INF    = std::numeric_limits<double>::infinity();

    for (int d = 0; d < n_dims; ++d) {
        const arma::vec col = points.col(d);
        const double mu = arma::mean(col);
        out(d, 0) = mu;
        if (n == 1) {
            // t_crit = Inf and se = 0 would give NaN; explicitly set ±Inf
            out(d, 1) = -INF;
            out(d, 2) =  INF;
        } else {
            const double se = arma::stddev(col) / std::sqrt(static_cast<double>(n));
            out(d, 1) = mu - t_crit * se;
            out(d, 2) = mu + t_crit * se;
        }
    }
    return out;
}

/**
 * @brief Tukey-fence outlier interval for a group of ENA unit points.
 *
 * For each dimension d the half-width is computed as:
 * @code
 *   half_width[d] = iqr_factor * IQR(points[:, d])
 * @endcode
 * and the interval [−half_width, +half_width] is symmetric around zero,
 * matching rENA's formula:
 * @code
 *   outlier.interval.values = matrix(...) * c(-1, 1)
 * @endcode
 * where the matrix is built from `c(IQR(dim1), IQR(dim2), ...) * iqr_factor`.
 *
 * IQR uses R's default type-7 / Hyndman-Fan #7 quantile, identical to
 * `numpy.percentile(method="linear")`.
 *
 * @param[in] points      ENA unit points for a single group, n_units × n_dims.
 * @param[in] iqr_factor  Multiplier applied to the IQR (default 1.5, the
 *                        standard Tukey fence threshold).
 * @returns An n_dims × 2 matrix whose columns are [lower, upper].
 *          When n == 0 all entries are NaN.
 *          When n == 1, IQR == 0, so lower == upper == 0.
 */
inline arma::mat outlier_ci(const arma::mat& points, double iqr_factor = 1.5) {
    const int n_dims = static_cast<int>(points.n_cols);
    arma::mat out(n_dims, 2, arma::fill::zeros);

    if (points.n_rows == 0) {
        out.fill(std::numeric_limits<double>::quiet_NaN());
        return out;
    }

    for (int d = 0; d < n_dims; ++d) {
        const double half_width = iqr_factor * detail::iqr(points.col(d));
        out(d, 0) = -half_width;
        out(d, 1) =  half_width;
    }
    return out;
}

/// @}

// ---------------------------------------------------------------------------
/// @name Node-position solvers
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Multiobjective least-squares node positions for undirected ENA.
 *
 * Half of each adjacency-vector entry (line weight) is distributed to each of
 * its two endpoint nodes, building a per-unit weight matrix.  Each row is then
 * L1-normalised.  An overdetermined system is solved per dimension:
 * @code
 *   (W^T W) X = W^T T
 * @endcode
 * where W is the normalised weight matrix and T contains the rotated unit
 * points.
 *
 * @param[in] adj_mats  Upper-triangular adjacency vectors stacked row-wise,
 *                      n_units × tri_size.
 * @param[in] t         Rotated unit points, n_units × n_dims.
 * @param[in] num_dims  Number of dimensions to solve for.
 * @returns A @ref NodePositions struct containing `nodes`, `centroids`,
 *          `weights`, and `points`.
 *
 * @note Equivalent to `lws_lsq_positions()` in rENA/ena.cpp.
 */
inline NodePositions node_positions(arma::mat adj_mats, arma::mat t,
                                        int num_dims) {
    if (!adj_mats.is_finite() || !t.is_finite())
        throw std::invalid_argument("node_positions: input matrices must not contain NaN or Inf");
    int tri_size  = adj_mats.n_cols;
    int num_nodes = static_cast<int>(
        std::pow(std::ceil(std::sqrt(static_cast<double>(2 * tri_size))), 2.0)
    ) - (2 * tri_size);
    int row_count = adj_mats.n_rows;

    arma::mat weights(row_count, num_nodes, arma::fill::zeros);
    for (int k = 0; k < row_count; k++) {
        arma::rowvec curr = adj_mats.row(k);
        int z = 0;
        for (int x = 0; x < num_nodes - 1; x++) {
            for (int y = 0; y <= x; y++) {
                weights(k, x + 1) += 0.5 * curr[z];
                weights(k, y)     += 0.5 * curr[z];
                z++;
            }
        }
    }
    for (int k = 0; k < row_count; k++) {
        double len = arma::accu(arma::abs(weights.row(k)));
        if (len < 0.0001) len = 0.0001;
        weights.row(k) /= len;
    }

    arma::mat ssX(num_dims, num_nodes, arma::fill::zeros);
    arma::mat ssA = weights.t() * weights;
    for (int i = 0; i < num_dims; i++) {
        arma::mat ssb = weights.t() * t.col(i);
        ssX.row(i) = qe::linalg::solve_spd(ssA, ssb).t();
    }

    NodePositions r;
    r.nodes     = ssX.t();
    r.centroids = (ssX * weights.t()).t();
    r.weights   = weights;
    r.points    = t;
    return r;
}

/**
 * @brief Least-squares node positions for directed (ordered) ENA.
 *
 * Each row of @p line_weights is an n_nodes × n_nodes directed weight matrix
 * stored in row-major order.  The per-unit node weight for node x accumulates
 * the full row weight plus all column weights directed at x from other nodes.
 * Each row is L1-normalised before solving.
 *
 * When @p combine_pairs is `false` (the standard directed case) the system:
 * @code
 *   (W^T W) X = W^T P
 * @endcode
 * is solved directly with W = normalised weight matrix and P = @p points.
 *
 * When @p combine_pairs is `true`, adjacent row pairs (ground row k and
 * response row k+1) are summed before solving:
 * @code
 *   W_combined[k/2] = W[k] + W[k+1]
 *   P_combined[k/2] = P[k] + P[k+1]
 * @endcode
 * The system is then solved on the combined matrices, but centroids are
 * projected back using the original (un-combined) weight matrix so that
 * every unit retains its own centroid.
 *
 * @param[in] line_weights  Directed adjacency vectors, n_units × (n_nodes²).
 * @param[in] points        Rotated unit points, n_units × n_dims.
 * @param[in] num_dims      Number of dimensions to solve for.
 * @param[in] combine_pairs If `true`, sum adjacent row pairs before solving
 *                          (ground + response model).  Default `false`.
 * @returns A @ref NodePositions struct containing `nodes`, `centroids`,
 *          `weights`, and `points`.
 *
 * @note Equivalent to `directed_node_positions()` in rENA/ena.cpp when
 *       @p combine_pairs is `false`.
 * @note Equivalent to `directed_node_positions_with_ground_response_added()`
 *       in rENA/ena.cpp when @p combine_pairs is `true`.
 *
 * @warning When @p combine_pairs is `true`, @p line_weights must have an even
 *          number of rows (row_count must be even); odd row counts result in
 *          an out-of-bounds access when forming the combined matrices.
 */
inline NodePositions directed_node_positions(arma::mat line_weights,
                                              arma::mat points, int num_dims,
                                              bool combine_pairs = false) {
    if (!line_weights.is_finite() || !points.is_finite())
        throw std::invalid_argument("directed_node_positions: input matrices must not contain NaN or Inf");
    int num_nodes = static_cast<int>(
        std::ceil(std::sqrt(static_cast<double>(line_weights.n_cols)))
    );
    int row_count = line_weights.n_rows;

    arma::mat nw(row_count, num_nodes, arma::fill::zeros);
    for (int k = 0; k < row_count; k++) {
        arma::mat curr = line_weights.row(k);
        int z = 0;
        for (int x = 0; x < num_nodes; x++)
            for (int y = 0; y < num_nodes; y++) {
                nw(k, x) += curr[z];
                if (x != y) nw(k, y) += curr[z];
                z++;
            }
    }
    for (int k = 0; k < row_count; k++) {
        double len = arma::accu(arma::abs(nw.row(k)));
        if (len < 0.0001) len = 0.0001;
        nw.row(k) /= len;
    }

    if (combine_pairs) {
        // Combine paired rows (ground at k, response at k+1)
        arma::mat nw_added(row_count / 2, num_nodes, arma::fill::zeros);
        arma::mat pts_added(row_count / 2, num_dims, arma::fill::zeros);
        for (int k = 0; k < row_count; k += 2) {
            nw_added.row(k / 2)  = nw.row(k)     + nw.row(k + 1);
            pts_added.row(k / 2) = points.row(k) + points.row(k + 1);
        }

        arma::mat ssX(num_dims, num_nodes, arma::fill::zeros);
        arma::mat ssA = nw_added.t() * nw_added;
        for (int i = 0; i < num_dims; i++) {
            arma::mat ssb = nw_added.t() * pts_added.col(i);
            ssX.row(i) = qe::linalg::solve_spd(ssA, ssb).t();
        }

        NodePositions r;
        r.nodes     = ssX.t();
        r.centroids = (ssX * nw.t()).t();
        r.weights   = nw;
        r.points    = points;
        return r;
    }

    arma::mat ssX(num_dims, num_nodes, arma::fill::zeros);
    arma::mat ssA = nw.t() * nw;
    for (int i = 0; i < num_dims; i++) {
        arma::mat ssb = nw.t() * points.col(i);
        ssX.row(i) = qe::linalg::solve_spd(ssA, ssb).t();
    }

    NodePositions r;
    r.nodes     = ssX.t();
    r.centroids = (ssX * nw.t()).t();
    r.weights   = nw;
    r.points    = points;
    return r;
}

/// @}

// ---------------------------------------------------------------------------
/// @name Two-group comparison statistics
/// @{
// ---------------------------------------------------------------------------

/**
 * @brief Per-dimension statistics from a two-group comparison.
 *
 * Matches the output of rENA-api's `group.stats()`:
 *
 * - **Parametric** (Welch two-sample t-test): `t`, `df`, `pvalue_t`,
 *   `cohens_d` (pooled-SD formula), `means`, `sds`.
 * - **Non-parametric** (Wilcoxon rank-sum / Mann-Whitney U):
 *   `U` (= R's W statistic for group 1), `pvalue_u` (normal approximation
 *   with tie correction and continuity correction), `effect_r`
 *   (rank-biserial: `1 − 2*U / (n1*n2)`), `medians`.
 *
 * All per-dimension vectors have length `n_dims`; `means`, `sds`, and
 * `medians` are `2 × n_dims` matrices with row 0 = group 1, row 1 = group 2.
 * Entries are NaN when a statistic is undefined (e.g. n < 2 for parametric
 * tests, or n = 0 for either group).
 */
struct GroupStats {
    int n1;       ///< Sample size of group 1.
    int n2;       ///< Sample size of group 2.

    // Parametric
    arma::vec t;          ///< Welch two-sample t-statistics.
    arma::vec df;         ///< Welch–Satterthwaite degrees of freedom.
    arma::vec pvalue_t;   ///< Two-tailed p-values from the Welch t-test.
    arma::vec cohens_d;   ///< Cohen's d (pooled-SD formula, group1 − group2).
    arma::mat means;      ///< 2 × n_dims: means; row 0 = group1, row 1 = group2.
    arma::mat sds;        ///< 2 × n_dims: sample standard deviations.

    // Non-parametric
    arma::vec U;          ///< Wilcoxon rank-sum U for group 1 (= R's W statistic).
    arma::vec pvalue_u;   ///< Two-tailed p-values (normal approx, tie + continuity correction).
    arma::vec effect_r;   ///< Rank-biserial effect: 1 − 2·U / (n1·n2).
    arma::mat medians;    ///< 2 × n_dims: medians; row 0 = group1, row 1 = group2.
};

/// @cond INTERNAL
namespace detail {

// Two-tailed p-value for a Welch t-statistic.
// Uses the relation P(T≤t|df) = 1 − I_x(df/2, 1/2)/2, x = df/(df+t²).
inline double t_test_pvalue(double t_stat, double df) {
    if (!std::isfinite(t_stat) || df <= 0.0)
        return std::numeric_limits<double>::quiet_NaN();
    const double x = df / (df + t_stat * t_stat);
    return betai(df / 2.0, 0.5, x);   // == 2 * P(T_df > |t_stat|)
}

// Wilcoxon rank-sum U (= R's W for group 1) and two-tailed p-value.
// Normal approximation with tie correction and continuity correction,
// matching R's wilcox.test(..., exact=FALSE, correct=TRUE).
struct WilcoxResult { double U; double p; };

inline WilcoxResult wilcox_ranksum(const arma::vec& g1, const arma::vec& g2) {
    const int    n1  = static_cast<int>(g1.n_elem);
    const int    n2  = static_cast<int>(g2.n_elem);
    const int    n   = n1 + n2;
    const double NaN = std::numeric_limits<double>::quiet_NaN();

    if (n1 == 0 || n2 == 0) return {NaN, NaN};

    // Merge both groups into a single array, tagging each element's origin.
    arma::vec      vals(n);
    std::vector<int> grp(n);
    for (int i = 0; i < n1; ++i) { vals[i]      = g1[i]; grp[i]      = 0; }
    for (int i = 0; i < n2; ++i) { vals[n1 + i] = g2[i]; grp[n1 + i] = 1; }

    arma::uvec ord = arma::sort_index(vals);

    // Assign average ranks for tied values; accumulate tie correction T = Σ(t³−t).
    arma::vec rnk(n);
    double    T_ties = 0.0;
    arma::uword i = 0;
    while (i < static_cast<arma::uword>(n)) {
        arma::uword j = i;
        while (j + 1 < static_cast<arma::uword>(n) &&
               vals[ord[j + 1]] == vals[ord[j]])
            ++j;
        const double avg = 0.5 * (static_cast<double>(i) + static_cast<double>(j)) + 1.0;
        for (arma::uword k = i; k <= j; ++k) rnk[ord[k]] = avg;
        const double t = static_cast<double>(j - i + 1);
        T_ties += t * (t * t - 1.0);   // t³ − t
        i = j + 1;
    }

    // W = rank-sum for group 1; U = W − n1*(n1+1)/2 (R's wilcox.test statistic).
    double W = 0.0;
    for (int k = 0; k < n; ++k)
        if (grp[k] == 0) W += rnk[k];
    const double U = W - static_cast<double>(n1) * (n1 + 1) / 2.0;

    // Normal approximation: Var[U] = n1·n2/12 · ((n+1) − T_ties/(n·(n−1)))
    const double mu_U  = static_cast<double>(n1) * n2 / 2.0;
    const double tie_adj = (n > 1) ? T_ties / (static_cast<double>(n) * (n - 1)) : 0.0;
    double var_U = static_cast<double>(n1) * n2 / 12.0
                   * (static_cast<double>(n + 1) - tie_adj);
    if (var_U < 0.0) var_U = 0.0;

    double p;
    if (var_U == 0.0) {
        p = (U == mu_U) ? 1.0 : 0.0;
    } else {
        // Continuity correction: |U − mu| − 0.5
        double z = (std::abs(U - mu_U) - 0.5) / std::sqrt(var_U);
        if (z < 0.0) z = 0.0;
        p = std::erfc(z / std::sqrt(2.0));   // = 2*(1−Φ(z))
        p = std::min(p, 1.0);
    }

    return {U, p};
}

} // namespace detail (group stats helpers)
/// @endcond

/**
 * @brief Compute per-dimension two-group comparison statistics.
 *
 * For each dimension @p d of the input matrices the function computes:
 *
 * **Parametric (Welch two-sample t-test)**
 *  - t-statistic and Welch–Satterthwaite degrees of freedom
 *  - Two-tailed p-value via the regularised incomplete beta function
 *  - Cohen's d using the pooled-SD formula
 *    (`(μ1−μ2) / √(((n1−1)s1² + (n2−1)s2²)/(n1+n2−2))`)
 *  - Per-group means and sample standard deviations
 *
 * **Non-parametric (Wilcoxon rank-sum / Mann-Whitney U)**
 *  - U statistic for group 1 (identical to R's `wilcox.test` W statistic)
 *  - Two-tailed p-value via the normal approximation with tie correction
 *    and continuity correction (matching R's `wilcox.test(..., correct=TRUE)`)
 *  - Rank-biserial effect size: `1 − 2·U / (n1·n2)`
 *    (rENA-api's `nonparam.effect` formula)
 *  - Per-group medians
 *
 * @param[in] g1  Unit points for group 1, n1 × n_dims.
 * @param[in] g2  Unit points for group 2, n2 × n_dims.
 * @returns A @ref GroupStats struct.
 *
 * @note Parametric statistics are NaN when either group has fewer than 2 units
 *       (no sample variance).  All statistics are NaN when either group is empty.
 * @note The Wilcoxon p-value uses the normal approximation regardless of sample
 *       size; R switches to an exact test for small (n ≤ ~50) tie-free samples,
 *       so p-values may differ slightly in those cases.
 * @note Equivalent to `group.stats()` in rENA-api, moved to the C++ layer so
 *       all language bindings share a single implementation.
 */
inline GroupStats group_stats(const arma::mat& g1, const arma::mat& g2) {
    const int    n1     = static_cast<int>(g1.n_rows);
    const int    n2     = static_cast<int>(g2.n_rows);
    const int    n_dims = static_cast<int>(g1.n_cols);
    const double NaN    = std::numeric_limits<double>::quiet_NaN();

    GroupStats out;
    out.n1 = n1;
    out.n2 = n2;

    out.t        = arma::vec(n_dims); out.t.fill(NaN);
    out.df       = arma::vec(n_dims); out.df.fill(NaN);
    out.pvalue_t = arma::vec(n_dims); out.pvalue_t.fill(NaN);
    out.cohens_d = arma::vec(n_dims); out.cohens_d.fill(NaN);
    out.means    = arma::mat(2, n_dims); out.means.fill(NaN);
    out.sds      = arma::mat(2, n_dims); out.sds.fill(NaN);
    out.U        = arma::vec(n_dims); out.U.fill(NaN);
    out.pvalue_u = arma::vec(n_dims); out.pvalue_u.fill(NaN);
    out.effect_r = arma::vec(n_dims); out.effect_r.fill(NaN);
    out.medians  = arma::mat(2, n_dims); out.medians.fill(NaN);

    if (n1 == 0 || n2 == 0) return out;

    for (int d = 0; d < n_dims; ++d) {
        const arma::vec c1 = g1.col(d);
        const arma::vec c2 = g2.col(d);

        // ── Means and medians ───────────────────────────────────────────────
        const double m1 = arma::mean(c1);
        const double m2 = arma::mean(c2);
        out.means(0, d) = m1;
        out.means(1, d) = m2;
        out.medians(0, d) = arma::median(c1);
        out.medians(1, d) = arma::median(c2);

        // ── Standard deviations (ddof = 1, matching R's sd()) ───────────────
        const double s1 = (n1 > 1) ? arma::stddev(c1) : NaN;
        const double s2 = (n2 > 1) ? arma::stddev(c2) : NaN;
        out.sds(0, d) = s1;
        out.sds(1, d) = s2;

        // ── Welch two-sample t-test ──────────────────────────────────────────
        if (n1 > 1 && n2 > 1 && std::isfinite(s1) && std::isfinite(s2)) {
            const double v1n = s1 * s1 / n1;
            const double v2n = s2 * s2 / n2;
            const double se  = std::sqrt(v1n + v2n);
            if (se > 0.0) {
                const double t_stat = (m1 - m2) / se;
                const double df_val = (v1n + v2n) * (v1n + v2n)
                                    / (v1n * v1n / (n1 - 1) + v2n * v2n / (n2 - 1));
                out.t(d)        = t_stat;
                out.df(d)       = df_val;
                out.pvalue_t(d) = detail::t_test_pvalue(t_stat, df_val);
            }
        }

        // ── Cohen's d (pooled SD) ────────────────────────────────────────────
        if (n1 + n2 > 2 && std::isfinite(s1) && std::isfinite(s2)) {
            const double pooled_var = ((n1 - 1) * s1 * s1 + (n2 - 1) * s2 * s2)
                                    / (n1 + n2 - 2);
            const double pooled_sd  = std::sqrt(pooled_var);
            out.cohens_d(d) = (pooled_sd > 0.0) ? (m1 - m2) / pooled_sd : 0.0;
        }

        // ── Wilcoxon rank-sum ────────────────────────────────────────────────
        const detail::WilcoxResult wr = detail::wilcox_ranksum(c1, c2);
        out.U(d)        = wr.U;
        out.pvalue_u(d) = wr.p;
        if (std::isfinite(wr.U))
            out.effect_r(d) = 1.0 - 2.0 * wr.U / (static_cast<double>(n1) * n2);
    }

    return out;
}

/// @}

} // namespace qe

#endif // LIBQE_MODELING_HPP

/**
 * @file normalization.hpp
 * @brief Row-wise normalization utilities for ENA connection networks.
 */
#ifndef LIBQE_NORMALIZATION_HPP
#define LIBQE_NORMALIZATION_HPP

#include <armadillo>
#include <cmath>

namespace qe {

/**
 * @brief Normalize each row to unit L2 norm (project onto the unit hypersphere).
 *
 * Rows whose L2 norm is zero are left unchanged.
 *
 * @param m Input matrix (n_units × n_connections), row-major.
 * @returns Matrix of the same shape with each row divided by its L2 norm.
 * @note Equivalent to @c fun_sphere_norm() in rENA/ena.cpp.
 */
inline arma::mat normalize_networks(arma::mat m) {
    arma::mat out(m.n_rows, m.n_cols, arma::fill::zeros);
    for (arma::uword r = 0; r < m.n_rows; r++) {
        double len = arma::norm(m.row(r), 2);
        if (std::isfinite(len) && len > 0.0) out.row(r) = m.row(r) / len;
    }
    return out;
}

/**
 * @brief Scale every entry by the reciprocal of the largest row L2 norm.
 *
 * Preserves relative magnitudes across rows; does not normalise each row
 * individually.  If the largest row norm is zero, the matrix is returned
 * unchanged.
 *
 * @param m Input matrix (n_units × n_connections), modified in place.
 * @returns Scaled matrix (same shape as @p m).
 * @note Equivalent to @c fun_skip_sphere_norm() in rENA/ena.cpp.
 */
inline arma::mat scale_networks(arma::mat m) {
    double largest = 0.0;
    for (arma::uword r = 0; r < m.n_rows; r++) {
        double len = arma::norm(m.row(r), 2);
        if (std::isfinite(len)) largest = std::max(largest, len);
    }
    if (largest > 0.0) m /= largest;
    return m;
}

} // namespace qe

#endif // LIBQE_NORMALIZATION_HPP

/**
 * @file rotation.hpp
 * @brief Rotation routines for Epistemic Network Analysis (ENA):
 *        SVD, deflation, orthogonal-SVD, complete (generalized), and
 *        means rotation.  All routines are designed to match the
 *        numerical output of the reference R implementation in rENA.
 */

#ifndef LIBQE_ROTATION_HPP
#define LIBQE_ROTATION_HPP

#include "linalg_fallback.hpp"
#include <armadillo>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace qe {

/// @name Return types
/// @{

/**
 * @brief Aggregated result returned by every rotation routine.
 *
 * The layout mirrors the output of R's `prcomp()`: `rotation` holds the
 * right-singular vectors (principal axes), `eigenvalues` holds the
 * explained-variance values (sdev^2 in R parlance), and `column_names`
 * carries the human-readable axis labels used downstream for network plots
 * and loadings tables.
 */
struct RotationResult {
    arma::mat                rotation;      ///< p x p orthogonal matrix; column j is rotation axis j.
    arma::vec                eigenvalues;   ///< Length-p vector of sdev^2 values, matching rENA's convention.
    std::vector<std::string> column_names;  ///< Axis labels, e.g. "MR1", "SVD2", "GMR1".
};

/**
 * @brief A pair of 0-based row-index vectors identifying two groups within a
 *        points matrix.
 *
 * Used by means_rotation() to specify which rows belong to group A and which
 * belong to group B for each successive mean-difference axis.
 */
struct GroupPair {
    arma::uvec a;  ///< Row indices of group A (0-based).
    arma::uvec b;  ///< Row indices of group B (0-based).
};

/// @}

/// @name SVD rotation
/// @{

/**
 * @brief Compute a full SVD rotation of a (pre-centered) points matrix.
 *
 * Produces the right-singular vectors V as the rotation matrix and scales
 * each singular value to an eigenvalue (sdev^2) compatible with rENA's
 * `prcomp()` output.
 *
 * @param points  n x p data matrix.  The caller is responsible for
 *                centering upstream; this function applies no centering.
 *
 * @returns A RotationResult where:
 *   - `rotation`    = V  (p x p full SVD; trailing columns span the null
 *                         space when the matrix is rank-deficient).
 *   - `eigenvalues[j]` = s[j]^2 / max(1, n - 1)  (== sdev^2 in R).
 *   - `column_names`   = {"SVD1", "SVD2", ..., "SVDp"}.
 *
 * @note Equivalent to `prcomp(points, retx=FALSE, scale=FALSE,
 *       center=FALSE, tol=0)` in rENA/R.
 *
 * @note Sign convention: none.  Signs follow the underlying LAPACK SVD,
 *       matching rENA's long-standing behavior.  A deterministic sign rule
 *       (e.g. svd_flip) may be added later as an opt-in flag.
 */
inline RotationResult ena_svd(const arma::mat& points) {
    if (!points.is_finite())
        throw std::invalid_argument("ena_svd: input matrix must not contain NaN or Inf");
    arma::mat U, V;
    arma::vec s;
    qe::linalg::svd(U, s, V, points);

    const arma::uword p     = points.n_cols;
    const double      denom = points.n_rows > 1
                                  ? static_cast<double>(points.n_rows - 1)
                                  : 1.0;

    arma::vec eigenvalues(p, arma::fill::zeros);
    const arma::uword k = std::min<arma::uword>(s.n_elem, p);
    for (arma::uword j = 0; j < k; ++j) {
        eigenvalues(j) = (s(j) * s(j)) / denom;
    }

    std::vector<std::string> labels(p);
    for (arma::uword j = 0; j < p; ++j) {
        labels[j] = "SVD" + std::to_string(j + 1);
    }
    return {V, eigenvalues, std::move(labels)};
}

/// @}

/// @name Deflation
/// @{

/**
 * @brief Project a data matrix onto the hyperplane orthogonal to a given axis.
 *
 * Computes `data - (data * axis) * axis^T`, effectively removing the
 * component of every row of `data` that lies along `axis`.
 *
 * @param data  n x p data matrix.
 * @param axis  Unit-norm column vector of length p.  The caller is
 *              responsible for normalizing `axis` before calling this
 *              function; no normalization is performed internally.
 *
 * @returns n x p matrix with the projection onto `axis` subtracted out.
 */
inline arma::mat deflate(const arma::mat& data, const arma::vec& axis) {
    return data - (data * axis) * axis.t();
}

/// @}

/// @name Orthogonal SVD rotation
/// @{

/**
 * @brief Orthonormalize named axes via QR, then complete the rotation with
 *        an SVD of the data projected onto the orthogonal complement.
 *
 * The named axes in `weights` are orthonormalized via a full QR
 * decomposition.  The first k columns of the resulting Q become the
 * "fixed" part of the rotation; the remaining p - k columns define a
 * complementary subspace onto which `data` is projected, and a further SVD
 * fills those trailing rotation columns.
 *
 * @param data          n x p data matrix (caller-centered if required).
 * @param weights       p x k matrix of named axis directions.  Column norms
 *                      need not be 1; QR handles normalization internally.
 * @param named_labels  Length-k vector of labels for the first k rotation
 *                      axes.  Trailing axes receive labels "SVD{k+1}" ..
 *                      "SVDp".
 *
 * @returns A RotationResult containing the combined p x p rotation,
 *          eigenvalues for the SVD-filled trailing axes (indices k..p-1;
 *          the first k eigenvalues are zero), and axis labels.
 *
 * @note Equivalent to `orthogonal_svd()` in rENA/R/ena.rotate.by.mean.R:
 * @code
 *   Q     = qr.Q(qr(weights), complete = TRUE)   # p x p
 *   X_bar = data %*% Q[, (k+1):p]                # n x (p - k)
 *   V     = prcomp(X_bar)$rotation
 *   out   = cbind(Q[, 1:k], Q[, (k+1):p] %*% V)
 * @endcode
 *
 * @note The named axes in the OUTPUT are the orthonormalized Q columns, NOT
 *       the original `weights` columns.  Use complete_rotation() if you need
 *       to preserve input axes verbatim.
 *
 * @throws std::runtime_error if `data.n_cols != weights.n_rows` or
 *         `named_labels.size() != weights.n_cols`.
 */
inline RotationResult orthogonal_svd(
    const arma::mat&                data,
    const arma::mat&                weights,
    const std::vector<std::string>& named_labels) {
    const arma::uword p = weights.n_rows;
    const arma::uword k = weights.n_cols;
    if (data.n_cols != p) {
        throw std::runtime_error(
            "orthogonal_svd: data.n_cols must equal weights.n_rows");
    }
    if (named_labels.size() != k) {
        throw std::runtime_error(
            "orthogonal_svd: named_labels.size() must equal weights.n_cols");
    }
    if (k == 0) {
        return ena_svd(data);
    }

    arma::mat Q, R;
    qe::linalg::qr_full(Q, R, weights);                 // Q is p x p

    arma::mat rotation(p, p);
    arma::vec eigenvalues(p, arma::fill::zeros);
    rotation.cols(0, k - 1) = Q.cols(0, k - 1);

    if (k < p) {
        arma::mat Q_comp = Q.cols(k, p - 1);            // p x (p - k)
        arma::mat X_bar  = data * Q_comp;               // n x (p - k)
        RotationResult inner = ena_svd(X_bar);
        rotation.cols(k, p - 1) = Q_comp * inner.rotation;
        for (arma::uword j = 0; j < (p - k); ++j) {
            eigenvalues(k + j) = inner.eigenvalues(j);
        }
    }

    std::vector<std::string> labels(p);
    for (arma::uword j = 0; j < k; ++j) labels[j] = named_labels[j];
    for (arma::uword j = k; j < p; ++j) {
        labels[j] = "SVD" + std::to_string(j + 1);
    }
    return {rotation, eigenvalues, std::move(labels)};
}

/// @}

/// @name Complete (generalized) rotation
/// @{

/**
 * @brief Keep named axes verbatim and fill remaining axes from an SVD of the
 *        data after parallel deflation by all named axes.
 *
 * Unlike orthogonal_svd(), this routine preserves the caller-supplied
 * `named_axes` exactly (no QR re-orthonormalization).  The complementary
 * axes are obtained by deflating `data` simultaneously by all named axes and
 * then running ena_svd() on the result.
 *
 * The deflation is **parallel**: every projection is subtracted from the
 * original `data`, not from a progressively deflated copy:
 * @code
 *   defA = data - data * (named_axes * named_axes^T)
 * @endcode
 * This matches `defA <- A - A %*% v1 %*% t(v1) - A %*% v2 %*% t(v2)` in
 * rENA line-for-line.  For mutually orthogonal axes the result is identical
 * to sequential deflation; for non-orthogonal axes it differs.
 *
 * @param data          n x p data matrix.
 * @param named_axes    p x k matrix of named rotation axes.  Each column
 *                      must be unit-norm; orthonormality between columns is
 *                      NOT required.
 * @param named_labels  Length-k vector of labels for the first k axes.
 *                      Trailing axes receive labels "SVD{k+1}" .. "SVDp".
 *                      Conventional labels for generalized rotation are
 *                      "GMR1", "GMR2", then "SVD{k+1}" onward, but labels
 *                      are entirely caller-supplied.
 *
 * @returns A RotationResult where columns 0..k-1 of `rotation` equal
 *          `named_axes` verbatim and columns k..p-1 come from the SVD of
 *          the deflated data.
 *
 * @note Equivalent to the tail of `ena.rotate.by.generalized` in rENA
 *       (canonical version: commit 2c079126 on rENA `origin/main`).
 *
 * @throws std::runtime_error if `data.n_cols != named_axes.n_rows` or
 *         `named_labels.size() != named_axes.n_cols`.
 */
inline RotationResult complete_rotation(
    const arma::mat&                data,
    const arma::mat&                named_axes,
    const std::vector<std::string>& named_labels) {
    const arma::uword p = named_axes.n_rows;
    const arma::uword k = named_axes.n_cols;
    if (data.n_cols != p) {
        throw std::runtime_error(
            "complete_rotation: data.n_cols must equal named_axes.n_rows");
    }
    if (named_labels.size() != k) {
        throw std::runtime_error(
            "complete_rotation: named_labels.size() must equal named_axes.n_cols");
    }
    if (k == 0) {
        return ena_svd(data);
    }

    arma::mat defA = data - data * named_axes * named_axes.t();
    RotationResult inner = ena_svd(defA);

    arma::mat rotation(p, p);
    arma::vec eigenvalues(p, arma::fill::zeros);
    rotation.cols(0, k - 1) = named_axes;
    if (k < p) {
        rotation.cols(k, p - 1) = inner.rotation.cols(0, p - k - 1);
        for (arma::uword j = 0; j < (p - k); ++j) {
            eigenvalues(k + j) = inner.eigenvalues(j);
        }
    }

    std::vector<std::string> labels(p);
    for (arma::uword j = 0; j < k; ++j) labels[j] = named_labels[j];
    for (arma::uword j = k; j < p; ++j) {
        labels[j] = "SVD" + std::to_string(j + 1);
    }
    return {rotation, eigenvalues, std::move(labels)};
}

/// @}

/// @name Means rotation
/// @{

/**
 * @brief Compute a means rotation from one or more group pairs.
 *
 * For each GroupPair, computes the normalized mean-difference vector on the
 * progressively deflated data and stacks the resulting unit-norm axes into a
 * weights matrix.  The final rotation is completed by orthogonal_svd().
 * The input `points` is mean-centered before any axes are computed, matching
 * rENA's `scale(data, scale=FALSE, center=TRUE)`.
 *
 * @param points  n x p matrix of network positions (raw, not pre-centered).
 * @param pairs   Ordered list of GroupPair objects.  Each pair contributes
 *                one "MR" axis in sequence.  Must be non-empty.
 *
 * @returns A RotationResult with `column_names` = {"MR1", ..., "MRm",
 *          "SVD{m+1}", ..., "SVDp"} where m = pairs.size().
 *
 * @note Equivalent to `ena.rotate.by.mean()` in rENA/R/ena.rotate.by.mean.R.
 *       The progressive deflation (each axis is computed on the data after
 *       removing all previous axes) is the key numerical step.
 *
 * @warning There is NO guard against a zero-norm mean-difference vector.
 *          If the two group means are identical on the current deflated data,
 *          the axis will contain NaN values (division by zero).  This mirrors
 *          a latent bug in rENA and is preserved verbatim for numerical
 *          fidelity; a safe opt-in guard will be added in a future release.
 *
 * @throws std::runtime_error if `pairs` is empty.
 */
inline RotationResult means_rotation(const arma::mat&              points,
                                      const std::vector<GroupPair>& pairs) {
    if (pairs.empty()) {
        throw std::runtime_error("Unable to rotate without 2 groups.");
    }

    arma::mat data     = points.each_row() - arma::mean(points, 0);
    arma::mat deflated = data;
    arma::mat weights(data.n_cols, pairs.size(), arma::fill::zeros);
    std::vector<std::string> labels;
    labels.reserve(pairs.size());

    for (std::size_t i = 0; i < pairs.size(); ++i) {
        const arma::rowvec mean_a = arma::mean(deflated.rows(pairs[i].a), 0);
        const arma::rowvec mean_b = arma::mean(deflated.rows(pairs[i].b), 0);
        const arma::vec    diff   = (mean_a - mean_b).t();
        const double       norm   = std::sqrt(arma::accu(diff % diff));
        const arma::vec    axis   = diff / norm;        // no zero-norm guard

        deflated        = deflate(deflated, axis);
        weights.col(i)  = axis;
        labels.push_back("MR" + std::to_string(i + 1));
    }

    return orthogonal_svd(deflated, weights, labels);
}

/// @}

}  // namespace qe

#endif  // LIBQE_ROTATION_HPP

/** @file stability.hpp
 *  @brief Stability metrics and distance-distance matrix correlation for parameter tuning.
 */
#ifndef LIBQE_STABILITY_HPP
#define LIBQE_STABILITY_HPP

#include <armadillo>
#include <cmath>
#include <algorithm>

namespace qe {

/** @brief Compute Pearson correlation between pairwise Euclidean distance structures of two matrices.
 *
 *  Evaluates dist-dist correlation: r(dist(X), dist(Y)), measuring whether relative
 *  neighborhoods and distances are preserved across parameter changes (e.g. window vs lookback).
 *
 *  @param[in] X Numeric matrix (n_rows × n_cols).
 *  @param[in] Y Numeric matrix (n_rows × n_cols).
 *
 *  @returns Pearson correlation coefficient in [-1, 1].
 */
inline double dist_dist_correlation(const arma::mat& X, const arma::mat& Y) {
    size_t n = X.n_rows;
    if (n != Y.n_rows || n < 3) return 0.0;

    size_t n_pairs = n * (n - 1) / 2;
    arma::vec dx(n_pairs);
    arma::vec dy(n_pairs);

    size_t idx = 0;
    for (size_t i = 0; i < n; ++i) {
        for (size_t j = i + 1; j < n; ++j) {
            arma::rowvec diff_x = X.row(i) - X.row(j);
            arma::rowvec diff_y = Y.row(i) - Y.row(j);
            dx[idx] = std::sqrt(arma::dot(diff_x, diff_x));
            dy[idx] = std::sqrt(arma::dot(diff_y, diff_y));
            idx++;
        }
    }

    double mean_x = arma::mean(dx);
    double mean_y = arma::mean(dy);

    arma::vec zx = dx - mean_x;
    arma::vec zy = dy - mean_y;

    double sxx = arma::dot(zx, zx);
    double syy = arma::dot(zy, zy);
    double sxy = arma::dot(zx, zy);

    if (sxx <= 1e-12 || syy <= 1e-12) return 0.0;
    return sxy / std::sqrt(sxx * syy);
}

} // namespace qe

#endif // LIBQE_STABILITY_HPP

/** @file trajectory_distance.hpp
 *  @brief Integrated Euclidean distance between trajectory curves and lagged distance metrics.
 */
#ifndef LIBQE_TRAJECTORY_DISTANCE_HPP
#define LIBQE_TRAJECTORY_DISTANCE_HPP

#include <armadillo>
#include <vector>
#include <cmath>
#include <algorithm>
#include "trajectory.hpp"

namespace qe {

struct GaussLegendreRule {
    arma::vec nodes;
    arma::vec weights;
};

/** @brief Obtain exact 32-point Gauss-Legendre quadrature nodes and weights on [-1, 1] via Golub-Welsch. */
inline const GaussLegendreRule& get_gl_rule_32() {
    static const GaussLegendreRule rule = []() {
        int N = 32;
        arma::mat J(N, N, arma::fill::zeros);
        for (int k = 1; k < N; ++k) {
            double b = static_cast<double>(k) / std::sqrt(4.0 * k * k - 1.0);
            J(k, k - 1) = b;
            J(k - 1, k) = b;
        }
        arma::vec evals;
        arma::mat evecs;
        arma::eig_sym(evals, evecs, J);
        arma::vec w = 2.0 * arma::square(evecs.row(0).t());
        return GaussLegendreRule{evals, w};
    }();
    return rule;
}

/** @brief Compute integrated Euclidean distance between two 2D parametric polynomial curves.
 *
 *  Computes integral_{t_start}^{t_end} ||F_a(t) - F_b(t)|| dt
 *  using 32-point Gauss-Legendre quadrature.
 *
 *  @param[in] coeffs_ax  Coefficients of x_a(t).
 *  @param[in] coeffs_ay  Coefficients of y_a(t).
 *  @param[in] coeffs_bx  Coefficients of x_b(t).
 *  @param[in] coeffs_by  Coefficients of y_b(t).
 *  @param[in] t_start    Start of integration interval (default 0.0).
 *  @param[in] t_end      End of integration interval (default 1.0).
 *
 *  @returns Integrated distance scalar.
 */
inline double integrated_curve_distance(
    const arma::vec& coeffs_ax,
    const arma::vec& coeffs_ay,
    const arma::vec& coeffs_bx,
    const arma::vec& coeffs_by,
    double t_start = 0.0,
    double t_end = 1.0
) {
    if (t_end <= t_start) return 0.0;

    const auto& rule = get_gl_rule_32();
    double mid = 0.5 * (t_start + t_end);
    double half_len = 0.5 * (t_end - t_start);
    double total = 0.0;

    for (size_t i = 0; i < rule.nodes.n_elem; ++i) {
        double t_eval = mid + half_len * rule.nodes[i];
        double xa = eval_poly_scalar(coeffs_ax, t_eval);
        double ya = eval_poly_scalar(coeffs_ay, t_eval);
        double xb = eval_poly_scalar(coeffs_bx, t_eval);
        double yb = eval_poly_scalar(coeffs_by, t_eval);
        double dx = xa - xb;
        double dy = ya - yb;
        double dist = std::sqrt(dx * dx + dy * dy);

        total += rule.weights[i] * dist;
    }

    return half_len * total;
}

/** @brief Compute lagged curve distance: 1/(1-lag) * integral_{lag}^1 ||F_follower(t) - F_leader(t - lag)|| dt. */
inline double lagged_curve_distance(
    const arma::vec& coeffs_fol_x,
    const arma::vec& coeffs_fol_y,
    const arma::vec& coeffs_ldr_x,
    const arma::vec& coeffs_ldr_y,
    double lag = 0.0
) {
    if (lag < 0.0 || lag >= 0.999) return 0.0;
    double t_start = lag;
    double t_end = 1.0;
    const auto& rule = get_gl_rule_32();
    double mid = 0.5 * (t_start + t_end);
    double half_len = 0.5 * (t_end - t_start);
    double total = 0.0;

    for (size_t i = 0; i < rule.nodes.n_elem; ++i) {
        double t_fol = mid + half_len * rule.nodes[i];
        double t_ldr = t_fol - lag;

        double x_fol = eval_poly_scalar(coeffs_fol_x, t_fol);
        double y_fol = eval_poly_scalar(coeffs_fol_y, t_fol);
        double x_ldr = eval_poly_scalar(coeffs_ldr_x, t_ldr);
        double y_ldr = eval_poly_scalar(coeffs_ldr_y, t_ldr);

        double dx = x_fol - x_ldr;
        double dy = y_fol - y_ldr;
        double dist = std::sqrt(dx * dx + dy * dy);

        total += rule.weights[i] * dist;
    }

    double integral = half_len * total;
    double interval = 1.0 - lag;
    return (interval > 1e-6) ? (integral / interval) : 0.0;
}

/** @brief Compute pairwise integrated distance matrix across K polynomial trajectory curves. */
inline arma::mat pairwise_trajectory_distance_matrix(
    const std::vector<arma::vec>& all_coeffs_x,
    const std::vector<arma::vec>& all_coeffs_y
) {
    size_t K = all_coeffs_x.size();
    arma::mat D(K, K, arma::fill::zeros);

    for (size_t i = 0; i < K; ++i) {
        for (size_t j = i + 1; j < K; ++j) {
            double dist = integrated_curve_distance(
                all_coeffs_x[i], all_coeffs_y[i],
                all_coeffs_x[j], all_coeffs_y[j]
            );
            D(i, j) = dist;
            D(j, i) = dist;
        }
    }

    return D;
}

} // namespace qe

#endif // LIBQE_TRAJECTORY_DISTANCE_HPP

/** @file trajectory_following.hpp
 *  @brief Who-follows-whom leader-follower turn lag analysis along shared timeline.
 */
#ifndef LIBQE_TRAJECTORY_FOLLOWING_HPP
#define LIBQE_TRAJECTORY_FOLLOWING_HPP

#include <armadillo>
#include <vector>
#include <cmath>
#include <algorithm>
#include <limits>

namespace qe {

/** @brief Result of signed turn-lag sweep across delta in [-max_lag, max_lag]. */
struct SignedTurnLagResult {
    int best_lag = 0;             ///< delta* that minimizes mean distance
    double min_mean_distance = 0.0;
    arma::ivec lags;              ///< Range of evaluated lags
    arma::vec mean_distances;     ///< Mean distance for each lag
    arma::ivec valid_counts;      ///< Number of paired turns for each lag
};

/** @brief Compute mean distance between agent A at turn T and agent B at turn T - delta.
 *
 *  @param[in] pts_a    n_a × 2 coordinates for agent A.
 *  @param[in] pts_b    n_b × 2 coordinates for agent B.
 *  @param[in] times_a  Length n_a vector of calendar/turn indices for A.
 *  @param[in] times_b  Length n_b vector of calendar/turn indices for B.
 *  @param[in] delta    Signed lag in turns (delta > 0: A compares to B's past => A follows B;
 *                      delta < 0: B compares to A's past => B follows A; delta = 0: sync).
 *
 *  @returns Pair of (mean_distance, valid_pairs_count).
 */
inline std::pair<double, int> signed_turn_lag_distance(
    const arma::mat& pts_a,
    const arma::mat& pts_b,
    const arma::vec& times_a,
    const arma::vec& times_b,
    int delta
) {
    int n_a = static_cast<int>(pts_a.n_rows);
    int n_b = static_cast<int>(pts_b.n_rows);
    if (n_a == 0 || n_b == 0) return {0.0, 0};

    double sum_dist = 0.0;
    int count = 0;

    for (int i = 0; i < n_a; ++i) {
        double t_a = times_a[i];
        double target_tb = t_a - static_cast<double>(delta);

        // Find the latest turn for agent B that occurred at or before target_tb
        int best_b_idx = -1;
        for (int j = 0; j < n_b; ++j) {
            if (times_b[j] <= target_tb) {
                best_b_idx = j;
            } else {
                break;
            }
        }

        if (best_b_idx >= 0) {
            double dx = pts_a(i, 0) - pts_b(best_b_idx, 0);
            double dy = pts_a(i, 1) - pts_b(best_b_idx, 1);
            sum_dist += std::sqrt(dx * dx + dy * dy);
            count++;
        }
    }

    double mean_d = (count > 0) ? (sum_dist / count) : 0.0;
    return {mean_d, count};
}

/** @brief Sweep signed turn lags in [-max_lag, max_lag] to determine leader-follower dynamics. */
inline SignedTurnLagResult best_signed_turn_lag(
    const arma::mat& pts_a,
    const arma::mat& pts_b,
    const arma::vec& times_a,
    const arma::vec& times_b,
    int max_lag = 15
) {
    int k_max = std::max(1, max_lag);
    int n_lags = 2 * k_max + 1;

    SignedTurnLagResult res;
    res.lags.set_size(n_lags);
    res.mean_distances.set_size(n_lags);
    res.valid_counts.set_size(n_lags);

    double min_dist = std::numeric_limits<double>::infinity();
    int best_delta = 0;

    for (int idx = 0; idx < n_lags; ++idx) {
        int delta = -k_max + idx;
        res.lags[idx] = delta;

        auto [mean_d, count] = signed_turn_lag_distance(pts_a, pts_b, times_a, times_b, delta);
        res.mean_distances[idx] = mean_d;
        res.valid_counts[idx] = count;

        if (count > 0 && mean_d < min_dist) {
            min_dist = mean_d;
            best_delta = delta;
        }
    }

    res.best_lag = best_delta;
    res.min_mean_distance = (min_dist == std::numeric_limits<double>::infinity()) ? 0.0 : min_dist;
    return res;
}

} // namespace qe

#endif // LIBQE_TRAJECTORY_FOLLOWING_HPP

/** @file trajectory.hpp
 *  @brief Parametric trajectory curve fitting, LOOCV-2D degree selection,
 *         and differential geometry (velocity, acceleration, curvature, turns).
 */
#ifndef LIBQE_TRAJECTORY_HPP
#define LIBQE_TRAJECTORY_HPP

#include <armadillo>
#include <vector>
#include <string>
#include <cmath>
#include <algorithm>
#include <limits>

namespace qe {

/** @brief Polynomial curve fit result in 2D epistemic space. */
struct PolyCurveFit {
    int degree = 1;
    arma::vec coeffs_x;     ///< Raw coefficients [c0, ..., cd] for x(t) = sum(ck * t^k)
    arma::vec coeffs_y;     ///< Raw coefficients [c0, ..., cd] for y(t) = sum(ck * t^k)
    std::string basis = "orthogonal"; ///< Basis used during fitting: "orthogonal" (R stats::poly-like) or "raw".
    arma::vec basis_alpha;  ///< Orthogonal-basis recurrence alpha values (empty for raw basis).
    arma::vec basis_norm2;  ///< Orthogonal-basis recurrence squared norms, matching stats::poly coefs$norm2.
    double cv_error = 0.0;  ///< 2D LOOCV error
    double aic = 0.0;       ///< Joint bivariate AIC
    arma::vec t;            ///< Evaluated or input time points
    arma::mat fitted;       ///< n × 2 fitted coordinates
};

inline arma::mat build_raw_poly_design(const arma::vec& t_vec, int degree) {
    int n = static_cast<int>(t_vec.n_elem);
    arma::mat X(n, degree + 1);
    for (int i = 0; i < n; ++i) {
        double cur_t = t_vec[i];
        double p = 1.0;
        for (int k = 0; k <= degree; ++k) {
            X(i, k) = p;
            p *= cur_t;
        }
    }
    return X;
}

inline arma::vec poly_multiply_by_linear(const arma::vec& coeffs, double shift) {
    arma::vec out(coeffs.n_elem + 1, arma::fill::zeros);
    for (size_t k = 0; k < coeffs.n_elem; ++k) {
        out[k] += -shift * coeffs[k];
        out[k + 1] += coeffs[k];
    }
    return out;
}

inline arma::mat orthogonal_poly_basis_from_coefs(
    const arma::vec& t_vec,
    int degree,
    const arma::vec& alpha,
    const arma::vec& norm2
) {
    int n = static_cast<int>(t_vec.n_elem);
    arma::mat Z(n, degree + 1, arma::fill::ones);
    if (degree >= 1) {
        Z.col(1) = t_vec - alpha[0];
    }
    for (int i = 2; i <= degree; ++i) {
        Z.col(i) = (t_vec - alpha[i - 1]) % Z.col(i - 1) -
                   (norm2[i] / norm2[i - 1]) * Z.col(i - 2);
    }
    for (int i = 0; i <= degree; ++i) {
        Z.col(i) /= std::sqrt(norm2[i + 1]);
    }
    return Z.cols(1, degree);
}

inline arma::mat orthogonal_basis_raw_coefficients(
    int degree,
    const arma::vec& alpha,
    const arma::vec& norm2
) {
    arma::mat raw(degree + 1, degree + 1, arma::fill::zeros);
    raw(0, 0) = 1.0;
    if (degree >= 1) {
        arma::vec p1(2);
        p1[0] = -alpha[0];
        p1[1] = 1.0;
        raw.submat(0, 1, 1, 1) = p1 / std::sqrt(norm2[2]);
    }
    arma::vec prev2(1);
    prev2[0] = 1.0;
    arma::vec prev1;
    if (degree >= 1) {
        prev1.set_size(2);
        prev1[0] = -alpha[0];
        prev1[1] = 1.0;
    }
    for (int i = 2; i <= degree; ++i) {
        arma::vec term = poly_multiply_by_linear(prev1, alpha[i - 1]);
        arma::vec padded_prev2(term.n_elem, arma::fill::zeros);
        padded_prev2.subvec(0, prev2.n_elem - 1) = prev2;
        arma::vec current = term - (norm2[i] / norm2[i - 1]) * padded_prev2;
        raw.submat(0, i, i, i) = current / std::sqrt(norm2[i + 1]);
        prev2 = prev1;
        prev1 = current;
    }
    return raw;
}

inline bool modified_gram_schmidt_qr(
    const arma::mat& X,
    arma::mat& Q,
    arma::mat& R
) {
    const int n = static_cast<int>(X.n_rows);
    const int p = static_cast<int>(X.n_cols);
    Q.zeros(n, p);
    R.zeros(p, p);

    for (int k = 0; k < p; ++k) {
        arma::vec v = X.col(k);
        for (int j = 0; j < k; ++j) {
            R(j, k) = arma::dot(Q.col(j), v);
            v -= R(j, k) * Q.col(j);
        }

        double norm = std::sqrt(arma::dot(v, v));
        if (norm <= 1e-12) {
            return false;
        }

        R(k, k) = norm;
        Q.col(k) = v / norm;
    }

    return true;
}

inline arma::mat build_orthogonal_poly_design(
    const arma::vec& t_vec,
    int degree,
    arma::vec& alpha,
    arma::vec& norm2
) {
    int n = static_cast<int>(t_vec.n_elem);
    if (degree < 1) {
        return arma::mat(n, 0);
    }
    arma::vec unique_t = arma::unique(t_vec);
    if (static_cast<size_t>(degree) >= unique_t.n_elem) {
        throw std::invalid_argument("orthogonal polynomial degree must be less than number of unique time points.");
    }

    double xbar = arma::mean(t_vec);
    arma::vec centered = t_vec - xbar;
    arma::mat X = build_raw_poly_design(centered, degree);

    arma::mat Q;
    arma::mat R;
    bool ok = modified_gram_schmidt_qr(X, Q, R);
    if (!ok) {
        throw std::invalid_argument("orthogonal polynomial degree must be less than number of unique time points.");
    }

    arma::mat Z = Q * arma::diagmat(R.diag());
    arma::vec z_norm2 = arma::sum(arma::square(Z), 0).t();

    alpha.set_size(degree);
    for (int i = 0; i < degree; ++i) {
        alpha[i] = arma::dot(centered, arma::square(Z.col(i))) / z_norm2[i] + xbar;
    }

    norm2.set_size(degree + 2);
    norm2[0] = 1.0;
    for (int i = 0; i <= degree; ++i) {
        norm2[i + 1] = z_norm2[i];
    }

    return orthogonal_poly_basis_from_coefs(t_vec, degree, alpha, norm2);
}

inline arma::mat build_poly_design(
    const arma::vec& t_vec,
    int degree,
    const std::string& basis,
    arma::vec& alpha,
    arma::vec& norm2,
    arma::mat& basis_raw_coeffs
) {
    int n = static_cast<int>(t_vec.n_elem);
    arma::mat X(n, degree + 1);
    X.col(0).ones();
    if (basis == "orthogonal") {
        arma::mat Z = build_orthogonal_poly_design(t_vec, degree, alpha, norm2);
        X.cols(1, degree) = Z;
        basis_raw_coeffs = orthogonal_basis_raw_coefficients(degree, alpha, norm2);
    } else if (basis == "raw") {
        X = build_raw_poly_design(t_vec, degree);
        alpha.reset();
        norm2.reset();
        basis_raw_coeffs.eye(degree + 1, degree + 1);
    } else {
        throw std::invalid_argument("basis must be 'orthogonal' or 'raw'.");
    }
    return X;
}

/** @brief Differential geometry profile along a trajectory curve. */
struct TrajectoryDerivatives {
    arma::vec t;
    arma::vec vx;           ///< dx/dt
    arma::vec vy;           ///< dy/dt
    arma::vec speed;        ///< sqrt(vx^2 + vy^2)
    arma::vec ax;           ///< d^2x/dt^2
    arma::vec ay;           ///< d^2y/dt^2
    arma::vec heading_rate; ///< omega = vx*ay - vy*ax (angular turning velocity)
    arma::vec curvature;    ///< kappa = omega / max(speed^3, eps)
};

/** @brief Evaluate a 1D polynomial sum(c_k * t^k) at a single time t. */
inline double eval_poly_scalar(const arma::vec& coeffs, double t) {
    double val = 0.0;
    double t_pow = 1.0;
    for (size_t k = 0; k < coeffs.n_elem; ++k) {
        val += coeffs[k] * t_pow;
        t_pow *= t;
    }
    return val;
}

/** @brief Evaluate polynomial 1st derivative sum(k * c_k * t^(k-1)) at a single time t. */
inline double eval_poly_deriv1_scalar(const arma::vec& coeffs, double t) {
    double val = 0.0;
    double t_pow = 1.0;
    for (size_t k = 1; k < coeffs.n_elem; ++k) {
        val += static_cast<double>(k) * coeffs[k] * t_pow;
        t_pow *= t;
    }
    return val;
}

/** @brief Evaluate polynomial 2nd derivative sum(k*(k-1) * c_k * t^(k-2)) at a single time t. */
inline double eval_poly_deriv2_scalar(const arma::vec& coeffs, double t) {
    double val = 0.0;
    double t_pow = 1.0;
    for (size_t k = 2; k < coeffs.n_elem; ++k) {
        val += static_cast<double>(k * (k - 1)) * coeffs[k] * t_pow;
        t_pow *= t;
    }
    return val;
}

/** @brief Evaluate 2D parametric curve (x(t), y(t)) at multiple time points. */
inline arma::mat eval_poly_curve(
    const arma::vec& coeffs_x,
    const arma::vec& coeffs_y,
    const arma::vec& t_eval
) {
    size_t n = t_eval.n_elem;
    arma::mat out(n, 2);
    for (size_t i = 0; i < n; ++i) {
        double t = t_eval[i];
        out(i, 0) = eval_poly_scalar(coeffs_x, t);
        out(i, 1) = eval_poly_scalar(coeffs_y, t);
    }
    return out;
}

/** @brief Fit a 2D polynomial curve with automatic LOOCV-2D or AIC degree selection.
 *
 *  Given coordinates (x_i, y_i) and times t_i in [0, 1], constructs polynomial models
 *  for degrees d in [1, max_degree]. LOOCV error is computed via the closed-form hat-matrix formula:
 *  CV(d) = sum( (r_{x,i} / (1 - h_{ii}))^2 + (r_{y,i} / (1 - h_{ii}))^2 )
 *
 *  @param[in] points      n × 2 matrix of trajectory coordinates.
 *  @param[in] t           Vector of length n of strictly increasing times (default [0, 1] linearly spaced).
 *  @param[in] max_degree  Maximum polynomial degree to evaluate (default 3).
 *  @param[in] fixed_degree When >= 1, uses this exact degree without search.
 *  @param[in] criterion   "loocv" (default) or "aic".
 *
 *  @returns PolyCurveFit containing optimal degree, coefficients, errors, and fitted points.
 */
inline PolyCurveFit fit_poly_loocv2d(
    const arma::mat& points,
    const arma::vec& t = arma::vec(),
    int max_degree = 3,
    int fixed_degree = 0,
    const std::string& criterion = "loocv",
    const std::string& basis = "orthogonal"
) {
    int n = static_cast<int>(points.n_rows);
    if (n < 2 || points.n_cols < 2) {
        throw std::invalid_argument("fit_poly_loocv2d requires at least 2 points and 2 columns.");
    }

    // Default normalized time t in [0, 1]
    arma::vec t_vec;
    if (t.n_elem == static_cast<size_t>(n)) {
        t_vec = t;
    } else {
        t_vec = arma::linspace<arma::vec>(0.0, 1.0, n);
    }

    arma::vec x = points.col(0);
    arma::vec y = points.col(1);

    int deg_upper = std::min(max_degree, std::max(1, n - 2));
    if (fixed_degree >= 1) {
        deg_upper = std::min(fixed_degree, n - 1);
    }

    int best_deg = 1;
    double best_score = std::numeric_limits<double>::infinity();
    PolyCurveFit best_fit;

    int deg_start = (fixed_degree >= 1) ? fixed_degree : 1;

    for (int d = deg_start; d <= deg_upper; ++d) {
        int n_params = d + 1;
        if (n_params > n) break;

        arma::vec basis_alpha;
        arma::vec basis_norm2;
        arma::mat basis_raw_coeffs;
        arma::mat X = build_poly_design(t_vec, d, basis, basis_alpha, basis_norm2, basis_raw_coeffs);

        // Solve least squares: beta = (X^T X)^(-1) X^T y
        arma::mat XtX = X.t() * X;
        arma::mat XtX_inv;
        bool inv_ok = arma::inv(XtX_inv, XtX);
        if (!inv_ok) {
            // Pseudo-inverse fallback
            XtX_inv = arma::pinv(XtX);
        }

        arma::vec beta_x = XtX_inv * (X.t() * x);
        arma::vec beta_y = XtX_inv * (X.t() * y);

        arma::vec pred_x = X * beta_x;
        arma::vec pred_y = X * beta_y;

        arma::vec rx = x - pred_x;
        arma::vec ry = y - pred_y;

        // Hat matrix diagonal: h_ii = x_i^T (X^T X)^(-1) x_i
        arma::vec h(n);
        for (int i = 0; i < n; ++i) {
            arma::rowvec xi = X.row(i);
            h[i] = arma::as_scalar(xi * XtX_inv * xi.t());
            h[i] = std::min(0.9999, std::max(0.0, h[i]));
        }

        // 2D LOOCV error
        double cv_err = 0.0;
        for (int i = 0; i < n; ++i) {
            double denom = 1.0 - h[i];
            double rxi = rx[i] / denom;
            double ryi = ry[i] / denom;
            cv_err += (rxi * rxi + ryi * ryi);
        }

        // Bivariate AIC
        double rss = arma::dot(rx, rx) + arma::dot(ry, ry);
        double aic = static_cast<double>(n) * std::log(std::max(1e-12, rss / (2.0 * n))) + 2.0 * (2.0 * n_params);

        double score = (criterion == "aic") ? aic : cv_err;

        if (score < best_score || d == deg_start) {
            best_score = score;
            best_deg = d;
            best_fit.degree = d;
            best_fit.coeffs_x = basis_raw_coeffs * beta_x;
            best_fit.coeffs_y = basis_raw_coeffs * beta_y;
            best_fit.basis = basis;
            best_fit.basis_alpha = basis_alpha;
            best_fit.basis_norm2 = basis_norm2;
            best_fit.cv_error = cv_err;
            best_fit.aic = aic;
            best_fit.t = t_vec;
            best_fit.fitted = arma::join_rows(pred_x, pred_y);
        }
    }

    return best_fit;
}

/** @brief Compute trajectory velocities, speeds, accelerations, heading turn rates, and curvatures. */
inline TrajectoryDerivatives eval_trajectory_derivatives(
    const arma::vec& coeffs_x,
    const arma::vec& coeffs_y,
    const arma::vec& t_eval
) {
    size_t n = t_eval.n_elem;
    TrajectoryDerivatives d;
    d.t = t_eval;
    d.vx.set_size(n);
    d.vy.set_size(n);
    d.speed.set_size(n);
    d.ax.set_size(n);
    d.ay.set_size(n);
    d.heading_rate.set_size(n);
    d.curvature.set_size(n);

    for (size_t i = 0; i < n; ++i) {
        double t = t_eval[i];
        double vx = eval_poly_deriv1_scalar(coeffs_x, t);
        double vy = eval_poly_deriv1_scalar(coeffs_y, t);
        double ax = eval_poly_deriv2_scalar(coeffs_x, t);
        double ay = eval_poly_deriv2_scalar(coeffs_y, t);

        double spd = std::sqrt(vx * vx + vy * vy);
        double omega = vx * ay - vy * ax; // angular velocity / heading rate
        double spd3 = spd * spd * spd;
        double kappa = (spd3 > 1e-12) ? (omega / spd3) : 0.0;

        d.vx[i] = vx;
        d.vy[i] = vy;
        d.speed[i] = spd;
        d.ax[i] = ax;
        d.ay[i] = ay;
        d.heading_rate[i] = omega;
        d.curvature[i] = kappa;
    }

    return d;
}

} // namespace qe

#endif // LIBQE_TRAJECTORY_HPP

// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>
#include <libqe/libqe.hpp>

using namespace Rcpp;

// =============================================================================
// Adjacency utilities
// =============================================================================

//' Number of upper-triangle pairs for n codes
//' @param n Number of codes
//' @return Integer: n*(n-1)/2
//' @export
// [[Rcpp::export]]
int choose_two(int n) {
    return qe::choose_two(n);
}

//' Upper-triangle index pairs
//' @param len Number of codes (side length of square matrix)
//' @param row -1 = both rows, 0 = row indices only, 1 = col indices only
//' @export
// [[Rcpp::export]]
arma::umat connection_indices(int len, int row = -1) {
    return qe::connection_indices(len, row);
}

//' Pairwise products → upper-triangle vector
//' @param v Numeric vector of code values
//' @export
// [[Rcpp::export]]
arma::rowvec code_connections(arma::mat v) {
    return qe::code_connections(v);
}

//' Fold a directed (n*n) vector into an undirected upper-triangle vector
//' @param v Numeric vector of length n*n
//' @export
// [[Rcpp::export]]
arma::rowvec fold_directed_network(arma::vec v) {
    return qe::fold_directed_network(v);
}

//' Flatten an adjacency matrix to a connection vector
//' @param x Numeric matrix
//' @param full TRUE = full n*n (directed); FALSE = upper triangle (undirected)
//' @export
// [[Rcpp::export]]
arma::rowvec network_to_vector(arma::mat x, bool full = true) {
    return qe::network_to_vector(x, full);
}

//' Code-name pairs for upper-triangle positions ("A & B")
//' @param v Character vector of code names
//' @export
// [[Rcpp::export]]
std::vector<std::string> connection_names(std::vector<std::string> v) {
    return qe::connection_names(v);
}

// =============================================================================
// Normalization
// =============================================================================

//' Row-wise L2 (sphere) normalization
//' @param m Numeric matrix
//' @export
// [[Rcpp::export]]
arma::mat normalize_networks(arma::mat m) {
    return qe::normalize_networks(m);
}

//' Max-norm scaling (divide all rows by the largest row L2 norm)
//' @param m Numeric matrix
//' @export
// [[Rcpp::export]]
arma::mat scale_networks(arma::mat m) {
    return qe::scale_networks(m);
}

// =============================================================================
// Modeling
// =============================================================================

//' Center points (subtract column means)
//' @param values Numeric matrix
//' @export
// [[Rcpp::export]]
arma::mat center_points(arma::mat values) {
    return qe::center_points(values);
}

//' Confidence interval for the mean of a group of ENA unit points
//'
//' For each dimension computes a t-based CI:
//' \code{mean ± t(α/2, n-1) × (SD / sqrt(n))} where \code{α = 1 - conf_level}.
//'
//' @param points     Numeric matrix (units x dims) — one row per unit in the group
//' @param conf_level Confidence level (default 0.95)
//' @return Numeric matrix (n_dims x 3): columns are [mean, ci_lower, ci_upper].
//'   When \code{nrow(points) == 1} the CI bounds are \code{±Inf}.
//' @export
// [[Rcpp::export]]
arma::mat mean_ci(arma::mat points, double conf_level = 0.95) {
    return qe::mean_ci(points, conf_level);
}

//' Outlier interval based on IQR (Tukey fence) for a group of ENA unit points
//'
//' For each dimension: \code{lower = -iqr_factor * IQR}, \code{upper = +iqr_factor * IQR}.
//' Equivalent to rENA's:
//' \preformatted{
//'   oi <- c(IQR(pts[,1]), ...) * iqr_factor
//'   matrix(rep(oi, 2), ncol = n_dims, byrow = TRUE) * c(-1, 1)
//' }
//' IQR uses R's default \code{type = 7} quantile.
//'
//' @param points     Numeric matrix (units x dims) — one row per unit in the group
//' @param iqr_factor Multiplier applied to the IQR (default 1.5, the Tukey fence)
//' @return Numeric matrix (n_dims x 2): columns are [lower, upper].
//' @export
// [[Rcpp::export]]
arma::mat outlier_ci(arma::mat points, double iqr_factor = 1.5) {
    return qe::outlier_ci(points, iqr_factor);
}

//' Two-group comparison statistics
//'
//' Computes per-dimension parametric and non-parametric statistics for two
//' groups of ENA unit points.  Equivalent to \code{group.stats()} in rENA-api,
//' moved to the C++ layer so all language bindings share a single implementation.
//'
//' \strong{Parametric (Welch two-sample t-test):}
//' \itemize{
//'   \item \code{t}        — Welch t-statistics
//'   \item \code{parameter}— Welch–Satterthwaite degrees of freedom
//'   \item \code{pvalue}   — two-tailed p-values
//'   \item \code{effect}   — Cohen's d (pooled-SD: \code{(μ1−μ2)/pooled_sd})
//'   \item \code{mean}     — 2×n_dims matrix, row 1 = group1, row 2 = group2
//'   \item \code{std.dev}  — 2×n_dims matrix of sample standard deviations
//' }
//'
//' \strong{Non-parametric (Wilcoxon rank-sum):}
//' \itemize{
//'   \item \code{U}        — rank-sum U for group 1 (= R's \code{wilcox.test} W)
//'   \item \code{pvalue}   — two-tailed p-values (normal approx, tie + continuity correction)
//'   \item \code{effect}   — rank-biserial: \code{1 − 2·U / (n1·n2)}
//'   \item \code{median}   — 2×n_dims matrix of medians
//' }
//'
//' Entries are \code{NA} when a statistic is undefined (e.g. n < 2 for
//' parametric tests).
//'
//' @param g1 Numeric matrix (n1 × n_dims) — unit points for group 1
//' @param g2 Numeric matrix (n2 × n_dims) — unit points for group 2
//' @return Named list with \code{N}, \code{parametric}, and \code{nonparametric}
//'   sub-lists, matching the structure of rENA-api's \code{group.stats()}.
//' @export
// [[Rcpp::export]]
List group_stats(arma::mat g1, arma::mat g2) {
    qe::GroupStats s = qe::group_stats(g1, g2);
    return List::create(
        _("N") = IntegerVector::create(s.n1, s.n2),
        _("parametric") = List::create(
            _("t")         = s.t,
            _("parameter") = s.df,
            _("pvalue")    = s.pvalue_t,
            _("effect")    = s.cohens_d,
            _("mean")      = s.means,
            _("std.dev")   = s.sds
        ),
        _("nonparametric") = List::create(
            _("U")      = s.U,
            _("pvalue") = s.pvalue_u,
            _("effect") = s.effect_r,
            _("median") = s.medians
        )
    );
}

//' Pearson correlation with CI between ENA points and centroids
//' @param points  Numeric matrix (units x dims)
//' @param centroids Numeric matrix (units x dims)
//' @param conf_level Confidence level (default 0.95)
//' @export
// [[Rcpp::export]]
arma::mat ena_correlation(arma::mat points, arma::mat centroids,
                              double conf_level = 0.95) {
    return qe::ena_correlation(points, centroids, conf_level);
}

//' Least-squares node positions for undirected ENA
//' @param adj_mats Numeric matrix of line weights (units x connections)
//' @param t        Numeric matrix of rotated points (units x dims)
//' @param num_dims Number of dimensions
//' @return List with nodes, centroids, weights, points
//' @export
// [[Rcpp::export]]
List node_positions(arma::mat adj_mats, arma::mat t, int num_dims) {
    if (!adj_mats.is_finite() || !t.is_finite())
        stop("node_positions: input matrices must not contain NaN or Inf - "
             "filter or impute rows with non-finite values before calling");
    qe::NodePositions r = qe::node_positions(adj_mats, t, num_dims);
    return List::create(
        _("nodes")     = r.nodes,
        _("centroids") = r.centroids,
        _("weights")   = r.weights,
        _("points")    = r.points
    );
}

//' Least-squares node positions for directed ENA
//' @param line_weights Numeric matrix (units x connections)
//' @param points       Numeric matrix of rotated points (units x dims)
//' @param num_dims     Number of dimensions
//' @return List with nodes, centroids, weights, points
//' @export
// [[Rcpp::export]]
List directed_node_positions(arma::mat line_weights, arma::mat points,
                                 int num_dims) {
    if (!line_weights.is_finite() || !points.is_finite())
        stop("directed_node_positions: input matrices must not contain NaN or Inf - "
             "filter or impute rows with non-finite values before calling");
    qe::NodePositions r = qe::directed_node_positions(line_weights, points, num_dims);
    return List::create(
        _("nodes")     = r.nodes,
        _("centroids") = r.centroids,
        _("weights")   = r.weights,
        _("points")    = r.points
    );
}

//' Directed node positions with paired ground+response rows combined
//' @param line_weights Numeric matrix (units x connections)
//' @param points       Numeric matrix of rotated points (units x dims)
//' @param num_dims     Number of dimensions
//' @return List with nodes, centroids, weights, points
//' @export
// [[Rcpp::export]]
List directed_node_positions_combine_pairs(arma::mat line_weights,
                                               arma::mat points,
                                               int num_dims) {
    if (!line_weights.is_finite() || !points.is_finite())
        stop("directed_node_positions_combine_pairs: input matrices must not contain NaN or Inf - "
             "filter or impute rows with non-finite values before calling");
    qe::NodePositions r = qe::directed_node_positions(
        line_weights, points, num_dims, true);
    return List::create(
        _("nodes")     = r.nodes,
        _("centroids") = r.centroids,
        _("weights")   = r.weights,
        _("points")    = r.points
    );
}

// =============================================================================
// Accumulation
// =============================================================================

//' Core connection matrix for one ground+response pair
//' @param ground          Numeric row vector of ground (context) code values
//' @param response        Numeric row vector of response code values
//' @param response_weight Scalar weight applied to the response self-connection
//' @param ordered         TRUE = directed; FALSE = undirected
//' @export
// [[Rcpp::export]]
arma::mat connection_matrix(arma::rowvec ground, arma::rowvec response,
                                         double response_weight = 1.0,
                                         bool ordered = true) {
    return qe::connection_matrix(ground, response, response_weight, ordered);
}

//' Traditional stanza-window accumulation (rENA model)
//'
//' For each row k in a single conversation's code matrix, accumulates
//' connections over a back/forward window.
//'
//' \strong{Undirected} (\code{ordered = FALSE}, default): upper-triangle
//' co-occurrences with back/forward-reference corrections.
//' Returns \code{n_rows × choose_two(n_codes)}.
//'
//' \strong{Directed} (\code{ordered = TRUE}): focal row k as response,
//' sum of prior window rows as ground, using the directed connection formula.
//' \code{window_forward} is ignored. Returns \code{n_rows × n_codes²}.
//'
//' @param codes          Numeric matrix (rows = lines, cols = codes) for ONE conversation
//' @param window_back    Number of prior lines in window (default 1); use .Machine$integer.max for Inf
//' @param window_forward Number of subsequent lines (default 0; ignored when ordered = TRUE)
//' @param binary         If TRUE, binarise non-zero connection counts
//' @param ordered        If TRUE, return directed n² output; FALSE (default) returns upper-tri
//' @return Undirected: (n_rows × choose_two(n_codes)); directed: (n_rows × n_codes²)
//' @export
// [[Rcpp::export]]
arma::mat accumulate_stanza(arma::mat codes,
                            int window_back    = 1,
                            int window_forward = 0,
                            bool binary        = true,
                            bool ordered       = false) {
    return qe::accumulate_stanza(codes, window_back, window_forward, binary, ordered);
}

//' Compute a column-major linear index into a multi-dimensional array
//'
//' Equivalent to tma's flat_index(). Throws if lengths of `indices`
//' and `dims` differ.
//'
//' @param indices 0-based integer vector of per-dimension indices
//' @param dims    Integer vector of array dimensions
//' @return Scalar integer linear index
//' @export
// [[Rcpp::export]]
int flat_index(std::vector<int> indices, std::vector<int> dims) {
    return qe::flat_index(indices, dims);
}

//' Per-row upper-triangle co-occurrence matrix
//'
//' For each row in \code{codes}, computes the pairwise code products
//' (upper-triangle) and optionally binarizes. Output has
//' \code{choose(n_codes, 2)} columns.
//'
//' @param codes  Numeric matrix (rows = observations, cols = codes)
//' @param binary If TRUE, binarise non-zero co-occurrences (default TRUE)
//' @return Numeric matrix (n_rows x choose_two(n_codes))
//' @export
// [[Rcpp::export]]
arma::mat row_connections(arma::mat codes, bool binary = true) {
    return qe::row_connections(codes, binary);
}

//' Rolling backward window sum of a code matrix
//'
//' For each row k, sums the raw code values over rows
//' \code{[max(0, k - window_size + 1), k]}. Returns a matrix of the same
//' shape as \code{codes}. \code{window_size <= 0} is treated as 1
//' (current row only).
//'
//' @param codes       Numeric matrix (rows = observations, cols = codes)
//' @param window_size Number of rows to look back (default 1)
//' @return Numeric matrix (same dimensions as \code{codes})
//' @export
// [[Rcpp::export]]
arma::mat rolling_window_sum(arma::mat codes, int window_size = 1) {
    return qe::rolling_window_sum(codes, window_size);
}

//' Ground/response accumulation for one unit (tma model)
//'
//' @param codes      Numeric matrix for the full context (n_rows x n_codes)
//' @param unit_rows  0-based integer vector of rows belonging to this unit
//' @param decay_fn   R function mapping a numeric distance vector to weights
//' @param ordered    TRUE = directed; FALSE = undirected (upper-tri)
//' @return Numeric vector of connection counts
//' @export
// [[Rcpp::export]]
arma::rowvec accumulate_unit(arma::mat codes, std::vector<int> unit_rows,
                                  Function decay_fn, bool ordered = false) {
    auto cpp_decay = [&](arma::vec distances) -> arma::vec {
        NumericVector d = wrap(distances);
        NumericVector w = decay_fn(d);
        return as<arma::vec>(w);
    };
    return qe::accumulate_unit(codes, unit_rows, cpp_decay, ordered);
}

//' Ground/response accumulation for one unit — returns unit vector and per-row matrix
//'
//' Equivalent to tma's accumulate_network() but without R-env-var setup.
//' The decay function receives a distance vector (response=0, older rows
//' have larger values) and returns a weight vector of the same length.
//'
//' @param codes      Numeric matrix (n_rows x n_codes)
//' @param unit_rows  0-based integer vector of rows belonging to this unit
//' @param decay_fn   R function(distances) -> weights
//' @param ordered    TRUE = directed (full p^2); FALSE = undirected (upper-tri)
//' @return List with `networks` (vector) and `row_networks` (matrix)
//' @export
// [[Rcpp::export]]
List accumulate_unit_with_rows(arma::mat codes, std::vector<int> unit_rows,
                                   Function decay_fn, bool ordered = false) {
    auto cpp_decay = [&](int unit_row, arma::uvec ground_indices) -> arma::vec {
        arma::vec dists(ground_indices.n_elem);
        for (arma::uword k = 0; k < ground_indices.n_elem; ++k)
            dists[k] = static_cast<double>(unit_row - ground_indices[k]);
        return as<arma::vec>(wrap(decay_fn(wrap(dists))));
    };
    qe::UnitNetworks r = qe::accumulate_unit_with_rows(codes, unit_rows, cpp_decay, ordered);
    return List::create(
        _("networks")     = r.networks,
        _("row_networks") = r.row_networks
    );
}

//' Tensor-based multi-modal accumulation for one unit (tma model)
//'
//' Pure-C++ port of tma's apply_tensor().  Accepts 0-based index vectors
//' for sender/receiver/mode dimensions and a pre-converted integer context
//' lookup matrix.
//'
//' @param tensor         Numeric vector (column-major flat tensor)
//' @param dims           Integer vector of tensor dimensions
//' @param dims_sender    0-based sender axis indices
//' @param dims_receiver  0-based receiver axis indices
//' @param dims_mode      0-based mode axis indices
//' @param context_lookup Integer matrix (n_context_rows x n_factors), 0-based
//' @param unit_rows      0-based response-row indices for this unit
//' @param codes          Numeric matrix (n_context_rows x n_codes)
//' @param times          Numeric vector of timestamps per context row
//' @param ordered        TRUE = directed; FALSE = undirected
//' @return List with `connection_counts` (vector) and `row_connection_counts` (matrix)
//' @export
// [[Rcpp::export]]
List apply_tensor(arma::vec tensor,
                     std::vector<int> dims,
                     std::vector<int> dims_sender,
                     std::vector<int> dims_receiver,
                     std::vector<int> dims_mode,
                     arma::imat context_lookup,
                     std::vector<int> unit_rows,
                     arma::mat codes,
                     arma::vec times,
                     bool ordered = true) {
    qe::TensorNetworks r = qe::apply_tensor_unit(
        tensor, dims, dims_sender, dims_receiver, dims_mode,
        context_lookup, unit_rows, codes, times, ordered);
    return List::create(
        _("connection_counts")     = r.connection_counts,
        _("row_connection_counts") = r.row_connection_counts
    );
}

//' Apply the TMA window/weight tensor, additionally returning per-response-row
//' window membership.
//'
//' Identical to \code{apply_tensor()} but also returns, for each response
//' (reference) row, the in-window ground rows and their resolved window sizes.
//' Powers the webtool Data View per-modality window-span hover. Kept as a
//' separate export so \code{apply_tensor()}'s ABI (and tma's direct .Call to it)
//' is unchanged.
//'
//' @param tensor         Numeric vector (column-major flat tensor)
//' @param dims           Integer vector of tensor dimensions
//' @param dims_sender    0-based sender axis indices
//' @param dims_receiver  0-based receiver axis indices
//' @param dims_mode      0-based mode axis indices
//' @param context_lookup Integer matrix (n_context_rows x n_factors), 0-based
//' @param unit_rows      0-based response-row indices for this unit
//' @param codes          Numeric matrix (n_context_rows x n_codes)
//' @param times          Numeric vector of timestamps per context row
//' @param ordered        TRUE = directed; FALSE = undirected
//' @return List with \code{connection_counts}, \code{row_connection_counts},
//'   \code{row_window_members} (list of 1-based ground-row index vectors, in
//'   \code{unit_rows} order) and \code{row_window_wins} (parallel resolved
//'   window sizes).
//' @export
// [[Rcpp::export]]
List apply_tensor_members(arma::vec tensor,
                     std::vector<int> dims,
                     std::vector<int> dims_sender,
                     std::vector<int> dims_receiver,
                     std::vector<int> dims_mode,
                     arma::imat context_lookup,
                     std::vector<int> unit_rows,
                     arma::mat codes,
                     arma::vec times,
                     bool ordered = true) {
    qe::TensorNetworks r = qe::apply_tensor_unit(
        tensor, dims, dims_sender, dims_receiver, dims_mode,
        context_lookup, unit_rows, codes, times, ordered, /*return_members=*/true);
    List members(r.row_window_members.size());
    List wins(r.row_window_wins.size());
    for (size_t i = 0; i < r.row_window_members.size(); ++i) {
        IntegerVector iv(r.row_window_members[i].begin(), r.row_window_members[i].end());
        for (R_xlen_t k = 0; k < iv.size(); ++k) iv[k] += 1;   // 0-based -> 1-based
        members[i] = iv;
        wins[i] = NumericVector(r.row_window_wins[i].begin(), r.row_window_wins[i].end());
    }
    return List::create(
        _("connection_counts")     = r.connection_counts,
        _("row_connection_counts") = r.row_connection_counts,
        _("row_window_members")    = members,
        _("row_window_wins")       = wins
    );
}

// =============================================================================
// Rotation
// =============================================================================

// Pack a RotationResult into the list shape that matches rENA's ENARotationSet
// payload (rotation + eigenvalues + column names). Column names are attached
// to the rotation matrix as dimnames so downstream R code can index by them.
// Eigenvalues are returned as a plain numeric vector (not an Nx1 matrix) to
// match what rENA's `pcaResults$sdev^2` produces.
static List pack_rotation_result(const qe::RotationResult& r) {
    NumericMatrix rotation = wrap(r.rotation);
    CharacterVector col_names(r.column_names.begin(), r.column_names.end());
    rotation.attr("dimnames") = List::create(R_NilValue, col_names);
    NumericVector eigenvalues(r.eigenvalues.begin(), r.eigenvalues.end());
    return List::create(
        _("rotation")     = rotation,
        _("eigenvalues")  = eigenvalues,
        _("column_names") = col_names
    );
}

//' SVD rotation (matches prcomp(retx=F, scale=F, center=F, tol=0))
//'
//' Caller is responsible for centering upstream. Eigenvalues are stored as
//' \code{sdev^2} (variance) to match rENA's \code{ena.svd}.
//'
//' @param points Numeric matrix (n_units x n_dims)
//' @return List with \code{rotation} (n_dims x n_dims), \code{eigenvalues}
//'   (length n_dims, = sdev^2), and \code{column_names} ("SVD1", "SVD2", ...)
//' @export
// [[Rcpp::export]]
List ena_svd(arma::mat points) {
    if (!points.is_finite())
        stop("ena_svd: input matrix must not contain NaN or Inf - "
             "filter or impute rows with non-finite values before calling");
    return pack_rotation_result(qe::ena_svd(points));
}

//' Project a matrix onto the hyperplane orthogonal to a unit-norm axis
//'
//' Computes \code{data - (data \%*\% axis) \%*\% t(axis)}. The caller is
//' responsible for ensuring \code{axis} is unit-norm.
//'
//' @param data Numeric matrix (n_units x n_dims)
//' @param axis Numeric vector of length n_dims, unit-norm
//' @return Numeric matrix of the same shape as \code{data}
//' @export
// [[Rcpp::export]]
arma::mat deflate(arma::mat data, arma::vec axis) {
    return qe::deflate(data, axis);
}

//' Orthogonal SVD — orthonormalize named axes via QR, fill the rest from SVD
//'
//' Mirrors rENA's \code{orthogonal_svd()} in \code{ena.rotate.by.mean.R}:
//' the named axes in the output are the orthonormalized Q columns, not the
//' original \code{weights} columns. Use \code{complete_rotation} to keep
//' the named axes verbatim.
//'
//' @param data         Numeric matrix (n_units x n_dims)
//' @param weights      Numeric matrix (n_dims x k); columns are the named axes
//' @param named_labels Character vector of length k
//' @return List with \code{rotation}, \code{eigenvalues}, \code{column_names}
//' @export
// [[Rcpp::export]]
List orthogonal_svd(arma::mat data,
                        arma::mat weights,
                        std::vector<std::string> named_labels) {
    return pack_rotation_result(qe::orthogonal_svd(data, weights, named_labels));
}

//' Complete a rotation — keep named axes verbatim, fill remainder from SVD
//'
//' Mirrors the tail of \code{ena.rotate.by.generalized}: the named axes
//' appear in the output exactly as provided, and the trailing columns come
//' from an SVD of the data deflated by all named axes.
//'
//' @param data         Numeric matrix (n_units x n_dims)
//' @param named_axes   Numeric matrix (n_dims x k); columns must be unit-norm
//' @param named_labels Character vector of length k
//' @return List with \code{rotation}, \code{eigenvalues}, \code{column_names}
//' @export
// [[Rcpp::export]]
List complete_rotation(arma::mat data,
                           arma::mat named_axes,
                           std::vector<std::string> named_labels) {
    return pack_rotation_result(qe::complete_rotation(data, named_axes, named_labels));
}

//' Means rotation
//'
//' For each group pair, computes a normalized mean-difference axis on the
//' progressively-deflated data and finishes with \code{orthogonal_svd}.
//' The input is column-centered first, matching rENA's
//' \code{scale(data, scale=F, center=T)} at the top of \code{ena.rotate.by.mean}.
//'
//' Each element of \code{group_pairs} is a length-2 list \code{list(a, b)}
//' of 0-based row indices into \code{points}.
//'
//' @param points      Numeric matrix (n_units x n_dims)
//' @param group_pairs List of length k; each element is \code{list(a, b)}
//'   where \code{a} and \code{b} are 0-based integer index vectors
//' @return List with \code{rotation}, \code{eigenvalues}, \code{column_names}
//' @export
// [[Rcpp::export]]
List means_rotation(arma::mat points, List group_pairs) {
    std::vector<qe::GroupPair> pairs;
    pairs.reserve(group_pairs.size());
    for (R_xlen_t i = 0; i < group_pairs.size(); ++i) {
        List pair = group_pairs[i];
        if (pair.size() != 2) {
            stop("group_pairs[[%d]] must be a length-2 list(a, b)",
                 static_cast<int>(i + 1));
        }
        IntegerVector ra = pair[0];
        IntegerVector rb = pair[1];
        arma::uvec a(ra.size());
        arma::uvec b(rb.size());
        for (R_xlen_t j = 0; j < ra.size(); ++j) {
            a(j) = static_cast<arma::uword>(ra[j]);
        }
        for (R_xlen_t j = 0; j < rb.size(); ++j) {
            b(j) = static_cast<arma::uword>(rb[j]);
        }
        pairs.push_back({a, b});
    }
    return pack_rotation_result(qe::means_rotation(points, pairs));
}

// =============================================================================
// Generalized Means Rotation (GMR)
// =============================================================================

//' Generalized Means Rotation
//'
//' Computes a rotation axis that best represents a target variable's
//' contribution to the ENA point space, after controlling for covariates via
//' Lasso (coordinate-descent, k-fold CV).  Mirrors
//' \code{rENA::ena.rotate.by.generalized()}.
//'
//' The caller is responsible for building \code{x_model_matrix} (e.g. via
//' \code{model.matrix()}) and identifying \code{x1_cols} (0-based indices of
//' the target variable's columns, which receive \code{penalty_factor = 0}).
//' Categorical targets should be encoded as 0-based integers.
//'
//' @param V              Numeric matrix (n_units x n_connections) — ENA points.
//' @param x_model_matrix Numeric matrix (n_units x p) — model matrix for x axis.
//' @param x_target       Numeric vector length n — target variable (raw float
//'   for numeric; 0-based integer codes for categorical).
//' @param x1_cols        Integer vector of 0-based column indices in
//'   \code{x_model_matrix} that belong to the target variable (unpenalized).
//' @param x_categorical  Logical — TRUE if target is categorical.
//' @param x_n_groups     Integer — number of distinct groups (categorical only).
//' @param x_subset       Integer vector of 0-based row indices to subset for
//'   the x-axis GMR step (pass \code{integer(0)} to use all rows).
//' @param has_y          Logical — TRUE to compute a second GMR axis.
//' @param y_model_matrix Numeric matrix (n_units x p) — model matrix for y axis
//'   (ignored when \code{has_y = FALSE}).
//' @param y_target       Numeric vector — y-axis target (ignored when
//'   \code{has_y = FALSE}).
//' @param y1_cols        0-based column indices for y target (ignored when
//'   \code{has_y = FALSE}).
//' @param y_categorical  Logical (ignored when \code{has_y = FALSE}).
//' @param y_n_groups     Integer (ignored when \code{has_y = FALSE}).
//' @param n_lambda       Length of the Lasso lambda path (default 50).
//' @param k_folds        Cross-validation folds for lambda selection (default 5).
//' @param lasso_eps      \code{lambda_min = lasso_eps * lambda_max} (default 0.01).
//' @return List with \code{rotation}, \code{eigenvalues}, \code{column_names}.
//' @export
// [[Rcpp::export]]
List generalized_means_rotation(
    arma::mat V,
    arma::mat x_model_matrix,
    arma::vec x_target,
    IntegerVector x1_cols,
    bool x_categorical,
    int  x_n_groups,
    IntegerVector x_subset,
    bool has_y,
    arma::mat y_model_matrix,
    arma::vec y_target,
    IntegerVector y1_cols,
    bool y_categorical,
    int  y_n_groups,
    int  n_lambda  = 50,
    int  k_folds   = 5,
    double lasso_eps = 0.01
) {
    // ── Unpack x1_cols ───────────────────────────────────────────────────────
    arma::uvec x1(x1_cols.size());
    for (R_xlen_t i = 0; i < x1_cols.size(); ++i)
        x1(i) = static_cast<arma::uword>(x1_cols[i]);

    // ── Unpack x_subset (empty IntegerVector → use all rows) ─────────────────
    arma::uvec x_sub;
    if (x_subset.size() > 0) {
        x_sub.set_size(x_subset.size());
        for (R_xlen_t i = 0; i < x_subset.size(); ++i)
            x_sub(i) = static_cast<arma::uword>(x_subset[i]);
    }

    // ── Unpack y1_cols ───────────────────────────────────────────────────────
    arma::uvec y1(y1_cols.size());
    for (R_xlen_t i = 0; i < y1_cols.size(); ++i)
        y1(i) = static_cast<arma::uword>(y1_cols[i]);

    // ── Build params ─────────────────────────────────────────────────────────
    qe::GeneralizedRotationParams p;
    p.x_model_matrix = x_model_matrix;
    p.x_target       = x_target;
    p.x1_cols        = x1;
    p.x_categorical  = x_categorical;
    p.x_n_groups     = static_cast<arma::uword>(x_n_groups);
    p.x_subset       = x_sub;
    p.has_y          = has_y;
    p.y_model_matrix = y_model_matrix;
    p.y_target       = y_target;
    p.y1_cols        = y1;
    p.y_categorical  = y_categorical;
    p.y_n_groups     = static_cast<arma::uword>(y_n_groups);
    p.n_lambda       = n_lambda;
    p.k_folds        = k_folds;
    p.lasso_eps      = lasso_eps;

    return pack_rotation_result(qe::generalized_means_rotation(V, p));
}

// =============================================================================
// Door Temporal Pooling & Smoothing
// =============================================================================

//' Door lookback window pooling for a single block
//' @param block Numeric matrix of connection counts
//' @param lookback_size Window lookback size
//' @param aggregate_mean Logical: TRUE = mean, FALSE = sum
//' @param weighting_linear Logical: TRUE = linear weights, FALSE = equal
//' @param segment_ids Integer vector of segment IDs (0-indexed or 1-indexed)
//' @export
// [[Rcpp::export]]
arma::mat door_lookback_block(
    arma::mat block,
    int lookback_size = 20,
    bool aggregate_mean = false,
    bool weighting_linear = false,
    IntegerVector segment_ids = IntegerVector()
) {
    std::vector<int> segs(segment_ids.begin(), segment_ids.end());
    return qe::lookback_block(block, lookback_size, aggregate_mean, weighting_linear, segs);
}

//' Door EMA smoothing for a single block
//' @param block Numeric matrix of connection counts
//' @param alpha Smoothing factor in (0, 1]
//' @param segment_ids Integer vector of segment IDs
//' @export
// [[Rcpp::export]]
arma::mat door_ema_block(
    arma::mat block,
    double alpha = 0.1,
    IntegerVector segment_ids = IntegerVector()
) {
    std::vector<int> segs(segment_ids.begin(), segment_ids.end());
    return qe::ema_block(block, alpha, segs);
}

//' Door lookback window pooling across units
//' @param conn_counts Numeric matrix of connection counts
//' @param unit_row_indices List of integer vectors of row indices (0-indexed or 1-indexed R converted)
//' @param lookback_sizes Integer vector of lookback sizes per unit
//' @param aggregate_mean Logical: TRUE = mean, FALSE = sum
//' @param weighting_linear Logical: TRUE = linear weights, FALSE = equal
//' @param segment_ids Optional integer vector of segment IDs
//' @export
// [[Rcpp::export]]
arma::mat door_lookback(
    arma::mat conn_counts,
    List unit_row_indices,
    IntegerVector lookback_sizes,
    bool aggregate_mean = false,
    bool weighting_linear = false,
    IntegerVector segment_ids = IntegerVector()
) {
    std::vector<std::vector<int>> unit_rows;
    unit_rows.reserve(unit_row_indices.size());
    for (int i = 0; i < unit_row_indices.size(); ++i) {
        IntegerVector iv = unit_row_indices[i];
        unit_rows.push_back(std::vector<int>(iv.begin(), iv.end()));
    }
    std::vector<int> lbs(lookback_sizes.begin(), lookback_sizes.end());
    std::vector<int> segs(segment_ids.begin(), segment_ids.end());

    return qe::apply_door_lookback(conn_counts, unit_rows, lbs, aggregate_mean, weighting_linear, segs);
}

//' Door EMA smoothing across units
//' @param conn_counts Numeric matrix of connection counts
//' @param unit_row_indices List of integer vectors of row indices
//' @param alpha Smoothing factor
//' @param segment_ids Optional integer vector of segment IDs
//' @export
// [[Rcpp::export]]
arma::mat door_ema(
    arma::mat conn_counts,
    List unit_row_indices,
    double alpha = 0.1,
    IntegerVector segment_ids = IntegerVector()
) {
    std::vector<std::vector<int>> unit_rows;
    unit_rows.reserve(unit_row_indices.size());
    for (int i = 0; i < unit_row_indices.size(); ++i) {
        IntegerVector iv = unit_row_indices[i];
        unit_rows.push_back(std::vector<int>(iv.begin(), iv.end()));
    }
    std::vector<int> segs(segment_ids.begin(), segment_ids.end());

    return qe::apply_door_ema(conn_counts, unit_rows, alpha, segs);
}

// =============================================================================
// Trajectory Curve Fitting & Differential Geometry
// =============================================================================

//' Fit 2D parametric polynomial curve with LOOCV-2D or AIC degree selection
//' @param points Numeric matrix (n x 2) of trajectory coordinates
//' @param t Numeric vector of time values (length n)
//' @param max_degree Maximum polynomial degree (default 3)
//' @param fixed_degree Exact degree to use when >= 1
//' @param criterion Criterion for degree selection: "loocv" or "aic"
//' @param basis Polynomial basis used for fitting: "orthogonal" or "raw"
//' @export
// [[Rcpp::export]]
List fit_trajectory_poly(
    arma::mat points,
    NumericVector t = NumericVector(),
    int max_degree = 3,
    int fixed_degree = 0,
    std::string criterion = "loocv",
    std::string basis = "orthogonal"
) {
    arma::vec t_vec;
    if (t.size() > 0) {
        t_vec = arma::vec(t.begin(), t.size());
    }
    qe::PolyCurveFit fit = qe::fit_poly_loocv2d(points, t_vec, max_degree, fixed_degree, criterion, basis);

    return List::create(
        Named("degree")        = fit.degree,
        Named("coeffs_x")      = fit.coeffs_x,
        Named("coeffs_y")      = fit.coeffs_y,
        Named("basis")         = fit.basis,
        Named("basis_alpha")   = fit.basis_alpha,
        Named("basis_norm2")   = fit.basis_norm2,
        Named("cv_error")      = fit.cv_error,
        Named("aic")           = fit.aic,
        Named("t")             = fit.t,
        Named("fitted_points") = fit.fitted
    );
}

//' Evaluate 2D parametric polynomial curve at specified time points
//' @param coeffs_x Polynomial coefficients for X
//' @param coeffs_y Polynomial coefficients for Y
//' @param t_eval Numeric vector of time evaluation points
//' @export
// [[Rcpp::export]]
arma::mat eval_trajectory_curve(
    arma::vec coeffs_x,
    arma::vec coeffs_y,
    arma::vec t_eval
) {
    return qe::eval_poly_curve(coeffs_x, coeffs_y, t_eval);
}

//' Evaluate trajectory differential geometry (velocities, speed, curvature, turns)
//' @param coeffs_x Polynomial coefficients for X
//' @param coeffs_y Polynomial coefficients for Y
//' @param t_eval Numeric vector of time evaluation points
//' @export
// [[Rcpp::export]]
List eval_trajectory_derivatives(
    arma::vec coeffs_x,
    arma::vec coeffs_y,
    arma::vec t_eval
) {
    qe::TrajectoryDerivatives d = qe::eval_trajectory_derivatives(coeffs_x, coeffs_y, t_eval);
    return List::create(
        Named("t")            = d.t,
        Named("vx")           = d.vx,
        Named("vy")           = d.vy,
        Named("speed")        = d.speed,
        Named("ax")           = d.ax,
        Named("ay")           = d.ay,
        Named("heading_rate") = d.heading_rate,
        Named("curvature")    = d.curvature
    );
}

// =============================================================================
// Trajectory Distance & Following Dynamics
// =============================================================================

//' Integrated Euclidean distance between two trajectory polynomial curves
//' @param coeffs_ax Polynomial coeffs for Curve A X
//' @param coeffs_ay Polynomial coeffs for Curve A Y
//' @param coeffs_bx Polynomial coeffs for Curve B X
//' @param coeffs_by Polynomial coeffs for Curve B Y
//' @param t_start Integration start (default 0.0)
//' @param t_end Integration end (default 1.0)
//' @export
// [[Rcpp::export]]
double integrated_trajectory_distance(
    arma::vec coeffs_ax,
    arma::vec coeffs_ay,
    arma::vec coeffs_bx,
    arma::vec coeffs_by,
    double t_start = 0.0,
    double t_end = 1.0
) {
    return qe::integrated_curve_distance(coeffs_ax, coeffs_ay, coeffs_bx, coeffs_by, t_start, t_end);
}

//' Lagged curve distance between follower and leader
//' @param coeffs_fol_x Polynomial coeffs for Follower X
//' @param coeffs_fol_y Polynomial coeffs for Follower Y
//' @param coeffs_ldr_x Polynomial coeffs for Leader X
//' @param coeffs_ldr_y Polynomial coeffs for Leader Y
//' @param lag Time lag in [0, 1)
//' @export
// [[Rcpp::export]]
double lagged_trajectory_distance(
    arma::vec coeffs_fol_x,
    arma::vec coeffs_fol_y,
    arma::vec coeffs_ldr_x,
    arma::vec coeffs_ldr_y,
    double lag = 0.0
) {
    return qe::lagged_curve_distance(coeffs_fol_x, coeffs_fol_y, coeffs_ldr_x, coeffs_ldr_y, lag);
}

//' Pairwise integrated trajectory distance matrix
//' @param all_coeffs_x List of X coefficient vectors
//' @param all_coeffs_y List of Y coefficient vectors
//' @export
// [[Rcpp::export]]
arma::mat pairwise_trajectory_distance(
    List all_coeffs_x,
    List all_coeffs_y
) {
    std::vector<arma::vec> cx;
    std::vector<arma::vec> cy;
    cx.reserve(all_coeffs_x.size());
    cy.reserve(all_coeffs_y.size());
    for (int i = 0; i < all_coeffs_x.size(); ++i) {
        NumericVector nx = all_coeffs_x[i];
        NumericVector ny = all_coeffs_y[i];
        cx.push_back(arma::vec(nx.begin(), nx.size()));
        cy.push_back(arma::vec(ny.begin(), ny.size()));
    }
    return qe::pairwise_trajectory_distance_matrix(cx, cy);
}

//' Signed turn-lag distance between two agents at specific lag delta
//' @param pts_a Points matrix for agent A (n_a x 2)
//' @param pts_b Points matrix for agent B (n_b x 2)
//' @param times_a Turn times for agent A
//' @param times_b Turn times for agent B
//' @param delta Signed lag in turns
//' @export
// [[Rcpp::export]]
List signed_turn_lag(
    arma::mat pts_a,
    arma::mat pts_b,
    arma::vec times_a,
    arma::vec times_b,
    int delta
) {
    auto [mean_d, count] = qe::signed_turn_lag_distance(pts_a, pts_b, times_a, times_b, delta);
    return List::create(
        Named("mean_distance") = mean_d,
        Named("valid_count")   = count
    );
}

//' Sweep signed turn lags to find optimal leader-follower lag
//' @param pts_a Points matrix for agent A (n_a x 2)
//' @param pts_b Points matrix for agent B (n_b x 2)
//' @param times_a Turn times for agent A
//' @param times_b Turn times for agent B
//' @param max_lag Maximum turn lag to evaluate (default 15)
//' @export
// [[Rcpp::export]]
List sweep_signed_turn_lags(
    arma::mat pts_a,
    arma::mat pts_b,
    arma::vec times_a,
    arma::vec times_b,
    int max_lag = 15
) {
    qe::SignedTurnLagResult res = qe::best_signed_turn_lag(pts_a, pts_b, times_a, times_b, max_lag);
    return List::create(
        Named("best_lag")           = res.best_lag,
        Named("min_mean_distance")  = res.min_mean_distance,
        Named("lags")               = res.lags,
        Named("mean_distances")     = res.mean_distances,
        Named("valid_counts")       = res.valid_counts
    );
}

//' Distance-distance matrix correlation for stability analysis
//' @param X Numeric matrix
//' @param Y Numeric matrix
//' @export
// [[Rcpp::export]]
double dist_dist_correlation(
    arma::mat X,
    arma::mat Y
) {
    return qe::dist_dist_correlation(X, Y);
}

//' Cross-covariance decay (CCD) window-size estimation
//'
//' Estimates the ENA moving-window size from the half-life decay lag of the
//' noise-corrected Frobenius norm of pooled conversation cross-covariance
//' matrices.
//'
//' @param conversations A list of numeric code matrices, one per conversation
//'   (rows in sequence; all matrices share the same number of columns).
//' @param max_window Maximum lag to evaluate (default 20).
//' @param min_overlap Minimum overlapping rows (N - lag) required for a
//'   conversation to contribute at a given lag (default 10).
//'
//' @return A list with \code{window_size}, \code{peak_lag}, and the per-lag
//'   curves \code{lag}, \code{frob}, \code{frob_sq_unbiased},
//'   \code{frob_unbiased_signed}, and \code{total_weight}.
//' @export
// [[Rcpp::export]]
Rcpp::List ccd_window(
    Rcpp::List conversations,
    int max_window = 20,
    int min_overlap = 10
) {
    std::vector<arma::mat> convos;
    convos.reserve(conversations.size());
    for (R_xlen_t i = 0; i < conversations.size(); ++i) {
        convos.push_back(Rcpp::as<arma::mat>(conversations[i]));
    }

    qe::CCDResult res = qe::ccd_window(convos, max_window, min_overlap);

    return Rcpp::List::create(
        Named("window_size")          = res.window_size,
        Named("peak_lag")             = res.peak_lag,
        Named("lag")                  = res.lag,
        Named("frob")                 = res.frob,
        Named("frob_sq_unbiased")     = res.frob_sq_unbiased,
        Named("frob_unbiased_signed") = res.frob_unbiased_signed,
        Named("total_weight")         = res.total_weight
    );
}

1		#' @useDynLib libqe, .registration = TRUE
2		#' @importFrom Rcpp evalCpp
3		NULL
4
5		.onLoad <- function(libname, pkgname) {
6		# Nothing to initialise — libqe is a header + Rcpp library.
7	!	invisible(NULL)
8		}

1		/**
2		* @file adjacency.hpp
3		* @brief Combinatorics and vector/matrix utilities for ENA connection networks.
4		*
5		* Pure C++ / Armadillo — no Rcpp dependency.
6		* Include @c <RcppArmadillo.h> before this header when building inside an R package.
7		*/
8		#ifndef LIBQE_ADJACENCY_HPP
9		#define LIBQE_ADJACENCY_HPP
10
11		#include <armadillo>
12		#include <string>
13		#include <vector>
14		#include <cmath>
15
16		namespace qe {
17
18		// ---------------------------------------------------------------------------
19		/// @name Combinatorics
20		/// @{
21		// ---------------------------------------------------------------------------
22
23		/**
24		* @brief Number of unordered pairs from @p n items (n choose 2).
25		*
26		* @param n Number of codes/nodes.
27		* @returns @c (n * (n - 1)) / 2.
28		*/
29	67x	inline int choose_two(int n) {
30	67x	return (n * (n - 1)) / 2;
31		}
32
33		/**
34		* @brief Upper-triangle index pairs for a symmetric matrix of side @p len.
35		*
36		* @param len Side length of the square matrix (number of codes).
37		* @param row Controls what is returned:
38		* - @c -1 (default): 2 × k matrix of [row_idx; col_idx].
39		* - @c 0: row indices only (1 × k).
40		* - @c 1: column indices only (1 × k).
41		* @returns An @c arma::umat of shape 2×k or 1×k depending on @p row.
42		* @note Equivalent to @c triIndices() in rENA/ena.cpp and tma/code.cpp.
43		*/
44	9x	inline arma::umat connection_indices(int len, int row = -1) {
45	9x	int vS = choose_two(len);
46	9x	int s = 0;
47
48	9x	arma::umat vR(2, vS, arma::fill::zeros);
49	9x	arma::umat vRone(1, vS, arma::fill::zeros);
50
51	45x	for (int i = 2; i <= len; i++) {
52	145x	for (int j = 0; j < i - 1; j++) {
53	109x	vR(0, s) = j;
54	109x	vR(1, s) = i - 1;
55	115x	if (row == 0) vRone[s] = j;
56	109x	else if (row == 1) vRone[s] = i - 1;
57	109x	s++;
58		}
59		}
60	18x	return (row == -1) ? vR : vRone;
61		}
62
63		/// @}
64
65		// ---------------------------------------------------------------------------
66		/// @name Vector ↔ upper-triangle conversions
67		/// @{
68		// ---------------------------------------------------------------------------
69
70		/**
71		* @brief Compute pairwise products and return as a flat upper-triangle vector.
72		*
73		* For each pair @c (j, i) with @c j < i, outputs @c v[j] * v[i].
74		* The result has length @c choose_two(v.size()).
75		*
76		* @param v A 1 × p row vector (or p-element matrix).
77		* @returns A row vector of length @c choose_two(p).
78		* @note Equivalent to @c vector_to_ut() in rENA/ena.cpp.
79		*/
80	40x	inline arma::rowvec code_connections(arma::mat v) {
81	40x	int vL = v.size();
82	40x	int vS = choose_two(vL);
83	40x	arma::rowvec out(vS, arma::fill::zeros);
84	40x	int s = 0;
85	130x	for (int i = 2; i <= vL; i++) {
86	240x	for (int j = 0; j < i - 1; j++) {
87	450x	out[s] = v[j] * v[i - 1];
88	150x	s++;
89		}
90		}
91	40x	return out;
92		}
93
94		/**
95		* @brief Fold a directed n² vector into an undirected upper-triangle vector.
96		*
97		* Symmetric elements are summed: A→B + B→A for every pair.
98		*
99		* @param v Flat directed vector of length n², column-major (n = sqrt(v.size())).
100		* @returns Upper-triangle row vector of length @c choose_two(n).
101		* @note Equivalent to @c vector_to_summed_uppertri() in tma/code.cpp.
102		*/
103	6x	inline arma::rowvec fold_directed_network(arma::vec v) {
104	6x	int n = static_cast<int>(std::round(std::sqrt(static_cast<double>(v.size()))));
105	6x	int tris = choose_two(n);
106
107	6x	arma::mat m(v);
108	6x	m.reshape(n, n);
109	12x	m = m + m.t();
110
111	6x	arma::colvec flat = m.as_col();
112	6x	arma::uvec up_inds = arma::trimatu_ind(arma::size(m), 1);
113	12x	return arma::conv_to<arma::rowvec>::from(flat.elem(up_inds));
114		}
115
116		/**
117		* @brief Flatten an adjacency matrix to a vector.
118		*
119		* @param x Square adjacency matrix (n × n).
120		* @param full If @c true (default), returns the full n² vector (directed).
121		* If @c false, returns the upper-triangle only (undirected).
122		* @returns Row vector of length n² or @c choose_two(n).
123		* @note Equivalent to @c adjacency_matrix_to_vector() in tma/code.cpp.
124		*/
125	2x	inline arma::rowvec network_to_vector(arma::mat x, bool full = true) {
126	3x	if (full) return arma::vectorise(x).t();
127	3x	arma::mat combined = arma::trimatu(x, 1) + arma::trimatl(x, -1);
128	1x	arma::uvec up_inds = arma::trimatu_ind(arma::size(combined), 1);
129	2x	return arma::vectorise(combined(up_inds)).t();
130		}
131
132		/// @}
133
134		// ---------------------------------------------------------------------------
135		/// @name String utilities
136		/// @{
137		// ---------------------------------------------------------------------------
138
139		/**
140		* @brief Generate "A & B" pair names for every upper-triangle position.
141		*
142		* For code names @c {"A", "B", "C"} returns @c {"A & B", "A & C", "B & C"}.
143		*
144		* @param v Ordered list of code names (length p).
145		* @returns Vector of @c choose_two(p) label strings.
146		* @note Equivalent to @c svector_to_ut() in rENA/ena.cpp.
147		*/
148	2x	inline std::vector<std::string> connection_names(std::vector<std::string> v) {
149	2x	int vL = v.size();
150	2x	int vS = choose_two(vL);
151	2x	std::vector<std::string> out(vS);
152	2x	int s = 0;
153	7x	for (int i = 2; i <= vL; i++) {
154	14x	for (int j = 0; j < i - 1; j++) {
155	9x	out[s] = v[j] + " & " + v[i - 1];
156	9x	s++;
157		}
158		}
159	2x	return out;
160		}
161
162		/// @}
163
164		} // namespace qe
165
166		#endif // LIBQE_ADJACENCY_HPP

1		/** @file ccd.hpp
2		* @brief Cross-Covariance Decay (CCD) window-size estimation.
3		*
4		* Estimates the ENA moving-window size from the "discourse coherence length":
5		* the half-life decay lag of the noise-corrected Frobenius norm of pooled
6		* conversation cross-covariance matrices.
7		*
8		* Reference: Shaffer, D. W. & Cai, Z. (2026). Discourse Coherence Length:
9		* Noise-Corrected Covariance Estimation of Window-Size in Epistemic Network
10		* Analysis. International Conference on Quantitative Ethnography (ICQE26).
11		*
12		* This kernel is pure-numeric: it consumes conversation subsets that have
13		* already been extracted from the raw data (one code matrix per conversation,
14		* rows in sequence). Column selection and conversation splitting belong in
15		* each language wrapper (cf. accumulation.hpp), so that the C++ core never
16		* touches data-frame semantics.
17		*/
18		#ifndef LIBQE_CCD_HPP
19		#define LIBQE_CCD_HPP
20
21		#include <armadillo>
22		#include <cmath>
23		#include <vector>
24
25		namespace qe {
26
27		/** @brief Result of a CCD window-size estimation.
28		*
29		* The `*` vectors are all length `max_window + 1`, indexed by lag `0..max_window`.
30		* Entries for lags with no contributing conversation are NaN (weight 0).
31		*/
32		struct CCDResult {
33		int window_size; ///< Estimated half-life lag window size (>= 1).
34		int peak_lag; ///< Lag of the peak corrected cross-covariance norm.
35		arma::vec lag; ///< Evaluated lags, 0..max_window.
36		arma::vec frob; ///< Uncorrected pooled Frobenius norm per lag.
37		arma::vec frob_sq_unbiased; ///< Noise-corrected squared Frobenius norm per lag.
38		arma::vec frob_unbiased_signed; ///< sign * sqrt(\|corrected sq\|) per lag.
39		arma::vec total_weight; ///< Pooled overlap weight per lag.
40		};
41
42		/** @brief Estimate the ENA moving-window size via cross-covariance decay.
43		*
44		* For each lag @c l in @c 0..max_window, pools window-matched cross-covariance
45		* matrices across all conversation subsets, subtracts a trace-variance noise
46		* floor, and locates the first lag strictly after the peak where the corrected
47		* norm falls to half of its peak value.
48		*
49		* @param[in] conversations One code matrix per conversation (N_k rows × C codes,
50		* rows in sequence). All matrices must share the same
51		* number of columns.
52		* @param[in] max_window Maximum lag to evaluate. Default 20.
53		* @param[in] min_overlap Minimum overlapping rows (N_k - l) required for a
54		* conversation to contribute at lag l. Default 10.
55		*
56		* @returns A @ref CCDResult. When no conversation contributes at any positive
57		* lag (all-NaN or non-positive corrected covariance), the window size
58		* defaults to 1 and @c peak_lag to 0.
59		*/
60	!	inline CCDResult ccd_window(
61		const std::vector<arma::mat>& conversations,
62		int max_window = 20,
63		int min_overlap = 10
64		) {
65	!	if (max_window < 0) max_window = 0;
66	!	const int n_lags = max_window + 1;
67
68	!	CCDResult out;
69	!	out.lag = arma::regspace<arma::vec>(0, max_window);
70	!	out.frob = arma::vec(n_lags).fill(arma::datum::nan);
71	!	out.frob_sq_unbiased = arma::vec(n_lags).fill(arma::datum::nan);
72	!	out.frob_unbiased_signed = arma::vec(n_lags).fill(arma::datum::nan);
73	!	out.total_weight = arma::zeros<arma::vec>(n_lags);
74	!	out.window_size = 1;
75	!	out.peak_lag = 0;
76
77		// Determine code count from the first non-empty conversation.
78	!	arma::uword C = 0;
79	!	for (const arma::mat& X : conversations) {
80	!	if (X.n_cols > 0) { C = X.n_cols; break; }
81		}
82	!	if (C == 0) return out; // nothing to work with -> window 1
83
84		// 1. Cross-covariance curves & noise floor across lags.
85	!	for (int i = 0; i < n_lags; ++i) {
86	!	const int l = i;
87	!	arma::mat Cov_sum = arma::zeros<arma::mat>(C, C);
88	!	double total_weight = 0.0;
89	!	double noise_floor_sq_weighted_sum = 0.0;
90
91	!	for (const arma::mat& X_k : conversations) {
92	!	const int N_k = static_cast<int>(X_k.n_rows);
93	!	const int weight = N_k - l;
94	!	if (weight < min_overlap) continue;
95
96		// A = rows [l, N_k-1], B = rows [0, N_k-1-l]; both have `weight` rows.
97	!	const arma::mat A = X_k.rows(l, N_k - 1);
98	!	const arma::mat B = X_k.rows(0, N_k - 1 - l);
99
100	!	const arma::rowvec p_A = arma::mean(A, 0);
101	!	const arma::rowvec p_B = arma::mean(B, 0);
102
103		// Window-matched cross-covariance: crossprod(B, A)/w - outer(p_B, p_A).
104	!	const arma::mat Cov_k = (B.t() * A) / weight - p_B.t() * p_A;
105
106		// Trace-variance noise floor (n-1 denominator, matching R's var()).
107		const double tr_var_A =
108	!	arma::accu(arma::square(A.each_row() - p_A)) / (weight - 1);
109		const double tr_var_B =
110	!	arma::accu(arma::square(B.each_row() - p_B)) / (weight - 1);
111	!	const double var_k_sq = tr_var_A * tr_var_B;
112
113	!	Cov_sum += weight * Cov_k;
114	!	total_weight += weight;
115	!	noise_floor_sq_weighted_sum += weight * var_k_sq;
116		}
117
118	!	if (total_weight > 0.0) {
119	!	const arma::mat Cov_pooled = Cov_sum / total_weight;
120	!	const double frob_sq_pooled = arma::accu(arma::square(Cov_pooled));
121	!	const double noise_floor_sq_pooled =
122	!	noise_floor_sq_weighted_sum / (total_weight * total_weight);
123	!	const double frob_sq_unbiased_pooled =
124		frob_sq_pooled - noise_floor_sq_pooled;
125
126	!	const double sgn = (frob_sq_unbiased_pooled > 0.0) ? 1.0
127	!	: (frob_sq_unbiased_pooled < 0.0) ? -1.0 : 0.0;
128
129	!	out.frob[i] = std::sqrt(frob_sq_pooled);
130	!	out.frob_sq_unbiased[i] = frob_sq_unbiased_pooled;
131	!	out.frob_unbiased_signed[i] = sgn * std::sqrt(std::abs(frob_sq_unbiased_pooled));
132	!	out.total_weight[i] = total_weight;
133		}
134		}
135
136		// 2. Half-life decay lag detection over the signed corrected norm.
137	!	const arma::vec& f = out.frob_unbiased_signed;
138
139		// Peak among lags > 0 with finite, positive corrected norm.
140	!	int peak_lag = -1;
141	!	double max_f = -arma::datum::inf;
142	!	for (int i = 1; i < n_lags; ++i) {
143	!	if (std::isfinite(f[i]) && f[i] > 0.0 && f[i] > max_f) {
144	!	max_f = f[i];
145	!	peak_lag = i; // lag == index
146		}
147		}
148
149	!	if (peak_lag < 0) {
150		// No positive corrected covariance anywhere -> default window 1.
151	!	out.window_size = 1;
152	!	out.peak_lag = 0;
153	!	return out;
154		}
155	!	out.peak_lag = peak_lag;
156
157		// First lag at or after the peak whose corrected norm has decayed to <= 50%.
158		// NaN comparisons are false, so NaN lags are skipped (matching R's which()).
159	!	int decay_lag = -1;
160	!	for (int i = peak_lag; i < n_lags; ++i) {
161	!	if (f[i] / max_f <= 0.5) { decay_lag = i; break; }
162		}
163
164	!	if (decay_lag >= 0) {
165	!	out.window_size = decay_lag;
166		} else {
167		// No decay below half-life within max_window: fall back to the last lag
168		// (at or after the peak) with a finite corrected norm.
169	!	int last_valid = -1;
170	!	for (int i = peak_lag; i < n_lags; ++i) {
171	!	if (std::isfinite(f[i])) last_valid = i;
172		}
173	!	out.window_size = (last_valid >= 0) ? last_valid : 1;
174		}
175
176	!	return out;
177		}
178
179		} // namespace qe
180
181		#endif // LIBQE_CCD_HPP

1		/** @file door.hpp
2		* @brief Door temporal pooling and smoothing functions for trajectory modeling.
3		*
4		* Provides sliding lookback window pooling (with uniform or linear weighting,
5		* sum or mean aggregation, and optional segment boundary resets) as well as
6		* Exponential Moving Average (EMA) smoothing over consecutive time steps.
7		*/
8		#ifndef LIBQE_DOOR_HPP
9		#define LIBQE_DOOR_HPP
10
11		#include <armadillo>
12		#include <vector>
13		#include <string>
14		#include <cmath>
15		#include <algorithm>
16
17		namespace qe {
18
19		/** @brief Apply a sliding lookback window pool over a single unit's connection matrix.
20		*
21		* For each row @c i, pools rows in the window [@c max(0, i - lookback_size + 1), @c i].
22		* If @p segment_ids is non-empty, the window does not cross rows with a different segment ID.
23		*
24		* @param[in] block Matrix of connection counts for one unit (n_rows × n_features).
25		* @param[in] lookback_size Number of rows in the lookback window (>= 1).
26		* @param[in] aggregate_mean When @c true, averages rows; when @c false, sums rows.
27		* @param[in] weighting_linear When @c true, applies linearly increasing weights
28		* (oldest row in window has weight 1, current row has weight k).
29		* When @c false, all rows in window have equal weight 1.
30		* @param[in] segment_ids Optional vector of segment/activity IDs (length n_rows).
31		* Window pooling resets at segment transitions.
32		*
33		* @returns A matrix of shape n_rows × n_features containing door-smoothed values.
34		*/
35	5x	inline arma::mat lookback_block(
36		const arma::mat& block,
37		int lookback_size = 20,
38		bool aggregate_mean = false,
39		bool weighting_linear = false,
40		const std::vector<int>& segment_ids = {}
41		) {
42	5x	int n_rows = static_cast<int>(block.n_rows);
43	5x	int n_cols = static_cast<int>(block.n_cols);
44	5x	arma::mat out(n_rows, n_cols, arma::fill::zeros);
45
46		if (n_rows == 0 \|\| n_cols == 0) return out;
47	5x	int k_max = std::max(1, lookback_size);
48	5x	bool has_segments = (!segment_ids.empty() && static_cast<int>(segment_ids.size()) >= n_rows);
49
50	25x	for (int i = 0; i < n_rows; ++i) {
51	20x	int start_idx = std::max(0, i - k_max + 1);
52
53	20x	if (has_segments) {
54	4x	int cur_seg = segment_ids[i];
55	10x	for (int s = i; s >= start_idx; --s) {
56	8x	if (segment_ids[s] != cur_seg) {
57	2x	start_idx = s + 1;
58	2x	break;
59		}
60		}
61		}
62
63	20x	int win_len = i - start_idx + 1;
64		if (win_len <= 0) continue;
65
66	20x	arma::rowvec acc(n_cols, arma::fill::zeros);
67	20x	arma::rowvec denom(n_cols, arma::fill::zeros);
68
69	56x	for (int j = start_idx; j <= i; ++j) {
70	36x	double w = 1.0;
71	36x	if (weighting_linear) {
72		// Linear weight: 1 for earliest in window, win_len for current row i
73	!	w = static_cast<double>(j - start_idx + 1);
74		}
75	108x	for (int col = 0; col < n_cols; ++col) {
76	72x	double value = block(j, col);
77	72x	if (!std::isnan(value)) {
78	66x	acc[col] += w * value;
79	132x	denom[col] += w;
80		}
81		}
82		}
83
84	20x	if (aggregate_mean) {
85	24x	for (int col = 0; col < n_cols; ++col) {
86		acc[col] = (denom[col] > 0.0) ? (acc[col] / denom[col]) : arma::datum::nan;
87		}
88		}
89
90	40x	out.row(i) = acc;
91		}
92
93	5x	return out;
94		}
95
96		/** @brief Apply Exponential Moving Average (EMA) smoothing over a single unit's connection matrix.
97		*
98		* Computes S_0 = Y_0, S_t = alpha * Y_t + (1 - alpha) * S_{t-1}.
99		* If @p segment_ids is non-empty, S_t resets to Y_t when transitioning to a new segment.
100		*
101		* @param[in] block Matrix of connection counts for one unit (n_rows × n_features).
102		* @param[in] alpha Smoothing parameter alpha in (0, 1].
103		* @param[in] segment_ids Optional vector of segment/activity IDs (length n_rows).
104		*
105		* @returns A matrix of shape n_rows × n_features containing EMA-smoothed values.
106		*/
107	2x	inline arma::mat ema_block(
108		const arma::mat& block,
109		double alpha = 0.1,
110		const std::vector<int>& segment_ids = {}
111		) {
112	2x	int n_rows = static_cast<int>(block.n_rows);
113	2x	int n_cols = static_cast<int>(block.n_cols);
114	2x	arma::mat out(n_rows, n_cols, arma::fill::zeros);
115
116		if (n_rows == 0 \|\| n_cols == 0) return out;
117	2x	double a = std::max(1e-6, std::min(1.0, alpha));
118		bool has_segments = (!segment_ids.empty() && static_cast<int>(segment_ids.size()) >= n_rows);
119
120	2x	out.row(0) = block.row(0);
121
122	6x	for (int i = 1; i < n_rows; ++i) {
123		bool reset = has_segments && (segment_ids[i] != segment_ids[i - 1]);
124	4x	if (reset) {
125	!	out.row(i) = block.row(i);
126		} else {
127	12x	for (int col = 0; col < n_cols; ++col) {
128	8x	double cur = block(i, col);
129	8x	if (std::isnan(cur)) {
130	4x	cur = out(i - 1, col);
131		}
132	24x	out(i, col) = a * cur + (1.0 - a) * out(i - 1, col);
133		}
134		}
135		}
136
137	2x	return out;
138		}
139
140		/** @brief Apply lookback door pooling to an entire connection matrix across multiple units.
141		*
142		* @param[in] conn_counts Complete connection matrix (all units).
143		* @param[in] unit_row_indices Vector of row-index vectors, one per unit (0-indexed).
144		* @param[in] lookback_sizes Vector of lookback sizes, one per unit.
145		* @param[in] aggregate_mean When @c true, compute mean; when @c false, compute sum.
146		* @param[in] weighting_linear When @c true, use linear weighting.
147		* @param[in] segment_ids Optional vector of segment IDs for every row in conn_counts.
148		*
149		* @returns Smoothed connection matrix with identical dimensions to @p conn_counts.
150		*/
151	!	inline arma::mat apply_door_lookback(
152		const arma::mat& conn_counts,
153		const std::vector<std::vector<int>>& unit_row_indices,
154		const std::vector<int>& lookback_sizes,
155		bool aggregate_mean = false,
156		bool weighting_linear = false,
157		const std::vector<int>& segment_ids = {}
158		) {
159	!	arma::mat out = conn_counts;
160	!	size_t n_units = unit_row_indices.size();
161	!	bool has_segments = (!segment_ids.empty() && segment_ids.size() >= conn_counts.n_rows);
162
163	!	for (size_t u = 0; u < n_units; ++u) {
164	!	const auto& rows = unit_row_indices[u];
165	!	if (rows.empty()) continue;
166
167	!	int k = (u < lookback_sizes.size()) ? lookback_sizes[u] : 20;
168	!	int n_u = static_cast<int>(rows.size());
169
170	!	arma::mat block(n_u, conn_counts.n_cols);
171	!	std::vector<int> block_segs(n_u, 0);
172
173	!	for (int i = 0; i < n_u; ++i) {
174	!	int r = rows[i];
175	!	block.row(i) = conn_counts.row(r);
176	!	if (has_segments) block_segs[i] = segment_ids[r];
177		}
178
179		arma::mat smoothed = lookback_block(block, k, aggregate_mean, weighting_linear,
180	!	has_segments ? block_segs : std::vector<int>{});
181
182	!	for (int i = 0; i < n_u; ++i) {
183	!	out.row(rows[i]) = smoothed.row(i);
184		}
185		}
186
187	!	return out;
188		}
189
190		/** @brief Apply EMA door smoothing to an entire connection matrix across multiple units.
191		*/
192	!	inline arma::mat apply_door_ema(
193		const arma::mat& conn_counts,
194		const std::vector<std::vector<int>>& unit_row_indices,
195		double alpha = 0.1,
196		const std::vector<int>& segment_ids = {}
197		) {
198	!	arma::mat out = conn_counts;
199	!	size_t n_units = unit_row_indices.size();
200	!	bool has_segments = (!segment_ids.empty() && segment_ids.size() >= conn_counts.n_rows);
201
202	!	for (size_t u = 0; u < n_units; ++u) {
203	!	const auto& rows = unit_row_indices[u];
204	!	if (rows.empty()) continue;
205
206	!	int n_u = static_cast<int>(rows.size());
207	!	arma::mat block(n_u, conn_counts.n_cols);
208	!	std::vector<int> block_segs(n_u, 0);
209
210	!	for (int i = 0; i < n_u; ++i) {
211	!	int r = rows[i];
212	!	block.row(i) = conn_counts.row(r);
213	!	if (has_segments) block_segs[i] = segment_ids[r];
214		}
215
216	!	arma::mat smoothed = ema_block(block, alpha, has_segments ? block_segs : std::vector<int>{});
217
218	!	for (int i = 0; i < n_u; ++i) {
219	!	out.row(rows[i]) = smoothed.row(i);
220		}
221		}
222
223	!	return out;
224		}
225
226		} // namespace qe
227
228		#endif // LIBQE_DOOR_HPP

1		/**
2		* @file generalized_rotation.hpp
3		* @brief Generalized Means Rotation (GMR) with Lasso-based covariate adjustment.
4		*
5		* @details
6		* Mirrors rENA's `ena.rotate.by.generalized()` + `gmr()` + `get_x1_main_effect()`.
7		* Authoritative reference: commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901
8		* ("Fixed bugs in rotation functions") by Zhiqiang Cai.
9		*
10		* Algorithm overview
11		*
12		* 1. Compute x_vector: the rotation axis that best represents the target
13		* variable's contribution to the ENA point space, after controlling for
14		* covariates via Lasso (`lasso_x1_contribution`).
15		* 2. Deflate V by x_vector → defA = V - Vxx'.
16		* 3. Optionally compute secondary axis x1: the normalised group-mean
17		* difference in the original V space (not defA) for groups 0 vs 1.
18		* x1 stays empty when the target is not categorical or there are not
19		* exactly 2 groups — there is no SVD fallback.
20		* 4. Compute y_vector from a GMR call on defA, or from SVD of defA.
21		* Explicitly orthogonalise y_vector against x_vector, and against x1
22		* when x1 is available.
23		* 5. Double-deflate V by x_vector and y_vector, then call `complete_rotation()`
24		* to fill remaining dimensions from the SVD of the doubly-deflated space.
25		*
26		* Caller responsibilities
27		*
28		* - Build the model matrix (`X_raw`) from the data frame, including dummy
29		* variables for categorical predictors and any desired interaction terms.
30		* No formula parsing is done here.
31		* - Identify `x1_cols`: the 0-based column indices in `X_raw` that correspond to
32		* the target variable (these are unpenalized in the Lasso).
33		* - Encode categorical targets as 0-based integer codes in `x_target`.
34		* - Set `x_n_groups` to the number of distinct groups when `x_categorical=true`.
35		*
36		* The R wrapper (`rENA::ena.rotate.by.generalized`) handles these steps and
37		* calls `libqe::generalized_means_rotation()` with the prepared matrices.
38		*/
39
40		#ifndef LIBQE_GENERALIZED_ROTATION_HPP
41		#define LIBQE_GENERALIZED_ROTATION_HPP
42
43		#include "linalg_fallback.hpp" // qe::linalg::eig_sym (works w/o LAPACK)
44		#include "rotation.hpp" // ena_svd, complete_rotation, RotationResult
45		#include "lasso.hpp" // lasso_x1_contribution
46
47		#include <armadillo>
48		#include <stdexcept>
49		#include <string>
50		#include <vector>
51
52		namespace qe {
53
54		/// @name Between-group scatter
55		/// @{
56
57		/**
58		* @brief Compute the between-group scatter matrix S_B.
59		*
60		* @details
61		* Computes S_B = Σ_g n_g * (μ_g − μ)(μ_g − μ)' where μ is the grand mean.
62		*
63		* @note Mirrors rENA's `compute_SB()` (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
64		*
65		* @param[in] V n×p matrix of ENA points (or covariate-adjusted contribution).
66		* @param[in] labels n×1 0-based integer group labels.
67		* @param[in] n_groups Number of distinct groups.
68		*
69		* @returns p×p symmetric between-group scatter matrix.
70		*/
71	4x	inline arma::mat between_group_scatter(
72		const arma::mat& V,
73		const arma::uvec& labels,
74		arma::uword n_groups
75		) {
76	4x	const arma::uword p = V.n_cols;
77	4x	arma::rowvec mu = arma::mean(V, 0);
78	4x	arma::mat SB(p, p, arma::fill::zeros);
79
80	12x	for (arma::uword g = 0; g < n_groups; ++g) {
81	8x	arma::uvec idx = arma::find(labels == g);
82	8x	if (idx.is_empty()) continue;
83	21x	arma::vec diff = (arma::mean(V.rows(idx), 0) - mu).t();
84	14x	SB += static_cast<double>(idx.n_elem) * diff * diff.t();
85		}
86	8x	return SB;
87		}
88
89		/// @}
90
91		/// @name Rotation direction
92		/// @{
93
94		/**
95		* @brief Compute the unit rotation direction from a covariate-adjusted contribution matrix.
96		*
97		* @details
98		* Numeric target:
99		* Centres the target, then computes β = (t_c' t_c)^{-1} t_c' Vx (OLS slope from
100		* lm(Vx ~ target)). Mirrors R: `model = lm(Vx ~ target); beta = coefficients[2,]`.
101		*
102		* Categorical target:
103		* Computes the between-group scatter matrix S_B and returns its leading
104		* eigenvector (the direction of maximum between-group variance).
105		* Mirrors R: `sb = compute_SB(Vx, target); r = svd(sb)$v[,1]`.
106		*
107		* @note Mirrors rENA's `gmr()` direction computation
108		* (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
109		*
110		* @param[in] Vx n×q covariate-adjusted contribution matrix.
111		* @param[in] target n×1 target variable (numeric float, or 0-based integer codes).
112		* @param[in] categorical `true` if the target is categorical (uses between-group scatter);
113		* `false` for a numeric target (uses OLS slope).
114		* @param[in] n_groups Number of distinct groups; only used when `categorical` is `true`.
115		*
116		* @returns Unit vector (p×1) representing the rotation direction.
117		*
118		* @throws std::runtime_error If `categorical` is `true` and `n_groups < 2`.
119		* @throws std::runtime_error If the numeric target has zero variance after centering.
120		* @warning Throws `std::runtime_error` if the resulting direction has zero norm,
121		* which can occur when the Lasso-adjusted contribution collapses to zero.
122		*/
123	15x	inline arma::vec gmr_direction(
124		const arma::mat& Vx, // n×q covariate-adjusted contribution
125		const arma::vec& target, // n×1 (numeric float, or 0-based int codes)
126		bool categorical,
127		arma::uword n_groups = 0
128		) {
129	15x	arma::vec r;
130
131	15x	if (categorical) {
132	4x	if (n_groups < 2)
133		throw std::runtime_error(
134	!	"gmr_direction: x_n_groups must be >= 2 for categorical target");
135	4x	arma::uvec labels = arma::conv_to<arma::uvec>::from(target);
136	4x	arma::mat SB = between_group_scatter(Vx, labels, n_groups);
137
138	4x	arma::vec eigval; arma::mat eigvec;
139	4x	qe::linalg::eig_sym(eigval, eigvec, SB); // ascending order
140	8x	r = eigvec.col(eigvec.n_cols - 1); // largest eigenvalue → last col
141		} else {
142		// Centre target to get the OLS slope (mirrors R's lm() intercept handling)
143	22x	arma::vec tc = target - arma::mean(target);
144	11x	double tTt = arma::dot(tc, tc);
145	11x	if (tTt < 1e-12)
146		throw std::runtime_error(
147	!	"gmr_direction: target has zero variance after centering");
148	11x	arma::rowvec beta = (tc.t() * Vx) / tTt; // 1×q slope vector
149	11x	r = beta.t();
150		}
151
152	15x	double norm = arma::norm(r);
153	15x	if (norm < 1e-12)
154	!	throw std::runtime_error("gmr_direction: resulting direction has zero norm");
155	30x	return r / norm;
156		}
157
158		/// @}
159
160		/**
161		* @brief Result of one GMR axis step.
162		*/
163		struct GMRStepResult {
164		arma::vec direction; ///< Unit vector in connection space (p×1).
165		};
166
167		/// @name GMR axis step
168		/// @{
169
170		/**
171		* @brief Compute the rotation direction for one axis of the Generalized Means Rotation.
172		*
173		* @details
174		* Applies optional row subsetting, performs Lasso-based covariate adjustment
175		* (or a simple OLS projection when all columns are target columns), then
176		* delegates to `gmr_direction()` to produce a unit direction vector.
177		*
178		* When `x1_cols.n_elem >= p` (all columns are target columns and there are no
179		* penalized covariates), the function falls back to a simple OLS projection:
180		* Vx = fitted values from lm(V ~ t).
181		*
182		* @note Mirrors rENA's `gmr()` step logic
183		* (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
184		*
185		* @param[in] V n×q matrix of ENA points (full data; subset applied internally).
186		* @param[in] X_model n×p model matrix (caller-built; target cols + covariate cols).
187		* @param[in] target n×1 target variable (raw float or 0-based integer codes).
188		* @param[in] x1_cols 0-based indices of target columns in `X_model` (unpenalized in Lasso).
189		* @param[in] categorical `true` if the target is categorical.
190		* @param[in] n_groups Number of distinct groups (only used when `categorical` is `true`).
191		* @param[in] subset Optional 0-based row indices to use; empty means use all rows.
192		* @param[in] n_lambda Length of the Lasso lambda path (forwarded to `lasso_x1_contribution`).
193		* @param[in] k_folds Number of cross-validation folds (forwarded to `lasso_x1_contribution`).
194		* @param[in] lasso_eps lambda_min = lasso_eps * lambda_max (forwarded to `lasso_x1_contribution`).
195		*
196		* @returns `GMRStepResult` containing the unit direction vector.
197		*
198		* @throws std::runtime_error Propagated from `gmr_direction()` if the direction
199		* is degenerate or the target has zero variance.
200		*/
201	15x	inline GMRStepResult gmr_step(
202		const arma::mat& V,
203		const arma::mat& X_model,
204		const arma::vec& target,
205		const arma::uvec& x1_cols,
206		bool categorical,
207		arma::uword n_groups,
208		const arma::uvec& subset,
209		int n_lambda = 50,
210		int k_folds = 5,
211		double lasso_eps = 0.01
212		) {
213	15x	const arma::uword n = V.n_rows;
214	15x	const arma::uword p = X_model.n_cols;
215
216		// ── Select rows ──────────────────────────────────────────────────────────
217	15x	arma::uvec rows = subset.is_empty()
218	15x	? arma::regspace<arma::uvec>(0, n - 1)
219	15x	: subset;
220
221	15x	const arma::mat V_sub = V.rows(rows);
222	15x	const arma::mat X_sub = X_model.rows(rows);
223	30x	const arma::vec t_sub = target.rows(rows);
224
225		// ── Lasso-adjusted contribution Vx ───────────────────────────────────────
226	15x	arma::mat Vx_sub;
227	15x	if (x1_cols.n_elem >= p) {
228		// All columns are target columns (no penalized covariates):
229		// fall back to simple OLS projection (Vx = fitted values from lm(V~t))
230	30x	arma::vec tc = t_sub - arma::mean(t_sub);
231	15x	double tTt = arma::dot(tc, tc);
232	15x	if (tTt > 1e-12) {
233	14x	arma::rowvec b1 = (tc.t() * V_sub) / tTt;
234	14x	Vx_sub = tc * b1;
235		} else {
236	1x	Vx_sub.zeros(V_sub.n_rows, V_sub.n_cols);
237		}
238		} else {
239	!	arma::vec pf(p, arma::fill::ones);
240	!	for (arma::uword j : x1_cols) pf(j) = 0.0;
241
242	!	Vx_sub = lasso_x1_contribution(
243	!	X_sub, V_sub, x1_cols, pf, n_lambda, k_folds, lasso_eps);
244		}
245
246		// ── Rotation direction ───────────────────────────────────────────────────
247	15x	arma::vec direction = gmr_direction(Vx_sub, t_sub, categorical, n_groups);
248
249	30x	return {direction};
250		}
251
252		/// @}
253
254		/**
255		* @brief Parameters for `generalized_means_rotation()`.
256		*
257		* @details
258		* Groups all inputs needed to specify both the X and Y rotation axes, along
259		* with Lasso tuning knobs. The Y axis is optional: when `has_y` is `false`
260		* the second axis is taken from the leading SVD of the x-deflated space.
261		*/
262		struct GeneralizedRotationParams {
263		/// @name X axis (required)
264		/// @{
265		arma::mat x_model_matrix; ///< n×p model matrix for the X axis.
266		arma::vec x_target; ///< n×1 target variable (raw float or 0-based integer codes).
267		arma::uvec x1_cols; ///< 0-based target column indices in `x_model_matrix`.
268		bool x_categorical = false; ///< `true` if the X target is categorical.
269		arma::uword x_n_groups = 0; ///< Number of distinct groups (used when `x_categorical` is `true`).
270		arma::uvec x_subset; ///< Optional row subset (0-based); empty means use all rows.
271		/// @}
272
273		/// @name Y axis (optional)
274		/// @{
275		/// When `has_y` is `false`, the Y axis is taken from the leading SVD of the x-deflated space.
276		bool has_y = false; ///< `true` if a Y-axis GMR target is provided.
277		arma::mat y_model_matrix; ///< n×p model matrix for the Y axis.
278		arma::vec y_target; ///< n×1 target variable for the Y axis.
279		arma::uvec y1_cols; ///< 0-based target column indices in `y_model_matrix`.
280		bool y_categorical = false; ///< `true` if the Y target is categorical.
281		arma::uword y_n_groups = 0; ///< Number of distinct groups for the Y axis.
282		/// The Y axis always uses all rows (no subset parameter).
283		/// @}
284
285		/// @name Lasso tuning
286		/// @{
287		int n_lambda = 50; ///< Length of the lambda path.
288		int k_folds = 5; ///< Number of cross-validation folds.
289		double lasso_eps = 0.01; ///< lambda_min = lasso_eps * lambda_max.
290		/// @}
291		};
292
293		/// @name Generalized Means Rotation
294		/// @{
295
296		/**
297		* @brief Full Generalized Means Rotation (GMR) producing a complete q×q rotation matrix.
298		*
299		* @details
300		* C++/Armadillo implementation of `ena.rotate.by.generalized()`, matching the
301		* authoritative R implementation in commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901
302		* by Zhiqiang Cai.
303		*
304		* Steps performed:
305		*
306		* 1. X axis via GMR — calls `gmr_step()` on V with the X-axis parameters.
307		* 2. Deflate V — computes defA = V − V·x·x' for subsequent Y-axis work.
308		* 3. Secondary axis x1 — when the X target is categorical with ≥ 2 groups,
309		* computes x1 = normalize(mean(V[group0]) − mean(V[group1])) using the
310		* original V (not defA) and the full-length target (all n rows).
311		* x1 is used only to orthogonalise y_vector; it does not deflate defA.
312		* There is no SVD fallback: x1 stays empty if the condition is not met.
313		* 4. Y axis — either GMR on defA (when `params.has_y` is `true`) or the
314		* leading SVD direction of defA.
315		* 5. Orthogonalise y_vector — mirrors R:
316		* `y_vector <- y_vector - (t(y_vector) %% x_vector) x_vector`
317		* and, when x1 is available,
318		* `y_vector <- y_vector - (t(y_vector) %% x1) x1`,
319		* followed by safe normalisation.
320		* 6. Complete with SVD — double-deflates V by both x_vector and y_vector,
321		* then calls `complete_rotation()` to fill remaining dimensions.
322		*
323		* @note Mirrors rENA's `ena.rotate.by.generalized()`
324		* (commit 46776a1981a90b3a3b2861ed1010e9dbb7acf901).
325		*
326		* @param[in] V n×q matrix of ENA points (line weights; caller normalises/
327		* centres upstream).
328		* @param[in] params Fully populated `GeneralizedRotationParams` struct.
329		*
330		* @returns `RotationResult` with:
331		* - `rotation` q×q full rotation matrix (column j = axis j).
332		* - `eigenvalues` q×1 variance per SVD-filled axis (named axes carry 0).
333		* - `column_names` "GMR1", "GMR2" or "SVD2", "SVD3", …, "SVD{q}".
334		*
335		* @throws std::runtime_error If `gmr_direction()` produces a degenerate direction
336		* for the X or Y axis.
337		* @throws std::runtime_error If y_vector has zero norm after orthogonalisation
338		* (can happen when defA is rank-deficient or perfectly aligned with x_vector).
339		*/
340	13x	inline RotationResult generalized_means_rotation(
341		const arma::mat& V,
342		const GeneralizedRotationParams& params
343		) {
344		// ── Step 1: X axis via GMR ────────────────────────────────────────────
345		GMRStepResult x_gmr = gmr_step(
346		V,
347	13x	params.x_model_matrix,
348	13x	params.x_target,
349	13x	params.x1_cols,
350	13x	params.x_categorical,
351	13x	params.x_n_groups,
352	13x	params.x_subset,
353	13x	params.n_lambda,
354	13x	params.k_folds,
355	13x	params.lasso_eps);
356
357	13x	const arma::vec x_vector = x_gmr.direction;
358
359		// ── Step 2: Deflate V by x_vector ─────────────────────────────────────
360		// defA is used only for y_vector computation; x1 uses original V.
361	26x	const arma::mat defA = V - V * x_vector * x_vector.t();
362
363		// ── Step 3: Secondary axis x1 ─────────────────────────────────────────
364		// Mirrors R: x1 = normalize(mean(V[group0]) - mean(V[group1])).
365		//
366		// Key points matching the authoritative R:
367		// - Uses original V (not defA)
368		// - Uses the full-length target (all n rows), not just the subset
369		// - Groups are always 0 (first) and 1 (second) in 0-based encoding
370		// - NO fallback: x1 stays empty when categorical/2-group condition fails
371		// - x1 does NOT deflate defA; it is only used to orthogonalise y_vector
372	13x	arma::vec x1;
373
374	13x	if (params.x_categorical && params.x_n_groups >= 2) {
375	4x	arma::uvec labels = arma::conv_to<arma::uvec>::from(params.x_target);
376	4x	arma::uvec g0 = arma::find(labels == 0);
377	4x	arma::uvec g1 = arma::find(labels == 1);
378
379		if (!g0.is_empty() && !g1.is_empty()) {
380	4x	arma::rowvec m0 = arma::mean(V.rows(g0), 0); // original V, not defA
381	4x	arma::rowvec m1 = arma::mean(V.rows(g1), 0);
382	4x	arma::vec diff = (m0 - m1).t();
383	4x	double dlen = arma::norm(diff);
384	8x	if (dlen > 1e-10) x1 = diff / dlen;
385		}
386		}
387
388		// ── Step 4: Y axis ────────────────────────────────────────────────────
389		// GMR on defA (deflated by x_vector only), or SVD of defA.
390	13x	arma::vec y_vector;
391	13x	std::string y_name = "SVD2";
392
393	13x	if (params.has_y) {
394		GMRStepResult y_gmr = gmr_step(
395		defA,
396	2x	params.y_model_matrix,
397	2x	params.y_target,
398	2x	params.y1_cols,
399	2x	params.y_categorical,
400	2x	params.y_n_groups,
401		{}, // no row subset for y axis
402	2x	params.n_lambda,
403	2x	params.k_folds,
404	2x	params.lasso_eps);
405	2x	y_vector = y_gmr.direction;
406	2x	y_name = "GMR2";
407		} else {
408	11x	RotationResult svd_r = ena_svd(defA);
409	11x	y_vector = svd_r.rotation.col(0);
410		}
411
412		// ── Step 5: Orthogonalise y_vector ────────────────────────────────────
413		// Mirrors R (authoritative commit):
414		// y_vector <- y_vector - (t(y_vector) %% x_vector) x_vector
415		// if (!is.null(x1)) y_vector <- y_vector - (t(y_vector) %% x1) x1
416		// y_vector <- safe_normalize(y_vector)
417	26x	y_vector -= arma::dot(y_vector, x_vector) * x_vector;
418	13x	if (!x1.is_empty())
419	8x	y_vector -= arma::dot(y_vector, x1) * x1;
420
421	13x	double yn = arma::norm(y_vector);
422	13x	if (yn < 1e-12)
423		throw std::runtime_error(
424	!	"generalized_means_rotation: y_vector has zero norm after orthogonalization");
425	13x	y_vector /= yn;
426
427		// ── Step 6: Combine and complete with SVD ─────────────────────────────
428		// Double-deflate the original data (mirrors R):
429		// defA <- A - A %% x %% t(x) - A %% y %% t(y)
430		// then pass to prcomp (equivalent: SVD of defA_final).
431		arma::mat defA_final = V
432	13x	- V * x_vector * x_vector.t()
433	26x	- V * y_vector * y_vector.t();
434
435	13x	arma::mat named_axes = arma::join_rows(x_vector, y_vector);
436	65x	return complete_rotation(defA_final, named_axes, {"GMR1", y_name});
437		}
438
439		/// @}
440
441		} // namespace qe
442		#endif // LIBQE_GENERALIZED_ROTATION_HPP

1		/**
2		* @file lasso.hpp
3		* @brief Coordinate-descent Lasso for multivariate response with k-fold CV.
4		*
5		* Mirrors `glmnet(x, y, family = "mgaussian", penalty.factor = ...)` with
6		* k-fold cross-validation to select lambda, and warm-starts along the lambda
7		* path for efficiency.
8		*
9		* The public entry point is qe::lasso_x1_contribution(), which returns an
10		* n×q matrix of fitted values attributable to the "target" predictors at the
11		* CV-selected lambda. It is used by `generalized_means_rotation` to isolate
12		* the target variable's contribution to the ENA point space after controlling
13		* for covariates.
14		*
15		* Penalty factors (`pf`):
16		* - `0.0` — predictor is always included (forced in; no shrinkage)
17		* - `1.0` — standard L1 penalty
18		*
19		* The caller is responsible for building the model matrix (main effects and/or
20		* interactions); this function accepts a numeric matrix, not a formula.
21		*/
22
23		#ifndef LIBQE_LASSO_HPP
24		#define LIBQE_LASSO_HPP
25
26		#include <armadillo>
27		#include <cmath>
28		#include <limits>
29		#include <stdexcept>
30
31		namespace qe {
32
33		/**
34		* @brief Soft-threshold (shrinkage) operator used by coordinate descent.
35		*
36		* @details Computes the scalar soft-threshold function:
37		* \f$ S(z,\gamma) = \operatorname{sign}(z)\,\max(\|z\|-\gamma,\,0) \f$
38		*
39		* @param[in] z Input value.
40		* @param[in] gamma Threshold amount (must be non-negative for meaningful use).
41		* @returns Soft-thresholded value; exactly zero when `\|z\| <= gamma`.
42		*/
43	!	inline double soft_threshold(double z, double gamma) noexcept {
44	!	return z > gamma ? z - gamma : z < -gamma ? z + gamma : 0.0;
45		}
46
47		/**
48		* @brief Single-response Lasso via cyclic coordinate descent at a fixed lambda.
49		*
50		* @details Fits a Lasso regression for a single response vector using cyclic
51		* coordinate descent (naive updates). The partial-residual update for
52		* predictor \f$j\f$ is:
53		* \f[
54		* z_j = \frac{X_j^\top r}{n} + \frac{X_j^\top X_j}{n}\,\beta_j
55		* \f]
56		* \f[
57		* \beta_j \leftarrow \frac{S\!\left(z_j,\;\lambda\,\mathit{pf}_j\right)}{X_j^\top X_j / n}
58		* \f]
59		* where \f$r = y - X\beta\f$ is the residual maintained incrementally.
60		* Predictors with \f$\mathit{pf}_j = 0\f$ skip the soft-threshold and receive
61		* the unconstrained OLS update instead.
62		*
63		* @param[in] X n×p column-standardized predictor matrix.
64		* @param[in] y n×1 response vector (centering recommended but not required).
65		* @param[in] lambda Regularization strength, on the standardized predictor scale.
66		* @param[in] pf p×1 penalty factors (0 = unpenalized / forced in, 1 = standard L1).
67		* @param[in] beta0 Warm-start coefficient vector; pass an empty `arma::vec{}`
68		* for a cold start from zero.
69		* @param[in] max_iter Maximum number of coordinate-descent passes (default 1000).
70		* @param[in] tol Convergence tolerance on the maximum coefficient change
71		* per pass (default 1e-7).
72		* @returns p×1 coefficient vector on the standardized predictor scale.
73		*/
74	!	inline arma::vec lasso_cd(
75		const arma::mat& X,
76		const arma::vec& y,
77		double lambda,
78		const arma::vec& pf,
79		arma::vec beta0 = {},
80		int max_iter = 1000,
81		double tol = 1e-7
82		) {
83	!	const arma::uword p = X.n_cols;
84	!	const double n = static_cast<double>(X.n_rows);
85
86		// Precompute X_j'X_j / n for each predictor
87	!	arma::vec xTx(p);
88	!	for (arma::uword j = 0; j < p; ++j)
89	!	xTx(j) = arma::dot(X.col(j), X.col(j)) / n;
90
91	!	arma::vec beta = beta0.is_empty()
92	!	? arma::vec(p, arma::fill::zeros)
93	!	: beta0;
94	!	arma::vec r = y - X * beta; // residual, maintained incrementally
95
96	!	for (int it = 0; it < max_iter; ++it) {
97	!	double max_delta = 0.0;
98
99	!	for (arma::uword j = 0; j < p; ++j) {
100	!	if (xTx(j) < 1e-14) continue; // skip zero-variance predictor
101
102		// Partial correlation statistic
103	!	double z = arma::dot(X.col(j), r) / n + xTx(j) * beta(j);
104
105	!	double b_new = (pf(j) < 1e-12)
106	!	? z / xTx(j) // forced in: OLS
107	!	: soft_threshold(z, lambda * pf(j)) / xTx(j); // L1
108
109	!	double delta = b_new - beta(j);
110	!	if (std::abs(delta) > max_delta) max_delta = std::abs(delta);
111	!	r -= X.col(j) * delta;
112	!	beta(j) = b_new;
113		}
114
115	!	if (max_delta < tol) break;
116		}
117	!	return beta;
118		}
119
120		/**
121		* @brief Compute the smallest lambda that drives all penalized coefficients to zero.
122		*
123		* @details For column-standardized `X` and column-centered `Yc`, lambda_max is:
124		* \f[
125		* \lambda_{\max} = \max_{\substack{j:\,\mathit{pf}_j > 0 \\ k}} \frac{\|X_j^\top Y_{c,k}\|}{n}
126		* \f]
127		* Above this value every penalized coefficient is exactly zero (unpenalized
128		* predictors are unaffected and remain in via OLS regardless of lambda).
129		*
130		* @param[in] X n×p column-standardized predictor matrix.
131		* @param[in] Yc n×q column-centered response matrix.
132		* @param[in] pf p×1 penalty factors; predictors with `pf(j) < 1e-12` are skipped.
133		* @returns Scalar lambda_max value.
134		*/
135	!	inline double lasso_lambda_max(
136		const arma::mat& X, // column-standardized
137		const arma::mat& Yc, // column-centered responses
138		const arma::vec& pf
139		) {
140	!	const double n = static_cast<double>(X.n_rows);
141	!	double lmax = 0.0;
142	!	for (arma::uword j = 0; j < X.n_cols; ++j) {
143	!	if (pf(j) < 1e-12) continue; // skip unpenalized predictors
144	!	for (arma::uword k = 0; k < Yc.n_cols; ++k) {
145	!	double v = std::abs(arma::dot(X.col(j), Yc.col(k))) / n;
146	!	if (v > lmax) lmax = v;
147		}
148		}
149	!	return lmax;
150		}
151
152		/**
153		* @brief Cross-validated multivariate Lasso; returns fitted values for target columns only.
154		*
155		* @details Fits a multivariate Lasso (one independent single-response problem per
156		* response column) along a log-spaced lambda path from lambda_max down to
157		* `eps * lambda_max`. Lambda is selected by k-fold cross-validation
158		* (mean squared error, averaged over folds and response columns), and the
159		* model is refit on the full data at the selected lambda using warm-starts
160		* along the path.
161		*
162		* Preprocessing performed internally:
163		* - `X_raw` is column-standardized to zero mean and unit variance (zero-variance
164		* columns are left as-is to avoid division by zero).
165		* - `Y` is column-centered to absorb the intercept.
166		* - Coefficients are unstandardized before the contribution is computed.
167		*
168		* Fold assignment: row \f$i\f$ is assigned to fold \f$i \bmod k\f$
169		* (deterministic, order-dependent).
170		*
171		* Return value: The n×q matrix \f$X_{\text{raw},x_1} \hat{B}_{x_1}\f$
172		* represents the portion of `Y` (mean-removed) that is linearly explained by
173		* the target predictors after penalizing covariate coefficients. A zero
174		* matrix is returned when the target predictors carry no signal at any lambda
175		* (e.g., the target is constant after subsetting). Constant shifts cancel
176		* in the downstream between-group scatter (categorical) and slope regression
177		* (numeric), so no `ymu` adjustment is applied to the return value.
178		*
179		* @param[in] X_raw n×p raw model matrix built by the caller (main effects
180		* and/or interactions); no formula parsing is done here.
181		* @param[in] Y n×q ENA point matrix (multivariate response).
182		* @param[in] x1_cols 0-based column indices in `X_raw` that belong to the
183		* target variable; these receive `pf = 0` (always included).
184		* @param[in] pf p×1 penalty factors: 0 for `x1_cols`, 1 elsewhere.
185		* The caller constructs this; non-standard weightings are
186		* supported.
187		* @param[in] n_lambda Length of the lambda path (default 50).
188		* @param[in] k_folds Number of cross-validation folds (default 5; clamped to
189		* `[2, n/2]`).
190		* @param[in] eps Ratio `lambda_min / lambda_max`; controls how far down
191		* the path to search (default 0.01).
192		* @returns n×q matrix of fitted values attributable to `x1_cols` at the
193		* CV-selected lambda (lambda.min), on the original `Y` scale.
194		*/
195	!	inline arma::mat lasso_x1_contribution(
196		const arma::mat& X_raw,
197		const arma::mat& Y,
198		const arma::uvec& x1_cols,
199		const arma::vec& pf,
200		int n_lambda = 50,
201		int k_folds = 5,
202		double eps = 0.01
203		) {
204	!	const arma::uword n = X_raw.n_rows;
205	!	const arma::uword p = X_raw.n_cols;
206	!	const arma::uword q = Y.n_cols;
207
208		// Guard: need at least 2 folds and 2 observations per fold
209	!	k_folds = std::max(2, std::min(k_folds, static_cast<int>(n) / 2));
210
211		// ── Standardize X (μ=0, σ=1); guard zero-variance predictors ────────────
212	!	arma::rowvec xmu = arma::mean(X_raw, 0);
213	!	arma::rowvec xstd = arma::stddev(X_raw, 0, 0);
214	!	xstd.elem(arma::find(xstd < 1e-12)).ones(); // avoid divide-by-zero
215	!	arma::mat X = (X_raw.each_row() - xmu).each_row() / xstd;
216
217		// ── Center Y (absorbs intercept into the standardized-scale fit) ─────────
218	!	arma::rowvec ymu = arma::mean(Y, 0);
219	!	arma::mat Yc = Y.each_row() - ymu;
220
221		// ── Build lambda path: lambda_max → lambda_min on log scale ─────────────
222	!	double lmax = lasso_lambda_max(X, Yc, pf);
223	!	if (lmax < 1e-12) // target uncorrelated with Y after standardization
224	!	return arma::mat(n, q, arma::fill::zeros);
225
226	!	arma::vec lambdas = arma::exp(
227	!	arma::linspace(std::log(lmax),
228		std::log(eps * lmax),
229	!	n_lambda));
230
231		// ── Stratified k-fold CV (fold assignment: row i → fold i % k_folds) ────
232	!	arma::uvec fold_id(n);
233	!	for (arma::uword i = 0; i < n; ++i)
234	!	fold_id(i) = i % static_cast<arma::uword>(k_folds);
235
236	!	arma::vec cv_err(n_lambda, arma::fill::zeros);
237
238	!	for (int fold = 0; fold < k_folds; ++fold) {
239	!	arma::uvec tr = arma::find(fold_id != static_cast<arma::uword>(fold));
240	!	arma::uvec va = arma::find(fold_id == static_cast<arma::uword>(fold));
241
242	!	const arma::mat X_tr = X.rows(tr), X_va = X.rows(va);
243	!	const arma::mat Y_tr = Yc.rows(tr), Y_va = Yc.rows(va);
244	!	const double n_va_q = static_cast<double>(va.n_elem * q);
245
246		// Warm-start along path for this fold (start cold at lambda_max)
247	!	arma::mat bw(p, q, arma::fill::zeros);
248	!	for (int l = 0; l < n_lambda; ++l) {
249	!	arma::mat b(p, q);
250	!	for (arma::uword k = 0; k < q; ++k)
251	!	b.col(k) = lasso_cd(X_tr, Y_tr.col(k), lambdas(l), pf, bw.col(k));
252	!	bw = b;
253
254	!	arma::mat err = Y_va - X_va * b;
255	!	cv_err(l) += arma::accu(err % err) / n_va_q;
256		}
257		}
258	!	cv_err /= static_cast<double>(k_folds);
259
260		// ── Refit on full data at lambda_min, warming along the path ─────────────
261	!	arma::uword best = cv_err.index_min();
262	!	arma::mat bfull(p, q, arma::fill::zeros);
263	!	for (arma::uword l = 0; l <= best; ++l) {
264	!	arma::mat b(p, q);
265	!	for (arma::uword k = 0; k < q; ++k)
266	!	b.col(k) = lasso_cd(X, Yc.col(k), lambdas(l), pf, bfull.col(k));
267	!	bfull = b;
268		}
269
270		// ── Unstandardize: beta_orig[j] = beta_std[j] / xstd[j] ─────────────────
271	!	for (arma::uword j = 0; j < p; ++j)
272	!	bfull.row(j) /= xstd(j);
273
274		// ── Return fitted contribution from x1 columns only ──────────────────────
275		// This is the portion of Y (mean-removed) linearly explained by the target
276		// predictors after covariate adjustment. Constant shifts cancel in both
277		// the between-group scatter (categorical) and slope regression (numeric)
278		// that follow, so no ymu adjustment is needed here.
279	!	return X_raw.cols(x1_cols) * bfull.rows(x1_cols);
280		}
281
282		} // namespace qe
283		#endif // LIBQE_LASSO_HPP

1		/** @file linalg_fallback.hpp
2		* @brief Thin LAPACK-free implementations of SVD, QR, and symmetric
3		* eigendecomposition for use in WASM / no-LAPACK builds.
4		*
5		* When @c ARMA_USE_LAPACK is defined these wrappers delegate directly to the
6		* standard Armadillo routines (@c arma::svd, @c arma::qr, @c arma::eig_sym),
7		* so there is zero overhead in normal builds. When LAPACK is not available
8		* (e.g. Emscripten WASM builds) the wrappers fall back to pure C++
9		* implementations that work for the small matrices typical in ENA
10		* (p ≤ ~100 connection dimensions).
11		*
12		* All routines live in namespace @c qe::linalg and operate on @c arma::mat
13		* (double precision).
14		*
15		* Fallback algorithms:
16		* - SVD — one-sided Jacobi iterations on \f$A^T\f$ (Demmel & Veselić 1992)
17		* - QR — modified Gram-Schmidt
18		* - eig_sym — cyclic Jacobi (classical Jacobi sweep)
19		*/
20
21		#ifndef LIBQE_LINALG_FALLBACK_HPP
22		#define LIBQE_LINALG_FALLBACK_HPP
23
24		#include <armadillo>
25		#include <cmath>
26		#include <limits>
27
28		namespace qe {
29		namespace linalg {
30
31		/// @cond INTERNAL
32		// ── helpers ───────────────────────────────────────────────────────────────────
33
34		namespace detail {
35
36		/// @brief Modified Gram-Schmidt orthonormalization applied in-place.
37		///
38		/// Columns of @p Q are successively orthonormalized against all previously
39		/// processed columns. A column whose norm falls below 1e-14 after projection
40		/// is left unchanged (not zeroed) to avoid introducing NaNs.
41		///
42		/// @param Q Matrix whose columns are orthonormalized in-place (modified).
43		inline void mgs(arma::mat& Q) {
44		const arma::uword p = Q.n_cols;
45		for (arma::uword j = 0; j < p; ++j) {
46		for (arma::uword i = 0; i < j; ++i)
47		Q.col(j) -= arma::dot(Q.col(i), Q.col(j)) * Q.col(i);
48		double n = arma::norm(Q.col(j));
49		if (n > 1e-14) Q.col(j) /= n;
50		}
51		}
52
53		/// @brief One cyclic Jacobi sweep on a p×p symmetric matrix.
54		///
55		/// Iterates over all super-diagonal index pairs @c (q,r) and applies a
56		/// two-sided Jacobi rotation that annihilates @c A(q,r). The same rotation
57		/// is accumulated in @p V so that, starting with @c V = I, the full sequence
58		/// of sweeps yields @c V = eigenvector matrix.
59		///
60		/// @param A Symmetric p×p matrix, updated in-place toward diagonal form.
61		/// @param V Rotation accumulator (p×p); initialize to identity to recover
62		/// eigenvectors.
63		///
64		/// @note Algorithm follows the classical cyclic-by-rows Jacobi scheme.
65		/// Off-diagonal entries no larger than 1e-14 × (\|a_qq\| + \|a_rr\|) are
66		/// skipped for numerical stability. The `<=` (not `<`) guard correctly
67		/// handles the all-zero case (a_qq = a_rr = a_qr = 0) that would
68		/// otherwise produce NaN via 0/0 in the theta computation.
69		inline void jacobi_sweep(arma::mat& A, arma::mat& V) {
70		const arma::uword p = A.n_rows;
71		for (arma::uword q = 0; q < p - 1; ++q) {
72		for (arma::uword r = q + 1; r < p; ++r) {
73		double aqq = A(q, q), arr = A(r, r), aqr = A(q, r);
74		if (std::abs(aqr) <= 1e-14 * (std::abs(aqq) + std::abs(arr)))
75		continue;
76		double theta = (arr - aqq) / (2.0 * aqr);
77		double t = (theta >= 0.0)
78		? 1.0 / (theta + std::sqrt(1.0 + theta * theta))
79		: 1.0 / (theta - std::sqrt(1.0 + theta * theta));
80		double c = 1.0 / std::sqrt(1.0 + t * t); ///< cosine of Jacobi rotation angle
81		double s = t * c; ///< sine of Jacobi rotation angle
82		// Apply Jacobi rotation to A (symmetric update)
83		double new_qq = aqq - t * aqr;
84		double new_rr = arr + t * aqr;
85		A(q, q) = new_qq; A(r, r) = new_rr; A(q, r) = A(r, q) = 0.0;
86		for (arma::uword k = 0; k < p; ++k) {
87		if (k == q \|\| k == r) continue;
88		double akq = A(k, q), akr = A(k, r);
89		A(k, q) = A(q, k) = c * akq - s * akr;
90		A(k, r) = A(r, k) = s * akq + c * akr;
91		}
92		// Accumulate rotation in V
93		for (arma::uword k = 0; k < p; ++k) {
94		double vkq = V(k, q), vkr = V(k, r);
95		V(k, q) = c * vkq - s * vkr;
96		V(k, r) = s * vkq + c * vkr;
97		}
98		}
99		}
100		}
101
102		} // namespace detail
103
104		/// @endcond
105
106		/// @name Eigendecomposition
107		/// @{
108
109		// ── eig_sym ───────────────────────────────────────────────────────────────────
110
111		/// @brief Eigendecomposition of a real symmetric matrix.
112		///
113		/// Computes eigenvalues and eigenvectors of the symmetric p×p matrix @p S and
114		/// returns them sorted in ascending eigenvalue order.
115		///
116		/// @param[out] eigval p×1 vector of eigenvalues, ascending.
117		/// @param[out] eigvec p×p matrix whose columns are the corresponding
118		/// orthonormal eigenvectors.
119		/// @param[in] S Real symmetric p×p input matrix (not modified).
120		///
121		/// @warning With @c ARMA_USE_LAPACK defined this calls @c arma::eig_sym
122		/// directly. Without LAPACK the fallback runs up to @c 50*p cyclic
123		/// Jacobi sweeps; convergence is guaranteed for well-conditioned
124		/// matrices but may be slow for p > 200.
125		///
126		/// @note Fallback algorithm: classical cyclic-by-rows Jacobi iteration.
127		/// Convergence is declared when the Frobenius norm of the strictly
128		/// off-diagonal part drops below 1e-13.
129	4x	inline void eig_sym(arma::vec& eigval, arma::mat& eigvec, const arma::mat& S) {
130		#if defined(ARMA_USE_LAPACK)
131	4x	arma::eig_sym(eigval, eigvec, S);
132		#else
133		const arma::uword p = S.n_rows;
134		arma::mat A = S; ///< working copy of S
135		eigvec.eye(p, p); ///< accumulate rotations; starts as identity
136
137		for (int sweep = 0; sweep < 50 * static_cast<int>(p); ++sweep) {
138		// Check off-diagonal convergence
139		double off = 0.0;
140		for (arma::uword q = 0; q < p - 1; ++q)
141		for (arma::uword r = q + 1; r < p; ++r)
142		off += A(q, r) * A(q, r);
143		if (std::sqrt(off) < 1e-13) break;
144		detail::jacobi_sweep(A, eigvec);
145		}
146		eigval = A.diag();
147
148		// Sort ascending
149		arma::uvec idx = arma::sort_index(eigval);
150		eigval = eigval(idx);
151		eigvec = eigvec.cols(idx);
152		#endif
153		}
154
155		/// @}
156
157		/// @name QR Factorization
158		/// @{
159
160		// ── qr_full ───────────────────────────────────────────────────────────────────
161
162		/// @brief Full (square) QR factorization of an n×p matrix.
163		///
164		/// Produces a square n×n orthogonal factor @p Q and an n×p upper-trapezoidal
165		/// factor @p R such that @c A = Q*R. This matches the convention of
166		/// @c arma::qr(), which returns the full (non-economy) @p Q.
167		///
168		/// @param[out] Q n×n orthogonal matrix.
169		/// @param[out] R n×p upper-trapezoidal matrix.
170		/// @param[in] A n×p input matrix (not modified).
171		///
172		/// @warning With @c ARMA_USE_LAPACK defined this calls @c arma::qr directly.
173		/// Without LAPACK the fallback constructs the full ONB by seeding
174		/// columns beyond @c A.n_cols with standard basis vectors and then
175		/// applying modified Gram-Schmidt; numerical quality degrades for
176		/// nearly rank-deficient inputs.
177		///
178		/// @note The caller (@c orthogonal_svd) depends on @p Q being n×n so that
179		/// @c Q.cols(k, n-1) spans the null space of @c A^T.
180	8x	inline void qr_full(arma::mat& Q, arma::mat& R, const arma::mat& A) {
181		#if defined(ARMA_USE_LAPACK)
182	8x	arma::qr(Q, R, A);
183		#else
184		// Build a full n×n orthogonal matrix.
185		// Step 1: thin QR via modified Gram-Schmidt (gives first k=A.n_cols columns).
186		// Step 2: extend to a full ONB using random vectors + re-orthogonalisation.
187		const arma::uword n = A.n_rows, p = A.n_cols;
188
189		// Start with all n columns of A padded with random vectors if n > p.
190		arma::mat Qfull(n, n, arma::fill::zeros);
191		// Copy A's columns first
192		for (arma::uword j = 0; j < p && j < n; ++j)
193		Qfull.col(j) = A.col(j);
194		// Fill remaining columns with standard basis vectors (or random fallback)
195		for (arma::uword j = p; j < n; ++j) {
196		Qfull(j, j) = 1.0; ///< standard basis vector e_j used as seed
197		}
198		// Full MGS orthonormalization
199		detail::mgs(Qfull);
200		Q = Qfull;
201
202		// R is n×p (upper trapezoidal); only first p rows are non-zero
203		R = Q.t() * A;
204		// Zero the strictly lower-triangular part of R[0:p, 0:p]
205		for (arma::uword j = 0; j < p; ++j)
206		for (arma::uword i = j + 1; i < n; ++i)
207		R(i, j) = 0.0;
208		#endif
209		}
210
211		/// @}
212
213		/// @name Singular Value Decomposition
214		/// @{
215
216		// ── svd ───────────────────────────────────────────────────────────────────────
217
218		/// @brief Full SVD: @c A ≈ U * diag(s) * V', singular values descending.
219		///
220		/// Decomposes the n×p matrix @p A into left singular vectors @p U, singular
221		/// values @p s, and right singular vectors @p V.
222		///
223		/// @param[out] U n×min(n,p) matrix of left singular vectors.
224		/// @param[out] s min(n,p)×1 vector of singular values, descending.
225		/// @param[out] V p×p matrix of right singular vectors (all p columns).
226		/// @param[in] A n×p input matrix (not modified).
227		///
228		/// @warning With @c ARMA_USE_LAPACK defined this calls @c arma::svd directly.
229		/// Without LAPACK the fallback computes @c A^T*A, calls @c eig_sym,
230		/// and derives @p U via @c U[:,j] = A*V[:,j]/s_j. Columns with
231		/// @c s_j < 1e-14 are zeroed. Accuracy is governed by the
232		/// conditioning of @c A^T*A (squares the condition number of @p A).
233		///
234		/// @note Fallback algorithm: eigendecomposition of \f$A^T A\f$ via cyclic
235		/// Jacobi (Demmel & Veselić 1992). Suitable for the small, moderately
236		/// conditioned matrices typical in ENA (p ≤ ~100).
237	40x	inline void svd(arma::mat& U, arma::vec& s, arma::mat& V, const arma::mat& A) {
238		#if defined(ARMA_USE_LAPACK)
239	40x	arma::svd(U, s, V, A);
240		#else
241		// One-sided Jacobi on A^T:
242		// We work with B = A^T (p×n). The right singular vectors of A are the
243		// left singular vectors of B (= columns of U_B after Jacobi on B^T B).
244		// We use the eigendecomposition of A^T * A to get V and s.
245		const arma::uword n = A.n_rows, p = A.n_cols;
246		const arma::uword k = std::min(n, p);
247
248		arma::mat AtA = A.t() * A; ///< p×p symmetric Gram matrix
249		arma::vec eigval;
250		arma::mat eigvec;
251		eig_sym(eigval, eigvec, AtA);
252
253		// Eigenvalues from eig_sym are ascending; we want descending.
254		arma::uvec idx = arma::sort_index(eigval, "descend");
255		eigval = eigval(idx);
256		eigvec = eigvec.cols(idx);
257
258		// Singular values (non-negative square roots)
259		s.set_size(k);
260		for (arma::uword j = 0; j < k; ++j)
261		s(j) = std::sqrt(std::max(0.0, eigval(j)));
262
263		// V = eigvec (right singular vectors of A are eigvectors of A^T*A)
264		V = eigvec;
265
266		// U = A * V / s (left singular vectors, only first k columns needed)
267		U.set_size(n, k);
268		for (arma::uword j = 0; j < k; ++j) {
269		if (s(j) > 1e-14)
270		U.col(j) = A * V.col(j) / s(j);
271		else
272		U.col(j).zeros();
273		}
274		#endif
275		}
276
277		// ── svd_right ─────────────────────────────────────────────────────────────────
278
279		/// @brief Singular values and right singular vectors only (left vectors skipped).
280		///
281		/// A lighter-weight variant of @c svd() for callers that only need @p s and
282		/// @p V. Under LAPACK the full @c arma::svd is still called internally; the
283		/// savings come in WASM builds where computing @p U is avoided.
284		///
285		/// @param[out] s min(n,p)×1 vector of singular values, descending.
286		/// @param[out] V p×p matrix of right singular vectors.
287		/// @param[in] A n×p input matrix (not modified).
288		///
289		/// @warning With @c ARMA_USE_LAPACK defined this still invokes the full
290		/// @c arma::svd; the left-vector output is discarded.
291		inline void svd_right(arma::vec& s, arma::mat& V, const arma::mat& A) {
292		#if defined(ARMA_USE_LAPACK)
293		arma::mat U;
294		arma::svd(U, s, V, A);
295		#else
296		arma::mat U_unused;
297		svd(U_unused, s, V, A);
298		#endif
299		}
300
301		/// @}
302
303		/// @name Linear System Solve
304		/// @{
305
306		// ── solve_spd ─────────────────────────────────────────────────────────────────
307
308		/// @brief Solve @c A*x = b for symmetric positive (semi-)definite @p A.
309		///
310		/// Returns the n×k solution matrix @p x given the n×n system matrix @p A and
311		/// the n×k right-hand side @p b. Near-zero eigenvalues are handled gracefully,
312		/// making this a pseudo-inverse solve safe for rank-deficient Gram matrices.
313		///
314		/// @param[in] A Real symmetric positive (semi-)definite n×n matrix.
315		/// @param[in] b n×k right-hand side matrix.
316		/// @param[in] tol Relative threshold below which eigenvalues are treated as
317		/// zero. Eigenvalues with @c \|λ\| < tol * max(\|λ\|) are
318		/// inverted as zero (i.e. those directions are projected out).
319		/// Default: 1e-12.
320		/// @returns n×k solution matrix @p x satisfying @c A*x ≈ b.
321		///
322		/// @warning With @c ARMA_USE_LAPACK defined this delegates to
323		/// @c arma::solve(A, b, arma::solve_opts::fast), which uses a
324		/// Cholesky factorization and does not handle rank deficiency.
325		/// Without LAPACK the eigendecomposition-based pseudo-inverse is used,
326		/// which is safe for rank-deficient inputs but slower.
327		///
328		/// @note Without LAPACK: @c A = V * diag(λ) * V^T is computed via @c eig_sym;
329		/// the solution is @c x = V * diag(1/λ) * V^T * b with small λ skipped.
330	10x	inline arma::mat solve_spd(const arma::mat& A, const arma::mat& b,
331		double tol = 1e-12) {
332		#if defined(ARMA_USE_LAPACK)
333	10x	return arma::solve(A, b, arma::solve_opts::fast);
334		#else
335		arma::vec eigval;
336		arma::mat eigvec;
337		eig_sym(eigval, eigvec, A); ///< A = eigvec * diag(eigval) * eigvec^T
338
339		const double thresh = tol * std::abs(eigval.max());
340		arma::vec inv_eigval(eigval.n_elem, arma::fill::zeros);
341		for (arma::uword i = 0; i < eigval.n_elem; ++i)
342		if (std::abs(eigval(i)) > thresh)
343		inv_eigval(i) = 1.0 / eigval(i);
344
345		// x = eigvec * diag(inv_eigval) * eigvec^T * b
346		return eigvec * arma::diagmat(inv_eigval) * eigvec.t() * b;
347		#endif
348		}
349
350		/// @}
351
352		} // namespace linalg
353		} // namespace qe
354
355		#endif // LIBQE_LINALG_FALLBACK_HPP

1		/**
2		* @file normalization.hpp
3		* @brief Row-wise normalization utilities for ENA connection networks.
4		*/
5		#ifndef LIBQE_NORMALIZATION_HPP
6		#define LIBQE_NORMALIZATION_HPP
7
8		#include <armadillo>
9		#include <cmath>
10
11		namespace qe {
12
13		/**
14		* @brief Normalize each row to unit L2 norm (project onto the unit hypersphere).
15		*
16		* Rows whose L2 norm is zero are left unchanged.
17		*
18		* @param m Input matrix (n_units × n_connections), row-major.
19		* @returns Matrix of the same shape with each row divided by its L2 norm.
20		* @note Equivalent to @c fun_sphere_norm() in rENA/ena.cpp.
21		*/
22	4x	inline arma::mat normalize_networks(arma::mat m) {
23	4x	arma::mat out(m.n_rows, m.n_cols, arma::fill::zeros);
24	14x	for (arma::uword r = 0; r < m.n_rows; r++) {
25	10x	double len = arma::norm(m.row(r), 2);
26	18x	if (std::isfinite(len) && len > 0.0) out.row(r) = m.row(r) / len;
27		}
28	4x	return out;
29		}
30
31		/**
32		* @brief Scale every entry by the reciprocal of the largest row L2 norm.
33		*
34		* Preserves relative magnitudes across rows; does not normalise each row
35		* individually. If the largest row norm is zero, the matrix is returned
36		* unchanged.
37		*
38		* @param m Input matrix (n_units × n_connections), modified in place.
39		* @returns Scaled matrix (same shape as @p m).
40		* @note Equivalent to @c fun_skip_sphere_norm() in rENA/ena.cpp.
41		*/
42	3x	inline arma::mat scale_networks(arma::mat m) {
43	3x	double largest = 0.0;
44	10x	for (arma::uword r = 0; r < m.n_rows; r++) {
45	7x	double len = arma::norm(m.row(r), 2);
46	7x	if (std::isfinite(len)) largest = std::max(largest, len);
47		}
48	3x	if (largest > 0.0) m /= largest;
49	6x	return m;
50		}
51
52		} // namespace qe
53
54		#endif // LIBQE_NORMALIZATION_HPP

1		/**
2		* @file rotation.hpp
3		* @brief Rotation routines for Epistemic Network Analysis (ENA):
4		* SVD, deflation, orthogonal-SVD, complete (generalized), and
5		* means rotation. All routines are designed to match the
6		* numerical output of the reference R implementation in rENA.
7		*/
8
9		#ifndef LIBQE_ROTATION_HPP
10		#define LIBQE_ROTATION_HPP
11
12		#include "linalg_fallback.hpp"
13		#include <armadillo>
14		#include <stdexcept>
15		#include <string>
16		#include <utility>
17		#include <vector>
18
19		namespace qe {
20
21		/// @name Return types
22		/// @{
23
24		/**
25		* @brief Aggregated result returned by every rotation routine.
26		*
27		* The layout mirrors the output of R's `prcomp()`: `rotation` holds the
28		* right-singular vectors (principal axes), `eigenvalues` holds the
29		* explained-variance values (sdev^2 in R parlance), and `column_names`
30		* carries the human-readable axis labels used downstream for network plots
31		* and loadings tables.
32		*/
33		struct RotationResult {
34		arma::mat rotation; ///< p x p orthogonal matrix; column j is rotation axis j.
35		arma::vec eigenvalues; ///< Length-p vector of sdev^2 values, matching rENA's convention.
36		std::vector<std::string> column_names; ///< Axis labels, e.g. "MR1", "SVD2", "GMR1".
37		};
38
39		/**
40		* @brief A pair of 0-based row-index vectors identifying two groups within a
41		* points matrix.
42		*
43		* Used by means_rotation() to specify which rows belong to group A and which
44		* belong to group B for each successive mean-difference axis.
45		*/
46		struct GroupPair {
47		arma::uvec a; ///< Row indices of group A (0-based).
48		arma::uvec b; ///< Row indices of group B (0-based).
49		};
50
51		/// @}
52
53		/// @name SVD rotation
54		/// @{
55
56		/**
57		* @brief Compute a full SVD rotation of a (pre-centered) points matrix.
58		*
59		* Produces the right-singular vectors V as the rotation matrix and scales
60		* each singular value to an eigenvalue (sdev^2) compatible with rENA's
61		* `prcomp()` output.
62		*
63		* @param points n x p data matrix. The caller is responsible for
64		* centering upstream; this function applies no centering.
65		*
66		* @returns A RotationResult where:
67		* - `rotation` = V (p x p full SVD; trailing columns span the null
68		* space when the matrix is rank-deficient).
69		* - `eigenvalues[j]` = s[j]^2 / max(1, n - 1) (== sdev^2 in R).
70		* - `column_names` = {"SVD1", "SVD2", ..., "SVDp"}.
71		*
72		* @note Equivalent to `prcomp(points, retx=FALSE, scale=FALSE,
73		* center=FALSE, tol=0)` in rENA/R.
74		*
75		* @note Sign convention: none. Signs follow the underlying LAPACK SVD,
76		* matching rENA's long-standing behavior. A deterministic sign rule
77		* (e.g. svd_flip) may be added later as an opt-in flag.
78		*/
79	40x	inline RotationResult ena_svd(const arma::mat& points) {
80	40x	if (!points.is_finite())
81	!	throw std::invalid_argument("ena_svd: input matrix must not contain NaN or Inf");
82	40x	arma::mat U, V;
83	40x	arma::vec s;
84	40x	qe::linalg::svd(U, s, V, points);
85
86	40x	const arma::uword p = points.n_cols;
87	80x	const double denom = points.n_rows > 1
88		? static_cast<double>(points.n_rows - 1)
89		: 1.0;
90
91	40x	arma::vec eigenvalues(p, arma::fill::zeros);
92	40x	const arma::uword k = std::min<arma::uword>(s.n_elem, p);
93	198x	for (arma::uword j = 0; j < k; ++j) {
94	474x	eigenvalues(j) = (s(j) * s(j)) / denom;
95		}
96
97	40x	std::vector<std::string> labels(p);
98	198x	for (arma::uword j = 0; j < p; ++j) {
99	158x	labels[j] = "SVD" + std::to_string(j + 1);
100		}
101	40x	return {V, eigenvalues, std::move(labels)};
102		}
103
104		/// @}
105
106		/// @name Deflation
107		/// @{
108
109		/**
110		* @brief Project a data matrix onto the hyperplane orthogonal to a given axis.
111		*
112		* Computes `data - (data * axis) * axis^T`, effectively removing the
113		* component of every row of `data` that lies along `axis`.
114		*
115		* @param data n x p data matrix.
116		* @param axis Unit-norm column vector of length p. The caller is
117		* responsible for normalizing `axis` before calling this
118		* function; no normalization is performed internally.
119		*
120		* @returns n x p matrix with the projection onto `axis` subtracted out.
121		*/
122	16x	inline arma::mat deflate(const arma::mat& data, const arma::vec& axis) {
123	32x	return data - (data * axis) * axis.t();
124		}
125
126		/// @}
127
128		/// @name Orthogonal SVD rotation
129		/// @{
130
131		/**
132		* @brief Orthonormalize named axes via QR, then complete the rotation with
133		* an SVD of the data projected onto the orthogonal complement.
134		*
135		* The named axes in `weights` are orthonormalized via a full QR
136		* decomposition. The first k columns of the resulting Q become the
137		* "fixed" part of the rotation; the remaining p - k columns define a
138		* complementary subspace onto which `data` is projected, and a further SVD
139		* fills those trailing rotation columns.
140		*
141		* @param data n x p data matrix (caller-centered if required).
142		* @param weights p x k matrix of named axis directions. Column norms
143		* need not be 1; QR handles normalization internally.
144		* @param named_labels Length-k vector of labels for the first k rotation
145		* axes. Trailing axes receive labels "SVD{k+1}" ..
146		* "SVDp".
147		*
148		* @returns A RotationResult containing the combined p x p rotation,
149		* eigenvalues for the SVD-filled trailing axes (indices k..p-1;
150		* the first k eigenvalues are zero), and axis labels.
151		*
152		* @note Equivalent to `orthogonal_svd()` in rENA/R/ena.rotate.by.mean.R:
153		* @code
154		* Q = qr.Q(qr(weights), complete = TRUE) # p x p
155		* X_bar = data %*% Q[, (k+1):p] # n x (p - k)
156		* V = prcomp(X_bar)$rotation
157		* out = cbind(Q[, 1:k], Q[, (k+1):p] %*% V)
158		* @endcode
159		*
160		* @note The named axes in the OUTPUT are the orthonormalized Q columns, NOT
161		* the original `weights` columns. Use complete_rotation() if you need
162		* to preserve input axes verbatim.
163		*
164		* @throws std::runtime_error if `data.n_cols != weights.n_rows` or
165		* `named_labels.size() != weights.n_cols`.
166		*/
167	10x	inline RotationResult orthogonal_svd(
168		const arma::mat& data,
169		const arma::mat& weights,
170		const std::vector<std::string>& named_labels) {
171	10x	const arma::uword p = weights.n_rows;
172	10x	const arma::uword k = weights.n_cols;
173	10x	if (data.n_cols != p) {
174		throw std::runtime_error(
175	1x	"orthogonal_svd: data.n_cols must equal weights.n_rows");
176		}
177	9x	if (named_labels.size() != k) {
178		throw std::runtime_error(
179	1x	"orthogonal_svd: named_labels.size() must equal weights.n_cols");
180		}
181	8x	if (k == 0) {
182	!	return ena_svd(data);
183		}
184
185	8x	arma::mat Q, R;
186	8x	qe::linalg::qr_full(Q, R, weights); // Q is p x p
187
188	8x	arma::mat rotation(p, p);
189	8x	arma::vec eigenvalues(p, arma::fill::zeros);
190	24x	rotation.cols(0, k - 1) = Q.cols(0, k - 1);
191
192	8x	if (k < p) {
193	16x	arma::mat Q_comp = Q.cols(k, p - 1); // p x (p - k)
194	8x	arma::mat X_bar = data * Q_comp; // n x (p - k)
195	8x	RotationResult inner = ena_svd(X_bar);
196	16x	rotation.cols(k, p - 1) = Q_comp * inner.rotation;
197	32x	for (arma::uword j = 0; j < (p - k); ++j) {
198	48x	eigenvalues(k + j) = inner.eigenvalues(j);
199		}
200		}
201
202	8x	std::vector<std::string> labels(p);
203	22x	for (arma::uword j = 0; j < k; ++j) labels[j] = named_labels[j];
204	32x	for (arma::uword j = k; j < p; ++j) {
205	24x	labels[j] = "SVD" + std::to_string(j + 1);
206		}
207	16x	return {rotation, eigenvalues, std::move(labels)};
208		}
209
210		/// @}
211
212		/// @name Complete (generalized) rotation
213		/// @{
214
215		/**
216		* @brief Keep named axes verbatim and fill remaining axes from an SVD of the
217		* data after parallel deflation by all named axes.
218		*
219		* Unlike orthogonal_svd(), this routine preserves the caller-supplied
220		* `named_axes` exactly (no QR re-orthonormalization). The complementary
221		* axes are obtained by deflating `data` simultaneously by all named axes and
222		* then running ena_svd() on the result.
223		*
224		* The deflation is parallel: every projection is subtracted from the
225		* original `data`, not from a progressively deflated copy:
226		* @code
227		* defA = data - data * (named_axes * named_axes^T)
228		* @endcode
229		* This matches `defA <- A - A %% v1 %% t(v1) - A %% v2 %% t(v2)` in
230		* rENA line-for-line. For mutually orthogonal axes the result is identical
231		* to sequential deflation; for non-orthogonal axes it differs.
232		*
233		* @param data n x p data matrix.
234		* @param named_axes p x k matrix of named rotation axes. Each column
235		* must be unit-norm; orthonormality between columns is
236		* NOT required.
237		* @param named_labels Length-k vector of labels for the first k axes.
238		* Trailing axes receive labels "SVD{k+1}" .. "SVDp".
239		* Conventional labels for generalized rotation are
240		* "GMR1", "GMR2", then "SVD{k+1}" onward, but labels
241		* are entirely caller-supplied.
242		*
243		* @returns A RotationResult where columns 0..k-1 of `rotation` equal
244		* `named_axes` verbatim and columns k..p-1 come from the SVD of
245		* the deflated data.
246		*
247		* @note Equivalent to the tail of `ena.rotate.by.generalized` in rENA
248		* (canonical version: commit 2c079126 on rENA `origin/main`).
249		*
250		* @throws std::runtime_error if `data.n_cols != named_axes.n_rows` or
251		* `named_labels.size() != named_axes.n_cols`.
252		*/
253	16x	inline RotationResult complete_rotation(
254		const arma::mat& data,
255		const arma::mat& named_axes,
256		const std::vector<std::string>& named_labels) {
257	16x	const arma::uword p = named_axes.n_rows;
258	16x	const arma::uword k = named_axes.n_cols;
259	16x	if (data.n_cols != p) {
260		throw std::runtime_error(
261	!	"complete_rotation: data.n_cols must equal named_axes.n_rows");
262		}
263	16x	if (named_labels.size() != k) {
264		throw std::runtime_error(
265	!	"complete_rotation: named_labels.size() must equal named_axes.n_cols");
266		}
267	16x	if (k == 0) {
268	!	return ena_svd(data);
269		}
270
271	32x	arma::mat defA = data - data * named_axes * named_axes.t();
272	16x	RotationResult inner = ena_svd(defA);
273
274	16x	arma::mat rotation(p, p);
275	16x	arma::vec eigenvalues(p, arma::fill::zeros);
276	32x	rotation.cols(0, k - 1) = named_axes;
277	16x	if (k < p) {
278	48x	rotation.cols(k, p - 1) = inner.rotation.cols(0, p - k - 1);
279	53x	for (arma::uword j = 0; j < (p - k); ++j) {
280	74x	eigenvalues(k + j) = inner.eigenvalues(j);
281		}
282		}
283
284	16x	std::vector<std::string> labels(p);
285	47x	for (arma::uword j = 0; j < k; ++j) labels[j] = named_labels[j];
286	53x	for (arma::uword j = k; j < p; ++j) {
287	37x	labels[j] = "SVD" + std::to_string(j + 1);
288		}
289	32x	return {rotation, eigenvalues, std::move(labels)};
290		}
291
292		/// @}
293
294		/// @name Means rotation
295		/// @{
296
297		/**
298		* @brief Compute a means rotation from one or more group pairs.
299		*
300		* For each GroupPair, computes the normalized mean-difference vector on the
301		* progressively deflated data and stacks the resulting unit-norm axes into a
302		* weights matrix. The final rotation is completed by orthogonal_svd().
303		* The input `points` is mean-centered before any axes are computed, matching
304		* rENA's `scale(data, scale=FALSE, center=TRUE)`.
305		*
306		* @param points n x p matrix of network positions (raw, not pre-centered).
307		* @param pairs Ordered list of GroupPair objects. Each pair contributes
308		* one "MR" axis in sequence. Must be non-empty.
309		*
310		* @returns A RotationResult with `column_names` = {"MR1", ..., "MRm",
311		* "SVD{m+1}", ..., "SVDp"} where m = pairs.size().
312		*
313		* @note Equivalent to `ena.rotate.by.mean()` in rENA/R/ena.rotate.by.mean.R.
314		* The progressive deflation (each axis is computed on the data after
315		* removing all previous axes) is the key numerical step.
316		*
317		* @warning There is NO guard against a zero-norm mean-difference vector.
318		* If the two group means are identical on the current deflated data,
319		* the axis will contain NaN values (division by zero). This mirrors
320		* a latent bug in rENA and is preserved verbatim for numerical
321		* fidelity; a safe opt-in guard will be added in a future release.
322		*
323		* @throws std::runtime_error if `pairs` is empty.
324		*/
325	8x	inline RotationResult means_rotation(const arma::mat& points,
326		const std::vector<GroupPair>& pairs) {
327	8x	if (pairs.empty()) {
328	1x	throw std::runtime_error("Unable to rotate without 2 groups.");
329		}
330
331	14x	arma::mat data = points.each_row() - arma::mean(points, 0);
332	7x	arma::mat deflated = data;
333	7x	arma::mat weights(data.n_cols, pairs.size(), arma::fill::zeros);
334	7x	std::vector<std::string> labels;
335	7x	labels.reserve(pairs.size());
336
337	19x	for (std::size_t i = 0; i < pairs.size(); ++i) {
338	24x	const arma::rowvec mean_a = arma::mean(deflated.rows(pairs[i].a), 0);
339	24x	const arma::rowvec mean_b = arma::mean(deflated.rows(pairs[i].b), 0);
340	24x	const arma::vec diff = (mean_a - mean_b).t();
341	12x	const double norm = std::sqrt(arma::accu(diff % diff));
342	12x	const arma::vec axis = diff / norm; // no zero-norm guard
343
344	12x	deflated = deflate(deflated, axis);
345	24x	weights.col(i) = axis;
346	12x	labels.push_back("MR" + std::to_string(i + 1));
347		}
348
349	14x	return orthogonal_svd(deflated, weights, labels);
350		}
351
352		/// @}
353
354		} // namespace qe
355
356		#endif // LIBQE_ROTATION_HPP

1		/** @file stability.hpp
2		* @brief Stability metrics and distance-distance matrix correlation for parameter tuning.
3		*/
4		#ifndef LIBQE_STABILITY_HPP
5		#define LIBQE_STABILITY_HPP
6
7		#include <armadillo>
8		#include <cmath>
9		#include <algorithm>
10
11		namespace qe {
12
13		/** @brief Compute Pearson correlation between pairwise Euclidean distance structures of two matrices.
14		*
15		* Evaluates dist-dist correlation: r(dist(X), dist(Y)), measuring whether relative
16		* neighborhoods and distances are preserved across parameter changes (e.g. window vs lookback).
17		*
18		* @param[in] X Numeric matrix (n_rows × n_cols).
19		* @param[in] Y Numeric matrix (n_rows × n_cols).
20		*
21		* @returns Pearson correlation coefficient in [-1, 1].
22		*/
23	1x	inline double dist_dist_correlation(const arma::mat& X, const arma::mat& Y) {
24	1x	size_t n = X.n_rows;
25		if (n != Y.n_rows \|\| n < 3) return 0.0;
26
27	1x	size_t n_pairs = n * (n - 1) / 2;
28	1x	arma::vec dx(n_pairs);
29	1x	arma::vec dy(n_pairs);
30
31	1x	size_t idx = 0;
32	11x	for (size_t i = 0; i < n; ++i) {
33	55x	for (size_t j = i + 1; j < n; ++j) {
34	180x	arma::rowvec diff_x = X.row(i) - X.row(j);
35	180x	arma::rowvec diff_y = Y.row(i) - Y.row(j);
36	90x	dx[idx] = std::sqrt(arma::dot(diff_x, diff_x));
37	45x	dy[idx] = std::sqrt(arma::dot(diff_y, diff_y));
38	45x	idx++;
39		}
40		}
41
42	1x	double mean_x = arma::mean(dx);
43	1x	double mean_y = arma::mean(dy);
44
45	1x	arma::vec zx = dx - mean_x;
46	1x	arma::vec zy = dy - mean_y;
47
48	1x	double sxx = arma::dot(zx, zx);
49	1x	double syy = arma::dot(zy, zy);
50	1x	double sxy = arma::dot(zx, zy);
51
52		if (sxx <= 1e-12 \|\| syy <= 1e-12) return 0.0;
53	1x	return sxy / std::sqrt(sxx * syy);
54		}
55
56		} // namespace qe
57
58		#endif // LIBQE_STABILITY_HPP

1		/** @file trajectory_distance.hpp
2		* @brief Integrated Euclidean distance between trajectory curves and lagged distance metrics.
3		*/
4		#ifndef LIBQE_TRAJECTORY_DISTANCE_HPP
5		#define LIBQE_TRAJECTORY_DISTANCE_HPP
6
7		#include <armadillo>
8		#include <vector>
9		#include <cmath>
10		#include <algorithm>
11		#include "trajectory.hpp"
12
13		namespace qe {
14
15		struct GaussLegendreRule {
16		arma::vec nodes;
17		arma::vec weights;
18		};
19
20		/** @brief Obtain exact 32-point Gauss-Legendre quadrature nodes and weights on [-1, 1] via Golub-Welsch. */
21	2x	inline const GaussLegendreRule& get_gl_rule_32() {
22	1x	static const GaussLegendreRule rule = []() {
23	1x	int N = 32;
24	1x	arma::mat J(N, N, arma::fill::zeros);
25	32x	for (int k = 1; k < N; ++k) {
26	31x	double b = static_cast<double>(k) / std::sqrt(4.0 * k * k - 1.0);
27	31x	J(k, k - 1) = b;
28	62x	J(k - 1, k) = b;
29		}
30	1x	arma::vec evals;
31	1x	arma::mat evecs;
32	1x	arma::eig_sym(evals, evecs, J);
33	1x	arma::vec w = 2.0 * arma::square(evecs.row(0).t());
34	2x	return GaussLegendreRule{evals, w};
35		}();
36	2x	return rule;
37		}
38
39		/** @brief Compute integrated Euclidean distance between two 2D parametric polynomial curves.
40		*
41		* Computes integral_{t_start}^{t_end} \|\|F_a(t) - F_b(t)\|\| dt
42		* using 32-point Gauss-Legendre quadrature.
43		*
44		* @param[in] coeffs_ax Coefficients of x_a(t).
45		* @param[in] coeffs_ay Coefficients of y_a(t).
46		* @param[in] coeffs_bx Coefficients of x_b(t).
47		* @param[in] coeffs_by Coefficients of y_b(t).
48		* @param[in] t_start Start of integration interval (default 0.0).
49		* @param[in] t_end End of integration interval (default 1.0).
50		*
51		* @returns Integrated distance scalar.
52		*/
53	2x	inline double integrated_curve_distance(
54		const arma::vec& coeffs_ax,
55		const arma::vec& coeffs_ay,
56		const arma::vec& coeffs_bx,
57		const arma::vec& coeffs_by,
58		double t_start = 0.0,
59		double t_end = 1.0
60		) {
61		if (t_end <= t_start) return 0.0;
62
63	2x	const auto& rule = get_gl_rule_32();
64	2x	double mid = 0.5 * (t_start + t_end);
65	2x	double half_len = 0.5 * (t_end - t_start);
66	2x	double total = 0.0;
67
68	66x	for (size_t i = 0; i < rule.nodes.n_elem; ++i) {
69	64x	double t_eval = mid + half_len * rule.nodes[i];
70	64x	double xa = eval_poly_scalar(coeffs_ax, t_eval);
71	64x	double ya = eval_poly_scalar(coeffs_ay, t_eval);
72	64x	double xb = eval_poly_scalar(coeffs_bx, t_eval);
73	64x	double yb = eval_poly_scalar(coeffs_by, t_eval);
74	64x	double dx = xa - xb;
75	64x	double dy = ya - yb;
76	64x	double dist = std::sqrt(dx * dx + dy * dy);
77
78	128x	total += rule.weights[i] * dist;
79		}
80
81	2x	return half_len * total;
82		}
83
84		/** @brief Compute lagged curve distance: 1/(1-lag) * integral_{lag}^1 \|\|F_follower(t) - F_leader(t - lag)\|\| dt. */
85	!	inline double lagged_curve_distance(
86		const arma::vec& coeffs_fol_x,
87		const arma::vec& coeffs_fol_y,
88		const arma::vec& coeffs_ldr_x,
89		const arma::vec& coeffs_ldr_y,
90		double lag = 0.0
91		) {
92	!	if (lag < 0.0 \|\| lag >= 0.999) return 0.0;
93	!	double t_start = lag;
94	!	double t_end = 1.0;
95	!	const auto& rule = get_gl_rule_32();
96	!	double mid = 0.5 * (t_start + t_end);
97	!	double half_len = 0.5 * (t_end - t_start);
98	!	double total = 0.0;
99
100	!	for (size_t i = 0; i < rule.nodes.n_elem; ++i) {
101	!	double t_fol = mid + half_len * rule.nodes[i];
102	!	double t_ldr = t_fol - lag;
103
104	!	double x_fol = eval_poly_scalar(coeffs_fol_x, t_fol);
105	!	double y_fol = eval_poly_scalar(coeffs_fol_y, t_fol);
106	!	double x_ldr = eval_poly_scalar(coeffs_ldr_x, t_ldr);
107	!	double y_ldr = eval_poly_scalar(coeffs_ldr_y, t_ldr);
108
109	!	double dx = x_fol - x_ldr;
110	!	double dy = y_fol - y_ldr;
111	!	double dist = std::sqrt(dx * dx + dy * dy);
112
113	!	total += rule.weights[i] * dist;
114		}
115
116	!	double integral = half_len * total;
117	!	double interval = 1.0 - lag;
118	!	return (interval > 1e-6) ? (integral / interval) : 0.0;
119		}
120
121		/** @brief Compute pairwise integrated distance matrix across K polynomial trajectory curves. */
122	!	inline arma::mat pairwise_trajectory_distance_matrix(
123		const std::vector<arma::vec>& all_coeffs_x,
124		const std::vector<arma::vec>& all_coeffs_y
125		) {
126	!	size_t K = all_coeffs_x.size();
127	!	arma::mat D(K, K, arma::fill::zeros);
128
129	!	for (size_t i = 0; i < K; ++i) {
130	!	for (size_t j = i + 1; j < K; ++j) {
131	!	double dist = integrated_curve_distance(
132	!	all_coeffs_x[i], all_coeffs_y[i],
133	!	all_coeffs_x[j], all_coeffs_y[j]
134		);
135	!	D(i, j) = dist;
136	!	D(j, i) = dist;
137		}
138		}
139
140	!	return D;
141		}
142
143		} // namespace qe
144
145		#endif // LIBQE_TRAJECTORY_DISTANCE_HPP

1		/** @file trajectory_following.hpp
2		* @brief Who-follows-whom leader-follower turn lag analysis along shared timeline.
3		*/
4		#ifndef LIBQE_TRAJECTORY_FOLLOWING_HPP
5		#define LIBQE_TRAJECTORY_FOLLOWING_HPP
6
7		#include <armadillo>
8		#include <vector>
9		#include <cmath>
10		#include <algorithm>
11		#include <limits>
12
13		namespace qe {
14
15		/** @brief Result of signed turn-lag sweep across delta in [-max_lag, max_lag]. */
16		struct SignedTurnLagResult {
17		int best_lag = 0; ///< delta* that minimizes mean distance
18		double min_mean_distance = 0.0;
19		arma::ivec lags; ///< Range of evaluated lags
20		arma::vec mean_distances; ///< Mean distance for each lag
21		arma::ivec valid_counts; ///< Number of paired turns for each lag
22		};
23
24		/** @brief Compute mean distance between agent A at turn T and agent B at turn T - delta.
25		*
26		* @param[in] pts_a n_a × 2 coordinates for agent A.
27		* @param[in] pts_b n_b × 2 coordinates for agent B.
28		* @param[in] times_a Length n_a vector of calendar/turn indices for A.
29		* @param[in] times_b Length n_b vector of calendar/turn indices for B.
30		* @param[in] delta Signed lag in turns (delta > 0: A compares to B's past => A follows B;
31		* delta < 0: B compares to A's past => B follows A; delta = 0: sync).
32		*
33		* @returns Pair of (mean_distance, valid_pairs_count).
34		*/
35	9x	inline std::pair<double, int> signed_turn_lag_distance(
36		const arma::mat& pts_a,
37		const arma::mat& pts_b,
38		const arma::vec& times_a,
39		const arma::vec& times_b,
40		int delta
41		) {
42	9x	int n_a = static_cast<int>(pts_a.n_rows);
43	9x	int n_b = static_cast<int>(pts_b.n_rows);
44		if (n_a == 0 \|\| n_b == 0) return {0.0, 0};
45
46	9x	double sum_dist = 0.0;
47	9x	int count = 0;
48
49	54x	for (int i = 0; i < n_a; ++i) {
50	45x	double t_a = times_a[i];
51	45x	double target_tb = t_a - static_cast<double>(delta);
52
53		// Find the latest turn for agent B that occurred at or before target_tb
54	45x	int best_b_idx = -1;
55	216x	for (int j = 0; j < n_b; ++j) {
56	382x	if (times_b[j] <= target_tb) {
57	171x	best_b_idx = j;
58		} else {
59	20x	break;
60		}
61		}
62
63	45x	if (best_b_idx >= 0) {
64	84x	double dx = pts_a(i, 0) - pts_b(best_b_idx, 0);
65	84x	double dy = pts_a(i, 1) - pts_b(best_b_idx, 1);
66	42x	sum_dist += std::sqrt(dx * dx + dy * dy);
67	42x	count++;
68		}
69		}
70
71		double mean_d = (count > 0) ? (sum_dist / count) : 0.0;
72	9x	return {mean_d, count};
73		}
74
75		/** @brief Sweep signed turn lags in [-max_lag, max_lag] to determine leader-follower dynamics. */
76	1x	inline SignedTurnLagResult best_signed_turn_lag(
77		const arma::mat& pts_a,
78		const arma::mat& pts_b,
79		const arma::vec& times_a,
80		const arma::vec& times_b,
81		int max_lag = 15
82		) {
83	1x	int k_max = std::max(1, max_lag);
84	1x	int n_lags = 2 * k_max + 1;
85
86	1x	SignedTurnLagResult res;
87	1x	res.lags.set_size(n_lags);
88	1x	res.mean_distances.set_size(n_lags);
89	1x	res.valid_counts.set_size(n_lags);
90
91	1x	double min_dist = std::numeric_limits<double>::infinity();
92	1x	int best_delta = 0;
93
94	10x	for (int idx = 0; idx < n_lags; ++idx) {
95	9x	int delta = -k_max + idx;
96	9x	res.lags[idx] = delta;
97
98	9x	auto [mean_d, count] = signed_turn_lag_distance(pts_a, pts_b, times_a, times_b, delta);
99	9x	res.mean_distances[idx] = mean_d;
100	9x	res.valid_counts[idx] = count;
101
102	9x	if (count > 0 && mean_d < min_dist) {
103	5x	min_dist = mean_d;
104	5x	best_delta = delta;
105		}
106		}
107
108	1x	res.best_lag = best_delta;
109		res.min_mean_distance = (min_dist == std::numeric_limits<double>::infinity()) ? 0.0 : min_dist;
110	1x	return res;
111		}
112
113		} // namespace qe
114
115		#endif // LIBQE_TRAJECTORY_FOLLOWING_HPP

1		/** @file trajectory.hpp
2		* @brief Parametric trajectory curve fitting, LOOCV-2D degree selection,
3		* and differential geometry (velocity, acceleration, curvature, turns).
4		*/
5		#ifndef LIBQE_TRAJECTORY_HPP
6		#define LIBQE_TRAJECTORY_HPP
7
8		#include <armadillo>
9		#include <vector>
10		#include <string>
11		#include <cmath>
12		#include <algorithm>
13		#include <limits>
14
15		namespace qe {
16
17		/** @brief Polynomial curve fit result in 2D epistemic space. */
18		struct PolyCurveFit {
19		int degree = 1;
20		arma::vec coeffs_x; ///< Raw coefficients [c0, ..., cd] for x(t) = sum(ck * t^k)
21		arma::vec coeffs_y; ///< Raw coefficients [c0, ..., cd] for y(t) = sum(ck * t^k)
22		std::string basis = "orthogonal"; ///< Basis used during fitting: "orthogonal" (R stats::poly-like) or "raw".
23		arma::vec basis_alpha; ///< Orthogonal-basis recurrence alpha values (empty for raw basis).
24		arma::vec basis_norm2; ///< Orthogonal-basis recurrence squared norms, matching stats::poly coefs$norm2.
25		double cv_error = 0.0; ///< 2D LOOCV error
26		double aic = 0.0; ///< Joint bivariate AIC
27		arma::vec t; ///< Evaluated or input time points
28		arma::mat fitted; ///< n × 2 fitted coordinates
29		};
30
31	5x	inline arma::mat build_raw_poly_design(const arma::vec& t_vec, int degree) {
32	5x	int n = static_cast<int>(t_vec.n_elem);
33	5x	arma::mat X(n, degree + 1);
34	50x	for (int i = 0; i < n; ++i) {
35	45x	double cur_t = t_vec[i];
36	45x	double p = 1.0;
37	183x	for (int k = 0; k <= degree; ++k) {
38	138x	X(i, k) = p;
39	138x	p *= cur_t;
40		}
41		}
42	5x	return X;
43		}
44
45	5x	inline arma::vec poly_multiply_by_linear(const arma::vec& coeffs, double shift) {
46	5x	arma::vec out(coeffs.n_elem + 1, arma::fill::zeros);
47	17x	for (size_t k = 0; k < coeffs.n_elem; ++k) {
48	24x	out[k] += -shift * coeffs[k];
49	36x	out[k + 1] += coeffs[k];
50		}
51	5x	return out;
52		}
53
54	4x	inline arma::mat orthogonal_poly_basis_from_coefs(
55		const arma::vec& t_vec,
56		int degree,
57		const arma::vec& alpha,
58		const arma::vec& norm2
59		) {
60	4x	int n = static_cast<int>(t_vec.n_elem);
61	4x	arma::mat Z(n, degree + 1, arma::fill::ones);
62	4x	if (degree >= 1) {
63	12x	Z.col(1) = t_vec - alpha[0];
64		}
65	9x	for (int i = 2; i <= degree; ++i) {
66	30x	Z.col(i) = (t_vec - alpha[i - 1]) % Z.col(i - 1) -
67	25x	(norm2[i] / norm2[i - 1]) * Z.col(i - 2);
68		}
69	17x	for (int i = 0; i <= degree; ++i) {
70	39x	Z.col(i) /= std::sqrt(norm2[i + 1]);
71		}
72	12x	return Z.cols(1, degree);
73		}
74
75	4x	inline arma::mat orthogonal_basis_raw_coefficients(
76		int degree,
77		const arma::vec& alpha,
78		const arma::vec& norm2
79		) {
80	4x	arma::mat raw(degree + 1, degree + 1, arma::fill::zeros);
81	4x	raw(0, 0) = 1.0;
82	4x	if (degree >= 1) {
83	4x	arma::vec p1(2);
84	12x	p1[0] = -alpha[0];
85	4x	p1[1] = 1.0;
86	12x	raw.submat(0, 1, 1, 1) = p1 / std::sqrt(norm2[2]);
87		}
88	4x	arma::vec prev2(1);
89	4x	prev2[0] = 1.0;
90	4x	arma::vec prev1;
91	4x	if (degree >= 1) {
92	4x	prev1.set_size(2);
93	12x	prev1[0] = -alpha[0];
94	4x	prev1[1] = 1.0;
95		}
96	9x	for (int i = 2; i <= degree; ++i) {
97	10x	arma::vec term = poly_multiply_by_linear(prev1, alpha[i - 1]);
98	5x	arma::vec padded_prev2(term.n_elem, arma::fill::zeros);
99	10x	padded_prev2.subvec(0, prev2.n_elem - 1) = prev2;
100	25x	arma::vec current = term - (norm2[i] / norm2[i - 1]) * padded_prev2;
101	20x	raw.submat(0, i, i, i) = current / std::sqrt(norm2[i + 1]);
102	5x	prev2 = prev1;
103	5x	prev1 = current;
104		}
105	8x	return raw;
106		}
107
108	4x	inline bool modified_gram_schmidt_qr(
109		const arma::mat& X,
110		arma::mat& Q,
111		arma::mat& R
112		) {
113	4x	const int n = static_cast<int>(X.n_rows);
114	4x	const int p = static_cast<int>(X.n_cols);
115	4x	Q.zeros(n, p);
116	4x	R.zeros(p, p);
117
118	17x	for (int k = 0; k < p; ++k) {
119	26x	arma::vec v = X.col(k);
120	29x	for (int j = 0; j < k; ++j) {
121	48x	R(j, k) = arma::dot(Q.col(j), v);
122	64x	v -= R(j, k) * Q.col(j);
123		}
124
125	13x	double norm = std::sqrt(arma::dot(v, v));
126	13x	if (norm <= 1e-12) {
127	!	return false;
128		}
129
130	26x	R(k, k) = norm;
131	26x	Q.col(k) = v / norm;
132		}
133
134	4x	return true;
135		}
136
137	4x	inline arma::mat build_orthogonal_poly_design(
138		const arma::vec& t_vec,
139		int degree,
140		arma::vec& alpha,
141		arma::vec& norm2
142		) {
143	4x	int n = static_cast<int>(t_vec.n_elem);
144	4x	if (degree < 1) {
145	!	return arma::mat(n, 0);
146		}
147	4x	arma::vec unique_t = arma::unique(t_vec);
148	4x	if (static_cast<size_t>(degree) >= unique_t.n_elem) {
149	!	throw std::invalid_argument("orthogonal polynomial degree must be less than number of unique time points.");
150		}
151
152	4x	double xbar = arma::mean(t_vec);
153	4x	arma::vec centered = t_vec - xbar;
154	4x	arma::mat X = build_raw_poly_design(centered, degree);
155
156	4x	arma::mat Q;
157	4x	arma::mat R;
158	4x	bool ok = modified_gram_schmidt_qr(X, Q, R);
159	4x	if (!ok) {
160	!	throw std::invalid_argument("orthogonal polynomial degree must be less than number of unique time points.");
161		}
162
163	8x	arma::mat Z = Q * arma::diagmat(R.diag());
164	4x	arma::vec z_norm2 = arma::sum(arma::square(Z), 0).t();
165
166	4x	alpha.set_size(degree);
167	13x	for (int i = 0; i < degree; ++i) {
168	36x	alpha[i] = arma::dot(centered, arma::square(Z.col(i))) / z_norm2[i] + xbar;
169		}
170
171	4x	norm2.set_size(degree + 2);
172	4x	norm2[0] = 1.0;
173	17x	for (int i = 0; i <= degree; ++i) {
174	39x	norm2[i + 1] = z_norm2[i];
175		}
176
177	4x	return orthogonal_poly_basis_from_coefs(t_vec, degree, alpha, norm2);
178		}
179
180	5x	inline arma::mat build_poly_design(
181		const arma::vec& t_vec,
182		int degree,
183		const std::string& basis,
184		arma::vec& alpha,
185		arma::vec& norm2,
186		arma::mat& basis_raw_coeffs
187		) {
188	5x	int n = static_cast<int>(t_vec.n_elem);
189	5x	arma::mat X(n, degree + 1);
190	5x	X.col(0).ones();
191	5x	if (basis == "orthogonal") {
192	4x	arma::mat Z = build_orthogonal_poly_design(t_vec, degree, alpha, norm2);
193	8x	X.cols(1, degree) = Z;
194	4x	basis_raw_coeffs = orthogonal_basis_raw_coefficients(degree, alpha, norm2);
195	5x	} else if (basis == "raw") {
196	1x	X = build_raw_poly_design(t_vec, degree);
197	1x	alpha.reset();
198	1x	norm2.reset();
199	1x	basis_raw_coeffs.eye(degree + 1, degree + 1);
200		} else {
201	!	throw std::invalid_argument("basis must be 'orthogonal' or 'raw'.");
202		}
203	5x	return X;
204		}
205
206		/** @brief Differential geometry profile along a trajectory curve. */
207		struct TrajectoryDerivatives {
208		arma::vec t;
209		arma::vec vx; ///< dx/dt
210		arma::vec vy; ///< dy/dt
211		arma::vec speed; ///< sqrt(vx^2 + vy^2)
212		arma::vec ax; ///< d^2x/dt^2
213		arma::vec ay; ///< d^2y/dt^2
214		arma::vec heading_rate; ///< omega = vxay - vyax (angular turning velocity)
215		arma::vec curvature; ///< kappa = omega / max(speed^3, eps)
216		};
217
218		/** @brief Evaluate a 1D polynomial sum(c_k * t^k) at a single time t. */
219	296x	inline double eval_poly_scalar(const arma::vec& coeffs, double t) {
220	296x	double val = 0.0;
221	296x	double t_pow = 1.0;
222	956x	for (size_t k = 0; k < coeffs.n_elem; ++k) {
223	660x	val += coeffs[k] * t_pow;
224	660x	t_pow *= t;
225		}
226	296x	return val;
227		}
228
229		/** @brief Evaluate polynomial 1st derivative sum(k * c_k * t^(k-1)) at a single time t. */
230	6x	inline double eval_poly_deriv1_scalar(const arma::vec& coeffs, double t) {
231	6x	double val = 0.0;
232	6x	double t_pow = 1.0;
233	15x	for (size_t k = 1; k < coeffs.n_elem; ++k) {
234	9x	val += static_cast<double>(k) * coeffs[k] * t_pow;
235	9x	t_pow *= t;
236		}
237	6x	return val;
238		}
239
240		/** @brief Evaluate polynomial 2nd derivative sum(k(k-1) c_k * t^(k-2)) at a single time t. */
241	6x	inline double eval_poly_deriv2_scalar(const arma::vec& coeffs, double t) {
242	6x	double val = 0.0;
243	6x	double t_pow = 1.0;
244	9x	for (size_t k = 2; k < coeffs.n_elem; ++k) {
245	3x	val += static_cast<double>(k * (k - 1)) * coeffs[k] * t_pow;
246	3x	t_pow *= t;
247		}
248	6x	return val;
249		}
250
251		/** @brief Evaluate 2D parametric curve (x(t), y(t)) at multiple time points. */
252	2x	inline arma::mat eval_poly_curve(
253		const arma::vec& coeffs_x,
254		const arma::vec& coeffs_y,
255		const arma::vec& t_eval
256		) {
257	2x	size_t n = t_eval.n_elem;
258	2x	arma::mat out(n, 2);
259	22x	for (size_t i = 0; i < n; ++i) {
260	20x	double t = t_eval[i];
261	20x	out(i, 0) = eval_poly_scalar(coeffs_x, t);
262	40x	out(i, 1) = eval_poly_scalar(coeffs_y, t);
263		}
264	2x	return out;
265		}
266
267		/** @brief Fit a 2D polynomial curve with automatic LOOCV-2D or AIC degree selection.
268		*
269		* Given coordinates (x_i, y_i) and times t_i in [0, 1], constructs polynomial models
270		* for degrees d in [1, max_degree]. LOOCV error is computed via the closed-form hat-matrix formula:
271		* CV(d) = sum( (r_{x,i} / (1 - h_{ii}))^2 + (r_{y,i} / (1 - h_{ii}))^2 )
272		*
273		* @param[in] points n × 2 matrix of trajectory coordinates.
274		* @param[in] t Vector of length n of strictly increasing times (default [0, 1] linearly spaced).
275		* @param[in] max_degree Maximum polynomial degree to evaluate (default 3).
276		* @param[in] fixed_degree When >= 1, uses this exact degree without search.
277		* @param[in] criterion "loocv" (default) or "aic".
278		*
279		* @returns PolyCurveFit containing optimal degree, coefficients, errors, and fitted points.
280		*/
281	3x	inline PolyCurveFit fit_poly_loocv2d(
282		const arma::mat& points,
283		const arma::vec& t = arma::vec(),
284		int max_degree = 3,
285		int fixed_degree = 0,
286		const std::string& criterion = "loocv",
287		const std::string& basis = "orthogonal"
288		) {
289	3x	int n = static_cast<int>(points.n_rows);
290	3x	if (n < 2 \|\| points.n_cols < 2) {
291	!	throw std::invalid_argument("fit_poly_loocv2d requires at least 2 points and 2 columns.");
292		}
293
294		// Default normalized time t in [0, 1]
295	3x	arma::vec t_vec;
296	3x	if (t.n_elem == static_cast<size_t>(n)) {
297	3x	t_vec = t;
298		} else {
299	!	t_vec = arma::linspace<arma::vec>(0.0, 1.0, n);
300		}
301
302	3x	arma::vec x = points.col(0);
303	3x	arma::vec y = points.col(1);
304
305	3x	int deg_upper = std::min(max_degree, std::max(1, n - 2));
306	3x	if (fixed_degree >= 1) {
307	2x	deg_upper = std::min(fixed_degree, n - 1);
308		}
309
310	3x	int best_deg = 1;
311	3x	double best_score = std::numeric_limits<double>::infinity();
312	3x	PolyCurveFit best_fit;
313
314	3x	int deg_start = (fixed_degree >= 1) ? fixed_degree : 1;
315
316	8x	for (int d = deg_start; d <= deg_upper; ++d) {
317	5x	int n_params = d + 1;
318		if (n_params > n) break;
319
320	5x	arma::vec basis_alpha;
321	5x	arma::vec basis_norm2;
322	5x	arma::mat basis_raw_coeffs;
323	5x	arma::mat X = build_poly_design(t_vec, d, basis, basis_alpha, basis_norm2, basis_raw_coeffs);
324
325		// Solve least squares: beta = (X^T X)^(-1) X^T y
326	5x	arma::mat XtX = X.t() * X;
327	5x	arma::mat XtX_inv;
328	5x	bool inv_ok = arma::inv(XtX_inv, XtX);
329	5x	if (!inv_ok) {
330		// Pseudo-inverse fallback
331	!	XtX_inv = arma::pinv(XtX);
332		}
333
334	5x	arma::vec beta_x = XtX_inv * (X.t() * x);
335	5x	arma::vec beta_y = XtX_inv * (X.t() * y);
336
337	5x	arma::vec pred_x = X * beta_x;
338	5x	arma::vec pred_y = X * beta_y;
339
340	10x	arma::vec rx = x - pred_x;
341	5x	arma::vec ry = y - pred_y;
342
343		// Hat matrix diagonal: h_ii = x_i^T (X^T X)^(-1) x_i
344	5x	arma::vec h(n);
345	50x	for (int i = 0; i < n; ++i) {
346	90x	arma::rowvec xi = X.row(i);
347	90x	h[i] = arma::as_scalar(xi * XtX_inv * xi.t());
348	90x	h[i] = std::min(0.9999, std::max(0.0, h[i]));
349		}
350
351		// 2D LOOCV error
352	5x	double cv_err = 0.0;
353	50x	for (int i = 0; i < n; ++i) {
354	45x	double denom = 1.0 - h[i];
355	45x	double rxi = rx[i] / denom;
356	45x	double ryi = ry[i] / denom;
357	45x	cv_err += (rxi * rxi + ryi * ryi);
358		}
359
360		// Bivariate AIC
361	5x	double rss = arma::dot(rx, rx) + arma::dot(ry, ry);
362	5x	double aic = static_cast<double>(n) * std::log(std::max(1e-12, rss / (2.0 * n))) + 2.0 * (2.0 * n_params);
363
364		double score = (criterion == "aic") ? aic : cv_err;
365
366	5x	if (score < best_score \|\| d == deg_start) {
367	3x	best_score = score;
368	3x	best_deg = d;
369	3x	best_fit.degree = d;
370	3x	best_fit.coeffs_x = basis_raw_coeffs * beta_x;
371	3x	best_fit.coeffs_y = basis_raw_coeffs * beta_y;
372	3x	best_fit.basis = basis;
373	3x	best_fit.basis_alpha = basis_alpha;
374	3x	best_fit.basis_norm2 = basis_norm2;
375	3x	best_fit.cv_error = cv_err;
376	3x	best_fit.aic = aic;
377	3x	best_fit.t = t_vec;
378	3x	best_fit.fitted = arma::join_rows(pred_x, pred_y);
379		}
380		}
381
382	6x	return best_fit;
383		}
384
385		/** @brief Compute trajectory velocities, speeds, accelerations, heading turn rates, and curvatures. */
386	1x	inline TrajectoryDerivatives eval_trajectory_derivatives(
387		const arma::vec& coeffs_x,
388		const arma::vec& coeffs_y,
389		const arma::vec& t_eval
390		) {
391	1x	size_t n = t_eval.n_elem;
392	1x	TrajectoryDerivatives d;
393	1x	d.t = t_eval;
394	1x	d.vx.set_size(n);
395	1x	d.vy.set_size(n);
396	1x	d.speed.set_size(n);
397	1x	d.ax.set_size(n);
398	1x	d.ay.set_size(n);
399	1x	d.heading_rate.set_size(n);
400	1x	d.curvature.set_size(n);
401
402	4x	for (size_t i = 0; i < n; ++i) {
403	3x	double t = t_eval[i];
404	3x	double vx = eval_poly_deriv1_scalar(coeffs_x, t);
405	3x	double vy = eval_poly_deriv1_scalar(coeffs_y, t);
406	3x	double ax = eval_poly_deriv2_scalar(coeffs_x, t);
407	3x	double ay = eval_poly_deriv2_scalar(coeffs_y, t);
408
409	3x	double spd = std::sqrt(vx * vx + vy * vy);
410	3x	double omega = vx * ay - vy * ax; // angular velocity / heading rate
411	3x	double spd3 = spd * spd * spd;
412		double kappa = (spd3 > 1e-12) ? (omega / spd3) : 0.0;
413
414	3x	d.vx[i] = vx;
415	3x	d.vy[i] = vy;
416	3x	d.speed[i] = spd;
417	3x	d.ax[i] = ax;
418	3x	d.ay[i] = ay;
419	3x	d.heading_rate[i] = omega;
420	6x	d.curvature[i] = kappa;
421		}
422
423	1x	return d;
424		}
425
426		} // namespace qe
427
428		#endif // LIBQE_TRAJECTORY_HPP