"""Write Prismatic (https://prism-em.com) input files."""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pymatgen.core import IStructure, Structure


class Prismatic:
    """Write Prismatic (https://prism-em.com) input files.
    This is designed for STEM image simulation.
    """

    def __init__(self, structure: Structure | IStructure, comment: str = "Generated by pymatgen") -> None:
        """Initialize a Prismatic.

        Args:
            structure: pymatgen Structure
            comment (str): comment.
        """
        self.structure = structure
        self.comment = comment

    def to_str(self) -> str:
        """Get a Prismatic XYZ string representation.

        This is similar to XYZ format but has specific requirements for extra
        fields, headers, etc.
        """
        lattice = self.structure.lattice
        lines = [self.comment, " ".join(map(str, lattice.lengths))]
        for site in self.structure:
            for sp, occu in site.species.items():
                x, y, z = site.coords
                lines.append(f"{sp.Z} {x} {y} {z} {occu} {site.properties.get('thermal_sigma', 0)}")

        lines.append("-1")

        return "\n".join(lines)


# ----------------------------------------------------------------------------
# pymatgen.io.registry plugin: Structure -> Prismatic (write-only)
# ----------------------------------------------------------------------------


def _prismatic_write_str(structure, **kwargs) -> str:
    return Prismatic(structure, **kwargs).to_str()


def _register_formats() -> None:
    from pymatgen.io.registry import StructureFormat, register_structure_format

    register_structure_format(
        StructureFormat(
            name="prismatic",
            patterns=("*prismatic*",),
            write_str=_prismatic_write_str,
        )
    )


_register_formats()
