import torch

from pi05mini.flow import FlowMatching


def _expand_batch_time(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    """Expand a per-example time tensor over non-batch dimensions."""
    return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))


def test_training_batch_obeys_linear_flow_equations() -> None:
    flow = FlowMatching()
    target = torch.tensor(
        [
            [[1.0, -2.0], [0.5, 3.0]],
            [[-1.0, 0.0], [2.5, -0.5]],
        ]
    )
    generator = torch.Generator().manual_seed(123)

    x_t, t, target_velocity = flow.training_batch(target, generator=generator)

    assert x_t.shape == target.shape
    assert target_velocity.shape == target.shape
    assert t.shape[0] == target.shape[0]
    assert torch.all((t >= 0.0) & (t <= 1.0))

    t_view = _expand_batch_time(t, target)
    # v = target - noise, so the sampled noise is exactly recoverable.
    noise = target - target_velocity
    expected_x_t = (1.0 - t_view) * noise + t_view * target
    torch.testing.assert_close(x_t, expected_x_t)


def test_training_batch_is_reproducible_with_generator() -> None:
    flow = FlowMatching()
    target = torch.randn(3, 5, 2)

    batch_a = flow.training_batch(
        target, generator=torch.Generator().manual_seed(99)
    )
    batch_b = flow.training_batch(
        target, generator=torch.Generator().manual_seed(99)
    )

    for tensor_a, tensor_b in zip(batch_a, batch_b):
        torch.testing.assert_close(tensor_a, tensor_b)


def test_loss_is_mean_squared_velocity_error() -> None:
    flow = FlowMatching()
    target_velocity = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
    prediction = target_velocity + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])

    # Squared errors are [1, 1, 4, 0], whose mean is 1.5.
    torch.testing.assert_close(flow.loss(prediction, target_velocity), torch.tensor(1.5))
    torch.testing.assert_close(
        flow.loss(target_velocity, target_velocity), torch.tensor(0.0)
    )


def test_euler_sampler_reaches_fixed_target_for_exact_vector_field() -> None:
    flow = FlowMatching()
    target = torch.tensor(
        [
            [[0.25, -0.5], [1.0, 2.0]],
            [[-2.0, 0.75], [0.0, 0.5]],
        ]
    )

    def exact_model(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        t_view = _expand_batch_time(t, x_t)
        # Straight conditional path to target. The sampler evaluates t < 1.
        return (target.to(x_t) - x_t) / (1.0 - t_view)

    torch.manual_seed(5)
    sampled = flow.sample(
        exact_model,
        shape=tuple(target.shape),
        device=target.device,
        steps=8,
    )

    torch.testing.assert_close(sampled, target, atol=1e-5, rtol=1e-5)