# ezga/utils/molecule_blacklist.py
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import numpy as np

from .chemistry import pair_cutoff, get_positions_species, neighbors_within


# ============================================================
# Performance / geometry knobs
# ============================================================
CUTOFF_MAX: float = 4.0                  # global radius for candidate neighbors
MAX_NEIGHBORS_FALLBACK: int = 96         # only used in fallback radius builder
CELL_MIN: float = 1e-9


# ============================================================
# Best-effort: discover cell + pbc
# ============================================================
def _try_get_cell_pbc(apm) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
    cell = None
    pbc = None

    for attr in ("cell", "Cell", "CELL"):
        if hasattr(apm, attr):
            try:
                cell = np.asarray(getattr(apm, attr), dtype=float)
            except Exception:
                cell = None
            break

    for attr in ("pbc", "PBC"):
        if hasattr(apm, attr):
            try:
                pbc = np.asarray(getattr(apm, attr), dtype=bool).reshape(3)
            except Exception:
                pbc = None
            break

    if cell is not None and cell.shape == (3,):
        cell = np.diag(cell)

    if cell is not None:
        cell = np.asarray(cell, dtype=float)
        if cell.shape != (3, 3):
            cell = None

    return cell, pbc


def _is_orthorhombic(cell: np.ndarray, tol: float = 1e-8) -> bool:
    off = cell - np.diag(np.diag(cell))
    return float(np.max(np.abs(off))) < tol


def _wrap_orthorhombic(dr: np.ndarray, L: np.ndarray, pbc: np.ndarray) -> np.ndarray:
    out = dr.copy()
    for a in range(3):
        if pbc[a]:
            out[..., a] -= L[a] * np.round(out[..., a] / L[a])
    return out


# ============================================================
# Radius graph (CSR)
# ============================================================
@dataclass
class RadiusGraph:
    R: np.ndarray                 # (N,3)
    S: Tuple[str, ...]            # (N,)
    indptr: np.ndarray            # (N+1,)
    indices: np.ndarray           # (E,)
    distances: np.ndarray         # (E,)
    cell: Optional[np.ndarray]    # (3,3) or None
    pbc: Optional[np.ndarray]     # (3,) or None

    @property
    def N(self) -> int:
        return int(len(self.S))


def build_radius_graph(apm, cutoff: float = CUTOFF_MAX) -> RadiusGraph:
    """
    Builds an undirected radius graph using the structure's kdtree or cKDTree fallback.
    """
    R, S = get_positions_species(apm)
    if not isinstance(R, np.ndarray):
        R = np.asarray(R, dtype=float)
    S = tuple(str(s) for s in S)
    N = len(S)

    cell, pbc = _try_get_cell_pbc(apm)
    
    # Use existing kdtree if available
    if hasattr(apm, 'kdtree') and apm.kdtree is not None:
        tree = apm.kdtree
    else:
        # Fallback to plain cKDTree
        from scipy.spatial import cKDTree
        if cell is not None and pbc is not None and all(pbc):
            box = cell if cell.ndim == 1 else np.diag(cell)
            tree = cKDTree(R, boxsize=box)
        else:
            tree = cKDTree(R)

    # Get dist map directly to avoid distance recalculation and handle PBC naturally.
    cutoff_f = float(cutoff)
    if hasattr(tree, 'sparse_distance_matrix2'):
        dist_map = tree.sparse_distance_matrix2(tree, cutoff_f)
    else:
        dist_map = tree.sparse_distance_matrix(tree, cutoff_f, output_type='dict')
    
    # dist_map is {(i, j): dist, ...}. Filter to uniquely i < j.
    ei, ej, ed = [], [], []
    for (i, j), d in dist_map.items():
        if i < j:
            ei.append(i)
            ej.append(j)
            ed.append(float(d))
            
    if not ei:
        indptr = np.zeros(N + 1, dtype=np.int64)
        return RadiusGraph(R=R, S=S, indptr=indptr,
                           indices=np.empty(0, dtype=np.int64),
                           distances=np.empty(0, dtype=float),
                           cell=cell, pbc=pbc)

    i_idx = np.array(ei, dtype=np.int64)
    j_idx = np.array(ej, dtype=np.int64)
    dists = np.array(ed, dtype=float)
    
    # Symmetric edges
    si = np.concatenate([i_idx, j_idx])
    sj = np.concatenate([j_idx, i_idx])
    sd = np.concatenate([dists, dists])

    # Sort by source index for CSR
    order = np.lexsort((sj, si))
    si, sj, sd = si[order], sj[order], sd[order]
    
    indptr = np.zeros(N + 1, dtype=np.int64)
    np.add.at(indptr, si + 1, 1)
    indptr = np.cumsum(indptr)
    
    return RadiusGraph(R=R, S=S, indptr=indptr, indices=sj, distances=sd, cell=cell, pbc=pbc)



