92def accumulate(
 93    data: pd.DataFrame,
 94    units: str,
 95    conversations: str,
 96    codes: List[str],
 97    window_size: int = 4,
 98    window_forward: int = 0,
 99    binary: bool = True,
100) -> ENAAccumulation:
101    """Accumulate ENA co-occurrence networks without fitting a rotation.
102
103    Runs the stanza-window accumulation step only, producing raw per-unit
104    adjacency vectors.  The result can be passed directly to
105    :meth:`ENA.fit` in place of a raw DataFrame, or used as-is for custom
106    downstream analysis.
107
108    Parameters
109    ----------
110    data : pd.DataFrame
111    units : str
112        Column whose values identify units of analysis.
113    conversations : str
114        Column whose values segment the data into conversations.
115    codes : list[str]
116        Code column names to include in the co-occurrence model.
117    window_size : int
118        Number of lines to look back (default 4).
119    window_forward : int
120        Number of lines to look forward (default 0).
121    binary : bool
122        If True (default), binarise co-occurrence counts.
123
124    Returns
125    -------
126    ENAAccumulation
127
128    Examples
129    --------
130    Basic accumulation::
131
132        accum = accumulate(rs, "unit_key", "convo_key", CODES, window_size=4)
133        print(accum.connection_counts_.shape)   # (48, 15)
134
135    Inspect raw networks before modeling::
136
137        import pandas as pd
138        df = pd.DataFrame(accum.connection_counts_,
139                          index=accum.unit_labels_,
140                          columns=accum.connection_names_)
141
142    Pass to ENA for modeling::
143
144        model = ENA().fit(accum)
145    """
146    n_codes       = len(codes)
147    n_connections = n_codes * (n_codes - 1) // 2
148
149    unit_labels: List[str] = list(dict.fromkeys(data[units].tolist()))
150    n_units = len(unit_labels)
151    unit_index = {label: i for i, label in enumerate(unit_labels)}
152
153    raw_networks = np.zeros((n_units, n_connections), dtype=np.float64)
154
155    for _, conv_df in data.groupby(conversations, sort=False):
156        codes_mat = np.ascontiguousarray(
157            conv_df[codes].to_numpy(dtype=np.float64)
158        )
159        co_occ = _acc.accumulate_stanza(codes_mat, window_size, window_forward, binary)
160        for row_idx, unit_label in enumerate(conv_df[units].tolist()):
161            raw_networks[unit_index[unit_label]] += co_occ[row_idx]
162
163    # Column-major upper-triangle order matching stanza_window output.
164    # Separator is " & " (with spaces) to match R's paste(..., collapse = " & ").
165    connection_names = [
166        f"{codes[i]} & {codes[j]}"
167        for j in range(1, n_codes)
168        for i in range(j)
169    ]
170
171    # Build unit-level metadata: one representative row per unit
172    meta_cols = [c for c in data.columns if c != units and c not in codes]
173    meta = (
174        data[meta_cols + [units]]
175        .drop_duplicates(subset=units)
176        .set_index(units)
177        .reindex(unit_labels)
178    )
179
180    return ENAAccumulation(
181        networks=raw_networks,
182        units=unit_labels,
183        codes=list(codes),
184        connection_names=connection_names,
185        meta=meta,
186        source_call={
187            "data":           data,
188            "units":          units,
189            "conversations":  conversations,
190            "codes":          list(codes),
191            "window_size":    window_size,
192            "window_forward": window_forward,
193            "binary":         binary,
194        },
195    )

