Source code for denia.schema

"""Campaign schema: a machine-actionable description of an experimental table.

The schema is the contract between the agentic ingestion layer and the
classical active-learning core. Whether it is produced by an LLM agent
(:class:`denia.agents.SchemaAgent`) or by heuristic inference
(:func:`denia.ingest.infer_schema`), downstream components only ever see
this structure.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional




class ColumnRole(str, Enum):
    """Role of a column in an experimental campaign."""

    CATEGORY_INPUT = "category_input"      # e.g. element identity, preparation method
    FRACTION_INPUT = "fraction_input"      # e.g. mol% paired with a category column
    NUMERIC_INPUT = "numeric_input"        # e.g. temperature, pressure, contact time
    TARGET = "target"                      # objective(s) to optimize
    AUXILIARY = "auxiliary"                # measured but not optimized (e.g. conversions)
    METADATA = "metadata"                  # ids, references, free text
    IGNORE = "ignore"





@dataclass
class ColumnSpec:
    """Specification of a single column."""

    name: str
    role: ColumnRole
    unit: Optional[str] = None
    partner: Optional[str] = None  # for FRACTION_INPUT: the category column it quantifies
    maximize: bool = True          # for TARGET columns
    notes: str = ""





@dataclass
class CampaignSchema:
    """Full schema of an experimental campaign table."""

    columns: List[ColumnSpec] = field(default_factory=list)
    provenance: str = "heuristic"  # "heuristic" | "agent" | "user"

    # ------------------------------------------------------------------ helpers
    def _by_role(self, *roles: ColumnRole) -> List[str]:
        return [c.name for c in self.columns if c.role in roles]

    @property
    def targets(self) -> List[str]:
        return self._by_role(ColumnRole.TARGET)

    @property
    def numeric_inputs(self) -> List[str]:
        return self._by_role(ColumnRole.NUMERIC_INPUT)

    @property
    def category_inputs(self) -> List[str]:
        return self._by_role(ColumnRole.CATEGORY_INPUT)

    @property
    def fraction_inputs(self) -> List[str]:
        return self._by_role(ColumnRole.FRACTION_INPUT)

    @property
    def inputs(self) -> List[str]:
        return self._by_role(
            ColumnRole.CATEGORY_INPUT, ColumnRole.FRACTION_INPUT, ColumnRole.NUMERIC_INPUT
        )

    def spec(self, name: str) -> ColumnSpec:
        for c in self.columns:
            if c.name == name:
                return c
        raise KeyError(name)



    def composition_pairs(self) -> List[tuple]:
        """(category_column, fraction_column) pairs, e.g. ('Cation 1', 'Cation 1 mol%')."""
        pairs = []
        for c in self.columns:
            if c.role is ColumnRole.FRACTION_INPUT and c.partner:
                pairs.append((c.partner, c.name))
        return pairs


    # ------------------------------------------------------------------ io
    def to_dict(self) -> Dict:
        return {
            "provenance": self.provenance,
            "columns": [
                {
                    "name": c.name,
                    "role": c.role.value,
                    "unit": c.unit,
                    "partner": c.partner,
                    "maximize": c.maximize,
                    "notes": c.notes,
                }
                for c in self.columns
            ],
        }

    @classmethod
    def from_dict(cls, d: Dict) -> "CampaignSchema":
        cols = [
            ColumnSpec(
                name=c["name"],
                role=ColumnRole(c["role"]),
                unit=c.get("unit"),
                partner=c.get("partner"),
                maximize=c.get("maximize", True),
                notes=c.get("notes", ""),
            )
            for c in d.get("columns", [])
        ]
        return cls(columns=cols, provenance=d.get("provenance", "user"))

    def summary(self) -> str:
        lines = [f"CampaignSchema (provenance: {self.provenance})"]
        for role in ColumnRole:
            names = self._by_role(role)
            if names:
                lines.append(f"  {role.value:>16}: {', '.join(names)}")
        return "\n".join(lines)