src/module_registry.py

"""Capability-based module registry for the TIDE evidence pipeline.

The pipeline is a set of *modules*. Each module declares a **capability**
(the kind of evidence it produces), an **eligibility** predicate (when it is
allowed to run for a given trial), and a **precedence** (how strongly it should
be preferred when several modules offer the same capability).

Selection is deliberately the module-level analogue of how the historical
comparator selects rows:

    capability match  ->  eligibility gate  ->  precedence ranking  ->  fallback

For each capability, every registered provider is evaluated against the trial
profile. Ineligible providers are skipped with a recorded reason. Among the
eligible providers the highest-precedence one becomes the *primary* (its result
is what the report consumes); lower-precedence eligible providers are recorded
as *superseded* fallbacks. If no provider is eligible the capability is simply
absent — nothing is fabricated.

This is what lets a new module (e.g. Layla's validated publication-likelihood
model) drop in without touching the report or the UI: it registers as a
higher-precedence provider of the ``publication_outlook`` capability, and the
selector prefers it automatically the moment its eligibility predicate passes.
"""

from __future__ import annotations

import shutil
from collections import OrderedDict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable

from modules.pubtime import summarize_from_archive
from summarizer import summarize_with_llm, template_summary


# --------------------------------------------------------------------------- #
# Module contract
# --------------------------------------------------------------------------- #

# applies(profile, ctx) -> (eligible, reason)
Predicate = Callable[[dict[str, Any], dict[str, Any]], "tuple[bool, str]"]
# run(profile, ctx) -> {"status": str, "result": Any}
Runner = Callable[[dict[str, Any], dict[str, Any]], dict[str, Any]]


@dataclass(frozen=True)
class ModuleSpec:
    name: str
    capability: str
    runtime: str  # "python" | "r"
    precedence: int
    description: str
    provenance: str
    applies: Predicate
    run: Runner
    supported_domains: tuple[str, ...] = field(default=())


def _domain_gate(supported: tuple[str, ...]) -> Predicate:
    def gate(profile: dict[str, Any], ctx: dict[str, Any]) -> tuple[bool, str]:
        domain = profile.get("domain", "")
        if not supported or domain in supported:
            return True, "eligible"
        return False, f"domain '{domain}' not covered by this module ({', '.join(supported)})."

    return gate


def _always(_profile: dict[str, Any], _ctx: dict[str, Any]) -> tuple[bool, str]:
    return True, "eligible"


# --------------------------------------------------------------------------- #
# Module implementations
# --------------------------------------------------------------------------- #

def _run_protocol_completeness(profile: dict[str, Any], _ctx: dict[str, Any]) -> dict[str, Any]:
    sections = profile.get("protocol_sections", {})
    total = sum(section["total"] for section in sections.values())
    filled = sum(section["filled"] for section in sections.values())
    ratio = round(filled / total, 3) if total else 0
    weakest = sorted(
        (
            {
                "section": name,
                "filled": section["filled"],
                "total": section["total"],
                "missing": section["missing"],
            }
            for name, section in sections.items()
        ),
        key=lambda item: (item["filled"] / item["total"]) if item["total"] else 0,
    )
    return {
        "status": "ok",
        "result": {
            "filled_fields": filled,
            "total_fields": total,
            "completion_ratio": ratio,
            "weakest_sections": weakest[:3],
        },
    }


