"""Unit tests for CompiledDeepAttentionForward (fast attribution compiled path).

Uses mock modules — no real model loading needed.
"""
import numpy as np
import pytest
import torch
import torch.nn as nn
from unittest.mock import MagicMock


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_N_TOK = 2048
_N_HEADS = 16
_D = 1024


class _MockEmbeddings(nn.Module):
    """Returns all-zero hidden states — mimics VJEPA2Embeddings."""
    def forward(self, pixel_values_videos):
        B = pixel_values_videos.shape[0]
        return torch.zeros(B, _N_TOK, _D)


class _MockBlock(nn.Module):
    """
    Fake VJEPA2Layer.

    block(hidden)                       → (hidden,)   [output_attentions=False]
    block(hidden, output_attentions=True) → (hidden, attn)
    """
    def forward(self, hidden: torch.Tensor, output_attentions: bool = False, **kwargs):
        if output_attentions:
            attn = torch.rand(1, _N_HEADS, _N_TOK, _N_TOK,
                              dtype=hidden.dtype, device=hidden.device)
            return (hidden, attn)
        return (hidden,)


@pytest.fixture
def embeddings():
    return _MockEmbeddings()


@pytest.fixture
def pixel_values():
    """Dummy (1, 16, 3, 256, 256) input tensor."""
    return torch.zeros(1, 16, 3, 256, 256)


# ---------------------------------------------------------------------------
# CompiledDeepAttentionForward
# ---------------------------------------------------------------------------

def test_output_shape(embeddings, pixel_values):
    """Output is (n_late_blocks, n_heads, 256, 256)."""
    from badas.analysis.attribution import CompiledDeepAttentionForward
    early = [_MockBlock() for _ in range(2)]
    late  = [_MockBlock() for _ in range(3)]
    fwd = CompiledDeepAttentionForward(embeddings, early, late)
    out = fwd(pixel_values)
    assert out.shape == (3, _N_HEADS, 256, 256), f"unexpected shape {out.shape}"


def test_early_blocks_called_without_attention(embeddings, pixel_values):
    """Early blocks are called without output_attentions=True."""
    from badas.analysis.attribution import CompiledDeepAttentionForward

    called_with_attention = []

    class _TrackingBlock(_MockBlock):
        def forward(self, hidden, output_attentions=False, **kw):
            called_with_attention.append(output_attentions)
            return super().forward(hidden, output_attentions, **kw)

    early = [_TrackingBlock() for _ in range(2)]
    late  = [_MockBlock() for _ in range(2)]
    fwd = CompiledDeepAttentionForward(embeddings, early, late)
    called_with_attention.clear()
    fwd(pixel_values)
    # First 2 calls are early blocks
    assert called_with_attention[:2] == [False, False]


def test_late_blocks_called_with_attention(embeddings, pixel_values):
    """Late blocks are called with output_attentions=True."""
    from badas.analysis.attribution import CompiledDeepAttentionForward

    called_with_attention = []

    class _TrackingBlock(_MockBlock):
        def forward(self, hidden, output_attentions=False, **kw):
            called_with_attention.append(output_attentions)
            return super().forward(hidden, output_attentions, **kw)

    early = [_MockBlock() for _ in range(2)]
    late  = [_TrackingBlock() for _ in range(2)]
    fwd = CompiledDeepAttentionForward(embeddings, early, late)
    called_with_attention.clear()
    fwd(pixel_values)
    assert called_with_attention == [True, True]


def test_slices_to_last_t_region(embeddings, pixel_values):
    """Returned attention has shape [..., 256, 256] (last-T slice)."""
    from badas.analysis.attribution import (
        CompiledDeepAttentionForward, _LAST_T_START, _LAST_T_END,
    )
    late = [_MockBlock() for _ in range(1)]
    fwd = CompiledDeepAttentionForward(embeddings, [], late)
    out = fwd(pixel_values)
    expected_len = _LAST_T_END - _LAST_T_START  # 256
    assert out.shape[-2:] == (expected_len, expected_len)


