tests/test_dynamics.py

"""
tests/test_dynamics.py — Phase 9 LagPredictor dynamics model tests
===================================================================
Covers CLAUDE.md §test_dynamics.py requirements:

  ✓ Basic forward pass — model accepts valid (16,) input without error
  ✓ Output shape — forward returns (batch, 1) for batched input
  ✓ Output range — predict_single returns float in (0.0, 1.0) [Sigmoid]
  ✓ MSE decreases over 10 gradient batches — model learns on synthetic data
  ✓ build_input_vector produces a (16,) tensor from valid obs+action pair
  ✓ store_transition + buffer_size work correctly
  ✓ train_step returns None when buffer is below BATCH_SIZE
"""
from __future__ import annotations

import random

import pytest
import torch

from dynamics_model import (
    LagPredictor,
    MultiObsPredictor,
    build_input_vector,
    build_full_obs_target_vector,
    BATCH_SIZE,
    INPUT_DIM,
    OUTPUT_DIM,
    MULTI_OBS_OUTPUT_DIM,
    MULTI_OBS_BATCH_SIZE,
)
from unified_gateway import AEPOAction


# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def model() -> LagPredictor:
    """Fresh LagPredictor with empty replay buffer."""
    return LagPredictor()


def _make_obs_normalized(kafka_lag: float = 0.1) -> dict[str, float]:
    """Build a minimal valid normalized obs dict."""
    return {
        "transaction_type":        0.0,
        "risk_score":              0.2,
        "adversary_threat_level":  0.1,
        "system_entropy":          0.3,
        "kafka_lag":               kafka_lag,
        "api_latency":             0.05,
        "rolling_p99":             0.04,
        "db_connection_pool":      0.5,
        "bank_api_status":         0.0,
        "merchant_tier":           0.0,
    }


def _make_action() -> AEPOAction:
    """Build a safe default action."""
    return AEPOAction(
        risk_decision=0,
        crypto_verify=1,
        infra_routing=0,
        db_retry_policy=0,
        settlement_policy=0,
        app_priority=2,
    )


# ---------------------------------------------------------------------------
# Test 1 — build_input_vector produces (16,) float32 tensor
# ---------------------------------------------------------------------------

def test_build_input_vector_shape_and_dtype():
    """build_input_vector must return a float32 Tensor of shape (16,)."""
    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)

    assert isinstance(x, torch.Tensor), "build_input_vector must return a torch.Tensor"
    assert x.shape == (INPUT_DIM,), f"Expected shape ({INPUT_DIM},), got {x.shape}"
    assert x.dtype == torch.float32, f"Expected float32, got {x.dtype}"


def test_build_input_vector_values_in_range():
    """All 16 input values must be in [0.0, 1.0] after normalization."""
    obs = _make_obs_normalized(kafka_lag=0.95)
    action = AEPOAction(
        risk_decision=2,    # max value for 3-choice field
        crypto_verify=1,
        infra_routing=2,    # max value for 3-choice field
        db_retry_policy=1,
        settlement_policy=1,
        app_priority=2,     # max value for 3-choice field
    )
    x = build_input_vector(obs, action)

    assert float(x.min()) >= 0.0, f"Min value {x.min()} < 0.0"
    assert float(x.max()) <= 1.0, f"Max value {x.max()} > 1.0"


# ---------------------------------------------------------------------------
# Test 2 — forward pass shape
# ---------------------------------------------------------------------------

def test_forward_single_input(model: LagPredictor):
    """Forward pass on a (1, 16) batch must return shape (1, 1)."""
    x = torch.rand(1, INPUT_DIM)
    out = model(x)
    assert out.shape == (1, OUTPUT_DIM), f"Expected (1, {OUTPUT_DIM}), got {out.shape}"


def test_forward_batched_input(model: LagPredictor):
    """Forward pass on a (32, 16) batch must return shape (32, 1)."""
    x = torch.rand(32, INPUT_DIM)
    out = model(x)
    assert out.shape == (32, OUTPUT_DIM), (
        f"Expected (32, {OUTPUT_DIM}), got {out.shape}"
    )


# ---------------------------------------------------------------------------
# Test 3 — output range (Sigmoid guarantees (0, 1))
# ---------------------------------------------------------------------------

def test_predict_single_returns_float_in_unit_interval(model: LagPredictor):
    """predict_single must return a Python float strictly in (0.0, 1.0)."""
    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)

    result = model.predict_single(x)

    assert isinstance(result, float), f"Expected float, got {type(result)}"
    assert 0.0 < result < 1.0, (
        f"Sigmoid output {result} out of (0, 1) — model architecture broken"
    )