# ============================================================
# Bond graph (CSR) derived from radius graph via cutoff matrix
# ============================================================
_CUTOFF_CACHE = {}

@dataclass
class BondGraph:
    rg: RadiusGraph
    # species encoding (per-structure)
    species: Tuple[str, ...]          # unique symbols in this structure
    sym2id: Dict[str, int]
    sid: np.ndarray                  # (N,) int16
    cutoff_mat: np.ndarray           # (M,M) float32
    bindptr: np.ndarray              # (N+1,)
    bindices: np.ndarray             # (Eb,)
    bdist: np.ndarray                # (Eb,)

    @property
    def N(self) -> int:
        return self.rg.N

    @property
    def M(self) -> int:
        return int(len(self.species))


def build_bond_graph(rg: RadiusGraph) -> BondGraph:
    S = rg.S
    uniq = tuple(sorted(set(S)))
    sym2id = {s: i for i, s in enumerate(uniq)}
    sid = np.array([sym2id[s] for s in S], dtype=np.int16)

    M = len(uniq)
    key = tuple(uniq)
    if key not in _CUTOFF_CACHE:
        cutoff_mat = np.zeros((M, M), dtype=np.float32)
        for i, a in enumerate(uniq):
            for j, b in enumerate(uniq):
                cutoff_mat[i, j] = float(pair_cutoff(a, b))
        _CUTOFF_CACHE[key] = cutoff_mat
    else:
        cutoff_mat = _CUTOFF_CACHE[key]

    N = rg.N
    if N == 0:
        return BondGraph(
            rg=rg, species=uniq, sym2id=sym2id, sid=sid,
            cutoff_mat=cutoff_mat,
            bindptr=np.zeros(1, dtype=np.int64),
            bindices=np.empty(0, dtype=np.int64),
            bdist=np.empty(0, dtype=float)
        )

    # Vectorized bond filter
    si = np.repeat(np.arange(N), rg.indptr[1:] - rg.indptr[:-1])
    sj = rg.indices
    sd = rg.distances

    sid_i = sid[si]
    sid_j = sid[sj]

    is_bond = sd <= cutoff_mat[sid_i, sid_j]

    bsi = si[is_bond]
    bsj = sj[is_bond]
    bd = sd[is_bond]

    bindptr = np.zeros(N + 1, dtype=np.int64)
    np.add.at(bindptr, bsi + 1, 1)
    bindptr = np.cumsum(bindptr)

    return BondGraph(
        rg=rg,
        species=uniq,
        sym2id=sym2id,
        sid=sid,
        cutoff_mat=cutoff_mat,
        bindptr=bindptr,
        bindices=bsj,
        bdist=bd,
    )


# ============================================================
# Environment (precomputed once, O(N+E))
# ============================================================
@dataclass
class Env:
    """
    Fast, generic atom environment derived from BondGraph.
    Everything is species-agnostic.
    """
    bg: BondGraph
    # per-atom features
    z: np.ndarray                 # (N,)
    coord: np.ndarray             # (N,) int16
    neigh_counts: np.ndarray      # (N,M) int16  (count of bonded neighbor species ids)
    atoms_by_species: List[np.ndarray]

    def count(self, i: int, sym: str) -> int:
        j = self.bg.sym2id.get(sym)
        if j is None:
            return 0
        return int(self.neigh_counts[i, j])

    def species(self, i: int) -> str:
        return self.bg.rg.S[i]


def build_env(bg: BondGraph) -> Env:
    N = bg.N
    M = bg.M
    z = bg.rg.R[:, 2].astype(float)
    coord = (bg.bindptr[1:] - bg.bindptr[:-1]).astype(np.int16)

    neigh_counts = np.zeros((N, M), dtype=np.int16)
    if N > 0:
        rows = np.repeat(np.arange(N), coord)
        cols = bg.sid[bg.bindices]
        np.add.at(neigh_counts, (rows, cols), 1)

    atoms_by_species = [
        np.where(bg.sid == i)[0].astype(np.int64)
        for i in range(M)
    ]

    return Env(bg=bg, z=z, coord=coord, neigh_counts=neigh_counts, atoms_by_species=atoms_by_species)


