# Copyright (c) SHRY Development Team.
# Distributed under the terms of the MIT License.

"""Core operations, pattern generation, etc."""

# python modules
import collections
import functools
import itertools
import logging
import math
import multiprocessing
import random
import sys
import threading
from typing import OrderedDict

# python modules
import numpy as np
import spglib
import sympy
import tqdm
from pymatgen.analysis.ewald import EwaldSummation
from pymatgen.core.composition import Composition
from pymatgen.io.cif import CifBlock, CifWriter
from pymatgen.symmetry.analyzer import SpacegroupOperations
from pymatgen.symmetry.structure import SymmetrizedStructure
from pymatgen.util.coord import find_in_coord_list_pbc
from pymatgen.util.string import transformation_to_string
from scipy.special import comb
from sympy.utilities.iterables import multiset_permutations
from tabulate import tabulate
from pymatgen.core import Structure

# shry modules
from . import const
from .patches import _apply_pymatgen_patches
from .cif_io import _structure_to_mcif_string
from .cif_io import _structure_to_mcif_string
from .cif_io import parse_cif_string_to_structure

# shry version control
try:
    from ._version import version as shry_version
except (ModuleNotFoundError, ImportError):
    shry_version = "unknown"

np.seterr(all="raise")
np.set_printoptions(linewidth=const.LINEWIDTH, threshold=16)
np.set_printoptions(linewidth=1000, threshold=sys.maxsize)

spglib.OLD_ERROR_HANDLING = False

_apply_pymatgen_patches()


# Global variables for multiprocessing workers
_worker_pm = None


def _init_worker_invarless(perms, bits, bit_perm, nix):
    """Initialize worker process with PatternMaker data (invarless search)."""
    global _worker_pm
    _worker_pm = {
        "perms": perms,
        "bits": bits,
        "bit_perm": bit_perm,
        "nix": nix,
    }


def _worker_search_invarless(args):
    """Worker function for parallel invarless subtree search."""
    import time

    initial_state, stop = args
    pattern, aut, pbs = initial_state

    start_time = time.time()
    perms = _worker_pm["perms"]
    bit_perm = _worker_pm["bit_perm"]
    nix = _worker_pm["nix"]

    results = []
    stack = [(pattern, aut, pbs)]
    iterations = 0

    while stack:
        iterations += 1
        pattern, aut, pbs = stack.pop()
        if pattern.size == stop:
            results.append((aut, pattern))
            continue

        # Tree is expanded by adding un-added sites
        leaf_mask = np.ones(nix, dtype="bool")
        leaf_mask[pattern] = False
        leaf_array = np.flatnonzero(leaf_mask)

        # Discard symmetry duplicates from the remaining leaves
        if aut.size > 1 and leaf_mask.sum() > 1:
            leaf_reps = perms[np.ix_(aut, leaf_array)].min(axis=0)
            leaf_indices = np.searchsorted(leaf_array, leaf_reps)
            uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
            uniq_mask[leaf_indices] = True
        else:
            uniq_mask = np.ones(leaf_array.shape, dtype="bool")

        # Insertion location
        loci = pattern.searchsorted(leaf_array)

        for i in np.flatnonzero(uniq_mask):
            x = leaf_array[i]
            j = loci[i]
            _pbs = bit_perm[:, x] + pbs

            _i = np.concatenate((pattern[:j], [x], pattern[j:]))
            _aut = np.flatnonzero(_pbs == _pbs[-1])

            # Compute canonical parent.
            m = np.argmin(_pbs)
            can_pattern = perms[m, _i]
            discard_i = np.where(can_pattern == can_pattern.max())[0]

            # If the new site is discarded then tree parent == canonical parent.
            if j == discard_i:
                stack.append((_i, _aut, _pbs))
            # Check if tree parent is related to canonical parent.
            elif _i[discard_i] in perms[_aut, x]:
                stack.append((_i, _aut, _pbs))

    elapsed = time.time() - start_time
    import os

    pid = os.getpid()
    # Log subtree processing info
    import logging

    logging.info(f"  Worker {pid}: {len(results)} patterns, {iterations} iterations, {elapsed:.2f}s")

    return results


def _init_worker_invar(perms, bits, bit_perm, nix, invar, t_kind, shuffle):
    """Initialize worker process with PatternMaker data (invar search)."""
    global _worker_pm
    _worker_pm = {
        "perms": perms,
        "bits": bits,
        "bit_perm": bit_perm,
        "nix": nix,
        "invar": invar,
        "t_kind": t_kind,
        "shuffle": shuffle,
    }


def _worker_get_subobj_ts_sum(pattern, leaf_array, subobj_ts, invar):
    """Calculate subobject Ts for sum t_kind."""
    new_rows = invar[np.ix_(leaf_array, pattern)]
    new_rows_sums = new_rows.sum(axis=1, keepdims=True)
    part_new_row_sums = subobj_ts + new_rows
    return np.concatenate((part_new_row_sums, new_rows_sums), axis=1)


def _worker_get_subobj_ts_det(pattern, leaf_array, subobj_ts, invar):
    """Calculate subobject Ts for det t_kind."""
    subobj_ts_list = []
    for leaf in leaf_array:
        new_pattern = np.append(pattern, leaf)
        pattern_invar = invar[np.ix_(new_pattern, new_pattern)]
        det = np.abs(np.linalg.det(np.atleast_2d(pattern_invar)))
        subobj_ts_list.append(np.array([det]))
    return np.array(subobj_ts_list)


def _worker_search_invar(args):
    """Worker function for parallel invar subtree search."""
    initial_state, stop = args

    perms = _worker_pm["perms"]
    bit_perm = _worker_pm["bit_perm"]
    nix = _worker_pm["nix"]
    invar = _worker_pm["invar"]
    t_kind = _worker_pm["t_kind"]
    shuffle = _worker_pm["shuffle"]

    # Select appropriate get_subobj_ts function
    if t_kind in ("sum", "plsum"):
        get_subobj_ts = lambda p, l, s: _worker_get_subobj_ts_sum(p, l, s, invar)
    elif t_kind == "det":
        get_subobj_ts = lambda p, l, s: _worker_get_subobj_ts_det(p, l, s, invar)
    else:
        raise RuntimeError(f'Unrecognized T kind "{t_kind}"')

    # Select appropriate mask functions based on t_kind
    if t_kind == "sum":
        get_not_reject_mask = lambda delta_t: ~(delta_t < 0).any(axis=1)
        get_accept_mask = lambda delta_t: (delta_t > 0).all(axis=1)
        get_mins = lambda subobj_ts: np.flatnonzero(subobj_ts == subobj_ts.min())
    else:  # plsum or det (float comparison)
        get_not_reject_mask = lambda delta_t: ~((delta_t < 0.0) & ~np.isclose(delta_t, 0.0)).any(axis=1)
        get_accept_mask = lambda delta_t: ((delta_t > 0.0) & ~np.isclose(delta_t, 0.0)).all(axis=1)
        get_mins = lambda subobj_ts: np.flatnonzero(np.isclose(subobj_ts, subobj_ts.min()))

    results = []
    stack = [initial_state]

    while stack:
        subobj_ts, pattern, aut, pbs = stack.pop()
        if pattern.size == stop:
            results.append((aut, pattern))
            continue

        # Invert index
        leaf_mask = np.ones(nix, dtype="bool")
        leaf_mask[pattern] = False
        leaf_array = np.flatnonzero(leaf_mask)

        # Calculate subobject Ts for all leaves
        leaf_subobj_ts = get_subobj_ts(pattern, leaf_array, subobj_ts)

        # Reject all leaves where any T is smaller than the new row's T
        delta_t = leaf_subobj_ts[:, :-1] - leaf_subobj_ts[:, -1:]
        not_reject_mask = get_not_reject_mask(delta_t)
        if not not_reject_mask.any():
            continue

        # Discard symmetry duplicates from the remaining leaves
        if aut.size > 1 and not_reject_mask.sum() > 1:
            not_reject_leaf = leaf_array[not_reject_mask]
            leaf_reps = perms[np.ix_(aut, not_reject_leaf)].min(axis=0)
            leaf_indices = leaf_array.searchsorted(leaf_reps)
            uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
            uniq_mask[leaf_indices] = True
        else:
            uniq_mask = not_reject_mask

        # Insertion location
        loci = pattern.searchsorted(leaf_array)

        accept_mask = get_accept_mask(delta_t)
        accept_mask &= uniq_mask
        accepts = leaf_array[accept_mask]

        # Iterate
        if shuffle:
            indices = np.flatnonzero(uniq_mask)
            np.random.shuffle(indices)
        else:
            indices = np.flatnonzero(uniq_mask)

        for i in indices:
            x = leaf_array[i]
            j = loci[i]
            _pbs = bit_perm[:, x] + pbs

            _subobj_ts = leaf_subobj_ts[i].copy()
            _subobj_ts[j:] = np.concatenate((_subobj_ts[-1:], _subobj_ts[j:-1]))

            _i = np.concatenate((pattern[:j], [x], pattern[j:]))
            _aut = np.flatnonzero(_pbs == _pbs[-1])

            if x in accepts:
                stack.append((_subobj_ts, _i, _aut, _pbs))
                continue

            # Compute canonical parent.
            ts_min_i = get_mins(_subobj_ts)
            _sub = _i[ts_min_i]
            m = np.argmin(_pbs)
            can_pattern = perms[m, _sub]
            discard_i = ts_min_i[can_pattern == can_pattern.max()]

            # If the new site is discarded then tree parent == canonical parent.
            if j == discard_i:
                stack.append((_subobj_ts, _i, _aut, _pbs))
            # Check if tree parent is related to canonical parent.
            elif _i[discard_i] in perms[_aut, x]:
                stack.append((_subobj_ts, _i, _aut, _pbs))

    return results


def _structure_to_spglib_cell(structure):
    """
    Convert pymatgen Structure to spglib cell tuple.

    Args:
        structure: pymatgen Structure object

    Returns:
        tuple: (lattice, positions, numbers) for spglib
    """
    lattice = structure.lattice.matrix
    positions = np.array([site.frac_coords for site in structure])

    # Get atomic numbers
    # For ordered sites, use Z directly
    # For disordered sites, use the first element (will be handled during enumeration)
    numbers = []
    for site in structure:
        if site.is_ordered:
            numbers.append(site.specie.Z)
        else:
            # For disordered sites, use the first element
            numbers.append(list(site.species.keys())[0].Z)
    numbers = np.array(numbers, dtype=int)

    return (lattice, positions, numbers)


def _structure_requests_msg(structure) -> bool:
    """Return True if the Structure should use magnetic symmetry (MSG)."""
    props = getattr(structure, "properties", None) or {}
    mode = props.get("_shry_symmetry_mode")
    if mode is not None:
        return mode == "msg"
    # Fallback (option a): if magnetic moments are present, use MSG.
    try:
        if "magmom" in getattr(structure, "site_properties", {}):
            return True
        for site in structure:
            mm = site.properties.get("magmom")
            if mm is not None and not np.allclose(np.array(mm, dtype=float), 0.0):
                return True
    except Exception:
        pass
    return False


class _StringWriter:
    """Minimal writer wrapper compatible with ScriptHelper.write (needs __str__)."""

    def __init__(self, text: str):
        self._text = text

    def __str__(self) -> str:
        return self._text


