generation

@dataclass
class RDKitGeneration[InputType: RDMolMolecule, OutputType: RecursiveMoleculeList](
    StructureGeneration, RDKitMaker[InputType, OutputType]
):
    """Generate 3D structures using RDKit distance geometry methods.

    This class uses RDKit's distance geometry embedding algorithms (ETKDG family)
    to generate 3D conformers from molecular graphs. The methods use distance
    bounds derived from experimental torsion angle preferences and small ring
    conformations to produce chemically reasonable 3D structures.

    The implementation supports all RDKit embedding parameters, allowing fine
    control over the generation process including optimization settings, pruning
    criteria, and random seed control.

    Units:
        Pass a float in the listed unit or a pint Quantity (e.g. ``jfchemistry.ureg``
        or ``jfchemistry.Q_``):

        - prune_rms_thresh: [Å]

    Attributes:
        name: Name of the job (default: "rdKit Generation").
        method: Distance geometry method to use:
            - "ETKDGv3": ETKDG version 3 (default, recommended)
            - "ETKDGv2": ETKDG version 2
            - "ETKDG": Original ETKDG method
            - "ETDG": Experimental torsion distance geometry
            - "KDG": Basic distance geometry
            - "srETKDGv3": Small ring ETKDG version 3
        basin_thresh: Energy threshold for basin hopping (default: None). Accepts float.
        bounds_mat_force_scaling: Scaling factor for distance bounds (default: None).
        box_size_mult: Box size multiplier for embedding (default: None).
        clear_confs: Clear existing conformers before embedding (default: None).
        embed_fragments_separately: Embed molecular fragments separately (default: None).
        enable_sequential_random_seeds: Use sequential random seeds (default: None).
        enforce_chirality: Enforce stereochemistry from molecular graph (default: None).
        force_trans_amides: Force amide bonds to trans configuration (default: None).
        ignore_smoothing_failures: Continue if smoothing fails (default: None).
        max_iterations: Maximum optimization iterations per conformer (default: None).
        num_threads: Number of parallel threads for embedding (default: 1).
        num_zero_fail: Number of failures before reporting (default: None).
        only_heavy_atoms_for_rms: Use only heavy atoms for RMSD pruning (default: None).
        optimizer_force_tol: Force tolerance for optimization (default: None).
        prune_rms_thresh: RMSD threshold for conformer pruning [Å] (default: None). \
            Accepts float or pint Quantity.
        rand_neg_eig: Randomize negative eigenvalues (default: None).
        random_seed: Random seed for reproducibility (default: None).
        symmetrize_conjugated_terminal_groups_for_pruning: Symmetrize terminal groups
            for RMSD calculations (default: None).
        timeout: Timeout [seconds] for embedding (default: None).
        track_failures: Track embedding failures (default: None).
        use_basic_knowledge: Use basic chemical knowledge (default: None).
        use_exp_torsion_angle_prefs: Use experimental torsion preferences (default: None).
        use_macrocycle_14config: Use 1,4 distance bounds for macrocycles (default: None).
        use_macrocycle_torsions: Use macrocycle torsion preferences (default: None).
        use_random_coords: Start from random coordinates (default: None).
        use_small_ring_torsions: Use small ring torsion preferences (default: None).
        use_symmetry_for_pruning: Use symmetry when pruning conformers (default: None).
        num_conformers: Number of conformers to generate (default: 1).

    """

    # Name of the job
    name: str = "rdKit Generation"
    # Number of conformers to generate
    num_conformers: Annotated[int, "positive"] = field(
        default=1,
        metadata={"description": "the number of conformers to generate"},
    )
    # Method to use for generating the structure
    method: Literal["ETDG", "ETKDG", "ETKDGv2", "ETKDGv3", "KDG", "srETKDGv3"] = field(
        default="ETKDGv3",
        metadata={"description": "The method to use for generating the structure"},
    )
    basin_thresh: Optional[float | Quantity] = field(
        default=None,
        metadata={
            "description": "set the basin threshold for the DGeom force field. \
                Accepts float or pint Quantity.",
            "unit": "kcal/mol",
        },
    )
    bounds_mat_force_scaling: Optional[float] = field(
        default=None,
        metadata={
            "description": "scale the weights of the atom pair distance restraints relative \
                to the other types of restraints."
        },
    )
    box_size_mult: Optional[float] = field(
        default=None,
        metadata={"description": "determines the size of the box used for random coordinates"},
    )
    clear_confs: Optional[bool] = field(
        default=None, metadata={"description": "clear existing conformers"}
    )
    embed_fragments_separately: Optional[bool] = field(
        default=None, metadata={"description": "embed fragments separately"}
    )
    enable_sequential_random_seeds: Optional[bool] = field(
        default=None,
        metadata={
            "description": "handle random number seeds so that \
                conformer generation can be restarted"
        },
    )
    enforce_chirality: Optional[bool] = field(
        default=None,
        metadata={"description": "enforce correct chirality if chiral centers are present"},
    )
    force_trans_amides: Optional[bool] = field(
        default=None, metadata={"description": "force trans amide bonds"}
    )
    ignore_smoothing_failures: Optional[bool] = field(
        default=None,
        metadata={
            "description": "try and embed the molecule if if triangle smoothing of the \
                bounds matrix fails"
        },
    )
    max_iterations: Optional[Annotated[int, "positive"]] = field(
        default=None,
        metadata={"description": "the maximum number of iterations to perform for each conformer"},
    )
    num_threads: Annotated[int, "positive"] = field(
        default=1,
        metadata={"description": "the number of threads to use for the embedding"},
    )
    num_zero_fail: Optional[Annotated[int, "positive"]] = field(
        default=None,
        metadata={"description": "fail embedding if we have at least this many zero eigenvalues"},
    )
    only_heavy_atoms_for_rms: Optional[bool] = field(
        default=None,
        metadata={"description": "Only consider heavy atoms when doing RMS filtering"},
    )
    optimizer_force_tol: Optional[float | Quantity] = field(
        default=None,
        metadata={
            "description": "the force tolerance during distance-geometry minimization. \
                Accepts float or pint Quantity.",
            "unit": "eV/Å",
        },
    )
    prune_rms_thresh: Optional[float | Quantity] = field(
        default=None,
        metadata={
            "description": "keep only conformations at least this far apart [Å]. \
                Accepts float in [Å] or pint Quantity.",
            "unit": "Å",
        },
    )
    rand_neg_eig: Optional[bool] = field(
        default=None,
        metadata={
            "description": "if the embedding yields a negative eigenvalue, pick coordinates that \
                correspond to this component at random"
        },
    )
    random_seed: Optional[Annotated[int, "positive"]] = field(
        default=None,
        metadata={
            "description": "if the embedding yields a negative eigenvalue, pick coordinates that \
                correspond to this component at random"
        },
    )
    symmetrize_conjugated_terminal_groups_for_pruning: Optional[bool] = field(
        default=None,
        metadata={"description": "symmetrize terminal groups for RMSD pruning"},
    )
    timeout: Optional[Annotated[int, "positive"]] = field(
        default=None,
        metadata={
            "description": "maximum time [seconds] to generate a conformer for a \
                single molecule fragment. If set to 0, no timeout is set"
        },
    )
    track_failures: Optional[bool] = field(
        default=None,
        metadata={"description": "keep track of which checks during the embedding process fail"},
    )
    use_basic_knowledge: Optional[bool] = field(
        default=None,
        metadata={
            "description": "impose basic-knowledge constraints \
            such as flat rings"
        },
    )
    use_exp_torsion_angle_prefs: Optional[bool] = field(
        default=None,
        metadata={"description": "impose experimental torsion angle preferences"},
    )
    use_macrocycle_14config: Optional[bool] = field(
        default=None,
        metadata={"description": "use the 1-4 distance bounds from ETKDGv3"},
    )
    use_macrocycle_torsions: Optional[bool] = field(
        default=None,
        metadata={"description": "impose macrocycle torsion angle preferences"},
    )
    use_random_coords: Optional[bool] = field(
        default=None,
        metadata={
            "description": "start the embedding from random coordinates \
                instead of using eigenvalues of the distance matrix"
        },
    )
    use_small_ring_torsions: Optional[bool] = field(
        default=None,
        metadata={"description": "impose small ring torsion angle preferences"},
    )
    use_symmetry_for_pruning: Optional[bool] = field(
        default=None,
        metadata={
            "description": "use molecule symmetry when doing the RMSD pruning. \
                Note that this option automatically also sets onlyHeavyAtomsForRMS to true."
        },
    )

    def __post_init__(self):
        """Normalize unit-bearing attributes."""
        if self.prune_rms_thresh is not None and isinstance(self.prune_rms_thresh, Quantity):
            object.__setattr__(
                self, "prune_rms_thresh", to_magnitude(self.prune_rms_thresh, "angstrom")
            )
        if self.basin_thresh is not None and isinstance(self.basin_thresh, Quantity):
            object.__setattr__(self, "basin_thresh", self.basin_thresh.magnitude)
        if self.optimizer_force_tol is not None and isinstance(self.optimizer_force_tol, Quantity):
            object.__setattr__(self, "optimizer_force_tol", self.optimizer_force_tol.magnitude)
        super().__post_init__()

    def _operation(
        self, input: InputType, **kwargs
    ) -> tuple[OutputType | list[OutputType], Properties | list[Properties] | None]:
        """Generate 3D structure(s) using RDKit distance geometry embedding.

        Embeds 3D coordinates into the molecule using the specified ETKDG method
        and parameters. Automatically configures the RDKit embedding parameters
        from the instance attributes and generates the requested number of conformers.

        The method:
        1. Creates RDKit embedding parameters object for the selected method
        2. Transfers all non-None attributes to the parameters
        3. Embeds conformers using distance geometry
        4. Converts conformers to Pymatgen Molecule objects
        5. Sets charge and spin multiplicity based on formal charge

        Args:
            input: RDKit molecule without 3D coordinates.
            **kwargs: Additional kwargs to pass to the operation.

        Returns:
            Tuple containing:
                - List of Pymatgen Molecule objects with 3D coordinates
                - Empty dictionary (no additional properties)

        Examples:
            >>> from rdkit import Chem # doctest: +SKIP
            >>> from jfchemistry import RDMolMolecule # doctest: +SKIP
            >>> from jfchemistry.generation import RDKitGeneration # doctest: +SKIP
            >>>
            >>> # Create molecule from SMILES
            >>> mol = Chem.MolFromSmiles("CCO") # doctest: +SKIP
            >>> rdmol = RDMolMolecule(Chem.AddHs(mol)) # doctest: +SKIP
            >>>
            >>> # Generate 10 conformers
            >>> gen = RDKitGeneration(num_conformers=10, random_seed=42) # doctest: +SKIP
            >>> structures, props = gen.operation(rdmol) # doctest: +SKIP
            >>> print(f"Generated {len(structures)} conformers") # doctest: +SKIP
        """
        params = getattr(rdDistGeom, self.method)()
        param_keys = [x[0] for x in inspect.getmembers(params)]
        for key, value in vars(self).items():

            def _to_camel_case(s: str) -> str:
                parts = s.split("_")
                return parts[0] + "".join(word.capitalize() for word in parts[1:])

            camel_key = _to_camel_case(key)
            if camel_key not in param_keys or value is None:
                continue
            setattr(params, camel_key, value)
        rdDistGeom.EmbedMultipleConfs(input, self.num_conformers, params)
        molecules: list[Molecule] = []
        for confId in range(input.GetNumConformers()):
            molecule: Molecule = Molecule.from_str(
                rdmolfiles.MolToXYZBlock(input, confId=int(confId)),
                fmt="xyz",
            )
            molecule.to("input.xyz", fmt="xyz")
            charge = rdmolops.GetFormalCharge(input)
            spin = int(2 * (abs(charge) // 2) + 1)
            molecule.set_charge_and_spin(charge, spin)
            molecules.append(molecule)
        if self.num_conformers == 1:
            return cast("OutputType", molecules[0]), Properties()
        else:
            return cast("OutputType", molecules), Properties()