# ============================================================
# Generic predicates + selector (compiled)
# ============================================================
PredicateFn = Callable[[Env, np.ndarray], np.ndarray]
# Given env and an array of candidate indices, return mask (boolean) or filtered indices.


@dataclass(frozen=True)
class CompiledRule:
    tag: str
    deny: bool                 # True => structure is rejected if any match
    remove: bool               # True => remove matching atoms and continue
    # Which atoms to consider (species filter); None => all atoms
    species_ids: Optional[np.ndarray]
    # predicate pipeline: applied to candidate indices
    predicates: Tuple[PredicateFn, ...]

    def run(self, apm, env: Env) -> Tuple[bool, Optional[np.ndarray]]:
        """
        Returns (triggered, idx_to_remove_if_remove).
        If deny and any match => triggered True.
        If remove and any match => return (False, idx) without mutating apm yet.
        """
        N = env.bg.N
        if self.species_ids is None:
            cand = np.arange(N, dtype=np.int64)
        else:
            if len(self.species_ids) > 0:
                cand = np.concatenate([env.atoms_by_species[i] for i in self.species_ids])
            else:
                cand = np.array([], dtype=np.int64)

        if cand.size == 0:
            return (False, None)

        idx = cand
        for pred in self.predicates:
            m = pred(env, idx)
            # pred returns filtered indices (preferred)
            idx = m
            if idx.size == 0:
                return (False, None)

        if self.remove:
            return (False, idx)

        if self.deny:
            return (True, None)

        # If neither deny nor remove, treat as "tag-only deny" behavior (conservative)
        return (True, None)


# ============================================================
# Predicate builders (vectorized on index arrays)
# ============================================================
def _pred_coord_eq(k: int) -> PredicateFn:
    k = int(k)
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        return idx[env.coord[idx] == k]
    return f

def _pred_coord_min(k: int) -> PredicateFn:
    k = int(k)
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        return idx[env.coord[idx] >= k]
    return f

def _pred_coord_max(k: int) -> PredicateFn:
    k = int(k)
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        return idx[env.coord[idx] <= k]
    return f

def _pred_neighbor_count(sym: str, op: str, val: int) -> PredicateFn:
    op = str(op)
    val = int(val)
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        j = env.bg.sym2id.get(sym)
        if j is None:
            # if species not present, counts are all 0
            counts = np.zeros(idx.size, dtype=np.int16)
        else:
            counts = env.neigh_counts[idx, j]
        if op == "==":
            keep = counts == val
        elif op == ">=":
            keep = counts >= val
        elif op == "<=":
            keep = counts <= val
        elif op == ">":
            keep = counts > val
        elif op == "<":
            keep = counts < val
        else:
            raise ValueError(f"Unsupported op '{op}'")
        return idx[keep]
    return f

def _pred_neighbors_only(allowed_syms: Sequence[str]) -> PredicateFn:
    allowed_syms = tuple(str(s) for s in allowed_syms)
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        M = env.bg.M
        allowed = np.zeros(M, dtype=bool)
        for s in allowed_syms:
            j = env.bg.sym2id.get(s)
            if j is not None:
                allowed[j] = True

        counts = env.neigh_counts[idx]              # (k,M)
        # Any neighbor in disallowed species => reject
        disallowed = ~allowed
        bad = (counts[:, disallowed].sum(axis=1) > 0)
        return idx[~bad]
    return f

def _pred_neighbors_same_species() -> PredicateFn:
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        sid = env.bg.sid
        counts = env.neigh_counts[idx]  # (k,M)
        own = sid[idx]
        own_counts = counts[np.arange(idx.size), own]
        # all bonded neighbors are of same species => coord == own_counts
        keep = own_counts == env.coord[idx]
        return idx[keep]
    return f

def _pred_z_outside(zmin: float, zmax: float) -> PredicateFn:
    zmin = float(zmin)
    zmax = float(zmax)
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        z = env.z[idx]
        return idx[(z < zmin) | (z > zmax)]
    return f