def _structure_to_spglib_msg_cell(structure):
    """Convert pymatgen Structure to spglib MsgCell tuple.

    spglib>=2 uses MsgCell = (lattice, positions, numbers, magmoms).
    magmoms is passed as a (N, 3) array (crystal axis components).

    Notes:
        - For disordered sites, we choose a representative atomic number
        (highest occupancy; tie-broken deterministically).
        - Magnetic moments are taken from site.properties['magmom'] if present;
        otherwise from site.properties['_atom_site_moment_by_label'] (first nonzero);
        otherwise zeros.
    """

    lattice = structure.lattice.matrix
    positions = np.array([site.frac_coords for site in structure])

    numbers = []
    magmoms = []

    for site in structure:
        # Representative atomic number
        if site.is_ordered:
            rep_specie = site.specie
        else:
            # Choose the majority-occupancy specie deterministically.
            items = list(site.species.items())
            # items: [(Element/Specie, occ), ...]
            items.sort(key=lambda t: (-float(t[1]), str(getattr(t[0], "symbol", t[0]))))
            rep_specie = items[0][0]
        try:
            numbers.append(int(rep_specie.Z))
        except Exception:
            numbers.append(int(list(site.species.keys())[0].Z))

        # Representative magnetic moment
        mm = site.properties.get("magmom")
        if mm is None:
            mom_map = site.properties.get("_atom_site_moment_by_label") or {}
            mm = next(
                (np.array(v, dtype=float) for v in mom_map.values() if not np.allclose(np.array(v, dtype=float), 0.0)),
                np.zeros(3, dtype=float),
            )
        mm = np.array(mm, dtype=float)
        if mm.shape == ():  # scalar -> treat as x-component
            mm = np.array([float(mm), 0.0, 0.0], dtype=float)
        elif mm.shape != (3,):
            mm = np.zeros(3, dtype=float)
        magmoms.append(mm)

    numbers = np.array(numbers, dtype=int)
    magmoms = np.array(magmoms, dtype=float)
    return (lattice, positions, numbers, magmoms)


def _get_magnetic_symmetry_operations_from_spglib(
    structure,
    symprec=1e-5,
    angle_tolerance=-1.0,
):
    """Get magnetic symmetry operations (MSG) from spglib.

    Returns a list of pymatgen SymmOp for coordinate mapping.
    Time-reversal flags are handled internally by spglib and are not
    encoded in SymmOp.
    """
    from pymatgen.core.operations import SymmOp

    cell = _structure_to_spglib_msg_cell(structure)
    symmetry = spglib.get_magnetic_symmetry(cell, symprec=symprec, angle_tolerance=angle_tolerance)
    if symmetry is None:
        raise RuntimeError("Couldn't find magnetic symmetry.")

    symmops = []
    for rotation, translation in zip(symmetry["rotations"], symmetry["translations"], strict=False):
        affine_matrix = np.eye(4)
        affine_matrix[:3, :3] = rotation
        affine_matrix[:3, 3] = translation
        symmops.append(SymmOp(affine_matrix))
    return symmops


def _get_symmetry_operations_from_spglib(structure, symprec=1e-5, angle_tolerance=-1.0):
    """
    Get symmetry operations from spglib.

    Args:
        structure: pymatgen Structure object
        symprec: Symmetry precision
        angle_tolerance: Angle tolerance

    Returns:
        list: List of SymmOp objects
    """
    from pymatgen.core.operations import SymmOp

    cell = _structure_to_spglib_cell(structure)
    symmetry = spglib.get_symmetry(cell, symprec=symprec, angle_tolerance=angle_tolerance)

    if symmetry is None:
        raise RuntimeError("Couldn't find symmetry.")

    # Convert spglib symmetry operations to pymatgen SymmOp objects
    symmops = []
    for rotation, translation in zip(symmetry["rotations"], symmetry["translations"]):
        # Create affine transformation matrix
        affine_matrix = np.eye(4)
        affine_matrix[:3, :3] = rotation
        affine_matrix[:3, 3] = translation
        symmops.append(SymmOp(affine_matrix))

    return symmops


class _PatchedSymmetrizedStructure(SymmetrizedStructure):
    """
    Fixed site_properties display
    """

    def __str__(self):
        outs = [
            "SymmetrizedStructure",
            f"Full Formula ({self.composition.formula})",
            f"Reduced Formula: {self.composition.reduced_formula}",
        ]

        def to_s(x):
            return f"{x:0.6f}"

        outs.append("abc   : " + " ".join([to_s(i).rjust(10) for i in self.lattice.abc]))
        outs.append("angles: " + " ".join([to_s(i).rjust(10) for i in self.lattice.angles]))
        if self._charge:
            if self._charge >= 0:
                outs.append(f"Overall Charge: +{self.charge}")
            else:
                outs.append(f"Overall Charge: -{self._charge}")
        outs.append(f"Sites ({len(self)})")
        data = []
        props = self.site_properties  # This should be updated!
        keys = sorted(props.keys())
        for i, sites in enumerate(self.equivalent_sites):
            site = sites[0]
            row = [str(i), site.species_string]
            row.extend([to_s(j) for j in site.frac_coords])
            row.append(self.wyckoff_symbols[i])
            for k in keys:
                row.append(site.properties[k])  # This line
            data.append(row)
        outs.append(
            tabulate(
                data,
                headers=["#", "SP", "a", "b", "c", "Wyckoff"] + keys,
            )
        )
        return "\n".join(outs)


def _get_symmetrized_structure_from_spglib(structure, symprec=1e-5, angle_tolerance=-1.0):
    """
    Get symmetrized structure using spglib directly.

    Args:
        structure: pymatgen Structure object
        symprec: Symmetry precision
        angle_tolerance: Angle tolerance

    Returns:
        PatchedSymmetrizedStructure object
    """
    cell = _structure_to_spglib_cell(structure)
    dataset = spglib.get_symmetry_dataset(cell, symprec=symprec, angle_tolerance=angle_tolerance)

    if dataset is None:
        raise RuntimeError("Couldn't find symmetry.")

    # Get symmetry operations
    symmops = _get_symmetry_operations_from_spglib(structure, symprec, angle_tolerance)

    # Create SpacegroupOperations object
    sg = SpacegroupOperations(
        dataset.international,
        dataset.number,
        symmops,
    )

    return _PatchedSymmetrizedStructure(structure, sg, dataset.equivalent_atoms, dataset.wyckoffs)


class _AltCifBlock(CifBlock):
    """
    CifBlock but optimized for when there are
    many, slightly different blocks:

    - Pre-format fields on init, value update, and from_str
    - Cache unchanged strings
    - Loop to list instead of to string
    """

    def _loop_to_list(self, loop):
        s = ["loop_"] + [f" {l}" for l in loop]

        # Join the fields
        lines = "\n".join("  ".join(x) for x in zip(*[self.data[k] for k in loop])).replace(";", "\n")

        # Check for maximum lengths
        for line in lines.split("\n"):
            if len(line) < self.max_len:
                s.append(line)
            else:
                sublines = [line[i : i + self.max_len] for i in range(0, len(line), self.max_len)]
                s.extend(sublines)
        return s

    def __str__(self):
        """
        Returns the cif string for the data block
        """
        s = [f"data_{self.header}"]
        keys = self.data.keys()
        written = []
        for k in keys:
            if k in written:
                continue
            # can this be optimized (won't affect much though)
            for l in self.loops:
                # search for a corresponding loop
                if k in l:
                    loop_id = "loop_\n " + "\n ".join(l)
                    if self.string_cache[loop_id] is not None:
                        s.extend(self.string_cache[loop_id])
                    else:
                        looplist = self._loop_to_list(l)
                        s.extend(looplist)
                        self.string_cache[loop_id] = looplist
                    written.extend(l)
                    break
            if k not in written:
                # k didn't belong to a loop
                if self.string_cache[k] is not None:
                    s.extend(self.string_cache[k])
                else:
                    v = self.data[k]
                    if len(k) + len(v) + 3 < self.max_len:
                        s.append(f"{k}   {v}")
                    else:
                        s.extend([k, v])
        return "\n".join(s)

    def __setitem__(self, key, value):
        for l in self.loops:
            if key in l:
                self.data[key] = [self._format_field(x) for x in value]
                loop_id = "loop_\n " + "\n ".join(l)
                self.string_cache[loop_id] = None
                return
        self.data[key] = self._format_field(value)
        self.string_cache[key] = None

    def __init__(self, data, loops, header):
        super().__init__(data, loops, header)
        # Upstream renamed maxlen -> max_len; assume upstream provides it.
        if hasattr(self, "max_len"):
            self.max_len = self.max_len
        # Pre-format everything.
        formatted = []
        self.string_cache = OrderedDict()
        for k in self.data.keys():
            if k in formatted:
                continue
            for l in self.loops:
                if k in l:
                    for _k in l:
                        self.data[_k] = list(map(self._format_field, self.data[_k]))
                    loop_id = "loop_\n " + "\n ".join(l)
                    self.string_cache[loop_id] = self._loop_to_list(l)
                    formatted.extend(l)
                    break
            if k not in formatted:
                self.data[k] = self._format_field(self.data[k])
                # Cache preparation
                v = self.data[k]
                s = []
                if len(k) + len(v) + 3 < self.max_len:
                    s.append(f"{k}   {v}")
                else:
                    s.extend([k, v])
                self.string_cache[k] = s


class NeedSupercellError(Exception):
    """
    Just so that this error will raise a distinguishable exception
    """

    ...


class TooBigError(Exception):
    """
    When the number of substitution is too large
    """

    ...


def _rec_asc(a, n, m, k, length):
    """
    Recursive part of the thing below
    """
    x = m
    while 2 * x <= n and k < length:
        a[k - 1] = x
        yield from _rec_asc(a, n - x, x, k + 1, length)
        x += 1
    a[k - 1] = n
    # Pad
    b = a[:k]
    b += [0] * (length - k)
    # Permute
    for z in multiset_permutations(b):
        yield z


@functools.lru_cache(None)
def _aR_array(n, length=None):
    """
    aR() cache.
    """
    a = [1] * n
    if length is None:
        length = n
    return np.array(list(_rec_asc(a, n, 1, 1, length)))


@functools.lru_cache(None)
def _multinomial_coeff(a):
    """
    Get multinomial coefficient of (sum(a), (a_1, a_2, ...))
    """
    return functools.reduce(
        lambda a, b: a * b,
        [comb(cx, x, exact=True) for cx, x in zip(np.cumsum(a), a)],
    )