38class ENAAccumulation:
39    """Result of :func:`accumulate`: raw adjacency vectors per unit.
40
41    Attributes
42    ----------
43    connection_counts_ : np.ndarray, shape (n_units, n_connections)
44        Raw (un-normalised) co-occurrence counts summed per unit.
45        Corresponds to R's ``enadata$adjacency.vectors`` /
46        ``set$connection.counts``.
47    unit_labels_ : list[str]
48        Unit labels in first-appearance order.
49        Corresponds to R's ``set$model$unit.labels``.
50    codes_ : list[str]
51        Code names in model order.
52    connection_names_ : list[str]
53        Labels for each connection column, e.g. ``"Data & Technical Constraints"``.
54        Order matches the column-major upper-triangle used by ``stanza_window``.
55    meta : pd.DataFrame
56        One row per unit (unit label as index).  Contains all non-code,
57        non-unit-key columns from the original data, deduplicated per unit.
58        Corresponds to R's ``set$meta.data``.
59    source_call : dict | None
60        The keyword arguments passed to :func:`accumulate` that produced this
61        object (``data``, ``units``, ``conversations``, ``codes``,
62        ``window_size``, ``window_forward``, ``binary``).  Corresponds to R's
63        ``ENAAccumulation$`_function.call```; retained so the accumulation can
64        be rebuilt at other window sizes (see :func:`pyena.tune_window_size`).
65        ``None`` when the object was constructed directly rather than via
66        :func:`accumulate`.
67    """
68
69    def __init__(
70        self,
71        networks: np.ndarray,
72        units: List[str],
73        codes: List[str],
74        connection_names: List[str],
75        meta: pd.DataFrame,
76        source_call: Optional[dict] = None,
77    ) -> None:
78        self.connection_counts_ = networks
79        self.unit_labels_       = units
80        self.codes_             = codes
81        self.connection_names_  = connection_names
82        self.meta               = meta
83        self.source_call        = source_call
84
85    def __repr__(self) -> str:
86        return (
87            f"<ENAAccumulation {len(self.unit_labels_)} units × "
88            f"{len(self.connection_names_)} connections>"
89        )

 79class ENA:
 80    """Standard Epistemic Network Analysis pipeline.
 81
 82    Can be used in three equivalent styles:
 83
 84    **One-liner** (accumulate + model in a single call)::
 85
 86        model = ENA().fit(rs, "unit_key", "convo_key", CODES)
 87
 88    **Constructor style** (data up front, options at fit time)::
 89
 90        model = ENA(rs, "unit_key", "convo_key", CODES).fit()
 91        model = ENA(rs, "unit_key", "convo_key", CODES).fit(rotation=mean_rotation(g1, g2))
 92
 93    **Chain style** (mirrors the R pipe)::
 94
 95        model = ENA().accumulate(rs, "unit_key", "convo_key", CODES).fit()
 96        model = (ENA()
 97                   .accumulate(rs, "unit_key", "convo_key", CODES, window_size=8)
 98                   .fit(rotation=generalized_rotation(x_var)))
 99
100    Attributes set after fitting (= R's ena.set fields, flat)
101    ---------------------------------------------------------
102    connection_counts_      raw adjacency vectors (n_units × n_connections)
103    line_weights_           sphere-normed adjacency vectors (= R set$line.weights)
104    points_                 projected unit positions (= R set$points)
105    rotation_matrix_        rotation matrix truncated to dims (= R set$rotation.matrix)
106    meta_data_              unit metadata DataFrame (= R set$meta.data)
107    centroids_              LWS centroid positions (= R model$centroids)
108    variance_               variance explained per dimension (= R model$variance)
109    unit_labels_            unit label strings (= R model$unit.labels)
110    points_for_projection_  centered normed networks (= R model$points.for.projection)
111    rotation_nodes_         code/node positions (= R rotation$nodes)
112    rotation_eigenvalues_   singular values from SVD, or None for other rotations
113    rotation_center_vec_    centering vector (= R rotation$center.vec)
114    rotation_adjacency_key_ [[codeI, codeJ], ...] per connection
115    codes_                  code name strings (= R rotation$codes)
116    connection_names_       connection label strings
117    full_rotation_          full rotation matrix before truncation (Python-specific)
118    weights_                LWS weights (Python-specific)
119    column_classes_         R S3 class annotation per matrix field (Python-specific)
120    accum_                  ENAAccumulation used to build this model
121    """
122
123    def __init__(
124        self,
125        data: Optional[Union[pd.DataFrame, ENAAccumulation]] = None,
126        units: Optional[str] = None,
127        conversations: Optional[str] = None,
128        codes: Optional[List[str]] = None,
129        window_size: int = 4,
130        window_forward: int = 0,
131        binary: bool = True,
132    ) -> None:
133        """Optionally provide data up front; call .fit() to run the model."""
134        if data is not None:
135            self.accumulate(
136                data, units, conversations, codes,
137                window_size=window_size,
138                window_forward=window_forward,
139                binary=binary,
140            )
141
142    def accumulate(
143        self,
144        data: Union[pd.DataFrame, ENAAccumulation],
145        units: Optional[str] = None,
146        conversations: Optional[str] = None,
147        codes: Optional[List[str]] = None,
148        window_size: int = 4,
149        window_forward: int = 0,
150        binary: bool = True,
151    ) -> "ENA":
152        """Run the accumulation step and store the result.
153
154        Returns ``self`` so you can chain directly into ``.fit()``.
155
156        Parameters
157        ----------
158        data : pd.DataFrame | ENAAccumulation
159            Raw data or a pre-built accumulation.
160        units : str
161            Column identifying units of analysis.
162        conversations : str
163            Column segmenting conversations.
164        codes : list[str]
165            Code column names.
166        window_size : int
167            Lines back for stanza window (default 4).
168        window_forward : int
169            Lines forward (default 0).
170        binary : bool
171            Binarise co-occurrences (default True).
172
173        Examples
174        --------
175        Chain into fit::
176
177            model = ENA().accumulate(rs, "unit_key", "convo_key", CODES).fit()
178        """
179        if isinstance(data, ENAAccumulation):
180            self.accum_ = data
181        else:
182            if units is None or conversations is None or codes is None:
183                raise ValueError(
184                    "accumulate() requires units, conversations, and codes "
185                    "when data is a DataFrame."
186                )
187            self.accum_ = _accumulate(
188                data, units, conversations, codes,
189                window_size=window_size,
190                window_forward=window_forward,
191                binary=binary,
192            )
193        return self
194
195    def fit(
196        self,
197        data: Optional[Union[pd.DataFrame, ENAAccumulation]] = None,
198        units: Optional[str] = None,
199        conversations: Optional[str] = None,
200        codes: Optional[List[str]] = None,
201        window_size: int = 4,
202        window_forward: int = 0,
203        binary: bool = True,
204        dims: int = 2,
205        norm: str = "sphere",
206        rotation=None,
207    ) -> "ENA":
208        """Fit an ENA model.
209
210        Parameters
211        ----------
212        data : pd.DataFrame | ENAAccumulation | None
213            * ``None``          — use the accumulation already stored by the
214              constructor or a prior :meth:`accumulate` call.
215            * ``ENAAccumulation`` — use this pre-built accumulation directly.
216            * ``pd.DataFrame``  — accumulate and model in one step (``units``,
217              ``conversations``, and ``codes`` are then required).
218        dims : int
219            Number of dimensions to retain (default 2).
220        norm : str
221            ``"sphere"`` (default) or ``"skip_sphere"``.
222        rotation : None | np.ndarray | callable
223            * ``None``        — default SVD rotation.
224            * ``np.ndarray``  — pre-computed (n_connections × k) matrix.
225            * callable        — factory from :func:`mean_rotation`,
226              :func:`generalized_rotation`, :func:`regression_rotation`, or
227              :func:`regression_rotation_2`.
228        """
229        if norm not in ("sphere", "skip_sphere"):
230            raise ValueError(f"norm must be 'sphere' or 'skip_sphere', got {norm!r}")
231
232        # ── accumulation ─────────────────────────────────────────────────────
233        if data is None:
234            if not hasattr(self, "accum_"):
235                raise ValueError(
236                    "No data provided. Either pass data to fit(), call "
237                    ".accumulate() first, or pass data to ENA()."
238                )
239            accum = self.accum_
240        else:
241            self.accumulate(data, units, conversations, codes,
242                            window_size=window_size,
243                            window_forward=window_forward,
244                            binary=binary)
245            accum = self.accum_
246
247        raw_networks = accum.connection_counts_
248
249        # ── normalization ────────────────────────────────────────────────────
250        if norm == "sphere":
251            normed = normalization.normalize_networks(raw_networks)
252        else:
253            normed = normalization.scale_networks(raw_networks)
254
255        # ── centering ────────────────────────────────────────────────────────
256        # Center only non-zero rows (rENA center.align.to.origin=TRUE default).
257        # center_vec = column means of non-zero rows (= R's rotation$center.vec).
258        non_zero   = normed.sum(axis=1) != 0
259        center_vec = normed[non_zero].mean(axis=0) if non_zero.any() else np.zeros(normed.shape[1])
260
261        centered = np.zeros_like(normed)
262        if non_zero.any():
263            centered[non_zero] = modeling.center_points(
264                np.ascontiguousarray(normed[non_zero])
265            )
266
267        # ── rotation ─────────────────────────────────────────────────────────
268        rotation_eigenvalues = None
269        if rotation is None:
270            # Default: SVD rotation
271            _, s, Vt = np.linalg.svd(centered, full_matrices=False)
272            full_rot = Vt.T   # (n_connections × min(n_units, n_connections))
273            rotation_eigenvalues = s
274        elif isinstance(rotation, np.ndarray):
275            full_rot = rotation
276        elif callable(rotation):
277            full_rot = rotation(centered)
278        else:
279            raise ValueError(
280                f"rotation must be None, np.ndarray, or callable, got {type(rotation)}"
281            )
282
283        rotation_matrix = full_rot[:, :dims]
284        t = centered @ rotation_matrix   # projected unit positions (= R's set$points)
285
286        # ── node positions (LWS) ─────────────────────────────────────────────
287        node_positions = modeling.node_positions(
288            np.ascontiguousarray(normed), np.ascontiguousarray(t), dims
289        )
290
291        # ── derived fields ────────────────────────────────────────────────────
292        variance      = _compute_variance(t)
293        adjacency_key = _build_adjacency_key(list(accum.codes_))
294
295        # ── assign all output fields ──────────────────────────────────────────
296        self.accum_                  = accum
297
298        # top-level (= R's set$...)
299        self.connection_counts_      = raw_networks          # set$connection.counts
300        self.line_weights_           = normed                # set$line.weights
301        self.points_                 = node_positions.points # set$points
302        self.rotation_matrix_        = rotation_matrix       # set$rotation.matrix
303        self.meta_data_              = accum.meta            # set$meta.data
304
305        # model sub-fields (= R's set$model$...)
306        self.centroids_              = node_positions.centroids  # model$centroids (LWS)
307        self.variance_               = variance                  # model$variance
308        self.unit_labels_            = accum.unit_labels_        # model$unit.labels
309        self.points_for_projection_  = centered                  # model$points.for.projection
310
311        # rotation sub-fields (= R's set$rotation$...)
312        self.rotation_nodes_         = node_positions.nodes      # rotation$nodes
313        self.rotation_eigenvalues_   = rotation_eigenvalues      # rotation$eigenvalues
314        self.rotation_center_vec_    = center_vec                # rotation$center.vec
315        self.rotation_adjacency_key_ = adjacency_key             # rotation$adjacency.key
316        self.codes_                  = accum.codes_              # rotation$codes
317        self.connection_names_       = accum.connection_names_
318
319        # Python-specific extras
320        self.full_rotation_          = full_rot               # full pre-truncation rotation
321        self.weights_                = node_positions.weights # LWS weights
322
323        # Column class annotations (mirrors R's S3 class tags per column)
324        self.column_classes_         = {
325            'connection_counts':          'ena.co.occurrence',
326            'line_weights':               'ena.co.occurrence',
327            'points':                     'ena.dimension',
328            'points_for_projection':      'ena.co.occurrence',
329            'rotation_matrix':            'ena.dimension',
330            'rotation_nodes':             'ena.dimension',
331        }
332        return self
333
334    def conf_ints(self, points: Optional[np.ndarray] = None,
335                  conf_level: float = 0.95) -> np.ndarray:
336        """Per-dimension t-based confidence intervals around the column means.
337
338        Parameters
339        ----------
340        points : np.ndarray | None
341            Units × dims matrix. Defaults to ``self.points_``.
342        conf_level : float
343            Confidence level (default 0.95).
344
345        Returns
346        -------
347        np.ndarray, shape (dims, 3)
348            Columns: mean, lower CI, upper CI — one row per dimension.
349        """
350        pts = self.points_ if points is None else points
351        return modeling.mean_ci(np.ascontiguousarray(pts, dtype=np.float64),
352                                conf_level)
353
354    def outlier_ints(self, points: Optional[np.ndarray] = None,
355                     iqr_factor: float = 1.5) -> np.ndarray:
356        """Per-dimension Tukey-fence outlier intervals (Q1-k*IQR, Q3+k*IQR).
357
358        Parameters
359        ----------
360        points : np.ndarray | None
361            Units × dims matrix. Defaults to ``self.points_``.
362        iqr_factor : float
363            IQR multiplier k (default 1.5).
364
365        Returns
366        -------
367        np.ndarray, shape (dims, 2)
368            Columns: lower fence, upper fence — one row per dimension.
369        """
370        pts = self.points_ if points is None else points
371        return modeling.outlier_ci(np.ascontiguousarray(pts, dtype=np.float64),
372                                   iqr_factor)
373
374    def compare_groups(self, g1_mask: np.ndarray,
375                       g2_mask: np.ndarray) -> "modeling.GroupStatsResult":
376        """Per-dimension parametric and non-parametric two-group statistics.
377
378        Parameters
379        ----------
380        g1_mask : np.ndarray of bool
381            Boolean index selecting group 1 rows from ``self.points_``.
382        g2_mask : np.ndarray of bool
383            Boolean index selecting group 2 rows from ``self.points_``.
384
385        Returns
386        -------
387        GroupStatsResult
388            Fields: n1, n2, t, df, pvalue_t, cohens_d, means, sds,
389            U, pvalue_u, effect_r, medians — each length dims.
390        """
391        g1 = np.ascontiguousarray(self.points_[g1_mask], dtype=np.float64)
392        g2 = np.ascontiguousarray(self.points_[g2_mask], dtype=np.float64)
393        return modeling.group_stats(g1, g2)
394
395        return self