def _pred_z_between(zmin: float, zmax: float) -> PredicateFn:
    zmin = float(zmin)
    zmax = float(zmax)
    def f(env: Env, idx: np.ndarray) -> np.ndarray:
        z = env.z[idx]
        return idx[(z >= zmin) & (z <= zmax)]
    return f


# ============================================================
# User DSL (chainable)
# ============================================================
class AtomQuery:
    """
    Chainable query builder:
        forbid(atoms("H").coord_eq(1).neighbors_only(["H"]))
    """
    def __init__(self, species: Optional[Sequence[str]] = None):
        self._species = None if species is None else tuple(str(s) for s in species)
        self._preds: List[PredicateFn] = []

    # --- selection ---
    @property
    def species(self) -> Optional[Tuple[str, ...]]:
        return self._species

    # --- predicates ---
    def coord_eq(self, k: int) -> "AtomQuery":
        self._preds.append(_pred_coord_eq(k))
        return self

    def coord_min(self, k: int) -> "AtomQuery":
        self._preds.append(_pred_coord_min(k))
        return self

    def coord_max(self, k: int) -> "AtomQuery":
        self._preds.append(_pred_coord_max(k))
        return self

    def neighbors(self, sym: str, op: str, val: int) -> "AtomQuery":
        self._preds.append(_pred_neighbor_count(sym, op, val))
        return self

    def neighbors_eq(self, sym: str, val: int) -> "AtomQuery":
        return self.neighbors(sym, "==", val)

    def neighbors_min(self, sym: str, val: int) -> "AtomQuery":
        return self.neighbors(sym, ">=", val)

    def neighbors_max(self, sym: str, val: int) -> "AtomQuery":
        return self.neighbors(sym, "<=", val)

    def neighbors_only(self, allowed_syms: Sequence[str]) -> "AtomQuery":
        self._preds.append(_pred_neighbors_only(allowed_syms))
        return self

    def neighbors_same_species(self) -> "AtomQuery":
        self._preds.append(_pred_neighbors_same_species())
        return self

    def z_outside(self, zmin: float, zmax: float) -> "AtomQuery":
        self._preds.append(_pred_z_outside(zmin, zmax))
        return self

    def z_between(self, zmin: float, zmax: float) -> "AtomQuery":
        self._preds.append(_pred_z_between(zmin, zmax))
        return self

    # --- compile ---
    def compile(self, env: Env) -> Tuple[Optional[np.ndarray], Tuple[PredicateFn, ...]]:
        if self._species is None:
            species_ids = None
        else:
            ids = []
            for s in self._species:
                j = env.bg.sym2id.get(s)
                if j is not None:
                    ids.append(int(j))
            species_ids = None if len(ids) == 0 else np.array(sorted(set(ids)), dtype=np.int16)
        return species_ids, tuple(self._preds)


def atoms(*species: str) -> AtomQuery:
    if len(species) == 0:
        return AtomQuery(None)
    return AtomQuery(species)


# ============================================================
# Rule constructors (deny / remove)
# ============================================================
@dataclass(frozen=True)
class RuleSpec:
    action: str                 # "DENY" or "REMOVE"
    query: AtomQuery
    tag: str


def deny(query: AtomQuery, tag: str) -> RuleSpec:
    return RuleSpec(action="DENY", query=query, tag=str(tag))

def remove(query: AtomQuery, tag: str) -> RuleSpec:
    return RuleSpec(action="REMOVE", query=query, tag=str(tag))

# Backward-ish naming
forbid = deny


