51class MolecularNetwork(NetworkVisualizeMixin):
  52    """Build a molecular network around a reference library.
  53
  54    Initialization is **lazy** — no similarity computation happens until you
  55    call one of the computation methods.
  56
  57    Parameters
  58    ----------
  59    fe_lib : ms_entropy.FlashEntropySearch or None, optional
  60        Pre-built FlashEntropy search instance (built from the reference
  61        library).  Required for query-vs-library and library-vs-library
  62        stages.  Pass ``None`` (default) to build a query-only network
  63        (stage 1 only).  When provided, fragment-tolerance parameters are
  64        automatically extracted from this library unless overridden by
  65        *ms2_tolerance_da*.
  66    search_type : str
  67        FlashEntropy search mode: ``"identity"``, ``"open"``, or
  68        ``"neutral_loss"``.  Default ``"open"`` (typical DDA networking;
  69        use ``"identity"`` when precursor filtering is required).
  70    additional_similarities : list of str, optional
  71        Extra similarity metrics to compute alongside entropy similarity.
  72        Currently supported: ``["cosine"]``.  Default ``["cosine"]``.
  73        Cosine is gated by *entropy_threshold_low* (best-effort secondary score).
  74    similarity_thresholds : dict, optional
  75        Mapping of metric name → edge threshold.  Edges are created in each
  76        metric's network when score >= threshold.  Metrics not listed use
  77        a default of 0.5.
  78        Example: ``{"entropy_similarity": 0.5, "cosine": 0.6}``
  79    ms1_tolerance_da : float, optional
  80        Precursor m/z tolerance (Da) used for ``"identity"`` search to filter
  81        library candidates by precursor m/z.  Ignored for ``"open"`` and
  82        ``"neutral_loss"`` (precursor filter is disabled for those modes).
  83        If None (default), uses 0.01 Da.
  84    ms2_tolerance_da : float, optional
  85        Fragment m/z tolerance (Da) used for spectrum cleaning and similarity
  86        scoring.  Priority: explicit kwarg > extracted from *fe_lib* >
  87        0.01 Da fallback.  Override this when *fe_lib* is ``None`` or when
  88        you need a tolerance different from the library's built-in value.
  89    entropy_threshold_low : float or None, optional
  90        Minimum entropy similarity score required to trigger additional metric
  91        computation (e.g. cosine).  If None (default), set to half the lowest
  92        non-entropy similarity threshold, or 0.25 when no thresholds are given.
  93    settings : SpectralSimilaritySearchSettings, optional
  94        Encapsulated spectral-similarity / networking defaults (for example
  95        ``lcms.parameters.mass_spectrum["ms2"].spectral_similarity_search`` or
  96        :func:`~corems.encapsulation.factory.parameters.settings_from_lcms`).
  97        When provided, fills any constructor fields left as ``None`` from this
  98        object.  Explicit keyword arguments always override *settings*.  Stage-1
  99        temporary FE index build uses ``settings.as_fe_kwargs()`` when
 100        *fe_kwargs* is not passed to :meth:`query_vs_library` /
 101        :meth:`run_query_vs_query_only`.  Stage-3 defaults
 102        (``hydrate_library_similarities``, ``library_similarity_threshold``)
 103        are also taken from *settings* when those method arguments are omitted.
 104
 105    Attributes
 106    ----------
 107    similarity_matrices : dict of str → SimilarityMatrix
 108        One :class:`~corems.molecular_networking.similarity_matrix.SimilarityMatrix`
 109        per metric (``"entropy_similarity"``, ``"cosine"``, …).
 110    stage_query_query_done : bool
 111        True after query-vs-query similarities have been computed (stage 1).
 112    stage_query_library_done : bool
 113        True after query-vs-library similarities have been computed (stage 2).
 114    stage_library_library_done : bool
 115        True after matched library-vs-library similarities have been computed
 116        (stage 3).
 117
 118    Notes
 119    -----
 120    **Tolerance resolution order** (highest priority first):
 121
 122    1. Explicit *ms2_tolerance_da* kwarg passed to this constructor.
 123    2. *settings.ms2_tolerance_da* when *settings* is provided and the kwarg is
 124       omitted.
 125    3. Value extracted from ``fe_lib.entropy_search.max_ms2_tolerance_in_da``.
 126    4. Hard-coded fallback of 0.01 Da (used when *fe_lib* is ``None`` and no
 127       explicit value is given).
 128
 129    Prefer building *fe_lib* with the same
 130    :class:`~corems.encapsulation.factory.processingSetting.SpectralSimilaritySearchSettings`
 131    instance used here (via library ``settings=`` / ``settings_from_lcms``) so
 132    FlashEntropy index tolerances match networking.
 133
 134    **Precursor filter behaviour by search type:**
 135
 136    - ``"identity"`` : library candidates are filtered by precursor m/z within
 137      *ms1_tolerance_da* before scoring.
 138    - ``"open"`` : precursor filter is disabled; all library entries are scored.
 139    - ``"neutral_loss"`` : precursor filter is disabled; matching occurs in
 140      neutral-loss mass space (``precursor_mz − fragment_mz``), not on the
 141      precursor itself.
 142
 143    Cosine scoring (when requested) is single-threaded: unique spectra are
 144    pre-sorted once, then pairs are scored sequentially.
 145    """
 146
 147    def __init__(
 148        self,
 149        fe_lib=None,
 150        search_type: str | None = None,
 151        additional_similarities: list[str] | None = None,
 152        similarity_thresholds: dict[str, float] | None = None,
 153        ms1_tolerance_da: float | None = None,
 154        ms2_tolerance_da: float | None = None,
 155        entropy_threshold_low: float | None = None,
 156        settings=None,
 157    ):
 158        if settings is not None:
 159            if search_type is None:
 160                search_type = settings.search_type
 161            if additional_similarities is None:
 162                additional_similarities = list(settings.additional_similarities)
 163            if similarity_thresholds is None:
 164                similarity_thresholds = dict(settings.similarity_thresholds)
 165            if ms1_tolerance_da is None:
 166                ms1_tolerance_da = settings.ms1_tolerance_da
 167            if ms2_tolerance_da is None:
 168                ms2_tolerance_da = settings.ms2_tolerance_da
 169            if entropy_threshold_low is None:
 170                entropy_threshold_low = settings.entropy_threshold_low
 171            self._settings = settings
 172        else:
 173            self._settings = None
 174
 175        if search_type is None:
 176            search_type = "open"
 177        if additional_similarities is None:
 178            additional_similarities = ["cosine"]
 179        if similarity_thresholds is None:
 180            similarity_thresholds = {}
 181
 182        if entropy_threshold_low is None:
 183            non_entropy_thresholds = [
 184                v for k, v in similarity_thresholds.items() if k != "entropy_similarity"
 185            ]
 186            if non_entropy_thresholds:
 187                entropy_threshold_low = min(non_entropy_thresholds) / 2
 188            else:
 189                entropy_threshold_low = 0.25  # Placeholder that won't be used if no additional similarities are computed
 190
 191        self.fe_lib = fe_lib
 192        self.search_type = search_type
 193        self.additional_similarities = list(additional_similarities)
 194        self.similarity_thresholds = similarity_thresholds
 195
 196        # Build the engine (wraps fe_lib + search parameters)
 197        # When fe_lib is provided, tolerance parameters are extracted from it.
 198        # When fe_lib is None (queries-only use), SimilarityEngine will fall back
 199        # to its internal defaults and only query-vs-query stage is allowed.
 200        self._engine = SimilarityEngine(
 201            fe_lib=fe_lib,
 202            search_type=search_type,
 203            additional_similarities=additional_similarities,
 204            ms1_tolerance_da=ms1_tolerance_da,
 205            ms2_tolerance_da=ms2_tolerance_da,
 206            entropy_threshold_low=entropy_threshold_low,
 207            settings=settings,
 208        )
 209
 210        # One SimilarityMatrix per metric – empty until query_vs_library() is called
 211        all_metrics = ["entropy_similarity"] + list(additional_similarities)
 212        self.similarity_matrices: dict[str, SimilarityMatrix] = {
 213            m: SimilarityMatrix(metric_name=m) for m in all_metrics
 214        }
 215
 216        # Accumulate query spectra (kept separate from the library FE index)
 217        self._all_query_spectra: list = []
 218        self._all_query_ids: list[str] = []
 219        self._all_query_precursor_mzs: list[float | None] = []
 220
 221        # Stage flags – enforce order:
 222        #   1) query-vs-query
 223        #   2) query-vs-library
 224        #   3) matched library-vs-library
 225        self.stage_query_query_done: bool = False
 226        self.stage_query_library_done: bool = False
 227        self.stage_library_library_done: bool = False
 228
 229        # Whether query_vs_library has been run (disabled repeated runs until dropped)
 230        self._has_queries_run: bool = False
 231
 232        # Cache for stage 2 entropy pairs (used by stage 3)
 233        self._stage2_entropy_pairs: dict[tuple[str, str], float] | None = None
 234
 235        # Cache of computed network clustering artifacts keyed by metric.
 236        self._network_clusters: dict[str, dict] = {}
 237
 238    @staticmethod
 239    def prepare_query_spectra_from_lcms_object(
 240        lcms_obj,
 241        mf_ids: set | None = None,
 242    ) -> tuple[list, list[str], list[float]]:
 243        """Build query lists from an ``LCMSBase`` object's mass features.
 244
 245        Iterates ``lcms_obj.mass_features`` and collects each feature's
 246        ``best_ms2`` spectrum, skipping features with no MS2 or an empty
 247        ``mz_exp``.  If *mf_ids* is provided, only those mass features are
 248        included.
 249
 250        Parameters
 251        ----------
 252        lcms_obj : LCMSBase
 253            LCMS object whose mass features have been populated with MS2
 254            spectra (e.g. via :meth:`add_associated_ms2_dda`).
 255        mf_ids : set, optional
 256            If provided, only mass features whose id is in this set are
 257            included.
 258
 259        Returns
 260        -------
 261        (query_spectra, query_ids, query_precursor_mzs)
 262            Three parallel lists ready to pass to :meth:`query_vs_library`
 263            or :meth:`run_query_vs_query_only`.
 264        """
 265        query_spectra: list = []
 266        query_ids: list[str] = []
 267        query_precursor_mzs: list[float] = []
 268        for mf_id, mf in lcms_obj.mass_features.items():
 269            if mf_ids is not None and mf_id not in mf_ids:
 270                continue
 271            ms2 = mf.best_ms2
 272            if ms2 is None:
 273                continue
 274            if not hasattr(ms2, "mz_exp") or len(ms2.mz_exp) == 0:
 275                continue
 276            query_spectra.append(ms2)
 277            query_ids.append(str(mf_id))
 278            query_precursor_mzs.append(float(mf.mz))
 279        return query_spectra, query_ids, query_precursor_mzs
 280
 281    @staticmethod
 282    def prepare_query_spectra_from_lcms_collection(
 283        lcms_collection,
 284        cluster_ids: set | None = None,
 285        representative_metric: str | None = None,
 286    ) -> tuple[list, list[str], list[float]]:
 287        """Build query lists from representative mass features in an ``LCMSCollection``.
 288
 289        Uses :meth:`~corems.mass_spectra.calc.lc_calc.LCMSCollectionCalculations.get_representative_mass_features_for_all_clusters`
 290        to select one mass feature per consensus cluster, then collects each
 291        feature's ``best_ms2`` from the corresponding sample object.
 292
 293        Representative mass features must already be loaded with MS2 spectra
 294        (e.g. via ``process_consensus_features(load_representatives=True,
 295        add_ms2=True, ...)`` or ``reload_representative_mass_features(add_ms2=True)``).
 296        Clusters whose representative has no usable MS2 are skipped.
 297
 298        Parameters
 299        ----------
 300        lcms_collection : LCMSCollection
 301            Collection with consensus clusters and loaded representative
 302            mass features (with MS2 where available).
 303        cluster_ids : set, optional
 304            If provided, only representatives for these cluster IDs are
 305            included.
 306        representative_metric : str, optional
 307            Metric passed to
 308            :meth:`~corems.mass_spectra.calc.lc_calc.LCMSCollectionCalculations.get_representative_mass_features_for_all_clusters`.
 309            If ``None``, uses the collection parameter
 310            ``consensus_representative_metric``.
 311
 312        Returns
 313        -------
 314        (query_spectra, query_ids, query_precursor_mzs)
 315            Three parallel lists ready to pass to :meth:`query_vs_library`
 316            or :meth:`run_query_vs_query_only`.  Query IDs are collection-level
 317            ``coll_mf_id`` strings (``"{sample_id}_{mf_id}"``).
 318
 319        Raises
 320        ------
 321        AttributeError
 322            If *lcms_collection* does not expose representative lookup or
 323            sample indexing (not an ``LCMSCollection``-like object).
 324        """
 325        if not hasattr(
 326            lcms_collection, "get_representative_mass_features_for_all_clusters"
 327        ):
 328            raise AttributeError(
 329                "lcms_collection must provide get_representative_mass_features_for_all_clusters "
 330                "(expected an LCMSCollection)."
 331            )
 332
 333        representatives = (
 334            lcms_collection.get_representative_mass_features_for_all_clusters(
 335                representative_metric=representative_metric
 336            )
 337        )
 338        if cluster_ids is not None:
 339            representatives = representatives[
 340                representatives["cluster"].isin(cluster_ids)
 341            ]
 342
 343        query_spectra: list = []
 344        query_ids: list[str] = []
 345        query_precursor_mzs: list[float] = []
 346
 347        for _, row in representatives.iterrows():
 348            sample_id = int(row["sample_id"])
 349            mf_id = row["mf_id"]
 350            coll_mf_id = str(row["coll_mf_id"])
 351
 352            try:
 353                sample = lcms_collection[sample_id]
 354            except (IndexError, KeyError, TypeError):
 355                continue
 356
 357            mass_features = getattr(sample, "mass_features", None) or {}
 358            mf = mass_features.get(mf_id)
 359            if mf is None:
 360                # mf_id from dataframes may not match dict key type (int/str/np)
 361                try:
 362                    mf = mass_features.get(int(mf_id))
 363                except (TypeError, ValueError):
 364                    mf = None
 365            if mf is None:
 366                try:
 367                    mf = mass_features.get(str(mf_id))
 368                except (TypeError, ValueError):
 369                    mf = None
 370            if mf is None:
 371                continue
 372
 373            ms2 = getattr(mf, "best_ms2", None)
 374            if ms2 is None:
 375                continue
 376            if not hasattr(ms2, "mz_exp") or len(ms2.mz_exp) == 0:
 377                continue
 378
 379            query_spectra.append(ms2)
 380            query_ids.append(coll_mf_id)
 381            query_precursor_mzs.append(float(mf.mz))
 382
 383        return query_spectra, query_ids, query_precursor_mzs
 384
 385    # ── Internal helpers ──────────────────────────────────────────────────────
 386
 387    def _threshold_for(self, metric: str) -> float:
 388        """Return the edge threshold for *metric*."""
 389        return self.similarity_thresholds.get(metric, 0.5)
 390
 391    @staticmethod
 392    def library_node_id(lib_idx: int) -> str:
 393        """Return the internal network node ID for a library spectrum index."""
 394        return f"{_LIB_NODE_PREFIX}{int(lib_idx)}"
 395
 396    @staticmethod
 397    def library_index_from_node_id(node_id: str) -> int | None:
 398        """Parse a library matrix index from an internal node ID, if any.
 399
 400        Accepts the canonical ``lib:<idx>`` form. Bare integer strings are
 401        also accepted for backward compatibility with older matrices.
 402        """
 403        if node_id is None:
 404            return None
 405        s = str(node_id)
 406        if s.startswith(_LIB_NODE_PREFIX):
 407            try:
 408                return int(s[len(_LIB_NODE_PREFIX) :])
 409            except ValueError:
 410                return None
 411        try:
 412            return int(s)
 413        except (TypeError, ValueError):
 414            return None
 415
 416    def _update_matrices(self, new_scores: dict[str, dict[tuple[str, str], float]]):
 417        """Write new similarity scores into the appropriate SimilarityMatrix."""
 418        for metric, pairs in new_scores.items():
 419            mat = self.similarity_matrices.get(metric)
 420            if mat is None:
 421                continue
 422            for (id1, id2), score in pairs.items():
 423                mat.set_similarity(id1, id2, score)
 424        # Finalise all matrices for efficient querying
 425        for mat in self.similarity_matrices.values():
 426            mat.finalise()
 427
 428    def _prepare_queries(
 429        self,
 430        query_spectra: list,
 431        query_ids: list[str],
 432        query_precursor_mzs: list[float | None] | None,
 433    ) -> None:
 434        """Validate and store query spectra, IDs, and precursor m/z values.
 435
 436        Resets all stage flags and the stage-2 entropy-pair cache so that a
 437        fresh computation can proceed.  Called internally by
 438        :meth:`run_query_vs_query_only` and :meth:`query_vs_library`.
 439
 440        Parameters
 441        ----------
 442        query_spectra : list
 443            Spectrum objects with ``.mz_exp`` and ``.abundance`` attributes.
 444        query_ids : list of str
 445            Unique string ID for each spectrum.  Must be the same length as
 446            *query_spectra*.
 447        query_precursor_mzs : list of float or None, or None
 448            Precursor m/z for each spectrum.  Required (non-None list) for
 449            ``"identity"`` and ``"neutral_loss"`` search types.  Pass ``None``
 450            to auto-fill with ``None`` values for ``"open"`` search.
 451
 452        Raises
 453        ------
 454        ValueError
 455            If *query_spectra* is empty, lengths are mismatched, or precursor
 456            m/z values are missing/None for a precursor-required search type.
 457        """
 458        n_query = len(query_spectra)
 459        if n_query == 0:
 460            raise ValueError("No query spectra provided.")
 461
 462        if len(query_ids) != n_query:
 463            raise ValueError(
 464                f"query_ids length ({len(query_ids)}) must match query_spectra length ({n_query})."
 465            )
 466
 467        # Validate precursor_mzs requirement
 468        if self.search_type in _PRECURSOR_REQUIRED:
 469            if query_precursor_mzs is None:
 470                raise ValueError(
 471                    f"query_precursor_mzs is required for "
 472                    f"search_type='{self.search_type}'. "
 473                    "Provide a list of precursor m/z values, one per query spectrum."
 474                )
 475            if len(query_precursor_mzs) != n_query:
 476                raise ValueError(
 477                    f"query_precursor_mzs length ({len(query_precursor_mzs)}) "
 478                    f"must match query_spectra length ({n_query})."
 479                )
 480            if any(pmz is None for pmz in query_precursor_mzs):
 481                raise ValueError(
 482                    f"query_precursor_mzs contains None values for search_type='{self.search_type}'. "
 483                    "Provide a valid precursor m/z for every query spectrum."
 484                )
 485        else:
 486            if query_precursor_mzs is None:
 487                query_precursor_mzs = [None] * n_query
 488            elif len(query_precursor_mzs) != n_query:
 489                raise ValueError(
 490                    f"query_precursor_mzs length ({len(query_precursor_mzs)}) "
 491                    f"must match query_spectra length ({n_query})."
 492                )
 493
 494        # Store provided queries
 495        self._all_query_spectra = list(query_spectra)
 496        self._all_query_ids = list(query_ids)
 497        self._all_query_precursor_mzs = list(query_precursor_mzs)
 498
 499        # Register query IDs now
 500        for mat in self.similarity_matrices.values():
 501            mat.register_spectra(self._all_query_ids)
 502
 503        # Reset stage flags and caches
 504        self.stage_query_query_done = False
 505        self.stage_query_library_done = False
 506        self.stage_library_library_done = False
 507        self._stage2_entropy_pairs = None
 508
 509    def _run_query_vs_query(
 510        self,
 511        fe_kwargs: dict | None = None,
 512    ) -> None:
 513        """Stage 1 – compute query-vs-query similarities.
 514
 515        Builds a temporary FlashEntropy index from the stored query spectra
 516        and computes all-vs-all pairwise similarities within that set.
 517        Results are written into :attr:`similarity_matrices` and
 518        :attr:`stage_query_query_done` is set to ``True``.
 519
 520        Parameters
 521        ----------
 522        fe_kwargs : dict, optional
 523            Extra keyword arguments forwarded to
 524            :class:`ms_entropy.FlashEntropySearch` when building the
 525            temporary query index.
 526
 527        Raises
 528        ------
 529        RuntimeError
 530            If no queries have been prepared (i.e. :meth:`_prepare_queries`
 531            has not been called).
 532        """
 533        if not self._all_query_spectra:
 534            raise RuntimeError(
 535                "No queries prepared. Call run_query_vs_query_only() or query_vs_library() first."
 536            )
 537
 538        if fe_kwargs is None and self._settings is not None:
 539            fe_kwargs = self._settings.as_fe_kwargs()
 540
 541        query_fe_lib = self._engine.build_fe_index_from_spectra(
 542            spectra=self._all_query_spectra,
 543            precursor_mzs=self._all_query_precursor_mzs,
 544            fe_kwargs=fe_kwargs,
 545        )
 546        qq_scores = self._engine.compute_all_vs_all_with_lib(
 547            spectra=self._all_query_spectra,
 548            spectrum_ids=self._all_query_ids,
 549            precursor_mzs=self._all_query_precursor_mzs,
 550            fe_lib_override=query_fe_lib,
 551        )
 552        self._update_matrices(qq_scores)
 553        self.stage_query_query_done = True
 554
 555    def _run_query_vs_library(
 556        self,
 557    ) -> tuple[dict[tuple[str, str], float], dict[tuple[str, str], float]]:
 558        """Stage 2 – compute query-vs-library similarities.
 559
 560        Searches each stored query spectrum against the reference FE library
 561        (``self.fe_lib``).  Entropy similarity scores are stored for all
 562        library entries with score > 0; cosine scores are computed for pairs
 563        above :attr:`~SimilarityEngine.entropy_threshold_low`.  Results are
 564        written into :attr:`similarity_matrices`,
 565        :attr:`stage_query_library_done` is set to ``True``, and the entropy
 566        pairs are cached in :attr:`_stage2_entropy_pairs` for use by stage 3.
 567
 568        Returns
 569        -------
 570        entropy_pairs : dict mapping (query_id, lib_idx_str) → float
 571            All query-vs-library entropy similarity scores > 0.
 572        cosine_pairs : dict mapping (query_id, lib_idx_str) → float
 573            Cosine scores for pairs above the entropy threshold.
 574
 575        Raises
 576        ------
 577        RuntimeError
 578            If stage 1 has not been completed or if ``self.fe_lib`` is ``None``.
 579        """
 580        if not self.stage_query_query_done:
 581            raise RuntimeError(
 582                "query-vs-query stage must be run before query-vs-library stage."
 583            )
 584        if self.fe_lib is None:
 585            raise RuntimeError(
 586                "query-vs-library stage requires a reference FE library (fe_lib). "
 587                "Construct MolecularNetwork with fe_lib or use run_query_vs_query_only for query-only networks."
 588            )
 589
 590        combined, lib_size = self._engine.search_queries_against_library(
 591            query_spectra=self._all_query_spectra,
 592            query_ids=self._all_query_ids,
 593            query_precursor_mzs=self._all_query_precursor_mzs,
 594            format_library_id=self.library_node_id,
 595        )
 596
 597        lib_ids = [self.library_node_id(i) for i in range(lib_size)]
 598        for mat in self.similarity_matrices.values():
 599            mat.register_spectra(lib_ids)
 600
 601        self._update_matrices(combined)
 602
 603        entropy_pairs = combined.get("entropy_similarity", {})
 604        cosine_pairs = combined.get("cosine", {})
 605
 606        self.stage_query_library_done = True
 607        self._stage2_entropy_pairs = entropy_pairs
 608
 609        return entropy_pairs, cosine_pairs
 610
 611    def _run_library_vs_library(
 612        self,
 613        entropy_pairs: dict[tuple[str, str], float],
 614        library_similarity_threshold: float,
 615    ) -> None:
 616        """Stage 3 – compute library-vs-library similarities for matched entries.
 617
 618        Identifies the subset of library spectra that scored ≥
 619        *library_similarity_threshold* against any query in stage 2, then
 620        computes all-vs-all pairwise similarities within that subset.
 621        Results are written into :attr:`similarity_matrices` and
 622        :attr:`stage_library_library_done` is set to ``True``.
 623
 624        Parameters
 625        ----------
 626        entropy_pairs : dict mapping (query_id, lib_idx_str) → float
 627            The entropy similarity pairs returned by :meth:`_run_query_vs_library`.
 628            Used to identify which library indices exceeded the threshold.
 629        library_similarity_threshold : float
 630            Minimum entropy similarity score a library entry must have against
 631            any query to be included in this stage.
 632
 633        Raises
 634        ------
 635        RuntimeError
 636            If stage 2 has not been completed.
 637        """
 638        if not self.stage_query_library_done:
 639            raise RuntimeError(
 640                "query-vs-library stage must be run before library-vs-library stage."
 641            )
 642
 643        matched_lib_indices: list[int] = sorted(
 644            {
 645                lib_idx
 646                for (q_id, lib_id), score in entropy_pairs.items()
 647                if q_id in set(self._all_query_ids)
 648                and score >= library_similarity_threshold
 649                and (lib_idx := self.library_index_from_node_id(lib_id)) is not None
 650            }
 651        )
 652
 653        n_matched = len(matched_lib_indices)
 654
 655        if n_matched > 1:
 656            matched_lib_ids = [self.library_node_id(i) for i in matched_lib_indices]
 657            ll_scores = self._engine.compute_library_vs_library_filtered(
 658                library_indices=matched_lib_indices,
 659                spectrum_ids=matched_lib_ids,
 660            )
 661            self._update_matrices(ll_scores)
 662
 663        self.stage_library_library_done = True
 664
 665    def _export_node_id(self, node_id: str, query_id_set: set[str]) -> str:
 666        """Return node ID for file export.
 667
 668        Query IDs are preserved as-is. Library nodes stored internally as
 669        ``lib:<idx>`` (or bare index strings) are mapped to FlashEntropy entry
 670        ``spectra_id`` when available (fallback to ``id``, then the internal
 671        node ID).
 672        """
 673        if node_id in query_id_set:
 674            return node_id
 675
 676        lib_idx = self.library_index_from_node_id(node_id)
 677        if lib_idx is None or lib_idx < 0:
 678            return node_id
 679
 680        if self.fe_lib is None:
 681            return node_id
 682
 683        try:
 684            lib_entry = self.fe_lib[lib_idx]
 685        except (IndexError, KeyError, TypeError):
 686            return node_id
 687
 688        if not isinstance(lib_entry, dict):
 689            return node_id
 690
 691        lib_id = lib_entry.get("spectra_id")
 692        if lib_id is None:
 693            lib_id = lib_entry.get("id")
 694        if lib_id is None:
 695            return node_id
 696
 697        lib_id_str = str(lib_id)
 698        return lib_id_str if lib_id_str else node_id
 699
 700    # ── Public API ────────────────────────────────────────────────────────────
 701
 702    def run_query_vs_query_only(
 703        self,
 704        query_spectra: list,
 705        query_ids: list[str],
 706        query_precursor_mzs: list[float | None] | None = None,
 707        fe_kwargs: dict | None = None,
 708    ) -> None:
 709        """Compute query-vs-query similarities (stage 1 only).
 710
 711        Builds a temporary FlashEntropy index from *query_spectra* and
 712        computes all-vs-all pairwise similarities within that set.  Results
 713        are stored in :attr:`similarity_matrices`.
 714
 715        After this call, :attr:`stage_query_query_done` is ``True``.  If a
 716        reference library was provided at construction time you may continue
 717        with :meth:`run_query_vs_library_stage`.
 718
 719        Parameters
 720        ----------
 721        query_spectra : list
 722            Spectrum objects with ``.mz_exp`` and ``.abundance`` attributes.
 723        query_ids : list of str
 724            Unique string ID for each spectrum.
 725        query_precursor_mzs : list of float or None, optional
 726            Precursor m/z for each spectrum.  Required (non-None list) for
 727            ``"identity"`` and ``"neutral_loss"`` search types.
 728        fe_kwargs : dict, optional
 729            Extra keyword arguments forwarded to
 730            :class:`ms_entropy.FlashEntropySearch` when building the
 731            temporary query index.
 732
 733        Raises
 734        ------
 735        RuntimeError
 736            If queries have already been run (call :meth:`drop_queries` first).
 737        ValueError
 738            If *query_spectra* is empty, lengths are mismatched, or precursor
 739            m/z values are missing for a precursor-required search type.
 740        """
 741        if self._has_queries_run:
 742            raise RuntimeError(
 743                "query_vs_library has already been run — call drop_queries() to clear queries and results before running again."
 744            )
 745
 746        self._prepare_queries(
 747            query_spectra=query_spectra,
 748            query_ids=query_ids,
 749            query_precursor_mzs=query_precursor_mzs,
 750        )
 751        self._run_query_vs_query(fe_kwargs=fe_kwargs)
 752        self._has_queries_run = True
 753
 754    def run_query_vs_library_stage(self) -> None:
 755        """Continue from stage 1 and compute query-vs-library similarities (stage 2).
 756
 757        Searches each stored query spectrum against the reference FE library
 758        and stores entropy + cosine scores in :attr:`similarity_matrices`.
 759        After this call, :attr:`stage_query_library_done` is ``True``.
 760
 761        Stage-3 library subset selection uses
 762        *library_similarity_threshold* on :meth:`run_library_vs_library_stage`
 763        (or the same keyword on :meth:`query_vs_library`), not this method.
 764
 765        Raises
 766        ------
 767        RuntimeError
 768            If stage 1 has not been completed (call
 769            :meth:`run_query_vs_query_only` first) or if ``fe_lib`` is
 770            ``None``.
 771        """
 772        if not self._has_queries_run:
 773            raise RuntimeError(
 774                "No queries have been run yet — call run_query_vs_query_only() or query_vs_library() first."
 775            )
 776
 777        entropy_pairs, _ = self._run_query_vs_library()
 778        self._stage2_entropy_pairs = entropy_pairs
 779
 780    def run_library_vs_library_stage(
 781        self,
 782        library_similarity_threshold: float = 0.3,
 783    ) -> None:
 784        """Continue from stage 2 and compute library-vs-library similarities (stage 3).
 785
 786        Identifies library spectra that scored ≥ *library_similarity_threshold*
 787        against any query in stage 2, then computes all-vs-all pairwise
 788        similarities within that subset.  Results are stored in
 789        :attr:`similarity_matrices` and :attr:`stage_library_library_done`
 790        is set to ``True``.
 791
 792        Parameters
 793        ----------
 794        library_similarity_threshold : float, optional
 795            Minimum entropy similarity score (from stage 2) required for a
 796            library spectrum to be included in this stage.  Default 0.3.
 797
 798        Raises
 799        ------
 800        RuntimeError
 801            If stage 2 has not been completed (call
 802            :meth:`run_query_vs_library_stage` first).
 803        """
 804        if self._stage2_entropy_pairs is None:
 805            raise RuntimeError(
 806                "query-vs-library stage must be run before library-vs-library stage."
 807            )
 808
 809        self._run_library_vs_library(
 810            entropy_pairs=self._stage2_entropy_pairs,
 811            library_similarity_threshold=library_similarity_threshold,
 812        )
 813
 814    def query_vs_library(
 815        self,
 816        query_spectra: list,
 817        query_ids: list[str],
 818        query_precursor_mzs: list[float | None] | None = None,
 819        fe_kwargs: dict | None = None,
 820        *,
 821        hydrate_library_similarities: bool | None = None,
 822        library_similarity_threshold: float | None = None,
 823    ) -> None:
 824        """Run all three stages in one call (convenience wrapper).
 825
 826        Executes stages 1–3 sequentially:
 827
 828        1. **Query-vs-query** – builds a temporary FlashEntropy index from
 829           *query_spectra* and computes all-vs-all pairwise similarities.
 830        2. **Query-vs-library** – searches each query against the reference
 831           FE library (``self.fe_lib``).
 832        3. **Library-vs-library** *(optional)* – when
 833           *hydrate_library_similarities* is ``True``, computes pairwise
 834           similarities for the subset of library spectra that matched any
 835           query above *library_similarity_threshold*.
 836
 837        All results are stored in :attr:`similarity_matrices`.
 838
 839        Parameters
 840        ----------
 841        query_spectra : list
 842            Spectrum objects with ``.mz_exp`` and ``.abundance`` attributes.
 843        query_ids : list of str
 844            Unique string ID for each spectrum.
 845        query_precursor_mzs : list of float or None, optional
 846            Precursor m/z for each spectrum.  Required (non-None list) for
 847            ``"identity"`` and ``"neutral_loss"`` search types.
 848        fe_kwargs : dict, optional
 849            Extra keyword arguments forwarded to
 850            :class:`ms_entropy.FlashEntropySearch` when building the
 851            temporary query index (stage 1). When None and *settings* was
 852            provided at construction, uses ``settings.as_fe_kwargs()``.
 853        hydrate_library_similarities : bool, optional
 854            When ``True``, run stage 3 (library-vs-library). Default from
 855            *settings* or ``False``.
 856        library_similarity_threshold : float, optional
 857            Minimum entropy similarity score (stage 2) required for a library
 858            spectrum to be included in stage 3. Default from *settings* or 0.3.
 859
 860        Raises
 861        ------
 862        RuntimeError
 863            If queries have already been run (call :meth:`drop_queries` first)
 864            or if ``fe_lib`` is ``None`` (required for stages 2 and 3).
 865        ValueError
 866            If *query_spectra* is empty, lengths are mismatched, or precursor
 867            m/z values are missing for a precursor-required search type.
 868        """
 869        if self._has_queries_run:
 870            raise RuntimeError(
 871                "query_vs_library has already been run — call drop_queries() to clear queries and results before running again."
 872            )
 873
 874        if hydrate_library_similarities is None:
 875            hydrate_library_similarities = (
 876                self._settings.hydrate_library_similarities
 877                if self._settings is not None
 878                else False
 879            )
 880        if library_similarity_threshold is None:
 881            library_similarity_threshold = (
 882                self._settings.library_similarity_threshold
 883                if self._settings is not None
 884                else 0.3
 885            )
 886
 887        self._prepare_queries(
 888            query_spectra=query_spectra,
 889            query_ids=query_ids,
 890            query_precursor_mzs=query_precursor_mzs,
 891        )
 892        self._run_query_vs_query(fe_kwargs=fe_kwargs)
 893        entropy_pairs, _ = self._run_query_vs_library()
 894
 895        if hydrate_library_similarities:
 896            self._run_library_vs_library(
 897                entropy_pairs=entropy_pairs,
 898                library_similarity_threshold=library_similarity_threshold,
 899            )
 900
 901        self._has_queries_run = True
 902
 903    def get_network_edges(
 904        self, metric: str = "entropy_similarity"
 905    ) -> list[tuple[str, str, float]]:
 906        """Return all edges above the similarity threshold for *metric*.
 907
 908        Parameters
 909        ----------
 910        metric : str
 911            Similarity metric to query.  Default ``"entropy_similarity"``.
 912
 913        Returns
 914        -------
 915        list of (id1, id2, score)
 916            Upper-triangle pairs with score >= threshold.
 917        """
 918        mat = self._get_matrix(metric)
 919        threshold = self._threshold_for(metric)
 920        return mat.get_pairs_above_threshold(threshold)
 921
 922    def get_spectrum_neighbors(
 923        self,
 924        spectrum_id: str,
 925        metric: str = "entropy_similarity",
 926    ) -> list[tuple[str, float]]:
 927        """Return all neighbors of *spectrum_id* above the threshold.
 928
 929        Parameters
 930        ----------
 931        spectrum_id : str
 932            Query spectrum ID.
 933        metric : str
 934            Similarity metric to query.  Default ``"entropy_similarity"``.
 935
 936        Returns
 937        -------
 938        list of (neighbor_id, score)
 939            Sorted by score descending.
 940        """
 941        mat = self._get_matrix(metric)
 942        threshold = self._threshold_for(metric)
 943
 944        if spectrum_id not in mat._id_to_idx:
 945            return []
 946
 947        idx = mat._id_to_idx[spectrum_id]
 948        if mat._matrix is None:
 949            return []
 950
 951        mat.finalise()
 952        row = mat._matrix.getrow(idx)
 953        cx = row.tocoo()
 954        neighbors = []
 955        for j, v in zip(cx.col, cx.data):
 956            if v >= threshold and j != idx:
 957                neighbors.append((mat._idx_to_id[j], float(v)))
 958        neighbors.sort(key=lambda x: x[1], reverse=True)
 959        return neighbors
 960
 961    def get_network_stats(self, metric: str = "entropy_similarity") -> dict[str, Any]:
 962        """Return basic network statistics for *metric*.
 963
 964        Parameters
 965        ----------
 966        metric : str
 967            Similarity metric to query.  Default ``"entropy_similarity"``.
 968
 969        Returns
 970        -------
 971        dict with keys:
 972            - ``n_nodes`` : number of registered spectra
 973            - ``n_edges`` : number of edges above threshold
 974            - ``avg_degree`` : average node degree
 975            - ``density`` : fraction of possible edges that exist
 976            - ``threshold`` : the threshold used
 977            - ``metric`` : the metric name
 978        """
 979        mat = self._get_matrix(metric)
 980        threshold = self._threshold_for(metric)
 981        edges = mat.get_pairs_above_threshold(threshold)
 982        n_nodes = mat.n_spectra
 983        n_edges = len(edges)
 984        max_edges = n_nodes * (n_nodes - 1) / 2 if n_nodes > 1 else 1
 985        degrees = {sid: 0 for sid in mat.spectrum_ids}
 986        for id1, id2, _ in edges:
 987            degrees[id1] += 1
 988            degrees[id2] += 1
 989        avg_degree = np.mean(list(degrees.values())) if degrees else 0.0
 990        return {
 991            "metric": metric,
 992            "threshold": threshold,
 993            "n_nodes": n_nodes,
 994            "n_edges": n_edges,
 995            "avg_degree": float(avg_degree),
 996            "density": n_edges / max_edges if max_edges > 0 else 0.0,
 997        }
 998
 999    # ── File output ───────────────────────────────────────────────────────────