def test_forward_output_in_unit_interval_on_random_inputs(model: LagPredictor):
    """Sigmoid output must be in (0, 1) for any random input (100 random trials)."""
    for _ in range(100):
        x = torch.rand(1, INPUT_DIM)
        out = model(x)
        val = float(out.item())
        assert 0.0 < val < 1.0, (
            f"Output {val} outside (0, 1) for random input — Sigmoid not applied"
        )


# ---------------------------------------------------------------------------
# Test 4 — MSE decreases over 10 gradient batches (model learns)
# ---------------------------------------------------------------------------

def test_mse_decreases_over_10_batches():
    """
    Train the model on a simple linear target and verify that the average
    loss over the last 10 gradient steps is lower than over the first 10.

    Synthetic target: next_lag = kafka_lag * 1.1 + 0.02 (clamped to [0, 1]).
    We run 50 gradient steps total and compare first-10-avg vs last-10-avg to
    avoid false failures from per-batch noise with a fresh random model.
    torch.manual_seed is set so weight init is deterministic regardless of
    where in the test suite this test runs.
    """
    torch.manual_seed(42)
    random.seed(42)
    model = LagPredictor()

    # ── Seed the replay buffer with 400 synthetic transitions ───────────────
    for _ in range(400):
        kafka_lag_norm = random.uniform(0.0, 0.8)
        obs = _make_obs_normalized(kafka_lag=kafka_lag_norm)
        action = _make_action()
        x = build_input_vector(obs, action)
        target = min(1.0, kafka_lag_norm * 1.1 + 0.02)
        model.store_transition(x, target)

    # ── Collect losses over 50 gradient steps ────────────────────────────────
    losses: list[float] = []
    for _ in range(50):
        loss = model.train_step()
        if loss is not None:
            losses.append(loss)

    assert len(losses) >= 20, (
        f"Expected at least 20 loss samples, got {len(losses)} — buffer too small?"
    )

    first_avg = sum(losses[:10]) / 10
    last_avg = sum(losses[-10:]) / 10

    assert last_avg < first_avg, (
        f"MSE did not decrease: first-10-avg={first_avg:.6f} last-10-avg={last_avg:.6f}. "
        "Model may not be learning — check optimizer or loss function."
    )


# ---------------------------------------------------------------------------
# Test 5 — store_transition and buffer_size
# ---------------------------------------------------------------------------

def test_store_transition_increments_buffer(model: LagPredictor):
    """Each store_transition call must increment buffer_size by 1."""
    assert model.buffer_size() == 0

    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)

    model.store_transition(x, 0.15)
    assert model.buffer_size() == 1

    model.store_transition(x, 0.20)
    assert model.buffer_size() == 2


def test_buffer_capacity_evicts_old_transitions(model: LagPredictor):
    """Buffer must not exceed REPLAY_CAPACITY (deque maxlen eviction)."""
    from dynamics_model import REPLAY_CAPACITY

    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)

    for _ in range(REPLAY_CAPACITY + 50):
        model.store_transition(x, 0.1)

    assert model.buffer_size() == REPLAY_CAPACITY, (
        f"Buffer exceeded capacity: {model.buffer_size()} > {REPLAY_CAPACITY}"
    )


# ---------------------------------------------------------------------------
# Test 6 — train_step returns None below BATCH_SIZE
# ---------------------------------------------------------------------------

def test_train_step_returns_none_below_batch_size(model: LagPredictor):
    """train_step() must return None when buffer has fewer than BATCH_SIZE items."""
    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)

    for _ in range(BATCH_SIZE - 1):
        model.store_transition(x, 0.1)

    result = model.train_step()
    assert result is None, (
        f"Expected None with {BATCH_SIZE - 1} transitions, got {result}"
    )


def test_train_step_returns_float_at_batch_size(model: LagPredictor):
    """train_step() must return a non-negative float once buffer ≥ BATCH_SIZE."""
    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)

    for _ in range(BATCH_SIZE):
        model.store_transition(x, 0.15)

    result = model.train_step()
    assert result is not None, "train_step() returned None with full batch"
    assert isinstance(result, float), f"Expected float, got {type(result)}"
    assert result >= 0.0, f"MSE loss must be non-negative, got {result}"


# ---------------------------------------------------------------------------
# MultiObsPredictor tests (Fix 10.1 — full observation world model)
# ---------------------------------------------------------------------------

@pytest.fixture
def multi_model() -> MultiObsPredictor:
    """Fresh MultiObsPredictor with empty replay buffer."""
    return MultiObsPredictor()