class Substitutor:
    """
    Enumerate symmetrically unique ordered structures from a disordered structure.

    Since the PatternMaker instances are saved,
    it is advantageous to reuse the same Substitutor
    for symmetrically similar structures.

    This includes different subsitute concentrations
    from the same basic structure, among others.

    Attributes:
        disorder_groups (dict): Mapping of grouping key (orbit identifier) to
            tuples of disordered ``PeriodicSite`` objects that are substituted
            together.

    Args:
        structure (pymatgen.core.Structure or shry.LabeledStructure):
            Input structure. Both plain pymatgen Structure and SHRY LabeledStructure
            instances are supported.
        symprec (float): Symmetry precision.
        angle_tolerance (float): Angle tolerance for symmetry search.
        atol (float): Absolute tolerance used in internal floating-point comparisons.
        groupby (function): Function to group disordered sites.
            Defaults to lambda x: x.properties["_atom_site_label"], with the x loops
            over all PeriodicSites within the Structure.

            The proper function depends on what you want; do you want all sites of
            a certain species or crystallographic orbit to be subsituted together?

            Fallback to crystallographic orbit when failed.
        shuffle (bool): If True, randomize branch traversal order during pattern search.
        seed (int): Random seed used when ``shuffle=True``.
        no_dmat (bool): Whether or not to use distance matrix as
            permutation invariant (faster, default=True)
        t_kind (str): Invariant mode used by PatternMaker (e.g. ``sum``, ``plsum``, ``det``).
        cache (bool or None): By default(=None), cache patterns when involving either
            multiple orbits or multiple species, but otherwise don't.
            Caching allows "reuse" of previously generated patterns,
            at the cost of memory.
            Set to False if memory is limited,
            but note that pattern generation will be much slower.
        n_jobs (int): Number of worker processes for pattern search.
            ``1`` (default) disables process-based parallel search and is the
            safest option for script usage on spawn-based platforms.
            ``-1`` uses all available CPU cores.
    """

    __slots__ = (
        "_symprec",
        "_angle_tolerance",
        "_groupby",
        "_shuffle",
        "_seed",
        "_atol",
        "_symmops",
        "_pms",
        "_get_polya",
        "disorder_groups",
        "_group_dmat",
        "_group_perms",
        "_group_indices",
        "_group_bits",
        "_group_bit_perm",
        "_structure",
        "_sample",
        "_cache",
        "_no_dmat",
        "_t_kind",
        "_segmenter",
        "_charset",
        "_template_cifwriter",
        "_template_structure",
        "_rg",
        "_n_jobs",
    )

    def __init__(
        self,
        structure,
        symprec=const.DEFAULT_SYMPREC,
        angle_tolerance=const.DEFAULT_ANGLE_TOLERANCE,
        atol=const.DEFAULT_ATOL,
        groupby=None,
        shuffle=False,
        seed=const.DEFAULT_SEED,
        no_dmat=const.DEFAULT_NO_DMAT,
        t_kind=const.DEFAULT_T_KIND,
        cache=None,  # "True", "False", "None" (default)
        n_jobs=1,
    ):
        self._symprec = symprec
        self._angle_tolerance = angle_tolerance
        self._no_dmat = no_dmat
        self._t_kind = t_kind
        self._n_jobs = n_jobs
        # TODO: These two are consequential to self._structure, so should be property
        self._atol = atol
        self._shuffle = shuffle
        self._seed = seed
        if groupby is None:
            self._groupby = lambda x: x.properties["_atom_site_label"]
        else:
            self._groupby = groupby

        # random.Random instance.
        # The default RNG is used in many other modules,
        # so setting seed does not make sense there.
        rg = random.Random()
        rg.seed(self._seed)
        self._rg = rg

        self._cache = cache
        self._symmops = None
        self._pms = dict()
        self._get_polya = _get_polya

        self.disorder_groups = dict()
        self._group_dmat = dict()
        self._group_perms = dict()
        self._group_indices = dict()
        self._group_bits = dict()
        self._group_bit_perm = dict()
        # TODO: Decouple output-format specific attributes.
        self._template_cifwriter = None
        self._template_structure = None

        self.structure = structure

    @property
    def structure(self):
        """
        Input structure currently held by this Substitutor.

        Returns:
            pymatgen.core.Structure: Internal working copy of the input structure.
        """
        return self._structure

    @property
    def groupby(self):
        """
        Grouping function used to partition disorder sites into substitution orbits.

        Returns:
            Callable: Function that maps each PeriodicSite to a grouping key.
        """
        return self._groupby

    @groupby.setter
    def groupby(self, groupby):
        """Set grouping function and rebuild disorder grouping metadata.

        Args:
            groupby (Callable): Function mapping each PeriodicSite to a grouping key.
        """
        self._groupby = groupby
        # Reload division part
        self.structure = self.structure

    @structure.setter
    def structure(self, structure):
        """Set input structure and rebuild all symmetry/disorder derived caches.

        Args:
            structure (pymatgen.core.Structure): New input structure.

        Raises:
            RuntimeError: If symmetry cannot be determined or occupancy is invalid.
            NeedSupercellError: If site multiplicities cannot fit integer composition.
        """
        logging.info("\nSetting Substitutor with Structure")
        logging.info(f"{structure}")

        # TODO: Avoid these kinds of manual invocations.
        self.disorder_groups.clear()
        self._group_dmat.clear()
        self._group_perms.clear()
        self._group_indices.clear()
        self._group_bits.clear()
        self._group_bit_perm.clear()
        self._template_cifwriter = None
        self._template_structure = None

        # Read.
        self._structure = structure.copy()

        # Get spglib cell and symmetry dataset (SG or MSG)
        use_msg = _structure_requests_msg(self._structure)
        logging.info("Using magnetic symmetry (MSG)" if use_msg else "Using crystallographic symmetry (SG)")
        if use_msg:
            cell = _structure_to_spglib_msg_cell(self._structure)
            dataset = spglib.get_magnetic_symmetry_dataset(
                cell,
                symprec=self._symprec,
                angle_tolerance=self._angle_tolerance,
            )
        else:
            cell = _structure_to_spglib_cell(self._structure)
            dataset = spglib.get_symmetry_dataset(cell, symprec=self._symprec, angle_tolerance=self._angle_tolerance)

        if dataset is None:
            raise RuntimeError("Couldn't find symmetry.")

        # Get symmetry operations
        try:
            if use_msg:
                self._symmops = _get_magnetic_symmetry_operations_from_spglib(
                    self._structure, symprec=self._symprec, angle_tolerance=self._angle_tolerance
                )
            else:
                self._symmops = _get_symmetry_operations_from_spglib(
                    self._structure, symprec=self._symprec, angle_tolerance=self._angle_tolerance
                )
        except Exception as exc:
            raise RuntimeError("Couldn't find symmetry.") from exc

        if use_msg:
            uni_number = getattr(dataset, "uni_number", None)
            msg_type = getattr(dataset, "msg_type", None)
            logging.info(f"Magnetic space group: uni_number={uni_number}, msg_type={msg_type}")
        else:
            logging.info(f"Space group: {dataset.hall} ({dataset.number})")
        logging.info(f"Total {len(self._symmops)} symmetry operations")

        # Get symmetrized structure for logging (SG only)
        if not use_msg:
            symm_struct = _get_symmetrized_structure_from_spglib(
                self._structure, symprec=self._symprec, angle_tolerance=self._angle_tolerance
            )
            logging.info(symm_struct)

        equivalent_atoms = dataset.equivalent_atoms

        # Identify and label disorder sites.
        # Index is used for later recoloring the sites.
        logging.info("\nProcessing disorder sites.")
        disorder_sites = []
        for i, site in enumerate(self._structure):
            site.properties["index"] = i
            site.properties["equivalent_atoms"] = equivalent_atoms[i]
            if not site.is_ordered:
                disorder_sites.append(site)
                # Ad hoc fix: if occupancy is less than 1, stop.
                # TODO: Automatic vacancy handling
                if not np.isclose(site.species.num_atoms, 1.0, atol=self._atol):
                    logging.warning(f"The occupancy of the site {site.species} is {site.species.num_atoms}.")
                    logging.warning(f"This should be 1 within the torelance, atol={self._atol}.")
                    logging.warning("If you want to consider vacancy sites, please add pseudo atoms.")
                    raise RuntimeError("The sum of number of occupancies is not 1.")
        if not disorder_sites:
            logging.warning("No disorder sites found within the Structure.")

        try:
            disorder_sites.sort(key=self._groupby)
        except KeyError:
            # Fallback to crystallographic orbit.
            self._groupby = lambda x: x.properties["equivalent_atoms"]
            disorder_sites.sort(key=self._groupby)

        for orbit, sites in itertools.groupby(disorder_sites, key=self._groupby):
            # Can it fit?
            sites = tuple(sites)
            composition = sites[0].species.to_int_dict()
            integer_formula = "".join(e + str(a) for e, a in composition.items())
            formula_unit_sum = sum(composition.values())
            if len(sites) % formula_unit_sum:
                raise NeedSupercellError(
                    (
                        f"Can't fit {integer_formula} "
                        f"within {len(sites)} sites "
                        f"(enlarge by {formula_unit_sum / len(sites):.4f}x). "
                        f"If the integer composition ({integer_formula}) is not what you expected, "
                        "please try to increase the precision of the occupancy "
                        "e.g., 1/3 = 0.3333 -> 1/3 = 0.3333333333333."
                    )
                )
            self.disorder_groups[orbit] = sites

            # DMAT
            indices = [x.properties["index"] for x in sites]
            self._group_indices[orbit] = indices
            group_dmat = self._structure.distance_matrix[np.ix_(indices, indices)]
            self._group_dmat[orbit] = self.ordinalize(group_dmat, atol=self._atol)

            # PERM
            coords = [x.frac_coords for x in sites]
            group_perms = []

            for symmop in tqdm.tqdm(
                self._symmops,
                desc="Building symmetry list",
                **const.TQDM_CONF,
                disable=const.DISABLE_PROGRESSBAR,
            ):
                o_coords = symmop.operate_multi(coords)
                atol = const.FIND_COORDS_LIST_ATOL_INIT
                step = const.FIND_COORDS_LIST_ATOL_MUL_STEP_INIT
                is_up = True

                for _ in range(const.FIND_COORDS_LIST_ATTEMPTS):
                    matches = [find_in_coord_list_pbc(coords, coord, atol=atol) for coord in o_coords]
                    # Sometimes the operated coordinates fell
                    # outside of the symmetry tolerance.
                    # On the other hand, if the tolerance is too low,
                    # one site can be matched with multiple (wrong) sites.
                    #
                    # Robust: Attempt to guess proper tolerance
                    # when one of the above case was detected.
                    # Increment/decrement tolerance with fixed steps,
                    # then binary search if the tolerance becomes too low
                    # after initially too high, and vice versa.
                    # Raise exception if can't find proper tolerance after 10 attempts
                    match_lengths = [len(x) for x in matches]
                    if all(x == 1 for x in match_lengths):
                        break
                    elif any(not x for x in match_lengths):
                        # Too strict
                        if not is_up:
                            step /= 2
                            is_up = True
                        atol *= step
                    elif any(x > 1 for x in match_lengths):
                        # Too loose
                        if is_up:
                            step /= 2
                            is_up = False
                        atol /= step
                else:
                    raise RuntimeError("Failed to build symmetry list.")

                indices = [x[0] for x in matches]
                group_perms.append(indices)
            group_perms = np.array(group_perms)
            self._group_perms[orbit] = group_perms
            row_members = np.sort(group_perms, axis=1)
            if not (row_members == row_members[0]).all():
                raise RuntimeError("Invalid permutation(s) generated.")

            bits = [2**i for i in range(group_perms.shape[1])]
            self._group_bits[orbit] = bits

            try:
                bit_perm = np.array(
                    [[2 ** int(i) for i in row] for row in group_perms],
                    dtype=int,
                )
            except OverflowError:
                bit_perm = np.array(
                    [[2 ** int(i) for i in row] for row in group_perms],
                    dtype=object,
                )
            self._group_bit_perm[orbit] = bit_perm

        # If not explicitly set, turn on caching if there are multiple orbits/colors,
        # and turn off otherwise
        if self._cache is None:
            da = list(self._disorder_amounts().values())
            if not da:  # no disorder
                pass
            else:
                if len(da) > 1 or len(da[0]) > 2:
                    self._cache = True
                else:
                    self._cache = False

        # Letter-related functions
        self._segmenter = self._disorder_amounts().values()
        n_segments = sum(len(x) for x in self._segmenter)
        self._charset = [chr(97 + i) for i in range(n_segments)]

    @staticmethod
    def ordinalize(array, atol=const.DEFAULT_ATOL):
        """
        Ordinalize array elements to the specified absolute tolerance.

        Args:
            array (np.array): The array to be ordinalized.
            atol (float): Absolute tolerance. Defaults to 1e-5.

        Returns:
            np.array: The ordinalized array.
        """
        flat = array.flatten()
        dflat = []
        points = [flat[0]]
        for element in flat:
            equal = 0
            for i, point in enumerate(points):
                dist = np.abs(point - element)
                if dist < atol:
                    equal = i
                    break
            else:
                equal = len(points)
                points.append(element)
                dflat.append(equal)
                continue
            dflat.append(equal)

        return np.array(dflat).reshape(array.shape)

    @property
    def pattern_makers(self):
        """
        Access cached PatternMaker instances.

        Returns:
            dict: Mapping from PatternMaker label to PatternMaker instance.
        """
        return self._pms

    @pattern_makers.setter
    def pattern_makers(self, pms):
        """Replace cached PatternMaker instances.

        Args:
            pms (dict): Mapping from PatternMaker label to PatternMaker instance.
        """
        self._pms = pms

    def _sorted_compositions(self):
        """
        Integer amount of sites substituted to meet desired specification.
        Ordered to minimize the size of the substituted part.
        """
        orbit_compositions = dict()
        for orbit, sites in self.disorder_groups.items():
            ref = sites[0]
            composition = dict(
                sorted(
                    ref.species.inted_composition.items(),
                    key=lambda x: x[1],
                    reverse=True,
                )
            )
            fu = int(sum(composition.values()))
            # Note: __init__ checks makes sure this has no remainder.
            fus = len(sites) // fu

            orbit_compositions[orbit] = {e: fus * int(a) for e, a in composition.items()}
        return orbit_compositions

    def _disorder_elements(self):
        return {orbit: tuple(x.keys()) for orbit, x in self._sorted_compositions().items()}

    def _disorder_amounts(self):
        return {orbit: tuple(x.values()) for orbit, x in self._sorted_compositions().items()}

    def make_patterns(self):
        """
        Generate raw substitution patterns for all disordered sites.

        Returns:
            Iterator[tuple[np.ndarray, list[np.ndarray]]]:
                Pairs of ``(aut, pattern)`` where ``aut`` is the automorphism index
                array and ``pattern`` is the orbit-wise substitution index pattern.
        """

        def rscum(iterable):
            """Cumulative sum of an iterable, but skip the last element, and reverse."""
            cum = 0
            cums = []
            for e in iterable[:-1]:
                cum += e
                cums.append(cum)
            for cum in cums[::-1]:
                yield cum

        def nocache_get_pm(subperm, dmat):
            pm = _PatternMaker(
                subperm,
                invar=dmat,
                enumerator_collection=self._get_polya,
                t_kind=self._t_kind,
                shuffle=self._shuffle,
                rg=self._rg,
                n_jobs=self._n_jobs,
            )
            row_map = pm.get_row_map()
            index_map = pm.get_index_map()
            return pm, row_map, index_map

        def cached_get_pm(subperm, dmat):
            label = _PatternMaker.get_label(subperm)
            if label in self._pms:
                pm = self._pms[label]
                row_map, index_map = pm.update_index(subperm)
            else:
                pm = _PatternMaker(
                    subperm,
                    invar=dmat,
                    enumerator_collection=self._get_polya,
                    t_kind=self._t_kind,
                    cache=self._cache,
                    shuffle=self._shuffle,
                    rg=self._rg,
                    n_jobs=self._n_jobs,
                )
                row_map = pm.get_row_map()
                index_map = pm.get_index_map()
                self._pms[label] = pm
            return pm, row_map, index_map

        def maker_recurse_unit(aut, pattern, orbit, amount):
            """PatternMaker aut/pattern generation recursion unit."""
            group_perms = self._group_perms[orbit]
            group_dmat = self._group_dmat[orbit]

            # TODO: Do pattern "plus" at the end.
            subpattern = pattern[-1]
            subperm = group_perms[np.ix_(aut, subpattern)]
            if not self._no_dmat:
                dmat = group_dmat[np.ix_(subpattern, subpattern)]
            else:
                dmat = None

            pm, row_map, index_map = get_pm(subperm, dmat)
            for _aut, _subpattern in pm.ap(amount):
                _aut = np.sort(row_map[_aut])
                _subpattern = index_map[_subpattern]
                yield [aut[_aut], pattern + [_subpattern]]

        def maker_recurse_c(aut, pattern, orbit, chain):
            if len(chain) > 0:
                amount = chain.pop()
                for aut, pattern in maker_recurse_unit(aut, pattern, orbit, amount):
                    _chain = chain.copy()
                    yield from maker_recurse_c(aut, pattern, orbit, _chain)
            else:
                yield aut, pattern

        def maker_recurse_o(aut, pattern, ochain):
            if len(ochain) > 0:
                orbit, sites = ochain.pop()

                chain = list(rscum(self._disorder_amounts()[orbit][::-1]))[::-1]
                indices = np.arange(len(sites))

                for aut, pattern in maker_recurse_c(aut, pattern + [indices], orbit, chain):
                    # TODO: something cheaper?
                    _ochain = ochain.copy()
                    yield from maker_recurse_o(aut, pattern, _ochain)
            else:
                yield aut, pattern

        logging.info("Making patterns.")
        # Safety kill
        count = self.count()
        if count > const.MAX_IRREDUCIBLE:
            raise TooBigError(f"({count} irreducible expected)")

        # To make the branches shallow.
        if self._cache:
            get_pm = cached_get_pm
        else:
            get_pm = nocache_get_pm

        for aut, pattern in maker_recurse_o(
            np.arange(len(self._symmops)),
            [],
            # TODO: Temporary fix.
            list(self.disorder_groups.items())[::-1],
        ):
            yield (aut, pattern)

    def total_count(self):
        """
        Return the total number of combinatorial substitutions (without symmetry reduction).

        Returns:
            int: Total number of possible substitutions.
        """
        ocount = (_multinomial_coeff(x) for x in self._disorder_amounts().values())
        return functools.reduce(lambda x, y: x * y, ocount, 1)

    def count(self):
        """
        Return the number of symmetry-inequivalent substitution patterns.

        Returns:
            int: Number of unique patterns after symmetry reduction.
        """
        logging.info(f"\nCounting unique patterns for {self.structure.formula}")

        if len(self._symmops):
            enumerator = self._get_polya(
                # orbit_symmetries,
                list(self._group_perms.values()),
                len(self._symmops),
            )
        else:  # No-permutations-supplied
            enumerator = self._get_polya([])
        count = enumerator.count(tuple(self._disorder_amounts().values()))
        return count

    def quantities(self, q, symprec=None):
        """
        Generate selected quantities for each symmetry-inequivalent configuration.

        Args:
            q (list[str]): Keys to include in each yielded packet.
                Valid keys are ``"cifwriter"``, ``"weight"``, ``"letter"``,
                ``"ewald"``, and ``"structure"``.
            symprec (float | None): Optional symmetry precision used for CIF/Ewald
                generation in helper methods.

        Yields:
            dict: Dictionary containing requested quantities for one configuration.
        """
        is_c = "cifwriter" in q
        is_s = "structure" in q
        is_w = "weight" in q
        is_l = "letter" in q
        is_e = "ewald" in q

        packet = collections.defaultdict(lambda: None)
        for a, p in self.make_patterns():
            packet.clear()
            if is_c:
                packet["cifwriter"] = self._get_cifwriter(p, symprec)
            if is_s:
                packet["structure"] = self._get_structure(p)
            if is_e:
                packet["ewald"] = self._get_ewald(p, symprec)
            if is_w:
                packet["weight"] = self._get_weight(a)
            if is_l:
                packet["letter"] = self._get_letters(p)
            # filter on packet!
            # naive: if filter(packet["ewald"]): yield else skip
            yield packet

    def letters(self):
        """
        Generate letter-encoded substitution patterns.

        Yields:
            str: Compact letter string representing one substitution pattern.
        """
        for _, p in self.make_patterns():
            yield self._get_letters(p)

    def weights(self):
        """
        Generate symmetry weights for all unique patterns.

        Yields:
            int: Weight of one pattern (group-size divided by automorphism size).
        """
        for a, _ in self.make_patterns():
            yield self._get_weight(a)

    def cifwriters(self, symprec=None):
        """
        Generate CIF writers (or magCIF string writers in MSG mode).

        Args:
            symprec (float | None): Optional symmetry precision used when building
                symmetry-reduced CIF blocks.

        Yields:
            CifWriter | _StringWriter: Writer object for each configuration.
        """
        for _, p in self.make_patterns():
            yield self._get_cifwriter(p, symprec)

    def structure_writers(self, symprec=None):
        """
        Generate ordered pymatgen structures for all unique patterns.

        Args:
            symprec (float | None): Unused. Kept for API compatibility with other
                writer generators.

        Yields:
            pymatgen.core.Structure: Ordered structure for each configuration.
        """
        # This one does not need symprec.
        # Just to keep the signature the same.
        del symprec
        for _, p in self.make_patterns():
            yield self._get_structure(p)

    def ase_atoms_writers(self, symprec=None):
        """Generate ASE Atoms objects for all unique patterns.

        Args:
            symprec (float | None): Unused. Kept for API compatibility.

        Yields:
            ase.Atoms: ASE representation of each ordered configuration.
        """
        from ase import Atoms

        def _from_pymatgen_struct_to_ase_atoms(structure: Structure) -> Atoms:
            return Atoms(
                symbols=[specie.symbol for specie in structure.species],
                positions=structure.cart_coords,
                cell=structure.lattice.matrix,
                pbc=True,
            )

        # This one does not need symprec.
        # Just to keep the signature the same.
        del symprec
        for _, p in self.make_patterns():
            yield _from_pymatgen_struct_to_ase_atoms(self._get_structure(p))

    def ewalds(self, symprec=None):
        """
        Generate Ewald energies for all unique patterns.

        Args:
            symprec (float | None): Optional symmetry precision forwarded to CIF generation.

        Yields:
            float: Ewald total energy for each configuration.
        """
        for _, p in self.make_patterns():
            yield self._get_ewald(p, symprec)

    def _get_letters(self, p):
        """
        (sub. data structure) turn list of list
        into list of some generic symbols from symbols_list.
        Used by `configurations()` and `cifwriter()`.
        Probably not the best idea. Graphical explanation:

        [[0, 1, 2], [0, 2], [0, 1], [1]]  # pattern
        +
        ((2, 1), (1, 1))  # segmenter
        +
        [a, b, c, d]  # symbols_list
        =
        [b, a, b, c, d]  # il
        """
        si = iter(self._charset)
        pi = iter(p)
        il = []
        for se in self._segmenter:
            sei = iter(se)
            word = [next(si)] * sum(se)
            next(pi)
            next(sei)
            # Iterate disoder sites
            for _ in sei:
                letter = next(si)
                # Iterate subpatterns
                for j in next(pi):
                    word[j] = letter
            il.extend(word)
        return "".join(il)

    def _get_weight(self, a):
        return len(self._symmops) // a.size

    def _get_cifwriter(self, p, symprec=None):
        """
        Return cifwriter for the given pattern.

        Args:
            p: Substitution pattern.
        """
        # TODO: Simplify or clarify.
        des = self._disorder_elements()
        orbits = des.keys()
        gis = self._group_indices

        # If the input requested MSG, emit magCIF output preserving moments.
        if _structure_requests_msg(self._structure):
            structure = self._get_structure(p)
            return _StringWriter(_structure_to_mcif_string(structure))

        # Build template CifWriter. First run only.
        if self._template_cifwriter is None:
            # In this template, all sites are ordered / single-occupied.
            template_structure = self._structure.copy()
            pi = iter(p)
            for orbit in orbits:
                indices = gis[orbit]
                de = des[orbit]
                for e in de:
                    subpattern = next(pi)
                    for i in subpattern:
                        template_structure.sites[indices[i]].species = e
            # NOTE: Ensure unique labels only for template creation to avoid CIF warnings.
            # This does not affect actual output labels, which are overwritten below.
            labels = [f"{site.species_string}{i}" for i, site in enumerate(template_structure.sites)]
            if "_atom_site_label" in template_structure.site_properties:
                template_structure.remove_site_property("_atom_site_label")
            template_structure.add_site_property("_atom_site_label", labels)
            for label, site in zip(labels, template_structure.sites):
                if hasattr(site, "label"):
                    site.label = label
            cifwriter = CifWriter(template_structure)

            # Use faster CifBlock implementation
            cfkey = cifwriter.cif_file.data.keys()
            cfkey = list(cfkey)[0]
            block = _AltCifBlock.from_str(str(cifwriter.cif_file.data[cfkey]))
            cifwriter.cif_file.data[cfkey] = block

            self._template_cifwriter = cifwriter
            self._template_structure = template_structure
        else:
            cifwriter = self._template_cifwriter
            template_structure = self._template_structure
            cfkey = cifwriter.cif_file.data.keys()
            cfkey = list(cfkey)[0]
            block = cifwriter.cif_file.data[cfkey]

        if symprec is None:
            type_symbol = block["_atom_site_type_symbol"].copy()
            label = block["_atom_site_label"].copy()

            # These two blocks are always "1" in the final structure
            block["_atom_site_symmetry_multiplicity"] = ["1"] * len(label)
            block["_atom_site_occupancy"] = ["1.0"] * len(label)

            pi = iter(p)
            for orbit in orbits:
                indices = gis[orbit]
                de = des[orbit]
                for e in de:
                    subpattern = next(pi)
                    for s in subpattern:
                        gi = indices[s]
                        type_symbol[gi] = str(e)
                        label[gi] = f"{e.symbol}{gi}"

            block["_atom_site_type_symbol"] = type_symbol
            block["_atom_site_label"] = label
        else:
            format_str = "{:.8f}"
            latt = template_structure.lattice.matrix
            positions = template_structure.frac_coords

            # this only actually
            cell_specie = list(set(x.species for x in template_structure))
            # Flattened list of species @ disorder sites
            specie = [y for x in des.values() for y in x]
            z_map = [cell_specie.index(Composition({specie[j]: 1})) for j in range(len(p))]
            zs = [cell_specie.index(x.species) for x in template_structure]

            pi = iter(p)
            zi = iter(z_map)
            for orbit in orbits:
                indices = gis[orbit]
                de = des[orbit]
                for e in de:
                    subpattern = next(pi)
                    z = next(zi)
                    for s in subpattern:
                        gi = indices[s]
                        zs[gi] = z

            space_group_data = spglib.get_symmetry_dataset(
                (latt, positions, zs),
                symprec=self._symprec,
                angle_tolerance=self._angle_tolerance,
            )

            rotations = space_group_data.rotations
            translations = space_group_data.translations
            equivalent_atoms = space_group_data.equivalent_atoms
            international = space_group_data.international
            number = space_group_data.number

            ops = [transformation_to_string(rot, trans, delim=", ") for rot, trans in zip(rotations, translations)]
            u, inv = np.unique(equivalent_atoms, return_inverse=True)
            equivalent_indices = [[] for _ in range(len(u))]
            for j, inv in enumerate(inv):
                equivalent_indices[inv].append(j)
            unique_indices = [
                (
                    sorted(
                        j,
                        key=lambda s: tuple(abs(x) for x in template_structure.sites[s].frac_coords),
                    )[0],
                    len(j),
                )
                for j in equivalent_indices
            ]
            unique_indices = sorted(
                unique_indices,
                key=lambda t: (
                    cell_specie[zs[t[0]]].average_electroneg,  # careful here
                    -t[1],
                    template_structure.sites[t[0]].a,
                    template_structure.sites[t[0]].b,
                    template_structure.sites[t[0]].c,
                ),
            )

            block["_symmetry_space_group_name_H-M"] = international
            block["_symmetry_Int_Tables_number"] = number
            block["_symmetry_equiv_pos_site_id"] = [str(i) for i in range(1, len(ops) + 1)]
            block["_symmetry_equiv_pos_as_xyz"] = ops

            block["_atom_site_type_symbol"] = []
            block["_atom_site_label"] = []
            block["_atom_site_symmetry_multiplicity"] = []
            block["_atom_site_fract_x"] = []
            block["_atom_site_fract_y"] = []
            block["_atom_site_fract_z"] = []
            block["_atom_site_occupancy"] = []
            count = 0
            for j, mult in unique_indices:
                # Careful: The structure species itself is not updated
                sp = cell_specie[zs[j]].elements[0]  # careful here
                site = template_structure.sites[j]
                block["_atom_site_type_symbol"].append(sp.__str__())
                block["_atom_site_label"].append(f"{sp.symbol}{count}")
                block["_atom_site_symmetry_multiplicity"].append(str(mult))
                block["_atom_site_fract_x"].append(format_str.format(site.a))
                block["_atom_site_fract_y"].append(format_str.format(site.b))
                block["_atom_site_fract_z"].append(format_str.format(site.c))
                block["_atom_site_occupancy"].append("1.0")
                count += 1

        return cifwriter

    def _get_structure(self, p):
        """
        Get Pymatgen structure for the given substitution pattern.
        """
        des = self._disorder_elements()
        orbits = des.keys()
        gis = self._group_indices

        structure = self._structure.copy()
        pi = iter(p)
        for orbit in orbits:
            indices = gis[orbit]
            de = des[orbit]
            for e in de:
                subpattern = next(pi)
                for i in subpattern:
                    site = structure.sites[indices[i]]
                    site.species = e

                    # Magnetic moment propagation (when template came from magCIF)
                    mom_map = site.properties.get("_atom_site_moment_by_label")
                    if mom_map:
                        # Determine element symbol for the new species (best effort)
                        symbol = None
                        try:
                            symbol = e.symbol  # Element
                        except Exception:
                            try:
                                elems = getattr(e, "elements", None)
                                if elems and len(elems) == 1:
                                    symbol = list(elems)[0].symbol
                            except Exception:
                                symbol = None

                        # Prefer transforming the base moment for this element label
                        # using the site's recorded magnetic symmetry operation.
                        moments_by_label = (getattr(structure, "properties", None) or {}).get("_atom_site_moments_by_label", {})

                        def _pick_base_label(sym: str | None) -> str | None:
                            if not sym:
                                return None
                            if sym in moments_by_label:
                                return sym
                            for k in sorted(moments_by_label.keys()):
                                if isinstance(k, str) and k.startswith(sym):
                                    return k
                            return None

                        def _transform_base_moment(base_vec):
                            op_aff = site.properties.get("_shry_magn_op_affine")
                            tr = site.properties.get("_shry_magn_time_reversal")
                            if op_aff is None or tr is None:
                                return base_vec
                            try:
                                r = np.array(op_aff, dtype=float)[:3, :3]
                                det = np.linalg.det(r)
                                det_sign = 1.0 if det >= 0 else -1.0
                                v = np.array(base_vec, dtype=float)
                                vv = det_sign * (r @ v)
                                vv = float(int(tr)) * vv
                                return (float(vv[0]), float(vv[1]), float(vv[2]))
                            except Exception:
                                return base_vec

                        chosen_label = _pick_base_label(symbol)
                        if chosen_label is not None:
                            base_vec = moments_by_label.get(chosen_label, (0.0, 0.0, 0.0))
                            vec = _transform_base_moment(base_vec)
                        else:
                            # Last resort: try the site's existing map.
                            chosen_label = next((k for k in mom_map.keys() if isinstance(k, str)), None)
                            vec = mom_map.get(chosen_label, (0.0, 0.0, 0.0)) if chosen_label else (0.0, 0.0, 0.0)

                        if chosen_label is not None:
                            site.properties["_atom_site_label"] = (chosen_label,)
                            site.properties["_atom_site_moment_by_label"] = {chosen_label: vec}
                            site.properties["_atom_site_occupancy_by_label"] = {chosen_label: 1.0}
                            site.properties["magmom"] = np.array(vec, dtype=float)
        return structure

    def _get_ewald(self, p, symprec):
        """
        Get ewald sums of a structure.
        """
        # TODO: less ad hoc implementation.
        cifwriter = self._get_cifwriter(p, symprec)
        structure = parse_cif_string_to_structure(
            str(cifwriter),
            primitive=False,
            sort=False,
            merge_tol=0.0,
            keep_oxidation_states=True,
        )
        try:
            if not np.isclose(structure.charge, 0.0):
                logging.warning(f"Unit cell is charged: (total charge = {structure.charge}).")
            return EwaldSummation(structure).total_energy
        except TypeError as exc:
            raise ValueError("Ewald summation required CIFs with defined oxidation states.") from exc


