src/semantic/regulators.py

"""
Regulator label classification.

classify_regulator_label() types a single label from a TF enrichment result
(e.g. CollecTRI output) into one of:
  - gene                  — individual HGNC-like gene symbol
  - TF_complex_or_family  — multi-gene complex or family reported as one label
  - unresolved_label      — could not be matched

Resolution order
----------------
1. Check resources/semantic/regulator_overrides.yaml (highest priority).
2. Check gene_symbols.normalize_gene_symbol() (seed set / future HGNC cache).
3. Return unresolved_label with a warning.

When to call this function
--------------------------
Call from result_annotation.annotate_result_entities() when analysis_role
is "regulator" (i.e. CollecTRI / TF enrichment results).
Do not call for pathway or gene set results — use gene_sets.py instead.
"""

from __future__ import annotations

import warnings
from functools import lru_cache
from pathlib import Path
from typing import Any

import yaml

from .gene_symbols import normalize_gene_symbol
from .schemas import EntityAnnotation

_RESOURCES_DIR = Path(__file__).parent.parent.parent / "resources" / "semantic"
_OVERRIDE_PATH = _RESOURCES_DIR / "regulator_overrides.yaml"


@lru_cache(maxsize=1)
def _load_regulator_overrides() -> dict[str, Any]:
    """
    Load and cache regulator_overrides.yaml.

    Keys are uppercased at load time so lookups are case-insensitive.
    Returns an empty dict if the file is missing (non-fatal).
    """
    if not _OVERRIDE_PATH.exists():
        warnings.warn(
            f"regulator_overrides.yaml not found at {_OVERRIDE_PATH}. "
            "Regulator override lookup will be skipped.",
            stacklevel=2,
        )
        return {}
    try:
        with open(_OVERRIDE_PATH, encoding="utf-8") as fh:
            raw = yaml.safe_load(fh) or {}
        return {k.upper(): v for k, v in raw.items()}
    except Exception as exc:
        warnings.warn(f"Failed to load regulator_overrides.yaml: {exc}", stacklevel=2)
        return {}


def classify_regulator_label(
    label: str,
    species: str = "human",
) -> EntityAnnotation:
    """
    Classify a single TF enrichment result label.

    Parameters
    ----------
    label:
        The regulator label as it appears in enrichment output (e.g. 'AP1',
        'HIF1A').  Case-insensitive.
    species:
        Target species for gene symbol resolution.  Default "human".

    Returns
    -------
    EntityAnnotation with entity_type set to one of:
        "gene"                 — individual HGNC-like gene
        "TF_complex_or_family" — complex or family label
        "unresolved_label"     — could not be matched
    """
    upper = label.strip().upper()
    overrides = _load_regulator_overrides()

    # ── 1. Override YAML (highest priority) ─────────────────────────────
    if upper in overrides:
        entry = overrides[upper]
        entity_type = entry.get("entity_type", "TF_complex_or_family")
        return EntityAnnotation(
            label=label,
            entity_type=entity_type,
            analysis_role="regulator",
            is_gene_symbol=False,
            normalized_symbol=None,
            normalized_name=entry.get("normalized_name"),
            members_or_related_genes=entry.get("members_or_related_genes", []),
            result_type=None,
            confidence="high",
            warnings=[],
        )

    # ── 2. Gene symbol seed set ──────────────────────────────────────────
    gene_result = normalize_gene_symbol(label, species=species)
    if gene_result["resolved"]:
        return EntityAnnotation(
            label=label,
            entity_type="gene",
            analysis_role="regulator",
            is_gene_symbol=True,
            normalized_symbol=gene_result["symbol"],
            normalized_name=None,
            members_or_related_genes=[],
            result_type=None,
            confidence=gene_result["confidence"],
            warnings=[],
        )

    # ── 3. Unresolved ────────────────────────────────────────────────────
    return EntityAnnotation(
        label=label,
        entity_type="unresolved_label",
        analysis_role="regulator",
        is_gene_symbol=False,
        normalized_symbol=None,
        normalized_name=None,
        members_or_related_genes=[],
        result_type=None,
        confidence="unresolved",
        warnings=[
            f"Label '{label}' could not be resolved to a gene symbol or known "
            "regulator complex. Check regulator_overrides.yaml or extend the "
            "gene symbol seed set in gene_symbols.py."
        ],
    )