conformers

@dataclass
class CRESTConformers[InputType: Molecule, OutputType: RecursiveMoleculeList](
    ConformerGeneration, PymatGenMaker[InputType, OutputType]
):
    """CREST conformer generation using metadynamics sampling.

    CREST (Conformer-Rotamer Ensemble Sampling Tool) performs automated
    conformational and rotameric searches using metadynamics simulations
    with GFN-xTB tight-binding methods. It efficiently explores conformational
    space to identify unique low-energy conformers.

    The implementation supports various metadynamics protocols and provides
    extensive control over optimization settings, energy calculations, and
    conformer filtering.

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

        - ewin: [kcal/mol]
        - ethr: [kcal/mol]
        - rthr: [Å]
        - bthr: [dimensionless]

    Attributes:
        name: Name of the job (default: "CREST Conformer Generation").
        runtype: Metadynamics protocol to use:
            - "imtd-gc": Iterative metadynamics with genetic crossing (default)
            - "nci-mtd": Non-covalent interaction metadynamics
            - "imtd-smtd": Iterative metadynamics with static metadynamics
        preopt: Pre-optimize structure before conformer search (default: True).
        multilevelopt: Use multi-level optimization (default: True).
        topo: Enable topology-based filtering (default: True).
        parallel: Number of parallel threads (default: 1).
        opt_engine: Optimization algorithm:
            - "ancopt": Approximate normal coordinate optimizer (default)
            - "rfo": Rational function optimizer

            - "gd": Gradient descent
        hess_update: Hessian update method:
            - "bfgs": BFGS update (default)
            - "powell": Powell update
            - "sd1": Steepest descent
            - "bofill": Bofill update
            - "schlegel": Schlegel update
        maxcycle: Maximum optimization cycles (default: None, auto).
        optlev: Optimization convergence level:
            - "crude", "vloose", "loose", "normal" (default), "tight", "vtight", "extreme"
        converge_e: Energy convergence threshold (default: None, auto).
        converge_g: Gradient convergence threshold (default: None, auto).
        freeze: Freeze constraints string (default: None).
        ewin: Energy window for conformer selection [kcal/mol] (default: 6.0).
        ethr: Energy threshold for duplicate detection [kcal/mol] (default: 0.05).
        rthr: RMSD threshold for structural similarity [Å] (default: 0.125).
        bthr: Rotational constant threshold for duplicate detection [dimensionless] (default: 0.01).
        calculation_energy_method: Method for energy calculations:
            - "gfn2" (default), "gfn1", "gfn0", "gfnff"
        calculation_energy_calcspace: Calculation space setting (default: None).
        calculation_energy_chrg: Charge for energy calculations (default: None, from structure).
        calculation_energy_uhf: Unpaired electrons for energy calc (default: None).
        calculation_energy_rdwbo: Read Wiberg bond orders (default: False).
        calculation_energy_rddip: Read dipole moments (default: False).
        calculation_energy_dipgrad: Compute dipole gradients (default: False).
        calculation_energy_gradfile: External gradient file (default: None).
        calculation_energy_gradtype: Gradient file type (default: None).
        calculation_dynamics_method: Method for metadynamics:
            - "gfn2" (default), "gfn1", "gfn0", "gfnff"
        calculation_dynamics_calcspace: Calculation space for dynamics (default: None).
        calculation_dynamics_chrg: Charge for dynamics (default: None, from structure).
        calculation_dynamics_uhf: Unpaired electrons for dynamics (default: None).
        calculation_dynamics_rdwbo: Read Wiberg bond orders in dynamics (default: False).
        calculation_dynamics_rddip: Read dipole in dynamics (default: False).
        calculation_dynamics_dipgrad: Compute dipole gradients in dynamics (default: False).
        calculation_dynamics_gradfile: External gradient file for dynamics (default: None).
        calculation_dynamics_gradtype: Gradient file type for dynamics (default: None).
        dynamics_dump_frequency: Dynamics trajectory dump frequency (default: 100.0).

    References:
        - CREST Documentation: https://crest-lab.github.io/crest-docs/
        - Pracht et al., PCCP 2020, 22, 7169-7192

    Examples:
        >>> from pymatgen.core import Molecule # doctest: +SKIP
        >>> from ase.build import molecule # doctest: +SKIP
        >>> from jfchemistry.conformers import CRESTConformers # doctest: +SKIP
        >>> mol = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
        >>> mol = mol.set_charge_and_spin(0, 1) # doctest: +SKIP
        >>> # Basic conformer search
        >>> conf_gen = CRESTConformers( # doctest: +SKIP
        ...     ewin=6.0, # doctest: +SKIP
        ...     calculation_energy_method="gfnff", # doctest: +SKIP
        ...     calculation_dynamics_method="gfnff" # doctest: +SKIP
        ... ) # doctest: +SKIP
        >>> structures, properties = conf_gen.operation(mol) # doctest: +SKIP
    """

    name: str = "CREST Conformer Generation"
    executable: str = field(default="crest", metadata={"description": "The CREST executable"})
    # Command line options
    solvation: Optional[SolvationType] = field(
        default=None, metadata={"description": "The solvation model to use for the calculation"}
    )
    charge: Optional[int] = field(
        default=None, metadata={"description": "The charge to use for the calculation"}
    )
    spin_multiplicity: Optional[int] = field(
        default=None, metadata={"description": "The spin multiplicity to use for the calculation"}
    )
    # General Settings Block
    threads: int = field(
        default=1, metadata={"description": "The number of threads to use for the calculation"}
    )
    runtype: Literal["imtd-gc", "nci-mtd", "imtd-smtd"] = field(
        default="imtd-gc", metadata={"description": "The run type to use for the calculation"}
    )
    preopt: bool = field(
        default=True,
        metadata={
            "description": "Whether to pre-optimize the structure before the conformer search"
        },
    )
    topo: bool = field(
        default=True, metadata={"description": "Whether to enable topology-based filtering"}
    )

    # Calculation Main Block
    opt_engine: Literal["ancopt", "rfo", "gd"] = field(
        default="ancopt",
        metadata={"description": "The optimization engine to use for the calculation"},
    )
    hess_update: Literal["bfgs", "powell", "sd1", "bofill", "schlegel"] = field(
        default="bfgs",
        metadata={"description": "The Hessian update method to use for the calculation"},
    )
    optlev: Literal["crude", "vloose", "loose", "normal", "tight", "vtight", "extreme"] = field(
        default="normal",
        metadata={"description": "The optimization convergence level to use for the calculation"},
    )

    # Optimization Calculation Block
    calculation_energy_method: Literal[
        "gfn2",
        "gfn1",
        "gfn0",
        "gfnff",
    ] = field(
        default="gfn2",
        metadata={"description": "The energy calculation method to use for the calculation"},
    )

    # Metadynamics Calculation Block
    calculation_dynamics_method: Literal[
        "gfn2",
        "gfn1",
        "gfn0",
        "gfnff",
    ] = field(
        default="gfn2",
        metadata={"description": "The metadynamics calculation method to use for the calculation"},
    )

    # Dynamics Block
    dynamics_dump_frequency: Optional[float] = field(
        default=100.0,
        metadata={
            "description": "The dynamics trajectory dump frequency (number of steps between dumps)"
        },
    )

    # CREGEN Block
    ewin: float | Quantity = field(
        default=6.0,
        metadata={
            "description": "The energy window for conformer selection [kcal/mol]. \
                Accepts float or pint Quantity.",
            "unit": "kcal/mol",
        },
    )
    ethr: float | Quantity = field(
        default=0.05,
        metadata={
            "description": "The energy threshold for duplicate detection [kcal/mol]. \
                Accepts float or pint Quantity.",
            "unit": "kcal/mol",
        },
    )
    rthr: float | Quantity = field(
        default=0.125,
        metadata={
            "description": "The RMSD threshold for structural similarity [Å]. \
                Accepts float or pint Quantity.",
            "unit": "Å",
        },
    )
    bthr: float | Quantity = field(
        default=0.01,
        metadata={
            "description": "The rotational constant threshold for duplicate detection \
                [dimensionless]. Accepts float or pint Quantity.",
            "unit": "dimensionless",
        },
    )

    # INTERNAL
    _input_dict: dict[str, Any] = field(
        default_factory=dict,
        metadata={"description": "Internal CREST input dictionary (built from public options)."},
    )
    _commands: list[str | int | float] = field(
        default_factory=list,
        metadata={"description": "Internal list of command-line arguments for CREST."},
    )
    _toml_filename: str = "crest.toml"
    _xyz_filename: str = "input.xyz"
    _properties_model: type[CRESTProperties] = CRESTProperties

    def __post_init__(self):
        """Normalize unit-bearing attributes."""
        if isinstance(self.ewin, Quantity):
            object.__setattr__(self, "ewin", to_magnitude(self.ewin, "kcal_per_mol"))
        if isinstance(self.ethr, Quantity):
            object.__setattr__(self, "ethr", to_magnitude(self.ethr, "kcal_per_mol"))
        if isinstance(self.rthr, Quantity):
            object.__setattr__(self, "rthr", to_magnitude(self.rthr, "angstrom"))
        if isinstance(self.bthr, Quantity):
            object.__setattr__(self, "bthr", to_magnitude(self.bthr, "dimensionless"))
        super().__post_init__()

    def _make_dict(self):
        """Make the dictionary for the CREST input."""
        self._input_dict["threads"] = self.threads
        self._input_dict["runtype"] = self.runtype
        self._input_dict["preopt"] = self.preopt
        self._input_dict["topo"] = self.topo
        self._input_dict["input"] = self._xyz_filename
        self._input_dict["calculation"] = {
            "opt_engine": self.opt_engine,
            "hess_update": self.hess_update,
            "optlev": self.optlev,
            "level": [
                {
                    "method": self.calculation_energy_method,
                },
                {
                    "method": self.calculation_dynamics_method,
                },
            ],
        }

        self._input_dict["cregen"] = {
            "ewin": self.ewin,
            "ethr": self.ethr,
            "rthr": self.rthr,
            "bthr": self.bthr,
        }

        self._input_dict["dynamics"] = {
            "active": [2],
            "dump": self.dynamics_dump_frequency,
        }

    def _write_toml(self):
        """Write the TOML file for the CREST input."""
        # Create a copy of the input nested dictionary without an key-value pairs with None values
        with open(self._toml_filename, "wb") as f:
            tomli_w.dump(self._input_dict, f)

    def _make_commands(self):
        """Make the CLI for the CREST input."""
        self._commands.append(self.executable)
        self._commands.append("--input")
        self._commands.append(self._toml_filename)
        if self.solvation is not None:
            self._commands.append(f"--{self.solvation[0]}")
            self._commands.append(self.solvation[1])
        if self.charge is not None:
            self._commands.append("--chrg")
            self._commands.append(str(self.charge))
        if self.spin_multiplicity is not None:
            self._commands.append("--uhf")
            self._commands.append(str(self.spin_multiplicity))

    def _run(self):
        """Run the CREST calculation."""
        # Save current working directory
        # original_dir = os.getcwd()
        subprocess.call(
            args=" ".join(str(x) for x in self._commands) + " > log.out",
            shell=True,
        )

        # # Create temporary directory and run crest there
        # with tempfile.TemporaryDirectory() as tmp_dir:
        #     # Copy input files to temp directory
        #     shutil.copy(self._xyz_filename, tmp_dir)
        #     shutil.copy(self._toml_filename, tmp_dir)

        #     # Change to temp directory
        #     os.chdir(tmp_dir)

        #     # Run crest command

        #     # Copy all files back to original directory
        #     for file in glob.glob("*"):
        #         shutil.copy(file, original_dir)

        #     # Change back to original directory
        #     os.chdir(original_dir)

        #     # Remove the temporary directory
        #     shutil.rmtree(tmp_dir)

    def _operation(
        self, input: InputType, **kwargs
    ) -> tuple[OutputType | list[OutputType], Properties | list[Properties] | None]:
        """Generate conformers using CREST metadynamics search.

        Performs a conformational search using CREST with the configured
        metadynamics protocol and GFN-xTB method. The calculation runs in
        a temporary directory and returns the unique low-energy conformers.

        Args:
            input: Input molecular structure with 3D coordinates. The
                structure's charge property is used if calculation charges
                are not explicitly set.
            **kwargs: Additional kwargs to pass to the operation.

        Returns:
            Tuple containing:
                - List of conformer structures sorted by energy
                - None (no additional properties returned)

        Examples:
            >>> from jfchemistry.conformers import CRESTConformers # doctest: +SKIP
            >>> from pymatgen.core import Molecule # doctest: +SKIP
            >>> from ase.build import molecule # doctest: +SKIP
            >>> mol = Molecule.from_ase_atoms(molecule("C2H6")) # doctest: +SKIP
            >>> mol = mol.set_charge_and_spin(0, 1) # doctest: +SKIP
            >>> gen = CRESTConformers(ewin=6.0, parallel=4) # doctest: +SKIP
            >>> conformers, props = gen.operation(mol) # doctest: +SKIP
        """
        # Write structures to sdf file
        input.to(self._xyz_filename, fmt="xyz")

        self.input = self._xyz_filename

        self._make_dict()
        self._write_toml()
        self._make_commands()
        self._run()

        conformers = XYZ.from_file("crest_conformers.xyz").all_molecules

        return (cast("list[OutputType]", conformers), self._properties_model())