class _PatternMaker:
    """
    Generate permutationally-unique patterns
    """

    __slots__ = (
        "_search",
        "ap",
        "_get_polya",
        "_patterns",
        "_auts",
        "_subobj_ts",
        "_get_subobj_ts",
        "_nix",
        "invar",
        "_perms",
        "_nperm",
        "_column_index",
        "_relabel_index",
        "_row_index",
        "_bits",
        "_bit_perm",
        "label",
        "_cache",
        "_sieve",
        "_get_not_reject_mask",
        "_get_accept_mask",
        "_get_mins",
        "_iterate",
        "_rg",
        "_n_jobs",
        "_t_kind",
        "_shuffle",
    )

    def __init__(
        self,
        perm_list,
        invar=None,
        enumerator_collection=None,
        t_kind=const.DEFAULT_T_KIND,
        cache=False,
        shuffle=False,
        rg=None,
        n_jobs=-1,
    ):
        # Store n_jobs parameter
        self._n_jobs = n_jobs

        # Store t_kind and shuffle for parallel execution
        self._t_kind = t_kind
        self._shuffle = shuffle

        # Algo select bits
        if invar is None:
            self._search = self._invarless_search
        else:
            self._search = self._invar_search

        if t_kind == "sum":
            self._get_subobj_ts = self._get_sum_subobj_ts
            self._get_not_reject_mask = self._get_not_reject_mask_int
            self._get_accept_mask = self._get_accept_mask_int
            self._get_mins = self._get_mins_int
        elif t_kind == "plsum":
            self._get_subobj_ts = self._get_sum_subobj_ts
            self._get_not_reject_mask = self._get_not_reject_mask_float
            self._get_accept_mask = self._get_accept_mask_float
            self._get_mins = self._get_mins_float

            # Modify invar
            self._sieve = [[3, 3]]
            self._fill_sieve(invar.max() + 1)
            invar = self._logprime(invar)
        elif t_kind == "det":
            self._get_subobj_ts = self._get_det_subobj_ts
            self._get_not_reject_mask = self._get_not_reject_mask_float
            self._get_accept_mask = self._get_accept_mask_float
            self._get_mins = self._get_mins_float

            # Modify invar
            # self._sieve = [[3, 3]]
            # self._fill_sieve(invar.max() + 1)
            # invar = self._logprime(invar)
        else:
            raise RuntimeError(f'Unrecognized T kind "{t_kind}"')

        self._cache = cache
        if cache:
            self.ap = self.cached_ap
        else:
            self.ap = self.nocache_ap

        # Optional random.Random instance.
        if rg is None:
            self._rg = random
        else:
            self._rg = rg

        # This is to avoid an additional branch in the generator itself.
        if shuffle:

            def _iterate(x):
                nonzeros = np.flatnonzero(x)
                self._rg.shuffle(nonzeros)
                for i in nonzeros:
                    yield i

            self._iterate = _iterate
        else:
            self._iterate = np.flatnonzero

        if enumerator_collection is None:
            self._get_polya = _get_polya
        elif callable(enumerator_collection):
            self._get_polya = enumerator_collection
        else:
            self._get_polya = enumerator_collection.get

        (
            column_index,
            relabel_index,
            row_index,
            indexed_perm_list,
        ) = _PatternMaker.reindex(perm_list)

        if invar is not None:
            self.invar = invar[np.ix_(column_index, column_index)]
        else:
            self.invar = invar

        self._nperm, self._nix = indexed_perm_list.shape

        # The confusing rotations;
        self._column_index = column_index
        self._relabel_index = relabel_index
        self._row_index = row_index
        self._perms = indexed_perm_list

        # Bits below is for set-like comparison upon arrays.
        # Use pure python's unbounded integer for unlimited array length.
        self._bits = [2**i for i in range(self._nix)]
        # Convert to list to avoid overflow
        # Will change type to object for very large integer
        # Careful: each are still numpy object type...
        try:
            self._bit_perm = np.array(
                [[2 ** int(i) for i in row] for row in indexed_perm_list],
                dtype=int,
            )
        except OverflowError:
            self._bit_perm = np.array(
                [[2 ** int(i) for i in row] for row in indexed_perm_list],
                dtype=object,
            )

        # Containers for search results.
        self._patterns = collections.defaultdict(list)
        self._patterns[0] = [{}]
        self._auts = collections.defaultdict(list)
        self._auts[0] = np.ones((self._nperm,), dtype="bool")
        self._subobj_ts = collections.defaultdict(list)
        self._subobj_ts[0] = [np.array([0])]

        self.label = self._perms.tobytes()

    @staticmethod
    def reindex(perm_list):
        """Standardize symmetry ordering for reuse (rough)"""
        perm_list = perm_list.copy()

        # Column sort
        stab_map = perm_list == perm_list[0]
        # TODO: column_index is not much used...
        column_index = np.lexsort(stab_map)
        perm_list = perm_list[:, column_index]

        # TODO: More intuitive if we do this first, then the previous one.
        # Relabel to match column position
        relabel_index = perm_list[0]
        relabel_element = np.vectorize({s: i for i, s in enumerate(relabel_index)}.get)
        try:
            perm_list = relabel_element(perm_list)
        except TypeError as exc:
            raise ValueError(f"\n{perm_list}\nRows must have same elements.") from exc

        # Row sort
        row_index = np.lexsort(perm_list.T)
        relabeled_perm_list = perm_list[row_index]

        return column_index, relabel_index, row_index, relabeled_perm_list

    def update_index(self, perm_list):
        """
        Update map for the output pattern.

        This is used to update internal/external map in case
        of similar permutations.

        Raises:
            ValueError: If the internal permutation is unrelated to the
                supplied one.
        """
        (
            column_index,
            relabel_index,
            row_index,
            relabeled_perm_list,
        ) = _PatternMaker.reindex(perm_list)

        if relabeled_perm_list.tobytes() != self._perms.tobytes():
            raise ValueError(
                "The supplied permutation does not match the internal one.",
            )

        # Update values.
        self._column_index = column_index
        self._relabel_index = relabel_index
        self._row_index = row_index
        return row_index, relabel_index

    @staticmethod
    def get_label(perm_list):
        """
        Rough equivalence. Byte string from a standardized
        representation of the given permutation
        """
        _, _, _, relabeled_perm_list = _PatternMaker.reindex(perm_list)
        return relabeled_perm_list.tobytes()

    def get_row_map(self):
        """
        Map between internal order of permutation vs. input permutation
        """
        return self._row_index

    def get_index_map(self):
        """
        Map between internal representation of elements vs. input
        """
        return self._relabel_index

    # TODO: cleanup
    def get_column_map(self):
        """
        Internal column shuffle.
        """
        return self._column_index

    def cached_ap(self, n):
        """
        Get patterns and automorphisms for the specified replacement amount (cached)
        """
        # Patterns are symmetrical
        _n = min(n, self._nix - n)
        # if not _n in self._patterns:
        if n not in self._patterns:
            if _n in self._patterns:
                inverter = np.arange(self._nix)
                self._auts[n] = self._auts[_n]
                self._patterns[n] = [np.setdiff1d(inverter, p) for p in self._patterns[_n]]
            else:
                starts = [i for i in self._patterns.keys() if i < _n]
                if not starts:
                    start = 0
                else:
                    start = max(starts)

                # "Mirror" patterns
                if _n != n:
                    inverter = np.arange(self._nix)
                    ap = [(a, np.setdiff1d(inverter, p)) for a, p in self._search(start=start, stop=_n)]
                else:
                    ap = [(a, p) for a, p in self._search(start=start, stop=_n)]
                self._auts[n], self._patterns[n] = zip(*ap)

        for a, p in zip(self._auts[n], self._patterns[n]):
            yield a, p

    def nocache_ap(self, n):
        """
        Get patterns and automorphisms for the specified replacement amount
        """
        # Patterns are symmetrical
        _n = min(n, self._nix - n)
        if _n != n:
            inverter = np.arange(self._nix)
            for a, p in self._search(stop=_n):
                rp = np.setdiff1d(inverter, p)
                yield a, rp
        else:
            for a, p in self._search(stop=_n):
                yield a, p

    def _fill_sieve(self, n):
        """
        Incremental sieve of eratosthenes to make first n prime numbers
        """
        i = self._sieve[-1][0]
        while len(self._sieve) < (n - 1):
            i += 2  # no even primes except 2
            stop = int(math.sqrt(i))
            for tup in self._sieve:
                prime, mult = tup
                if prime > stop:
                    self._sieve.append([i, i])
                    break
                while mult < i:
                    mult += prime
                if mult == i:
                    break
            else:
                self._sieve.append([i, i])

    @functools.lru_cache()
    def _primes_list(self):
        return [2] + [x[0] for x in self._sieve]

    def _logprime(self, array):
        primed = [self._primes_list()[x] for x in array.flatten()]
        return np.log(np.array(primed).reshape(array.shape))

    def _expand_one_level_invarless(self, states):
        """
        Expand given states by one level for invarless search.

        Args:
            states: List of (pattern, aut, bs) tuples

        Returns:
            List of expanded (pattern, aut, bs) tuples
        """
        expanded_states = []

        for pattern, aut, pbs in states:
            # Tree is expanded by adding un-added sites
            leaf_mask = np.ones(self._nix, dtype="bool")
            leaf_mask[pattern] = False
            leaf_array = np.flatnonzero(leaf_mask)

            # Discard symmetry duplicates from the remaining leaves
            if aut.size > 1 and leaf_mask.sum() > 1:
                leaf_reps = self._perms[np.ix_(aut, leaf_array)].min(axis=0)
                leaf_indices = np.searchsorted(leaf_array, leaf_reps)
                uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
                uniq_mask[leaf_indices] = True
            else:
                uniq_mask = np.ones(leaf_array.shape, dtype="bool")

            # Insertion location
            loci = pattern.searchsorted(leaf_array)

            for i in self._iterate(uniq_mask):
                x = leaf_array[i]
                j = loci[i]
                _pbs = self._bit_perm[:, x] + pbs

                _i = np.concatenate((pattern[:j], [x], pattern[j:]))
                _aut = np.flatnonzero(_pbs == _pbs[-1])

                # Compute canonical parent.
                m = np.argmin(_pbs)
                can_pattern = self._perms[m, _i]
                discard_i = np.where(can_pattern == can_pattern.max())[0]

                # If the new site is discarded then tree parent == canonical parent.
                if j == discard_i:
                    expanded_states.append((_i, _aut, _pbs))
                # Check if tree parent is related to canonical parent.
                elif _i[discard_i] in self._perms[_aut, x]:
                    expanded_states.append((_i, _aut, _pbs))

        return expanded_states

    def _search_subtree_invarless(self, initial_state, stop):
        """
        Search a subtree starting from initial_state (for parallel execution).

        Args:
            initial_state: (pattern, aut, bs) tuple
            stop: Maximum substitution size

        Returns:
            List of (aut, pattern) tuples
        """
        results = []
        stack = [initial_state]

        while stack:
            pattern, aut, pbs = stack.pop()
            if pattern.size == stop:
                results.append((aut, pattern))
                continue

            # Tree is expanded by adding un-added sites
            leaf_mask = np.ones(self._nix, dtype="bool")
            leaf_mask[pattern] = False
            leaf_array = np.flatnonzero(leaf_mask)

            # Discard symmetry duplicates from the remaining leaves
            if aut.size > 1 and leaf_mask.sum() > 1:
                leaf_reps = self._perms[np.ix_(aut, leaf_array)].min(axis=0)
                leaf_indices = np.searchsorted(leaf_array, leaf_reps)
                uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
                uniq_mask[leaf_indices] = True
            else:
                uniq_mask = np.ones(leaf_array.shape, dtype="bool")

            # Insertion location
            loci = pattern.searchsorted(leaf_array)

            for i in self._iterate(uniq_mask):
                x = leaf_array[i]
                j = loci[i]
                _pbs = self._bit_perm[:, x] + pbs

                _i = np.concatenate((pattern[:j], [x], pattern[j:]))
                _aut = np.flatnonzero(_pbs == _pbs[-1])

                # Compute canonical parent.
                m = np.argmin(_pbs)
                can_pattern = self._perms[m, _i]
                discard_i = np.where(can_pattern == can_pattern.max())[0]

                # If the new site is discarded then tree parent == canonical parent.
                if j == discard_i:
                    stack.append((_i, _aut, _pbs))
                # Check if tree parent is related to canonical parent.
                elif _i[discard_i] in self._perms[_aut, x]:
                    stack.append((_i, _aut, _pbs))

        return results

    def _expand_one_level_invar(self, states):
        """
        Expand given states by one level for invar search.

        Args:
            states: List of (subobj_ts, pattern, aut, bs) tuples

        Returns:
            List of expanded (subobj_ts, pattern, aut, bs) tuples
        """
        expanded_states = []

        for subobj_ts, pattern, aut, pbs in states:
            # Invert index
            leaf_mask = np.ones(self._nix, dtype="bool")
            leaf_mask[pattern] = False
            leaf_array = np.flatnonzero(leaf_mask)

            # Calculate subobject Ts for all leaves
            leaf_subobj_ts = self._get_subobj_ts(pattern, leaf_array, subobj_ts)

            # Reject all leaves where any T is smaller than the new row's T
            delta_t = leaf_subobj_ts[:, :-1] - leaf_subobj_ts[:, -1:]
            not_reject_mask = self._get_not_reject_mask(delta_t)
            if not not_reject_mask.any():
                continue

            # Discard symmetry duplicates from the remaining leaves
            if aut.size > 1 and not_reject_mask.sum() > 1:
                not_reject_leaf = leaf_array[not_reject_mask]
                leaf_reps = self._perms[np.ix_(aut, not_reject_leaf)].min(axis=0)
                leaf_indices = leaf_array.searchsorted(leaf_reps)
                uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
                uniq_mask[leaf_indices] = True
            else:
                uniq_mask = not_reject_mask

            # Insertion location
            loci = pattern.searchsorted(leaf_array)

            accept_mask = self._get_accept_mask(delta_t)
            accept_mask &= uniq_mask
            accepts = leaf_array[accept_mask]

            for i in self._iterate(uniq_mask):
                x = leaf_array[i]
                j = loci[i]
                _pbs = self._bit_perm[:, x] + pbs

                _subobj_ts = leaf_subobj_ts[i]
                _subobj_ts[j:] = np.concatenate((_subobj_ts[-1:], _subobj_ts[j:-1]))

                _i = np.concatenate((pattern[:j], [x], pattern[j:]))
                _aut = np.flatnonzero(_pbs == _pbs[-1])

                if x in accepts:
                    expanded_states.append((_subobj_ts, _i, _aut, _pbs))
                    continue

                # Compute canonical parent.
                ts_min_i = self._get_mins(_subobj_ts)
                _sub = _i[ts_min_i]
                m = np.argmin(_pbs)
                can_pattern = self._perms[m, _sub]
                discard_i = ts_min_i[can_pattern == can_pattern.max()]

                # If the new site is discarded then tree parent == canonical parent.
                if j == discard_i:
                    expanded_states.append((_subobj_ts, _i, _aut, _pbs))
                # Check if tree parent is related to canonical parent.
                elif _i[discard_i] in self._perms[_aut, x]:
                    expanded_states.append((_subobj_ts, _i, _aut, _pbs))

        return expanded_states

    def _search_subtree_invar(self, initial_state, stop):
        """
        Search a subtree with invariant starting from initial_state (for parallel execution).

        Args:
            initial_state: (subobj_ts, pattern, aut, bs) tuple
            stop: Maximum substitution size

        Returns:
            List of (aut, pattern) tuples
        """
        results = []
        stack = [initial_state]

        while stack:
            subobj_ts, pattern, aut, pbs = stack.pop()
            if pattern.size == stop:
                results.append((aut, pattern))
                continue

            # Invert index
            leaf_mask = np.ones(self._nix, dtype="bool")
            leaf_mask[pattern] = False
            leaf_array = np.flatnonzero(leaf_mask)

            # Calculate subobject Ts for all leaves
            leaf_subobj_ts = self._get_subobj_ts(pattern, leaf_array, subobj_ts)

            # Reject all leaves where any T is smaller than the new row's T
            delta_t = leaf_subobj_ts[:, :-1] - leaf_subobj_ts[:, -1:]
            not_reject_mask = self._get_not_reject_mask(delta_t)
            if not not_reject_mask.any():
                continue

            # Discard symmetry duplicates from the remaining leaves
            if aut.size > 1 and not_reject_mask.sum() > 1:
                not_reject_leaf = leaf_array[not_reject_mask]
                leaf_reps = self._perms[np.ix_(aut, not_reject_leaf)].min(axis=0)
                leaf_indices = leaf_array.searchsorted(leaf_reps)
                uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
                uniq_mask[leaf_indices] = True
            else:
                uniq_mask = not_reject_mask

            # Insertion location
            loci = pattern.searchsorted(leaf_array)

            accept_mask = self._get_accept_mask(delta_t)
            accept_mask &= uniq_mask
            accepts = leaf_array[accept_mask]

            for i in self._iterate(uniq_mask):
                x = leaf_array[i]
                j = loci[i]
                _pbs = self._bit_perm[:, x] + pbs

                _subobj_ts = leaf_subobj_ts[i]
                _subobj_ts[j:] = np.concatenate((_subobj_ts[-1:], _subobj_ts[j:-1]))

                _i = np.concatenate((pattern[:j], [x], pattern[j:]))
                _aut = np.flatnonzero(_pbs == _pbs[-1])

                if x in accepts:
                    stack.append((_subobj_ts, _i, _aut, _pbs))
                    continue

                # Compute canonical parent.
                ts_min_i = self._get_mins(_subobj_ts)
                _sub = _i[ts_min_i]
                m = np.argmin(_pbs)
                can_pattern = self._perms[m, _sub]
                discard_i = ts_min_i[can_pattern == can_pattern.max()]

                # If the new site is discarded then tree parent == canonical parent.
                if j == discard_i:
                    stack.append((_subobj_ts, _i, _aut, _pbs))
                # Check if tree parent is related to canonical parent.
                elif _i[discard_i] in self._perms[_aut, x]:
                    stack.append((_subobj_ts, _i, _aut, _pbs))

        return results

    def _invarless_search(self, start=0, stop=None):
        """
        Alternative definition of canonical parent without invariant (much slower).

        Args:
            stop (int): Maximum substitution size before stop
        """
        if stop is None:
            stop = self._nix // 2
        stop = min(stop, self._nix - stop)

        enumerator = self._get_polya([self._perms])
        n_pred = enumerator.count(((stop, self._nix - stop),))

        # Prepare initial states for (parallel) search
        initial_states = []
        if not start:
            # Fill first stack
            noniso_orbits = np.unique(self._perms.min(axis=0))
            for _i in noniso_orbits:
                pattern = np.array([_i])
                bitsum = self._bits[_i]
                patch = np.zeros(self._nix, dtype=int)
                patch[_i] = 1
                bs = self._bit_perm.dot(patch)
                o_bitsums = self._bit_perm.dot(patch)
                aut = np.flatnonzero(o_bitsums == bitsum)
                initial_states.append((pattern, aut, bs))
        else:
            patterns = self._patterns[start]
            auts = self._auts[start]
            for pattern, aut in zip(patterns, auts):
                bs = self._bit_perm[:, pattern].sum(axis=1)
                initial_states.append((pattern, aut, bs))

        # Adaptive depth expansion: expand as deep as needed for good load balancing
        n_cores = multiprocessing.cpu_count() if self._n_jobs == -1 else abs(self._n_jobs)
        min_splits = max(n_cores * 8, 32)  # Target: 8x cores
        max_depth = min(stop - 1, 8)  # Can expand up to stop-1, max 8
        current_depth = 1 if not start else start

        logging.info(f"Initial: {len(initial_states)} subtrees, target: {min_splits}, max_depth: {max_depth}, stop: {stop}")

        # Keep expanding until we have enough subtrees or reach max depth
        while self._n_jobs != 1 and len(initial_states) < min_splits and current_depth < max_depth and current_depth < stop:
            prev_count = len(initial_states)
            current_depth += 1
            logging.info(f"Expanding to depth {current_depth} (currently {prev_count} subtrees, target: {min_splits})")
            initial_states = self._expand_one_level_invarless(initial_states)
            new_count = len(initial_states)
            logging.info(f"After expansion: {new_count} subtrees (increased by {new_count - prev_count})")
            if len(initial_states) >= min_splits:
                logging.info(f"Target reached: {len(initial_states)} subtrees at depth {current_depth}")
                break

        if self._n_jobs != 1 and len(initial_states) < min_splits:
            logging.warning(f"Could not reach target splits: {len(initial_states)} < {min_splits} (depth limit reached)")

        # Parallel execution if n_jobs != 1 (with dynamic load balancing)
        n_splits = len(initial_states)
        if self._n_jobs != 1 and n_splits > 1:
            n_cores = multiprocessing.cpu_count() if self._n_jobs == -1 else abs(self._n_jobs)
            logging.info(f"Parallel search (dynamic): {n_splits} subtrees across {n_cores} cores (n_jobs={self._n_jobs})")

            pbar_subtrees = tqdm.tqdm(
                total=n_splits,
                desc=f"Processing subtrees ({stop}/{self._nix})",
                **const.TQDM_CONF,
                disable=const.DISABLE_PROGRESSBAR,
            )

            # Use multiprocessing Pool with imap_unordered for dynamic load balancing
            pool = multiprocessing.Pool(
                processes=n_cores,
                initializer=_init_worker_invarless,
                initargs=(self._perms, self._bits, self._bit_perm, self._nix),
            )

            try:
                # imap_unordered with chunksize=1 for optimal load balancing
                for result in pool.imap_unordered(
                    _worker_search_invarless, [(state, stop) for state in initial_states], chunksize=1
                ):
                    pbar_subtrees.update(1)
                    # Yield all results from this subtree
                    for aut, pattern in result:
                        yield aut, pattern
            finally:
                pool.close()
                pool.join()

            pbar_subtrees.close()
            tqdm.tqdm(disable=const.DISABLE_PROGRESSBAR).write("Done.")
        else:
            # Sequential execution (original code)
            if n_splits > 1:
                logging.info(f"Sequential search: {n_splits} subtrees (n_jobs={self._n_jobs})")
            else:
                logging.info(f"Sequential search: 1 subtree (parallelization not applicable)")

            pbar = tqdm.tqdm(
                total=n_pred,
                desc=f"Making patterns ({stop}/{self._nix})",
                **const.TQDM_CONF,
                disable=const.DISABLE_PROGRESSBAR,
            )

            stack = list(initial_states)
            while stack:
                pattern, aut, pbs = stack.pop()
                if pattern.size == stop:
                    yield aut, pattern
                    pbar.update()
                    continue
                # Tree is expanded by adding un-added sites
                leaf_mask = np.ones(self._nix, dtype="bool")
                leaf_mask[pattern] = False
                leaf_array = np.flatnonzero(leaf_mask)

                # Discard symmetry duplicates from the remaining leaves
                if aut.size > 1 and leaf_mask.sum() > 1:
                    leaf_reps = self._perms[np.ix_(aut, leaf_array)].min(axis=0)
                    leaf_indices = np.searchsorted(leaf_array, leaf_reps)
                    uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
                    uniq_mask[leaf_indices] = True
                else:
                    uniq_mask = np.ones(leaf_array.shape, dtype="bool")

                # Insertion location
                loci = pattern.searchsorted(leaf_array)

                for i in self._iterate(uniq_mask):
                    x = leaf_array[i]
                    j = loci[i]
                    _pbs = self._bit_perm[:, x] + pbs

                    _i = np.concatenate((pattern[:j], [x], pattern[j:]))
                    _aut = np.flatnonzero(_pbs == _pbs[-1])

                    # Compute canonical parent.
                    m = np.argmin(_pbs)
                    can_pattern = self._perms[m, _i]
                    discard_i = np.where(can_pattern == can_pattern.max())[0]

                    # If the new site is discarded then tree parent == canonical parent.
                    if j == discard_i:
                        stack.append((_i, _aut, _pbs))
                    # Check if tree parent is related to canonical parent.
                    elif _i[discard_i] in self._perms[_aut, x]:
                        stack.append((_i, _aut, _pbs))
            pbar.close()
            tqdm.tqdm(disable=const.DISABLE_PROGRESSBAR).write("Done.")

    def _get_sum_subobj_ts(self, pattern, leaf_array, subobj_ts):
        new_rows = self.invar[np.ix_(leaf_array, pattern)]
        new_rows_sums = new_rows.sum(axis=1, keepdims=True)
        part_new_row_sums = subobj_ts + new_rows

        return np.concatenate((part_new_row_sums, new_rows_sums), axis=1)

    @functools.lru_cache(None)
    def _get_det(self, pattern):
        pattern_invar = self.invar[np.ix_(pattern, pattern)]
        return np.abs(np.linalg.det(np.atleast_2d(pattern_invar)))
        # return np.linalg.det(np.atleast_2d(pattern_invar))

    def _get_det_subobj_ts(self, pattern, leaf_array, subobj_ts):
        tiled_pattern = np.tile(pattern, (leaf_array.size, 1))
        leaves = np.column_stack((tiled_pattern, leaf_array))

        subobj_ts = []
        for leaf in leaves:
            # NOTE: Be careful with the order of these combinations...
            leaf_subobjs = list(itertools.combinations(leaf, leaf.size - 1))
            leaf_subobjs_ts = [self._get_det(x) for x in leaf_subobjs[::-1]]
            subobj_ts.append(leaf_subobjs_ts)

        return np.array(subobj_ts)

    @staticmethod
    def _get_not_reject_mask_int(delta_t):
        return ~(delta_t < 0).any(axis=1)

    def _get_not_reject_mask_float(self, delta_t):
        # return np.ones(delta_t.shape[0], dtype=bool)
        return ~((delta_t < 0.0) & ~np.isclose(delta_t, 0.0)).any(axis=1)

    @staticmethod
    def _get_accept_mask_int(delta_t):
        return (delta_t > 0).all(axis=1)

    def _get_accept_mask_float(self, delta_t):
        # return np.zeros(delta_t.shape[0], dtype=bool)
        return ((delta_t > 0.0) & ~np.isclose(delta_t, 0.0)).all(axis=1)

    @staticmethod
    def _get_mins_int(subobj_ts):
        return np.flatnonzero(subobj_ts == subobj_ts.min())

    @staticmethod
    def _get_mins_float(subobj_ts):
        return np.flatnonzero(np.isclose(subobj_ts, subobj_ts.min()))

    def _invar_search(self, start=0, stop=None):
        """
        Combines invar and lexmax to determine canonical parent.

        Args:
            start: Start seach at this depth.
            stop: Stop search at this depth.
        """
        if stop is None:
            stop = self._nix // 2
        stop = min(stop, self._nix - stop)

        enumerator = self._get_polya([self._perms])
        n_pred = enumerator.count(((stop, self._nix - stop),))

        # Prepare initial states for (parallel) search
        initial_states = []
        if not start:
            # Populate the first layer.
            noniso_orbits = np.unique(self._perms.min(axis=0))
            for _i in noniso_orbits:
                pattern = np.array([_i])
                # TODO: refine
                bitsum = self._bits[_i]
                patch = np.zeros(self._nix, dtype=int)
                patch[_i] = 1
                bs = self._bit_perm.dot(patch)
                aut = np.flatnonzero(bs == bitsum)
                subobj_ts = np.array([0])
                initial_states.append((subobj_ts, pattern, aut, bs))
        else:
            patterns = self._patterns[start]
            auts = self._auts[start]
            for pattern, aut in zip(patterns, auts):
                bs = self._bit_perm[:, pattern].sum(axis=1)
                subobj_ts = self.invar[np.ix_(pattern, pattern)].sum(axis=1)
                initial_states.append((subobj_ts, pattern, aut, bs))

        # Adaptive depth expansion: expand as deep as needed for good load balancing
        n_cores = multiprocessing.cpu_count() if self._n_jobs == -1 else abs(self._n_jobs)
        min_splits = max(n_cores * 8, 32)  # Target: 8x cores
        max_depth = min(stop - 1, 8)  # Can expand up to stop-1, max 8
        current_depth = 1 if not start else start

        logging.info(f"Initial: {len(initial_states)} subtrees, target: {min_splits}, max_depth: {max_depth}, stop: {stop}")

        # Keep expanding until we have enough subtrees or reach max depth
        while self._n_jobs != 1 and len(initial_states) < min_splits and current_depth < max_depth and current_depth < stop:
            prev_count = len(initial_states)
            current_depth += 1
            logging.info(f"Expanding to depth {current_depth} (currently {prev_count} subtrees, target: {min_splits})")
            initial_states = self._expand_one_level_invar(initial_states)
            new_count = len(initial_states)
            logging.info(f"After expansion: {new_count} subtrees (increased by {new_count - prev_count})")
            if len(initial_states) >= min_splits:
                logging.info(f"Target reached: {len(initial_states)} subtrees at depth {current_depth}")
                break

        if self._n_jobs != 1 and len(initial_states) < min_splits:
            logging.warning(f"Could not reach target splits: {len(initial_states)} < {min_splits} (depth limit reached)")

        # Parallel execution if n_jobs != 1 (with dynamic load balancing)
        n_splits = len(initial_states)
        if self._n_jobs != 1 and n_splits > 1:
            n_cores = multiprocessing.cpu_count() if self._n_jobs == -1 else abs(self._n_jobs)
            logging.info(f"Parallel search (dynamic): {n_splits} subtrees across {n_cores} cores (n_jobs={self._n_jobs})")

            pbar_subtrees = tqdm.tqdm(
                total=n_splits,
                desc=f"Processing subtrees ({stop}/{self._nix})",
                **const.TQDM_CONF,
                disable=const.DISABLE_PROGRESSBAR,
            )

            # Use multiprocessing Pool with imap_unordered for dynamic load balancing
            pool = multiprocessing.Pool(
                processes=n_cores,
                initializer=_init_worker_invar,
                initargs=(self._perms, self._bits, self._bit_perm, self._nix, self.invar, self._t_kind, self._shuffle),
            )

            try:
                # imap_unordered with chunksize=1 for optimal load balancing
                for result in pool.imap_unordered(
                    _worker_search_invar, [(state, stop) for state in initial_states], chunksize=1
                ):
                    pbar_subtrees.update(1)
                    # Yield all results from this subtree
                    for aut, pattern in result:
                        yield aut, pattern
            finally:
                pool.close()
                pool.join()

            pbar_subtrees.close()
            tqdm.tqdm(disable=const.DISABLE_PROGRESSBAR).write("Done.")
        else:
            # Sequential execution (original code)
            if n_splits > 1:
                logging.info(f"Sequential search: {n_splits} subtrees (n_jobs={self._n_jobs})")
            else:
                logging.info(f"Sequential search: 1 subtree (parallelization not applicable)")

            pbar = tqdm.tqdm(
                total=n_pred,
                desc=f"Making patterns ({stop}/{self._nix})",
                **const.TQDM_CONF,
                disable=const.DISABLE_PROGRESSBAR,
            )

            stack = list(initial_states)
            while stack:
                subobj_ts, pattern, aut, pbs = stack.pop()
                if pattern.size == stop:
                    yield aut, pattern
                    pbar.update()
                    continue
                # Invert index
                leaf_mask = np.ones(self._nix, dtype="bool")
                leaf_mask[pattern] = False
                leaf_array = np.flatnonzero(leaf_mask)

                # Calculate subobject Ts for all leaves
                leaf_subobj_ts = self._get_subobj_ts(pattern, leaf_array, subobj_ts)

                # Reject all leaves where any T is smaller than the new row's T
                delta_t = leaf_subobj_ts[:, :-1] - leaf_subobj_ts[:, -1:]
                not_reject_mask = self._get_not_reject_mask(delta_t)
                if not not_reject_mask.any():
                    continue

                # Discard symmetry duplicates from the remaining leaves
                if aut.size > 1 and not_reject_mask.sum() > 1:
                    not_reject_leaf = leaf_array[not_reject_mask]
                    leaf_reps = self._perms[np.ix_(aut, not_reject_leaf)].min(axis=0)
                    leaf_indices = leaf_array.searchsorted(leaf_reps)
                    uniq_mask = np.zeros(leaf_array.shape, dtype="bool")
                    uniq_mask[leaf_indices] = True
                else:
                    uniq_mask = not_reject_mask

                # Insertion location TODO: without this possible?
                loci = pattern.searchsorted(leaf_array)

                accept_mask = self._get_accept_mask(delta_t)
                accept_mask &= uniq_mask
                accepts = leaf_array[accept_mask]

                for i in self._iterate(uniq_mask):
                    x = leaf_array[i]
                    j = loci[i]
                    _pbs = self._bit_perm[:, x] + pbs

                    _subobj_ts = leaf_subobj_ts[i]
                    _subobj_ts[j:] = np.concatenate((_subobj_ts[-1:], _subobj_ts[j:-1]))

                    _i = np.concatenate((pattern[:j], [x], pattern[j:]))
                    # NOTE: just in case I fail to consistently sort perm
                    # bitsum = sum([self._bits[y] for y in _i])
                    # _aut = np.flatnonzero(_pbs == bitsum)
                    _aut = np.flatnonzero(_pbs == _pbs[-1])

                    if x in accepts:
                        stack.append((_subobj_ts, _i, _aut, _pbs))
                        continue

                    # Compute canonical parent.
                    ts_min_i = self._get_mins(_subobj_ts)
                    _sub = _i[ts_min_i]
                    m = np.argmin(_pbs)
                    can_pattern = self._perms[m, _sub]
                    discard_i = ts_min_i[can_pattern == can_pattern.max()]

                    # If the new site is discarded then tree parent == canonical parent.
                    if j == discard_i:
                        stack.append((_subobj_ts, _i, _aut, _pbs))
                    # Check if tree parent is related to canonical parent.
                    elif _i[discard_i] in self._perms[_aut, x]:
                        stack.append((_subobj_ts, _i, _aut, _pbs))
            pbar.close()
            tqdm.tqdm(disable=const.DISABLE_PROGRESSBAR).write("Done.")
            # self._gen_flag[stop] = True