def test_empty_late_blocks_raises(embeddings):
    """Empty late_blocks raises ValueError at construction."""
    from badas.analysis.attribution import CompiledDeepAttentionForward
    with pytest.raises(ValueError, match="at least one late block"):
        CompiledDeepAttentionForward(embeddings, [_MockBlock()], [])


def test_no_early_blocks(embeddings, pixel_values):
    """Works with empty early_blocks (block_range starts at layer 0)."""
    from badas.analysis.attribution import CompiledDeepAttentionForward
    late = [_MockBlock() for _ in range(2)]
    fwd = CompiledDeepAttentionForward(embeddings, [], late)
    out = fwd(pixel_values)
    assert out.shape == (2, _N_HEADS, 256, 256)


# ---------------------------------------------------------------------------
# SpatialAttribution compiled path
# ---------------------------------------------------------------------------

def _make_mock_model(n_blocks=24):
    model = MagicMock()
    model.backbone.encoder.layer = nn.ModuleList([_MockBlock() for _ in range(n_blocks)])
    model.backbone.encoder.embeddings = _MockEmbeddings()
    model.parameters.return_value = iter([torch.zeros(1, dtype=torch.float16)])
    return model


def test_compiled_fwd_created_when_block_range_set():
    """_compiled_fwd is set when block_range is provided at init."""
    from badas.analysis.attribution import SpatialAttribution, CompiledDeepAttentionForward
    model = _make_mock_model()
    attr = SpatialAttribution(model, block_indices=(18, 21), block_range=(20, 24))
    assert attr._compiled_fwd is not None


def test_compiled_fwd_none_without_block_range():
    """_compiled_fwd stays None when block_range is None."""
    from badas.analysis.attribution import SpatialAttribution
    model = _make_mock_model()
    attr = SpatialAttribution(model, block_indices=(18, 21))
    assert attr._compiled_fwd is None


def test_compute_compiled_attention_returns_correct_shape():
    """compute_compiled_attention returns (16, 16) ndarray."""
    from badas.analysis.attribution import SpatialAttribution
    model = _make_mock_model()
    attr = SpatialAttribution(model, block_indices=(18, 21), block_range=(20, 24))
    dummy = torch.zeros(1, 16, 3, 256, 256, dtype=torch.float16)
    result = attr.compute_compiled_attention(dummy)
    assert isinstance(result, np.ndarray)
    assert result.shape == (16, 16)
    assert result.dtype == np.float32


def test_compute_compiled_attention_values_in_range():
    """compute_compiled_attention output is in [0, 1]."""
    from badas.analysis.attribution import SpatialAttribution
    model = _make_mock_model()
    attr = SpatialAttribution(model, block_indices=(18, 21), block_range=(20, 24))
    dummy = torch.zeros(1, 16, 3, 256, 256, dtype=torch.float16)
    result = attr.compute_compiled_attention(dummy)
    assert result.min() >= 0.0
    assert result.max() <= 1.0


def test_compute_compiled_attention_raises_without_block_range():
    """compute_compiled_attention raises RuntimeError if _compiled_fwd is None."""
    from badas.analysis.attribution import SpatialAttribution
    model = _make_mock_model()
    attr = SpatialAttribution(model, block_indices=(18, 21))
    dummy = torch.zeros(1, 16, 3, 256, 256, dtype=torch.float16)
    with pytest.raises(RuntimeError, match="block_range"):
        attr.compute_compiled_attention(dummy)


# ---------------------------------------------------------------------------
# _run_attribution dispatch
# ---------------------------------------------------------------------------

def _make_badas_ns():
    """Return a minimal BADAS-like namespace with mocked attribution."""
    from types import SimpleNamespace
    from badas.inference.inference import BADAS

    attr = MagicMock()
    attr.compute_compiled_attention.return_value = np.zeros((16, 16))
    attr.compute_deep_attention.return_value = np.zeros((16, 16))
    attr.get_bbox.return_value = None

    ns = SimpleNamespace(
        _attribution=attr,
        _use_fast_path=False,
        _hazard_descriptor=None,
    )
    ns._run_attribution = BADAS._run_attribution.__get__(ns)
    return ns


