src/uofa_cli/interrogate/init_wizard.py

"""Guided-setup helpers for `uofa interrogate init` (Addendum A14).

Pure, testable functions; the interactive prompting lives in
``commands/interrogate.py``. The hard rules (A14) are enforced here in code:

- **No silent scope.** ``build_scope`` records a provenance tag for every scope
  field (``extracted-from:<src>;confirmed-by-engineer`` or
  ``entered-by-engineer``); there is no path that emits a scope field without a
  provenance tag. The command's ``--yes`` is non-interactive but still requires a
  fully-provenanced ``--scope`` file — it does not fabricate or default scope.
- **Never fabricate reference.** Nothing here generates reference values; the
  reference is always a supplied input.
- **Smoke test before done (A14.3).** ``smoke_test_adapter`` exercises the
  generated adapter on one benchmark row and checks the output shape against the
  declared QoIs, so setup fails loudly rather than interrogation failing later.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

SUPPORTED_FORMATS = ("onnx", "torch", "sklearn", "tensorflow", "unknown")


def detect_model_format(model_path: Path) -> str:
    """Best-effort model-format detection by extension / directory markers."""
    path = Path(model_path)
    suffix = path.suffix.lower()
    if suffix == ".onnx":
        return "onnx"
    if suffix in (".pt", ".pth"):
        return "torch"
    if suffix in (".pkl", ".joblib"):
        return "sklearn"
    if path.is_dir() and (path / "saved_model.pb").exists():
        return "tensorflow"
    if suffix == ".pb":
        return "tensorflow"
    return "unknown"


# Per-format adapter bodies. Each defines `_load()` and the `predict` mapping.
# Where the load is deterministic (onnx/sklearn) the generated adapter is
# runnable immediately and the smoke test validates it; where it is not
# (torch needs the nn.Module class; tf varies) a clearly-marked TODO remains and
# the smoke test fails at setup with a pointer to complete it (A14.3).
_LOAD_BODY = {
    "onnx": (
        "import onnxruntime as ort\n"
        "        self._sess = ort.InferenceSession({model_path!r})\n"
        "        self._input_name = self._sess.get_inputs()[0].name"
    ),
    "sklearn": (
        "import joblib\n"
        "        self._model = joblib.load({model_path!r})"
    ),
    "torch": (
        "import torch  # TODO: import YOUR nn.Module subclass and construct it\n"
        "        # self._model = YourModule(); self._model.load_state_dict(torch.load({model_path!r}))\n"
        "        # self._model.eval()\n"
        "        raise NotImplementedError('Complete the torch model load, then re-run init')"
    ),
    "tensorflow": (
        "import tensorflow as tf\n"
        "        self._model = tf.saved_model.load({model_path!r})"
    ),
    "unknown": (
        "raise NotImplementedError('Unknown model format — implement load + predict, then re-run init')"
    ),
}

_PREDICT_BODY = {
    "onnx": (
        "out = self._sess.run(None, {{self._input_name: np.asarray(inputs, dtype=np.float32)}})[0]\n"
        "        out = np.asarray(out)"
    ),
    "sklearn": "out = np.asarray(self._model.predict(np.asarray(inputs, dtype=float)))",
    "torch": (
        "import torch\n"
        "        with torch.no_grad():\n"
        "            out = self._model(torch.as_tensor(np.asarray(inputs, dtype='float32'))).cpu().numpy()"
    ),
    "tensorflow": "out = np.asarray(self._model(np.asarray(inputs, dtype='float32')))",
    "unknown": "out = np.asarray(inputs)",
}


def generate_adapter_source(
    *, class_name: str, model_format: str, model_path: str,
    input_names: list[str], output_names: list[str],
) -> str:
    """Render a ModelAdapter subclass mapping inputs -> the declared QoIs."""
    fmt = model_format if model_format in _LOAD_BODY else "unknown"
    load = _LOAD_BODY[fmt].format(model_path=str(model_path))
    predict = _PREDICT_BODY[fmt]
    # Map output columns to QoI names; a single 1-D output maps to the first QoI.
    mapping_lines = []
    for index, name in enumerate(output_names):
        mapping_lines.append(
            f"            {name!r}: out[:, {index}] if out.ndim == 2 else out,"
        )
    mapping = "\n".join(mapping_lines) if mapping_lines else "            # TODO: map outputs to QoIs"
    return (
        '"""Auto-generated by `uofa interrogate init`. Review before interrogating.\n'
        f"Model format: {fmt}. Inputs (physical): {input_names}. QoIs: {output_names}.\n"
        '"""\n'
        "import numpy as np\n"
        "from uofa_cli.interrogate.adapter import ModelAdapter\n"
        "\n\n"
        f"class {class_name}(ModelAdapter):\n"
        "    def __init__(self):\n"
        f"        {load}\n"
        "\n"
        "    def predict(self, inputs):\n"
        f"        {predict}\n"
        "        return {\n"
        f"{mapping}\n"
        "        }\n"
    )


def build_scope(
    *,
    subject: dict,
    envelope_dimensions: list[dict],
    physics_constraints: list[dict],
    provenance: dict[str, str],
    evaluation_point: list[dict] | None = None,
    evaluation_region: list[dict] | None = None,
) -> dict:
    """Assemble the SIP scope config with a provenance tag for every field.

    ``provenance`` maps scope-field identifiers to a tag — caller MUST provide
    one per declared field (no silent default). Reference values are never part
    of scope.
    """
    scope: dict[str, Any] = {
        "subject": subject,
        "trainingEnvelope": {"dimensions": envelope_dimensions},
        "declaredPhysicsConstraint": physics_constraints,
        "scopeProvenance": dict(provenance),
    }
    if evaluation_point is not None:
        scope["evaluationPoint"] = {"coordinates": evaluation_point}
    if evaluation_region is not None:
        scope["evaluationRegion"] = {"dimensions": evaluation_region}
    return scope


def unprovenanced_scope_fields(scope: dict) -> list[str]:
    """Return declared scope fields that lack a provenance tag (must be empty).

    Enforces 'no silent scope': every envelope dimension, evaluation coordinate,
    and physics constraint must appear in scopeProvenance.
    """
    tags = scope.get("scopeProvenance", {})
    missing: list[str] = []
    for dim in scope.get("trainingEnvelope", {}).get("dimensions", []):
        key = f"trainingEnvelope.{dim.get('name')}"
        if key not in tags:
            missing.append(key)
    for coord in scope.get("evaluationPoint", {}).get("coordinates", []):
        key = f"evaluationPoint.{coord.get('name')}"
        if key not in tags:
            missing.append(key)
    for constraint in scope.get("declaredPhysicsConstraint", []):
        key = f"constraint.{constraint.get('constraintId')}"
        if key not in tags:
            missing.append(key)
    return missing


def smoke_test_adapter(adapter, benchmark_row, expected_output_names: list[str]) -> tuple[bool, str]:
    """Call predict on one row; check it returns a dict carrying the declared QoIs."""
    try:
        out = adapter.predict(benchmark_row)
    except Exception as exc:  # NotImplementedError from an incomplete template lands here
        return False, f"adapter.predict raised at setup: {exc}"
    if not isinstance(out, dict):
        return False, f"predict must return a dict of QoI->array, got {type(out).__name__}"
    missing = [name for name in expected_output_names if name not in out]
    if missing:
        return False, f"predict output is missing declared QoIs: {missing}"
    return True, "ok"