171def mean_rotation(group1, group2):
172    """
173    Mean rotation — axis 1 is the mean difference between two groups.
174
175    Mirrors rENA's ``ena.rotate.by.mean``.
176
177    Parameters
178    ----------
179    group1, group2 : array-like of bool (length n_units)
180        Boolean masks identifying the two groups. Applied to the centred
181        normalised networks in the order units appear in the fitted model.
182
183    Returns
184    -------
185    callable : ``rotation_fn(centered) -> (n_connections × n_connections)``
186
187    Example
188    -------
189    rotation = mean_rotation(
190        meta["Condition"] == "FirstGame",
191        meta["Condition"] == "SecondGame",
192    )
193    model = ENA().fit(..., rotation=rotation)
194    """
195    g1 = np.asarray(group1, dtype=bool)
196    g2 = np.asarray(group2, dtype=bool)
197
198    def rotation_fn(centered: np.ndarray) -> np.ndarray:
199        if g1.sum() == 0 or g2.sum() == 0:
200            raise ValueError("mean_rotation: each group must have at least one unit.")
201        diff = centered[g1].mean(axis=0) - centered[g2].mean(axis=0)
202        norm = np.linalg.norm(diff)
203        if norm < 1e-10:
204            raise ValueError(
205                "mean_rotation: group means are identical; cannot define axis 1."
206            )
207        v1 = diff / norm
208        return _build_full_rotation(centered, [v1])
209
210    return rotation_fn