@dataclass
class GOATConformers(ORCACalculator, ConformerGeneration):
    """Generate conformers using GOAT."""

    name: str = "GOAT Conformer Generation"
    goat: GoatKeywordsType = field(
        default="GOAT", metadata={"description": "The GOAT keyword to use for the calculation"}
    )
    _basename: str = "goat_conformer_generation"

    def _operation(self, molecule: Molecule) -> tuple[list[Molecule], ORCAProperties]:
        """Generate conformers using GOAT."""
        molecule.to("input.xyz", fmt="xyz")
        sk_list = super()._set_keywords()
        sk_list.append(getattr(Goat, self.goat.upper()))
        calc = super()._build_calculator(self._basename)
        calc.structure = Structure.from_xyz("input.xyz")
        super()._set_structure_charge_and_spin(
            calc, int(molecule.charge), int(molecule.spin_multiplicity)
        )
        super()._configure_calculator_input(calc, sk_list)
        calc.write_input()
        calc.run()
        output = calc.get_output()
        properties = super()._parse_output(output)
        conformers = XYZ.from_file(output.get_file(".xyz")).all_molecules
        return conformers, properties

@dataclass
class MMMCConformers[InputType: Molecule, OutputType: RecursiveMoleculeList](
    ConformerGeneration, PymatGenMaker[InputType, OutputType]
):
    """Generate conformers with the multiple minimum monte carlo method.

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

        - energy_window: [kcal/mol]
        - angle_step: [degrees]
        - rmsd_threshold: [Å]
    """

    name: str = "Multiple Minimum Monte Carlo Conformer Generation"
    optimizer: ASEOptimizer | TorchSimOptimizer = field(
        default_factory=lambda: ASEOptimizer,
        metadata={"description": "the calculator to use for the calculation"},
    )
    energy_window: float | Quantity = field(
        default=10.0,
        metadata={
            "description": "the energy window for the conformer ensemble [kcal/mol]. \
                Accepts float in [kcal/mol] or pint Quantity.",
            "unit": "kcal/mol",
        },
    )
    max_bonds_rotate: int = field(
        default=3,
        metadata={"description": "Maximum number of rotatable bonds to rotate in each step"},
    )
    max_attempts: int = field(
        default=1000,
        metadata={
            "description": "Maximum number of times to try and rotate dihedrals per iteration"
        },
    )
    angle_step: float | Quantity = field(
        default=30.0,
        metadata={
            "description": "Step size for bond rotation [degrees]. \
                Accepts float in [degrees] or pint Quantity.",
            "unit": "degrees",
        },
    )
    rmsd_threshold: float | Quantity = field(
        default=0.3,
        metadata={
            "description": "RMSD threshold for conformer selection [Å]. \
                Accepts float in [Å] or pint Quantity.",
            "unit": "Å",
        },
    )
    initial_optimization: bool = field(
        default=True,
        metadata={"description": "Whether to perform initial optimization of the structure"},
    )
    random_walk: bool = field(
        default=False, metadata={"description": "If True, use random walk for bond rotations"}
    )
    reduce_angle: bool = field(
        default=False,
        metadata={"description": " If True, reduce angle step size during the search"},
    )
    reduce_angle_every: int = field(
        default=50,
        metadata={
            "description": "Number of iterations to reduce angle step size.\
                Only applicable if reduce_angle is True"
        },
    )
    reduce_angle_by: int = field(
        default=2,
        metadata={
            "description": "Factor to reduce angle step size.\
                Only applicable if reduce_angle is True"
        },
    )
    only_heavy: bool = field(
        default=False,
        metadata={"description": "If True, only heavy atoms are considered for bond rotations"},
    )
    parallel: bool = field(default=False, metadata={"description": "If True, run MMMC in parallel"})
    num_cpus: int = field(
        default=1,
        metadata={
            "description": "Number of CPUs to use for MMMC. Only applicable if parallel is True"
        },
    )

    _filename: str = "molecule.xyz"
    _optimizer: Optional[ASEOptimizer | TorchSimOptimizer] = field(
        default=None,
        metadata={"description": "Cached optimizer instance (set internally after validation)."},
    )
    _properties_model: type[MMMCProperties] = MMMCProperties
    _output_model: type[MMMCOutput] = MMMCOutput

    def __post_init__(self):
        """Normalize unit-bearing attributes."""
        if isinstance(self.energy_window, Quantity):
            object.__setattr__(
                self, "energy_window", to_magnitude(self.energy_window, "kcal_per_mol")
            )
        if isinstance(self.angle_step, Quantity):
            object.__setattr__(self, "angle_step", to_magnitude(self.angle_step, "degree"))
        if isinstance(self.rmsd_threshold, Quantity):
            object.__setattr__(
                self, "rmsd_threshold", to_magnitude(self.rmsd_threshold, "angstrom")
            )
        super().__post_init__()

    def _operation(
        self, input: Molecule, **kwargs
    ) -> tuple[OutputType | list[OutputType], Properties | list[Properties]]:
        """Operation of the MMMC conformer generation."""
        # Write to XYZ file
        input.to(self._filename)
        # Create conformer
        conformer = Conformer(input_xyz=self._filename, charge=int(input.charge))
        # Create Optimizer
        if isinstance(self.optimizer, ASEOptimizer):
            atoms = self.optimizer.calculator._set_calculator(
                input.to_ase_atoms(),
                charge=int(input.charge),
                spin_multiplicity=int(input.spin_multiplicity),
            )
            calc = ASEOptimization(
                atoms.calc,
                optimizer=getattr(optimize, self.optimizer.optimizer),
                fmax=to_magnitude(self.optimizer.fmax, "eV/angstrom"),
                max_cycles=self.optimizer.steps,
            )
        elif isinstance(self.optimizer, TorchSimOptimizer):
            model = self.optimizer.calculator._get_model()
            calc = TorchSimCalculation(
                model=model,
                optimizer=ts.Optimizer.fire,
                max_cycles=self.optimizer.max_steps,
                device=self.optimizer.calculator.device,
            )
        # Build conformer ensemble
        conformer_ensemble = ConformerEnsemble(
            conformer=conformer,
            calc=calc,
            energy_window=to_magnitude(self.energy_window, "kcal_per_mol"),
            max_bonds_rotate=self.max_bonds_rotate,
            max_attempts=self.max_attempts,
            angle_step=to_magnitude(self.angle_step, "degree"),
            rmsd_threshold=to_magnitude(self.rmsd_threshold, "angstrom"),
            initial_optimization=self.initial_optimization,
            random_walk=self.random_walk,
            reduce_angle=self.reduce_angle,
            reduce_angle_every=self.reduce_angle_every,
            reduce_angle_by=self.reduce_angle_by,
            only_heavy=self.only_heavy,
            parallel=self.parallel,
            num_cpus=self.num_cpus,
        )
        # RUN
        conformer_ensemble.run_monte_carlo()
        # Get the best conformer
        if conformer.atoms is None:
            raise ValueError("Atoms are not set")
        if conformer.atoms is None or isinstance(conformer.atoms, list):
            raise ValueError("Atoms are not set")
        atom_symbols = list(conformer.atoms.get_chemical_symbols())
        molecules = [
            Molecule(species=atom_symbols, coords=coords)
            for coords in conformer_ensemble.final_ensemble
        ]
        from ase import io

        for _, m in enumerate(molecules):
            io.write("conformer.xyz", m.to_ase_atoms(), append=True)
        # Return the properties
        properties = [
            MMMCProperties(
                system=MMMCSystemProperties(
                    total_energy=SystemProperty(
                        name="total_energy", value=energy * ureg.kcal_per_mol
                    )
                )
            )
            for energy in conformer_ensemble.final_energies
        ]
        return (cast("OutputType", molecules), properties)