"""PRISM - Pruning Interface for Similar Molecules."""

import itertools
import math
from dataclasses import dataclass, field
from time import perf_counter
from typing import Any, Callable, Sequence

import numpy as np
from networkx import Graph, connected_components
from scipy.spatial.distance import cdist

from prism_pruner.algebra import get_inertia_moments
from prism_pruner.graph_manipulations import graphize
from prism_pruner.periodic_table import MASSES_TABLE
from prism_pruner.rmsd import rmsd_and_max
from prism_pruner.torsion_module import (
    _get_rotation_mask,
    get_angles,
    get_hydrogen_bonds,
    get_torsions,
    is_nondummy,
    rotationally_corrected_rmsd_and_max,
)
from prism_pruner.typing import (
    Array1D_bool,
    Array1D_float,
    Array1D_int,
    Array1D_str,
    Array2D_float,
    Array2D_int,
    Array3D_float,
)
from prism_pruner.utils import flatten, get_double_bonds_indices, rotate_dihedral, time_to_string


@dataclass
class PrunerConfig:
    """Configuration dataclass for Pruner."""

    structures: Array3D_float

    # Optional parameters that get initialized
    energies: Array1D_float = field(default_factory=lambda: np.array([]))
    max_dE: float = field(default=0.0)
    debugfunction: Callable[[str], None] | None = field(default=None)

    # Computed fields
    eval_calls: int = field(default=0, init=False)
    cache_calls: int = field(default=0, init=False)
    cache: set[tuple[int, int]] = field(default_factory=lambda: set(), init=False)

    def __post_init__(self) -> None:
        """Validate inputs and initialize computed fields."""
        # validate input types
        assert type(self.structures) is np.ndarray

        self.mask = np.ones(shape=(self.structures.shape[0],), dtype=np.bool_)

        if len(self.energies) != 0:
            assert self.max_dE > 0.0, (
                "If you provide energies, please also provide an appropriate energy window max_dE."
            )

        # Set defaults for optional parameters
        if len(self.energies) == 0:
            assert type(self.energies) is np.ndarray
            self.energies = np.zeros(self.structures.shape[0], dtype=float)

        assert len(self.energies) == len(self.structures), (
            "Please make sure that the energies "
            + "provided have the same len as the input structures."
        )

        if self.max_dE == 0.0:
            self.max_dE = 1.0

    def evaluate_sim(self, *args: Any, **kwargs: Any) -> bool:
        """Stub method - override in subclasses as needed."""
        raise NotImplementedError


@dataclass
class RMSDRotCorrPrunerConfig(PrunerConfig):
    """Configuration dataclass for Pruner."""

    atoms: Array1D_str = field(kw_only=True)
    max_rmsd: float = field(kw_only=True)
    max_dev: float = field(kw_only=True)
    angles: Sequence[Sequence[int]] = field(kw_only=True)
    torsions: Array2D_int = field(kw_only=True)
    graph: Graph = field(kw_only=True)
    heavy_atoms_only: bool = True
    single_atom_masks: list[Array1D_bool] = field(kw_only=True, default_factory=list)
    rotation_masks: list[Array1D_bool] = field(kw_only=True, default_factory=list)

    def __post_init__(self) -> None:
        """Add type enforcing to the parent's __post_init__."""
        super().__post_init__()

        # validate input types
        assert type(self.atoms) is np.ndarray
        assert type(self.graph) is Graph

    def evaluate_sim(self, i1: int, i2: int) -> bool:
        """Return whether the structures are similar."""
        rmsd, max_dev = rotationally_corrected_rmsd_and_max(
            self.structures[i1],
            self.structures[i2],
            atoms=self.atoms,
            torsions=self.torsions,
            graph=self.graph,
            angles=self.angles,
            # debugfunction=self.debugfunction, # lots of printout
            heavy_atoms_only=self.heavy_atoms_only,
            single_atom_masks=self.single_atom_masks or None,
            rotation_masks=self.rotation_masks or None,
        )

        if rmsd > self.max_rmsd:
            return False

        if max_dev > self.max_dev:
            return False

        return True