213def generalized_rotation(x_var, y_var=None, select_2_groups=None):
214    """
215    Generalised Means Rotation (GMR).
216
217    Mirrors rENA's ``ena.rotate.by.generalized``.
218
219    Parameters
220    ----------
221    x_var : array-like (n_units,)
222        Predictor for axis 1.  Numeric → OLS regression;
223        string/categorical → first eigenvector of between-group scatter SB.
224    y_var : array-like (n_units,) | None
225        Predictor for axis 2.  If None, axis 2 comes from SVD of deflated data.
226    select_2_groups : tuple(val1, val2) | None
227        When x_var is categorical, use the mean-difference between exactly
228        these two groups (instead of the first eigenvector of SB).
229
230    Returns
231    -------
232    callable : ``rotation_fn(centered) -> (n_connections × n_connections)``
233
234    Examples
235    --------
236    # Continuous predictor
237    rotation = generalized_rotation(meta["CONFIDENCE.Change"].astype(float))
238
239    # Categorical
240    rotation = generalized_rotation(meta["Condition"])
241
242    # Categorical, two explicit groups
243    rotation = generalized_rotation(
244        meta["Condition"],
245        select_2_groups=("FirstGame", "SecondGame"),
246    )
247    """
248    x_arr = np.asarray(x_var)
249    y_arr = np.asarray(y_var) if y_var is not None else None
250
251    def rotation_fn(centered: np.ndarray) -> np.ndarray:
252        if (
253            select_2_groups is not None
254            and not (
255                np.issubdtype(x_arr.dtype, np.floating)
256                or np.issubdtype(x_arr.dtype, np.integer)
257            )
258        ):
259            g1_val, g2_val = select_2_groups
260            diff = (
261                centered[x_arr == g1_val].mean(axis=0)
262                - centered[x_arr == g2_val].mean(axis=0)
263            )
264            norm = np.linalg.norm(diff)
265            v1 = diff / norm if norm > 1e-10 else diff
266        else:
267            v1 = _gmr(centered, x_arr)
268
269        axes = [v1]
270
271        if y_arr is not None:
272            deflated = _deflate(centered, v1)
273            v2 = _gmr(deflated, y_arr)
274            axes.append(v2)
275
276        return _build_full_rotation(centered, axes)
277
278    return rotation_fn

