import torch

from stimulus_synthesis.scoring.robust_transform import RobustTransformScorer, RobustTransformSpec, apply_robust_transform


class MeanScorer:
    def score(self, videos, target, **kwargs):
        return videos.mean(dim=(1, 2, 3, 4)).tolist()


def test_robust_transform_expands_draws_and_is_reproducible():
    videos = torch.linspace(0, 1, steps=2 * 3 * 8 * 8).reshape(1, 2, 3, 8, 8)
    spec = RobustTransformSpec(num_draws=4, crop_scale=0.8, gaussian_sigma=0.1)

    first = apply_robust_transform(videos, spec)
    second = apply_robust_transform(videos, spec)

    assert first.shape == (4, 2, 3, 8, 8)
    assert torch.equal(first, second)


def test_robust_transform_scorer_is_independent_of_batch_order():
    a = torch.zeros(2, 3, 8, 8)
    b = torch.ones(2, 3, 8, 8) * 0.5
    videos = torch.stack([a, b], dim=0)
    reversed_videos = torch.stack([b, a], dim=0)
    scorer = RobustTransformScorer(MeanScorer(), RobustTransformSpec(num_draws=4, crop_scale=0.8, gaussian_sigma=0.1))

    scores = scorer.score(videos, None)
    reversed_scores = scorer.score(reversed_videos, None)

    assert torch.allclose(torch.tensor(scores), torch.tensor(list(reversed(reversed_scores))))


def test_pipeline_default_score_transform_is_clean(monkeypatch):
    """Default scoring matches the canonical clean single pass (no robust augmentation)."""
    from stimulus_synthesis.pipeline import NevoPipeline
    from stimulus_synthesis.scoring.robust_transform import RobustTransformScorer

    class DummyEncoderScorer:
        def __init__(self, *args, **kwargs):
            pass

        def score(self, videos, target, **kwargs):
            return [0.0] * videos.shape[0]

    monkeypatch.setattr("stimulus_synthesis.pipeline.EncoderScorer", DummyEncoderScorer)
    pipe = NevoPipeline(text_to_image=object(), image_to_video=object())
    pipe._ensure_components(device="cpu")

    # Robust augmentation is disabled by default -> the encoder scorer is used directly.
    assert not isinstance(pipe.scorer, RobustTransformScorer)
    assert isinstance(pipe.scorer, DummyEncoderScorer)