Main_reasoning.py

"""
DFOD reasoning with weighted multimodal memory, unified score spaces,
optional metric fusion, and calibration-aware rejection.
"""

from __future__ import annotations

from typing import Any, Dict, List, Tuple

import numpy as np

from dfod_config import normalize_config
from Baseline_Reasoning.calibration import resolve_calibration
from Baseline_Reasoning.formatting import format_results
from Baseline_Reasoning.memory import build_memory_bank
from Baseline_Reasoning.rejection import apply_unknown_rejection
from Baseline_Reasoning.serialization import MemoryBundle, load_memory_bundle
from Baseline_Reasoning.similarity import compute_metric_score_details, fuse_metric_probabilities


def _finalize_scores(
    embeddings: np.ndarray,
    memory_bank: Dict[str, Any],
    config: Dict[str, Any],
) -> Tuple[np.ndarray, Dict[str, Dict[str, np.ndarray]], Dict[str, Any]]:
    metrics = config["reasoning"]["metrics"]
    details = compute_metric_score_details(
        embeddings=embeddings,
        memory_bank=memory_bank,
        metrics=metrics,
    )

    reasoning_cfg = config["reasoning"]
    primary_metric = str(reasoning_cfg["primary_metric"]).strip().lower()
    if reasoning_cfg.get("enable_metric_fusion", False) and len(details) > 1:
        final_scores = fuse_metric_probabilities(
            metric_details=details,
            metric_weights=reasoning_cfg.get("metric_weights", {}),
            temperature=float(reasoning_cfg.get("softmax_temperature", 1.0)),
        )
        provenance = {
            "score_metric": "fused",
            "primary_metric": primary_metric,
            "best_mode_indices": details[primary_metric]["best_mode_indices"],
        }
    else:
        final_scores = details[primary_metric]["class_scores"]
        provenance = {
            "score_metric": primary_metric,
            "primary_metric": primary_metric,
            "best_mode_indices": details[primary_metric]["best_mode_indices"],
        }
    return final_scores.astype(np.float32, copy=False), details, provenance


def run_reasoning(
    embeddings: np.ndarray,
    memory_path: str,
    config: Dict[str, object] | None = None,
) -> List[Dict[str, object]]:
    formatted, _, _, _ = run_reasoning_detailed(
        embeddings=embeddings,
        memory_path=memory_path,
        config=config,
    )
    return formatted


def _run_reasoning_with_bundle(
    embeddings: np.ndarray,
    memory_bundle: MemoryBundle,
    config: Dict[str, object] | None = None,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]], np.ndarray, List[str]]:
    cfg = normalize_config(config)
    E = np.asarray(embeddings, dtype=np.float32)
    if E.ndim != 2:
        raise ValueError(f"embeddings must be 2D, got shape={E.shape}")
    if E.shape[0] == 0:
        return [], [], np.zeros((0, 0), dtype=np.float32), []

    memory_bank = build_memory_bank(memory_bundle, config=cfg)
    final_scores, metric_details, provenance = _finalize_scores(E, memory_bank, cfg)
    class_names = list(memory_bank["class_names"])
    calibration = resolve_calibration(config=cfg, memory_bundle=memory_bundle)

    diagnostics = {
        "score_metric": provenance["score_metric"],
        "primary_metric": provenance["primary_metric"],
    }
    decisions = apply_unknown_rejection(
        scores=final_scores,
        class_names=class_names,
        config=cfg,
        calibration=calibration,
        diagnostics=diagnostics if cfg["diagnostics"].get("include", True) else None,
    )

    if decisions and cfg["diagnostics"].get("include", True):
        best_modes = provenance["best_mode_indices"]
        for row_idx, decision in enumerate(decisions):
            top1_idx = int(decision.get("top1_idx", -1))
            if top1_idx >= 0 and best_modes.shape[1] > top1_idx:
                decision["mode_index"] = int(best_modes[row_idx, top1_idx])
                decision["mode_id"] = int(memory_bank["mode_ids"][decision["mode_index"]])
            decision["metric_scores"] = {
                metric: float(detail["class_scores"][row_idx, top1_idx]) if top1_idx >= 0 else float("-inf")
                for metric, detail in metric_details.items()
            }

    formatted = format_results(decisions=decisions)
    return formatted, decisions, final_scores, class_names


def run_reasoning_bundle_detailed(
    embeddings: np.ndarray,
    memory_bundle: MemoryBundle,
    config: Dict[str, object] | None = None,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]], np.ndarray, List[str]]:
    return _run_reasoning_with_bundle(
        embeddings=embeddings,
        memory_bundle=memory_bundle,
        config=config,
    )


def run_reasoning_bundle(
    embeddings: np.ndarray,
    memory_bundle: MemoryBundle,
    config: Dict[str, object] | None = None,
) -> List[Dict[str, object]]:
    formatted, _, _, _ = run_reasoning_bundle_detailed(
        embeddings=embeddings,
        memory_bundle=memory_bundle,
        config=config,
    )
    return formatted


def run_reasoning_detailed(
    embeddings: np.ndarray,
    memory_path: str,
    config: Dict[str, object] | None = None,
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]], np.ndarray, List[str]]:
    bundle = load_memory_bundle(memory_path)
    return _run_reasoning_with_bundle(
        embeddings=embeddings,
        memory_bundle=bundle,
        config=config,
    )