src/infrastructure/model/mock_model_service.py

"""
infrastructure/model/mock_model_service.py
───────────────────────────────────────────
MockModelService — deterministic fake inference for testing and local dev.

Returns physiologically plausible (but fake) BP values without loading
any model weights or requiring a GPU.

Use case: run the full ETL pipeline locally with USE_MOCK_MODEL=true.
"""
from __future__ import annotations

import asyncio
import hashlib
import time

import numpy as np

from src.domain.entities.prediction import BPPrediction
from src.domain.interfaces.services.model_service import ModelService
from src.shared.constants import MODEL_VERSION_MOCK
from src.shared.logger import get_logger

logger = get_logger(__name__)

# Physiologically plausible reference ranges for mock output
_SBP_BASE = 115.0   # Normal systolic
_DBP_BASE = 75.0    # Normal diastolic
_JITTER = 15.0      # ± range around base values


class MockModelService(ModelService):
    """
    Deterministic mock model service for testing and local development.

    Given the same ``ppg_signal_id``, it always returns the same BP values
    (deterministic via hash of the ID) — making tests reproducible.
    """

    def __init__(self) -> None:
        self._loaded = False

    # ── ModelService interface ────────────────────────────────────────────────

    async def load_model(self) -> None:
        """Simulates model loading — just flips a flag."""
        if self._loaded:
            return
        logger.info("MockModelService.load_model() — simulating 100ms load delay")
        await asyncio.sleep(0.1)  # simulate loading latency
        self._loaded = True
        logger.info("MockModelService ready (mock mode).")

    async def predict(
        self,
        ppg_signal_id: str,
        segments: np.ndarray,
    ) -> BPPrediction:
        """
        Return deterministic fake BP values based on the signal ID hash.

        The hash ensures the same ID always gives the same output, making
        unit tests with MockModelService fully reproducible.

        Args:
            ppg_signal_id: UUID of the source PPGSignal.
            segments:      Preprocessed signal segments (shape: N × window).

        Returns:
            BPPrediction with fake but physiologically valid SBP/DBP.
        """
        if not self._loaded:
            await self.load_model()

        start = time.perf_counter()

        # Deterministic jitter: hash the signal ID to get a repeatable seed
        hash_bytes = hashlib.sha256(ppg_signal_id.encode()).digest()
        seed = int.from_bytes(hash_bytes[:4], "big")

        # Determine number of segments (always at least 1)
        n_segments = segments.shape[0] if len(segments.shape) > 1 and segments.shape[0] > 0 else 1

        # 1. Generate fake segments of size 224 to simulate JIT features computation
        fake_ppg = np.zeros((n_segments, 224), dtype=np.float32)
        fake_ecg = np.zeros((n_segments, 224), dtype=np.float32)

        # Populate with deterministic signal + noise + optional plateau
        for i in range(n_segments):
            seg_seed = (seed + i) & 0xffffffff
            rng_seg = np.random.default_rng(seg_seed)
            t = np.linspace(0, 2 * np.pi, 224)
            # i%3 == 0: noisy (high entropy), others: clean (low entropy)
            noise_lvl = 0.5 if (i % 3 == 0) else 0.05
            fake_ppg[i] = np.sin(t) + rng_seg.normal(0, noise_lvl, 224)
            fake_ecg[i] = np.sin(t * 2) + rng_seg.normal(0, noise_lvl, 224)

            # i%4 == 0: inject flat lines (plateau)
            if i % 4 == 0:
                fake_ppg[i, 50:70] = 0.5
                fake_ecg[i, 100:120] = -0.5

        # 2. Generate segment predictions
        sbp_preds = []
        dbp_preds = []
        for i in range(n_segments):
            seg_seed = (seed + i) & 0xffffffff
            rng_seg = np.random.default_rng(seg_seed)
            s = float(_SBP_BASE + rng_seg.uniform(-_JITTER, _JITTER))
            d = float(_DBP_BASE + rng_seg.uniform(-_JITTER / 2, _JITTER / 2))
            
            # Add variance to noisy windows
            if i % 3 == 0:
                s += rng_seg.uniform(-25, 25)
                d += rng_seg.uniform(-12, 12)
                # Clamp within physiological bounds
                s = max(80.0, min(200.0, s))
                d = max(40.0, min(120.0, d))
            
            if s <= d:
                s = d + 30.0

            sbp_preds.append(s)
            dbp_preds.append(d)

        sbp_preds_arr = np.array(sbp_preds)
        dbp_preds_arr = np.array(dbp_preds)

        # 3. Run Simulated Annealing (1000 steps)
        from src.infrastructure.processing.sa_helpers import run_simulated_annealing

        sa_result = run_simulated_annealing(
            ppg_segments=fake_ppg,
            ecg_segments=fake_ecg,
            sbp_preds=sbp_preds_arr,
            dbp_preds=dbp_preds_arr,
            n_steps=1000,
        )

        clean_indices = sa_result["clean_indices"]
        if len(clean_indices) > 0:
            sbp = float(np.mean(sbp_preds_arr[clean_indices]))
            dbp = float(np.mean(dbp_preds_arr[clean_indices]))
        else:
            sbp = float(np.mean(sbp_preds_arr))
            dbp = float(np.mean(dbp_preds_arr))

        # Assemble SA log dict
        sa_log = {
            "optimal_lo": sa_result["optimal_lo"],
            "optimal_hi": sa_result["optimal_hi"],
            "optimal_max_plateau": sa_result["optimal_max_plateau"],
            "best_loss": sa_result["best_loss"],
            "initial_loss": sa_result["initial_loss"],
            "n_total_segments": sa_result["n_total_segments"],
            "n_clean_segments": sa_result["n_clean_segments"],
            "yield_rate": sa_result["yield_rate"],
            "history": sa_result["history"],
        }

        # Simulate short inference time
        await asyncio.sleep(0.05)
        elapsed_ms = (time.perf_counter() - start) * 1000

        logger.info(
            "MockModelService.predict() — signal_id=%s segments=%d (clean=%d) "
            "SBP=%.1f DBP=%.1f (%.1f ms)",
            ppg_signal_id,
            n_segments,
            len(clean_indices),
            sbp,
            dbp,
            elapsed_ms,
        )

        return BPPrediction(
            ppg_signal_id=ppg_signal_id,
            predicted_sbp=round(sbp, 1),
            predicted_dbp=round(dbp, 1),
            predicted_ecg=fake_ecg.tolist(),
            model_version=self.model_version,
            inference_time_ms=round(elapsed_ms, 2),
            sa_log=sa_log,
        )

    def is_loaded(self) -> bool:
        return self._loaded

    @property
    def model_version(self) -> str:
        return MODEL_VERSION_MOCK