utilities

@dataclass
class CombineMolecules(
    PymatGenMaker[Molecule | list[Molecule], Molecule],
):
    """Combine one or more Molecule objects into a single Molecule.

    Only supports pymatgen Molecule inputs (not periodic Structure).
    Molecules are appended in the order provided, in the same coordinate frame.
    """

    name: str = "CombineMolecules"
    _ensemble: bool = field(
        default=True,
        metadata={"description": "Treat list as single input to combine (ensemble mode)."},
    )

    def _operation(
        self,
        input: Molecule | list[Molecule],
        **kwargs: object,
    ) -> _OpResult:
        """Combine molecules into one."""
        if input is None:
            raise TypeError("molecules cannot be None")
        if isinstance(input, Molecule):
            items = [input]
        else:
            if len(input) == 0:
                raise ValueError("molecules list cannot be empty")
            for i, mol in enumerate(input):
                if mol is None:
                    raise TypeError(f"molecules[{i}] cannot be None")
                if not isinstance(mol, Molecule):
                    raise TypeError(
                        f"CombineMolecules only supports Molecule inputs; "
                        f"got molecules[{i}] as {type(mol).__name__}"
                    )
            items = input

        atoms = items[0].to_ase_atoms()
        for mol in items[1:]:
            atoms = atoms + mol.to_ase_atoms()
        combined = Molecule.from_ase_atoms(atoms)
        return cast("_OpResult", (combined, None))

@dataclass
class PropertiesToDisk(PymatGenMaker):
    """Persist one or more properties objects to JSON on disk.

    If a list is provided, indexed suffixes are appended to the base filename:
    e.g. ``properties.json`` -> ``properties_0.json``, ``properties_1.json``.
    """

    name: str = "Properties To Disk"
    output_dir: str = "."
    filename: str = "properties.json"
    _output_model: type[Output] = Output

    @staticmethod
    def _normalize_filename(filename: str) -> str:
        return filename if filename.lower().endswith(".json") else f"{filename}.json"

    @jfchem_job()
    def make(self, properties: Properties | list[Properties]) -> Response[_output_model]:
        """Write properties object(s) to JSON file(s)."""
        out_dir = Path(self.output_dir)
        out_dir.mkdir(parents=True, exist_ok=True)

        base = Path(self._normalize_filename(self.filename))
        written: list[str] = []

        if isinstance(properties, list):
            stem = base.stem
            suffix = base.suffix
            for idx, prop in enumerate(properties):
                p = Properties.model_validate(prop, extra="allow", strict=False)
                out_path = out_dir / f"{stem}_{idx}{suffix}"
                out_path.write_text(p.model_dump_json(indent=2))
                written.append(str(out_path))
        else:
            p = Properties.model_validate(properties, extra="allow", strict=False)
            out_path = out_dir / base
            out_path.write_text(p.model_dump_json(indent=2))
            written.append(str(out_path))

        return Response(output=Output(files={"properties_json": written}))

@dataclass
class RotateMolecule(PymatGenMaker[Molecule, Molecule]):
    """Rotate a Molecule using a rotation matrix, axis-angle, or principal-axes alignment.

    Only supports pymatgen Molecule (not Structure). Modes:
    - "matrix": apply a user-provided 3x3 rotation matrix.
    - "axis_angle": rotate by angle [degrees] around an axis (e.g. "z" or [0,0,1]).
    - "principal_axes": align principal moments of inertia with lab axes (e.g. largest along z).

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

        - angle_deg: [degrees]

    Set rotation parameters as instance attributes; then call make(molecule) or make([mol1, mol2]).
    """

    name: str = "RotateMolecule"

    # Rotation parameters (set at construction or on the instance)
    mode: Literal["matrix", "axis_angle", "principal_axes"] = "principal_axes"
    rotation_matrix: RotationMatrix | None = None
    axis: AxisSpec | None = None
    angle_deg: float | Quantity | None = None
    axis_order: str = "largest_z"
    center: tuple[float, float, float] | None = None

    def __post_init__(self):
        """Normalize unit-bearing attributes."""
        if self.angle_deg is not None and isinstance(self.angle_deg, Quantity):
            object.__setattr__(self, "angle_deg", to_magnitude(self.angle_deg, "degree"))
        super().__post_init__()

    def _operation(self, input: Molecule, **kwargs: object) -> _OpResult:
        """Rotate a molecule using this instance's mode and rotation parameters."""
        if input is None or not isinstance(input, Molecule):
            raise TypeError(
                "RotateMolecule only supports Molecule inputs; "
                f"got {type(input).__name__ if input is not None else 'None'}"
            )
        if self.mode not in ("matrix", "axis_angle", "principal_axes"):
            raise ValueError(
                f"mode must be 'matrix', 'axis_angle', or 'principal_axes'; got {self.mode!r}"
            )
        atoms = input.to_ase_atoms()
        pos = np.asarray(atoms.get_positions())
        masses = np.asarray(atoms.get_masses())
        center_pt = (
            np.average(pos, axis=0, weights=masses)
            if self.center is None
            else np.array(self.center, dtype=float)
        )
        if center_pt.shape != (3,):
            raise ValueError("center must be a length-3 sequence (x, y, z)")
        opts = {
            "mode": self.mode,
            "rotation_matrix": self.rotation_matrix,
            "axis": self.axis,
            "angle_deg": self.angle_deg,
            "axis_order": self.axis_order,
        }
        pos_new = _apply_rotation_mode(atoms, pos, masses, center_pt, opts)
        atoms.set_positions(pos_new)
        return cast("_OpResult", (Molecule.from_ase_atoms(atoms), None))