# ============================================================
# Engine
# ============================================================
class ConstraintEngine:
    """
    Global, MD-ready, species-agnostic constraint engine.

    Pipeline per structure:
      1) Build RadiusGraph (cell-list if orthorhombic cell is available; else fallback)
      2) Build BondGraph using cutoff matrix (pair_cutoff for all present species)
      3) Build Env (coordination + neighbor species counts)
      4) Compile RuleSpecs into CompiledRules (species-indexed) for this structure
      5) Evaluate rules (early exit on DENY; REMOVE mutates apm and continues)
    """

    def __init__(
        self,
        rules: Iterable[RuleSpec] = (),
        cutoff_max: float = CUTOFF_MAX,
        fail_soft: bool = True,
    ):
        self.rules: List[RuleSpec] = list(rules)
        self.cutoff_max = float(cutoff_max)
        self.fail_soft = bool(fail_soft)

    def add(self, spec: RuleSpec) -> None:
        self.rules.append(spec)

    def contains(self, struct, *, remove_override: Optional[bool] = None) -> Tuple[bool, Optional[str]]:
        apm = struct.AtomPositionManager

        # Convergence loop: repeatedly build graphs and evaluate all rules.
        # If any atoms are marked for removal, we actually remove them in a single batch
        # at the end of the pass, then loop again to ensure the newly mutated structure
        # still satisfies all rules (no chain-reaction violations left behind).
        # We cap at 10 iterations as a safety net against pathological cyclic removals.
        max_iters = 10
        for _ in range(max_iters):
            try:
                rg = build_radius_graph(apm, cutoff=self.cutoff_max)
                bg = build_bond_graph(rg)
                env = build_env(bg)
            except Exception:
                if self.fail_soft:
                    return (False, None)
                raise

            # Compile and cache for this structure based on species present and overriding
            if not hasattr(self, '_compiled_cache'):
                self._compiled_cache = {}
                
            species_key = tuple(env.bg.species)
            cache_key = (species_key, remove_override)
            
            if cache_key not in self._compiled_cache:
                compiled_rules: List[CompiledRule] = []
                for spec in self.rules:
                    try:
                        species_ids, preds = spec.query.compile(env)

                        action = spec.action.upper()
                        do_remove = (action == "REMOVE")
                        do_deny = (action == "DENY")

                        # allow per-call override (useful to switch between "repair" and "reject")
                        if isinstance(remove_override, bool):
                            do_remove = bool(remove_override)
                            do_deny = not do_remove

                        compiled_rules.append(
                            CompiledRule(
                                tag=spec.tag,
                                deny=do_deny,
                                remove=do_remove,
                                species_ids=species_ids,
                                predicates=preds,
                            )
                        )
                    except Exception:
                        if not self.fail_soft:
                            raise
                        continue
                self._compiled_cache[cache_key] = compiled_rules

            compiled = self._compiled_cache[cache_key]

            N = env.bg.N
            alive = np.ones(N, dtype=bool)
            remove_idx = []

            # Evaluate in user order
            for rule in compiled:
                try:
                    # Evaluating rule doesn't mutate apm directly anymore
                    triggered, idx = rule.run(apm, env)
                    
                    if rule.deny and triggered:
                        return (True, rule.tag)

                    if rule.remove and idx is not None and idx.size > 0:
                        alive[idx] = False
                        remove_idx.append(idx)

                except Exception:
                    if not self.fail_soft:
                        raise
                    continue

            # If nothing was marked for removal, the structure has converged and is clean
            if not remove_idx:
                return (False, None)
                
            # Otherwise, perform the single batch removal and loop again
            apm.remove_atom(np.where(~alive)[0])

        # If we hit the max iterations, act conservatively and fail the structure
        return (True, "MAX_REMOVAL_ITERATIONS_EXCEEDED")


# ============================================================
# Example presets (modern replacements for old rules)
# ============================================================
def preset_free_diatomic(sym: str, *, tag: Optional[str] = None) -> RuleSpec:
    """
    Deny any "free diatomic" of a given species sym:
      sym atom with coord==1 and neighbor only sym.
    This generalizes H2, O2, N2, Cl2, etc.
    """
    t = tag or f"FREE_{sym}2"
    q = atoms(sym).coord_eq(1).neighbors_only([sym])
    return deny(q, t)

def preset_free_diatomic_any(*, tag: str = "FREE_DIATOMIC_ANY") -> RuleSpec:
    """
    Deny any free diatomic A2 for any species A:
      coord==1 and neighbors_same_species.
    """
    q = atoms().coord_eq(1).neighbors_same_species()
    return deny(q, tag)

def preset_outside_z(zmin: float, zmax: float, *, tag: str = "OUTSIDE_Z") -> RuleSpec:
    """
    Deny if any atom lies outside a z-window.
    """
    q = atoms().z_outside(zmin, zmax)
    return deny(q, tag)

def preset_flying_h(zmin: float = 3.0, zmax: float = 19.0, tag: str = "FLYING_H") -> RuleSpec:
    """
    Detects isolated H atoms flying in the vacuum.
    """
    q = atoms("H").z_outside(zmin, zmax).coord_eq(0)
    return deny(q, tag)

def preset_water(tag: str = "H2O") -> RuleSpec:
    """
    Detects H2O: O atom with exactly 2 H neighbors.
    """
    q = atoms("O").neighbors_eq("H", 2).coord_eq(2)
    return deny(q, tag)

