"""Smoke test for SpatialAttribution — uses a tiny random model to avoid
loading the real checkpoint (too slow for unit tests)."""

import numpy as np
import torch
import torch.nn as nn
import pytest
from badas.analysis.attribution import SpatialAttribution, _H_PATCHES, _W_PATCHES


# ---------------------------------------------------------------------------
# Minimal stub that matches the interface SpatialAttribution expects
# ---------------------------------------------------------------------------

class _FakeLayer(nn.Module):
    def __init__(self, d=64):
        super().__init__()
        self.linear = nn.Linear(d, d)
    def forward(self, x):
        return self.linear(x)


class _FakeEncoder(nn.Module):
    def __init__(self, n_layers=24, d=64):
        super().__init__()
        self.layer = nn.ModuleList([_FakeLayer(d) for _ in range(n_layers)])


class _FakeBackbone(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = _FakeEncoder()


class _FakeModelConfig:
    """Minimal config stub matching model.config.model interface."""
    use_future_prediction = False


class _FakeConfig:
    model = _FakeModelConfig()


class _FakeModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = _FakeBackbone()
        self._linear = nn.Linear(2048 * 64, 1)
        self.config = _FakeConfig()

    def forward(self, x):
        B = x.shape[0]
        # Use input to seed token tensor so gradients flow from input through layers
        tok = x.reshape(B, -1)[:, :2048 * 64].reshape(B, 2048, 64)
        for layer in self.backbone.encoder.layer:
            tok = layer(tok)
        return self._linear(tok.reshape(B, -1)).squeeze(-1)


def _make_attr(block_indices=(18, 21), n_smooth=1):
    model = _FakeModel().eval().float()
    return SpatialAttribution(model, block_indices=block_indices, n_smooth=n_smooth)


def test_compute_returns_correct_shape():
    attr = _make_attr()
    x = torch.randn(1, 16, 3, 256, 256)
    result = attr.compute(x)
    assert result.shape == (_H_PATCHES, _W_PATCHES), result.shape


def test_compute_output_in_unit_range():
    attr = _make_attr()
    x = torch.randn(1, 16, 3, 256, 256)
    result = attr.compute(x)
    assert result.min() >= 0.0
    assert result.max() <= 1.0 + 1e-6


def test_get_bbox_returns_none_for_flat_map():
    attr = _make_attr()
    flat = np.full((_H_PATCHES, _W_PATCHES), 0.1)
    assert attr.get_bbox(flat) is None


def test_get_bbox_returns_tuple_for_peaked_map():
    attr = _make_attr()
    peaked = np.zeros((_H_PATCHES, _W_PATCHES))
    peaked[6:10, 8:12] = 1.0
    bbox = attr.get_bbox(peaked)
    assert bbox is not None
    x1, y1, x2, y2 = bbox
    assert x1 < x2 and y1 < y2
    assert x1 >= 0 and y1 >= 0


def test_annotate_frame_returns_copy_not_inplace():
    attr = _make_attr()
    frame = np.zeros((256, 256, 3), dtype=np.uint8)
    out = attr.annotate_frame(frame, (10, 10, 100, 100), risk=0.8)
    assert out.shape == frame.shape
    assert out is not frame


def test_annotate_frame_none_bbox():
    attr = _make_attr()
    frame = np.zeros((256, 256, 3), dtype=np.uint8)
    out = attr.annotate_frame(frame, None, risk=0.1)
    assert out.shape == frame.shape


def test_compute_fast_returns_map_and_risk():
    attr = _make_attr()
    x = torch.randn(1, 16, 3, 256, 256)
    attr_map, risk = attr.compute_fast(x)
    assert attr_map.shape == (_H_PATCHES, _W_PATCHES)
    assert attr_map.min() >= 0.0
    assert attr_map.max() <= 1.0 + 1e-6
    assert 0.0 <= risk <= 1.0