@dataclass
class RMSDPrunerConfig(PrunerConfig):
    """Configuration dataclass for Pruner."""

    atoms: Array1D_str = field(kw_only=True)
    max_rmsd: float = field(kw_only=True)
    max_dev: float = field(kw_only=True)
    heavy_atoms_only: bool = True

    def __post_init__(self) -> None:
        """Add type enforcing to the parent's __post_init__."""
        super().__post_init__()

        # validate input types
        assert type(self.atoms) is np.ndarray

        # pre-compute heavy atom mask once
        if self.heavy_atoms_only:
            self._heavy_mask: Array1D_bool = self.atoms != "H"
        else:
            self._heavy_mask = np.ones(self.atoms.shape[0], dtype=np.bool_)

    def evaluate_sim(self, i1: int, i2: int) -> bool:
        """Return whether the structures are similar."""
        rmsd, max_dev = rmsd_and_max(
            self.structures[i1][self._heavy_mask],
            self.structures[i2][self._heavy_mask],
            center=True,
        )

        if rmsd > self.max_rmsd:
            return False

        if max_dev > self.max_dev:
            return False

        return True


@dataclass
class MOIPrunerConfig(PrunerConfig):
    """Configuration dataclass for Pruner."""

    masses: Array1D_float = field(kw_only=True)
    max_dev: float = 0.01

    def __post_init__(self) -> None:
        """Add type enforcing and moi_vecs to the parent's __post_init__."""
        super().__post_init__()

        # validate input types
        assert type(self.masses) is np.ndarray

        self.moi_vecs = {
            c: get_inertia_moments(
                coord,
                self.masses,
            )
            for c, coord in enumerate(self.structures)
        }

    def evaluate_sim(self, i1: int, i2: int) -> bool:
        """Return whether the structures are similar."""
        im_1 = self.moi_vecs[i1]
        im_2 = self.moi_vecs[i2]

        # compare the three MOIs via a Python loop:
        # apparently much faster than numpy array operations
        # for such a small array!
        for j in range(3):
            if np.abs(im_1[j] - im_2[j]) / im_1[j] >= self.max_dev:
                return False
        return True


def _main_compute_subrow(
    prunerconfig: PrunerConfig,
    in_mask: Array1D_bool,
    first_abs_index: int,
    num_str_in_row: int,
) -> bool:
    """Evaluate the similarity of a subrow of the similarity matrix.

    Return True if ref is similar to any
    structure in structures, returning at the first instance of a match.
    Ignores structures that are False (0) in in_mask and does not perform
    the comparison if the energy difference between the structures is less
    than self.max_dE. Saves dissimilar structural pairs (i.e. that evaluate to
    False (0)) by adding them to self.cache, avoiding redundant calculations.
    """
    i1 = first_abs_index

    # iterate over target structures
    for i in range(num_str_in_row):
        # only compare active structures
        if in_mask[i]:
            # check if we have performed this comparison already:
            # if so, we already know that those two structures are not similar,
            # since the in_mask attribute is not False for ref nor for i
            i2 = first_abs_index + 1 + i
            hash_value = (i1, i2)

            if hash_value in prunerconfig.cache:
                prunerconfig.cache_calls += 1
                continue

            # if we have not computed the value before, check if the two
            # structures have close enough energy before running the comparison
            elif (
                np.abs(prunerconfig.energies[i1] - prunerconfig.energies[i2]) < prunerconfig.max_dE
            ):
                # function will return True whether the structures are similar,
                # and will stop iterating on this row, returning
                prunerconfig.eval_calls += 1
                if prunerconfig.evaluate_sim(i1, i2):
                    return True

                # if structures are not similar, add the result to the
                # cache, because they will return here,
                # while similar structures are discarded and won't come back
                else:
                    prunerconfig.cache.add(hash_value)

            # if energy is not similar enough, also add to cache
            else:
                prunerconfig.cache.add(hash_value)

    return False


def _main_compute_row(
    prunerconfig: PrunerConfig,
    row_indices: Array1D_int,
    in_mask: Array1D_bool,
    first_abs_index: int,
) -> Array1D_bool:
    """Evaluate the similarity of a row of the similarity matrix.

    For a given set of structures, check if each is similar
    to any other after itself. Return a boolean mask to slice
    the array, only retaining the structures that are dissimilar.
    The inner subrow function caches computed non-similar pairs.

    """
    # initialize the result container
    out_mask = np.ones(shape=in_mask.shape, dtype=np.bool_)
    line_len = len(row_indices)

    # loop over the structures
    for i in range(line_len):
        # only check for similarity if the structure is active
        if in_mask[i]:
            # reject structure i if it is similar to any other after itself
            similar = _main_compute_subrow(
                prunerconfig,
                in_mask[i + 1 :],
                first_abs_index=first_abs_index + i,
                num_str_in_row=line_len - i - 1,
            )
            out_mask[i] = not similar

        else:
            out_mask[i] = False

    return out_mask