def _run_historical_comparator(profile: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    result = summarize_from_archive(profile, ctx["project_root"])
    return {"status": "ok", "result": result}


def _run_comparator_base_rate(profile: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    """Empirical publication outlook, derived from the matched historical cohort.

    This is an honest base rate ("of trials like yours, X% published"), not a
    validated per-trial prediction. The predictive-model slot below is reserved
    for a model that produces a true per-trial probability.
    """
    evidence = ctx["capabilities"].get("historical_comparator") or {}
    summary = evidence.get("summary", {})
    return {
        "status": "ok",
        "result": {
            "publication_likelihood": summary.get("publication_rate"),
            "results_reporting_likelihood": summary.get("results_reported_rate"),
            "basis_rows": evidence.get("used_rows"),
            "match_strategy": evidence.get("match_strategy"),
            "model_type": "empirical_base_rate",
            "provenance_label": "Historical comparator",
            "provenance_detail": (
                "Publication rate among matched historical trials in the PubTime dataset. "
                "This is an empirical base rate, not a validated per-trial prediction."
            ),
            "predictive_model": {
                "status": "reserved",
                "reason": (
                    "Layla's validated publication-likelihood model can register as a "
                    "higher-precedence provider of the 'publication_outlook' capability; "
                    "the selector will then prefer it automatically."
                ),
            },
        },
    }


def _predictive_model_available(_profile: dict[str, Any], _ctx: dict[str, Any]) -> tuple[bool, str]:
    # Reserved contract slot: Layla's trained model is not registered in this
    # runtime yet. Flip this to check for the model artifact / service once it
    # is integrated, and it will supersede the empirical base rate.
    return False, "Layla's validated predictive model is not yet registered in this runtime."


def _run_predictive_model(_profile: dict[str, Any], _ctx: dict[str, Any]) -> dict[str, Any]:  # pragma: no cover
    raise NotImplementedError(
        "Publication-likelihood model not integrated. Register the trained model "
        "and implement per-trial probability here."
    )


def _llm_configured(_profile: dict[str, Any], ctx: dict[str, Any]) -> tuple[bool, str]:
    from llm import build_client

    if build_client(ctx["project_root"] / ".env") is not None:
        return True, "eligible"
    return False, "OpenAI API key not configured; using deterministic summary."


def _run_llm_summary(profile: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    return {"status": "ok", "result": summarize_with_llm(profile, ctx["capabilities"], ctx["project_root"])}


def _run_template_summary(profile: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
    return {"status": "ok", "result": template_summary(profile, ctx["capabilities"])}


def _run_r_module_adapter(_profile: dict[str, Any], _ctx: dict[str, Any]) -> dict[str, Any]:
    rscript = shutil.which("Rscript")
    return {
        "status": "not_configured",
        "result": {
            "rscript_available": bool(rscript),
            "contract": "Imported projects under modules/ are read-only references, not runtime module folders.",
        },
    }


# --------------------------------------------------------------------------- #
# The registry
# --------------------------------------------------------------------------- #

REGISTRY: tuple[ModuleSpec, ...] = (
    ModuleSpec(
        name="protocol_completeness",
        capability="protocol_completeness",
        runtime="python",
        precedence=10,
        description="Checks whether major ClinicalTrials.gov-style protocol sections are filled.",
        provenance="TIDE built-in",
        applies=_always,
        run=_run_protocol_completeness,
    ),
    ModuleSpec(
        name="historical_comparator",
        capability="historical_comparator",
        runtime="python",
        precedence=10,
        description="Reads archived domain CSVs from the imported publication-likelihood/timeliness study.",
        provenance="PubTime (R-parity verified)",
        applies=_domain_gate(("cancer", "covid", "cvd")),
        run=_run_historical_comparator,
        supported_domains=("cancer", "covid", "cvd"),
    ),
    # ---- publication_outlook: two providers, selected by precedence ---- #
    ModuleSpec(
        name="publication_model",
        capability="publication_outlook",
        runtime="python",
        precedence=100,
        description="Validated per-trial publication-likelihood model (Layla's project).",
        provenance="Predictive model (reserved)",
        applies=_predictive_model_available,
        run=_run_predictive_model,
    ),
    ModuleSpec(
        name="comparator_base_rate",
        capability="publication_outlook",
        runtime="python",
        precedence=10,
        description="Empirical publication/results-reporting rate from the matched historical cohort.",
        provenance="Historical comparator",
        applies=_always,
        run=_run_comparator_base_rate,
    ),
    # ---- narrative_summary: LLM interpretation, template fallback ---- #
    ModuleSpec(
        name="llm_summary",
        capability="narrative_summary",
        runtime="python",
        precedence=100,
        description="LLM (OpenAI Responses API) interpretation of the raw module outputs into plain language.",
        provenance="OpenAI Responses API",
        applies=_llm_configured,
        run=_run_llm_summary,
    ),
    ModuleSpec(
        name="template_summary",
        capability="narrative_summary",
        runtime="python",
        precedence=10,
        description="Deterministic plain-language summary when no LLM is configured.",
        provenance="TIDE built-in",
        applies=_always,
        run=_run_template_summary,
    ),
    ModuleSpec(
        name="r_module_adapter",
        capability="r_runtime_adapter",
        runtime="r",
        precedence=10,
        description="No R runtime modules are registered in this MVP.",
        provenance="TIDE built-in",
        applies=_always,
        run=_run_r_module_adapter,
    ),
)


# --------------------------------------------------------------------------- #
# Selection + execution
# --------------------------------------------------------------------------- #

def run_registry(profile: dict[str, Any], project_root: Path) -> dict[str, Any]:
    """Select and run modules. Returns modules (by name), a capability index,
    and a human-readable selection trace."""
    ctx: dict[str, Any] = {"project_root": project_root, "capabilities": {}}

    by_capability: "OrderedDict[str, list[ModuleSpec]]" = OrderedDict()
    for spec in REGISTRY:
        by_capability.setdefault(spec.capability, []).append(spec)

    modules: dict[str, Any] = {}
    selection: list[dict[str, Any]] = []

    for capability, specs in by_capability.items():
        evaluated = [(spec, *spec.applies(profile, ctx)) for spec in specs]
        eligible = sorted(
            (item for item in evaluated if item[1]),
            key=lambda item: item[0].precedence,
            reverse=True,
        )
        primary_spec = eligible[0][0] if eligible else None

        candidates: list[dict[str, Any]] = []
        for spec, ok, reason in evaluated:
            role = "skipped"
            status = "skipped"
            result: Any = None
            if spec is primary_spec:
                role = "primary"
                payload = spec.run(profile, ctx)
                status = payload.get("status", "ok")
                result = payload.get("result")
                ctx["capabilities"][capability] = result
            elif ok:
                role = "superseded"
                reason = f"eligible but superseded by '{primary_spec.name}' (higher precedence)."

            modules[spec.name] = {
                "status": status,
                "language": spec.runtime,
                "capability": capability,
                "precedence": spec.precedence,
                "selection": role,
                "description": spec.description,
                "provenance": spec.provenance,
                "reason": reason,
                "result": result,
            }
            candidates.append({"module": spec.name, "selection": role, "reason": reason})

        selection.append(
            {
                "capability": capability,
                "primary": primary_spec.name if primary_spec else None,
                "candidates": candidates,
            }
        )

    return {
        "modules": modules,
        "capabilities": ctx["capabilities"],
        "pipeline": {"selection": selection},
    }


def run_modules(profile: dict[str, Any], project_root: Path) -> dict[str, Any]:
    """Backward-compatible entry point: the module envelopes keyed by name."""
    return run_registry(profile, project_root)["modules"]