281def regression_rotation(x_var, y_var=None):
282    """
283    Regression rotation — ENA networks as dependent variable.
284
285    Mirrors rENA's ``ena.rotate.by.hena.regression``.
286
287    Parameters
288    ----------
289    x_var : array-like (n_units,)
290        Predictor for axis 1 (numeric or 0/1-encoded categorical).
291    y_var : array-like (n_units,) | None
292        Predictor for axis 2.
293
294    Returns
295    -------
296    callable : ``rotation_fn(centered) -> (n_connections × n_connections)``
297
298    Example
299    -------
300    condition_binary = (meta["Condition"] == "FirstGame").astype(float)
301    rotation = regression_rotation(condition_binary)
302    model = ENA().fit(..., rotation=rotation)
303    """
304    x_arr = np.asarray(x_var, dtype=float)
305    y_arr = np.asarray(y_var, dtype=float) if y_var is not None else None
306
307    def rotation_fn(centered: np.ndarray) -> np.ndarray:
308        v1 = _regression_axis(centered, x_arr)
309        axes = [v1]
310
311        if y_arr is not None:
312            deflated = _deflate(centered, v1)
313            v2 = _regression_axis(deflated, y_arr)
314            axes.append(v2)
315
316        return _build_full_rotation(centered, axes)
317
318    return rotation_fn