class _Polya:
    """
    Perform operations related (weighed) Polya Enumeration Theorem.

    Includes enumeration (count) and calculation of cycle index
    and configuration generation function.

    Args:
        perms_list (list, tuple): List of permutations as (np.array).
            When specifying multiple substitution involving separated orbits,
            each of the separated group of sites should be assigned their own
            letter in the cycle index to avoid mixing of the substituted species.

            Rows of the arrays corresponds to a single permutation,
            where each element e permutes column c to e (c -> e).
            All arrays should contain the same number of rows.
    """

    def __init__(
        self,
        perm_list,
        group_size=None,
    ):
        # Used only for dividing the final weight
        if group_size is not None:
            self.group_size = group_size
        else:
            perm_sizes = set(x.shape[0] for x in perm_list)
            if not perm_sizes:
                self.group_size = 1
                logging.warning("No permutations supplied to enumerator.")
            else:
                self.group_size = max(perm_sizes)
        self._perm_list = perm_list
        self._expand_cache = dict()

    def label(self):
        """Use ci as label"""
        return sum(self.sym_ci().values()) / self.group_size

    @functools.lru_cache(None)
    def ci(self):
        """
        Returns the cycle index (as Counters of cycle lengths per permutation).

        Returns:
            dict: Dictionary of cycle index with integer index
                and Counters as a key-value pair.
        """
        # Use indexed base.
        cycle_index = collections.defaultdict(list)
        for permutations in self._perm_list:
            try:
                index_map = {s: i for i, s in enumerate(permutations[0])}
            except IndexError as e:
                logging.error("========================================")
                logging.error(self._perm_list)
                logging.error(permutations)
                logging.error("========================================")
                raise e
            # Iterate by row.
            for i, permutation in enumerate(permutations):
                # Build cycles.
                cycles = []
                visited = []
                for head in permutation:
                    if head in visited:
                        continue
                    cycle = []
                    while head not in cycle:
                        cycle.append(head)
                        visited.append(head)
                        try:
                            head = permutation[index_map[head]]
                        except (KeyError, IndexError) as exc:
                            logging.error("=============================")
                            logging.error(f"HEAD: {head}")
                            logging.error(f"IMAP: {index_map}")
                            logging.error(f"P: {permutation}")
                            logging.error(f"BP: {permutations}")
                            raise RuntimeError("Check permutation list.") from exc
                    cycles.append(cycle)
                counter = collections.Counter(len(cycle) for cycle in cycles)
                cycle_index[i].append(counter)
        return cycle_index

    def sym_ci(self):
        """
        Returns the symbolic cycle index.
        Uses sympy.IndexedBase to represent subscript.

        Conventionally, each term is in the form of (x_n**m),
        with (n) denoting cycle length and (m) number of
        cycles with such length for a permutation.

        Returns:
            dict: Dictionary of cycle index with integer index
                and sympy equations as a key-value pair.
        """
        cycle_index = self.ci()
        sym_cycle_index = dict()
        for j, permutations in enumerate(self._perm_list):
            symbol = sympy.IndexedBase(chr(97 + j))
            for i, _ in enumerate(permutations):
                sym_cycle_index.setdefault(i, sympy.Integer(1))
                for length, n in cycle_index[i][j].items():
                    sym_cycle_index[i] *= symbol[length] ** n
        return sym_cycle_index

    @functools.lru_cache(None)
    def count(self, amt_tuple):
        """
        Count unique pattern given.

        Args:
            amt_list (tuple): List of list, with the
                numbers in the inner list are coefficients
                of a chemical formula A_xB_y.

                In order to find results, the formula needs to be
                in the integer amount of the number of actual sites
                in the structure.

                Each elements addresses separate orbits.

        Returns:
            int: Number of unique patterns for the given amount.
        """

        def combine(a, b):
            assert len(a.shape) == 2 and len(b.shape) == 2
            ra, ca = a.shape
            rb, cb = b.shape
            assert ca == cb
            return (a[:, None] + b[None, :]).reshape(ra * rb, ca)

        # Should be cached too somehow
        # @functools.lru_cache(None)
        def exmul(arrays):
            """
            Calculate exponent coefficients "multiplied".
            """
            return functools.reduce(combine, arrays)

        logging.info(f"Counting unique patterns for {amt_tuple}.")

        # Padding
        joint_coeffs = np.array(sum(amt_tuple, ()))
        coeff_sums = [len(x) for x in amt_tuple]
        pads = [(sum(coeff_sums[:i]), sum(coeff_sums[i + 1 :])) for i in range(len(amt_tuple))]

        o_counts = []
        for o_cycles in self.ci().values():
            o_parts = [[_aR_array(cnum, len(color)) for cnum in cycles.values()] for cycles, color in zip(o_cycles, amt_tuple)]

            # Exponent values of each variables
            exps = [
                np.pad(clen * partition, [(0, 0), pad])
                for cycles, pad, partitions in zip(o_cycles, pads, o_parts)
                for clen, partition in zip(cycles, partitions)
            ]
            config_shape = [x.shape[0] for x in exps]

            mul_exps = exmul(exps)
            # Find matching coefficient; if none then 0
            match = np.flatnonzero((mul_exps == joint_coeffs).all(axis=1))
            match_i = np.array(np.unravel_index(match, config_shape)).T

            # TODO: Basically flatten; can be better written
            f_parts = [part for parts in o_parts for part in parts]
            counts = [
                functools.reduce(
                    lambda x, y: x * y,
                    [_multinomial_coeff(tuple(p[j])) for p, j in zip(f_parts, i)],
                )
                for i in match_i
            ]
            if not counts:
                counts = [0]
            o_counts.append(sum(counts))
        # Special case: no cycles means 1 unique object
        if not o_counts:
            return 1
        return int(sum(o_counts) / self.group_size)


def _serialize_perm_list(permutation_list):
    return tuple((arr.shape, arr.dtype.str, arr.tobytes()) for arr in permutation_list)


def _deserialize_perm_list(serialized):
    return [np.frombuffer(data, dtype=np.dtype(dtype)).reshape(shape) for shape, dtype, data in serialized]


@functools.lru_cache(None)
def _get_polya_cached(serialized_perms, group_size=None):
    perm_list = _deserialize_perm_list(serialized_perms)
    return _Polya(perm_list, group_size)


def _get_polya(permutation_list, group_size=None):
    """
    Return a cached Polya instance for the given permutation list.

    Uses an lru_cache by serializing permutation arrays into a hashable key.
    """
    serialized = _serialize_perm_list(permutation_list)
    return _get_polya_cached(serialized, group_size)