def _main_compute_group(
    prunerconfig: PrunerConfig,
    in_mask: Array1D_bool,
    k: int,
) -> Array1D_bool:
    """Evaluate the similarity of each chunk of the similarity matrix.

    Acts individually on k chunks of the structures array,
    returning the updated mask.
    """
    # initialize final result container
    out_mask = np.ones(shape=in_mask.shape, dtype=np.bool_)

    # calculate the size of each chunk
    chunksize = int(len(prunerconfig.structures) // k)

    # iterate over chunks (multithreading here?)
    for chunk in range(int(k)):
        first = chunk * chunksize
        if chunk == k - 1:
            last = len(prunerconfig.structures)
        else:
            last = chunksize * (chunk + 1)

        # get the structure indices chunk
        indices_chunk = np.arange(first, last, 1, dtype=int)

        # compare structures within that chunk and save results to the out_mask
        out_mask[first:last] = _main_compute_row(
            prunerconfig,
            indices_chunk,
            in_mask[first:last],
            first_abs_index=first,
        )
    return out_mask


def _run(prunerconfig: PrunerConfig) -> tuple[Array2D_float, Array1D_bool]:
    """Perform the similarity pruning.

    Remove similar structures by repeatedly grouping them into k
    subgroups and removing similar ones. A cache is present to avoid
    repeating RMSD computations.

    Similarity occurs for structures with both rmsd < self.max_rmsd and
    maximum absolute atomic deviation < self.max_dev.

    Sets the self.structures and the corresponding self.mask attributes.
    """
    start_t = perf_counter()

    # initialize the output mask
    out_mask = np.ones(shape=prunerconfig.structures.shape[0], dtype=np.bool_)
    prunerconfig.cache = set()

    # sort structures by ascending energy: this will have the effect of
    # having energetically similar structures end up in the same chunk
    # and therefore being pruned early
    if np.abs(prunerconfig.energies[-1]) > 0:
        sorting_indices = np.argsort(prunerconfig.energies)
        prunerconfig.structures = prunerconfig.structures[sorting_indices]
        prunerconfig.energies = prunerconfig.energies[sorting_indices]

    # split the structure array in subgroups and prune them internally
    for k in (
        500_000,
        200_000,
        100_000,
        50_000,
        20_000,
        10_000,
        5000,
        2000,
        1000,
        500,
        200,
        100,
        50,
        20,
        10,
        5,
        2,
        1,
    ):
        # choose only k values such that every subgroup
        # has on average at least twenty active structures in it
        if k == 1 or 20 * k < np.count_nonzero(out_mask):
            before = np.count_nonzero(out_mask)

            start_t_k = perf_counter()

            # compute similarities and get back the out_mask
            # and the pairings to be added to cache
            out_mask = _main_compute_group(
                prunerconfig,
                out_mask,
                k=k,
            )

            after = np.count_nonzero(out_mask)
            newly_discarded = before - after

            if prunerconfig.debugfunction is not None:
                elapsed = perf_counter() - start_t_k
                prunerconfig.debugfunction(
                    f"DEBUG: {prunerconfig.__class__.__name__} - k={k}, rejected {newly_discarded} "
                    + f"(keeping {after}/{len(out_mask)}), in {time_to_string(elapsed)}"
                )

    del prunerconfig.cache

    if prunerconfig.debugfunction is not None:
        elapsed = perf_counter() - start_t
        prunerconfig.debugfunction(
            f"DEBUG: {prunerconfig.__class__.__name__} - keeping "
            + f"{after}/{len(out_mask)} "
            + f"({time_to_string(elapsed)})"
        )

        if prunerconfig.eval_calls == 0:
            fraction = 0.0
        else:
            fraction = prunerconfig.cache_calls / (
                prunerconfig.eval_calls + prunerconfig.cache_calls
            )

        prunerconfig.debugfunction(
            f"DEBUG: {prunerconfig.__class__.__name__} - Used cached data "
            + f"{prunerconfig.cache_calls}/{prunerconfig.eval_calls + prunerconfig.cache_calls}"
            + f" times, {100 * fraction:.2f}% of total calls"
        )

    return prunerconfig.structures[out_mask], out_mask


def prune_by_rmsd(
    structures: Array3D_float,
    atoms: Array1D_str,
    max_rmsd: float = 0.25,
    max_dev: float | None = None,
    energies: Array1D_float | None = None,
    max_dE: float = 0.0,
    heavy_atoms_only: bool = True,
    debugfunction: Callable[[str], None] | None = None,
) -> tuple[Array3D_float, Array1D_bool]:
    """Remove duplicate structures using a heavy-atom RMSD metric.

    Remove similar structures by repeatedly grouping them into k
    subgroups and removing similar ones. A cache is present to avoid
    repeating RMSD computations.

    Similarity occurs for structures with both RMSD < max_rmsd and
    maximum deviation < max_dev. max_dev by default is 2 * max_rmsd.
    """
    if energies is None:
        energies = np.zeros(len(structures))
        max_dE = 1.0

    # set default max_dev if not provided
    max_dev = max_dev or 2 * max_rmsd

    N = len(structures)
    if N * N <= _BATCH_THRESHOLD:
        return _batch_rmsd_prune(
            structures=structures,
            atoms=atoms,
            max_rmsd=max_rmsd,
            max_dev=max_dev,
            energies=energies,
            max_dE=max_dE,
            heavy_atoms_only=heavy_atoms_only,
            debugfunction=debugfunction,
        )

    # Sequential fallback for large ensembles.
    prunerconfig = RMSDPrunerConfig(
        structures=structures,
        atoms=atoms,
        max_rmsd=max_rmsd,
        max_dev=max_dev,
        energies=energies,
        max_dE=max_dE,
        debugfunction=debugfunction,
        heavy_atoms_only=heavy_atoms_only,
    )
    return _run(prunerconfig)


def _batch_rmsd_prune(
    structures: Array3D_float,
    atoms: Array1D_str,
    max_rmsd: float,
    max_dev: float,
    energies: Array1D_float,
    max_dE: float,
    heavy_atoms_only: bool,
    debugfunction: Callable[[str], None] | None,
) -> tuple[Array3D_float, Array1D_bool]:
    """Batch vectorized RMSD pruning using numpy's batched SVD.

    Computes all N² pairwise Kabsch RMSDs in one shot using the formula:

        RMSD²(i, j) = ( ‖p_i‖² + ‖q_j‖² - 2·Σσ_k ) / M

    where p_i and q_j are centered heavy-atom coordinate sets and sigma_k are the
    singular values of p_i.T @ q_j (with Kabsch handedness correction on σ₃).
    """
    start_t = perf_counter()

    N = len(structures)
    heavy_mask: Array1D_bool = (
        atoms != "H" if heavy_atoms_only else np.ones(structures.shape[1], dtype=bool)
    )
    M = int(heavy_mask.sum())

    # Center each structure's heavy atoms at the origin
    heavy = structures[:, heavy_mask, :]  # (N, M, 3)
    centered = heavy - heavy.mean(axis=1, keepdims=True)  # (N, M, 3)

    # All N² pairwise 3*3 covariance matrices: cov[i,j] = centered[i].T @ centered[j]
    cov = np.einsum("ima,jmb->ijab", centered, centered)  # (N, N, 3, 3)

    # Batch SVD — singular values only
    S = np.linalg.svd(cov, compute_uv=False)  # (N, N, 3)
    det_sign = np.sign(np.linalg.det(cov))  # (N, N)
    adjusted_sum_S = S[..., 0] + S[..., 1] + det_sign * S[..., 2]  # (N, N)

    # RMSD² for all pairs
    norm_sq = np.sum(centered**2, axis=(-2, -1))  # (N,)
    rmsd_sq = (norm_sq[:, None] + norm_sq[None, :] - 2.0 * adjusted_sum_S) / M
    rmsd_sq = np.maximum(rmsd_sq, 0.0)
    rmsd_all = np.sqrt(rmsd_sq)  # (N, N)

    # Energy filter
    energy_diff = np.abs(energies[:, None] - energies[None, :])
    rmsd_all[energy_diff >= max_dE] = np.inf

    # max_dev verification for candidate pairs (upper triangle only)
    similar = np.zeros((N, N), dtype=bool)
    i_cands, j_cands = np.where(np.triu(rmsd_all < max_rmsd, k=1))
    for i, j in zip(i_cands.tolist(), j_cands.tolist(), strict=True):
        _, max_dev_val = rmsd_and_max(centered[i], centered[j], center=False)
        if max_dev_val < max_dev:
            similar[i, j] = True

    # Greedy sequential pruning
    keep = np.ones(N, dtype=bool)
    for j in range(1, N):
        if np.any(similar[:j, j] & keep[:j]):
            keep[j] = False

    if debugfunction is not None:
        elapsed = perf_counter() - start_t
        n_kept = int(keep.sum())
        debugfunction(
            f"DEBUG: RMSDPrunerConfig - rejected {N - n_kept} "
            f"(keeping {n_kept}/{N}), in {time_to_string(elapsed)}"
        )

    return structures[keep], keep


# Threshold: use batch approach when total work C*N² ≤ this many pair-combo evaluations.
# Above this, fall back to the sequential cache-based approach.
_BATCH_THRESHOLD = 10_000_000


def _batch_rot_corr_prune(
    temp_structures: Array3D_float,
    atoms: Array1D_str,
    torsions_ids: Array2D_int,
    angles: Sequence[Sequence[int]],
    rotation_masks: list[Array1D_bool],
    max_rmsd: float,
    max_dev: float,
    energies: Array1D_float,
    max_dE: float,
    heavy_atoms_only: bool,
    debugfunction: Callable[[str], None] | None,
) -> Array1D_bool:
    """Batch vectorized rot-corr RMSD pruning using numpy's batched SVD.

    Enumerates all C torsion-correction combinations upfront, applies them to all N
    structures, then computes every pairwise rotationally-corrected RMSD in one
    batched SVD call.  The RMSD is derived analytically from singular values:

        RMSD²(i, c, j) = ( ‖p_i‖² + ‖q_{c,j}‖² - 2·Σσ_k ) / M

    where p_i = centered ref structure i, q_{c,j} = centered corrected structure j
    under combo c, and the sign of σ₃ is flipped when det(cov) < 0 (Kabsch
    handedness correction).  The minimum RMSD over all combos is the
    rotationally-corrected RMSD for that pair.
    """
    start_t = perf_counter()

    N = len(temp_structures)
    natoms = temp_structures.shape[1]
    heavy_mask: Array1D_bool = atoms != "H" if heavy_atoms_only else np.ones(natoms, dtype=bool)
    M = int(heavy_mask.sum())

    # --- Step 1: generate all C correction combinations and apply to all structures ---
    combos = list(itertools.product(*angles))
    C = len(combos)

    corrected_all = np.empty((C, N, natoms, 3))
    for c, combo in enumerate(combos):
        corrected_all[c] = temp_structures
        for k, angle in enumerate(combo):
            if abs(angle) > 1e-3:
                for i in range(N):
                    rotate_dihedral(
                        corrected_all[c, i], torsions_ids[k], angle, mask=rotation_masks[k]
                    )

    # --- Step 2: center all heavy-atom coordinate sets at the origin ---
    ref_heavy = temp_structures[:, heavy_mask, :]  # (N, M, 3)
    centered_ref = ref_heavy - ref_heavy.mean(axis=1, keepdims=True)

    corr_heavy = corrected_all[:, :, heavy_mask, :]  # (C, N, M, 3)
    centered_corr = corr_heavy - corr_heavy.mean(axis=2, keepdims=True)

    # --- Step 3: batch all N*C*N pairwise 3*3 covariance matrices ---
    # cov[i,c,j,a,b] = Σ_m centered_ref[i,m,a] · centered_corr[c,j,m,b]
    #                 = centered_ref[i].T @ centered_corr[c,j]
    cov = np.einsum("ima,cjmb->icjab", centered_ref, centered_corr)  # (N, C, N, 3, 3)

    # --- Step 4: batch SVD (singular values only) + Kabsch handedness correction ---
    S = np.linalg.svd(cov, compute_uv=False)  # (N, C, N, 3)
    det_sign = np.sign(np.linalg.det(cov))  # (N, C, N)
    adjusted_sum_S = S[..., 0] + S[..., 1] + det_sign * S[..., 2]  # (N, C, N)

    # --- Step 5: RMSD² for every (ref i, combo c, target j) ---
    norm_ref_sq = np.sum(centered_ref**2, axis=(-2, -1))  # (N,)
    norm_corr_sq = np.sum(centered_corr**2, axis=(-2, -1))  # (C, N)
    rmsd_sq = (
        norm_ref_sq[:, None, None] + norm_corr_sq[None, :, :] - 2.0 * adjusted_sum_S
    ) / M  # (N, C, N)
    rmsd_sq = np.maximum(rmsd_sq, 0.0)

    # --- Step 6: best rotationally-corrected RMSD per pair ---
    best_c = rmsd_sq.argmin(axis=1)  # (N, N)
    rmsd_min = np.sqrt(rmsd_sq.min(axis=1))  # (N, N)

    # Energy filter: exclude pairs outside the energy window
    energy_diff = np.abs(energies[:, None] - energies[None, :])  # (N, N)
    rmsd_min[energy_diff >= max_dE] = np.inf

    # --- Step 7: max_dev verification for candidate similar pairs ---
    # Only the upper triangle matters (j gets removed only if similar to an earlier i).
    similar = np.zeros((N, N), dtype=bool)
    i_cands, j_cands = np.where(np.triu(rmsd_min < max_rmsd, k=1))
    for i, j in zip(i_cands.tolist(), j_cands.tolist(), strict=True):
        c_star = int(best_c[i, j])
        _, max_dev_val = rmsd_and_max(centered_ref[i], centered_corr[c_star, j], center=False)
        if max_dev_val < max_dev:
            similar[i, j] = True

    # --- Step 8: greedy sequential pruning ---
    keep = np.ones(N, dtype=bool)
    for j in range(1, N):
        if np.any(similar[:j, j] & keep[:j]):
            keep[j] = False

    if debugfunction is not None:
        elapsed = perf_counter() - start_t
        n_kept = int(keep.sum())
        debugfunction(
            f"DEBUG: RMSDRotCorrPrunerConfig - rejected {N - n_kept} "
            f"(keeping {n_kept}/{N}), in {time_to_string(elapsed)}"
        )

    return keep


def prune_by_rmsd_rot_corr(
    structures: Array3D_float,
    atoms: Array1D_str,
    graph: Graph,
    max_rmsd: float = 0.25,
    max_dev: float | None = None,
    energies: Array1D_float | None = None,
    max_dE: float = 0.0,
    heavy_atoms_only: bool = True,
    logfunction: Callable[[str], None] | None = None,
    debugfunction: Callable[[str], None] | None = None,
) -> tuple[Array3D_float, Array1D_bool]:
    """Remove duplicates using a heavy-atom RMSD metric, corrected for degenerate torsions.

    Remove similar structures by repeatedly grouping them into k
    subgroups and removing similar ones. A cache is present to avoid
    repeating RMSD computations.

    Similarity occurs for structures with both RMSD < max_rmsd and
    maximum deviation < max_dev. max_dev by default is 2 * max_rmsd.

    The RMSD and maximum deviation metrics used are the lowest ones
    of all the degenerate rotamers of the input structure.
    """
    # center structures
    temp_structures = np.array([s - s.mean(axis=0) for s in structures])
    ref = temp_structures[0]

    # get the number of molecular fragments
    subgraphs = list(connected_components(graph))

    # if they are more than two, give up on pruning by rot corr rmsd
    if len(subgraphs) > 2:
        return structures, np.ones(structures.shape[0], dtype=bool)

    # if they are two, we can add a fictitious bond between the closest
    # atoms on the two molecular fragment in the provided graph, and
    # then removing it before returning
    if len(subgraphs) == 2:
        subgraphs = [list(vals) for vals in connected_components(graph)]
        all_dists_array = cdist(ref[list(subgraphs[0])], ref[list(subgraphs[1])])
        min_d = np.min(all_dists_array)
        s1, s2 = np.where(all_dists_array == min_d)
        i1, i2 = subgraphs[0][s1[0]], subgraphs[1][s2[0]]
        graph.add_edge(i1, i2)

        if debugfunction is not None:
            debugfunction(
                f"DEBUG: prune_by_rmsd_rot_corr - temporarily added "
                f"edge {i1}-{i2} to the graph (will be removed before returning)"
            )

    # set default max_dev if not provided
    max_dev = max_dev or 2 * max_rmsd

    # add hydrogen bonds to molecular graph
    hydrogen_bonds = get_hydrogen_bonds(ref, atoms, graph)
    for hb in hydrogen_bonds:
        graph.add_edge(*hb)

    # keep an unraveled set of atoms in hbs
    flat_hbs = set(flatten(hydrogen_bonds))

    # get all rotable bonds in the molecule, including dummy rotations
    torsions = get_torsions(
        ref,
        graph,
        hydrogen_bonds=hydrogen_bonds,
        double_bonds=get_double_bonds_indices(ref, atoms),
        keepdummy=True,
        mode="symmetry",
    )

    # only keep dummy rotations (checking both directions)
    torsions = [
        t
        for t in torsions
        if not (is_nondummy(t.i2, t.i3, graph) and (is_nondummy(t.i3, t.i2, graph)))
    ]

    # since we only compute RMSD based on heavy atoms, discard
    # quadruplets that involve hydrogen atom as termini, unless
    # they are involved in hydrogen bonding
    torsions = [
        t
        for t in torsions
        if ("H" not in [atoms[i] for i in t.torsion])
        or (t.torsion[0] in flat_hbs or t.torsion[3] in flat_hbs)
    ]
    # get torsions angles
    angles = [get_angles(t, graph) for t in torsions]

    # Used specific directionality of torsions so that we always
    # rotate the dummy portion (the one attached to the last index)
    torsions_ids = np.asarray(
        [
            list(t.torsion) if is_nondummy(t.i2, t.i3, graph) else list(reversed(t.torsion))
            for t in torsions
        ]
    )

    # Set up final mask and cache
    final_mask = np.ones(structures.shape[0], dtype=bool)

    # Halt the run if there are too many structures or no subsymmetrical bonds
    if len(torsions_ids) == 0:
        if debugfunction is not None:
            debugfunction(
                "DEBUG: prune_by_rmsd_rot_corr - No subsymmetrical torsions found: skipping "
                "symmetry-corrected RMSD pruning"
            )

        return structures[final_mask], final_mask

    # Print out torsion information
    if logfunction is not None:
        logfunction("\n >> Dihedrals considered for rotamer corrections:")
        for i, (torsion, angle) in enumerate(zip(torsions_ids, angles, strict=False)):
            logfunction(
                " {:2s} - {:21s} : {}{}{}{} : {}-fold".format(
                    str(i + 1),
                    str(torsion),
                    atoms[torsion[0]],
                    atoms[torsion[1]],
                    atoms[torsion[2]],
                    atoms[torsion[3]],
                    len(angle),
                )
            )
        logfunction("\n")

    if energies is None:
        energies = np.array([])

    # Pre-compute rotation masks (used by both batch and sequential paths).
    rotation_masks = [_get_rotation_mask(graph, tuple(t)) for t in torsions_ids]

    N = len(temp_structures)
    C = math.prod(len(a) for a in angles)

    if C * N * N <= _BATCH_THRESHOLD:
        # Batch path: enumerate all C correction combinations, compute all N*C*N
        # pairwise RMSDs at once via numpy batched SVD.  Much faster for small-to-
        # medium ensembles because it replaces thousands of Python function calls
        # with a handful of vectorised numpy operations.
        mask = _batch_rot_corr_prune(
            temp_structures=temp_structures,
            atoms=atoms,
            torsions_ids=torsions_ids,
            angles=angles,
            rotation_masks=rotation_masks,
            max_rmsd=max_rmsd,
            max_dev=max_dev,
            energies=energies,
            max_dE=max_dE,
            heavy_atoms_only=heavy_atoms_only,
            debugfunction=debugfunction,
        )
    else:
        # Sequential fallback for large ensembles or many torsion combinations:
        # enumerate pairs one at a time with a cache to skip recomputation.
        single_atom_masks = []
        for t in torsions_ids:
            m = np.zeros(len(atoms), dtype=bool)
            m[t[3]] = True
            single_atom_masks.append(m)

        prunerconfig = RMSDRotCorrPrunerConfig(
            structures=temp_structures,
            atoms=atoms,
            energies=energies,
            max_dE=max_dE,
            graph=graph,
            torsions=torsions_ids,
            debugfunction=debugfunction,
            heavy_atoms_only=heavy_atoms_only,
            angles=angles,
            max_rmsd=max_rmsd,
            max_dev=max_dev,
            single_atom_masks=single_atom_masks,
            rotation_masks=rotation_masks,
        )
        _, mask = _run(prunerconfig)

    # remove the extra bond in the molecular graph
    if len(subgraphs) == 2:
        graph.remove_edge(i1, i2)

    # return the original coordinates (and not the temp)
    # to make sure they return untouched by the function
    return structures[mask], mask


def prune_by_moment_of_inertia(
    structures: Array3D_float,
    atoms: Array1D_str,
    max_deviation: float = 1e-2,
    energies: Array1D_float | None = None,
    max_dE: float = 0.0,
    debugfunction: Callable[[str], None] | None = None,
) -> tuple[Array3D_float, Array1D_bool]:
    """Remove duplicate structures using a moments of inertia-based metric.

    Remove duplicate structures (enantiomeric or rotameric) based on the
    moment of inertia on the principal axes. If all three deviate less than
    max_deviation percent from another one, the structure is removed from
    the ensemble (i.e. max_deviation = 0.1 is 10% relative deviation).
    """
    if energies is None:
        energies = np.array([])

    # set up PrunerConfig dataclass
    prunerconfig = MOIPrunerConfig(
        structures=structures,
        energies=energies,
        max_dE=max_dE,
        debugfunction=debugfunction,
        max_dev=max_deviation,
        masses=np.array([MASSES_TABLE[a] for a in atoms]),
    )

    return _run(prunerconfig)


def prune(
    structures: Array3D_float,
    atoms: Array1D_str,
    moi_pruning: bool = True,
    rmsd_pruning: bool = True,
    rot_corr_rmsd_pruning: bool = False,
    energies: Array1D_float | None = None,
    max_dE: float = 0.0,
    debugfunction: Callable[[str], None] | None = None,
    logfunction: Callable[[str], None] | None = None,
) -> tuple[Array3D_float, Array1D_bool]:
    """Remove duplicate structures.

    Chains the three main pruning modes on the
    input ensemble, unless prompted otherwise.

    Will only compare structures less than max_dE apart
    in energy, if energies and max_dE are provided.

    Note: will use automatic pruning thresholds.
    """
    if energies is None:
        energies = np.array([0.0 for _ in range(len(structures))])
        max_dE = 1.0

    active_ens = structures
    active_indices = np.arange(structures.shape[0])

    if moi_pruning:
        active_ens, mask = prune_by_moment_of_inertia(
            structures=active_ens,
            atoms=atoms,
            max_deviation=0.01,
            energies=energies,
            max_dE=max_dE,
            debugfunction=debugfunction,
        )
        energies = energies[mask]
        active_indices = active_indices[mask]

        # add space between different logs
        if debugfunction is not None:
            debugfunction("")

    if rmsd_pruning:
        active_ens, mask = prune_by_rmsd(
            structures=active_ens,
            atoms=atoms,
            max_rmsd=0.25,
            max_dev=0.5,
            energies=energies,
            max_dE=max_dE,
            debugfunction=debugfunction,
        )
        energies = energies[mask]
        active_indices = active_indices[mask]

        # add space between different logs
        if debugfunction is not None:
            debugfunction("")

    if rot_corr_rmsd_pruning:
        graph = graphize(atoms, active_ens[0])

        active_ens, mask = prune_by_rmsd_rot_corr(
            structures=active_ens,
            atoms=atoms,
            graph=graph,
            max_rmsd=0.25,
            max_dev=0.5,
            energies=energies,
            max_dE=max_dE,
            debugfunction=debugfunction,
            logfunction=logfunction,
        )
        active_indices = active_indices[mask]

    # now backtrack the effect of all the pruning
    # so that we know which structures we got rid of
    # and we can return the appropriate boolean mask
    cumulative_mask = np.zeros(structures.shape[0], dtype=np.bool_)
    cumulative_mask[active_indices] = True

    return active_ens, cumulative_mask