321def regression_rotation_2(x_var, y_var=None):
322    """
323    Regression rotation (reversed) — predictor as dependent variable.
324
325    Mirrors rENA's ``ena.rotate.by.hena.regression_2``.
326
327    Parameters
328    ----------
329    x_var : array-like (n_units,)
330        Response variable for axis 1.
331    y_var : array-like (n_units,) | None
332        Response variable for axis 2.
333
334    Returns
335    -------
336    callable : ``rotation_fn(centered) -> (n_connections × n_connections)``
337
338    Example
339    -------
340    condition_binary = (meta["Condition"] == "FirstGame").astype(float)
341    rotation = regression_rotation_2(condition_binary)
342    """
343    x_arr = np.asarray(x_var, dtype=float)
344    y_arr = np.asarray(y_var, dtype=float) if y_var is not None else None
345
346    def rotation_fn(centered: np.ndarray) -> np.ndarray:
347        v1 = _regression_axis_2(centered, x_arr)
348        axes = [v1]
349
350        if y_arr is not None:
351            deflated = _deflate(centered, v1)
352            v2 = _regression_axis_2(deflated, y_arr)
353            axes.append(v2)
354
355        return _build_full_rotation(centered, axes)
356
357    return rotation_fn

31def ena_space_dist_corr(
32    A: np.ndarray,
33    B: np.ndarray,
34    max_sample_size: int = 100_000,
35    random_state: Optional[Union[int, np.random.Generator]] = None,
36) -> float:
37    """Pearson correlation between the pairwise distances of two ENA spaces.
38
39    Computes the Euclidean distance between every pair of points within ``A``
40    and within ``B`` (using the *same* pairing for both), then returns the
41    Pearson correlation of the two distance vectors.  Because pairwise
42    distances are invariant to rotation/reflection of a space, this measures
43    how similar the two point configurations are up to an orthogonal transform
44    — exactly what is needed to compare ENA solutions across window sizes,
45    whose SVD axes may otherwise flip sign.
46
47    Mirrors R's :func:`ena_space_dist_corr`: exact for small spaces, sampled
48    (with replacement, self-pairs dropped) once the number of unique pairs
49    exceeds ``max_sample_size``.
50
51    Parameters
52    ----------
53    A, B : np.ndarray
54        Point matrices (rows are points).  Must share the same number of rows.
55    max_sample_size : int
56        Maximum number of pairwise distances to compute before switching to
57        sampling.  Default 100,000.
58    random_state : int | np.random.Generator | None
59        Seed or generator for the sampled path (ignored on the exact path).
60
61    Returns
62    -------
63    float
64        Pearson correlation of the paired distance vectors.
65    """
66    A = np.ascontiguousarray(A, dtype=np.float64)
67    B = np.ascontiguousarray(B, dtype=np.float64)
68    m = A.shape[0]
69
70    if m == 0 or B.shape[0] != m:
71        raise ValueError("The spaces must have the same non-zero number of rows.")
72
73    total_possible_pairs = m * (m - 1) // 2
74
75    if total_possible_pairs <= max_sample_size:
76        # Exact: all unique i<j pairs.
77        i, j = np.triu_indices(m, k=1)
78        dist_a = np.linalg.norm(A[i] - A[j], axis=1)
79        dist_b = np.linalg.norm(B[i] - B[j], axis=1)
80    else:
81        # Sample pairs with replacement, drop self-pairs (matches R).
82        rng = (random_state if isinstance(random_state, np.random.Generator)
83               else np.random.default_rng(random_state))
84        idx1 = rng.integers(0, m, size=max_sample_size)
85        idx2 = rng.integers(0, m, size=max_sample_size)
86        keep = idx1 != idx2
87        idx1, idx2 = idx1[keep], idx2[keep]
88        dist_a = np.linalg.norm(A[idx1] - A[idx2], axis=1)
89        dist_b = np.linalg.norm(B[idx1] - B[idx2], axis=1)
90
91    return float(np.corrcoef(dist_a, dist_b)[0, 1])

 96def tune_window_size(
 97    accum: ENAAccumulation,
 98    min_size: int = 1,
 99    max_size: int = 20,
100    cutoff: float = 0.95,
101) -> ENAAccumulation:
102    """Find the stability-plateau window size and rebuild the accumulation there.
103
104    Iterates the stanza window from ``min_size`` to ``max_size``, rebuilding the
105    accumulation and fitting a default (SVD) ENA model at each size.  Adjacent
106    window sizes are compared with :func:`ena_space_dist_corr` on their unit
107    points; the smallest window whose adjacent correlation reaches
108    ``cutoff * max(correlation)`` is selected.
109
110    Parameters
111    ----------
112    accum : ENAAccumulation
113        An accumulation produced by :func:`pyena.accumulate`.  Its
114        ``source_call`` is used to rebuild at each window size — pass an object
115        built via :func:`accumulate` (not one constructed directly).
116    min_size : int
117        Smallest window size to test (default 1).
118    max_size : int
119        Largest window size to test (default 20).
120    cutoff : float
121        Fraction of the maximum adjacent correlation used as the selection
122        threshold (default 0.95).
123
124    Returns
125    -------
126    ENAAccumulation
127        A new accumulation rebuilt at the selected window size (mirrors R,
128        which returns the rebuilt object).  The chosen size is available as
129        ``result.source_call["window_size"]``.
130    """
131    call = getattr(accum, "source_call", None)
132    if call is None:
133        raise ValueError(
134            "accum has no stored source_call; build it with pyena.accumulate() "
135            "to enable window-size tuning."
136        )
137
138    window_range = list(range(min_size, max_size + 1))
139    if len(window_range) < 2:
140        raise ValueError("max_size must be greater than min_size to compare windows.")
141
142    # Imported here to avoid a circular import (ena.py imports accumulation).
143    from .ena import ENA
144
145    # 1. Rebuild + fit at each window size, collecting the unit points.
146    all_points = []
147    for window_size in window_range:
148        new_accum = accumulate(
149            call["data"], call["units"], call["conversations"], call["codes"],
150            window_size=window_size,
151            window_forward=call["window_forward"],
152            binary=call["binary"],
153        )
154        model = ENA().fit(new_accum)
155        all_points.append(model.points_)
156
157    # 2. Adjacent-window distance-space correlations.
158    adj_correlations = np.array([
159        ena_space_dist_corr(all_points[i], all_points[i + 1])
160        for i in range(len(window_range) - 1)
161    ])
162
163    # 3. Smallest window crossing cutoff * max correlation.
164    max_corr = np.nanmax(adj_correlations)
165    threshold = cutoff * max_corr
166    crossers = np.where(adj_correlations >= threshold)[0]
167    best_idx = int(crossers[0]) if crossers.size else 0
168    best_window_size = window_range[best_idx]
169
170    # 4. Rebuild the accumulation at the selected window size.
171    return accumulate(
172        call["data"], call["units"], call["conversations"], call["codes"],
173        window_size=best_window_size,
174        window_forward=call["window_forward"],
175        binary=call["binary"],
176    )