def preset_isolated_carbon(tag: str = "C_STAR") -> RuleSpec:
    """
    Detects isolated C* on metal (no C or O neighbors).
    """
    q = atoms("C").neighbors_eq("C", 0).neighbors_eq("O", 0)
    return deny(q, tag)

def preset_co_gas(tag: str = "CO") -> RuleSpec:
    """
    Detects isolated CO molecule (C-O bond, both coord 1).
    """
    q = atoms("C").neighbors_eq("O", 1).coord_eq(1)
    return deny(q, tag)

def preset_co2_gas(tag: str = "CO2") -> RuleSpec:
    """
    Detects isolated CO2 molecule (C with 2 O neighbors, total coord 2).
    """
    q = atoms("C").neighbors_eq("O", 2).coord_eq(2)
    return deny(q, tag)

def preset_methane(tag: str = "CH4") -> RuleSpec:
    """
    Detects isolated Methane (C with 4 H neighbors).
    """
    q = atoms("C").neighbors_eq("H", 4).coord_eq(4)
    return deny(q, tag)

def preset_isolated_any(tag: str = "ISOLATED") -> RuleSpec:
    """
    Detects any atom with zero bonded neighbors.
    """
    q = atoms().coord_eq(0)
    return deny(q, tag)

def preset_all_oo(tag: str = "O2_ALL") -> RuleSpec:
    """
    Forbids ANY O-O bond, whether in free O2 or as a peroxide/superoxide fragment.
    """
    q = atoms("O").neighbors_min("O", 1)
    return deny(q, tag)

def preset_all_hh(tag: str = "H2_ALL") -> RuleSpec:
    """
    Forbids ANY H-H bond, whether in free H2 or as a bound fragment.
    """
    q = atoms("H").neighbors_min("H", 1)
    return deny(q, tag)

def preset_all_h2o(tag: str = "H2O_ALL") -> RuleSpec:
    """
    Forbids ANY Oxygen with 2 or more Hydrogen neighbors (free H2O or bound fragment).
    """
    q = atoms("O").neighbors_min("H", 2)
    return deny(q, tag)

# Registry of standard presets for the BlacklistDetector wrapper
_DEFAULT_PRESETS = {
    "H2":       lambda: preset_free_diatomic("H", tag="H2"),
    "O2":       lambda: preset_free_diatomic("O", tag="O2"),
    "H2O":      preset_water,
    "H2_ALL":   preset_all_hh,
    "O2_ALL":   preset_all_oo,
    "H2O_ALL":  preset_all_h2o,
    "HERO":     lambda: preset_outside_z(3.0, 19.0),
    "FLYING_H": preset_flying_h,
    "C_STAR":   preset_isolated_carbon,
    "CO":       preset_co_gas,
    "CO2":      preset_co2_gas,
    "CH4":      preset_methane,
    "ISOLATED": preset_isolated_any,
}


class BlacklistDetector:
    """
    Compatibility wrapper for ConstraintEngine.
    Enables backward compatibility with Population.py.
    """
    def __init__(self, blacklist: Iterable[str] = ()):
        self._engine = ConstraintEngine()
        tags = set(s.upper() for s in blacklist)
        
        for tag, factory in _DEFAULT_PRESETS.items():
            if tag in tags:
                self._engine.add(factory())
        
        self.remove = False

    # ---------------------------------------------------------
    # Public API 
    # ---------------------------------------------------------
    def contains(self, struct, remove: Optional[bool] = None) -> Tuple[bool, Optional[str]]:
        """
        Returns:
            (True, tag) if structure violates constraints
            (False, None) otherwise
        """
        res, tag = self._engine.contains(struct, remove_override=remove)
        if res:
            return True, tag
            
        custom_rules = getattr(self, '_custom_rules', None)
        if custom_rules:
            for ctag, func in custom_rules.items():
                try:
                    apm = struct.AtomPositionManager
                    if func(apm):
                        return True, ctag
                except Exception:
                    continue
                    
        return False, None

    def register(self, tag: str, func: Callable):
        """
        Legacy registration backward compatibility for custom functions.
        Accepts a function that takes AtomPositionManager as input and returns boolean.
        """
        if getattr(self, '_custom_rules', None) is None:
            self._custom_rules = {}
        self._custom_rules[tag] = func
