Source code for vibesketch.inputs.base

"""The input-writer interface.

A writer turns one species and one program-neutral job specification into
the input files of one program. It fails closed (AGENTS.md, ground rule 5):
before rendering anything it lists everything in the request it cannot
express, and if that list is not empty it raises :class:`WriterRefusal` with
the reasons instead of guessing a keyword. The vibe-qc QCSchema adapter's
stance is the model: reject what you cannot represent, never drop it.

Subclasses implement :meth:`InputWriter._refusals` and
:meth:`InputWriter._render`. :meth:`InputWriter.write` is the only public
entry point and always runs the checks first, so no subclass can skip them.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass
from typing import ClassVar

from vibesketch.document.model import Conformer, JobSpec, Species, Task

__all__ = [
    "InputBundle",
    "InputFile",
    "InputWriter",
    "WriterRefusal",
]




@dataclass(frozen=True)
class InputFile:
    """One file of a program input."""

    #: Suggested file name, without directories.
    name: str
    content: str
    media_type: str





@dataclass(frozen=True)
class InputBundle:
    """Everything a writer produced for one job."""

    program: str
    species: str
    files: tuple[InputFile, ...]





class WriterRefusal(ValueError):
    """The writer cannot express the request; ``reasons`` says what and why."""

    def __init__(self, program: str, reasons: Sequence[str]) -> None:
        self.program = program
        self.reasons = tuple(reasons)
        super().__init__(f"{program} cannot write this input:\n  " + "\n  ".join(self.reasons))



#: How each specification field reads in a refusal.
_FIELD_LABELS = {
    "method": "method",
    "basis": "basis set",
    "aux_basis": "auxiliary basis set",
    "approximation": "integral approximation",
    "dispersion": "dispersion correction",
    "solvent": "solvation model",
    "states": "number of excited states",
    "constraints": "geometry constraint",
    "resources": "resources",
}




class InputWriter(ABC):
    """Translate a job specification into one program's input, or refuse."""

    #: The registry name, also the value of ``Job.program``.
    name: ClassVar[str]
    #: One line for listings.
    description: ClassVar[str]
    #: The tasks this writer can express.
    tasks: ClassVar[frozenset[Task]]
    #: The specification fields, besides ``task``, this writer can express.
    fields: ClassVar[frozenset[str]] = frozenset()



    def refusals(self, species: Species, spec: JobSpec) -> list[str]:
        """Everything in the request this writer cannot express; empty if it can."""
        reasons: list[str] = []
        if spec.task not in self.tasks:
            supported = ", ".join(sorted(self.tasks))
            reasons.append(
                f"{self.name} cannot express task {spec.task!r}; it supports: {supported}"
            )
        for field in (*JobSpec.calculation_fields(), "resources"):
            value = getattr(spec, field)
            if value is not None and field not in self.fields:
                shown = (
                    value
                    if isinstance(value, (str, int))
                    else [c.model_dump(exclude_none=True) for c in value]
                    if isinstance(value, list)
                    else value.model_dump(exclude_none=True)
                )
                reasons.append(
                    f"{self.name} cannot express a {_FIELD_LABELS[field]} ({field}={shown!r})"
                )
        from vibesketch.document.solids import solid_refusal

        solid = solid_refusal(self.name, species)
        if solid:
            return [*reasons, *solid]  # a solid has no structure: nothing else applies (#142)
        reasons.extend(self._refusals(species, spec))
        return reasons




    def write(self, species: Species, spec: JobSpec) -> InputBundle:
        """The input files for ``species`` and ``spec``; raises :class:`WriterRefusal`."""
        reasons = self.refusals(species, spec)
        if reasons:
            raise WriterRefusal(self.name, reasons)
        return InputBundle(self.name, species.id, tuple(self._render(species, spec)))


    @abstractmethod
    def _refusals(self, species: Species, spec: JobSpec) -> list[str]:
        """Writer-specific reasons to refuse, beyond the task and fields."""

    @abstractmethod
    def _render(self, species: Species, spec: JobSpec) -> Sequence[InputFile]:
        """Produce the files; only called when there is nothing to refuse."""



def needs_conformer(writer: str, species: Species) -> list[str]:
    """The refusal for a species without 3D coordinates, if it has none."""
    if species.conformers:
        return []
    return [
        f"{writer} needs 3D coordinates, and species {species.id} has no conformer; "
        "embed it in 3D first"
    ]


def first_conformer(species: Species) -> Conformer:
    """The conformer writers use: the first, the one the document ranks highest."""
    return species.conformers[0]