@dataclass
class SaveToDisk[
    InputType: RecursiveStructureList | RecursiveMoleculeList,
    OutputType: RecursiveStructureList | RecursiveMoleculeList,
](PymatGenMaker[InputType, OutputType]):
    """Save a structure or molecule to disk; appends suffixes for lists.

    Accepts a single Structure/Molecule or a list. When given a list, each
    item is written to a file with an appended suffix (e.g. output_0.xyz,
    output_1.xyz). Format is inferred from the filename extension if not
    provided.

    Attributes:
        name: Maker name for jobflow.
        suffix_fmt: Format string for list indices, e.g. "_{i}" -> output_0.xyz.
            Must contain exactly one "{i}" placeholder.
        filename: Base filename for saving.
        fmt: Format to save as.
    """

    name: str = "SaveToDisk"
    suffix_fmt: str = "_{i}"
    filename: str | None = None
    fmt: FormatLiteral | None = None
    _ensemble: bool = True

    @staticmethod
    def _write_one(
        obj: Structure | Molecule,
        path: Path,
        fmt: str | None = None,
    ) -> str:
        """Write a single structure or molecule to path. Returns absolute path as string."""
        path = path.resolve()
        fmt = fmt or _infer_fmt(path)
        if fmt not in VALID_FORMATS:
            raise ValueError(f"Invalid format: {fmt}")
        obj.to(filename=str(path), fmt=cast("Any", fmt))
        return str(path)

    def _operation(self, input: InputType, **kwargs) -> tuple[OutputType | list[OutputType], None]:
        """Save structure(s) or molecule(s) to disk.

        Args:
            input: Single Structure/Molecule or list of them.
            **kwargs: Additional kwargs to pass to the operation.

        Returns:
            Response with output.structure (unchanged input) and output.files
            (list of written absolute paths).

        Raises:
            TypeError: If input is not a Structure/Molecule or list of them.
            ValueError: If filename is empty, list is empty, or suffix_fmt invalid.
        """
        if not isinstance(self.filename, str) or not self.filename.strip():
            raise ValueError("filename must be a non-empty string")
        if "{i}" not in self.suffix_fmt:
            raise ValueError(
                "suffix_fmt must contain exactly one '{i}' placeholder for list "
                f"indices; got {self.suffix_fmt!r}"
            )
        single = not isinstance(input, list)
        if single:
            if not isinstance(input, (Structure, Molecule)):
                raise TypeError(
                    f"input must be a Structure, Molecule, or list of them; "
                    f"got {type(input).__name__}"
                )
            items = [input]
        else:
            if len(input) == 0:
                raise ValueError("input list cannot be empty")
            for i, obj in enumerate(input):
                if not isinstance(obj, (Structure, Molecule)):
                    raise TypeError(
                        f"input[{i}] must be a Structure or Molecule; got {type(obj).__name__}"
                    )
            items = input

        base = Path(self.filename).resolve()

        if single:
            paths = [base]
        else:
            paths = _paths_for_list(base, len(items), self.suffix_fmt)

        written: list[str] = []
        for obj, path in zip(items, paths, strict=True):
            written.append(self._write_one(obj, path, fmt=self.fmt))

        return (cast("OutputType", input), None)

@dataclass
class TranslateMolecule(PymatGenMaker[Molecule, Molecule]):
    """Translate a Molecule by a vector or move its center to the origin.

    Only supports pymatgen Molecule (not Structure). Modes:
    - "vector": apply a user-provided translation (dx, dy, dz).
    - "center_of_mass": move center of mass to the origin (mass-weighted).
    - "center_of_geometry": move centroid (center of geometry) to the origin.

    Set mode and translation as instance attributes; then call make(molecule) or make([mol1, ...]).
    """

    name: str = "TranslateMolecule"

    mode: Literal["vector", "center_of_mass", "center_of_geometry"] = "center_of_mass"
    translation: TranslationVector | None = None

    def _operation(self, input: Molecule, **kwargs: object) -> _OpResult:
        """Translate a molecule using this instance's mode and translation."""
        if input is None or not isinstance(input, Molecule):
            raise TypeError(
                "TranslateMolecule only supports Molecule inputs; "
                f"got {type(input).__name__ if input is not None else 'None'}"
            )
        atoms = input.to_ase_atoms()
        pos = np.asarray(atoms.get_positions(), dtype=float)

        if self.mode == "vector":
            if self.translation is None:
                raise ValueError("mode 'vector' requires translation (dx, dy, dz)")
            vec = np.asarray(self.translation, dtype=float)
            if vec.shape != (3,):
                raise ValueError("translation must be a length-3 vector")
            pos_new = pos + vec
        elif self.mode == "center_of_mass":
            masses = np.asarray(atoms.get_masses())
            com = np.average(pos, axis=0, weights=masses)
            pos_new = pos - com
        elif self.mode == "center_of_geometry":
            centroid = np.mean(pos, axis=0)
            pos_new = pos - centroid
        else:
            raise ValueError(
                f"mode must be 'vector', 'center_of_mass', or 'center_of_geometry'; "
                f"got {self.mode!r}"
            )

        atoms.set_positions(pos_new)
        return cast("_OpResult", (Molecule.from_ase_atoms(atoms), None))