def test_run_attribution_uses_compiled_path_when_enabled():
    """_run_attribution calls compute_compiled_attention when _use_fast_path=True."""
    ns = _make_badas_ns()
    ns._use_fast_path = True
    tensor = torch.zeros(1, 16, 3, 256, 256)
    ns._run_attribution(tensor)
    ns._attribution.compute_compiled_attention.assert_called_once_with(tensor)
    ns._attribution.compute_deep_attention.assert_not_called()


def test_run_attribution_uses_slow_path_when_disabled():
    """_run_attribution calls compute_deep_attention when _use_fast_path=False."""
    ns = _make_badas_ns()
    ns._use_fast_path = False
    tensor = torch.zeros(1, 16, 3, 256, 256)
    ns._run_attribution(tensor)
    ns._attribution.compute_deep_attention.assert_called_once_with(tensor)
    ns._attribution.compute_compiled_attention.assert_not_called()


def test_run_attribution_passes_probability_to_hazard_descriptor():
    """_run_attribution passes actual probability (not threshold) to describe()."""
    from types import SimpleNamespace
    from badas.inference.inference import BADAS

    attr = MagicMock()
    attr.compute_deep_attention.return_value = np.zeros((16, 16))
    attr.get_bbox.return_value = (10, 10, 100, 100)  # non-None so describe() fires

    hd = MagicMock()
    hd.describe.return_value = "A truck ahead — brake!"

    ns = SimpleNamespace(
        _attribution=attr,
        _use_fast_path=False,
        _hazard_descriptor=hd,
        img_size=256,
    )
    ns._run_attribution = BADAS._run_attribution.__get__(ns)

    frame = np.zeros((720, 1280, 3), dtype=np.uint8)
    ns._run_attribution(torch.zeros(1, 16, 3, 256, 256), frame_bgr=frame, probability=0.87)

    hd.describe.assert_called_once()
    _, kwargs = hd.describe.call_args
    assert kwargs['risk'] == 0.87, f"Expected risk=0.87, got {kwargs['risk']}"
    assert kwargs['img_size'] == 256


def test_run_attribution_hazard_descriptor_skipped_without_frame():
    """_run_attribution does not call describe() when frame_bgr is None."""
    from types import SimpleNamespace
    from badas.inference.inference import BADAS

    attr = MagicMock()
    attr.compute_deep_attention.return_value = np.zeros((16, 16))
    attr.get_bbox.return_value = (10, 10, 100, 100)

    hd = MagicMock()
    ns = SimpleNamespace(
        _attribution=attr,
        _use_fast_path=False,
        _hazard_descriptor=hd,
        img_size=256,
    )
    ns._run_attribution = BADAS._run_attribution.__get__(ns)

    ns._run_attribution(torch.zeros(1, 16, 3, 256, 256), frame_bgr=None, probability=0.9)
    hd.describe.assert_not_called()


# ---------------------------------------------------------------------------
# _check_attribution_trigger
# ---------------------------------------------------------------------------

def _make_trigger_ns(threshold=0.5):
    """Minimal namespace for _check_attribution_trigger."""
    from types import SimpleNamespace
    from badas.inference.inference import BADAS
    ns = SimpleNamespace(_attr_threshold=threshold, _attr_triggered=False)
    ns._check_attribution_trigger = BADAS._check_attribution_trigger.__get__(ns)
    return ns


def test_trigger_fires_on_first_crossing():
    """Returns True exactly once when risk first crosses the threshold."""
    ns = _make_trigger_ns(threshold=0.5)
    assert ns._check_attribution_trigger(0.3) is False   # below
    assert ns._check_attribution_trigger(0.8) is True    # first crossing
    assert ns._check_attribution_trigger(0.9) is False   # already triggered


def test_trigger_resets_when_risk_drops():
    """After risk drops below threshold, the next crossing fires again."""
    ns = _make_trigger_ns(threshold=0.5)
    assert ns._check_attribution_trigger(0.8) is True    # fires
    assert ns._check_attribution_trigger(0.9) is False   # suppressed
    assert ns._check_attribution_trigger(0.3) is False   # drops — resets
    assert ns._check_attribution_trigger(0.7) is True    # fires again