1000
1001    def save_edge_list(
1002        self,
1003        path: str,
1004        metric: str = "entropy_similarity",
1005    ):
1006        """Save the edge list as a CSV file.
1007
1008        Parameters
1009        ----------
1010        path : str
1011            Output file path.
1012        metric : str
1013            Similarity metric.  Default ``"entropy_similarity"``.
1014        """
1015        edges = self.get_network_edges(metric=metric)
1016        query_id_set = set(self._all_query_ids)
1017        export_edges = [
1018            (
1019                self._export_node_id(id1, query_id_set),
1020                self._export_node_id(id2, query_id_set),
1021                score,
1022            )
1023            for id1, id2, score in edges
1024        ]
1025        df = pd.DataFrame(export_edges, columns=["id1", "id2", "score"])
1026        df.to_csv(path, index=False)
1027
1028    def save_similarity_matrix(
1029        self,
1030        path: str,
1031        metric: str = "entropy_similarity",
1032        threshold: float = 0.0,
1033    ):
1034        """Save the similarity matrix as a CSV file.
1035
1036        Parameters
1037        ----------
1038        path : str
1039            Output file path.
1040        metric : str
1041            Similarity metric.  Default ``"entropy_similarity"``.
1042        threshold : float
1043            Minimum score to include.  Default 0.0 (all stored pairs).
1044        """
1045        mat = self._get_matrix(metric)
1046        df = mat.to_dataframe(threshold=threshold)
1047        df.to_csv(path, index=False)
1048
1049    # ── Helpers ───────────────────────────────────────────────────────────────
1050
1051    def _get_matrix(self, metric: str) -> SimilarityMatrix:
1052        if metric not in self.similarity_matrices:
1053            raise KeyError(
1054                f"Unknown metric '{metric}'. "
1055                f"Available: {list(self.similarity_matrices.keys())}"
1056            )
1057        return self.similarity_matrices[metric]
1058
1059    def drop_queries(self):
1060        """Clear all stored queries and computed similarity results.
1061
1062        Resets internal query storage, stage flags, the stage-2 entropy-pair
1063        cache, clustering artifacts, and reinitialises all
1064        :attr:`similarity_matrices` to empty so that
1065        :meth:`run_query_vs_query_only` or :meth:`query_vs_library` can be
1066        called again with a consistent cleared state.
1067        """
1068        # Clear stored queries
1069        self._all_query_spectra = []
1070        self._all_query_ids = []
1071        self._all_query_precursor_mzs = []
1072        self._has_queries_run = False
1073        self._stage2_entropy_pairs = None
1074
1075        # Stage flags must match empty matrices (not "stale done")
1076        self.stage_query_query_done = False
1077        self.stage_query_library_done = False
1078        self.stage_library_library_done = False
1079
1080        # Recreate empty similarity matrices for each metric
1081        all_metrics = list(self.similarity_matrices.keys())
1082        self.similarity_matrices = {
1083            m: SimilarityMatrix(metric_name=m) for m in all_metrics
1084        }
1085
1086        # Drop any cached clustering artifacts tied to previous query results.
1087        self._network_clusters = {}
1088
1089    def __repr__(self) -> str:
1090        n_nodes = sum(
1091            mat.n_spectra for mat in self.similarity_matrices.values()
1092        ) // max(1, len(self.similarity_matrices))
1093        metrics = list(self.similarity_matrices.keys())
1094        return (
1095            f"MolecularNetwork(search_type='{self.search_type}', "
1096            f"n_spectra={n_nodes}, metrics={metrics})"
1097        )