import torch

from pi0mini.flow import FlowMatching


def _time_view(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
    return t.reshape(t.shape[0], *([1] * (target.ndim - 1)))


def test_training_batch_obeys_linear_path() -> None:
    flow = FlowMatching()
    action = torch.tensor([[[1.0, -2.0], [0.5, 3.0]], [[-1.0, 0.0], [2.5, -0.5]]])
    x_t, t, velocity = flow.training_batch(
        action, generator=torch.Generator().manual_seed(123)
    )

    noise = action - velocity
    expected = (1.0 - _time_view(t, action)) * noise + _time_view(t, action) * action
    torch.testing.assert_close(x_t, expected)
    assert x_t.shape == action.shape
    assert torch.all((t >= 0.0) & (t <= 1.0))


def test_training_batch_is_reproducible() -> None:
    action = torch.randn(3, 16, 2)
    flow = FlowMatching()
    a = flow.training_batch(action, generator=torch.Generator().manual_seed(9))
    b = flow.training_batch(action, generator=torch.Generator().manual_seed(9))
    for left, right in zip(a, b):
        torch.testing.assert_close(left, right)


def test_training_times_use_low_biased_shifted_beta() -> None:
    flow = FlowMatching()
    times = flow.sample_training_times(
        100_000,
        device="cpu",
        dtype=torch.float32,
        generator=torch.Generator().manual_seed(17),
    )

    assert torch.all(times >= 0.0)
    assert torch.all(times <= 0.999)
    # E[0.999 * (1 - Beta(1.5, 1))] = 0.3996.  This also guards
    # against accidentally reverting to a uniform or high-time-biased draw.
    assert abs(float(times.mean()) - 0.3996) < 0.005


def test_loss_is_velocity_mse() -> None:
    flow = FlowMatching()
    target = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]])
    prediction = target + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]])
    torch.testing.assert_close(flow.loss(prediction, target), torch.tensor(1.5))


def test_exact_vector_field_integrates_to_action() -> None:
    flow = FlowMatching(inference_steps=10)
    action = torch.tensor([[[0.25, -0.5], [1.0, 2.0]], [[-2.0, 0.75], [0.0, 0.5]]])

    def exact_field(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        return (action.to(x_t) - x_t) / (1.0 - _time_view(t, x_t))

    sample = flow.sample(exact_field, tuple(action.shape), action.device,
                         generator=torch.Generator().manual_seed(5))
    torch.testing.assert_close(sample, action, atol=1e-5, rtol=1e-5)