def test_multi_obs_forward_shape(multi_model: MultiObsPredictor) -> None:
    """Forward pass on (1, 16) batch must return shape (1, 10)."""
    x = torch.rand(1, INPUT_DIM)
    out = multi_model(x)
    assert out.shape == (1, MULTI_OBS_OUTPUT_DIM), (
        f"Expected (1, {MULTI_OBS_OUTPUT_DIM}), got {out.shape}"
    )


def test_multi_obs_output_in_unit_interval(multi_model: MultiObsPredictor) -> None:
    """Sigmoid output must guarantee all 10 values in (0, 1) for any input."""
    for _ in range(50):
        x = torch.rand(1, INPUT_DIM)
        out = multi_model(x)
        assert float(out.min()) > 0.0, "Sigmoid output below 0"
        assert float(out.max()) < 1.0, "Sigmoid output above 1"


def test_build_full_obs_target_vector_shape_and_range() -> None:
    """build_full_obs_target_vector must return a (10,) float32 tensor with values in [0,1]."""
    obs_norm = _make_obs_normalized()
    target = build_full_obs_target_vector(obs_norm)

    assert isinstance(target, torch.Tensor), "Must return torch.Tensor"
    assert target.shape == (MULTI_OBS_OUTPUT_DIM,), (
        f"Expected ({MULTI_OBS_OUTPUT_DIM},), got {target.shape}"
    )
    assert target.dtype == torch.float32, f"Expected float32, got {target.dtype}"
    assert float(target.min()) >= 0.0
    assert float(target.max()) <= 1.0


def test_multi_obs_predict_single_returns_dict(multi_model: MultiObsPredictor) -> None:
    """predict_single must return a dict with exactly 10 keys, all values in (0, 1)."""
    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)
    result = multi_model.predict_single(x)

    assert isinstance(result, dict), f"Expected dict, got {type(result)}"
    assert len(result) == MULTI_OBS_OUTPUT_DIM, (
        f"Expected {MULTI_OBS_OUTPUT_DIM} keys, got {len(result)}"
    )
    expected_keys = {
        "transaction_type", "risk_score", "adversary_threat_level",
        "system_entropy", "kafka_lag", "api_latency", "rolling_p99",
        "db_connection_pool", "bank_api_status", "merchant_tier",
    }
    assert set(result.keys()) == expected_keys
    for k, v in result.items():
        assert 0.0 < v < 1.0, f"predict_single['{k}'] = {v} outside (0, 1)"


def test_multi_obs_store_and_train_step(multi_model: MultiObsPredictor) -> None:
    """store_transition + train_step: buffer grows, loss returned at MULTI_OBS_BATCH_SIZE."""
    obs = _make_obs_normalized()
    action = _make_action()
    x = build_input_vector(obs, action)
    target = build_full_obs_target_vector(obs)

    assert multi_model.buffer_size() == 0
    assert multi_model.train_step() is None, "Should return None below batch size"

    for _ in range(MULTI_OBS_BATCH_SIZE):
        multi_model.store_transition(x, target)

    assert multi_model.buffer_size() == MULTI_OBS_BATCH_SIZE
    loss = multi_model.train_step()
    assert loss is not None, "train_step() must return float at full batch"
    assert isinstance(loss, float)
    assert loss >= 0.0, f"Weighted MSE loss must be non-negative, got {loss}"


def test_multi_obs_weighted_mse_loss_shape(multi_model: MultiObsPredictor) -> None:
    """weighted_mse_loss must return a scalar tensor."""
    pred = torch.rand(8, MULTI_OBS_OUTPUT_DIM)
    target = torch.rand(8, MULTI_OBS_OUTPUT_DIM)
    loss = multi_model.weighted_mse_loss(pred, target)
    assert loss.shape == torch.Size([]), f"Expected scalar, got shape {loss.shape}"
    assert float(loss.item()) >= 0.0


def test_multi_obs_mse_decreases_over_training() -> None:
    """MultiObsPredictor loss must trend down over 50 gradient steps."""
    torch.manual_seed(99)
    random.seed(99)
    model = MultiObsPredictor()

    obs = _make_obs_normalized(kafka_lag=0.3)
    action = _make_action()
    x = build_input_vector(obs, action)
    target = build_full_obs_target_vector(obs)

    for _ in range(400):
        model.store_transition(x, target)

    losses: list[float] = [l for _ in range(50) if (l := model.train_step()) is not None]
    assert len(losses) >= 20
    assert sum(losses[-10:]) / 10 < sum(losses[:10]) / 10, (
        "MultiObsPredictor MSE did not decrease — check weighted_mse_loss or optimizer"
    )