inputs

@dataclass
class PolymerInput(CoreMaker):
    """Polymer Input."""

    name: str = "Polymer Input"
    _output_model: type[PolymerInputOutput] = PolymerInputOutput
    _properties_model: type[Properties] = Properties

    @jfchem_job()
    def make(
        self, monomer: str, head: str | None = None, tail: str | None = None
    ) -> Response[_output_model]:
        """Make a polymer."""
        monomer_mol = rdmolops.AddHs(rdmolfiles.MolFromSmiles(monomer))
        head_mol = rdmolops.AddHs(rdmolfiles.MolFromSmiles(head)) if head else None
        tail_mol = rdmolops.AddHs(rdmolfiles.MolFromSmiles(tail)) if tail else None

        polymer = Polymer(
            head=RDMolMolecule(head_mol) if head_mol else None,
            monomer=RDMolMolecule(monomer_mol),
            tail=RDMolMolecule(tail_mol) if tail_mol else None,
        )
        return Response(output=PolymerInputOutput(structure=polymer))

@dataclass
class PubChemCID(MoleculeInput):
    """Retrieve molecules from PubChem database by compound ID.

    Downloads molecular structures from the PubChem database using the
    compound identifier (CID). The structure is retrieved in SDF format
    from PubChem's REST API.

    Attributes:
        name: Name of the input method (default: "PubChem CID Input").
        remove_salts: Inherited from MoleculeInput.
        add_hydrogens: Inherited from MoleculeInput.

    Examples:
        >>> from jfchemistry.inputs import PubChemCID
        >>>
        >>> # Retrieve ethanol (CID: 702)
        >>> pubchem = PubChemCID()
        >>> job = pubchem.make(702)
        >>> mol = job.output["structure"]
        >>>
        >>> # Retrieve without adding hydrogens
        >>> pubchem_no_h = PubChemCID(add_hydrogens=False)
        >>> job = pubchem_no_h.make(2244)  # Aspirin
    """

    name: str = "PubChem CID Input"

    def _get_structure(self, input: int | str) -> rdchem.Mol:
        """Fetch the structure from PubChem database.

        Downloads the molecular structure from PubChem using the compound ID
        and converts it to an RDKit Mol object.

        Args:
            input: PubChem compound ID (CID) as integer or string.

        Returns:
            RDKit Mol object from PubChem SDF data.
        """
        import requests
        from rdkit.Chem import rdmolfiles

        url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{input}/sdf"
        resp = requests.get(url)
        mol = rdmolfiles.MolFromMolBlock(resp.content.decode("utf-8"), removeHs=False)
        return mol

    @job(files="files", properties="properties")
    def make(self, input: int) -> Response[Output]:
        """Create a workflow job to retrieve a molecule from PubChem.

        Args:
            input: PubChem compound ID (CID).

        Returns:
            Response containing:
                - structure: RDMolMolecule from PubChem
                - files: MOL file representation

        """
        return super()._make(input)

@dataclass
class Smiles(MoleculeInput):
    """Create molecules from SMILES strings.

    Parses SMILES (Simplified Molecular Input Line Entry System) strings
    and converts them to RDKit Mol objects. Supports standard SMILES
    notation including stereochemistry and aromaticity.

    Attributes:
        name: Name of the input method (default: "SMILES Input").
        remove_salts: Inherited from MoleculeInput.
        add_hydrogens: Inherited from MoleculeInput.
    """

    name: str = "SMILES Input"

    def _get_structure(self, input: int | str) -> rdchem.Mol:
        """Parse SMILES string and convert to RDKit molecule.

        Args:
            input: SMILES string representation of the molecule.

        Returns:
            RDKit Mol object parsed from SMILES.

        Examples:
            >>> smiles = Smiles()
            >>> mol = smiles._get_structure("CCO")
            >>> mol.GetNumAtoms()  # Without hydrogens
            3
        """
        mol = rdmolfiles.MolFromSmiles(input)
        if mol is None:
            raise ValueError(f"Failed to parse SMILES string: {input}")
        return mol

    @job(files="files", properties="properties")
    def make(self, input: str) -> Response[Output]:
        """Create a workflow job to generate a molecule from SMILES.

        Args:
            input: SMILES string.

        Returns:
            Response containing:
                - structure: RDMolMolecule from SMILES
                - files: MOL file representation
        """
        resp = super()._make(input)
        return resp