packing

@dataclass
class PackmolPacking[InputType: list[Molecule], OutputType: Structure](
    PymatGenMaker[InputType, OutputType], StructurePacking
):
    """Pack molecules using Packmol.

    This class provides an interface to Packmol for packing molecules into
    simulation boxes. It supports two packing modes:
    - Box packing: Pack multiple copies of a molecule into a defined box
    - Fixed packing: Place molecules at fixed positions

    Units:
        Pass a float in the listed unit or a pint Quantity (e.g. ``jfchemistry.ureg``
        or ``jfchemistry.Q_``). For tuples (e.g. box_dimensions), each component
        may be a float or Quantity:

        - tolerance: [Å]
        - box_dimensions: [Å] (per component)
        - fixed_positions: [Å] (per component)
        - density: [g/cm³]

    Attributes:
        name: Name of the packing job (default: "Packmol Packing").
        packing_mode: Packing strategy to use:
            - "box": Pack molecules into a box (requires num_molecules and either
              box_dimensions or density)
            - "fixed": Place molecules at fixed positions (requires fixed_positions)
        box_dimensions: Box size [Å] as (x, y, z) tuple. Required for box mode
            if density is not specified.
        density: Target density [g/cm³]. If specified, box dimensions are automatically
            calculated for a cubic box. Cannot be used together with box_dimensions.
        num_molecules: Number of molecule copies to pack. Required for box mode.
        fixed_positions: List of (x, y, z) positions for fixed packing. Required for fixed mode.
        tolerance: Minimum distance between molecules [Å] (default: 2.0). \
            Accepts float or pint Quantity.
        packmol_path: Path to packmol executable (default: "packmol").
        filetype: Input/output file format (default: "xyz").
    """

    name: str = "Packmol Packing"
    packing_mode: Literal["box"] = field(
        default="box",
        metadata={"description": "the packing strategy to use"},
    )
    box_dimensions: Optional[tuple[float | Quantity, float | Quantity, float | Quantity]] = field(
        default=None,
        metadata={
            "description": "box size [Å] as (x, y, z) tuple. \
                Accepts float or pint Quantity per component.",
            "unit": "Å",
        },
    )
    density: Optional[float | Quantity] = field(
        default=None,
        metadata={
            "description": "target density [g/cm³] (alternative to box_dimensions). \
                Accepts float or pint Quantity.",
            "unit": "g/cm³",
        },
    )
    num_molecules: list[int] | int = field(
        default=1,
        metadata={"description": "number of molecule copies to pack"},
    )
    fixed_positions: Optional[list[tuple[float | Quantity, float | Quantity, float | Quantity]]] = (
        field(
            default=None,
            metadata={
                "description": "list of (x, y, z) positions [Å] for fixed packing. \
                    Accepts float or pint Quantity per component.",
                "unit": "Å",
            },
        )
    )
    tolerance: float | Quantity = field(
        default=2.0,
        metadata={
            "description": "minimum distance between molecules [Å]. \
                Accepts float or pint Quantity.",
            "unit": "Å",
        },
    )
    packmol_path: str = field(
        default="packmol",
        metadata={"description": "path to packmol executable"},
    )
    _filetype: str = "xyz"
    _structure_prefix: str = "structure"
    _ensemble: Literal[True] = True

    def __post_init__(self):
        """Post-initialization hook to make the output model."""
        if isinstance(self.num_molecules, int):
            self.num_molecules = [self.num_molecules]
        if isinstance(self.tolerance, Quantity):
            object.__setattr__(self, "tolerance", to_magnitude(self.tolerance, "angstrom"))
        if self.box_dimensions is not None:
            object.__setattr__(
                self,
                "box_dimensions",
                tuple(
                    to_magnitude(x, "angstrom") if isinstance(x, Quantity) else float(x)
                    for x in self.box_dimensions
                ),
            )
        if self.density is not None and isinstance(self.density, Quantity):
            object.__setattr__(self, "density", to_magnitude(self.density, "g/cm**3"))
        if self.fixed_positions is not None:
            object.__setattr__(
                self,
                "fixed_positions",
                [
                    tuple(
                        to_magnitude(v, "angstrom") if isinstance(v, Quantity) else float(v)
                        for v in pos
                    )
                    for pos in self.fixed_positions
                ],
            )

    def _calculate_box_dimensions_from_density(
        self, structures: InputType
    ) -> tuple[float, float, float]:
        """Calculate box dimensions from target density.

        Args:
            structures: Input molecule structure.

        Returns:
            Box dimensions as (x, y, z) tuple [Å] for a cubic box.

        Raises:
            ValueError: If num_molecules or density is not set.
        """
        if self.num_molecules is None:
            raise ValueError("num_molecules is required when using density")
        if self.density is None:
            raise ValueError("density is required for density-based box calculation")

        # Avogadro's number
        AVOGADRO = 6.02214076e23  # mol^-1
        total_mass = 0.0
        if isinstance(self.num_molecules, list):
            for num_molecule, structure in zip(self.num_molecules, structures, strict=True):
                # Calculate molecular weight in g/mol
                molecular_weight = structure.composition.weight.real  # g/mol
                total_mass += (num_molecule * molecular_weight) / AVOGADRO  # g
        else:
            # Calculate molecular weight in g/mol
            molecular_weight = structures[0].composition.weight.real  # g/mol
            total_mass = (self.num_molecules * molecular_weight) / AVOGADRO  # g
        # Calculate volume in cm^3
        volume_cm3 = total_mass / self.density  # cm^3

        # Convert to Angstrom^3 (1 cm^3 = 1e24 Angstrom^3)
        volume_ang3 = volume_cm3 * 1e24  # Angstrom^3

        # Calculate cubic box side length
        side_length = volume_ang3 ** (1.0 / 3.0)  # Angstrom

        self.box_dimensions = (side_length, side_length, side_length)
        return self.box_dimensions

    def _write_packmol_input(self, output_file: str, structure: InputType) -> str:
        """Generate packmol input file.

        Args:
            input_mol_file: Path to input molecule file.
            output_file: Path to output packed structure file.
            structure: Input molecule structure.

        Returns:
            Path to the generated packmol input file.

        Raises:
            ValueError: If required parameters are missing for the selected mode.
        """
        if self.packing_mode == "box":
            # At least one of box_dimensions or density must be set
            if self.box_dimensions is None:
                raise ValueError(
                    "Either box_dimensions or density must be specified for box packing mode"
                )
        elif self.packing_mode == "fixed":
            if self.fixed_positions is None:
                raise ValueError("fixed_positions is required for fixed packing mode")

        input_file = "packmol_input.inp"
        # Convert to absolute paths for packmol
        abs_output_file = os.path.abspath(output_file)

        with open(input_file, "w") as f:
            f.write(f"tolerance {self.tolerance}\n")
            f.write(f"filetype {self._filetype}\n")
            f.write(f'output "{abs_output_file}"\n')
            f.write("\n")

            if self.packing_mode == "box":
                if self.box_dimensions is None:
                    raise ValueError("box_dimensions is required for box packing mode")
                if self.num_molecules is None:
                    raise ValueError("num_molecules is required for box packing mode")
                for i, molecule in enumerate(structure):
                    if isinstance(self.num_molecules, list):
                        num_molecule = self.num_molecules[i]
                    else:
                        num_molecule = self.num_molecules
                    input_mol_file = os.path.abspath(
                        f"{self._structure_prefix}_{i}.{self._filetype}"
                    )
                    molecule.to(filename=input_mol_file, fmt=self._filetype)
                    f.write(f'structure "{input_mol_file}"\n')
                    f.write(f"  number {num_molecule}\n")
                    f.write(
                        f"  inside box 0. 0. 0. {self.box_dimensions[0]} \
                            {self.box_dimensions[1]} {self.box_dimensions[2]}\n"
                    )
                    f.write("end structure\n")
        return input_file

    def _run_packmol(self, input_file: str) -> None:
        """Execute packmol subprocess.

        Args:
            input_file: Path to packmol input file.

        Raises:
            RuntimeError: If packmol execution fails.
        """
        try:
            result = subprocess.run(
                [self.packmol_path, "-i", input_file],
                capture_output=False,
                text=True,
                check=False,
            )
            # Packmol returns non-zero exit codes even on success in some cases
            # Check if output file was created or if there's an actual error
            if result.returncode != 0:
                # Packmol writes errors to stdout
                error_output = result.stdout if result.stdout else result.stderr
                if "STOP" in error_output or "ERROR" in error_output.upper():
                    error_msg = (
                        f"Packmol execution failed (exit code {result.returncode}):\n{error_output}"
                    )
                    raise RuntimeError(error_msg)
        except FileNotFoundError:
            raise RuntimeError(
                f"Packmol executable not found at '{self.packmol_path}'. "
                "Please ensure packmol is installed and in your PATH."
            ) from None

    def _read_packed_structure(
        self, output_file: str, box_dimensions: Optional[tuple[float, float, float]] = None
    ) -> OutputType:
        """Read packmol output and convert to Structure.

        Args:
            output_file: Path to packmol output file.
            box_dimensions: Box dimensions to use for creating the lattice.
                If None, will use self.box_dimensions or calculate from density.

        Returns:
            Pymatgen Structure from the packed output.

        Raises:
            FileNotFoundError: If output file doesn't exist.
        """
        if not os.path.exists(output_file):
            raise FileNotFoundError(
                f"Packmol output file not found: {output_file}. "
                "Packmol may have failed to generate the packed structure."
            )

        # Read the packed structure
        if self._filetype == "xyz":
            mol = Molecule.from_file(output_file)
            # Convert to Structure for periodic systems (box mode)
            # For fixed mode, we can return as Molecule or Structure
            if self.packing_mode == "box":
                # Use provided box_dimensions or fall back to self.box_dimensions
                raw_dims = box_dimensions if box_dimensions is not None else self.box_dimensions
                if raw_dims is not None:
                    box_dimensions = cast(
                        "tuple[float, float, float]",
                        tuple(
                            to_magnitude(d, "angstrom") if isinstance(d, Quantity) else float(d)
                            for d in raw_dims
                        ),
                    )
                else:
                    box_dimensions = None
                if box_dimensions is not None:
                    # Create a Structure with the box as the lattice
                    from pymatgen.core import Lattice

                    # Create orthorhombic lattice with the box dimensions
                    lattice = Lattice.orthorhombic(
                        box_dimensions[0], box_dimensions[1], box_dimensions[2]
                    )
                    structure = Structure(
                        lattice=lattice,
                        species=mol.species,
                        coords=mol.cart_coords,
                        coords_are_cartesian=True,
                    )
                    return cast("OutputType", structure)
                else:
                    # Fallback: create a large enough lattice
                    from pymatgen.core import Lattice

                    max_coords = mol.cart_coords.max(axis=0)
                    min_coords = mol.cart_coords.min(axis=0)
                    box_size = max_coords - min_coords + 10.0  # Add padding
                    lattice = Lattice.cubic(max(box_size))
                    structure = Structure(
                        lattice=lattice,
                        species=mol.species,
                        coords=mol.cart_coords,
                        coords_are_cartesian=True,
                    )
                    return cast("OutputType", structure)
            else:
                # For fixed packing, return as Structure with no lattice
                from pymatgen.core import Lattice

                # Create a large enough lattice to contain all molecules
                max_coords = mol.cart_coords.max(axis=0)
                min_coords = mol.cart_coords.min(axis=0)
                box_size = max_coords - min_coords + 10.0  # Add padding
                lattice = Lattice.cubic(max(box_size))
                structure = Structure(
                    lattice=lattice,
                    species=mol.species,
                    coords=mol.cart_coords,
                    coords_are_cartesian=True,
                )
                return cast("OutputType", structure)
        else:
            # For other file types, try to read as Structure
            return cast("OutputType", Structure.from_file(output_file))

    def _operation(
        self, input: InputType, **kwargs
    ) -> tuple[OutputType | list[OutputType], Properties | list[Properties] | None]:
        """Pack a structure using Packmol.

        Args:
            input: The molecular structure to pack.
            **kwargs: Additional kwargs to pass to the operation.

        Returns:
            A tuple containing the packed structure and properties.

        Raises:
            ValueError: If required parameters are missing.
            RuntimeError: If packmol execution fails.
        """
        # Validate mode-specific parameters
        # Validation will be done in _write_packmol_input
        print(input)
        # Generate packmol output filename
        output_file = os.path.abspath(f"packed_structure.{self._filetype}")

        # Get box dimensions (either specified or calculated from density)
        box_dims: Optional[tuple[float | Quantity, float | Quantity, float | Quantity]] = None
        if self.packing_mode == "box":
            if self.density is not None:
                box_dims = self._calculate_box_dimensions_from_density(input)
            else:
                box_dims = self.box_dimensions

        # Write packmol input file (this may calculate box_dimensions from density)
        packmol_input = self._write_packmol_input(output_file, input)

        # Run packmol
        self._run_packmol(packmol_input)

        # Read packed structure (use absolute path)
        abs_output_file = os.path.abspath(output_file)
        if box_dims is not None:
            box_dims_normalized: tuple[float, float, float] = cast(
                "tuple[float, float, float]",
                tuple(
                    to_magnitude(d, "angstrom") if isinstance(d, Quantity) else float(d)
                    for d in box_dims
                ),
            )
        else:
            box_dims_normalized = None
        packed_structure = self._read_packed_structure(abs_output_file, box_dims_normalized)

        # Convert Structure to Molecule if needed (remove lattice for non-periodic representation)
        # For packed structures, we typically want to keep them as Structures,
        # but the base class signature requires Molecule. We'll convert to Molecule.
        # if isinstance(packed_structure, Structure):
        #     spin_multiplicity: int | None = None
        #     sp = getattr(packed_structure, "spin_multiplicity", None)
        #     if sp is not None:
        #         spin_multiplicity = int(sp)
        #     packed_molecule = Molecule(
        #         species=packed_structure.species,
        #         coords=packed_structure.cart_coords,
        #         charge=packed_structure.charge,
        #         spin_multiplicity=spin_multiplicity,
        #     )
        # else:
        #     packed_molecule = packed_structure

        # Prepare properties and convert to Properties object
        properties = self._get_properties(packed_structure)

        return packed_structure, properties

    def _get_properties(self, structure: Structure) -> Properties:
        """Get the properties of the packed structure.

        Args:
            structure: The packed structure.

        Returns:
            A dictionary containing packing metadata and structure properties.
        """
        properties: dict[str, Any] = {
            "packing_mode": self.packing_mode,
            "tolerance": self.tolerance,
            "num_atoms": len(structure),
        }

        if self.packing_mode == "box":
            properties["num_molecules"] = self.num_molecules

            # Get box dimensions from structure lattice or use specified/calculated values
            if hasattr(structure, "lattice") and structure.lattice is not None:
                # Extract box dimensions from lattice
                box_dims = (
                    structure.lattice.a,
                    structure.lattice.b,
                    structure.lattice.c,
                )
                properties["box_dimensions"] = box_dims
            elif self.box_dimensions is not None:
                box_dims = self.box_dimensions
                properties["box_dimensions"] = box_dims
            else:
                box_dims = None

            # Calculate density from packed structure
            if box_dims is not None and all(d > 0 for d in box_dims):
                volume = box_dims[0] * box_dims[1] * box_dims[2]
                # Calculate molecular weight from structure composition
                density = structure.density
                volume_cm3 = volume / 1e24  # Convert Angstrom^3 to cm^3
                properties["density"] = density if volume_cm3 > 0 else None
            else:
                properties["density"] = None

            # Include target density if it was specified
            if self.density is not None:
                properties["target_density"] = self.density  # g/cm^3
        elif self.packing_mode == "fixed" and self.fixed_positions is not None:
            properties["fixed_positions"] = self.fixed_positions
            properties["num_molecules"] = len(self.fixed_positions)

        return Properties(**properties)