app/embed_experiment.py

"""Domain-embedding experiment: probe models → re-index → eval hybrid fusion.

Compares BM25-only vs k-NN vs hybrid RRF / normalization on
``data/pediatric_eval.json`` after swapping the SentenceTransformer model.

Usage (from ``backend/``, OpenSearch up, index already populated once)::

    # Probe dims without touching the index
    PYTHONPATH=. python -m app.embed_experiment --probe

    # Eval current configured model (no re-embed)
    PYTHONPATH=. python -m app.embed_experiment --eval-only

    # Swap one model, recreate index if dim changes, force re-embed, eval
    PYTHONPATH=. python -m app.embed_experiment \\
        --model BAAI/bge-base-en-v1.5 --reindex

    # Sweep shortlist (slow — downloads + full re-embed per model)
    PYTHONPATH=. python -m app.embed_experiment --compare --reindex

    # Write JSON report
    PYTHONPATH=. python -m app.embed_experiment --eval-only -o /tmp/embed_eval.json

Switch production only if hybrid beats the current BM25 baseline
(top-3 ≈ 0.64 / MRR ≈ 0.62). Then set EMBEDDING_MODEL / EMBEDDING_DIM /
FUSION_TECHNIQUE in ``.env`` and force-refresh once.
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from .config import settings
from .eval_retrieval import run as run_eval
from .opensearch_client import ensure_index, get_client, index_mapping

logger = logging.getLogger(__name__)

# Refuse accidental wipe of the live Aiven codes index.
_PROD_HOST_MARKERS = ("aivencloud.com",)
_PROD_INDEX_NAMES = frozenset({"ohip_codes"})

# Candidate models that load cleanly via sentence-transformers.
# expected_dim is a hint; --probe / load measures the real size.
CANDIDATES: list[dict[str, Any]] = [
    {
        "id": "Alibaba-NLP/gte-multilingual-base",
        "expected_dim": 768,
        "notes": "Current default (general multilingual)",
        "tier": "baseline",
    },
    {
        "id": "BAAI/bge-base-en-v1.5",
        "expected_dim": 768,
        "notes": "Strong general English bi-encoder",
        "tier": "general",
    },
    {
        "id": "BAAI/bge-small-en-v1.5",
        "expected_dim": 384,
        "notes": "Faster / smaller; requires index dim change",
        "tier": "general",
    },
    {
        "id": "NeuML/pubmedbert-base-embeddings",
        "expected_dim": 768,
        "notes": "PubMedBERT medical embeddings",
        "tier": "medical",
    },
    {
        "id": "pritamdeka/S-PubMedBert-MS-MARCO",
        "expected_dim": 768,
        "notes": "PubMedBERT fine-tuned for retrieval (MS MARCO)",
        "tier": "medical",
    },
    {
        "id": "sentence-transformers/all-MiniLM-L6-v2",
        "expected_dim": 384,
        "notes": "Cheap sanity check (not domain-tuned)",
        "tier": "sanity",
    },
]

# Historical BM25-only bar on the 50-case pediatric set (README).
# Promote hybrid only if pooled metrics clear this bar AND beat this-run BM25,
# with no specialty top3 drop > MAX_SPECIALTY_DROP vs BM25 on the same run.
BM25_BASELINE = {"top3": 0.64, "mrr": 0.62}
MAX_SPECIALTY_DROP = 0.05


@dataclass
class ProbeResult:
    model_id: str
    dim: int
    expected_dim: int | None
    ok: bool
    error: str | None = None


def _clear_embed_cache() -> None:
    from . import embeddings  # noqa: PLC0415

    embeddings._model.cache_clear()


def configure_embedding_model(model_id: str, dim: int) -> None:
    """Mutate runtime settings and drop the cached SentenceTransformer."""
    settings.embedding_model = model_id
    settings.embedding_dim = dim
    _clear_embed_cache()
    logger.info("Configured EMBEDDING_MODEL=%s EMBEDDING_DIM=%s", model_id, dim)


def probe_model(model_id: str, expected_dim: int | None = None) -> ProbeResult:
    """Load a model once and measure embedding dimension."""
    try:
        from sentence_transformers import SentenceTransformer  # noqa: PLC0415

        logger.info("Probing %s …", model_id)
        model = SentenceTransformer(model_id, trust_remote_code=True)
        vec = model.encode(
            ["OHIP paediatric consultation fee schedule"],
            normalize_embeddings=True,
            convert_to_numpy=True,
        )[0]
        dim = int(vec.shape[-1])
        return ProbeResult(
            model_id=model_id,
            dim=dim,
            expected_dim=expected_dim,
            ok=expected_dim is None or dim == expected_dim,
        )
    except Exception as exc:  # noqa: BLE001
        return ProbeResult(
            model_id=model_id,
            dim=0,
            expected_dim=expected_dim,
            ok=False,
            error=str(exc),
        )


def index_vector_dim(client=None) -> int | None:
    """Return current ``code_vector`` dimension, or None if index missing."""
    client = client or get_client()
    idx = settings.opensearch_index
    if not client.indices.exists(index=idx):
        return None
    mapping = client.indices.get_mapping(index=idx)
    props = mapping[idx]["mappings"]["properties"]
    vec = props.get("code_vector") or {}
    return vec.get("dimension")


def assert_reindex_allowed(*, allow_prod: bool = False) -> None:
    """Block --reindex against production Aiven ``ohip_codes`` unless forced."""
    url = (settings.opensearch_url or "").lower()
    idx = settings.opensearch_index or ""
    is_prod_host = any(m in url for m in _PROD_HOST_MARKERS)
    is_prod_index = idx in _PROD_INDEX_NAMES
    if is_prod_host and is_prod_index and not allow_prod:
        raise SystemExit(
            "REFUSING --reindex: OPENSEARCH_URL looks like production Aiven "
            f"and OPENSEARCH_INDEX={idx!r}.\n"
            "Use a staging index (e.g. OPENSEARCH_INDEX=ohip_codes_staging) "
            "or local Docker OpenSearch.\n"
            "To override intentionally: pass --allow-prod-reindex "
            "(destructive; recreates live codes index)."
        )


def recreate_codes_index(*, confirm: bool = True) -> None:
    """Delete and recreate ``ohip_codes`` with the current embedding_dim."""
    client = get_client()
    idx = settings.opensearch_index
    if client.indices.exists(index=idx):
        if confirm:
            logger.warning("Deleting index '%s' to apply dim=%s", idx, settings.embedding_dim)
        client.indices.delete(index=idx)
    # ensure_index creates mapping from settings.embedding_dim
    ensure_index(client)
    # Belt-and-suspenders: if index somehow existed with wrong dim, recreate.
    current = index_vector_dim(client)
    if current != settings.embedding_dim:
        client.indices.delete(index=idx)
        client.indices.create(index=idx, body=index_mapping())
        logger.info("Recreated '%s' with dimension %s", idx, settings.embedding_dim)


def force_reembed() -> dict:
    from .ingestion import ingest  # noqa: PLC0415

    summary = ingest(force=True)
    logger.info("Force re-embed complete: %s", summary)
    return summary


def _best_hybrid(results: dict[str, dict]) -> tuple[str, dict] | None:
    hybrids = {k: v for k, v in results.items() if k.startswith("hybrid_")}
    if not hybrids:
        return None
    best_k = max(hybrids, key=lambda k: (hybrids[k]["top3"], hybrids[k]["mrr"]))
    return best_k, hybrids[best_k]


def _specialty_drop_ok(bm25: dict, hybrid: dict) -> tuple[bool, list[str]]:
    """True if no specialty loses more than MAX_SPECIALTY_DROP top3 vs BM25."""
    notes: list[str] = []
    bm25_specs = bm25.get("by_specialty") or {}
    hyb_specs = hybrid.get("by_specialty") or {}
    ok = True
    for spec, b in bm25_specs.items():
        h = hyb_specs.get(spec) or {}
        b3 = float(b.get("top3") or 0)
        h3 = float(h.get("top3") or 0)
        drop = b3 - h3
        notes.append(
            f"  specialty {spec}: bm25_top3={b3} hybrid_top3={h3} drop={drop:.3f}"
        )
        if drop > MAX_SPECIALTY_DROP + 1e-9:
            ok = False
    return ok, notes


def recommend(eval_out: dict) -> str:
    results = eval_out["results"]
    bm25 = results.get("bm25_only") or {}
    knn = results.get("knn_only") or {}
    best = _best_hybrid(results)
    sets = eval_out.get("sets") or {}
    lines = [
        f"Model: {settings.embedding_model} (dim={settings.embedding_dim})",
        f"Cases: n={eval_out.get('n_cases')} sets={sets}"
        + (f" smoke={eval_out['smoke']}" if eval_out.get("smoke") else ""),
        f"Index dim: {eval_out.get('index_dim')}  vectors_ok={eval_out.get('vectors_ok')}",
        f"BM25-only:  top3={bm25.get('top3')}  mrr={bm25.get('mrr')}",
        f"k-NN-only:  top3={knn.get('top3')}  mrr={knn.get('mrr')}",
    ]
    for note in eval_out.get("notes") or []:
        lines.append(f"NOTE: {note}")
    if eval_out.get("vectors_ok") is False:
        lines.append(
            "KEEP: FUSION_TECHNIQUE=bm25 — k-NN/hybrid skipped "
            "(embed dim ≠ index dim). Re-embed on staging before promoting hybrid."
        )
        return "\n".join(lines)
    if best:
        name, m = best
        lines.append(f"Best hybrid: {name}  top3={m['top3']}  mrr={m['mrr']}")
        drop_ok, drop_notes = _specialty_drop_ok(bm25, m)
        lines.extend(drop_notes)
        beats_bm25 = (m["top3"], m["mrr"]) > (
            bm25.get("top3", 0),
            bm25.get("mrr", 0),
        )
        beats_hist = (m["top3"], m["mrr"]) >= (
            BM25_BASELINE["top3"],
            BM25_BASELINE["mrr"],
        )
        if beats_bm25 and beats_hist and drop_ok:
            lines.append(
                "RECOMMEND: switch production to this model + "
                f"FUSION_TECHNIQUE matching '{name}' "
                "(hybrid beats BM25, clears historical bar, specialty drops OK)."
            )
        elif beats_bm25 and not drop_ok:
            lines.append(
                "KEEP: FUSION_TECHNIQUE=bm25 — hybrid beats pooled BM25 but "
                f"drops a specialty by more than {MAX_SPECIALTY_DROP:.0%} top3."
            )
        elif beats_bm25:
            lines.append(
                "CAUTION: hybrid beats BM25 on this run but is still below the "
                f"historical bar (top3≥{BM25_BASELINE['top3']}, "
                f"mrr≥{BM25_BASELINE['mrr']}). Keep investigating."
            )
        else:
            lines.append(
                "KEEP: FUSION_TECHNIQUE=bm25 — hybrid does not beat BM25-only yet."
            )
    return "\n".join(lines)


def print_table(title: str, eval_out: dict) -> None:
    print(f"\n{title}")
    print(f"n_cases={eval_out['n_cases']}  model={settings.embedding_model}  "
          f"dim={settings.embedding_dim}")
    print(f"{'config':<28} {'top1':>6} {'top3':>6} {'mrr':>6}")
    print("-" * 50)
    for cfg, m in eval_out["results"].items():
        print(f"{cfg:<28} {m['top1']:>6} {m['top3']:>6} {m['mrr']:>6}")
    print()
    print(recommend(eval_out))
    print()


def run_for_model(
    model_id: str,
    *,
    reindex: bool,
    expected_dim: int | None = None,
    eval_sets: list[str] | None = None,
    smoke: int | None = None,
    bm25_only: bool = False,
    allow_prod_reindex: bool = False,
) -> dict[str, Any]:
    probe = probe_model(model_id, expected_dim)
    if probe.error or probe.dim <= 0:
        return {
            "model_id": model_id,
            "probe": asdict(probe),
            "error": probe.error or "probe failed",
        }

    configure_embedding_model(model_id, probe.dim)
    ingest_summary = None
    if reindex:
        assert_reindex_allowed(allow_prod=allow_prod_reindex)
        current_dim = index_vector_dim()
        if current_dim != probe.dim:
            recreate_codes_index()
        else:
            # Same dim: still force re-embed so vectors match the new model.
            logger.info(
                "Index dim already %s — force re-embedding without delete",
                current_dim,
            )
            ensure_index(get_client())
        ingest_summary = force_reembed()
    else:
        current_dim = index_vector_dim()
        if current_dim is not None and current_dim != probe.dim:
            raise RuntimeError(
                f"Index dim is {current_dim} but model produces {probe.dim}. "
                "Re-run with --reindex."
            )

    eval_out = run_eval(sets=eval_sets, smoke=smoke, bm25_only=bm25_only)
    return {
        "model_id": model_id,
        "probe": asdict(probe),
        "ingest": ingest_summary,
        "eval": eval_out,
        "recommendation": recommend(eval_out),
    }


def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Domain embedding experiment + hybrid fusion re-eval"
    )
    p.add_argument(
        "--probe",
        action="store_true",
        help="Probe candidate model dimensions (no index changes)",
    )
    p.add_argument(
        "--eval-only",
        action="store_true",
        help="Run fusion eval with the currently configured embedding model",
    )
    p.add_argument(
        "--model",
        type=str,
        default=None,
        help="HuggingFace / local SentenceTransformer id to evaluate",
    )
    p.add_argument(
        "--compare",
        action="store_true",
        help="Sweep shortlisted candidates (use with --reindex)",
    )
    p.add_argument(
        "--tier",
        choices=("baseline", "general", "medical", "sanity", "all"),
        default="all",
        help="Filter --probe / --compare shortlist (default: all)",
    )
    p.add_argument(
        "--reindex",
        action="store_true",
        help="Recreate index if dim changes and force re-embed every code",
    )
    p.add_argument(
        "--allow-prod-reindex",
        action="store_true",
        help="Override safety check that blocks --reindex on Aiven ohip_codes",
    )
    p.add_argument(
        "--sets",
        type=str,
        default="pediatric,family",
        help="Comma-separated eval sets (default: pediatric,family)",
    )
    p.add_argument(
        "--smoke",
        type=int,
        default=None,
        help="Evaluate only N cases (round-robin across sets) for cheap iteration",
    )
    p.add_argument(
        "--bm25-only",
        action="store_true",
        help="Skip k-NN / hybrid eval (BM25 metrics only)",
    )
    p.add_argument(
        "-o",
        "--output",
        type=str,
        default=None,
        help="Write full JSON report to this path",
    )
    p.add_argument("-v", "--verbose", action="store_true")
    return p.parse_args(argv)


def _eval_sets_from_args(args: argparse.Namespace) -> list[str]:
    return [s.strip() for s in (args.sets or "").split(",") if s.strip()]


def _filtered_candidates(tier: str) -> list[dict[str, Any]]:
    if tier == "all":
        return list(CANDIDATES)
    return [c for c in CANDIDATES if c["tier"] == tier]


def main(argv: list[str] | None = None) -> int:
    args = _parse_args(argv)
    logging.basicConfig(
        level=logging.INFO if args.verbose else logging.WARNING,
        format="%(levelname)s %(name)s: %(message)s",
    )

    eval_sets = _eval_sets_from_args(args)
    report: dict[str, Any] = {
        "bm25_historical_baseline": BM25_BASELINE,
        "max_specialty_drop": MAX_SPECIALTY_DROP,
        "eval_sets": eval_sets,
        "smoke": args.smoke,
        "current_settings": {
            "embedding_model": settings.embedding_model,
            "embedding_dim": settings.embedding_dim,
            "fusion_technique": settings.fusion_technique,
            "index_dim": index_vector_dim(),
        },
    }

    if args.probe:
        rows = []
        print(f"\n{'model':<45} {'exp':>5} {'got':>5} {'ok':>4}")
        print("-" * 65)
        for c in _filtered_candidates(args.tier):
            pr = probe_model(c["id"], c.get("expected_dim"))
            rows.append(asdict(pr) | {"notes": c.get("notes"), "tier": c.get("tier")})
            status = "yes" if pr.ok and not pr.error else "NO"
            print(
                f"{c['id']:<45} {c.get('expected_dim') or '-':>5} "
                f"{pr.dim or '-':>5} {status:>4}"
            )
            if pr.error:
                print(f"  error: {pr.error[:200]}")
        report["probes"] = rows
        if args.output:
            Path(args.output).write_text(json.dumps(report, indent=2))
            print(f"\nWrote {args.output}")
        return 0

    if args.reindex:
        try:
            assert_reindex_allowed(allow_prod=args.allow_prod_reindex)
        except SystemExit as exc:
            print(str(exc), file=sys.stderr)
            return 2

    if args.eval_only:
        out = run_eval(sets=eval_sets, smoke=args.smoke, bm25_only=args.bm25_only)
        print_table("Retrieval eval (current model)", out)
        report["runs"] = [
            {
                "model_id": settings.embedding_model,
                "eval": out,
                "recommendation": recommend(out),
            }
        ]
        if args.output:
            Path(args.output).write_text(json.dumps(report, indent=2))
            print(f"Wrote {args.output}")
        return 0

    if args.compare:
        if not args.reindex:
            print(
                "WARNING: --compare without --reindex only works if every "
                "model shares the current index dim and vectors already match "
                "that model (unlikely). Prefer --compare --reindex.",
                file=sys.stderr,
            )
        runs = []
        for c in _filtered_candidates(args.tier):
            print(f"\n=== {c['id']} ({c['tier']}) ===")
            try:
                result = run_for_model(
                    c["id"],
                    reindex=args.reindex,
                    expected_dim=c.get("expected_dim"),
                    eval_sets=eval_sets,
                    smoke=args.smoke,
                    bm25_only=args.bm25_only,
                    allow_prod_reindex=args.allow_prod_reindex,
                )
            except Exception as exc:  # noqa: BLE001
                result = {"model_id": c["id"], "error": str(exc)}
                logger.exception("Failed on %s", c["id"])
            runs.append(result)
            if "eval" in result:
                print_table(f"Eval — {c['id']}", result["eval"])
            elif result.get("error"):
                print(f"ERROR: {result['error']}")

        # Summary leaderboard
        print("\n=== Leaderboard (by best hybrid top3, then mrr) ===")
        print(f"{'model':<45} {'bm25_t3':>8} {'best_hyb':>18} {'hyb_t3':>7} {'hyb_mrr':>7}")
        print("-" * 90)
        board = []
        for r in runs:
            if "eval" not in r:
                continue
            res = r["eval"]["results"]
            bm25_t3 = res.get("bm25_only", {}).get("top3")
            best = _best_hybrid(res)
            if not best:
                continue
            name, m = best
            board.append((r["model_id"], bm25_t3, name, m["top3"], m["mrr"]))
        board.sort(key=lambda t: (t[3], t[4]), reverse=True)
        for model_id, bm25_t3, name, t3, mrr in board:
            short = name.replace("hybrid_", "")
            print(
                f"{model_id:<45} {bm25_t3:>8} {short:>18} {t3:>7} {mrr:>7}"
            )

        report["runs"] = runs
        if args.output:
            Path(args.output).write_text(json.dumps(report, indent=2))
            print(f"\nWrote {args.output}")
        return 0

    if args.model:
        result = run_for_model(
            args.model,
            reindex=args.reindex,
            eval_sets=eval_sets,
            smoke=args.smoke,
            bm25_only=args.bm25_only,
            allow_prod_reindex=args.allow_prod_reindex,
        )
        if result.get("error"):
            print(f"ERROR: {result['error']}", file=sys.stderr)
            return 1
        print_table(f"Eval — {args.model}", result["eval"])
        report["runs"] = [result]
        if args.output:
            Path(args.output).write_text(json.dumps(report, indent=2))
            print(f"Wrote {args.output}")
        print(
            "\nTo keep this model in production, set in .env:\n"
            f"  EMBEDDING_MODEL={settings.embedding_model}\n"
            f"  EMBEDDING_DIM={settings.embedding_dim}\n"
            "  FUSION_TECHNIQUE=rrf   # or normalization / bm25 per recommendation\n"
            "then: curl -X POST 'http://localhost:8080/admin/refresh?force=true'\n"
            "Also recreate ohip_feedback / ohip_remittance if EMBEDDING_DIM changed."
        )
        return 0

    print(
        "Specify one of: --probe | --eval-only | --model ID | --compare\n"
        "See: python -m app.embed_experiment -h",
        file=sys.stderr,
    )
    return 2


if __name__ == "__main__":
    raise SystemExit(main())